WIP on InferableSameElementCompatibleShapeOperandsAndResultType

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

function_ref<LogicalResult(
         componentTypeFn,
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes);
 
 /// Verifies that the inferred result types match the actual result types for
 /// the op. Precondition: op implements InferTypeOpInterface.
 LogicalResult verifyInferredResultTypes(Operation *op);
+
+LogicalResult inferSameElementCompatibleShapeOperandsAndResultType(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes);
+LogicalResult
+verifyInferableSameElementCompatibleShapeOperandsAndResultType(Operation *op);
+bool verifyInferableSameElementCompatibleShapeOperandsAndResultType(
+    TypeRange lhs, TypeRange rhs);
+
 } // namespace detail
 
 namespace OpTrait {
 template <typename ConcreteType>
 class InferTensorType;
 } // namespace OpTrait
 } // namespace mlir
 
static_assert(
         ConcreteType::template hasTrait<InferTypeOpInterface::Trait>(),
         "requires InferTypeOpInterface to ensure succesful invocation");
     return ::mlir::detail::inferReturnTensorTypes(
         ConcreteType::inferReturnTypeComponents, context, location, operands,
         attributes, regions, inferredReturnTypes);
   }
 };
 
+/// This class provides verification for ops that are known to have the same
+/// element types but only compatible shapes, and still want some result type
+/// inference.
+template <typename ConcreteType>
+class InferableSameElementCompatibleShapeOperandsAndResultType
+    : public TraitBase<
+          ConcreteType,
+          InferableSameElementCompatibleShapeOperandsAndResultType> {
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    return ::mlir::detail::
+        verifyInferableSameElementCompatibleShapeOperandsAndResultType(op);
+  }
+  static LogicalResult
+  inferReturnTypes(MLIRContext *context, Optional<Location> location,
+                   ValueRange operands, DictionaryAttr attributes,
+                   RegionRange regions,
+                   SmallVectorImpl<Type> &inferredReturnTypes) {
+    return ::mlir::detail::inferSameElementCompatibleShapeOperandsAndResultType(
+        context, location, operands, attributes, regions, inferredReturnTypes);
+  }
+
+  static bool isCompatibleReturnTypes(TypeRange lhs, TypeRange rhs) {
+    return ::mlir::detail::
+        verifyInferableSameElementCompatibleShapeOperandsAndResultType(lhs,
+                                                                       rhs);
+  }
+};
+
 } // namespace OpTrait
 } // namespace mlir
 
 #endif // MLIR_INTERFACES_INFERTYPEOPINTERFACE_H_

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

let methods = [
     >
   ];
 }
 
 // Op has the same operand and result type.
 // TODO: Change from hard coded to utilizing type inference trait.
 def SameOperandsAndResultType : NativeOpTrait<"SameOperandsAndResultType">;
 
+// Op has the compatiable shapes for operand and result type, with result
+// type inference enabled.
+def InferableSameElementCompatibleShapeOperandsAndResultType : TraitList<[
+    NativeOpTrait<"InferableSameElementCompatibleShapeOperandsAndResultType">,
+    InferTypeOpInterface
+  ]>;
+
+
 #endif // MLIR_INFERTYPEOPINTERFACE

mlir/lib/Interfaces/InferTypeOpInterface.cpp

 //===- InferTypeOpInterface.cpp - Infer Type Interfaces ---------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the definitions of the infer op interfaces defined in
 // `InferTypeOpInterface.td`.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 
 namespace mlir {
 #include "mlir/Interfaces/InferTypeOpInterface.cpp.inc"
 } // namespace mlir
 
LogicalResult mlir::detail::verifyInferredResultTypes(Operation *op) {
   if (!retTypeFn.isCompatibleReturnTypes(inferredReturnTypes,
                                          op->getResultTypes()))
     return op->emitOpError("inferred type(s) ")
            << inferredReturnTypes
            << " are incompatible with return type(s) of operation "
            << op->getResultTypes();
   return success();
 }
+
+LogicalResult
+mlir::detail::inferSameElementCompatibleShapeOperandsAndResultType(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  inferredReturnTypes.push_back(operands.front().getType());
+  return success();
+}
+
+LogicalResult
+mlir::detail::verifyInferableSameElementCompatibleShapeOperandsAndResultType(
+    Operation *op) {
+  if (failed(OpTrait::impl::verifySameOperandsAndResultElementType(op)))
+    return failure();
+  SmallVector<Type, 8> types(op->getOperandTypes());
+  types.append(llvm::to_vector<4>(op->getResultTypes()));
+  if (failed(verifyCompatibleShapes(types)))
+    return op->emitOpError()
+           << "requires the same shape for all operands and results";
+  return success();
+}

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

 
 def OpWithInferTypeInterfaceOp : TEST_Op<"op_with_infer_type_if", [
     DeclareOpInterfaceMethods<InferTypeOpInterface,
         ["inferReturnTypeComponents"]>]> {
   let arguments = (ins AnyTensor, AnyTensor);
   let results = (outs AnyTensor);
 }
 
+def OpWithInferableSameElementCompatibleShapeOperandsAndResultType :
+  TEST_Op<"op_with_infer_type_compatible_shapes",
+          InferableSameElementCompatibleShapeOperandsAndResultType.traits> {
+  let arguments = (ins AnyTensor:$lhs, AnyTensor:$rhs);
+  let results = (outs AnyTensor);
+}
+
 def OpWithShapedTypeInferTypeInterfaceOp : TEST_Op<"op_with_shaped_type_infer_type_if",
       [InferTensorTypeWithReify]> {
   let arguments = (ins AnyTensor, AnyTensor);
   let results = (outs AnyTensor);
 }
 
 def OpWithResultShapeInterfaceOp : TEST_Op<"op_with_result_shape_interface",
       [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,

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

if (getOperation().getName() == "testCreateFunctions") {
       // Generate test patterns for each, but skip terminator.
       for (auto *op : llvm::makeArrayRef(ops).drop_back()) {
         // Test create method of each of the Op classes below. The resultant
         // output would be in reverse order underneath `op` from which
         // the attributes and regions are used.
         invokeCreateWithInferredReturnType<OpWithInferTypeInterfaceOp>(op);
         invokeCreateWithInferredReturnType<
             OpWithShapedTypeInferTypeInterfaceOp>(op);
+        invokeCreateWithInferredReturnType<
+            OpWithInferableSameElementCompatibleShapeOperandsAndResultType>(op);
       };
       return;
     }
     if (getOperation().getName() == "testReifyFunctions") {
       std::vector<Operation *> ops;
       // Collect ops to avoid triggering on inserted ops.
       for (auto &op : getOperation().getBody().front())
         if (isa<OpWithShapedTypeInferTypeInterfaceOp>(op))