[intel] Fix getTotalElemsPerThread failures

whitneywhtsang · whitneywhtsang · commit 1efa822bd5fd · 2024-11-17T01:08:53.000Z
Signed-off-by: Whitney Tsang &lt;whitney.tsang@intel.com&gt;
diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp
@@ -992,6 +992,10 @@ unsigned DotOperandEncodingAttr::getTotalElemsPerThread(ArrayRef<int64_t> shape,
       return amdWmmaParent.getTotalElemsPerThreadForOperand(
           shape, eltTy, getKWidth(), getOpIdx());
     }
+    if (auto dpasParent = mlir::dyn_cast<intel::DpasEncodingAttr>(mmaParent)) {
+      return dpasParent.getTotalElemsPerThreadForOperand(
+          shape, eltTy, getKWidth(), getOpIdx());
+    }
   }
   if (auto blockedLayout = mlir::dyn_cast<BlockedEncodingAttr>(getParent())) {
     auto shapePerCTA = getShapePerCTA(*this, shape);

Original file line number	Diff line number	Diff line change
`@@ -992,6 +992,10 @@ unsigned DotOperandEncodingAttr::getTotalElemsPerThread(ArrayRef<int64_t> shape,`
`992`	`992`	`return amdWmmaParent.getTotalElemsPerThreadForOperand(`
`993`	`993`	`shape, eltTy, getKWidth(), getOpIdx());`
`994`	`994`	`}`
	`995`	`+ if (auto dpasParent = mlir::dyn_cast<intel::DpasEncodingAttr>(mmaParent)) {`
	`996`	`+ return dpasParent.getTotalElemsPerThreadForOperand(`
	`997`	`+ shape, eltTy, getKWidth(), getOpIdx());`
	`998`	`+ }`
`995`	`999`	`}`
`996`	`1000`	`if (auto blockedLayout = mlir::dyn_cast<BlockedEncodingAttr>(getParent())) {`
`997`	`1001`	`auto shapePerCTA = getShapePerCTA(*this, shape);`