[Gluon] Fix layout mangling (#7020)

This fixes a few things to allow layout objects to be passed as
constexprs to other functions

Author: Mogball

python/test/gluon/test_frontend.py

@@ -539,3 +539,33 @@ def kernel():
         run_parser(kernel)
 
     assert "order must be a permutation of 0..(rank-1), but was [1]" in str(e.value.__cause__)
+
+
+@gluon.jit
+def smem_and_layout_user(smem, a: ttgl.constexpr):
+    pass
+
+
+def test_layout_mangling():
+
+    @gluon.jit
+    def kernel():
+        a: ttgl.constexpr = ttgl.SwizzledSharedLayout(1, 1, 1, [1, 0])
+        smem = ttgl.allocate_shared_memory(ttgl.int32, [32, 32], a)
+        smem_and_layout_user(smem, a)
+
+    expecttest.assert_expected_inline(
+        run_parser(kernel).str_nodebug(), """\
+#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}>
+#smem = #ttg.shared_memory
+module {
+  tt.func public @kernel() attributes {noinline = false} {
+    %0 = ttg.local_alloc : () -> !ttg.memdesc<32x32xi32, #shared, #smem, mutable>
+    tt.call @"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.return
+  }
+  tt.func private @"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.return
+  }
+}
+""")

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

@@ -167,7 +167,7 @@ def __neq__(self, other) -> bool:
         return not (self == other)
 
     def mangle(self) -> str:
-        shape_str = "_".join(self.shape)
+        shape_str = "_".join([str(s) for s in self.shape])
         return f"MD{self.element_ty.mangle()}S{shape_str}SL{self.layout.mangle()}LAS{self.alloc_shape}ASMD"
 
 
@@ -290,6 +290,7 @@ def full(shape, value, dtype, layout, _builder=None):
 def allocate_shared_memory(element_ty, shape, layout, value=None, _builder=None):
     element_ty = _unwrap_if_constexpr(element_ty)
     shape = _unwrap_if_constexpr(shape)
+    shape = [_unwrap_if_constexpr(s) for s in shape]
     layout = _unwrap_if_constexpr(layout)
     return semantic.allocate_shared(element_ty, shape, layout, value, _builder)
 

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

@@ -186,4 +186,10 @@ def _to_ir(self, builder):
         )
 
     def mangle(self) -> str:
-        return f"NVMMA_{self.swizzle_byte_width}_{self.element_bitwidth}_{self.transposed}_{self.fp4_padded}_NVMMA"
+
+        def stringify(x):
+            if x is None:
+                return ""
+            return "_".join(map(str, x))
+
+        return f"SSS_{self.vec}_{self.per_phase}_{self.max_phase}_{stringify(self.order)}_{stringify(self.ctas_per_cga)}_{stringify(self.cta_split_num)}_{stringify(self.cta_order)}_SSS"

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

@@ -39,6 +39,12 @@ def _to_ir(self, builder):
             cta_split_num,
         )
 
+    def mangle(self) -> str:
+        block_str = f"{self.block[0]}x{self.block[1]}"
+        unpacked_str = "U" if self.unpacked else "P"
+        cta_split_str = f"CS{self.cta_split_num[0]}x{self.cta_split_num[1]}" if self.cta_split_num else ""
+        return f"TL{block_str}{unpacked_str}{cta_split_str}TL"
+
 
 class tensor_memory_descriptor_type(base_type):
 
@@ -75,7 +81,7 @@ def __neq__(self, other) -> bool:
         return not (self == other)
 
     def mangle(self) -> str:
-        shape_str = "_".join(self.shape)
+        shape_str = "_".join([str(s) for s in self.shape])
         return f"MD{self.element_ty.mangle()}S{shape_str}SL{self.layout.mangle()}LAS{self.alloc_shape}ASMD"
 
 

python/triton/language/core.py

@@ -1308,6 +1308,9 @@ def _flatten_ir(self, handles: List[ir.value]):
         for v in self.values:
             v._flatten_ir(handles)
 
+    def __repr__(self):
+        return f"({' ,'.join(repr(x) for x in self.values)})"
+
 
 class slice: