From 78101edb9b58ecee931ebcc8a55929fd6d54d374 Mon Sep 17 00:00:00 2001 From: Stanley Winata Date: Tue, 26 Nov 2024 17:20:11 -0800 Subject: [PATCH] Refactor to oneliner based on Harsh comments Signed-off-by: Stanley Winata --- iree/turbine/kernel/wave/hoisting.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/iree/turbine/kernel/wave/hoisting.py b/iree/turbine/kernel/wave/hoisting.py index 140e896d..a2efbee4 100644 --- a/iree/turbine/kernel/wave/hoisting.py +++ b/iree/turbine/kernel/wave/hoisting.py @@ -43,13 +43,15 @@ def get_hoistable_ops( continue # Only handle case where we are not writing to the same memory. # Counterproof: we may expect different read if we write to same memory. - for memory_user in custom_node.memory.users: - if isinstance(get_custom(memory_user), Write): - continue + if any( + isinstance(get_custom(mem_user), Write) + for mem_user in custom_node.memory.users + ): + continue # Only hoist Read that is loop invariant. - dims_indexing = [ind.start for ind in custom_node.index.values()] - dim_depends_on_ivar = [ind.has(induction_variable) for ind in dims_indexing] - if any(dim_depends_on_ivar): + if any( + ind.start.has(induction_variable) for ind in custom_node.index.values() + ): continue hoistable_ops.append(custom_node) else: