[intel] Small fixes for dot operand properties

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

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

@@ -82,10 +82,10 @@ along the row (resp. col) dimension.
     SmallVector<unsigned> getShapeB() const;
     SmallVector<unsigned> getShapeC() const;
     SmallVector<int64_t> getDPASRepetitions(ArrayRef<int64_t> shape, int opIdx) const;
-    SmallVector<unsigned> getSizePerThreadForOperands(unsigned opIdx) const;
+    SmallVector<unsigned> getSizePerThreadForOperand(int kWidth,unsigned opIdx) const;
     SmallVector<unsigned> getElemsPerThreadForOperands(ArrayRef<int64_t> shape, Type eltTy, unsigned opIdx) const;
-    SmallVector<unsigned> getShapePerCTATileForDotOperands(ArrayRef<int64_t> shape, int opIdx) const;
-    unsigned getTotalElemsPerThreadForOperands(ArrayRef<int64_t> shape, Type eltTy, 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;
 
     bool supportReduction() const {
       return true;

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

@@ -247,7 +247,7 @@ DpasEncodingAttr::getDPASRepetitions(ArrayRef<int64_t> shape, int opIdx) const {
                                                   warpsPerCTA[rank - 1]))};
 }
 
-unsigned DpasEncodingAttr::getTotalElemsPerThreadForOperands(
+unsigned DpasEncodingAttr::getTotalElemsPerThreadForOperand(
     ArrayRef<int64_t> shape, mlir::Type eltTy, int kWidth, int opIdx) const {
   auto shapePerCTA = getShapePerCTA(*this, shape);
   auto rep = getDPASRepetitions(shapePerCTA, opIdx);
@@ -298,8 +298,8 @@ SmallVector<unsigned> DpasEncodingAttr::getThreadsPerWarp() const {
 }
 
 SmallVector<unsigned>
-DpasEncodingAttr::getShapePerCTATileForDotOperands(ArrayRef<int64_t> shape,
-                                                   int opIdx) const {
+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");
@@ -325,7 +325,7 @@ DpasEncodingAttr::getShapePerCTATileForDotOperands(ArrayRef<int64_t> shape,
 }
 
 SmallVector<unsigned>
-DpasEncodingAttr::getSizePerThreadForOperands(unsigned opIdx) const {
+DpasEncodingAttr::getSizePerThreadForOperand(int kWidth, unsigned opIdx) const {
   ArrayRef<unsigned> repCluster = getRepCluster();
   size_t rank = repCluster.size();
   assert((rank == 2 || rank == 3) && "unexpected rank number for Dpas layout");
@@ -367,7 +367,7 @@ DpasEncodingAttr::getSizePerThreadForOperands(unsigned opIdx) const {
 
 SmallVector<unsigned> DpasEncodingAttr::getElemsPerThreadForOperands(
     ArrayRef<int64_t> shape, Type eltTy, unsigned opIdx) const {
-  SmallVector<unsigned> sizePerThread = getSizePerThreadForOperands(opIdx);
+  SmallVector<unsigned> sizePerThread = getSizePerThreadForOperand(0, opIdx);
   SmallVector<int64_t> repetitions = getDPASRepetitions(shape, opIdx);
 
   size_t rank = shape.size();