[mlir] Add ResultRange::setTypes helper.

[silvas]

Summary

This makes some code much clearer. The corresponding helpers on
MultiResultTraitBase and Operation are added too.

Also, add a TODO about non-1:1 type conversions.

[silvas]

update

[rriddle]

Is this common enough to warrant an addition to the API? I looked at upstream usages as well as several downstream projects and found very few usages of setType. I'm mostly apprehensive given the relative danger that setType has.

[silvas]

In D89591#2340243, @rriddle wrote:

Is this common enough to warrant an addition to the API? I looked at upstream usages as well as several downstream projects and found very few usages of setType. I'm mostly apprehensive given the relative danger that setType has.

You're right. It's probably not that common. Abandoning this for now, especially since I abandoned the revision that originally motivated this.

mlir/include/mlir/IR/OpDefinition.h

result_type_iterator result_type_begin() {
     return this->getOperation()->result_type_begin();
   }
   result_type_iterator result_type_end() {
     return this->getOperation()->result_type_end();
   }
   result_type_range getResultTypes() {
     return this->getOperation()->getResultTypes();
   }
+
+  /// Sets the types of all results to `types`.
+  ///
+  /// Requires that the number of elements of `types` matches the number of
+  /// results.
+  void setResultTypes(TypeRange types) {
+    this->getOperation()->setResultTypes(types);
+  }
 };
 } // end namespace detail
 
 /// This class provides return value APIs for ops that are known to have a
 /// single result.
 template <typename ConcreteType>
 class OneResult : public TraitBase<ConcreteType, OneResult> {
 public:

mlir/include/mlir/IR/Operation.h

public:
 
   /// Support result type iteration.
   using result_type_iterator = result_range::type_iterator;
   using result_type_range = result_range::type_range;
   result_type_iterator result_type_begin() { return getResultTypes().begin(); }
   result_type_iterator result_type_end() { return getResultTypes().end(); }
   result_type_range getResultTypes();
 
+  /// Sets the types of all results to `types`.
+  ///
+  /// Requires that the number of elements of `types` matches the number of
+  /// results.
+  void setResultTypes(TypeRange types) { getResults().setTypes(types); }
+
   //===--------------------------------------------------------------------===//
   // Attributes
   //===--------------------------------------------------------------------===//
 
   // Operations may optionally carry a list of attributes that associate
   // constants to names.  Attributes may be dynamically added and removed over
   // the lifetime of an operation.
 

mlir/include/mlir/IR/OperationSupport.h

using indexed_accessor_range<ResultRange, Operation *, OpResult, OpResult,
                                OpResult>::indexed_accessor_range;
   ResultRange(Operation *op);
 
   /// Returns the types of the values within this range.
   using type_iterator = ArrayRef<Type>::iterator;
   using type_range = ArrayRef<Type>;
   type_range getTypes() const;
   auto getType() const { return getTypes(); }
+  /// Sets the types of values in this range to the given `types`.
+  ///
+  /// Requires that the number of elements of `types` matches the number of
+  /// elements in this range.
+  void setTypes(TypeRange types);
 
 private:
   /// See `llvm::indexed_accessor_range` for details.
   static OpResult dereference(Operation *op, ptrdiff_t index);
 
   /// Allow access to `dereference_iterator`.
   friend llvm::indexed_accessor_range<ResultRange, Operation *, OpResult,
                                       OpResult, OpResult>;

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

matchAndRewrite(ForOp op, ArrayRef<Value> operands,
     SmallVector<Type, 6> newResultTypes;
     for (auto type : op.getResultTypes()) {
       Type newType = typeConverter->convertType(type);
       if (!newType)
         return rewriter.notifyMatchFailure(op, "not a 1:1 type conversion");
       newResultTypes.push_back(newType);
     }
     rewriter.updateRootInPlace(op, [&] {
-      for (auto t : llvm::zip(op.getResults(), newResultTypes))
-        std::get<0>(t).setType(std::get<1>(t));
+      op.setResultTypes(newResultTypes);
       auto bodyArgs = op.getBody()->getArguments();
       for (auto t : llvm::zip(llvm::drop_begin(bodyArgs, 1), newResultTypes))
         std::get<0>(t).setType(std::get<1>(t));
     });
     return success();
   }
 };
 } // namespace
 
 namespace {
 class ConvertIfOpTypes : public OpConversionPattern<IfOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(IfOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
+    // TODO: Generalize this to any type conversion, not just 1:1.
+    //
+    // We need to implement something more sophisticated here that tracks which
+    // types convert to which other types and does the appropriate
+    // materialization logic.
+    // For example, it's possible that one result type converts to 0 types and
+    // another to 2 types, so newResultTypes would at least be the right size to
+    // not crash in the setResultTypes call, but then we would set the the wrong
+    // type on the SSA values! These edge cases are also why we cannot safely
+    // use the TypeConverter::convertTypes helper here.
     SmallVector<Type, 6> newResultTypes;
     for (auto type : op.getResultTypes()) {
       Type newType = typeConverter->convertType(type);
       if (!newType)
         return rewriter.notifyMatchFailure(op, "not a 1:1 type conversion");
       newResultTypes.push_back(newType);
     }
-    rewriter.updateRootInPlace(op, [&] {
-      for (auto t : llvm::zip(op.getResults(), newResultTypes))
-        std::get<0>(t).setType(std::get<1>(t));
-    });
+    rewriter.updateRootInPlace(op, [&] { op.setResultTypes(newResultTypes); });
     return success();
   }
 };
 } // namespace
 
 void mlir::scf::populateSCFStructuralTypeConversions(
     MLIRContext *context, TypeConverter &typeConverter,
     OwningRewritePatternList &patterns) {
   patterns.insert<ConvertIfOpTypes>(typeConverter, context);
   patterns.insert<ConvertForOpTypes>(typeConverter, context);
 }

mlir/lib/IR/OperationSupport.cpp

 
 ResultRange::ResultRange(Operation *op)
     : ResultRange(op, /*startIndex=*/0, op->getNumResults()) {}
 
 ArrayRef<Type> ResultRange::getTypes() const {
   return getBase()->getResultTypes().slice(getStartIndex(), size());
 }
 
+void ResultRange::setTypes(TypeRange types) {
+  assert(types.size() == size() && "mismatch in number of types!");
+  for (auto t : llvm::zip(*this, types))
+    std::get<0>(t).setType(std::get<1>(t));
+}
+
 /// See `llvm::indexed_accessor_range` for details.
 OpResult ResultRange::dereference(Operation *op, ptrdiff_t index) {
   return op->getResult(index);
 }
 
 //===----------------------------------------------------------------------===//
 // ValueRange