Add dotOperandDpasToLinearLayout (Tensor B) (#2422)

`dotOperandMfmaToLinearLayout` for tensor B is added in upstream commit triton-lang/triton@ff02a46 recently. CI: https://github.com/intel/intel-xpu-backend-for-triton/actions/runs/11170777152 Signed-off-by: Whitney Tsang <[email protected]>
intel · Oct 8, 2024 · c2570b7 · c2570b7
1 parent 549c9e1
commit c2570b7
Show file tree

Hide file tree

Showing 4 changed files with 168 additions and 161 deletions.
diff --git a/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp b/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp
@@ -1,5 +1,6 @@
 #include <vector>
 
+#include "intel/include/Dialect/TritonIntelGPU/IR/LinearLayoutConversions.h"
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/TritonGPU/IR/Attributes.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
@@ -827,6 +828,9 @@ DotOperandEncodingAttr::toLinearLayout(ArrayRef<int64_t> shape) const {
   if (auto mfmaLayout = llvm::dyn_cast<AMDMfmaEncodingAttr>(getParent())) {
     return dotOperandMfmaToLinearLayout(*this, shape);
   }
+  if (auto dpasLayout = llvm::dyn_cast<intel::DpasEncodingAttr>(getParent())) {
+    return dotOperandDpasToLinearLayout(*this, shape);
+  }
 
   return std::nullopt;
 }