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

format + remove debug prints

Author: alexbaden

third_party/intel/backend/compiler.py

@@ -250,7 +250,6 @@ def make_ttgir(mod, metadata, opt, properties):
             intel.passes.ttgpuir.add_rewrite_tensor_pointer(pm)
         intel.passes.ttgpuir.add_pipeline(pm, opt.num_stages, False)
 
-    
         passes.ttgpuir.add_optimize_thread_locality(pm)
         passes.ttgpuir.add_optimize_dot_operands(pm, True)
         passes.common.add_cse(pm)

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;
@@ -122,31 +123,23 @@ class BlockedToDPAS : public OpRewritePattern<tt::DotOp> {
         oldAType.getElementType().isFloat8E4M3FN())
       dpasElemBitWidths = 2 * dpasElemBitWidths;
 
-    // now we can get the order from the a defining op 
-
-    llvm::errs() << "oldAType: " << oldAType << "\n";
-    llvm::errs() << "oldBType: " << oldBType << "\n";
-
-    llvm::errs() << "a: " << a << "\n";
-    llvm::errs() << "a defining op: " << *a.getDefiningOp() << "\n";
-
     SmallVector<unsigned> order;
-    Operation* aOp = a.getDefiningOp();
+    Operation *aOp = a.getDefiningOp();
     if (isa<ttg::ConvertLayoutOp>(aOp)) {
       assert(aOp->getNumOperands() == 1);
       auto aLoad = aOp->getOperand(0);
-      order = triton::gpu::getOrder(cast<RankedTensorType>(aLoad.getType()).getEncoding());
+      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 
 
     SmallVector<unsigned> warpsPerTile =
-        getWarpsPerTile(dotOp, dpasCap, retShape, numWarps);
+        getWarpsPerTile(dotOp, dpasCap, retShape, numWarps, order);
     size_t rank = retShape.size();
     SmallVector<unsigned> repCluster(rank, 1);