fix order for DPAS

anmyachev · anmyachev · commit 1b6689e87280 · 2024-10-21T21:45:51.000Z
Signed-off-by: Anatoly Myachev &lt;anatoly.myachev@intel.com&gt;
diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp
@@ -299,6 +299,11 @@ SmallVector<unsigned> getOrder(Attribute layout) {
   }
   if (auto dotLayout = dyn_cast<DotOperandEncodingAttr>(layout)) {
     auto rank = getWarpsPerCTA(dotLayout.getParent()).size();
+    if (dyn_cast<intel::DpasEncodingAttr>(dotLayout.getParent())) {
+      SmallVector<unsigned> order(rank);
+      std::iota(order.rbegin(), order.rend(), 0);
+      return order;
+    }
     return getOrderForDotOperand(dotLayout.getOpIdx(), rank);
   }
   if (auto sliceLayout = dyn_cast<SliceEncodingAttr>(layout)) {

Original file line number	Diff line number	Diff line change
`@@ -299,6 +299,11 @@ SmallVector<unsigned> getOrder(Attribute layout) {`
`299`	`299`	`}`
`300`	`300`	`if (auto dotLayout = dyn_cast<DotOperandEncodingAttr>(layout)) {`
`301`	`301`	`auto rank = getWarpsPerCTA(dotLayout.getParent()).size();`
	`302`	`+ if (dyn_cast<intel::DpasEncodingAttr>(dotLayout.getParent())) {`
	`303`	`+ SmallVector<unsigned> order(rank);`
	`304`	`+ std::iota(order.rbegin(), order.rend(), 0);`
	`305`	`+ return order;`
	`306`	`+ }`
`302`	`307`	`return getOrderForDotOperand(dotLayout.getOpIdx(), rank);`
`303`	`308`	`}`
`304`	`309`	`if (auto sliceLayout = dyn_cast<SliceEncodingAttr>(layout)) {`