[mlir] Attach InferTypeOpInterface on SameOperandsAndResultType operations when possible

[rriddle]

Summary

This allows for inferring the result types of operations in certain situations by using the type of
an operand. This commit allowed for automatically supporting type inference for many more
operations with no additional effort, e.g. nearly all Arithmetic operations now support
result type inferrence with no additional changes.

Depends on D124580

[jpienaar]

Nice, thanks

mlir/lib/TableGen/Operator.cpp

336

We could also make this an OpTraitList and avoid need to inject here. But I can do that in follow up (that special casing Todo still applies)

mlir/include/mlir/Dialect/Math/IR/Math.h

 //===- Math.h - Math dialect --------------------------------------*- C++-*-==//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MATH_IR_MATH_H_
 #define MLIR_DIALECT_MATH_IR_MATH_H_
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // Math Dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Math/IR/MathOpsDialect.h.inc"

mlir/include/mlir/Dialect/Quant/QuantOps.h

 #define MLIR_DIALECT_QUANT_QUANTOPS_H_
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "llvm/Support/MathExtras.h"
 
 #include "mlir/Dialect/Quant/QuantOpsDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Quant/QuantOps.h.inc"
 
 #endif // MLIR_DIALECT_QUANT_QUANTOPS_H_

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

 //===- SparseTensor.h - Sparse tensor dialect -------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_
 #define MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TensorEncoding.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.h.inc"
 

mlir/lib/Dialect/Quant/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRQuant
   QuantOps.cpp
   QuantTypes.cpp
   TypeDetail.h
   TypeParser.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/QuantOps/IR
 
   DEPENDS
   MLIRQuantOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
+  MLIRInferTypeOpInterface
   MLIRSideEffectInterfaces
   MLIRSupport
   )

mlir/lib/TableGen/Operator.cpp

if (getNumResults() == 0)
     return;
 
   // Skip ops with variadic or optional results.
   if (getNumVariableLengthResults() > 0)
     return;
 
   // Skip cases currently being custom generated.
   // TODO: Remove special cases.
-  if (getTrait("::mlir::OpTrait::SameOperandsAndResultType"))
+  if (getTrait("::mlir::OpTrait::SameOperandsAndResultType")) {
+    // Check for a non-variable length operand to use as the type anchor.
+    auto *operandI = llvm::find_if(arguments, [](const Argument &arg) {
+      NamedTypeConstraint *operand = arg.dyn_cast<NamedTypeConstraint *>();
+      return operand && !operand->isVariableLength();
+    });
+    if (operandI == arguments.end())
       return;
 
+    // Map each of the result types to the anchor operation.
+    int operandIdx = operandI - arguments.begin();
+    resultTypeMapping.resize(getNumResults());
+    for (int i = 0; i < getNumResults(); ++i)
+      resultTypeMapping[i].emplace_back(operandIdx);
+
+    allResultsHaveKnownTypes = true;
+    traits.push_back(Trait::create(inferTrait->getDefInit()));
+    return;
+  }
+
   // We create equivalence classes of argument/result types where arguments
   // and results are mapped into the same index space and indices corresponding
   // to the same type are in the same equivalence class.
   llvm::EquivalenceClasses<int> ecs;
   resultTypeMapping.resize(getNumResults());
   // Captures the argument whose type matches a given result type. Preference
   // towards capturing operands first before attributes.
   auto captureMapping = [&](int i) {

mlir/test/Analysis/test-shape-fn-report.mlir

 // RUN: mlir-opt %s --test-shape-function-report -verify-diagnostics
 
 module attributes {shape.lib = [@shape_lib]} {
 
 // expected-remark@+1 {{associated shape function: same_result_shape}}
 func.func @tanh(%arg: tensor<10x20xf32>) -> tensor<10x20xf32>
     attributes {shape.function = @shape_lib::@same_result_shape} {
-  // expected-remark@+1 {{no associated way}}
+  // expected-remark@+1 {{implements InferType op interface}}
   %0 = math.tanh %arg : tensor<10x20xf32>
-  // expected-remark@+1 {{associated shape function: same_result_shape}}
+  // expected-remark@+1 {{implements InferType op interface}}
   %1 = "test.same_operand_result_type"(%0) : (tensor<10x20xf32>) -> tensor<10x20xf32>
   return %1 : tensor<10x20xf32>
 }
 
 // The shape function library with some local functions.
 shape.function_library @shape_lib {
   // Test shape function that returns the shape of input arg as result shape.
   func @same_result_shape(%arg: !shape.value_shape) -> !shape.shape {
     %0 = shape_of %arg : !shape.value_shape -> !shape.shape
     return %0 : !shape.shape
   }
 } mapping {
   test.same_operand_result_type = @same_result_shape
 }
 
 }

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

static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *,
           ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
           ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
       inferredReturnTypes.assign({operands[0].getType()});
       return ::mlir::success();
     }
    }];
 }
 
