Introduce 'getOrderForDotOperand' interface function

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td

@@ -804,6 +804,10 @@ def MmaEncodingTrait : AttrInterface<"MmaEncodingTrait"> {
       "getElemsPerThreadForOperands", (ins "ArrayRef<int64_t>":$tensorShape,
                                            "Type":$eltTy,
                                            "unsigned":$opIdx)>,
+
+    InterfaceMethod<"Return order of dimensions for dot operands.", "SmallVector<unsigned>",
+      "getOrderForDotOperand", (ins "unsigned":$opIdx,
+                                    "unsigned": $rank)>,
   ];
 }
 
@@ -926,6 +930,7 @@ V [ 0,4,8...60   1,5...61     2,6...62     3,7...63    ]   [ 128,132...188  129,
     unsigned getTotalElemsPerThreadForOperand(ArrayRef<int64_t> shape, Type eltTy, int kWidth, int opIdx) const;
     SmallVector<int64_t> getInstrShapeForOperand(int kWidth, int opIdx) const;
     SmallVector<int64_t> getRepForOperand(ArrayRef<int64_t> operandShape, int kWidth, int opIdx) const;
+    SmallVector<unsigned> getOrderForDotOperand(unsigned opIdx, unsigned rank) const;
 
     SmallVector<unsigned> getContigPerThread() {
       auto rank = getWarpsPerCTA().size();
@@ -1038,6 +1043,7 @@ Row |       warp 0                warp 2
     SmallVector<int64_t> getElemsPerInstrForOperands() const;
     SmallVector<int64_t> getRepForOperand(ArrayRef<int64_t> operandShape,
                                           Type elemType, int kWidth, int opIdx) const;
+    SmallVector<unsigned> getOrderForDotOperand(unsigned opIdx, unsigned rank) const;
     static SmallVector<unsigned> getMNKDimPerInstr();
 
     SmallVector<unsigned> getContigPerThread() {
@@ -1249,6 +1255,7 @@ For example, the matrix L corresponding to blockTileSize=[32,16] is:
     SmallVector<unsigned> getSizePerThreadForOperand(int kWidth, int opIdx) const;
     SmallVector<unsigned> getShapePerCTATileForOperand(ArrayRef<int64_t> shape, int kWidth, int opIdx) const;
     unsigned getTotalElemsPerThreadForOperand(ArrayRef<int64_t> shape, Type eltTy, int kWidth, int opIdx) const;
+    SmallVector<unsigned> getOrderForDotOperand(unsigned opIdx, unsigned rank) const;
 
     SmallVector<unsigned> getContigPerThread() {
       assert(isVolta() || isAmpere() || isHopper());

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -299,11 +299,11 @@ SmallVector<unsigned> getOrder(Attribute layout) {
   }
   if (auto dotLayout = dyn_cast<DotOperandEncodingAttr>(layout)) {
     auto rank = getWarpsPerCTA(dotLayout.getParent()).size();
-    if (dyn_cast<intel::DpasEncodingAttr>(dotLayout.getParent())) {
-      SmallVector<unsigned> order(rank);
-      std::iota(order.rbegin(), order.rend(), 0);
-      return order;
+    if (auto mmaParent = dyn_cast<MmaEncodingTrait>(dotLayout.getParent())) {
+      return mmaParent.getOrderForDotOperand(dotLayout.getOpIdx(), rank);
     }
+    // This branch had to be left because not all types
+    // inherit `MmaEncodingTrait` interface, for example `BlockedEncodingAttr`.
     return getOrderForDotOperand(dotLayout.getOpIdx(), rank);
   }
   if (auto sliceLayout = dyn_cast<SliceEncodingAttr>(layout)) {
@@ -837,6 +837,12 @@ unsigned AMDMfmaEncodingAttr::getTotalElemsPerThread(ArrayRef<int64_t> shape,
   return product<unsigned>(getElemsPerThread(shape, eltTy));
 }
 
+SmallVector<unsigned>
+AMDMfmaEncodingAttr::getOrderForDotOperand(unsigned opIdx,
+                                           unsigned rank) const {
+  return ::getOrderForDotOperand(opIdx, rank);
+}
+
 // Wmma encoding
 
 SmallVector<unsigned>
@@ -866,6 +872,12 @@ unsigned AMDWmmaEncodingAttr::getTotalElemsPerThread(ArrayRef<int64_t> shape,
   return product<unsigned>(getElemsPerThread(shape, eltTy));
 }
 
+SmallVector<unsigned>
+AMDWmmaEncodingAttr::getOrderForDotOperand(unsigned opIdx,
+                                           unsigned rank) const {
+  return ::getOrderForDotOperand(opIdx, rank);
+}
+
 SmallVector<unsigned>
 NvidiaMmaEncodingAttr::getElemsPerThread(ArrayRef<int64_t> shape,
                                          Type eltTy) const {
@@ -953,6 +965,12 @@ unsigned NvidiaMmaEncodingAttr::getTotalElemsPerThread(ArrayRef<int64_t> shape,
   return product<unsigned>(getElemsPerThread(shape, eltTy));
 }
 
+SmallVector<unsigned>
+NvidiaMmaEncodingAttr::getOrderForDotOperand(unsigned opIdx,
+                                             unsigned rank) const {
+  return ::getOrderForDotOperand(opIdx, rank);
+}
+
 //
 
 SmallVector<unsigned>

third_party/intel/include/Dialect/TritonIntelGPU/IR/TritonIntelGPUAttrDefs.td

@@ -84,6 +84,7 @@ along the row (resp. col) dimension.
     SmallVector<int64_t> getDPASRepetitions(ArrayRef<int64_t> shape, int opIdx) const;
     SmallVector<unsigned> getSizePerThreadForOperand(int kWidth,unsigned opIdx) const;
     SmallVector<unsigned> getElemsPerThreadForOperands(ArrayRef<int64_t> shape, Type eltTy, unsigned opIdx) const;
+    SmallVector<unsigned> getOrderForDotOperand(unsigned opIdx, unsigned rank) const;
     SmallVector<unsigned> getShapePerCTATileForOperand(ArrayRef<int64_t> shape, int kWidth, int opIdx) const;
     unsigned getTotalElemsPerThreadForOperand(ArrayRef<int64_t> shape, Type eltTy, int kWidth, int opIdx) const;
 

third_party/intel/lib/Dialect/TritonIntelGPU/IR/Dialect.cpp

@@ -297,6 +297,13 @@ SmallVector<unsigned> DpasEncodingAttr::getElemsPerThreadForOperands(
           static_cast<unsigned>(sizePerThread[1] * repetitions[1])};
 };
 
+SmallVector<unsigned>
+DpasEncodingAttr::getOrderForDotOperand(unsigned opIdx, unsigned rank) const {
+  SmallVector<unsigned> order(rank);
+  std::iota(order.rbegin(), order.rend(), 0);
+  return order;
+}
+
 SmallVector<unsigned> DpasEncodingAttr::getContigPerThread() {
   unsigned threadsPerWarp = getSubGroupSize();
   auto shapeC = getDPASInstShapeC();