[mlir] Enable specifying verify on OpInterface

[jpienaar]

Summary

Add method in ODS to specify verification for operations implementing a
OpInterface. Use this with infer type op interface to verify that the
inferred type matches the return type and remove special case in
TestPatterns.

This could also have been achieved by using OpInterfaceMethod but verify
seems pretty common and it is not an arbitrary method that just happened
to be named verifyTrait, so having it be defined in special way seems
appropriate/better documenting.

[merge_guards_bot]

Unit tests: pass. 62059 tests passed, 0 failed and 784 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

mlir/include/mlir/Analysis/InferTypeOpInterface.td

65

Is '$_op' documented somewhere?

mlir/include/mlir/IR/OpBase.td

1417

Why the underscore?

1445

Why the newlines?

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

160

nit: verify -> 'verify' field

[mehdi_amini]

Nice!

I was talking to someone last week about this, I just don't remember who :)

mlir/include/mlir/IR/OpBase.td

1446

spurious?

[jpienaar]

Quoted 'verify', removed spurious newlines.

[jpienaar]

Updated thanks

mlir/include/mlir/Analysis/InferTypeOpInterface.td

65

Yes, in OpBase and this matches placeholder we use in other code body replacements.

mlir/include/mlir/IR/OpBase.td

1417

To match the same placeholder in line 79. I was tempted to just document that op will refer to the op, but given we already have those placeholders, kept it consistent.

[merge_guards_bot]

Unit tests: pass. 62095 tests passed, 0 failed and 785 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

mlir/include/mlir/IR/OpBase.td

1417

This isn't consistent with the rest of the OpInterface methods though.

mlir/docs/OpDefinitions.md

 // Interfaces can optionally be wrapped inside DeclareOpInterfaceMethods. This
 // would result in autogenerating declarations for members `foo`, `bar` and
 // `fooStatic`. Methods with bodies are not declared inside the op
 // declaration but instead handled by the op interface trait directly.
 def OpWithInferTypeInterfaceOp : Op<...
     [DeclareOpInterfaceMethods<MyInterface>]> { ... }
 ```
 
+A verification method can also be specified on the `OpInterface` by setting
+`verify`. Setting `verify` results in the generated trait having a `verifyTrait`
+method that is applied to all operations implementing the trait.
+
 ### Builder methods
 
 For each operation, there are a few builders automatically generated based on
 the arguments and returns types. For example, given the following op definition:
 
 ```tablegen
 def MyOp : ... {
   let arguments = (ins

mlir/include/mlir/Analysis/InferTypeOpInterface.h

 
 private:
   ShapeStorageT dims;
   bool ranked;
   Type elementType;
   Attribute attr;
 };
 
-#include "mlir/Analysis/InferTypeOpInterface.h.inc"
-
 namespace detail {
 // Helper function to infer return tensor returns types given element and shape
 // inference function.
 //
 // TODO: Consider generating typedefs for trait member functions if this usage
 // becomes more common.
 LogicalResult inferReturnTensorTypes(
     function_ref<LogicalResult(
         MLIRContext *, Optional<Location> location, ValueRange operands,
         ArrayRef<NamedAttribute> attributes, RegionRange regions,
         SmallVectorImpl<ShapedTypeComponents> &retComponents)>
         componentTypeFn,
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     ArrayRef<NamedAttribute> attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferedReturnTypes);
+
+/// Verifies that the inferred result types match the actual result types for
+/// the op. Precondition: op implements InferTypeOpInterface.
+LogicalResult verifyInferredResultTypes(Operation *op);
 } // namespace detail
 
+#include "mlir/Analysis/InferTypeOpInterface.h.inc"
+
 namespace OpTrait {
 
 /// Tensor type inference trait that constructs a tensor from the infered
 /// shape and elemental types.
 /// Requires: Op implements functions of InferShapedTypeOpInterface.
 template <typename ConcreteType>
 class InferTensorType : public TraitBase<ConcreteType, InferTensorType> {
 public:

mlir/include/mlir/Analysis/InferTypeOpInterface.td

StaticInterfaceMethod<
       }],
       /*defaultImplementation=*/[{
         /// Returns whether two arrays are equal as strongest check for
         /// compatibility by default.
         return lhs == rhs;
       }]
     >,
   ];
