Cleanup ConvertToStream to accomodate llvm/llvm-project@3f136f7 (#19451)

The upstream change llvm/llvm-project@3f136f7
allows `ConvertToStream` to better handle the 1:N type conversion,
specifically the type conversion of a `tensor<...>` to
`!stream.resource<*>, index`. Now instead of trying to work around
`builtin.unrealized_conversion_cast`s the conversion can get the
converted values directly using the `OneToNAdaptor` and can also replace
a `tensor<..>` directly with multiple values using the
`ConversionPatternRewriter::replaceOpWithMultiple`.

These changes are required to drop the revert of
llvm/llvm-project#116470 in the IREE ToM. The
change drops these reverts as well.

Fixes #19448

---------

Signed-off-by: MaheshRavishankar <[email protected]>

Author: MaheshRavishankar

compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp

@@ -16,6 +16,14 @@ namespace mlir::iree_compiler {
 
 namespace {
 
+/// Flatten the given value ranges into a single vector of values.
+static SmallVector<Value> flattenValues(ArrayRef<ValueRange> values) {
+  SmallVector<Value> result;
+  for (const auto &vals : values)
+    llvm::append_range(result, vals);
+  return result;
+}
+
 // %1 = hal.tensor.import %0 : !hal.buffer_view -> tensor<4xf32>
 // ->
 // %1 = stream.tensor.import %0 : !hal.buffer_view ->
@@ -24,7 +32,7 @@ struct ConvertTensorImportOp
     : public AffinityOpConversionPattern<IREE::HAL::TensorImportOp> {
   using AffinityOpConversionPattern::AffinityOpConversionPattern;
   LogicalResult matchAndRewriteOnAffinity(
-      IREE::HAL::TensorImportOp op, OpAdaptor adaptor,
+      IREE::HAL::TensorImportOp op, OneToNOpAdaptor adaptor,
       IREE::Stream::AffinityAttr executionAffinityAttr,
       ConversionPatternRewriter &rewriter) const override {
     auto sourceType = op.getSource().getType();
@@ -42,9 +50,9 @@ struct ConvertTensorImportOp
       // mistake and it's better to know of a shape mismatch than just buffer
       // byte length difference.
       if (auto tensorType = llvm::dyn_cast<RankedTensorType>(targetType)) {
-        if (failed(buildEncodingAssertions(op.getLoc(), adaptor.getSource(),
-                                           op.getNameAttr(), tensorType,
-                                           op.getTargetDims(), rewriter))) {
+        if (failed(buildEncodingAssertions(
+                op.getLoc(), adaptor.getSource().front(), op.getNameAttr(),
+                tensorType, op.getTargetDims(), rewriter))) {
           return rewriter.notifyMatchFailure(op, "unsupported tensor type");
         }
       }
@@ -55,11 +63,12 @@ struct ConvertTensorImportOp
         IREE::Stream::Lifetime::External);
     Value resultSize = rewriter.create<IREE::Stream::TensorSizeOfOp>(
         op.getLoc(), rewriter.getIndexType(),
-        TypeAttr::get(op.getTarget().getType()), adaptor.getTargetDims(),
-        executionAffinityAttr);
+        TypeAttr::get(op.getTarget().getType()),
+        flattenValues(adaptor.getTargetDims()), executionAffinityAttr);
     Value resource = rewriter.create<IREE::Stream::TensorImportOp>(
-        op.getLoc(), resultType, adaptor.getSource(), TypeAttr::get(targetType),
-        adaptor.getTargetDims(), resultSize, executionAffinityAttr);
+        op.getLoc(), resultType, adaptor.getSource().front(),
+        TypeAttr::get(targetType), flattenValues(adaptor.getTargetDims()),
+        resultSize, executionAffinityAttr);
 
     // Await the fence, if needed. When not specified the resource is assumed to
     // be immediately available.
@@ -75,10 +84,11 @@ struct ConvertTensorImportOp
     }
 
     auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
-    rewriter.replaceOpWithNewOp<IREE::Stream::AsyncTransferOp>(
-        op, unknownType, resource, resultSize, resultSize,
+    Value newImport = rewriter.create<IREE::Stream::AsyncTransferOp>(
+        op.getLoc(), unknownType, resource, resultSize, resultSize,
         /*source_affinity=*/executionAffinityAttr,
         /*target_affinity=*/executionAffinityAttr);
+    rewriter.replaceOpWithMultiple(op, {{newImport, resultSize}});
     return success();
   }
 
@@ -125,7 +135,7 @@ struct ConvertTensorExportOp
     : public AffinityOpConversionPattern<IREE::HAL::TensorExportOp> {
   using AffinityOpConversionPattern::AffinityOpConversionPattern;
   LogicalResult matchAndRewriteOnAffinity(
-      IREE::HAL::TensorExportOp op, OpAdaptor adaptor,
+      IREE::HAL::TensorExportOp op, OneToNOpAdaptor adaptor,
       IREE::Stream::AffinityAttr executionAffinityAttr,
       ConversionPatternRewriter &rewriter) const override {
     auto sourceType = op.getSourceEncoding();
@@ -136,12 +146,12 @@ struct ConvertTensorExportOp
     }
 
     auto source =
-        transferTensorOperand(op.getLoc(), op.getSource(), adaptor.getSource(),
-                              executionAffinityAttr, rewriter);
+        transferTensorOperands(op.getLoc(), op.getSource(), adaptor.getSource(),
+                               executionAffinityAttr, rewriter);
 
     // Exporting a produced value - transfer our source value to an externally
     // usable resource and directly export it. This will cause an allocation.
-    auto exportSource = adaptor.getSource();
+    Value exportSource = adaptor.getSource().front();
     auto externalType = rewriter.getType<IREE::Stream::ResourceType>(
         IREE::Stream::Lifetime::External);
     if (source.resource.getType() != externalType) {
@@ -154,7 +164,8 @@ struct ConvertTensorExportOp
     // Export (stream resource to buffer view).
     rewriter.replaceOpWithNewOp<IREE::Stream::TensorExportOp>(
         op, targetType, exportSource, TypeAttr::get(sourceType),
-        adaptor.getSourceDims(), source.resourceSize, executionAffinityAttr);
+        flattenValues(adaptor.getSourceDims()), source.resourceSize,
+        executionAffinityAttr);
     return success();
   }
 };
@@ -174,19 +185,21 @@ struct ConvertTensorAliasOp
     : public AffinityOpConversionPattern<IREE::HAL::TensorAliasOp> {
   using AffinityOpConversionPattern::AffinityOpConversionPattern;
   LogicalResult matchAndRewriteOnAffinity(
-      IREE::HAL::TensorAliasOp op, OpAdaptor adaptor,
+      IREE::HAL::TensorAliasOp op, OneToNOpAdaptor adaptor,
       IREE::Stream::AffinityAttr executionAffinityAttr,
       ConversionPatternRewriter &rewriter) const override {
     auto sourceType = op.getSource().getType();
     auto source =
-        transferTensorOperand(op.getLoc(), op.getSource(), adaptor.getSource(),
-                              executionAffinityAttr, rewriter);
+        transferTensorOperands(op.getLoc(), op.getSource(), adaptor.getSource(),
+                               executionAffinityAttr, rewriter);
 
     // Query the target storage buffer length; we will only populate up to
     // what is required for the output.
+    SmallVector<Value> convertedSourceDims =
+        flattenValues(adaptor.getSourceDims());
     Value storageSize = rewriter.create<IREE::Stream::TensorSizeOfOp>(
         op.getLoc(), rewriter.getIndexType(),
-        TypeAttr::get(op.getSource().getType()), adaptor.getSourceDims(),
+        TypeAttr::get(op.getSource().getType()), convertedSourceDims,
         executionAffinityAttr);
 
     // Import the target storage as a resource that we can use as an update
@@ -195,8 +208,8 @@ struct ConvertTensorAliasOp
     auto externalType = rewriter.getType<IREE::Stream::ResourceType>(
         IREE::Stream::Lifetime::External);
     auto importOp = rewriter.create<IREE::Stream::TensorImportOp>(
-        op.getLoc(), externalType, adaptor.getStorage(),
-        TypeAttr::get(sourceType), adaptor.getSourceDims(), storageSize,
+        op.getLoc(), externalType, adaptor.getStorage().front(),
+        TypeAttr::get(sourceType), convertedSourceDims, storageSize,
         executionAffinityAttr);
 
     // Await the fence, if needed. When not specified the storage is assumed to
@@ -235,7 +248,7 @@ struct ConvertTensorAliasOp
           op.getLoc(), source.resource.getType(), result, source.resourceSize,
           source.resourceSize, executionAffinityAttr, executionAffinityAttr);
     }
-    rewriter.replaceOp(op, result);
+    rewriter.replaceOpWithMultiple(op, {{result, source.resourceSize}});
 
     return success();
   }
