[TLE][MTHREADS] Support transposed operands in tle.wgmma

python/triton/experimental/tle/language/gpu/core.py

@@ -813,7 +813,7 @@ def wgmma(
     trans_b = _require_wgmma_bool(trans_b, "trans_b")
     mthreads_enabled = mthreads_common.enabled()
     if mthreads_enabled:
-        mthreads_wgmma.validate_operands(a, b, acc, trans_a, trans_b)
+        a, b = mthreads_wgmma.prepare_operands(a, b, acc, trans_a, trans_b, _semantic)
     else:
         a, b = _canonicalize_wgmma_operands(a, b, trans_a, trans_b, _semantic)
 

python/triton/experimental/tle/language/gpu/mthreads/wgmma.py

@@ -20,6 +20,8 @@
 
 import triton.language as tl
 
+from .. import types as tle
+
 
 def _unwrap(value):
     return value.value if isinstance(value, tl.constexpr) else value
@@ -36,9 +38,7 @@ def mark_auto_shared_layout(builder, handle) -> None:
     builder.mark_musa_tle_auto_shared_layout(handle)
 
 
-def validate_operands(a, b, acc, trans_a: bool, trans_b: bool) -> None:
-    if trans_a or trans_b:
-        raise ValueError("initial mthreads TLE wgmma does not support trans_a/trans_b")
+def validate_operands(a, b, acc) -> None:
     for name, operand in (("a", a), ("b", b)):
         operand_type = getattr(operand, "type", None)
         if not hasattr(operand_type, "storage"):
@@ -56,6 +56,38 @@ def validate_operands(a, b, acc, trans_a: bool, trans_b: bool) -> None:
         raise ValueError("mthreads TLE wgmma requires an f32 accumulator")
 
 
+def _transpose_smem_operand(operand, semantic):
+    order = [1, 0]
+    handle = semantic.builder.create_memdesc_trans(operand.handle, order)
+    shape = [operand.type.shape[index] for index in order]
+
+    alloc_shape = operand.type.alloc_shape
+    leading_rank = len(alloc_shape) - len(operand.type.shape)
+    alloc_tail = alloc_shape[leading_rank:]
+    transposed_alloc_shape = alloc_shape[:leading_rank] + [alloc_tail[index] for index in order]
+
+    layout = operand.type.layout.make_permute(order)
+    return tle.buffered_tensor(
+        handle,
+        operand.dtype,
+        shape,
+        operand.type.storage,
+        layout,
+        semantic,
+        alloc_shape=transposed_alloc_shape,
+    )
+
+
+def prepare_operands(a, b, acc, trans_a: bool, trans_b: bool, semantic):
+    """Validate mthreads SQMMA operands and build descriptor transpose views."""
+    validate_operands(a, b, acc)
+    if trans_a:
+        a = _transpose_smem_operand(a, semantic)
+    if trans_b:
+        b = _transpose_smem_operand(b, semantic)
+    return a, b
+
+
 def validate_options(max_num_imprecise_acc: int, out_dtype) -> None:
     if max_num_imprecise_acc != 0:
         raise ValueError("mthreads TLE wgmma requires max_num_imprecise_acc=0")

third_party/mthreads/musa/include/TritonMUSACommon/MemDescUtils.h

@@ -504,6 +504,44 @@ inline Value materializeReshapedMemDescForTarget(
   return {};
 }
 
+inline void setInsertionPointAfterSameBlockDep(RewriterBase &rewriter,
+                                               Operation *anchor, Value dep) {
+  rewriter.setInsertionPoint(anchor);
+  if (Operation *def = dep.getDefiningOp())
+    if (def->getBlock() == anchor->getBlock() && anchor->isBeforeInBlock(def))
+      rewriter.setInsertionPointAfter(def);
+}
+
+inline void moveMemDescViewChainAfterDef(ArrayRef<Operation *> directUsers,
+                                         Operation *def) {
+  Block *block = def->getBlock();
+  SmallVector<Operation *> chain;
+  SmallPtrSet<Operation *, 8> seen;
+  SmallVector<Operation *> worklist(directUsers.begin(), directUsers.end());
+  while (!worklist.empty()) {
+    Operation *op = worklist.pop_back_val();
+    if (!seen.insert(op).second)
+      continue;
+    if (op->getBlock() != block)
+      continue;
+    if (!op->hasTrait<OpTrait::MemDescViewTrait>())
+      continue;
+    chain.push_back(op);
+    llvm::append_range(worklist, op->getUsers());
+  }
+  if (chain.empty())
+    return;
+  llvm::sort(chain,
+             [](Operation *a, Operation *b) { return a->isBeforeInBlock(b); });
+  Operation *anchor = def;
+  for (Operation *op : chain) {
+    if (anchor->isBeforeInBlock(op))
+      continue;
+    op->moveAfter(anchor);
+    anchor = op;
+  }
+}
+
 inline bool replaceTensorLocalAllocWithMemDesc(RewriterBase &rewriter,
                                                Operation *user,
                                                Value sourceMemDesc) {
@@ -514,13 +552,17 @@ inline bool replaceTensorLocalAllocWithMemDesc(RewriterBase &rewriter,
   if (!targetTy)
     return false;
   OpBuilder::InsertionGuard guard(rewriter);
-  rewriter.setInsertionPoint(localAlloc);
+  setInsertionPointAfterSameBlockDep(rewriter, localAlloc, sourceMemDesc);
   Value replacement =
       adaptMemDescValue(rewriter, localAlloc.getLoc(), sourceMemDesc, targetTy,
                         localAlloc.getOperation());
   if (!replacement)
     return false;
+  SmallVector<Operation *> users(localAlloc->getUsers().begin(),
+                                 localAlloc->getUsers().end());
   rewriter.replaceOp(localAlloc, replacement);
+  if (Operation *def = replacement.getDefiningOp())
+    moveMemDescViewChainAfterDef(users, def);
   return true;
 }
 
@@ -617,7 +659,7 @@ inline bool tryReplaceTensorUserWithMemDesc(RewriterBase &rewriter,
       if (!targetTy)
         continue;
       OpBuilder::InsertionGuard guard(rewriter);
-      rewriter.setInsertionPoint(localAlloc);
+      setInsertionPointAfterSameBlockDep(rewriter, localAlloc, sourceMemDesc);
       Value replacement = materializeTransformedMemDescForTarget(
           rewriter, transOp, sourceMemDesc, targetTy,
           localAlloc.getOperation());
@@ -640,7 +682,7 @@ inline bool tryReplaceTensorUserWithMemDesc(RewriterBase &rewriter,
       if (!targetTy)
         continue;
       OpBuilder::InsertionGuard guard(rewriter);
-      rewriter.setInsertionPoint(localAlloc);
+      setInsertionPointAfterSameBlockDep(rewriter, localAlloc, sourceMemDesc);
       Value replacement = materializeReshapedMemDescForTarget(
           rewriter, reshapeOp, sourceMemDesc, targetTy,
           localAlloc.getOperation());