[mlir] Avoid folding in OpBuilder::tryFold when types change

[rriddle]

Summary

This was missed when tightening fold restrictions in https://reviews.llvm.org/D95991.

mlir/lib/IR/Builders.cpp

Operation *OpBuilder::createOperation(const OperationState &state) {
   return insert(Operation::create(state));
 }
 
 /// Attempts to fold the given operation and places new results within
 /// 'results'. Returns success if the operation was folded, failure otherwise.
 /// Note: This function does not erase the operation on a successful fold.
 LogicalResult OpBuilder::tryFold(Operation *op,
                                  SmallVectorImpl<Value> &results) {
-  results.reserve(op->getNumResults());
+  ResultRange opResults = op->getResults();
+
+  results.reserve(opResults.size());
   auto cleanupFailure = [&] {
-    results.assign(op->result_begin(), op->result_end());
+    results.assign(opResults.begin(), opResults.end());
     return failure();
   };
 
   // If this operation is already a constant, there is nothing to do.
   if (matchPattern(op, m_Constant()))
     return cleanupFailure();
 
   // Check to see if any operands to the operation is constant and whether
   // the operation knows how to constant fold itself.
   SmallVector<Attribute, 4> constOperands(op->getNumOperands());
-  for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i)
+  for (unsigned i = 0, e = constOperands.size(); i != e; ++i)
     matchPattern(op->getOperand(i), m_Constant(&constOperands[i]));
 
   // Try to fold the operation.
   SmallVector<OpFoldResult, 4> foldResults;
   if (failed(op->fold(constOperands, foldResults)) || foldResults.empty())
     return cleanupFailure();
 
   // A temporary builder used for creating constants during folding.
   OpBuilder cstBuilder(context);
   SmallVector<Operation *, 1> generatedConstants;
 
   // Populate the results with the folded results.
   Dialect *dialect = op->getDialect();
-  for (auto &it : llvm::enumerate(foldResults)) {
+  for (auto it : llvm::zip(foldResults, opResults.getTypes())) {
+    Type expectedType = std::get<1>(it);
+
     // Normal values get pushed back directly.
-    if (auto value = it.value().dyn_cast<Value>()) {
+    if (auto value = std::get<0>(it).dyn_cast<Value>()) {
+      if (value.getType() != expectedType)
+        return cleanupFailure();
+
       results.push_back(value);
       continue;
     }
 
     // Otherwise, try to materialize a constant operation.
     if (!dialect)
       return cleanupFailure();
 
     // Ask the dialect to materialize a constant operation for this value.
-    Attribute attr = it.value().get<Attribute>();
-    auto *constOp = dialect->materializeConstant(
-        cstBuilder, attr, op->getResult(it.index()).getType(), op->getLoc());
+    Attribute attr = std::get<0>(it).get<Attribute>();
+    auto *constOp = dialect->materializeConstant(cstBuilder, attr, expectedType,
+                                                 op->getLoc());
     if (!constOp) {
       // Erase any generated constants.
       for (Operation *cst : generatedConstants)
         cst->erase();
       return cleanupFailure();
     }
     assert(matchPattern(constOp, m_Constant()));
 

mlir/test/Transforms/test-legalizer.mlir

builtin.module {
 
   func @create_unregistered_op_in_pattern() -> i32 {
     // expected-error@+1 {{failed to legalize operation 'test.illegal_op_g'}}
     %0 = "test.illegal_op_g"() : () -> (i32)
     "test.return"(%0) : (i32) -> ()
   }
 
 }
+
+// -----
+
+// The "passthrough_fold" folder will naively return its operand, but we don't
+// want to fold here because of the type mismatch.
+func @typemismatch(%arg: f32) -> i32 {
+  // expected-remark@+1 {{op 'test.passthrough_fold' is not legalizable}}
+  %0 = "test.passthrough_fold"(%arg) : (f32) -> (i32)
+  "test.return"(%0) : (i32) -> ()
+}