@@ -254,20 +267,22 @@ struct ConvertTensorBarrierOp
     : public AffinityAwareConversionPattern<IREE::HAL::TensorBarrierOp> {
   using AffinityAwareConversionPattern::AffinityAwareConversionPattern;
   LogicalResult
-  matchAndRewrite(IREE::HAL::TensorBarrierOp op, OpAdaptor adaptor,
+  matchAndRewrite(IREE::HAL::TensorBarrierOp op, OneToNOpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto timepointType = rewriter.getType<IREE::Stream::TimepointType>();
     IREE::Stream::AffinityAttr anyAffinityAttr;
     SmallVector<Value> signaledResources;
+    SmallVector<Value> signaledResourceSizes;
     SmallVector<Value> signaledTimepoints;
     for (auto [sourceTensor, sourceResource] :
          llvm::zip_equal(op.getSources(), adaptor.getSources())) {
-      auto source = resolveTensorOperand(op.getLoc(), sourceTensor,
-                                         sourceResource, rewriter);
+      auto source = resolveTensorOperands(op.getLoc(), sourceTensor,
+                                          sourceResource, rewriter);
       auto barrierOp = rewriter.create<IREE::Stream::TimepointBarrierOp>(
-          sourceResource.getLoc(), source.resource.getType(), timepointType,
-          source.resource, source.resourceSize, source.affinity);
+          sourceResource.front().getLoc(), source.resource.getType(),
+          timepointType, source.resource, source.resourceSize, source.affinity);
       signaledResources.push_back(barrierOp.getResult());
+      signaledResourceSizes.push_back(source.resourceSize);
       signaledTimepoints.push_back(barrierOp.getResultTimepoint());
 
       // When joining from multiple affinities we need to pick one to perform
@@ -283,7 +298,8 @@ struct ConvertTensorBarrierOp
     rewriter.create<IREE::Stream::TimepointChainExternalOp>(
         op.getLoc(), joinedTimepoint, ValueRange{adaptor.getSignalFence()},
         anyAffinityAttr);
-    rewriter.replaceOp(op, signaledResources);
+    replaceOpWithMultiple(op, signaledResources, signaledResourceSizes,
+                          rewriter);
     return success();
   }
 };

compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp

@@ -44,73 +44,25 @@ tryLookupResultAffinity(Value value,
   return affinityAnalysis->lookupResourceAffinity(value);
 }
 
-static std::pair<Value, Value>
-resolveTensorOperand(Location loc, Value convertedOperand, OpBuilder &builder) {
-  auto operandType = convertedOperand.getType();
-  if (llvm::isa<IREE::Stream::ResourceType>(operandType)) {
-    // Prior to https://reviews.llvm.org/D111620 this is the path we'd take;
-    // the tensor operands would be remapped into their new resource types.
-    // This is still possible during rewriting if we ourselves produce a new
-    // resource type, but the automatic materialization will go down the
-    // unrealized_conversion_cast path below.
-    return std::make_pair(convertedOperand,
-                          builder.createOrFold<IREE::Stream::ResourceSizeOp>(
-                              loc, builder.getIndexType(), convertedOperand));
-  } else if (auto castOp =
-                 convertedOperand
-                     .getDefiningOp<mlir::UnrealizedConversionCastOp>()) {
-    // We only have a single tensor type conversion and it expands to (resource,
-    // size) so that's all we look for here.
-    assert(castOp.getNumOperands() == 2 && "expected (resource, size)");
-    return std::make_pair(castOp.getOperand(0), castOp.getOperand(1));
-  }
-  assert(false &&
-         "unexpected operand; expected either a IREE::Stream::ResourceType or "
-         "the result of a mlir::UnrealizedConversionCastOp");
-  return std::make_pair(Value{}, Value{});
-}
-
-void expandResourceOperand(Location loc, Value operand,
-                           SmallVectorImpl<Value> &newOperands,
-                           OpBuilder &builder) {
-  if (llvm::isa<TensorType>(operand.getType())) {
-    auto [resource, resourceSize] = resolveTensorOperand(loc, operand, builder);
-    newOperands.push_back(resource);
-    newOperands.push_back(resourceSize);
-  } else if (llvm::isa<IREE::Stream::ResourceType>(operand.getType())) {
-    newOperands.push_back(operand);
-    newOperands.push_back(
-        builder.createOrFold<IREE::Stream::ResourceSizeOp>(loc, operand));
-  } else {
-    newOperands.push_back(operand);
-  }
-}
-
-SmallVector<Value> expandResourceOperands(Location loc, ValueRange operands,
-                                          ConversionPatternRewriter &rewriter) {
-  SmallVector<Value> expandedOperands;
-  expandedOperands.reserve(operands.size());
-  for (auto operand : operands) {
-    expandResourceOperand(loc, operand, expandedOperands, rewriter);
-  }
-  return expandedOperands;
-}
-
-ConvertedTensor resolveTensorOperand(
-    Location loc, Value originalOperand, Value convertedOperand,
+ConvertedTensor resolveTensorOperands(
+    Location loc, Value originalOperand, ValueRange convertedOperand,
     IREE::Stream::AffinityAnalysis *affinityAnalysis, OpBuilder &builder) {
-  auto [resource, resourceSize] =
-      resolveTensorOperand(loc, convertedOperand, builder);
+  assert(convertedOperand.size() == 2 &&
+         "expected tensor operands to be converted to `!stream.resource<*>, "
+         "index`");
   auto affinityAttr = affinityAnalysis->lookupResourceAffinity(originalOperand);
-  return {affinityAttr, resource, resourceSize};
+  return {affinityAttr, convertedOperand[0], convertedOperand[1]};
 }
 
-ConvertedTensor transferTensorOperand(
-    Location loc, Value originalOperand, Value convertedOperand,
+ConvertedTensor transferTensorOperands(
+    Location loc, Value originalOperand, ValueRange convertedOperand,
     IREE::Stream::AffinityAttr requiredAffinityAttr,
     IREE::Stream::AffinityAnalysis *affinityAnalysis, OpBuilder &builder) {
-  auto [resource, resourceSize] =
-      resolveTensorOperand(loc, convertedOperand, builder);
+  assert(convertedOperand.size() == 2 &&
+         "expected tensor operands to be converted to `!stream.resource<*>, "
+         "index`");
+  Value resource = convertedOperand[0];
+  Value resourceSize = convertedOperand[1];
   auto affinityAttr = affinityAnalysis->lookupResourceAffinity(originalOperand);
   if (affinityAttr != requiredAffinityAttr) {
     resource = builder.create<IREE::Stream::AsyncTransferOp>(
@@ -120,4 +72,25 @@ ConvertedTensor transferTensorOperand(
   return {requiredAffinityAttr, resource, resourceSize};
 }
 
+void replaceOpWithMultiple(Operation *op,
+                           ArrayRef<SmallVector<Value>> replacements,
+                           ConversionPatternRewriter &rewriter) {
+  auto r = llvm::map_to_vector(
+      replacements, [](ArrayRef<Value> v) -> ValueRange { return v; });
+  rewriter.replaceOpWithMultiple(op, r);
+}
+
+void replaceOpWithMultiple(Operation *op, ValueRange resources,
+                           ValueRange sizes,
+                           ConversionPatternRewriter &rewriter) {
+  SmallVector<SmallVector<Value>> replacements = llvm::map_to_vector(
+      llvm::zip_equal(resources, sizes), [](auto it) -> SmallVector<Value> {
+        if (std::get<1>(it)) {
+          return {std::get<0>(it), std::get<1>(it)};
+        }
+        return {std::get<0>(it)};
+      });
+  replaceOpWithMultiple(op, replacements, rewriter);
+}
+
 } // namespace mlir::iree_compiler

compiler/src/iree/compiler/Dialect/Stream/Conversion/PatternUtils.cpp

@@ -42,18 +42,11 @@ struct ConvertedTensor {
   Value resourceSize;
 };
 
-void expandResourceOperand(Location loc, Value convertedOperand,
-                           SmallVectorImpl<Value> &newOperands,
-                           OpBuilder &builder);
-SmallVector<Value> expandResourceOperands(Location loc,
-                                          ValueRange convertedOperands,
-                                          ConversionPatternRewriter &rewriter);
-
-ConvertedTensor resolveTensorOperand(
-    Location loc, Value originalOperand, Value convertedOperand,
+ConvertedTensor resolveTensorOperands(
+    Location loc, Value originalOperand, ValueRange convertedOperand,
     IREE::Stream::AffinityAnalysis *affinityAnalysis, OpBuilder &builder);
-ConvertedTensor transferTensorOperand(
-    Location loc, Value originalOperand, Value convertedOperand,
+ConvertedTensor transferTensorOperands(
+    Location loc, Value originalOperand, ValueRange convertedOperand,
     IREE::Stream::AffinityAttr requiredAffinityAttr,
     IREE::Stream::AffinityAnalysis *affinityAnalysis, OpBuilder &builder);
 
@@ -72,19 +65,19 @@ struct AffinityAwareConversionPattern : public OpConversionPattern<OpT> {
   }
 
 protected:
-  ConvertedTensor resolveTensorOperand(Location loc, Value originalOperand,
-                                       Value convertedOperand,
-                                       OpBuilder &builder) const {
-    return mlir::iree_compiler::resolveTensorOperand(
+  ConvertedTensor resolveTensorOperands(Location loc, Value originalOperand,
+                                        ValueRange convertedOperand,
+                                        OpBuilder &builder) const {
+    return mlir::iree_compiler::resolveTensorOperands(
         loc, originalOperand, convertedOperand, affinityAnalysis, builder);
   }
 
   ConvertedTensor
-  transferTensorOperand(Location loc, Value originalOperand,
-                        Value convertedOperand,
-                        IREE::Stream::AffinityAttr requiredAffinityAttr,
-                        OpBuilder &builder) const {
-    return mlir::iree_compiler::transferTensorOperand(
+  transferTensorOperands(Location loc, Value originalOperand,
+                         ValueRange convertedOperand,
+                         IREE::Stream::AffinityAttr requiredAffinityAttr,
+                         OpBuilder &builder) const {
+    return mlir::iree_compiler::transferTensorOperands(
         loc, originalOperand, convertedOperand, requiredAffinityAttr,
         affinityAnalysis, builder);
   }
@@ -110,13 +103,14 @@ struct AffinityOpConversionPattern
 
 protected:
   virtual LogicalResult matchAndRewriteOnAffinity(
-      OpT op, typename OpConversionPattern<OpT>::OpAdaptor adaptor,
+      OpT op, typename OpConversionPattern<OpT>::OneToNOpAdaptor adaptor,
       IREE::Stream::AffinityAttr executionAffinityAttr,
       ConversionPatternRewriter &rewriter) const = 0;
 
 private:
   LogicalResult
-  matchAndRewrite(OpT op, typename OpConversionPattern<OpT>::OpAdaptor adaptor,
+  matchAndRewrite(OpT op,
+                  typename OpConversionPattern<OpT>::OneToNOpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override final {
     auto executionAffinityAttr =
         tryLookupExecutionAffinity(op, this->getAffinityAnalysis());
@@ -125,6 +119,13 @@ struct AffinityOpConversionPattern
   }
 };
 
+void replaceOpWithMultiple(Operation *op,
+                           ArrayRef<SmallVector<Value>> replacements,
+                           ConversionPatternRewriter &rewriter);
+void replaceOpWithMultiple(Operation *op, ValueRange resources,
+                           ValueRange sizes,
+                           ConversionPatternRewriter &rewriter);
+
 } // namespace mlir::iree_compiler
 
 #endif // IREE_COMPILER_DIALECT_STREAM_CONVERSION_PATTERN_UTILS_H_