From d977705ce739b70db424405df493f57dd5cf4e1f Mon Sep 17 00:00:00 2001 From: Hongtao Yu Date: Fri, 20 Sep 2024 17:20:42 -0700 Subject: [PATCH] [AMD] Enable the loop unroller. --- python/test/unit/language/test_core.py | 6 ------ third_party/amd/backend/compiler.py | 1 + 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/python/test/unit/language/test_core.py b/python/test/unit/language/test_core.py index 3040af84d5c5..00016733df3c 100644 --- a/python/test/unit/language/test_core.py +++ b/python/test/unit/language/test_core.py @@ -5635,12 +5635,6 @@ def _kernel(dst, unroll_factor: tl.constexpr): tl.atomic_add(dst + pid, i + pid) def check_loop_unroll_count(ir, opStr, loop_unroll_factor): - - # Loop unrolling pass only working on Nvidia backend at the moment - if not is_cuda(): - assert f'tt.loop_unroll_factor = {loop_unroll_factor}' in ir - return - for line in ir.splitlines(): if opStr in line: loop_unroll_factor = loop_unroll_factor - 1 diff --git a/third_party/amd/backend/compiler.py b/third_party/amd/backend/compiler.py index 705547763696..78a3c21280b3 100644 --- a/third_party/amd/backend/compiler.py +++ b/third_party/amd/backend/compiler.py @@ -142,6 +142,7 @@ def make_ttir(mod, metadata, options): passes.common.add_cse(pm) passes.common.add_licm(pm) passes.common.add_symbol_dce(pm) + passes.ttir.add_loop_unroll(pm) pm.run(mod) return mod