+
+  let verify = [{
+    return detail::verifyInferredResultTypes($_op);
+  }];
 }
 
 def InferShapedTypeOpInterface : OpInterface<"InferShapedTypeOpInterface"> {
   let description = [{
     Interface to infer the components of a ShapedType returned by an operation
     that could be used during op construction, verification or shape inference.
 
     The components consists of element type, shape and raw attribute.

mlir/include/mlir/IR/OpBase.td

 //===----------------------------------------------------------------------===//
 
 // Marker used to identify the argument list for an op or interface method.
 def ins;
 
 // OpInterfaceTrait corresponds to a specific 'OpInterface' class defined in
 // C++. The purpose to wrap around C++ symbol string with this class is to make
 // interfaces specified for ops in TableGen less alien and more integrated.
-class OpInterfaceTrait<string name> : NativeOpTrait<""> {
+class OpInterfaceTrait<string name, code verifyBody = [{}]> : NativeOpTrait<""> {
   let trait = name # "::Trait";
+
+  // Specify the body of the verification function. `$_op` will be replaced with
+  // the operation being verified.
+  code verify = verifyBody;
 }
 
 // This class represents a single, optionally static, interface method.
 // Note: non-static interface methods have an implicit 'op' parameter
 // corresponding to an instance of the derived operation.
 class InterfaceMethod<string desc, string retTy, string methodName,
                       dag args = (ins), code methodBody = [{}],
                       code defaultImplementation = [{}]> {
class InterfaceMethod<string desc, string retTy, string methodName,
   // A dag of string that correspond to the arguments of the method.
   dag arguments = args;
 
   // An optional body to the method.
   code body = methodBody;
 
   // An optional default implementation of the method.
   code defaultBody = defaultImplementation;
 }
 
 // This class represents a single static interface method.
 class StaticInterfaceMethod<string desc, string retTy, string methodName,
                             dag args = (ins), code methodBody = [{}],
                             code defaultImplementation = [{}]>
     : InterfaceMethod<desc, retTy, methodName, args, methodBody,
                       defaultImplementation>;
 
 // OpInterface represents an interface regarding an op.

mlir/include/mlir/TableGen/OpInterfaces.h

public:
   StringRef getName() const;
 
   // Return the methods of this interface.
   ArrayRef<OpInterfaceMethod> getMethods() const;
 
   // Return the description of this method if it has one.
   llvm::Optional<StringRef> getDescription() const;
 
+  // Return the verify method body if it has one.
+  llvm::Optional<StringRef> getVerify() const;
+
 private:
   // The TableGen definition of this interface.
   const llvm::Record *def;
 
   // The methods of this interface.
   SmallVector<OpInterfaceMethod, 8> methods;
 };
 
 } // end namespace tblgen
 } // end namespace mlir
 
 #endif // MLIR_TABLEGEN_OPINTERFACES_H_

mlir/lib/Analysis/InferTypeOpInterface.cpp

if (shapeAndType.hasRank())
       inferedReturnTypes.push_back(RankedTensorType::get(
           shapeAndType.getDims(), shapeAndType.getElementType()));
     else
       inferedReturnTypes.push_back(
           UnrankedTensorType::get(shapeAndType.getElementType()));
   }
   return success();
 }
+
+LogicalResult mlir::detail::verifyInferredResultTypes(Operation *op) {
+  SmallVector<Type, 4> inferedReturnTypes;
+  auto retTypeFn = cast<InferTypeOpInterface>(op);
+  if (failed(retTypeFn.inferReturnTypes(op->getContext(), op->getLoc(),
+                                        op->getOperands(), op->getAttrs(),
+                                        op->getRegions(), inferedReturnTypes)))
+    return failure();
+  SmallVector<Type, 4> resultTypes(op->getResultTypes());
+  if (!retTypeFn.isCompatibleReturnTypes(inferedReturnTypes, resultTypes))
+    return op->emitOpError(
+        "inferred type incompatible with return type of operation");
+  return success();
+}

mlir/lib/TableGen/OpInterfaces.cpp

 // Return the methods of this interface.
 ArrayRef<OpInterfaceMethod> OpInterface::getMethods() const { return methods; }
 
 // Return the description of this method if it has one.
 llvm::Optional<StringRef> OpInterface::getDescription() const {
   auto value = def->getValueAsString("description");
   return value.empty() ? llvm::Optional<StringRef>() : value;
 }
+
+// Return the body for this method if it has one.
+llvm::Optional<StringRef> OpInterface::getVerify() const {
+  auto value = def->getValueAsString("verify");
+  return value.empty() ? llvm::Optional<StringRef>() : value;
+}

mlir/test/lib/TestDialect/TestPatterns.cpp

