Make the folder more robust against op fold() methods that generate a type mismatch

[mehdi_amini]

Summary

We could extend this with an interface to allow dialect to perform a type
conversion, but that would make the folder creating operation which isn't
the case at the moment, and isn't necessarily always desirable.

mlir/lib/Transforms/Utils/FoldUtils.cpp

LogicalResult OperationFolder::tryToFold(
 
   // Create the result constants and replace the results.
   auto *dialect = op->getDialect();
   for (unsigned i = 0, e = op->getNumResults(); i != e; ++i) {
     assert(!foldResults[i].isNull() && "expected valid OpFoldResult");
 
     // Check if the result was an SSA value.
     if (auto repl = foldResults[i].dyn_cast<Value>()) {
+      if (repl.getType() != op->getResult(i).getType())
+        return failure();
       results.emplace_back(repl);
       continue;
     }
 
     // Check to see if there is a canonicalized version of this constant.
     auto res = op->getResult(i);
     Attribute attrRepl = foldResults[i].get<Attribute>();
     if (auto *constOp =

mlir/test/Transforms/test-canonicalize.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='func(canonicalize)' | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline='func(canonicalize)' | FileCheck %s
 
 // CHECK-LABEL: func @remove_op_with_inner_ops_pattern
 func @remove_op_with_inner_ops_pattern() {
   // CHECK-NEXT: return
   "test.op_with_region_pattern"() ({
-    "foo.op_with_region_terminator"() : () -> ()
+    "test.op_with_region_terminator"() : () -> ()
   }) : () -> ()
   return
 }
 
 // CHECK-LABEL: func @remove_op_with_inner_ops_fold_no_side_effect
 func @remove_op_with_inner_ops_fold_no_side_effect() {
   // CHECK-NEXT: return
   "test.op_with_region_fold_no_side_effect"() ({
-    "foo.op_with_region_terminator"() : () -> ()
+    "test.op_with_region_terminator"() : () -> ()
   }) : () -> ()
   return
 }
 
 // CHECK-LABEL: func @remove_op_with_inner_ops_fold
 // CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32)
 func @remove_op_with_inner_ops_fold(%arg0 : i32) -> (i32) {
   // CHECK-NEXT: return %[[ARG_0]]
   %0 = "test.op_with_region_fold"(%arg0) ({
-    "foo.op_with_region_terminator"() : () -> ()
+    "test.op_with_region_terminator"() : () -> ()
   }) : (i32) -> (i32)
   return %0 : i32
 }
 
 // CHECK-LABEL: func @remove_op_with_variadic_results_and_folder
 // CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: i32, %[[ARG_1:[a-z0-9]*]]: i32)
 func @remove_op_with_variadic_results_and_folder(%arg0 : i32, %arg1 : i32) -> (i32, i32) {
   // CHECK-NEXT: return %[[ARG_0]], %[[ARG_1]]
func @test_commutative_multi(%arg0: i32, %arg1: i32) -> (i32, i32) {
   // CHECK-NEXT: %[[O0:.*]] = "test.op_commutative"(%[[ARG_0]], %[[ARG_1]], %[[C42]], %[[C43]]) : (i32, i32, i32, i32) -> i32
   %y = "test.op_commutative"(%c42_i32, %arg0, %arg1, %c43_i32) : (i32, i32, i32, i32) -> i32
 
   // CHECK-NEXT: %[[O1:.*]] = "test.op_commutative"(%[[ARG_0]], %[[ARG_1]], %[[C42]], %[[C43]]) : (i32, i32, i32, i32) -> i32
   %z = "test.op_commutative"(%arg0, %c42_i32, %c43_i32, %arg1): (i32, i32, i32, i32) -> i32
   // CHECK-NEXT: return %[[O0]], %[[O1]]
   return %y, %z: i32, i32
 }
+
+func @typemismatch() -> i32 {
+  %c42 = constant 42.0 : f32
+
+  // The "passthrough_fold" folder will naively return its operand, but we don't
+  // want to fold here because of the type mismatch.
+
+  // CHECK: "test.passthrough_fold"
+  %0 = "test.passthrough_fold"(%c42) : (f32) -> (i32)
+  return %0 : i32
+}

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

OpFoldResult TestOpInPlaceFold::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 1);
   if (operands.front()) {
     (*this)->setAttr("attr", operands.front());
     return getResult();
   }
   return {};
 }
 
+OpFoldResult TestPassthroughFold::fold(ArrayRef<Attribute> operands) {
+  return getOperand();
+}
+
 LogicalResult OpWithInferTypeInterfaceOp::inferReturnTypes(
     MLIRContext *, Optional<Location> location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
   if (operands[0].getType() != operands[1].getType()) {
     return emitOptionalError(location, "operand type mismatch ",
                              operands[0].getType(), " vs ",
                              operands[1].getType());

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

 }
 
 def TestOpInPlaceFold : TEST_Op<"op_in_place_fold"> {
   let arguments = (ins I32:$op, I32Attr:$attr);
   let results = (outs I32);
   let hasFolder = 1;
 }
 
+// An op that always fold itself.
+def TestPassthroughFold : TEST_Op<"passthrough_fold"> {
+  let arguments = (ins AnyType:$op);
+  let results = (outs AnyType);
+  let hasFolder = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Test Patterns (Symbol Binding)
 
 // Test symbol binding.
 def OpSymbolBindingA : TEST_Op<"symbol_binding_a", []> {
   let arguments = (ins I32:$operand, I64Attr:$attr);
   let results = (outs I32);
 }