Integrate llvm-project @9372a3b70cf3969dac2d1a14cf41358205944e60 (#17926)
Bumps llvm-project to
https://github.com/llvm/llvm-project/commits/266a5a9cb9daa96c1eeaebc18e10f5a37d638734
Still carrying revert:
https://github.com/iree-org/llvm-project/commit/9372a3b70cf3969dac2d1a14cf41358205944e60
https://github.com/llvm/llvm-project/pull/97903 Updated type conversion
argument materialization, so this PR includes minor bug fixes in Codegen
and Stream conversions after the change.
---------
Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
Co-authored-by: Matthias Springer <mspringer@nvidia.com>
diff --git a/compiler/plugins/input/StableHLO/Conversion/TypeConversion.cpp b/compiler/plugins/input/StableHLO/Conversion/TypeConversion.cpp
index 6b2ef75..947fcba 100644
--- a/compiler/plugins/input/StableHLO/Conversion/TypeConversion.cpp
+++ b/compiler/plugins/input/StableHLO/Conversion/TypeConversion.cpp
@@ -56,17 +56,20 @@
->getResult(0);
}
-std::optional<Value> scalarToTensor(OpBuilder &builder, Type /*type*/,
+std::optional<Value> scalarToTensor(OpBuilder &builder, Type type,
ValueRange inputs, Location loc) {
assert(inputs.size() == 1);
if (llvm::isa<ShapedType>(inputs.front().getType())) {
return std::nullopt;
}
- return builder
- .create<tensor::FromElementsOp>(
- loc, RankedTensorType::get({}, inputs.front().getType()),
- inputs.front())
- .getResult();
+ auto tensor =
+ builder
+ .create<tensor::FromElementsOp>(
+ loc, RankedTensorType::get({}, inputs.front().getType()),
+ inputs.front())
+ .getResult();
+ return builder.create<UnrealizedConversionCastOp>(loc, type, tensor)
+ .getResult(0);
}
} // namespace
@@ -77,7 +80,7 @@
addConversion(convertInteger);
addConversion(convertShapedType);
- addArgumentMaterialization(materializeCastFromIllegal);
+ addArgumentMaterialization(materializeCastToIllegal);
addSourceMaterialization(materializeCastToIllegal);
addTargetMaterialization(materializeCastFromIllegal);
}
diff --git a/compiler/src/iree/compiler/Codegen/Common/ConvertBf16ToUInt16Buffers.cpp b/compiler/src/iree/compiler/Codegen/Common/ConvertBf16ToUInt16Buffers.cpp
index 13958d5..cb8a32b 100644
--- a/compiler/src/iree/compiler/Codegen/Common/ConvertBf16ToUInt16Buffers.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/ConvertBf16ToUInt16Buffers.cpp
@@ -164,19 +164,22 @@
for (Region &r : op->getRegions()) {
Region *newRegion = state.addRegion();
rewriter.inlineRegionBefore(r, *newRegion, newRegion->begin());
- TypeConverter::SignatureConversion result(newRegion->getNumArguments());
+ }
+ Operation *newOp = rewriter.create(state);
+
+ for (Region &newRegion : newOp->getRegions()) {
+ TypeConverter::SignatureConversion result(newRegion.getNumArguments());
if (failed(getTypeConverter()->convertSignatureArgs(
- newRegion->getArgumentTypes(), result))) {
+ newRegion.getArgumentTypes(), result))) {
return rewriter.notifyMatchFailure(op,
"argument type conversion failed");
}
- rewriter.applySignatureConversion(&newRegion->front(), result,
+ rewriter.applySignatureConversion(&newRegion.front(), result,
typeConverter);
}
- Operation *newOp = rewriter.create(state);
rewriter.replaceOp(op, newOp->getResults());
return success();
}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
index 0942148..060aeb8 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
@@ -413,11 +413,11 @@
consumeTensorOperand(op.getLoc(), adaptor.getSource(), rewriter);
rewriter.replaceOpWithNewOp<IREE::Stream::AsyncCollectiveOp>(
- op, collectiveAttr, adaptor.getTarget(),
+ op, collectiveAttr, newTargetCast.resource,
/*target_size=*/newTargetCast.resourceSize,
/*target_offset=*/zeroOffset,
/*target_end=*/newTargetCast.resourceSize,
- /*target_length=*/newTargetCast.resourceSize, adaptor.getSource(),
+ /*target_length=*/newTargetCast.resourceSize, newSourceCast.resource,
/*source_size=*/newSourceCast.resourceSize,
/*source_offset=*/zeroOffset, /*source_end=*/newSourceCast.resourceSize,
/*source_length=*/newSourceCast.resourceSize, elementCount,
@@ -448,11 +448,11 @@
consumeTensorOperand(op.getLoc(), adaptor.getSource(), rewriter);
rewriter.replaceOpWithNewOp<IREE::Stream::AsyncCollectiveOp>(
- op, collectiveAttr, adaptor.getTarget(),
+ op, collectiveAttr, newTargetCast.resource,
/*target_size=*/newTargetCast.resourceSize,
/*target_offset=*/zeroOffset,
/*target_end=*/newTargetCast.resourceSize,
- /*target_length=*/newTargetCast.resourceSize, adaptor.getSource(),
+ /*target_length=*/newTargetCast.resourceSize, newSourceCast.resource,
/*source_size=*/newSourceCast.resourceSize,
/*source_offset=*/zeroOffset, /*source_end=*/newSourceCast.resourceSize,
/*source_length=*/newSourceCast.resourceSize, elementCount,
@@ -483,11 +483,11 @@
consumeTensorOperand(op.getLoc(), adaptor.getSource(), rewriter);
rewriter.replaceOpWithNewOp<IREE::Stream::AsyncCollectiveOp>(
- op, collectiveAttr, adaptor.getTarget(),
+ op, collectiveAttr, newTargetCast.resource,
/*target_size=*/newTargetCast.resourceSize,
/*target_offset=*/zeroOffset,
/*target_end=*/newTargetCast.resourceSize,
- /*target_length=*/newTargetCast.resourceSize, adaptor.getSource(),
+ /*target_length=*/newTargetCast.resourceSize, newSourceCast.resource,
/*source_size=*/newSourceCast.resourceSize,
/*source_offset=*/zeroOffset, /*source_end=*/newSourceCast.resourceSize,
/*source_length=*/newSourceCast.resourceSize, elementCount,
@@ -518,11 +518,11 @@
consumeTensorOperand(op.getLoc(), adaptor.getSource(), rewriter);
rewriter.replaceOpWithNewOp<IREE::Stream::AsyncCollectiveOp>(
- op, collectiveAttr, adaptor.getTarget(),
+ op, collectiveAttr, newTargetCast.resource,
/*target_size=*/newTargetCast.resourceSize,
/*target_offset=*/zeroOffset,
/*target_end=*/newTargetCast.resourceSize,
- /*target_length=*/newTargetCast.resourceSize, adaptor.getSource(),
+ /*target_length=*/newTargetCast.resourceSize, newSourceCast.resource,
/*source_size=*/newSourceCast.resourceSize,
/*source_offset=*/zeroOffset, /*source_end=*/newSourceCast.resourceSize,
/*source_length=*/newSourceCast.resourceSize, elementCount,
@@ -567,11 +567,11 @@
auto param = rewriter.create<arith::OrIOp>(op.getLoc(), hi, lo);
rewriter.replaceOpWithNewOp<IREE::Stream::AsyncCollectiveOp>(
- op, collectiveAttr, adaptor.getTarget(),
+ op, collectiveAttr, newTargetCast.resource,
/*target_size=*/newTargetCast.resourceSize,
/*target_offset=*/zeroOffset,
/*target_end=*/newTargetCast.resourceSize,
- /*target_length=*/newTargetCast.resourceSize, adaptor.getSource(),
+ /*target_length=*/newTargetCast.resourceSize, newSourceCast.resource,
/*source_size=*/newSourceCast.resourceSize,
/*source_offset=*/zeroOffset, /*source_end=*/newSourceCast.resourceSize,
/*source_length=*/newSourceCast.resourceSize, elementCount,
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/compiler_hints.mlir b/compiler/src/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/compiler_hints.mlir
index f12a2ad..c778fbf 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/compiler_hints.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/compiler_hints.mlir
@@ -2,9 +2,9 @@
// CHECK-LABEL: @optimizationBarrier
util.func public @optimizationBarrier(%arg0: tensor<i32>) -> tensor<i32> {
- // CHECK: stream.async.transfer
- // CHECK: %[[RESOURCE:.*]] = util.optimization_barrier %0
- // CHECK: %[[SIZE:.*]] = stream.resource.size %1 : !stream.resource<*>
+ // CHECK-SAME: %[[ARG0:.+]]: !stream.resource<*>
+ // CHECK: %[[RESOURCE:.*]] = util.optimization_barrier %[[ARG0]]
+ // CHECK: %[[SIZE:.*]] = stream.resource.size %[[RESOURCE]] : !stream.resource<*>
// CHECK: util.return %[[RESOURCE]], %[[SIZE]] : !stream.resource<*>, index
%0 = util.optimization_barrier %arg0 : tensor<i32>
util.return %0 : tensor<i32>
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
index 46c09fa..11873a2 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
@@ -181,6 +181,31 @@
}
};
+namespace {
+struct OptimizationBarrierOpConversion
+ : public OpConversionPattern<IREE::Util::OptimizationBarrierOp> {
+ using OpConversionPattern<
+ IREE::Util::OptimizationBarrierOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(IREE::Util::OptimizationBarrierOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ SmallVector<Value> newOperands;
+ for (Value v : adaptor.getOperands()) {
+ if (isa<TensorType>(v.getType())) {
+ newOperands.push_back(
+ consumeTensorOperand(op.getLoc(), v, rewriter).resource);
+ } else {
+ newOperands.push_back(v);
+ }
+ }
+ rewriter.replaceOpWithNewOp<IREE::Util::OptimizationBarrierOp>(op,
+ newOperands);
+ return success();
+ }
+};
+} // namespace
+
//===----------------------------------------------------------------------===//
// --iree-stream-conversion
//===----------------------------------------------------------------------===//
@@ -228,13 +253,16 @@
auto resourceSize = inputs[1];
assert(inputs.size() == 2 &&
"expecting 2 operands (resource + size)");
+ Value cast = builder
+ .create<IREE::Stream::AsyncTransferOp>(
+ loc, resourceValue.getType(), resourceValue,
+ resourceSize, resourceSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr)
+ .getResult();
return builder
- .create<IREE::Stream::AsyncTransferOp>(
- loc, resourceValue.getType(), resourceValue, resourceSize,
- resourceSize,
- /*source_affinity=*/nullptr,
- /*result_affinity=*/nullptr)
- .getResult();
+ .create<UnrealizedConversionCastOp>(loc, resultType, cast)
+ .getResult(0);
});
populateUtilConversionPatterns(context, conversionTarget, typeConverter,
@@ -252,6 +280,8 @@
conversionTarget.markUnknownOpDynamicallyLegal(
[&](Operation *op) -> bool { return !doesOperationNeedWrapping(op); });
patterns.insert<GenericResourcePattern>(context, typeConverter);
+ patterns.insert<OptimizationBarrierOpConversion>(typeConverter, context,
+ /*benefit=*/2);
// NOTE: we allow ops that we don't know about to allow custom dialects
// that don't need anything Stream-specific to pass through.
diff --git a/third_party/llvm-project b/third_party/llvm-project
index 1f11b9f..9372a3b 160000
--- a/third_party/llvm-project
+++ b/third_party/llvm-project
@@ -1 +1 @@
-Subproject commit 1f11b9fed2337ea24d137ff82fec75bddcd85b3c
+Subproject commit 9372a3b70cf3969dac2d1a14cf41358205944e60
