Allow fallback to legacy layout for DotOp layout (#4617)

The support of toLinearLayout of DotOp is not yet implemented upstream, so this PR allows fallback to legacy layout for DotOp instead of error out. When the support of toLinearLayout of DotOp is implemented, this change can be reverted. There is a general issue in Triton lowering for the nested layout slice->dot. This combination of the layout is not hit for NV and AMD backends while it is possible for Intel backend after some optimization. Signed-off-by: Whitney Tsang <[email protected]>
triton-lang · Aug 30, 2024 · e3d2f44 · e3d2f44
1 parent 698ec4e
commit e3d2f44
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp b/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp
@@ -649,9 +649,12 @@ SliceEncodingAttr::toLinearLayout(ArrayRef<int64_t> shape) const {
   parentShape.insert(parentShape.begin() + getDim(), 1);
   std::optional<LinearLayout> parentLL =
       triton::gpu::toLinearLayout(parentShape, getParent());
-  if (!parentLL.has_value())
+  if (!parentLL.has_value()) {
+    if (mlir::isa<DotOperandEncodingAttr>(getParent()))
+      return std::nullopt;
     llvm::report_fatal_error(
         "Failed to compute parent layout for slice layout.");
+  }
 
   // Remove dimension getDim() from the parent layout.
   //