[mlir] Add VerifyOpBeforeTraits trait

[springerm]

Summary

Add a new "marker" trait that allows users to specify the order in which ops are verified.

By default: verify traits, then the op
When VerifyOpBeforeTraits is attached: verify op, then the traits

VerifyOpBeforeTraits is useful when an interface verifier calls an InterfaceMethod. Without VerifyOpBeforeTraits, the op was not verified yet (and could be broken), so calling an InterfaceMethod is dangerous and could crash.

[nicolasvasilache]

Very nice, this solves some very concrete problems!

LGTM but please wait for @rriddle and @mehdi_amini to weigh in.

[rriddle]

This doesn't really feel like the right direction. It's really fragile to flip verification either before or after the operation. We often have "core" traits that we expect to be verified before we ever get to the operations "verify" method. A better path forward would be evolving the trait verification model, which there have been discussions of in the past on discourse (can't find the thread right now).

[mehdi_amini]

Agree with River: this seems like a workaround to me.
Here is the previous thinking about this: https://discourse.llvm.org/t/rfc-verification-order-in-mlir/4718 ; can you position your issue with respect to this?

mlir/include/mlir/IR/OpBase.td

 // Similar to AttrSizedOperandSegments, but used for results. The attribute
 // should be named as `result_segment_sizes`.
 def AttrSizedResultSegments  :
   NativeOpTrait<"AttrSizedResultSegments">, StructuralOpTrait;
 
 // Op attached regions have no arguments
 def NoRegionArguments : NativeOpTrait<"NoRegionArguments">, StructuralOpTrait;
 
+// Indicates that an op should be verified before its traits.
+def VerifyOpBeforeTraits : NativeOpTrait<"VerifyOpBeforeTraits">;
+
 //===----------------------------------------------------------------------===//
 // OpInterface definitions
 //===----------------------------------------------------------------------===//
 
 // Marker used to identify the argument list for an op or interface method.
 def ins;
 
 // This class represents a typed argument with optional default value for C

mlir/include/mlir/IR/OpDefinition.h

struct Tensorizable : public TraitBase<ConcreteType, Tensorizable> {
   }
 };
 
 /// Together, `Elementwise`, `Scalarizable`, `Vectorizable`, and `Tensorizable`
 /// provide an easy way for scalar operations to conveniently generalize their
 /// behavior to vectors/tensors, and systematize conversion between these forms.
 bool hasElementwiseMappableTraits(Operation *op);
 
