[mlir][ODS] Verify type constraint in `TypeAttrOf`

[zero9178]

Summary

The current implementation does not verify the type constraint, meaning that any type that happens to be of the same C++ type would pass the verifier.
E.g. a TypeAttrOf<I64> would happily accept a i32 since both satisfy isa<IntegerType>().

This patch fixes that by adding an optional type predicate parameter to TypeAttrBase that the type within TypeAttr has to satisfy. TypeAttrOf then simply passes the predicate of its type parameter as argument.

mlir/include/mlir/IR/OpBase.td

class TypedStrAttr<Type ty>
     : StringBasedAttr<CPred<"$_self.isa<::mlir::StringAttr>()">,
                             "string attribute"> {
   let valueType = ty;
 }
 
 // Base class for attributes containing types. Example:
 //   def IntTypeAttr : TypeAttrBase<"IntegerType", "integer type attribute">
 // defines a type attribute containing an integer type.
-class TypeAttrBase<string retType, string summary> :
+class TypeAttrBase<string retType, string summary,
+                        Pred typePred = CPred<"true">> :
     Attr<And<[
       CPred<"$_self.isa<::mlir::TypeAttr>()">,
       CPred<"$_self.cast<::mlir::TypeAttr>().getValue().isa<"
-            # retType # ">()">]>,
+            # retType # ">()">,
+      SubstLeaves<"$_self",
+                    "$_self.cast<::mlir::TypeAttr>().getValue()", typePred>]>,
     summary> {
   let storageType = [{ ::mlir::TypeAttr }];
   let returnType = retType;
   let valueType = NoneType;
   let convertFromStorage = "$_self.getValue().cast<" # retType # ">()";
 }
 
 def TypeAttr : TypeAttrBase<"::mlir::Type", "any type attribute"> {
   let constBuilderCall = "::mlir::TypeAttr::get($0)";
 }
 
 class TypeAttrOf<Type ty>
-   : TypeAttrBase<ty.cppClassName, "type attribute of " # ty.summary> {
+   : TypeAttrBase<ty.cppClassName, "type attribute of " # ty.summary,
+                    ty.predicate> {
   let constBuilderCall = "::mlir::TypeAttr::get($0)";
 }
 
 // The mere presence of unit attributes has a meaning.  Therefore, unit
 // attributes are always treated as optional and accessors to them return
 // "true" if the attribute is present and "false" otherwise.
 def UnitAttr : Attr<CPred<"$_self.isa<::mlir::UnitAttr>()">, "unit attribute"> {
   let storageType = [{ ::mlir::UnitAttr }];

mlir/test/IR/attribute.mlir

func.func @default_value_printing(%arg0 : i32) {
   // The attribute SHOULD NOT be printed because it is equal to the default
   // CHECK: test.default_value_print %arg0
   "test.default_value_print"(%arg0) {"value_with_default" = 0 : i32} : (i32) -> ()
   // The attribute SHOULD be printed because it is not equal to the default
   // CHECK: test.default_value_print {value_with_default = 1 : i32} %arg0
   "test.default_value_print"(%arg0) {"value_with_default" = 1 : i32} : (i32) -> ()
   return
 }
+
+// -----
+
+func.func @type_attr_of_fail() {
+    // expected-error @below {{failed to satisfy constraint: type attribute of 64-bit signless integer}}
+    test.type_attr_of i32
+    return
+}

mlir/test/lib/Dialect/Test/TestOps.td

 
 def TypedAttrOp : TEST_Op<"typed_attr"> {
   let arguments = (ins TypeAttr:$type, AnyAttr:$attr);
   let assemblyFormat = [{
     attr-dict $type `=` custom<AttrElideType>(ref($type), $attr)
   }];
 }
 
+def TypeAttrOfOp : TEST_Op<"type_attr_of"> {
+  let arguments = (ins TypeAttrOf<I64>:$type);
+  let assemblyFormat = [{
+    attr-dict $type
+  }];
+}
+
 def DenseArrayAttrOp : TEST_Op<"dense_array_attr"> {
   let arguments = (ins
     DenseBoolArrayAttr:$i1attr,
     DenseI8ArrayAttr:$i8attr,
     DenseI16ArrayAttr:$i16attr,
     DenseI32ArrayAttr:$i32attr,
     DenseI64ArrayAttr:$i64attr,
     DenseF32ArrayAttr:$f32attr,

mlir/test/mlir-tblgen/op-attribute.td

 // DEF: if (tblgen_bool_attr && !((tblgen_bool_attr.isa<::mlir::BoolAttr>())))
 // DEF: if (tblgen_i32_attr && !(((tblgen_i32_attr.isa<::mlir::IntegerAttr>())) && ((tblgen_i32_attr.cast<::mlir::IntegerAttr>().getType().isSignlessInteger(32)))))
 // DEF: if (tblgen_i64_attr && !(((tblgen_i64_attr.isa<::mlir::IntegerAttr>())) && ((tblgen_i64_attr.cast<::mlir::IntegerAttr>().getType().isSignlessInteger(64)))))
 // DEF: if (tblgen_f32_attr && !(((tblgen_f32_attr.isa<::mlir::FloatAttr>())) && ((tblgen_f32_attr.cast<::mlir::FloatAttr>().getType().isF32()))))
 // DEF: if (tblgen_f64_attr && !(((tblgen_f64_attr.isa<::mlir::FloatAttr>())) && ((tblgen_f64_attr.cast<::mlir::FloatAttr>().getType().isF64()))))
 // DEF: if (tblgen_str_attr && !((tblgen_str_attr.isa<::mlir::StringAttr>())))
 // DEF: if (tblgen_elements_attr && !((tblgen_elements_attr.isa<::mlir::ElementsAttr>())))
 // DEF: if (tblgen_function_attr && !((tblgen_function_attr.isa<::mlir::FlatSymbolRefAttr>())))
-// DEF: if (tblgen_some_type_attr && !(((tblgen_some_type_attr.isa<::mlir::TypeAttr>())) && ((tblgen_some_type_attr.cast<::mlir::TypeAttr>().getValue().isa<SomeType>()))))
+// DEF: if (tblgen_some_type_attr && !(((tblgen_some_type_attr.isa<::mlir::TypeAttr>())) && ((tblgen_some_type_attr.cast<::mlir::TypeAttr>().getValue().isa<SomeType>())) && ((true))))
 // DEF: if (tblgen_array_attr && !((tblgen_array_attr.isa<::mlir::ArrayAttr>())))
 // DEF: if (tblgen_some_attr_array && !(((tblgen_some_attr_array.isa<::mlir::ArrayAttr>())) && (::llvm::all_of(tblgen_some_attr_array.cast<::mlir::ArrayAttr>(), [&](::mlir::Attribute attr) { return attr && ((some-condition)); }))))
-// DEF: if (tblgen_type_attr && !(((tblgen_type_attr.isa<::mlir::TypeAttr>())) && ((tblgen_type_attr.cast<::mlir::TypeAttr>().getValue().isa<::mlir::Type>()))))
+// DEF: if (tblgen_type_attr && !(((tblgen_type_attr.isa<::mlir::TypeAttr>())) && ((tblgen_type_attr.cast<::mlir::TypeAttr>().getValue().isa<::mlir::Type>())) && ((true))))
 
 // Test common attribute kind getters' return types
 // ---
 
 // DEF: ::mlir::Attribute BOp::getAnyAttr()
 // DEF: bool BOp::getBoolAttr()
 // DEF: uint32_t BOp::getI32Attr()
 // DEF: uint64_t BOp::getI64Attr()