[Gluon] Add support for assume (#8394)

Author: ThomasRaoux

python/test/gluon/test_frontend.py

@@ -76,6 +76,32 @@ def test_convert_layout(target):
 """)
 
 
+@gluon.jit
+def assume_kernel(arg: tl.int32):
+    ttgl.assume(arg > 1)
+
+
+@pytest.mark.parametrize("target", ALL_TARGETS)
+def test_assume(target):
+    arg = 100
+    mod = run_parser(
+        assume_kernel,
+        *make_args(arg),
+        target=target,
+    )
+    expecttest.assert_expected_inline(
+        anonymize_ir(mod.str_nodebug()), """\
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "...", "ttg.threads-per-warp" = 32 : i32} {
+  tt.func public @assume_kernel(%arg0: i32) attributes {noinline = false} {
+    %c1_i32 = arith.constant 1 : i32
+    %0 = arith.cmpi sgt, %arg0, %c1_i32 : i32
+    llvm.intr.assume %0 : i1
+    tt.return
+  }
+}
+""")
+
+
 @filecheck_test
 @gluon.jit
 def test_histogram_frontend():

python/triton/experimental/gluon/language/__init__.py

@@ -34,6 +34,7 @@
     allocate_shared_memory,
     arange,
     associative_scan,
+    assume,
     atomic_add,
     atomic_and,
     atomic_cas,

python/triton/experimental/gluon/language/_core.py

@@ -100,6 +100,7 @@ def wrapper(*args, **kwargs):
 
 # Explicitly import forwarded Triton language symbols so mypy sees them.
 associative_scan = builtin(tl_core.associative_scan)
+assume = builtin(tl_core.assume)
 atomic_add = builtin(tl_core.atomic_add)
 atomic_and = builtin(tl_core.atomic_and)
 atomic_cas = builtin(tl_core.atomic_cas)