[Gluon] Infer memdesc_trans layout automatically (#7121)

Author: peterbell10

python/src/gluon_ir.cc

@@ -264,6 +264,12 @@ void init_gluon_ir(py::module &&m) {
              assert(ty.getEncoding());
              return layoutToGluon(ty.getEncoding());
            })
+      .def("get_gluon_layout_from_memdesc",
+           [](GluonOpBuilder &self, Value memdesc) -> py::object {
+             auto ty = dyn_cast<ttg::MemDescType>(memdesc.getType());
+             assert(ty.getEncoding());
+             return layoutToGluon(ty.getEncoding());
+           })
       .def("create_convert_layout",
            [](GluonOpBuilder &self, Type resultTy, Value value) -> Value {
              return self.create<ttg::ConvertLayoutOp>(resultTy, value);
@@ -296,9 +302,9 @@ void init_gluon_ir(py::module &&m) {
                                                        offsets);
            })
       .def("create_memdesc_trans",
-           [](GluonOpBuilder &self, Type resultType, Value src,
+           [](GluonOpBuilder &self, Value src,
               std::vector<int> &order) -> Value {
-             return self.create<ttg::MemDescTransOp>(resultType, src, order);
+             return self.create<ttg::MemDescTransOp>(src, order);
            })
       .def("create_memdesc_reshape",
            [](GluonOpBuilder &self, Type resultType, Value src) -> Value {

python/test/gluon/test_frontend.py

@@ -247,9 +247,11 @@ def shared_memory_cast_kernel():
     layout_a: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=64, transposed=False, element_bitwidth=8,
                                                       rank=2)
     layout_T: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=64, transposed=True, element_bitwidth=8,
-                                                      rank=2)
+                                                      rank=2, ctas_per_cga=[1, 1], cta_split_num=[1,
+                                                                                                  1], cta_order=[1, 0])
     smem = ttgl.allocate_shared_memory(ttgl.int8, [2, 256, 128], layout_a)
-    smem.subslice(0).permute((1, 0), layout_T)
+    perm = smem.subslice(0).permute((1, 0))
+    ttgl.static_assert(perm.type.layout == layout_T)
 
     layout_b: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=64, transposed=False, element_bitwidth=16,
                                                       rank=4, cta_order=[3, 2, 1, 0])

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

@@ -239,11 +239,9 @@ def subslice(self, index, shape=None, layout=None, _semantic: GluonSemantic = No
         return _semantic.memdesc_slice(self, index, shape, layout)
 
     @builtin
-    def permute(self, order, layout, _semantic: GluonSemantic) -> shared_memory_descriptor:
+    def permute(self, order, _semantic: GluonSemantic) -> shared_memory_descriptor:
         order = [_unwrap_if_constexpr(o) for o in order]
-        layout = _unwrap_if_constexpr(layout)
-
-        return _semantic.memdesc_trans(self, order, layout)
+        return _semantic.memdesc_trans(self, order)
 
     @builtin
     def reshape(self, shape, layout, _semantic: GluonSemantic) -> shared_memory_descriptor:

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

@@ -163,7 +163,7 @@ def memdesc_slice(self, mem_desc, index, shape, layout):
         offsets[0] = self._convert_elem_to_ir_value(index, require_i64=False)
         return self._memdesc_subview(mem_desc, offsets, shape, layout)
 
-    def memdesc_trans(self, mem_desc, order, layout):
+    def memdesc_trans(self, mem_desc, order):
         assert len(order) == len(
             mem_desc.shape), f"source rank ({mem_desc.rank}) and order length ({len(order)}) must match"
 
@@ -172,9 +172,10 @@ def memdesc_trans(self, mem_desc, order, layout):
         new_alloc_shape = alloc_shape[:len(alloc_shape) - mem_desc.rank]
         new_alloc_shape += [alloc_shape[len(alloc_shape) - mem_desc.rank:][i] for i in order]
 
-        ty = ttgl.shared_memory_descriptor_type(mem_desc.dtype, shape, layout, new_alloc_shape)
-        handle = self.builder.create_memdesc_trans(ty.to_ir(self.builder), mem_desc.handle, order)
-        return ttgl.shared_memory_descriptor(handle, **ty.__dict__)
+        handle = self.builder.create_memdesc_trans(mem_desc.handle, order)
+        layout = self.builder.get_gluon_layout_from_memdesc(handle)
+        return ttgl.shared_memory_descriptor(handle, element_ty=mem_desc.dtype, shape=shape, alloc_shape=alloc_shape,
+                                             layout=layout)
 
     def memdesc_reshape(self, mem_desc, shape, layout):
         ty = ttgl.shared_memory_descriptor_type(mem_desc.dtype, shape, layout, mem_desc.type.alloc_shape)