-// Single variadic arg with SameOperandsAndResultType and InferTypeOpInterface.
-// Tests suppression of ambiguous build methods for operations with
-// SameOperandsAndResultType and InferTypeOpInterface.
-def TableGenBuildOp5 : TableGenBuildInferReturnTypeBaseOp<
-    "tblgen_build_5", [SameOperandsAndResultType]>;
-
 // Op with InferTypeOpInterface and regions.
-def TableGenBuildOp6 : TableGenBuildInferReturnTypeBaseOp<
-    "tblgen_build_6", [InferTypeOpInterface]> {
+def TableGenBuildOp5 : TableGenBuildInferReturnTypeBaseOp<
+    "tblgen_build_5", [InferTypeOpInterface]> {
   let regions = (region AnyRegion:$body);
 }
 
 //===----------------------------------------------------------------------===//
 // Test BufferPlacement
 //===----------------------------------------------------------------------===//
 
 def GetTupleElementOp: TEST_Op<"get_tuple_element"> {

mlir/test/mlir-tblgen/op-decl-and-defs.td

 
 // Check default value for collective params builder. Check that other builders
 // are generated as well.
 def NS_HCollectiveParamsOp : NS_Op<"op_collective_params", []> {
   let arguments = (ins AnyType:$a);
   let results = (outs AnyType:$b);
 }
 
-// CHECK_LABEL: class NS_HCollectiveParamsOp :
+// CHECK_LABEL: class HCollectiveParamsOp :
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type b, ::mlir::Value a);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a);
 // CHECK: static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {})
 
 // Check suppression of "separate arg, separate result" build method for an op
 // with single variadic arg and single variadic result (since it will be
 // ambiguous with the collective params build method).
 def NS_HCollectiveParamsSuppress0Op : NS_Op<"op_collective_suppress0", []> {
   let arguments = (ins Variadic<I32>:$a);
   let results = (outs Variadic<I32>:$b);
 }
 
-// CHECK_LABEL: class NS_HCollectiveParamsSuppress0Op :
+// CHECK_LABEL: class HCollectiveParamsSuppress0Op :
 // CHECK-NOT: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange b, ::mlir::ValueRange a);
 // CHECK: static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 
 // Check suppression of "separate arg, collective result" build method for an op
 // with single variadic arg and non variadic result (since it will be
 // ambiguous with the collective params build method).
 def NS_HCollectiveParamsSuppress1Op : NS_Op<"op_collective_suppress1", []> {
   let arguments = (ins Variadic<I32>:$a);
   let results = (outs I32:$b);
 }
 
-// CHECK_LABEL: class NS_HCollectiveParamsSuppress1Op :
+// CHECK_LABEL: class HCollectiveParamsSuppress1Op :
 // CHECK-NOT: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange b, ::mlir::ValueRange a);
 // CHECK: static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 
 // Check suppression of "separate arg, collective result" build method for an op
 // with single variadic arg and > 1 variadic result (since it will be
 // ambiguous with the collective params build method). Note that "separate arg,
 // separate result" build method should be generated in this case as its not
 // ambiguous with the collective params build method.
 def NS_HCollectiveParamsSuppress2Op : NS_Op<"op_collective_suppress2", [SameVariadicResultSize]> {
   let arguments = (ins Variadic<I32>:$a);
   let results = (outs Variadic<I32>:$b, Variadic<F32>:$c);
 }
-// CHECK_LABEL: class NS_HCollectiveParamsSuppress2Op :
+// CHECK_LABEL: class HCollectiveParamsSuppress2Op :
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange b, ::mlir::TypeRange c, ::mlir::ValueRange a);
 // CHECK-NOT: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange b, ::mlir::ValueRange a);
 // CHECK: static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 
 // Check default value of `attributes` for the `genUseOperandAsResultTypeCollectiveParamBuilder` builder
 def NS_IOp : NS_Op<"op_with_same_operands_and_result_types_trait", [SameOperandsAndResultType]> {
   let arguments = (ins AnyType:$a, AnyType:$b);
   let results = (outs AnyType:$r);
 }
-// CHECK_LABEL: class NS_IOp :
+// CHECK_LABEL: class IOp :
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::mlir::Value a, ::mlir::Value b);
+// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b);
 // CHECK: static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
-// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 
 // Check default value of `attributes` for the `genInferredTypeCollectiveParamBuilder` builder
 def NS_JOp : NS_Op<"op_with_InferTypeOpInterface_interface", [DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let arguments = (ins AnyType:$a, AnyType:$b);
   let results = (outs AnyType:$r);
 }
