the pipeline pass now generates and accepts valid IR

2022-06-07 19:34:59 +08:00
parent 560e29229b
commit 7b09b5f9e9
4 changed files with 82 additions and 37 deletions
--- a/lib/Dialect/TritonGPU/IR/Dialect.cpp
+++ b/lib/Dialect/TritonGPU/IR/Dialect.cpp
@@ -276,6 +276,17 @@ static Type getPointeeType(Type type) {
 }
 }

+static LogicalResult verify(CopyAsyncOp op) {
+  Type resType = op.getResult().getType();
+  if (auto tensorType = resType.dyn_cast<RankedTensorType>()) {
+    Attribute encoding = tensorType.getEncoding();
+    if (!encoding.isa<TritonGPUSharedEncodingAttr>())
+      return op.emitOpError("copy_async should return a shared memory tensor");
+  } else
+    return op.emitOpError("copy_async should return a tensor");
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "triton/Dialect/TritonGPU/IR/Ops.cpp.inc"