[Mlir-commits] [mlir] [mlir][irdl] Fix missing verifier in irdl.parametric (PR #92700)

[Théo Degioanni]

https://github.com/Moxinilian created https://github.com/llvm/llvm-project/pull/92700

The parametric op was not checking the symbol it points to is a type or attribute. This PR also fixes a small bug where an invalid IRDL file would not end processing in mlir-opt.

>From 3856182a2bb122ec6c148c9d491439c6680850be Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Th=C3=A9o=20Degioanni?=
 <theo.degioanni.llvm.deluge062 at simplelogin.fr>
Date: Sun, 19 May 2024 18:48:26 +0100
Subject: [PATCH] fix missing verifier in IRDL parametric

---
 mlir/include/mlir/Dialect/IRDL/IR/IRDLOps.td |  3 +-
 mlir/lib/Dialect/IRDL/IR/IRDL.cpp            | 31 +++++++++++++++-----
 mlir/lib/Tools/mlir-opt/MlirOptMain.cpp      |  2 ++
 mlir/test/Dialect/IRDL/invalid.irdl.mlir     | 17 ++++++++++-
 4 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Dialect/IRDL/IR/IRDLOps.td b/mlir/include/mlir/Dialect/IRDL/IR/IRDLOps.td
index aa6a8e93c0288..d2765dec420ac 100644
--- a/mlir/include/mlir/Dialect/IRDL/IR/IRDLOps.td
+++ b/mlir/include/mlir/Dialect/IRDL/IR/IRDLOps.td
@@ -503,7 +503,8 @@ def IRDL_BaseOp : IRDL_ConstraintOp<"base",
 }
 
 def IRDL_ParametricOp : IRDL_ConstraintOp<"parametric",
-    [ParentOneOf<["TypeOp", "AttributeOp", "OperationOp"]>, Pure]> {
+    [ParentOneOf<["TypeOp", "AttributeOp", "OperationOp"]>,
+     DeclareOpInterfaceMethods<SymbolUserOpInterface>, Pure]> {
   let summary = "Constraints an attribute/type base and its parameters";
   let description = [{
     `irdl.parametric` defines a constraint that accepts only a single type
diff --git a/mlir/lib/Dialect/IRDL/IR/IRDL.cpp b/mlir/lib/Dialect/IRDL/IR/IRDL.cpp
index 4eae2b03024c2..222eca5baebe8 100644
--- a/mlir/lib/Dialect/IRDL/IR/IRDL.cpp
+++ b/mlir/lib/Dialect/IRDL/IR/IRDL.cpp
@@ -132,22 +132,37 @@ LogicalResult BaseOp::verify() {
   return success();
 }
 
+static LogicalResult
+checkSymbolIsTypeOrAttribute(SymbolTableCollection &symbolTable,
+                             Operation *source, SymbolRefAttr symbol) {
+  Operation *targetOp = symbolTable.lookupNearestSymbolFrom(source, symbol);
+  if (!targetOp)
+    return source->emitOpError() << "symbol '" << symbol << "' not found";
+
+  if (!isa<TypeOp, AttributeOp>(targetOp))
+    return source->emitOpError() << "'" << symbol
+                                 << "' does not refer to a type or attribute "
+                                    "definition (refers to '"
+                                 << targetOp->getName() << "')";
+
+  return success();
+}
+
 LogicalResult BaseOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   std::optional<SymbolRefAttr> baseRef = getBaseRef();
   if (!baseRef)
     return success();
 
-  TypeOp typeOp = symbolTable.lookupNearestSymbolFrom<TypeOp>(*this, *baseRef);
-  if (typeOp)
-    return success();
+  return checkSymbolIsTypeOrAttribute(symbolTable, *this, *baseRef);
+}
 
-  AttributeOp attrOp =
-      symbolTable.lookupNearestSymbolFrom<AttributeOp>(*this, *baseRef);
-  if (attrOp)
+LogicalResult
+ParametricOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  std::optional<SymbolRefAttr> baseRef = getBaseType();
+  if (!baseRef)
     return success();
 
-  return emitOpError() << "'" << *baseRef
-                       << "' does not refer to a type or attribute definition";
+  return checkSymbolIsTypeOrAttribute(symbolTable, *this, *baseRef);
 }
 
 /// Parse a value with its variadicity first. By default, the variadicity is
diff --git a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
index 44c5e9826f3b7..a1b2893a973b1 100644
--- a/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
+++ b/mlir/lib/Tools/mlir-opt/MlirOptMain.cpp
@@ -266,6 +266,8 @@ LogicalResult loadIRDLDialects(StringRef irdlFile, MLIRContext &ctx) {
 
   // Parse the input file.
   OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, &ctx));
+  if (!module)
+    return failure();
 
   // Load IRDL dialects.
   return irdl::loadDialects(module.get());
diff --git a/mlir/test/Dialect/IRDL/invalid.irdl.mlir b/mlir/test/Dialect/IRDL/invalid.irdl.mlir
index d62bb498a7ad9..a63b0df2c5da2 100644
--- a/mlir/test/Dialect/IRDL/invalid.irdl.mlir
+++ b/mlir/test/Dialect/IRDL/invalid.irdl.mlir
@@ -6,7 +6,7 @@ func.func private @foo()
 
 irdl.dialect @testd {
   irdl.type @type {
-    // expected-error at +1 {{'@foo' does not refer to a type or attribute definition}}
+    // expected-error at +1 {{symbol '@foo' not found}}
     %0 = irdl.base @foo
     irdl.parameters(%0)
   }
@@ -41,3 +41,18 @@ irdl.dialect @testd {
     irdl.parameters(%0)
   }
 }
+
+// -----
+
+irdl.dialect @invalid_parametric {
+  irdl.operation @foo {
+    // expected-error at +1 {{'@not_a_type_or_attr' does not refer to a type or attribute definition}}
+    %param = irdl.parametric @not_a_type_or_attr<>
+    irdl.results(%param)
+  }
+
+  irdl.operation @not_a_type_or_attr {
+    %param = irdl.is i1
+    irdl.results(%param)
+  }
+}