use A matrix layout order when determining dpas order in accelerate matmul

Author: alexbaden

third_party/intel/lib/TritonIntelGPUTransforms/AccelerateMatmul.cpp

@@ -31,7 +31,8 @@ namespace {
 SmallVector<unsigned>
 getWarpsPerTile(tt::DotOp dotOp,
                 ttg::intel::DpasEncodingAttr::DPASCapability dpasCap,
-                const ArrayRef<int64_t> shape, unsigned numWarps) {
+                const ArrayRef<int64_t> shape, unsigned numWarps, const SmallVector<unsigned>& order) {
+
   auto filter = [&dotOp](Operation *op) {
     return op->getParentRegion() == dotOp->getParentRegion();
   };
@@ -63,7 +64,7 @@ getWarpsPerTile(tt::DotOp dotOp,
   uint32_t colRowRatio =
       ceil<uint32_t>(dpasCap.executionSize, dpasCap.repeatCount);
 
-  int rowDim = rank - 2, colDim = rank - 1;
+  int rowDim = order[rank - 2], colDim = order[rank - 1];
   do {
     if (ret[rowDim] * ret[colDim] >= numWarps)
       break;
@@ -138,15 +139,18 @@ class BlockedToDPAS : public OpRewritePattern<tt::DotOp> {
       order = triton::gpu::getOrder(cast<RankedTensorType>(aLoad.getType()).getEncoding());
     } else {
       assert(isa<tt::LoadOp>(aOp) && "expecting load input to DPAS");
-      order = triton::gpu::getOrder(cast<RankedTensorType>(aLoad.getType()).getEncoding());
+      assert(aOp->getNumResults() == 1);
+      auto ret = aOp->getResult(0);
+      order = triton::gpu::getOrder(cast<RankedTensorType>(ret.getType()).getEncoding());
     }
-    // order = triton::gpu::getOrder(a.getDefiningOp().getEncoding());
     llvm::errs() << "a load order: " << order[0] << ", " << order[1] << "\n";
-
-    // now find the fast changing dimension from the order 
+    #if 0
+    const bool aIsTransposed = order.size() == 2 && order[0] == 0 && order[1] == 1;
+    llvm::errs() << "Transposed? " << aIsTransposed << "\n";
+    #endif
 
     SmallVector<unsigned> warpsPerTile =
-        getWarpsPerTile(dotOp, dpasCap, retShape, numWarps);
+        getWarpsPerTile(dotOp, dpasCap, retShape, numWarps, order);
     size_t rank = retShape.size();
     SmallVector<unsigned> repCluster(rank, 1);