fix order for DPAS

anmyachev · anmyachev · commit dfd526b946f8 · 2024-10-28T14:53:13.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
@@ -304,6 +304,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
`@@ -304,6 +304,11 @@ SmallVector<unsigned> getOrder(Attribute layout) {`
`304`	`304`	`}`
`305`	`305`	`if (auto dotLayout = dyn_cast<DotOperandEncodingAttr>(layout)) {`
`306`	`306`	`auto rank = getWarpsPerCTA(dotLayout.getParent()).size();`
	`307`	`+ if (dyn_cast<intel::DpasEncodingAttr>(dotLayout.getParent())) {`
	`308`	`+ SmallVector<unsigned> order(rank);`
	`309`	`+ std::iota(order.rbegin(), order.rend(), 0);`
	`310`	`+ return order;`
	`311`	`+ }`
`307`	`312`	`return getOrderForDotOperand(dotLayout.getOpIdx(), rank);`
`308`	`313`	`}`
`309`	`314`	`if (auto sliceLayout = dyn_cast<SliceEncodingAttr>(layout)) {`