[Frontend][Gluon] Make static_assert actually work (#7168)

* Added `ttgl.static_assert` (and `static_print`)
* Fix binop where RHS is constexpr but LHS is not
* Fix `unwrap_if_constexpr` to recursively unwrap

---------

Co-authored-by: peterbell10 <[email protected]>

Author: Mogball

python/test/gluon/test_frontend.py

@@ -13,7 +13,7 @@
 from triton._filecheck import filecheck_test, run_parser
 import triton.language as tl
 from triton._internal_testing import is_cuda
-from triton.compiler.errors import CompilationError
+from triton.compiler.errors import CompilationError, CompileTimeAssertionFailure
 
 TARGET_PAT = re.compile('ttg.target = "[^"]*"')
 
@@ -604,10 +604,10 @@ def kernel():
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "...", "ttg.threads-per-warp" = 32 : i32} {
   tt.func public @kernel() attributes {noinline = false} {
     %0 = ttg.local_alloc : () -> !ttg.memdesc<32x32xi32, #shared, #smem, mutable>
-    tt.call @"test_frontend.smem_and_layout_user__MDi32S32_32SLSSS_1_1_1_constexpr[1]_constexpr[0]____SSSLAS[32, 32]ASMD__(1,)cconstexpr_SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=(constexpr_1_ ,constexpr_0_), ctas_per_cga=None, cta_split_num=None, cta_order=None)_"(%0) : (!ttg.memdesc<32x32xi32, #shared, #smem, mutable>) -> ()
+    tt.call @"test_frontend.smem_and_layout_user__MDi32S32_32SLSSS_1_1_1_1_0____SSSLAS[32, 32]ASMD__(1,)cconstexpr_SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=(1 ,0), ctas_per_cga=None, cta_split_num=None, cta_order=None)_"(%0) : (!ttg.memdesc<32x32xi32, #shared, #smem, mutable>) -> ()
     tt.return
   }
-  tt.func private @"test_frontend.smem_and_layout_user__MDi32S32_32SLSSS_1_1_1_constexpr[1]_constexpr[0]____SSSLAS[32, 32]ASMD__(1,)cconstexpr_SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=(constexpr_1_ ,constexpr_0_), ctas_per_cga=None, cta_split_num=None, cta_order=None)_"(%arg0: !ttg.memdesc<32x32xi32, #shared, #smem, mutable>) attributes {noinline = false} {
+  tt.func private @"test_frontend.smem_and_layout_user__MDi32S32_32SLSSS_1_1_1_1_0____SSSLAS[32, 32]ASMD__(1,)cconstexpr_SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=(1 ,0), ctas_per_cga=None, cta_split_num=None, cta_order=None)_"(%arg0: !ttg.memdesc<32x32xi32, #shared, #smem, mutable>) attributes {noinline = false} {
     tt.return
   }
 }
@@ -855,7 +855,7 @@ def test_tensor_permute():
 def test_split_join():
     # CHECK: [[BLOCKED:#.*]] = #ttg.blocked<{sizePerThread = [2], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
     # CHECK: [[BLOCKED1:#.*]] = #ttg.blocked<{sizePerThread = [2, 2], threadsPerWarp = [32, 1], warpsPerCTA = [4, 1], order = [1, 0]}>
-    layout: ttgl.constexpr = ttgl.BlockedLayout([2], [32], [4], [0])
+    layout: ttgl.constexpr = ttgl.BlockedLayout([2], [32], [4], [0], [1], [1], [0])
     a = ttgl.full([128], 1, ttgl.int32, layout)
     b = ttgl.full([128], 2, ttgl.int32, layout)
     # CHECK: tt.join {{.*}} : tensor<128xi32, [[BLOCKED]]> -> tensor<128x2xi32, [[BLOCKED1]]>
@@ -883,6 +883,16 @@ def test_tensor_reshape():
     ttgl.static_assert(v.type.layout == expect_layout)
 
 
+@gluon.jit
+def static_assert_kernel():
+    ttgl.static_assert(False)
+
+
+def test_static_assert():
+    with pytest.raises(CompileTimeAssertionFailure):
+        run_parser(static_assert_kernel)
+
+
 @filecheck_test
 @gluon.jit
 def test_zeros():

python/triton/compiler/code_generator.py

@@ -637,6 +637,8 @@ def _apply_binary_method(self, method_name, lhs, rhs):
         if _is_triton_tensor(rhs):
             reverse_method_name = re.sub(r"__(.*)__", r"__r\1__", method_name)
             return getattr(rhs, reverse_method_name)(lhs, _semantic=self.semantic)
+        if not isinstance(lhs, (constexpr, language.tuple)) and isinstance(rhs, constexpr):
+            lhs = constexpr(lhs)
         return getattr(lhs, method_name)(rhs)
 
     def visit_BinOp(self, node):
@@ -1457,9 +1459,12 @@ def ret(self, node: ast.Call):
 
         return ret
 
+    from ..experimental.gluon import language as ttgl
     statically_implemented_functions: Dict[object, Callable[[ast.Call], Any]] = {
         language.core.static_assert: execute_static_assert,
         language.core.static_print: static_executor(print),
+        ttgl.static_assert: execute_static_assert,
+        ttgl.static_print: static_executor(print),
         int: static_executor(int),
         len: static_executor(len),
     }

python/triton/compiler/compiler.py

@@ -10,8 +10,6 @@
 from ..runtime.cache import get_cache_manager, get_dump_manager, get_override_manager
 from ..runtime.driver import driver
 from ..tools.disasm import get_sass
-# TODO: this shouldn't be here
-from .code_generator import ast_to_ttir
 from pathlib import Path
 import re
 import functools
@@ -81,6 +79,7 @@ def hash(self):
         return hashlib.sha256(key.encode("utf-8")).hexdigest()
 
     def make_ir(self, options, codegen_fns, module_map, context):
+        from .code_generator import ast_to_ttir
         return ast_to_ttir(self.fn, self, context=context, options=options, codegen_fns=codegen_fns,
                            module_map=module_map)
 

python/triton/experimental/gluon/_runtime.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 import triton
-from triton.compiler.code_generator import ast_to_ttir
 from triton.compiler.compiler import ASTSource
 from triton.backends.compiler import Language
 from triton.runtime.jit import JITFunction
@@ -19,6 +18,7 @@ def __init__(self, fn, signature, constexprs=None, attrs=None) -> None:
 
     def make_ir(self, options, codegen_fns, module_map, context):
         from triton.compiler.compiler import make_backend
+        from triton.compiler.code_generator import ast_to_ttir
 
         builder = ir.builder(context)
         module = builder.create_module()

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

@@ -54,6 +54,7 @@
     "reshape",
     "split",
     "static_assert",
+    "static_print",
     "store",
     "to_tensor",
     "where",

python/triton/language/core.py

@@ -364,6 +364,12 @@ def wrapper(*args, _semantic=None, **kwargs):
 
 
 def _unwrap_if_constexpr(o):
+    if isinstance(o, list):
+        return [_unwrap_if_constexpr(x) for x in o]
+    if isinstance(o, builtins.tuple):
+        return builtins.tuple(_unwrap_if_constexpr(x) for x in o)
+    if isinstance(o, tuple):
+        return tuple(_unwrap_if_constexpr(x) for x in o)
     return o.value if isinstance(o, constexpr) else o