Add a new interface for 3rd party Triton GPU dialect to verify the Do…

…tOp layout.
intel · Dec 18, 2024 · cbf3c59 · cbf3c59
1 parent a7dcf4c
commit cbf3c59
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 7 deletions.
diff --git a/include/triton/Dialect/Triton/IR/Dialect.h b/include/triton/Dialect/Triton/IR/Dialect.h
@@ -76,6 +76,11 @@ class DialectInferLayoutInterface
   virtual LogicalResult
   verifyDotOpEncodingCompatibility(Operation *op, Attribute operandEncodingA,
                                    Attribute operandEncodingB) const = 0;
+
+  // Verify the dotOp layout encoding is legal if it uses 3rd party Triton GPU
+  // dialect attribute as parent.
+  virtual LogicalResult verifyDotOpEncoding(unsigned opIdx, Attribute parent,
+                                            unsigned kWidth) const = 0;
 };
 
 class DialectVerifyTensorLayoutInterface

diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp
@@ -1189,20 +1189,22 @@ LogicalResult DotOperandEncodingAttr::verify(
     return success();
   }
 
-  if (auto parentAttr = mlir::dyn_cast<intel::DpasEncodingAttr>(parent)) {
-    if (kWidth != parentAttr.getOpsPerChannel())
-      return emitError() << "ttg.dot_op kWidth parameter must match the "
-                            "parent's opsPerChannel";
-    return success();
-  }
-
   if (auto parentAttr = mlir::dyn_cast<intel::WarpEncodingAttr>(parent)) {
     if (kWidth != 0)
       return emitError() << "ttg.dot_op kWidth parameter is not supported "
                             "when the parent is a warp layout";
     return success();
   }
 
+  if (auto parentAttr = mlir::dyn_cast<MmaEncodingTrait>(parent)) {
+    Dialect &dialect = parentAttr.getDialect();
+    auto interface = mlir::cast<DialectInferLayoutInterface>(&dialect);
+    if (interface->verifyDotOpEncoding(opIdx, parent, kWidth).failed())
+      return emitError() << "ttg.dot_op is invalid with parent layout: "
+                         << parent;
+    return success();
+  }
+
   if (auto parentAttr = mlir::dyn_cast<BlockedEncodingAttr>(parent)) {
     if (kWidth != 0)
       return emitError() << "ttg.dot_op kWidth parameter is not supported "

diff --git a/third_party/intel/lib/Dialect/TritonIntelGPU/IR/Dialect.cpp b/third_party/intel/lib/Dialect/TritonIntelGPU/IR/Dialect.cpp
@@ -639,6 +639,20 @@ struct TritonIntelGPUInferLayoutInterface
     return success();
   }
 
+  LogicalResult verifyDotOpEncoding(unsigned opIdx, Attribute parent,
+                                    unsigned kWidth) const override {
+    ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError;
+
+    if (auto parentAttr = mlir::dyn_cast<intel::DpasEncodingAttr>(parent)) {
+      if (kWidth != parentAttr.getOpsPerChannel())
+        return emitError() << "ttg.dot_op kWidth parameter must match the "
+                              "parent's opsPerChannel";
+      return success();
+    }
+
+    return emitError() << "ttg.dot_op unknown parent layout: " << parent;
+  }
+
   LogicalResult
   inferReshapeOpNoReorderEncoding(ArrayRef<int64_t> srcShape, Attribute srcEnc,
                                   ArrayRef<int64_t> dstShape, Attribute &dstEnc,