if (getFunction().getName() == "testCreateFunctions") {
         // output would be in reverse order underneath `op` from which
         // the attributes and regions are used.
         invokeCreateWithInferedReturnType<OpWithInferTypeInterfaceOp>(op);
         invokeCreateWithInferedReturnType<OpWithShapedTypeInferTypeInterfaceOp>(
             op);
       };
       return;
     }
-
-    // Verification check.
-    // TODO: Move to ops that implement type infer interface.
-    getFunction().walk([this](Operation *op) -> void {
-      auto retTypeFn = dyn_cast<InferTypeOpInterface>(op);
-      if (!retTypeFn)
-        return;
-      auto *context = &getContext();
-      SmallVector<Type, 2> inferedReturnTypes;
-      if (failed(retTypeFn.inferReturnTypes(
-              context, op->getLoc(), op->getOperands(), op->getAttrs(),
-              op->getRegions(), inferedReturnTypes)))
-        return;
-      SmallVector<Type, 1> resultTypes(op->getResultTypes());
-      if (!retTypeFn.isCompatibleReturnTypes(inferedReturnTypes, resultTypes)) {
-        op->emitOpError(
-            "inferred type incompatible with return type of operation");
-        return;
-      }
-    });
   }
 };
 } // end anonymous namespace
 
 static mlir::PassRegistration<TestReturnTypeDriver>
     rt_pass("test-return-type", "Run return type functions");
 
 //===----------------------------------------------------------------------===//

mlir/test/mlir-tblgen/return-types.mlir

 // CHECK-SAME: (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
 // CHECK: "test.op_with_infer_type_if"
 // CHECK-SAME: (tensor<20xi32>, tensor<20xi32>) -> tensor<20xi32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: testReturnTypeOpInterface
 func @testReturnTypeOpInterface(%arg0 : tensor<10xf32>) {
   // expected-error@+1 {{incompatible with return type}}
   %bad = "test.op_with_infer_type_if"(%arg0, %arg0) : (tensor<10xf32>, tensor<10xf32>) -> tensor<*xf32>
   return
 }
 
 // -----
 
-// CHECK-LABEL: testReturnTypeOpInterface
 func @testReturnTypeOpInterfaceMismatch(%arg0 : tensor<10xf32>, %arg1 : tensor<20xf32>) {
   // expected-error@+1 {{operand type mismatch}}
   %bad = "test.op_with_infer_type_if"(%arg0, %arg1) : (tensor<10xf32>, tensor<20xf32>) -> tensor<*xf32>
   return
 }

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

 //===- OpInterfacesGen.cpp - MLIR op interface utility generator ----------===//
 //
 // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // OpInterfacesGen generates definitions for operation interfaces.
 //
 //===----------------------------------------------------------------------===//
 
 #include "DocGenUtilities.h"
 #include "mlir/Support/STLExtras.h"
+#include "mlir/TableGen/Format.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "mlir/TableGen/OpInterfaces.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
static void emitTraitDecl(OpInterface &interface, raw_ostream &os,
                           StringRef interfaceTraitsName) {
   os << "  template <typename ConcreteOp>\n  "
      << llvm::formatv("struct Trait : public OpInterface<{0},"
                       " detail::{1}>::Trait<ConcreteOp> {{\n",
                       interfaceName, interfaceTraitsName);
 
   // Insert the default implementation for any methods.
   for (auto &method : interface.getMethods()) {
+    // Flag interface methods named verifyTrait.
+    if (method.getName() == "verifyTrait")
+      PrintFatalError(
+          formatv("'verifyTrait' method cannot be specified as interface "
+                  "method for '{0}'; set 'verify' on OpInterfaceTrait instead",
+                  interfaceName));
     auto defaultImpl = method.getDefaultImplementation();
     if (!defaultImpl)
       continue;
 
     os << "  " << (method.isStatic() ? "static " : "") << method.getReturnType()
        << " ";
     emitMethodNameAndArgs(method, os, /*addOperationArg=*/false);
     os << " {\n" << defaultImpl.getValue() << "  }\n";
   }
 
+  tblgen::FmtContext traitCtx;
+  traitCtx.withOp("op");
+  if (auto verify = interface.getVerify()) {
+    os << "  static LogicalResult verifyTrait(Operation* op) {\n"
+       << tblgen::tgfmt(*verify, &traitCtx) << "\n  }\n";
+  }
+
   os << "  };\n";
 }
 
 static void emitInterfaceDecl(OpInterface &interface, raw_ostream &os) {
   StringRef interfaceName = interface.getName();
   auto interfaceTraitsName = (interfaceName + "InterfaceTraits").str();
 
   // Emit the traits struct containing the concept and model declarations.