[AMD] Enable the loop unroller (#4773)

triton-lang · Sep 21, 2024 · 1e06252 · 1e06252
1 parent 35a8a00
commit 1e06252
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 6 deletions.
diff --git 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
@@ -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