[mlir][ods] Handle DeclareOpInterfaceMethods<InferTypeOpInterface> in formatgen

[jpienaar]

Summary

Previously it would not consider ops with DeclareOpInterfaceMethods<InferTypeOpInterface> as having the InferTypeOpInterface
interfaces added. The OpInterface nested inside DeclareOpInterfaceMethods is not retained so that one could query it, so check for the the C++ class directly (a bit raw/low level).

[mehdi_amini]

Your description does not seem correct to me, can you revise?

[jpienaar]

In D116572#3218761, @mehdi_amini wrote:

Your description does not seem correct to me, can you revise?

Seems I omitted some words indeed. Revised.

[mehdi_amini]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

2357

Should we have a method op.hasCppInterface("::mlir::InferTypeOpInterface") or something like that?

[jpienaar]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

2357

Good idea, I'll do that in follow up (I'm not crazy about the C++ names , so might consider if there is a different way to do this/some unwrapping like is done for OpTraitList).

mlir/test/lib/Dialect/Test/TestDialect.cpp

delete static_cast<TestOpEffectInterfaceFallback *>(
       fallbackEffectOpInterfaces);
 }
 
 Operation *TestDialect::materializeConstant(OpBuilder &builder, Attribute value,
                                             Type type, Location loc) {
   return builder.create<TestOpConstant>(loc, type, value);
 }
 
+::mlir::LogicalResult FormatInferType2Op::inferReturnTypes(
+    ::mlir::MLIRContext *context, ::llvm::Optional<::mlir::Location> location,
+    ::mlir::ValueRange operands, ::mlir::DictionaryAttr attributes,
+    ::mlir::RegionRange regions,
+    ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
+  inferredReturnTypes.assign({::mlir::IntegerType::get(context, 16)});
+  return ::mlir::success();
+}
+
 void *TestDialect::getRegisteredInterfaceForOp(TypeID typeID,
                                                OperationName opName) {
   if (opName.getIdentifier() == "test.unregistered_side_effect_op" &&
       typeID == TypeID::get<TestEffectOpInterface>())
     return fallbackEffectOpInterfaces;
   return nullptr;
 }
 

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

static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
           ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
           ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
       inferredReturnTypes.assign({::mlir::IntegerType::get(context, 16)});
       return ::mlir::success();
     }
    }];
 }
 
+// Check that formatget supports DeclareOpInterfaceMethods.
+def FormatInferType2Op : TEST_Op<"format_infer_type2", [DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  let results = (outs AnyType);
+  let assemblyFormat = "attr-dict";
+}
+
 // Base class for testing mixing allOperandTypes, allOperands, and
 // inferResultTypes.
 class FormatInferAllTypesBaseOp<string mnemonic, list<OpTrait> traits = []>
     : TEST_Op<mnemonic, [InferTypeOpInterface] # traits> {
   let arguments = (ins Variadic<AnyType>:$args);
   let results = (outs Variadic<AnyType>:$outs);
   let extraClassDeclaration = [{
     static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,

mlir/test/mlir-tblgen/op-format.mlir

 // CHECK: test.format_types_match_context %[[I64]] : i64
 %ignored_res6 = test.format_types_match_context %i64 : i64
 
 //===----------------------------------------------------------------------===//
 // InferTypeOpInterface type inference
 //===----------------------------------------------------------------------===//
 
 // CHECK: test.format_infer_type
-%ignored_res7 = test.format_infer_type
+%ignored_res7a = test.format_infer_type
+
+// CHECK: test.format_infer_type2
+%ignored_res7b = test.format_infer_type2
 
 // CHECK: test.format_infer_type_all_operands_and_types(%[[I64]], %[[I32]]) : i64, i32
 %ignored_res8:2 = test.format_infer_type_all_operands_and_types(%i64, %i32) : i64, i32
 
 // CHECK: test.format_infer_type_all_types_one_operand(%[[I64]], %[[I32]]) : i64, i32
 %ignored_res9:2 = test.format_infer_type_all_types_one_operand(%i64, %i32) : i64, i32
 
 // CHECK: test.format_infer_type_all_types_two_operands(%[[I64]], %[[I32]]) (%[[I64]], %[[I32]]) : i64, i32, i64, i32

mlir/tools/mlir-tblgen/OpFormatGen.cpp

if (def.isSubClassOf("AllTypesMatch")) {
       handleAllTypesMatchConstraint(def.getValueAsListOfStrings("values"),
                                     variableTyResolver);
     } else if (def.getName() == "SameTypeOperands") {
       handleSameTypesConstraint(variableTyResolver, /*includeResults=*/false);
     } else if (def.getName() == "SameOperandsAndResultType") {
       handleSameTypesConstraint(variableTyResolver, /*includeResults=*/true);
     } else if (def.isSubClassOf("TypesMatchWith")) {
       handleTypesMatchConstraint(variableTyResolver, def);
-    } else if (def.getName() == "InferTypeOpInterface" &&
-               !op.allResultTypesKnown()) {
+    } else if (!op.allResultTypesKnown()) {
+      // This doesn't check the name directly to handle
+      //    DeclareOpInterfaceMethods<InferTypeOpInterface>
+      // and the like.
+      // TODO: Add hasCppInterface check.
+      if (auto name = def.getValueAsOptionalString("cppClassName")) {
+        if (*name == "InferTypeOpInterface" &&
+            def.getValueAsString("cppNamespace") == "::mlir")
           canInferResultTypes = true;
       }
     }
+  }
 
   // Verify the state of the various operation components.
   if (failed(verifyAttributes(loc)) ||
       failed(verifyResults(loc, variableTyResolver)) ||
       failed(verifyOperands(loc, variableTyResolver)) ||
       failed(verifyRegions(loc)) || failed(verifySuccessors(loc)))
     return ::mlir::failure();