Clean the getShapePerCTATile and getShapePerCTATileForOperand interface in DPAS layout.

Author: chengjunlu

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

@@ -84,7 +84,6 @@ 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> getShapePerCTATileForOperand(ArrayRef<int64_t> shape, int kWidth, int opIdx) const;
     SmallVector<unsigned> getRepOrderForOperand(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

@@ -147,18 +147,6 @@ SmallVector<unsigned> DpasEncodingAttr::getSizePerThread() const {
   return res;
 }
 
-SmallVector<unsigned>
-DpasEncodingAttr::getShapePerCTATile(ArrayRef<int64_t> tensorShape) const {
-  auto shapeC = getShapeC();
-  SmallVector<unsigned> warpsPerCTA = getWarpsPerCTA();
-  size_t rank = shapeC.size();
-  SmallVector<unsigned> shapePerCTATile(rank);
-  llvm::transform(
-      llvm::zip_equal(shapeC, warpsPerCTA), shapePerCTATile.begin(),
-      [](auto entry) { return std::get<0>(entry) * std::get<1>(entry); });
-  return shapePerCTATile;
-}
-
 SmallVector<unsigned> DpasEncodingAttr::getRepOrder() const {
   llvm::report_fatal_error("NYI. DpasEncodingAttr::getRepOrder");
 }
@@ -174,7 +162,14 @@ DpasEncodingAttr::getElemsPerThread(ArrayRef<int64_t> shape, Type eltTy) const {
   assert((rank == 2 || rank == 3) && "Unexpected rank of mma layout");
 
   SmallVector<unsigned> elemsPerThread(rank, 1);
-  auto shapePerCTATile = getShapePerCTATile(shape);
+
+  auto shapeC = getShapeC();
+  SmallVector<unsigned> warpsPerCTA = getWarpsPerCTA();
+  SmallVector<unsigned> shapePerCTATile(rank);
+  llvm::transform(
+      llvm::zip_equal(shapeC, warpsPerCTA), shapePerCTATile.begin(),
+      [](auto entry) { return std::get<0>(entry) * std::get<1>(entry); });
+
   unsigned tilesRow =
       ceil<unsigned>(shape[rank - 2], shapePerCTATile[rank - 2]);
   unsigned tilesCol =
@@ -306,33 +301,6 @@ SmallVector<unsigned> DpasEncodingAttr::getThreadsPerWarp() const {
   return res;
 }
 
-SmallVector<unsigned>
-DpasEncodingAttr::getShapePerCTATileForOperand(ArrayRef<int64_t> shape,
-                                               int kWidth, int opIdx) const {
-  auto parentShapePerCTATile = getShapePerCTATile(shape);
-  size_t rank = parentShapePerCTATile.size();
-  assert((rank == 2 || rank == 3) && "unexpected rank number for Dpas layout");
-  if (opIdx == 0) {
-    auto shapeA = getShapeA();
-    return (rank == 2)
-               ? SmallVector<unsigned>{parentShapePerCTATile[0], shapeA[1]}
-               : SmallVector<unsigned>{parentShapePerCTATile[0],
-                                       parentShapePerCTATile[rank - 2],
-                                       shapeA[rank - 1]};
-  }
-
-  if (opIdx == 1) {
-    auto shapeB = getShapeB();
-    return (rank == 2)
-               ? SmallVector<unsigned>{shapeB[0], parentShapePerCTATile[1]}
-               : SmallVector<unsigned>{parentShapePerCTATile[0],
-                                       shapeB[rank - 2],
-                                       parentShapePerCTATile[rank - 1]};
-  }
-
-  llvm::report_fatal_error("DotOperandEncodingAttr opIdx must be 0 or 1");
-}
-
 SmallVector<unsigned>
 DpasEncodingAttr::getSizePerThreadForOperand(int kWidth, unsigned opIdx) const {
   ArrayRef<unsigned> repCluster = getRepCluster();

third_party/intel/lib/TritonIntelGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -248,8 +248,8 @@ struct ConvertLayoutOpConversion
     SmallVector<unsigned> outNumCTAsEachRep(rank);
     SmallVector<unsigned> inNumCTAs(rank);
     SmallVector<unsigned> outNumCTAs(rank);
-    auto srcShapePerCTATile = getShapePerCTATile(srcLayout, srcTy.getShape());
-    auto dstShapePerCTATile = getShapePerCTATile(dstLayout, shape);
+    auto srcShapePerCTATile = getShapePerCTATile(srcLayout);
+    auto dstShapePerCTATile = getShapePerCTATile(dstLayout);
     auto shapePerCTA = getShapePerCTA(srcLayout, shape);
 
     for (unsigned d = 0; d < rank; ++d) {

third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -36,7 +36,7 @@ Value redundantDataMask(Type valueTy, ConversionPatternRewriter &rewriter,
     auto threadsPerWarp = triton::gpu::getThreadsPerWarp(layout);
     auto warpsPerCTA = triton::gpu::getWarpsPerCTA(layout);
     auto order = triton::gpu::getOrder(layout);
-    auto shapePerCTATile = triton::gpu::getShapePerCTATile(layout, shape);
+    auto shapePerCTATile = triton::gpu::getShapePerCTATile(layout);
     Value warpSize = LLVM::intel::getModuleWarpSize(rewriter, loc);
     Value laneId = urem(tid, warpSize);
     Value warpId = udiv(tid, warpSize);

third_party/intel/lib/TritonIntelGPUToLLVM/ReduceOpToLLVM.cpp

@@ -408,7 +408,7 @@ struct ReduceOpConversion
         auto resultIndices = ::intel::emitIndices(loc, rewriter, targetInfo,
                                                   resultLayout, resultTy, true);
         auto resultShape = resultTy.getShape();
-        auto resultCTATile = getShapePerCTATile(resultLayout, resultShape);
+        auto resultCTATile = getShapePerCTATile(resultLayout);
         assert(resultIndices.size() == resultElems);
 
         SmallVector<Value> resultVals(resultElems);