+/// The presence of this trait indicates an op should be verified before its
+/// traits.
+template <typename ConcreteType>
+struct VerifyOpBeforeTraits
+    : public TraitBase<ConcreteType, VerifyOpBeforeTraits> {};
+
 } // namespace OpTrait
 
 //===----------------------------------------------------------------------===//
 // Internal Trait Utilities
 //===----------------------------------------------------------------------===//
 
 namespace op_definition_impl {
 //===----------------------------------------------------------------------===//
private:
   /// Implementation of `PopulateDefaultAttrsFn` OperationName hook.
   static OperationName::PopulateDefaultAttrsFn getPopulateDefaultAttrsFn() {
     return ConcreteType::populateDefaultAttrs;
   }
   /// Implementation of `VerifyInvariantsFn` OperationName hook.
   static LogicalResult verifyInvariants(Operation *op) {
     static_assert(hasNoDataMembers(),
                   "Op class shouldn't define new data members");
+    if (ConcreteType::template hasTrait<OpTrait::VerifyOpBeforeTraits>())
+      return failure(
+          failed(cast<ConcreteType>(op).verify()) ||
+          failed(
+              op_definition_impl::verifyTraits<Traits<ConcreteType>...>(op)));
     return failure(
         failed(op_definition_impl::verifyTraits<Traits<ConcreteType>...>(op)) ||
         failed(cast<ConcreteType>(op).verify()));
   }
   static OperationName::VerifyInvariantsFn getVerifyInvariantsFn() {
     return static_cast<LogicalResult (*)(Operation *)>(&verifyInvariants);
   }
   /// Implementation of `VerifyRegionInvariantsFn` OperationName hook.
   static LogicalResult verifyRegionInvariants(Operation *op) {
     static_assert(hasNoDataMembers(),
                   "Op class shouldn't define new data members");
+    if (ConcreteType::template hasTrait<OpTrait::VerifyOpBeforeTraits>())
+      return failure(
+          failed(cast<ConcreteType>(op).verifyRegions()) ||
+          failed(
+              op_definition_impl::verifyRegionTraits<Traits<ConcreteType>...>(
+                  op)));
     return failure(
         failed(op_definition_impl::verifyRegionTraits<Traits<ConcreteType>...>(
             op)) ||
         failed(cast<ConcreteType>(op).verifyRegions()));
   }
   static OperationName::VerifyRegionInvariantsFn getVerifyRegionInvariantsFn() {
     return static_cast<LogicalResult (*)(Operation *)>(&verifyRegionInvariants);
   }

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

 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 namespace mlir {
 class DLTIDialect;
 class RewritePatternSet;
+namespace detail {
+LogicalResult verifyTestPositiveValueOpInterface(Operation *op);
+} // namespace detail
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
 #include "TestOpInterfaces.h.inc"
 #include "TestOpsDialect.h.inc"

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

void TestReflectBoundsOp::inferResultRanges(
   Builder b(ctx);
   setUminAttr(b.getIndexAttr(range.umin().getZExtValue()));
   setUmaxAttr(b.getIndexAttr(range.umax().getZExtValue()));
   setSminAttr(b.getIndexAttr(range.smin().getSExtValue()));
   setSmaxAttr(b.getIndexAttr(range.smax().getSExtValue()));
   setResultRanges(getResult(), range);
 }
 
+//===----------------------------------------------------------------------===//
+// PositiveIntAttrWithInterfaceOp
+//===----------------------------------------------------------------------===//
+
+int64_t PositiveIntAttrWithInterfaceOp::getIntValue() {
+  return getIntAttrAttr().cast<IntegerAttr>().getInt();
+}
+
+LogicalResult PositiveIntAttrWithInterfaceOp::verify() {
+  return success(isa<IntegerAttr>(getIntAttrAttr()));
+}
+
+//===----------------------------------------------------------------------===//
+// TestPositiveValueOpInterface
+//===----------------------------------------------------------------------===//
+
+LogicalResult mlir::detail::verifyTestPositiveValueOpInterface(Operation *op) {
+  return success(cast<mlir::TestPositiveValueOpInterface>(op).getIntValue() >
+                 0);
+}
+
 #include "TestOpEnums.cpp.inc"
 #include "TestOpInterfaces.cpp.inc"
 #include "TestTypeInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

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

 class TestEffect<string effectName>
     : SideEffect<TestEffectOpInterface, effectName, DefaultResource>;
 
 class TestEffects<list<TestEffect> effects = []>
    : SideEffectsTraitBase<TestEffectOpInterface, effects>;
 
 def TestConcreteEffect : TestEffect<"TestEffects::Concrete">;
 
+def TestPositiveValueOpInterface : OpInterface<"TestPositiveValueOpInterface"> {
+  let cppNamespace = "::mlir";
+  let methods = [
+    InterfaceMethod<"Returns the int value associated with this op.",
+      "int64_t", "getIntValue", (ins)>,
+  ];
+  let verify = [{ return detail::verifyTestPositiveValueOpInterface($_op); }];
+}
+
 #endif // MLIR_TEST_DIALECT_TEST_INTERFACES

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

 def FloatAttrOp : TEST_Op<"float_attrs"> {
   // TODO: Clean up the OpBase float type and attribute selectors so they
   // can express all of the types.
   let arguments = (ins
     AnyAttr:$float_attr
   );
 }
 
+def PositiveIntAttrWithInterfaceOp : TEST_Op<"positive_int_attr_with_interface",
+    [DeclareOpInterfaceMethods<TestPositiveValueOpInterface>,
+     VerifyOpBeforeTraits]> {
+  // This could be an I64Attr. But we check this in the hand-written verifier
+  // to test VerifyOpBeforeTraits. The hand-written op verifier should be
+  // called before the VerifyOpBeforeTraits verifier.
+  let arguments = (ins AnyAttr:$int_attr);
+  let hasVerifier = 1;
+}
+
 def I32EnumAttrOp : TEST_Op<"i32_enum_attr"> {
   let arguments = (ins SomeI32Enum:$attr);
   let results = (outs I32:$val);
 }
 
 def I64EnumAttrOp : TEST_Op<"i64_enum_attr"> {
   let arguments = (ins SomeI64Enum:$attr);
   let results = (outs I32:$val);

mlir/unittests/IR/InterfaceTest.cpp

 //===- InterfaceTest.cpp - Test interfaces --------------------------------===//
 //
 // Part of the LLVM 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OwningOpRef.h"
+#include "mlir/IR/Verifier.h"
 #include "gtest/gtest.h"
 
 #include "../../test/lib/Dialect/Test/TestAttributes.h"
 #include "../../test/lib/Dialect/Test/TestDialect.h"
 #include "../../test/lib/Dialect/Test/TestTypes.h"
 
 using namespace mlir;
 using namespace test;
TEST(InterfaceTest, TypeInterfaceDenseMapKey) {
   auto type3 = builder.getType<test::TestTypeWithLayoutType>(3);
   typeSet.insert(type1);
   typeSet.insert(type2);
   typeSet.erase(type1);
   EXPECT_FALSE(typeSet.contains(type1));
   EXPECT_TRUE(typeSet.contains(type2));
   EXPECT_FALSE(typeSet.contains(type3));
 }
+
+TEST(InterfaceTest, VerifyOpFirst) {
+  MLIRContext context;
+  context.loadDialect<test::TestDialect>();
+
+  // Build an invalid op.
+  OpBuilder builder(&context);
+  Location loc = builder.getUnknownLoc();
+  SmallVector<NamedAttribute> attrs;
+  attrs.push_back(NamedAttribute(builder.getStringAttr("int_attr"),
+                                 builder.getStringAttr("foo")));
+  Operation *op = builder.create(
+      loc,
+      builder.getStringAttr(PositiveIntAttrWithInterfaceOp::getOperationName()),
+      /*operands=*/{}, /*types=*/{}, attrs);
+  // The verifier fails but it does not crash. (It would crash if we were to
+  // verify the interface first.)
+  LogicalResult status = verify(op);
+  EXPECT_FALSE(succeeded(status));
+}