-// CHECK_LABEL: class NS_JOp :
+// CHECK_LABEL: class JOp :
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::mlir::Value a, ::mlir::Value b);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b);
 // CHECK: static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 
 // Test usage of TraitList getting flattened during emission.
 def NS_KOp : NS_Op<"k_op", [IsolatedFromAbove,

 // CHECK-LABEL: NS::KWithTraitOp declarations
 // CHECK:       class KWithTraitOp : public ::mlir::Op<KWithTraitOp
 // CHECK-SAME:      SomeNamespace::TestTrait
 
 def NS_LOp : NS_Op<"op_with_same_operands_and_result_types_unwrapped_attr", [SameOperandsAndResultType]> {
   let arguments = (ins AnyType:$a, AnyType:$b, I32Attr:$attr1);
   let results = (outs AnyType:$r);
 }
-// CHECK_LABEL: class NS_LOp :
+// CHECK_LABEL: class LOp :
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::mlir::Value a, ::mlir::Value b, ::mlir::IntegerAttr attr1);
+// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b, ::mlir::IntegerAttr attr1);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b, ::mlir::IntegerAttr attr1);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::mlir::Value a, ::mlir::Value b, uint32_t attr1);
+// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b, uint32_t attr1);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b, uint32_t attr1);
 // CHECK: static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
-// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b, ::mlir::IntegerAttr attr1);
-// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b, uint32_t attr1);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 
 
 // Test that type defs have the proper namespaces when used as a constraint.
 // ---
 
 def Test_Dialect2 : Dialect {
   let name = "test";

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

def OpB : NS_Op<"same_input_output_type_op", [SameOperandsAndResultType]> {
   let arguments = (ins I32:$x);
   let results = (outs I32:$y);
 }
 
 // CHECK-LABEL: OpB definitions
 // CHECK: void OpB::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type y, ::mlir::Value x)
 // CHECK:   odsState.addTypes(y);
 // CHECK: void OpB::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value x)
-// CHECK:   odsState.addTypes({x.getType()});
+// CHECK:   ::llvm::SmallVector<::mlir::Type, 2> inferredReturnTypes;
+// CHECK:   if (::mlir::succeeded(OpB::inferReturnTypes(odsBuilder.getContext(),
+// CHECK:             odsState.location, odsState.operands,
+// CHECK:             odsState.attributes.getDictionary(odsState.getContext()),
+// CHECK:             /*regions=*/{}, inferredReturnTypes)))
+// CHECK:     odsState.addTypes(inferredReturnTypes);
 
 def OpC : NS_Op<"three_normal_result_op", []> {
   let results = (outs I32:$x, /*unnamed*/I32, I32:$z);
 }
 
 // CHECK-LABEL: OpC definitions
 // CHECK:       void OpC::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type x, ::mlir::Type resultType1, ::mlir::Type z)
 // CHECK-NEXT:   odsState.addTypes(x)

mlir/unittests/TableGen/OpBuildGen.cpp

TEST_F(OpBuildGenTest,
   verifyOp(op, {i32Ty, f32Ty}, {*cstI32}, noAttrs);
 
   // Test collective params build method, with attributes.
   op = builder.create<test::TableGenBuildOp3>(loc, TypeRange{i32Ty, f32Ty},
                                               ValueRange{*cstI32}, attrs);
   verifyOp(op, {i32Ty, f32Ty}, {*cstI32}, attrs);
 }
 
-// The next 2 tests test supression of ambiguous build methods for ops that
+// The next test checks supression of ambiguous build methods for ops that
 // have a single variadic input, and single non-variadic result, and which
 // support the SameOperandsAndResultType trait and and optionally the
 // InferOpTypeInterface interface. For such ops, the ODS framework generates
 // build methods with no result types as they are inferred from the input types.
 TEST_F(OpBuildGenTest, BuildMethodsSameOperandsAndResultTypeSuppression) {
   testSingleVariadicInputInferredType<test::TableGenBuildOp4>();
 }
 
-TEST_F(
-    OpBuildGenTest,
-    BuildMethodsSameOperandsAndResultTypeAndInferOpTypeInterfaceSuppression) {
-  testSingleVariadicInputInferredType<test::TableGenBuildOp5>();
-}
-
 TEST_F(OpBuildGenTest, BuildMethodsRegionsAndInferredType) {
-  auto op = builder.create<test::TableGenBuildOp6>(
+  auto op = builder.create<test::TableGenBuildOp5>(
       loc, ValueRange{*cstI32, *cstF32}, /*attributes=*/noAttrs);
   ASSERT_EQ(op->getNumRegions(), 1u);
   verifyOp(op, {i32Ty}, {*cstI32, *cstF32}, noAttrs);
 }
 
 } // namespace mlir