Adding the head of the stream dialect conversion pipeline. (#7526)
Fixing the breakages from upstream dialect conversion changes.
See #7439 / #7520.
This reverts commit 80a3c85474aafaacb8a2e2c6faf698e9c63154c6.
diff --git a/iree/compiler/Dialect/Stream/Conversion/BUILD b/iree/compiler/Dialect/Stream/Conversion/BUILD
new file mode 100644
index 0000000..0a187e7
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/BUILD
@@ -0,0 +1,27 @@
+# Copyright 2021 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+cc_library(
+ name = "Conversion",
+ srcs = [
+ "PatternUtils.cpp",
+ ],
+ hdrs = [
+ "PatternUtils.h",
+ ],
+ deps = [
+ "//iree/compiler/Dialect/Stream/IR",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:StandardOps",
+ "@llvm-project//mlir:Transforms",
+ ],
+)
diff --git a/iree/compiler/Dialect/Stream/Conversion/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/CMakeLists.txt
new file mode 100644
index 0000000..ed6cbea
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/CMakeLists.txt
@@ -0,0 +1,28 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# iree/compiler/Dialect/Stream/Conversion/BUILD #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_cc_library(
+ NAME
+ Conversion
+ HDRS
+ "PatternUtils.h"
+ SRCS
+ "PatternUtils.cpp"
+ DEPS
+ MLIRIR
+ MLIRStandard
+ MLIRTransforms
+ iree::compiler::Dialect::Stream::IR
+ PUBLIC
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/BUILD b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/BUILD
new file mode 100644
index 0000000..d317206
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/BUILD
@@ -0,0 +1,31 @@
+# Copyright 2021 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+cc_library(
+ name = "FlowToStream",
+ srcs = [
+ "ConvertFlowToStream.cpp",
+ ],
+ hdrs = [
+ "ConvertFlowToStream.h",
+ ],
+ deps = [
+ "//iree/compiler/Dialect/Flow/IR",
+ "//iree/compiler/Dialect/HAL/IR",
+ "//iree/compiler/Dialect/Shape/IR",
+ "//iree/compiler/Dialect/Stream/Conversion",
+ "//iree/compiler/Dialect/Stream/IR",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:StandardOps",
+ "@llvm-project//mlir:TensorDialect",
+ ],
+)
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/CMakeLists.txt
new file mode 100644
index 0000000..d3b2d39
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/CMakeLists.txt
@@ -0,0 +1,32 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# iree/compiler/Dialect/Stream/Conversion/FlowToStream/BUILD #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_cc_library(
+ NAME
+ FlowToStream
+ HDRS
+ "ConvertFlowToStream.h"
+ SRCS
+ "ConvertFlowToStream.cpp"
+ DEPS
+ MLIRIR
+ MLIRStandard
+ MLIRTensor
+ iree::compiler::Dialect::Flow::IR
+ iree::compiler::Dialect::HAL::IR
+ iree::compiler::Dialect::Shape::IR
+ iree::compiler::Dialect::Stream::Conversion
+ iree::compiler::Dialect::Stream::IR
+ PUBLIC
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/ConvertFlowToStream.cpp b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/ConvertFlowToStream.cpp
new file mode 100644
index 0000000..bcb1fb8
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/ConvertFlowToStream.cpp
@@ -0,0 +1,467 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Stream/Conversion/FlowToStream/ConvertFlowToStream.h"
+
+#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
+#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/Shape/IR/Builders.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
+#include "iree/compiler/Dialect/Stream/Conversion/PatternUtils.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+namespace {
+
+// Inserts a sizeof calculation for the given tensor value type and dims.
+// This should only be used to produce sizes for values produced by an op; the
+// size of operands must be queried from the input resource.
+static Value buildResultSizeOf(Location loc, Value tensorValue,
+ ValueRange dynamicDims,
+ ConversionPatternRewriter &rewriter) {
+ // TODO(benvanik): see if we can stash this on the side to avoid expensive
+ // materialization of a bunch of redundant IR.
+ return rewriter.createOrFold<IREE::Stream::TensorSizeOfOp>(
+ loc, rewriter.getIndexType(), TypeAttr::get(tensorValue.getType()),
+ dynamicDims, /*affinity=*/nullptr);
+}
+
+// hal.tensor.cast is inserted by frontends to ensure that ABI types are HAL
+// buffer views. We need to map those to the stream import/export equivalents as
+// the cast has special meaning when we are dealing with asynchronous values.
+//
+// %1 = hal.tensor.cast %0 : !hal.buffer_view -> tensor<4xf32>
+// ->
+// %1 = stream.tensor.import %0 : !hal.buffer_view ->
+// tensor<4xf32> in !stream.resource<*>
+struct ConvertHALTensorCastOp
+ : public OpConversionPattern<IREE::HAL::TensorCastOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::HAL::TensorCastOp op, ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::HAL::TensorCastOpAdaptor operands(newOperands,
+ op->getAttrDictionary());
+ if (op.source().getType().isa<IREE::HAL::BufferViewType>()) {
+ // Import (buffer view to stream resource).
+ auto resultType = rewriter.getType<IREE::Stream::ResourceType>(
+ IREE::Stream::Lifetime::External);
+ auto resultSize = buildResultSizeOf(op.getLoc(), op.target(),
+ operands.target_dims(), rewriter);
+ auto newOp = rewriter.create<IREE::Stream::TensorImportOp>(
+ op.getLoc(), resultType, operands.source(),
+ TypeAttr::get(op.target().getType()), operands.target_dims(),
+ resultSize,
+ /*affinity=*/nullptr);
+
+ auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
+ rewriter.replaceOpWithNewOp<IREE::Stream::AsyncTransferOp>(
+ op, unknownType, newOp.result(), resultSize, resultSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr);
+
+ } else if (op.target().getType().isa<IREE::HAL::BufferViewType>()) {
+ auto source =
+ consumeTensorOperand(op.getLoc(), operands.source(), rewriter);
+ auto externalType = rewriter.getType<IREE::Stream::ResourceType>(
+ IREE::Stream::Lifetime::External);
+ auto exportSource = operands.source();
+ if (source.resource.getType() != externalType) {
+ exportSource = rewriter.create<IREE::Stream::AsyncTransferOp>(
+ op.getLoc(), externalType, source.resource, source.resourceSize,
+ source.resourceSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr);
+ }
+
+ // Export (stream resource to buffer view).
+ rewriter.replaceOpWithNewOp<IREE::Stream::TensorExportOp>(
+ op, op.target().getType(), exportSource,
+ TypeAttr::get(op.source().getType()), operands.source_dims(),
+ source.resourceSize,
+ /*affinity=*/nullptr);
+ } else {
+ return rewriter.notifyMatchFailure(op, "unsupported HAL cast conversion");
+ }
+ return success();
+ }
+};
+
+// Reshapes become clones here to preserve shape information (which may become
+// actual transfers depending on source/target shape) - they'll be elided if not
+// needed.
+struct ConvertTensorReshapeOp
+ : public OpConversionPattern<IREE::Flow::TensorReshapeOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::TensorReshapeOp op, ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Flow::TensorReshapeOpAdaptor operands(newOperands,
+ op->getAttrDictionary());
+ auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
+ auto source =
+ consumeTensorOperand(op.getLoc(), operands.source(), rewriter);
+ auto resultSize =
+ buildResultSizeOf(op.getLoc(), op.result(), op.result_dims(), rewriter);
+ rewriter.replaceOpWithNewOp<IREE::Stream::TensorCloneOp>(
+ op, unknownType, source.resource, op.source().getType(),
+ op.source_dims(), source.resourceSize, op.result().getType(),
+ operands.result_dims(), resultSize,
+ /*affinity=*/nullptr);
+ return success();
+ }
+};
+
+struct ConvertTensorSplatOp
+ : public OpConversionPattern<IREE::Flow::TensorSplatOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::TensorSplatOp op, ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Flow::TensorSplatOpAdaptor operands(newOperands);
+ auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
+ auto resultSize =
+ buildResultSizeOf(op.getLoc(), op.result(), op.result_dims(), rewriter);
+ rewriter.replaceOpWithNewOp<IREE::Stream::TensorSplatOp>(
+ op, unknownType, operands.value(), op.result().getType(),
+ operands.result_dims(), resultSize,
+ /*affinity=*/nullptr);
+ return success();
+ }
+};
+
+struct ConvertTensorCloneOp
+ : public OpConversionPattern<IREE::Flow::TensorCloneOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::TensorCloneOp op, ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Flow::TensorCloneOpAdaptor operands(newOperands);
+ auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
+ auto operand =
+ consumeTensorOperand(op.getLoc(), operands.operand(), rewriter);
+ rewriter.replaceOpWithNewOp<IREE::Stream::TensorCloneOp>(
+ op, unknownType, operand.resource, op.operand().getType(),
+ op.operand_dims(), operand.resourceSize, op.result().getType(),
+ operands.operand_dims(), operand.resourceSize,
+ /*affinity=*/nullptr);
+ return success();
+ }
+};
+
+struct ConvertTensorSliceOp
+ : public OpConversionPattern<IREE::Flow::TensorSliceOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::TensorSliceOp op, ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Flow::TensorSliceOpAdaptor operands(newOperands,
+ op->getAttrDictionary());
+ auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
+ auto source =
+ consumeTensorOperand(op.getLoc(), operands.source(), rewriter);
+ auto resultSize =
+ buildResultSizeOf(op.getLoc(), op.result(), op.result_dims(), rewriter);
+ rewriter.replaceOpWithNewOp<IREE::Stream::TensorSliceOp>(
+ op, unknownType, source.resource, op.source().getType(),
+ op.source_dims(), source.resourceSize, operands.start_indices(),
+ operands.lengths(), op.result().getType(), operands.result_dims(),
+ resultSize,
+ /*affinity=*/nullptr);
+ return success();
+ }
+};
+
+struct ConvertTensorUpdateOp
+ : public OpConversionPattern<IREE::Flow::TensorUpdateOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::TensorUpdateOp op, ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Flow::TensorUpdateOpAdaptor operands(newOperands,
+ op->getAttrDictionary());
+ auto update =
+ consumeTensorOperand(op.getLoc(), operands.update(), rewriter);
+ auto target =
+ consumeTensorOperand(op.getLoc(), operands.target(), rewriter);
+ rewriter.replaceOpWithNewOp<IREE::Stream::TensorUpdateOp>(
+ op, target.resource.getType(), target.resource, op.target().getType(),
+ operands.target_dims(), target.resourceSize, operands.start_indices(),
+ update.resource, op.update().getType(), op.update_dims(),
+ update.resourceSize,
+ /*affinity=*/nullptr);
+ return success();
+ }
+};
+
+struct ConvertTensorLoadOp
+ : public OpConversionPattern<IREE::Flow::TensorLoadOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::TensorLoadOp op, ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Flow::TensorLoadOpAdaptor operands(newOperands,
+ op->getAttrDictionary());
+ auto resultType = getTypeConverter()->convertType(op.result().getType());
+ auto source =
+ consumeTensorOperand(op.getLoc(), operands.source(), rewriter);
+
+ auto stagingType = rewriter.getType<IREE::Stream::ResourceType>(
+ IREE::Stream::Lifetime::Staging);
+ auto loadSource = source.resource;
+ if (source.resource.getType() != stagingType) {
+ loadSource = rewriter.createOrFold<IREE::Stream::AsyncTransferOp>(
+ op.getLoc(), stagingType, source.resource, source.resourceSize,
+ source.resourceSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr);
+ }
+
+ rewriter.replaceOpWithNewOp<IREE::Stream::TensorLoadOp>(
+ op, resultType, loadSource, op.source().getType(), op.source_dims(),
+ source.resourceSize, operands.indices());
+ return success();
+ }
+};
+
+struct ConvertTensorStoreOp
+ : public OpConversionPattern<IREE::Flow::TensorStoreOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::TensorStoreOp op, ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Flow::TensorStoreOpAdaptor operands(newOperands,
+ op->getAttrDictionary());
+ auto target =
+ consumeTensorOperand(op.getLoc(), operands.target(), rewriter);
+
+ auto stagingType = rewriter.getType<IREE::Stream::ResourceType>(
+ IREE::Stream::Lifetime::Staging);
+ auto storeTarget = target.resource;
+ if (target.resource.getType() != stagingType) {
+ storeTarget = rewriter.createOrFold<IREE::Stream::AsyncTransferOp>(
+ op.getLoc(), stagingType, storeTarget, target.resourceSize,
+ target.resourceSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr);
+ }
+
+ auto newOp = rewriter.create<IREE::Stream::TensorStoreOp>(
+ op.getLoc(), storeTarget.getType(), storeTarget, op.target().getType(),
+ operands.target_dims(), target.resourceSize, operands.indices(),
+ operands.value());
+
+ Value newResult = newOp.result();
+ if (target.resource.getType() != stagingType) {
+ newResult = rewriter.createOrFold<IREE::Stream::AsyncTransferOp>(
+ op.getLoc(), target.resource.getType(), newResult,
+ target.resourceSize, target.resourceSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr);
+ }
+ rewriter.replaceOp(op, {newResult});
+
+ return success();
+ }
+};
+
+struct ConvertDispatchOp : public OpConversionPattern<IREE::Flow::DispatchOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::DispatchOp op, ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Flow::DispatchOpAdaptor operands(newOperands,
+ op->getAttrDictionary());
+
+ // Query and resolve all operands and their sizes.
+ SmallVector<Value> dispatchOperands;
+ SmallVector<Value> dispatchOperandSizes;
+ for (auto oldNewOperand : llvm::zip(op.operands(), operands.operands())) {
+ auto oldOperand = std::get<0>(oldNewOperand);
+ auto newOperand = std::get<1>(oldNewOperand);
+ if (oldOperand.getType().isa<ShapedType>()) {
+ auto newOperandCast =
+ consumeTensorOperand(op.getLoc(), newOperand, rewriter);
+ newOperand = newOperandCast.resource;
+ dispatchOperandSizes.push_back(newOperandCast.resourceSize);
+ }
+ dispatchOperands.push_back(newOperand);
+ }
+
+ // Construct result sizes or reuse tied operand sizes from above.
+ SmallVector<Value> resultSizes;
+ SmallVector<Type> resultTypes;
+ auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
+ auto tiedOperandBase = op.getTiedOperandsIndexAndLength().first;
+ for (auto result : llvm::enumerate(op.results())) {
+ auto oldResultType = result.value().getType();
+ if (!oldResultType.isa<ShapedType>()) {
+ resultTypes.push_back(getTypeConverter()->convertType(oldResultType));
+ continue;
+ }
+ auto tiedOperand = op.getTiedResultOperandIndex(result.index());
+ if (tiedOperand.hasValue()) {
+ auto operandIndex = tiedOperand.getValue() - tiedOperandBase;
+ resultSizes.push_back(dispatchOperandSizes[operandIndex]);
+ resultTypes.push_back(dispatchOperands[operandIndex].getType());
+ } else {
+ auto resultDynamicDims = Shape::buildOrFindDynamicDimsForValue(
+ op.getLoc(), result.value(), rewriter);
+ resultSizes.push_back(buildResultSizeOf(op.getLoc(), result.value(),
+ resultDynamicDims, rewriter));
+ resultTypes.push_back(unknownType);
+ }
+ }
+
+ rewriter.replaceOpWithNewOp<IREE::Stream::AsyncDispatchOp>(
+ op, resultTypes, operands.workgroup_count(), operands.entry_point(),
+ dispatchOperands, dispatchOperandSizes, resultSizes,
+ operands.tied_operands(),
+ /*affinity=*/nullptr);
+ return success();
+ }
+};
+
+static SmallVector<Value> makeBindingDynamicDims(
+ Location loc, IREE::Flow::DispatchTensorType tensorType, BlockArgument arg,
+ OpBuilder &builder) {
+ if (tensorType.hasStaticShape()) return {};
+
+ // We can expect its first user to be a tie_shape op to associate
+ // concrete dimension values. Originally we have such information
+ // maintained in the flow ops handling dynamic tensors. But during
+ // flow executable outlining, such information is transfered to
+ // tie_shape ops.
+ //
+ // HACK: this is disgusting - we should carry this from the flow level in the
+ // right way such that we don't need to make this assumption.
+ auto tieShapeOp = dyn_cast<IREE::Flow::DispatchTieShapeOp>(*arg.user_begin());
+ assert(tieShapeOp && "missing flow tie shape for dynamic value");
+ builder.setInsertionPointAfter(tieShapeOp.shape().getDefiningOp());
+
+ // Get the SSA values for all dynamic dimensions.
+ SmallVector<Value> dynamicDims;
+ dynamicDims.reserve(tensorType.getNumDynamicDims());
+ for (int i = 0; i < tensorType.getRank(); ++i) {
+ if (!tensorType.isDynamicDim(i)) continue;
+ dynamicDims.push_back(builder.create<Shape::RankedDimOp>(
+ tieShapeOp.getLoc(), tieShapeOp.shape(), i));
+ }
+ assert(dynamicDims.size() == tensorType.getNumDynamicDims());
+
+ return dynamicDims;
+}
+
+struct ConvertExecutableOp
+ : public OpConversionPattern<IREE::Flow::ExecutableOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Flow::ExecutableOp flowOp, ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ // flow.executable -> stream.executable
+ auto streamOp = rewriter.create<IREE::Stream::ExecutableOp>(
+ flowOp.getLoc(), flowOp.sym_name());
+ streamOp.setVisibility(flowOp.getVisibility());
+ streamOp->setDialectAttrs(flowOp->getDialectAttrs());
+ rewriter.setInsertionPointToStart(&streamOp.body().front());
+
+ // flow.dispatch.entry -> stream.executable.entry_point
+ for (auto entryOp : flowOp.getOps<IREE::Flow::DispatchEntryOp>()) {
+ auto newOp = rewriter.create<IREE::Stream::ExecutableExportOp>(
+ entryOp.getLoc(), entryOp.sym_name(), entryOp.function_refAttr());
+ newOp->setDialectAttrs(entryOp->getDialectAttrs());
+ }
+
+ // Move the original nested module body into the new executable directly.
+ auto moduleOp = rewriter.cloneWithoutRegions(flowOp.getInnerModule());
+ streamOp.getInnerModule().body().takeBody(flowOp.getInnerModule().body());
+
+ // Update the entry point signatures in the module.
+ // Dispatch tensor arguments become bindings and all others are preserved as
+ // operands. Note that we only touch public (exported) functions.
+ for (auto funcOp : moduleOp.getOps<mlir::FuncOp>()) {
+ if (!funcOp.isPublic()) continue;
+
+ SmallVector<Type> newTypes;
+ newTypes.reserve(funcOp.getNumArguments());
+ assert(funcOp.getNumResults() == 0 && "flow dispatches have no results");
+
+ rewriter.setInsertionPointToStart(&funcOp.front());
+ auto zero = rewriter.create<arith::ConstantIndexOp>(funcOp.getLoc(), 0);
+ for (auto arg : funcOp.front().getArguments()) {
+ auto oldType = arg.getType();
+ if (auto tensorType =
+ oldType.dyn_cast<IREE::Flow::DispatchTensorType>()) {
+ // Now a binding.
+ auto newType = rewriter.getType<IREE::Stream::BindingType>();
+ newTypes.push_back(newType);
+ auto dynamicDims =
+ makeBindingDynamicDims(arg.getLoc(), tensorType, arg, rewriter);
+ auto subspanOp = rewriter.create<IREE::Stream::BindingSubspanOp>(
+ arg.getLoc(), tensorType, arg, zero, dynamicDims);
+ arg.replaceAllUsesExcept(subspanOp.result(), subspanOp);
+ arg.setType(newType);
+ } else {
+ // Preserved - will eventually be a push constants.
+ newTypes.push_back(oldType);
+ }
+ }
+
+ funcOp.setType(rewriter.getFunctionType(newTypes, {}));
+ }
+
+ rewriter.replaceOp(flowOp, {});
+ return success();
+ }
+};
+
+} // namespace
+
+void populateFlowToStreamConversionPatterns(
+ MLIRContext *context, TypeConverter &typeConverter,
+ OwningRewritePatternList &patterns) {
+ typeConverter.addConversion(
+ [](IREE::HAL::BufferViewType type) { return type; });
+ patterns.insert<ConvertHALTensorCastOp>(typeConverter, context);
+
+ patterns
+ .insert<ConvertTensorReshapeOp, ConvertTensorSplatOp,
+ ConvertTensorCloneOp, ConvertTensorSliceOp, ConvertTensorUpdateOp,
+ ConvertTensorLoadOp, ConvertTensorStoreOp>(typeConverter,
+ context);
+ patterns.insert<ConvertDispatchOp>(typeConverter, context);
+ patterns.insert<ConvertExecutableOp>(typeConverter, context);
+}
+
+void populateFlowToStreamConversionPatterns(
+ MLIRContext *context, ConversionTarget &conversionTarget,
+ TypeConverter &typeConverter, OwningRewritePatternList &patterns) {
+ // Disallow all flow ops besides the ones we pass through (today).
+ // We don't have a stream-equivalent of several of the dispatch-level flow
+ // ops as the codegen backends directly touch them and so long as we have both
+ // paths we can't cut over. Once we convert the flow.executable to a
+ // stream.executable we ignore the contents and cross our fingers.
+ conversionTarget.addIllegalDialect<IREE::Flow::FlowDialect>();
+ conversionTarget.addLegalOp<IREE::Stream::ExecutableOp>();
+ conversionTarget.markOpRecursivelyLegal<IREE::Stream::ExecutableOp>();
+
+ conversionTarget.addDynamicallyLegalOp<IREE::HAL::TensorCastOp>(
+ [&](IREE::HAL::TensorCastOp op) {
+ return typeConverter.isLegal(op.source().getType()) &&
+ typeConverter.isLegal(op.target().getType());
+ });
+
+ populateFlowToStreamConversionPatterns(context, typeConverter, patterns);
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/ConvertFlowToStream.h b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/ConvertFlowToStream.h
new file mode 100644
index 0000000..0ea90e6
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/ConvertFlowToStream.h
@@ -0,0 +1,31 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_DIALECT_STREAM_CONVERSION_FLOWTOSTREAM_CONVERTFLOWTOSTREAM_H_
+#define IREE_COMPILER_DIALECT_STREAM_CONVERSION_FLOWTOSTREAM_CONVERTFLOWTOSTREAM_H_
+
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+// Populates conversion patterns that perform flow->stream conversion.
+// These patterns ensure that nested types are run through the provided
+// |typeConverter|.
+void populateFlowToStreamConversionPatterns(MLIRContext *context,
+ TypeConverter &typeConverter,
+ OwningRewritePatternList &patterns);
+void populateFlowToStreamConversionPatterns(MLIRContext *context,
+ ConversionTarget &conversionTarget,
+ TypeConverter &typeConverter,
+ OwningRewritePatternList &patterns);
+
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_DIALECT_STREAM_CONVERSION_FLOWTOSTREAM_CONVERTFLOWTOSTREAM_H_
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD
new file mode 100644
index 0000000..1f3d342
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD
@@ -0,0 +1,30 @@
+# Copyright 2021 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+iree_lit_test_suite(
+ name = "lit",
+ srcs = enforce_glob(
+ [
+ "cast_ops.mlir",
+ "dispatch_ops.mlir",
+ "tensor_ops.mlir",
+ ],
+ include = ["*.mlir"],
+ ),
+ data = [
+ "//iree/tools:IreeFileCheck",
+ "//iree/tools:iree-opt",
+ ],
+)
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt
new file mode 100644
index 0000000..1974566
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/CMakeLists.txt
@@ -0,0 +1,25 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/BUILD #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_lit_test_suite(
+ NAME
+ lit
+ SRCS
+ "cast_ops.mlir"
+ "dispatch_ops.mlir"
+ "tensor_ops.mlir"
+ DATA
+ iree::tools::IreeFileCheck
+ iree::tools::iree-opt
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/cast_ops.mlir b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/cast_ops.mlir
new file mode 100644
index 0000000..17f9a6a
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/cast_ops.mlir
@@ -0,0 +1,34 @@
+// RUN: iree-opt -split-input-file -iree-stream-conversion -canonicalize %s | IreeFileCheck %s
+
+// CHECK-LABEL: @importBufferView
+// CHECK-SAME: (%[[VIEW:.+]]: !hal.buffer_view)
+// CHECK-SAME: -> (!stream.resource<*>, index)
+func @importBufferView(%view: !hal.buffer_view) -> tensor<?x?x4xf32> {
+ // CHECK-DAG: %[[DIM0:.+]] = hal.buffer_view.dim{{.+}}[0]
+ %dim0 = hal.buffer_view.dim<%view : !hal.buffer_view>[0] : index
+ // CHECK-DAG: %[[DIM1:.+]] = hal.buffer_view.dim{{.+}}[1]
+ %dim1 = hal.buffer_view.dim<%view : !hal.buffer_view>[1] : index
+ // CHECK-DAG: %[[SIZE:.+]] = stream.tensor.sizeof tensor<?x?x4xf32>{%[[DIM0]], %[[DIM1]]} : index
+ // CHECK: %[[RESOURCE:.+]] = stream.tensor.import %[[VIEW]] : !hal.buffer_view ->
+ // CHECK-SAME: tensor<?x?x4xf32>{%[[DIM0]], %[[DIM1]]} in !stream.resource<external>{%[[SIZE]]}
+ // CHECK-NEXT: %[[RESULT:.+]] = stream.async.transfer %[[RESOURCE]] :
+ // CHECK-SAME: !stream.resource<external>{%[[SIZE]]} -> !stream.resource<*>{%[[SIZE]]}
+ %0 = hal.tensor.cast %view : !hal.buffer_view -> tensor<?x?x4xf32>{%dim0, %dim1}
+ // CHECK: return %[[RESULT]], %[[SIZE]] : !stream.resource<*>, index
+ return %0 : tensor<?x?x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @exportBufferView
+// CHECK-SAME: (%[[TENSOR:.+]]: !stream.resource<*>, %[[SIZE:.+]]: index, %[[DIM0:.+]]: index, %[[DIM1:.+]]: index)
+func @exportBufferView(%tensor: tensor<?x?x4xf32>, %dim0: index, %dim1: index) -> !hal.buffer_view {
+ // CHECK: %[[VIEW:.+]] = stream.async.transfer %[[TENSOR]] :
+ // CHECK-SAME: !stream.resource<*>{%[[SIZE]]} -> !stream.resource<external>{%[[SIZE]]}
+ // CHECK-NEXT: %[[RESULT:.+]] = stream.tensor.export %[[VIEW]] :
+ // CHECK-SAME: tensor<?x?x4xf32>{%[[DIM0]], %[[DIM1]]} in !stream.resource<external>{%[[SIZE]]}
+ // CHECK-SAME: -> !hal.buffer_view
+ %0 = hal.tensor.cast %tensor : tensor<?x?x4xf32>{%dim0, %dim1} -> !hal.buffer_view
+ // CHECK: return %[[RESULT]]
+ return %0 : !hal.buffer_view
+}
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/dispatch_ops.mlir b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/dispatch_ops.mlir
new file mode 100644
index 0000000..af4204b
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/dispatch_ops.mlir
@@ -0,0 +1,32 @@
+// RUN: iree-opt -split-input-file -iree-stream-conversion -canonicalize %s | IreeFileCheck %s
+
+// CHECK-LABEL: @dispatch
+// CHECK-SAME: (%[[INPUT:.+]]: !stream.resource<*>, %[[INPUT_SIZE:.+]]: index, %[[DIM1:.+]]: index, %[[DIM3:.+]]: index)
+func @dispatch(%input: tensor<7x?x24x?xf32>, %dim1: index, %dim3: index) -> tensor<?x?x1024xf32> {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ // CHECK: %[[RESULT_SIZE:.+]] = stream.tensor.sizeof tensor<?x?x1024xf32>{%[[DIM1]], %[[DIM3]]}
+ // CHECK: %[[RESULT:.+]] = stream.async.dispatch @ex::@entry[%c1, %c2, %c3](%[[INPUT]]) :
+ // CHECK-SAME: (!stream.resource<*>{%[[INPUT_SIZE]]}) -> !stream.resource<*>{%[[RESULT_SIZE]]}
+ %0 = flow.dispatch @ex::@entry[%c1, %c2, %c3](%input) : (tensor<7x?x24x?xf32>{%dim1, %dim3}) -> tensor<?x?x1024xf32>{%dim1, %dim3}
+ // return %[[RESULT]], %[[RESULT_SIZE]] : !stream.resource<*>, index
+ return %0 : tensor<?x?x1024xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @tiedDispatch
+// CHECK-SAME: (%[[INPUT0:.+]]: !stream.resource<*>, %[[INPUT0_SIZE:.+]]: index, %[[INPUT1:.+]]: !stream.resource<*>, %[[INPUT1_SIZE:.+]]: index)
+func @tiedDispatch(%input0: tensor<i32>, %input1: tensor<2x3xi32>) -> tensor<3x9xi32> {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ // CHECK: %[[T_SIZE:.+]] = stream.tensor.sizeof tensor<3x9xi32> : index
+ // CHECK: %[[T:.+]] = stream.async.dispatch @ex::@entry0[%c1, %c2, %c3](%[[INPUT0]]) : (!stream.resource<*>{%[[INPUT0_SIZE]]}) -> !stream.resource<*>{%[[T_SIZE]]}
+ %0 = flow.dispatch @ex::@entry0[%c1, %c2, %c3](%input0) : (tensor<i32>) -> tensor<3x9xi32>
+ // CHECK: %[[RESULT:.+]] = stream.async.dispatch @ex::@entry1[%c1, %c2, %c3](%[[INPUT1]], %[[T]]) : (!stream.resource<*>{%[[INPUT1_SIZE]]}, !stream.resource<*>{%[[T_SIZE]]}) -> %[[T]]{%[[T_SIZE]]}
+ %1 = flow.dispatch @ex::@entry1[%c1, %c2, %c3](%input1, %0) : (tensor<2x3xi32>, tensor<3x9xi32>) -> %0
+ // CHECK: return %[[RESULT]], %[[T_SIZE]] : !stream.resource<*>, index
+ return %1 : tensor<3x9xi32>
+}
diff --git a/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir
new file mode 100644
index 0000000..629232f
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/FlowToStream/test/tensor_ops.mlir
@@ -0,0 +1,117 @@
+// RUN: iree-opt -split-input-file -iree-stream-conversion %s | IreeFileCheck %s
+
+// CHECK-LABEL: @tensorReshapePassThrough
+// CHECK-SAME: (%[[INPUT:.+]]: !stream.resource<*>, %[[INPUT_SIZE:.+]]: index)
+func @tensorReshapePassThrough(%input: tensor<5x24x48xf32>) -> tensor<30x2x96xf32> {
+ // CHECK: %[[RESULT_SIZE:.+]] = stream.tensor.sizeof tensor<30x2x96xf32> : index
+ // CHECK: %[[RESULT:.+]] = stream.tensor.clone %[[INPUT]] : tensor<5x24x48xf32> in !stream.resource<*>{%[[INPUT_SIZE]]} -> tensor<30x2x96xf32> in !stream.resource<*>{%[[RESULT_SIZE]]}
+ %0 = flow.tensor.reshape %input : tensor<5x24x48xf32> -> tensor<30x2x96xf32>
+ // CHECK: return %[[RESULT]], %[[RESULT_SIZE]] : !stream.resource<*>, index
+ return %0 : tensor<30x2x96xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @tensorReshapeWithSingleUse
+// CHECK-SAME: (%[[INPUT:.+]]: !stream.resource<*>, %[[INPUT_SIZE:.+]]: index)
+func @tensorReshapeWithSingleUse(%input: tensor<5x24x48xf32>) -> tensor<30x2x96xf32> {
+ // CHECK: %[[RESULT_SIZE:.+]] = stream.tensor.sizeof tensor<30x2x96xf32> : index
+ // CHECK: %[[RESHAPE:.+]] = stream.tensor.clone %[[INPUT]] : tensor<5x24x48xf32> in !stream.resource<*>{%[[INPUT_SIZE]]} -> tensor<30x2x96xf32> in !stream.resource<*>{%[[RESULT_SIZE]]}
+ %0 = flow.tensor.reshape %input : tensor<5x24x48xf32> -> tensor<30x2x96xf32>
+ // CHECK: %[[RESULT:.+]] = stream.tensor.clone %[[RESHAPE]] : tensor<30x2x96xf32> in !stream.resource<*>{%[[RESULT_SIZE]]} -> tensor<30x2x96xf32> in !stream.resource<*>{%[[RESULT_SIZE]]}
+ %1 = flow.tensor.clone %0 : tensor<30x2x96xf32>
+ // CHECK: return %[[RESULT]], %[[RESULT_SIZE]] : !stream.resource<*>, index
+ return %1 : tensor<30x2x96xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @tensorReshapeWithMultipleUses
+// CHECK-SAME: (%[[INPUT:.+]]: !stream.resource<*>, %[[INPUT_SIZE:.+]]: index)
+func @tensorReshapeWithMultipleUses(%input: tensor<5x24x48xf32>)
+ -> (tensor<60x2x48xf32>, tensor<30x2x96xf32>) {
+ // CHECK: %[[T0:.+]] = stream.tensor.clone %[[INPUT]] : tensor<5x24x48xf32> in !stream.resource<*>{%[[INPUT_SIZE]]} -> tensor<5x24x48xf32> in !stream.resource<*>{%[[INPUT_SIZE]]}
+ %1 = flow.tensor.clone %input : tensor<5x24x48xf32>
+ // CHECK: %[[T1_SIZE:.+]] = stream.tensor.sizeof tensor<60x2x48xf32> : index
+ // CHECK: %[[T1:.+]] = stream.tensor.clone %[[INPUT]] : tensor<5x24x48xf32> in !stream.resource<*>{%[[INPUT_SIZE]]} -> tensor<60x2x48xf32> in !stream.resource<*>{%[[T1_SIZE]]}
+ %2 = flow.tensor.reshape %input : tensor<5x24x48xf32> -> tensor<60x2x48xf32>
+ // CHECK: %[[T2:.+]] = stream.tensor.clone %[[T1]] : tensor<60x2x48xf32> in !stream.resource<*>{%[[T1_SIZE]]} -> tensor<60x2x48xf32> in !stream.resource<*>{%[[T1_SIZE]]}
+ %3 = flow.tensor.clone %2 : tensor<60x2x48xf32>
+ // CHECK: %[[T3_SIZE:.+]] = stream.tensor.sizeof tensor<30x2x96xf32> : index
+ // CHECK: %[[T3:.+]] = stream.tensor.clone %[[T0]] : tensor<5x24x48xf32> in !stream.resource<*>{%[[INPUT_SIZE]]} -> tensor<30x2x96xf32> in !stream.resource<*>{%[[T3_SIZE]]}
+ %4 = flow.tensor.reshape %1 : tensor<5x24x48xf32> -> tensor<30x2x96xf32>
+ // CHECK: return %[[T2]], %[[T1_SIZE]], %[[T3]], %[[T3_SIZE]] : !stream.resource<*>, index, !stream.resource<*>, index
+ return %3, %4 : tensor<60x2x48xf32>, tensor<30x2x96xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @tensorSplat
+// CHECK-SAME: (%[[VALUE:.+]]: i8, %[[DIM0:.+]]: index)
+func @tensorSplat(%value: i8, %dim0: index) -> tensor<?x128xi8> {
+ // CHECK: %[[T_SIZE:.+]] = stream.tensor.sizeof tensor<?x128xi8>{%[[DIM0]]} : index
+ // CHECK: %[[T:.+]] = stream.tensor.splat %[[VALUE]] : i8 -> tensor<?x128xi8>{%[[DIM0]]} in !stream.resource<*>{%[[T_SIZE]]}
+ %0 = flow.tensor.splat %value : tensor<?x128xi8>{%dim0}
+ // CHECK: return %[[T]], %[[T_SIZE]]
+ return %0 : tensor<?x128xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @tensorSlice
+// CHECK-SAME: (%[[INPUT:.+]]: !stream.resource<*>, %[[INPUT_SIZE:.+]]: index)
+func @tensorSlice(%input : tensor<5x24x48xf32>) -> tensor<3x24x48xf32> {
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %c24 = arith.constant 24 : index
+ %c48 = arith.constant 48 : index
+ // CHECK: %[[T_SIZE:.+]] = stream.tensor.sizeof tensor<3x24x48xf32> : index
+ // CHECK: %[[T:.+]] = stream.tensor.slice %[[INPUT]][%c2, %c0, %c0 for %c3, %c24, %c48] : tensor<5x24x48xf32> in !stream.resource<*>{%[[INPUT_SIZE]]} -> tensor<3x24x48xf32> in !stream.resource<*>{%[[T_SIZE]]}
+ %0 = flow.tensor.slice %input[%c2, %c0, %c0 for %c3, %c24, %c48] : tensor<5x24x48xf32> -> tensor<3x24x48xf32>
+ // CHECK: return %[[T]], %[[T_SIZE]] : !stream.resource<*>, index
+ return %0 : tensor<3x24x48xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @tensorUpdate
+// CHECK-SAME: (%[[UPDATE:.+]]: !stream.resource<*>, %[[UPDATE_SIZE:.+]]: index, %[[TARGET:.+]]: !stream.resource<*>, %[[TARGET_SIZE:.+]]: index)
+func @tensorUpdate(%update : tensor<1x1x10xf32>, %target : tensor<5x1x10xf32>) -> tensor<5x1x10xf32> {
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ // CHECK: %[[T:.+]] = stream.tensor.update %[[UPDATE]], %[[TARGET]][%c4, %c1, %c1] : tensor<1x1x10xf32> in !stream.resource<*>{%[[UPDATE_SIZE]]} -> tensor<5x1x10xf32> in %[[TARGET]] as !stream.resource<*>{%[[TARGET_SIZE]]}
+ %0 = flow.tensor.update %update, %target[%c4, %c1, %c1] : tensor<1x1x10xf32> -> %target as tensor<5x1x10xf32>
+ // CHECK: return %[[T]], %[[TARGET_SIZE]] : !stream.resource<*>, index
+ return %0 : tensor<5x1x10xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @tensorLoad
+// CHECK-SAME: (%[[SOURCE:.+]]: !stream.resource<*>, %[[SOURCE_SIZE:.+]]: index)
+func @tensorLoad(%source : tensor<2x3xi32>) -> i32 {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ // CHECK: %[[T0:.+]] = stream.async.transfer %[[SOURCE]] : !stream.resource<*>{%[[SOURCE_SIZE]]} -> !stream.resource<staging>{%[[SOURCE_SIZE]]}
+ // CHECK: %[[T1:.+]] = stream.tensor.load %[[T0]][%c0, %c1] : tensor<2x3xi32> in !stream.resource<staging>{%[[SOURCE_SIZE]]} -> i32
+ %0 = flow.tensor.load %source[%c0, %c1] : tensor<2x3xi32>
+ // CHECK: return %[[T1]]
+ return %0 : i32
+}
+
+// -----
+
+// CHECK-LABEL: @tensorStore
+// CHECK-SAME: (%[[TARGET:.+]]: !stream.resource<*>, %[[TARGET_SIZE:.+]]: index)
+func @tensorStore(%target : tensor<2x3xi32>) -> tensor<2x3xi32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c9 = arith.constant 9 : i32
+ // CHECK: %[[T0:.+]] = stream.async.transfer %[[TARGET]] : !stream.resource<*>{%[[TARGET_SIZE]]} -> !stream.resource<staging>{%[[TARGET_SIZE]]}
+ // CHECK: %[[T1:.+]] = stream.tensor.store %c9_i32, %[[T0]][%c0, %c1] : i32 -> tensor<2x3xi32> in %[[T0]] as !stream.resource<staging>{%[[TARGET_SIZE]]}
+ // CHECK: %[[T2:.+]] = stream.async.transfer %[[T1]] : !stream.resource<staging>{%[[TARGET_SIZE]]} -> !stream.resource<*>{%[[TARGET_SIZE]]}
+ %0 = flow.tensor.store %c9, %target[%c0, %c1] : tensor<2x3xi32>
+ // CHECK: return %[[T2]]
+ return %0 : tensor<2x3xi32>
+}
diff --git a/iree/compiler/Dialect/Stream/Conversion/PatternUtils.cpp b/iree/compiler/Dialect/Stream/Conversion/PatternUtils.cpp
new file mode 100644
index 0000000..202040f
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/PatternUtils.cpp
@@ -0,0 +1,46 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Stream/Conversion/PatternUtils.h"
+
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+ConvertedTensor consumeTensorOperand(Location loc, Value operand,
+ OpBuilder &builder) {
+ auto operandType = operand.getType();
+ if (operandType.isa<IREE::Stream::ResourceType>()) {
+ // 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 {
+ operand,
+ builder.createOrFold<IREE::Stream::ResourceSizeOp>(
+ loc, builder.getIndexType(), operand),
+ };
+ } else if (auto castOp =
+ operand.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 {
+ castOp.getOperand(0),
+ castOp.getOperand(1),
+ };
+ }
+ llvm_unreachable(
+ "unexpected operand; should have either been converted from tensor or "
+ "not");
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Stream/Conversion/PatternUtils.h b/iree/compiler/Dialect/Stream/Conversion/PatternUtils.h
new file mode 100644
index 0000000..996efd2
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/PatternUtils.h
@@ -0,0 +1,35 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_DIALECT_STREAM_CONVERSION_PATTERN_UTILS_H_
+#define IREE_COMPILER_DIALECT_STREAM_CONVERSION_PATTERN_UTILS_H_
+
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+// https://reviews.llvm.org/D111620 broke 1->N type expansion during dialect
+// conversion. It inserts unrealized_conversion_casts but then passes the
+// illegal source dialect types for pattern operands, meaning that even though
+// we say tensors are illegal the patterns get the new remapped values as
+// tensors. This, naturally, breaks everything. To work around this we have this
+// helper that tries to peek through the unrealized_conversion_casts and get out
+// the actual values we expected to see from the conversion (and did before that
+// change).
+struct ConvertedTensor {
+ Value resource;
+ Value resourceSize;
+};
+ConvertedTensor consumeTensorOperand(Location loc, Value operand,
+ OpBuilder &builder);
+
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_DIALECT_STREAM_CONVERSION_PATTERN_UTILS_H_
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/BUILD b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/BUILD
new file mode 100644
index 0000000..cf4978f
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/BUILD
@@ -0,0 +1,36 @@
+# Copyright 2021 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+cc_library(
+ name = "StandardToStream",
+ srcs = [
+ "ConvertConstantOps.cpp",
+ "ConvertStandardToStream.cpp",
+ "ConvertStructuralOps.cpp",
+ ],
+ hdrs = [
+ "ConvertStandardToStream.h",
+ ],
+ deps = [
+ "//iree/compiler/Dialect/Stream/Conversion",
+ "//iree/compiler/Dialect/Stream/IR",
+ "@llvm-project//llvm:Support",
+ "@llvm-project//mlir:ArithmeticDialect",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:MemRefDialect",
+ "@llvm-project//mlir:Pass",
+ "@llvm-project//mlir:Shape",
+ "@llvm-project//mlir:StandardOps",
+ "@llvm-project//mlir:TensorDialect",
+ "@llvm-project//mlir:Transforms",
+ ],
+)
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/CMakeLists.txt
new file mode 100644
index 0000000..3cea8dd
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/CMakeLists.txt
@@ -0,0 +1,37 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# iree/compiler/Dialect/Stream/Conversion/StandardToStream/BUILD #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_cc_library(
+ NAME
+ StandardToStream
+ HDRS
+ "ConvertStandardToStream.h"
+ SRCS
+ "ConvertConstantOps.cpp"
+ "ConvertStandardToStream.cpp"
+ "ConvertStructuralOps.cpp"
+ DEPS
+ LLVMSupport
+ MLIRArithmetic
+ MLIRIR
+ MLIRMemRef
+ MLIRPass
+ MLIRShape
+ MLIRStandard
+ MLIRTensor
+ MLIRTransforms
+ iree::compiler::Dialect::Stream::Conversion
+ iree::compiler::Dialect::Stream::IR
+ PUBLIC
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertConstantOps.cpp b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertConstantOps.cpp
new file mode 100644
index 0000000..2a69d27
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertConstantOps.cpp
@@ -0,0 +1,63 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStandardToStream.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace {
+
+struct ConvertTensorConstantOp : public OpConversionPattern<arith::ConstantOp> {
+ public:
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ arith::ConstantOp constantOp, llvm::ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ // Only handle tensor types - other arith.constant types (like i32) are
+ // ignored.
+ if (!constantOp.getType().isa<TensorType>()) return failure();
+
+ Type constantType = IREE::Stream::ResourceType::get(
+ getContext(), IREE::Stream::Lifetime::Constant);
+ auto newOp = rewriter.create<IREE::Stream::TensorConstantOp>(
+ constantOp.getLoc(), constantType,
+ constantOp.value().cast<ElementsAttr>(),
+ TypeAttr::get(constantOp.getType()),
+ /*result_encoding_dims=*/ValueRange{},
+ /*affinity=*/nullptr);
+
+ Type unknownType = IREE::Stream::ResourceType::get(getContext());
+ auto constantSize = rewriter.createOrFold<IREE::Stream::ResourceSizeOp>(
+ constantOp.getLoc(), rewriter.getIndexType(), newOp.result());
+ rewriter.replaceOpWithNewOp<IREE::Stream::AsyncTransferOp>(
+ constantOp, unknownType, newOp.result(), constantSize, constantSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr);
+ return success();
+ }
+};
+
+} // namespace
+
+void populateStandardConstantToStreamPatterns(
+ MLIRContext *context, ConversionTarget &conversionTarget,
+ TypeConverter &typeConverter, OwningRewritePatternList &patterns) {
+ conversionTarget.addDynamicallyLegalOp<arith::ConstantOp>(
+ [](arith::ConstantOp op) { return !op.getType().isa<TensorType>(); });
+
+ patterns.insert<ConvertTensorConstantOp>(typeConverter, context);
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStandardToStream.cpp b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStandardToStream.cpp
new file mode 100644
index 0000000..9534d53
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStandardToStream.cpp
@@ -0,0 +1,48 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStandardToStream.h"
+
+#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+void populateStandardConstantToStreamPatterns(
+ MLIRContext *context, ConversionTarget &conversionTarget,
+ TypeConverter &typeConverter, OwningRewritePatternList &patterns);
+
+void populateStandardStructuralToStreamPatterns(
+ MLIRContext *context, ConversionTarget &conversionTarget,
+ TypeConverter &typeConverter, OwningRewritePatternList &patterns);
+
+void populateStandardToStreamConversionPatterns(
+ MLIRContext *context, ConversionTarget &conversionTarget,
+ TypeConverter &typeConverter, OwningRewritePatternList &patterns) {
+ typeConverter.addConversion([](IndexType type) { return type; });
+ typeConverter.addConversion([](IntegerType type) { return type; });
+ typeConverter.addConversion([](FloatType type) { return type; });
+
+ // Ensure all shape related ops are fully converted as we should no longer
+ // have any types they are valid to be used on after this conversion.
+ conversionTarget.addIllegalOp<memref::DimOp>();
+ conversionTarget.addIllegalOp<mlir::RankOp>();
+ conversionTarget.addIllegalOp<tensor::DimOp>();
+
+ populateStandardConstantToStreamPatterns(context, conversionTarget,
+ typeConverter, patterns);
+ populateStandardStructuralToStreamPatterns(context, conversionTarget,
+ typeConverter, patterns);
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStandardToStream.h b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStandardToStream.h
new file mode 100644
index 0000000..12997c6
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStandardToStream.h
@@ -0,0 +1,26 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_DIALECT_STREAM_CONVERSION_STANDARDTOSTREAM_CONVERTSTANDARDTOSTREAM_H_
+#define IREE_COMPILER_DIALECT_STREAM_CONVERSION_STANDARDTOSTREAM_CONVERTSTANDARDTOSTREAM_H_
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+// Populates conversion patterns that perform standard/builtin->stream
+// conversion. These patterns ensure that nested types are run through the
+// provided |typeConverter|.
+void populateStandardToStreamConversionPatterns(
+ MLIRContext *context, ConversionTarget &conversionTarget,
+ TypeConverter &typeConverter, OwningRewritePatternList &patterns);
+
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_DIALECT_STREAM_CONVERSION_STANDARDTOSTREAM_CONVERTSTANDARDTOSTREAM_H_
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStructuralOps.cpp b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStructuralOps.cpp
new file mode 100644
index 0000000..6a050c5
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStructuralOps.cpp
@@ -0,0 +1,263 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Stream/Conversion/PatternUtils.h"
+#include "iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStandardToStream.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace {
+
+struct FuncOpSignatureConversion : public OpConversionPattern<mlir::FuncOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ mlir::FuncOp funcOp, llvm::ArrayRef<Value> operands,
+ ConversionPatternRewriter &rewriter) const override {
+ auto &typeConverter = *getTypeConverter();
+
+ // Convert the input signature types.
+ // TODO(benvanik): dynamic shapes by passing in tensor dynamic dims.
+ auto originalType = funcOp.getType();
+ TypeConverter::SignatureConversion newSignature(
+ originalType.getNumInputs());
+ for (auto argType : llvm::enumerate(originalType.getInputs())) {
+ if (failed(typeConverter.convertSignatureArg(
+ argType.index(), argType.value(), newSignature))) {
+ return failure();
+ }
+ }
+ SmallVector<Type, 4> newResultTypes;
+ if (failed(typeConverter.convertTypes(originalType.getResults(),
+ newResultTypes))) {
+ return failure();
+ }
+
+ // Replace function.
+ auto newFuncOp = rewriter.cloneWithoutRegions(funcOp);
+ newFuncOp.getBlocks().clear();
+ rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
+ newFuncOp.end());
+ newFuncOp.setType(rewriter.getFunctionType(newSignature.getConvertedTypes(),
+ newResultTypes));
+ if (failed(rewriter.convertRegionTypes(&newFuncOp.getBody(), typeConverter,
+ &newSignature))) {
+ return failure();
+ }
+
+ rewriter.eraseOp(funcOp);
+ return success();
+ }
+};
+
+static SmallVector<Value> expandResourceOperands(
+ Location loc, ValueRange operands, ConversionPatternRewriter &rewriter) {
+ SmallVector<Value> expandedOperands;
+ expandedOperands.reserve(operands.size());
+ for (auto operand : operands) {
+ if (operand.getType().isa<TensorType>()) {
+ auto value = consumeTensorOperand(loc, operand, rewriter);
+ expandedOperands.push_back(value.resource);
+ expandedOperands.push_back(value.resourceSize);
+ } else if (operand.getType().isa<IREE::Stream::ResourceType>()) {
+ expandedOperands.push_back(operand);
+ expandedOperands.push_back(
+ rewriter.createOrFold<IREE::Stream::ResourceSizeOp>(loc, operand));
+ } else {
+ expandedOperands.push_back(operand);
+ }
+ }
+ return expandedOperands;
+}
+
+struct CallOpConversion : public OpConversionPattern<mlir::CallOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ mlir::CallOp op, llvm::ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ // Expand any resource operands to resource + size.
+ auto expandedOperands =
+ expandResourceOperands(op.getLoc(), newOperands, rewriter);
+
+ // Expand any resource results to resource + size.
+ SmallVector<Type> expandedTypes;
+ struct Result {
+ size_t originalIndex;
+ size_t newIndex;
+ Type newType;
+ };
+ SmallVector<Result> resultMap;
+ for (auto originalType : llvm::enumerate(op.getResultTypes())) {
+ SmallVector<Type> newTypes;
+ if (failed(getTypeConverter()->convertType(originalType.value(),
+ newTypes))) {
+ return rewriter.notifyMatchFailure(op,
+ "unable to convert result types");
+ }
+ resultMap.push_back(
+ Result{originalType.index(), expandedTypes.size(), newTypes.front()});
+ expandedTypes.append(newTypes);
+ }
+
+ // Create a new call that takes the expanded input operands and returns the
+ // expanded output results. We can't directly replace the original call as
+ // the result counts differ.
+ auto callOp = rewriter.create<mlir::CallOp>(op.getLoc(), expandedTypes,
+ op.callee(), expandedOperands);
+
+ // Tie all resource results together so we end up with 1:1 results with the
+ // original op.
+ SmallVector<Value> results;
+ for (auto result : resultMap) {
+ if (result.newType.isa<IREE::Stream::ResourceType>()) {
+ auto oldType = op.getResult(result.originalIndex).getType();
+ auto resource = callOp.getResult(result.newIndex + 0);
+ auto resourceSize = callOp.getResult(result.newIndex + 1);
+ results.push_back(rewriter
+ .create<mlir::UnrealizedConversionCastOp>(
+ op.getLoc(), TypeRange{oldType},
+ ValueRange{resource, resourceSize})
+ .getResult(0));
+ } else {
+ results.push_back(callOp.getResult(result.newIndex));
+ }
+ }
+ rewriter.replaceOp(op, results);
+
+ return success();
+ }
+};
+
+struct ReturnOpConversion : public OpConversionPattern<mlir::ReturnOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ mlir::ReturnOp op, llvm::ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ // Expand any resource operands to resource + size.
+ auto expandedOperands =
+ expandResourceOperands(op.getLoc(), newOperands, rewriter);
+ rewriter.replaceOpWithNewOp<mlir::ReturnOp>(op, expandedOperands);
+ return success();
+ }
+};
+
+struct BranchOpConversion : public OpConversionPattern<mlir::BranchOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ mlir::BranchOp op, llvm::ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ mlir::BranchOpAdaptor operands(newOperands);
+ // Expand any resource operands to resource + size.
+ auto expandedOperands =
+ expandResourceOperands(op.getLoc(), operands.destOperands(), rewriter);
+ rewriter.replaceOpWithNewOp<mlir::BranchOp>(op, op.dest(),
+ expandedOperands);
+ return success();
+ }
+};
+
+struct CondBranchOpConversion : public OpConversionPattern<mlir::CondBranchOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ mlir::CondBranchOp op, llvm::ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ mlir::CondBranchOpAdaptor operands(newOperands,
+ op.getOperation()->getAttrDictionary());
+ // Expand any resource operands to resource + size.
+ auto trueDestOperands = expandResourceOperands(
+ op.getLoc(), operands.trueDestOperands(), rewriter);
+ auto falseDestOperands = expandResourceOperands(
+ op.getLoc(), operands.falseDestOperands(), rewriter);
+ rewriter.replaceOpWithNewOp<mlir::CondBranchOp>(
+ op, operands.condition(), op.trueDest(), trueDestOperands,
+ op.falseDest(), falseDestOperands);
+ return success();
+ }
+};
+
+struct SelectOpConversion : public OpConversionPattern<mlir::SelectOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(
+ mlir::SelectOp op, mlir::SelectOp::Adaptor operands,
+ ConversionPatternRewriter &rewriter) const override {
+ // Only handle selects where the operands are tensors (resources).
+ if (!op.true_value().getType().isa<TensorType>()) return failure();
+ auto trueOperand =
+ consumeTensorOperand(op.getLoc(), operands.true_value(), rewriter);
+ auto falseOperand =
+ consumeTensorOperand(op.getLoc(), operands.false_value(), rewriter);
+ auto resourceSelectOp = rewriter.create<mlir::SelectOp>(
+ op.getLoc(), operands.condition(), trueOperand.resource,
+ falseOperand.resource);
+ auto sizeSelectOp = rewriter.create<mlir::SelectOp>(
+ op.getLoc(), operands.condition(), trueOperand.resourceSize,
+ falseOperand.resourceSize);
+ rewriter.replaceOpWithNewOp<mlir::UnrealizedConversionCastOp>(
+ op, operands.true_value().getType(),
+ ValueRange{resourceSelectOp.result(), sizeSelectOp.result()});
+ return success();
+ }
+};
+
+} // namespace
+
+void populateStandardStructuralToStreamPatterns(
+ MLIRContext *context, ConversionTarget &conversionTarget,
+ TypeConverter &typeConverter, OwningRewritePatternList &patterns) {
+ conversionTarget.addLegalOp<mlir::ModuleOp>();
+
+ // We need to rewrite certain types on operands/results so use the default
+ // dynamic legality checker to force any ops using such types to run through
+ // our patterns.
+
+ conversionTarget.addDynamicallyLegalOp<mlir::FuncOp>([&](mlir::FuncOp op) {
+ return typeConverter.isSignatureLegal(op.getType()) &&
+ typeConverter.isLegal(&op.getBody());
+ });
+ conversionTarget.addDynamicallyLegalOp<mlir::CallOp>([&](mlir::CallOp op) {
+ return llvm::all_of(
+ op.getOperandTypes(),
+ [&](Type type) { return typeConverter.isLegal(type); }) &&
+ llvm::all_of(op.getResultTypes(),
+ [&](Type type) { return typeConverter.isLegal(type); });
+ });
+ conversionTarget.addDynamicallyLegalOp<mlir::ReturnOp>(
+ [&](mlir::ReturnOp op) {
+ return llvm::all_of(op.getOperandTypes(), [&](Type type) {
+ return typeConverter.isLegal(type);
+ });
+ });
+
+ conversionTarget.addDynamicallyLegalOp<mlir::BranchOp>(
+ [&](mlir::BranchOp op) {
+ return llvm::all_of(op.getOperandTypes(), [&](Type type) {
+ return typeConverter.isLegal(type);
+ });
+ });
+ conversionTarget.addDynamicallyLegalOp<mlir::CondBranchOp>(
+ [&](mlir::CondBranchOp op) {
+ return llvm::all_of(op.getOperandTypes(), [&](Type type) {
+ return typeConverter.isLegal(type);
+ });
+ });
+
+ patterns
+ .insert<FuncOpSignatureConversion, CallOpConversion, ReturnOpConversion,
+ BranchOpConversion, CondBranchOpConversion, SelectOpConversion>(
+ typeConverter, context);
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/BUILD b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/BUILD
new file mode 100644
index 0000000..ca9fdbc
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/BUILD
@@ -0,0 +1,29 @@
+# Copyright 2021 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+iree_lit_test_suite(
+ name = "lit",
+ srcs = enforce_glob(
+ [
+ "constant_ops.mlir",
+ "structural_ops.mlir",
+ ],
+ include = ["*.mlir"],
+ ),
+ data = [
+ "//iree/tools:IreeFileCheck",
+ "//iree/tools:iree-opt",
+ ],
+)
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/CMakeLists.txt
new file mode 100644
index 0000000..32bd14d
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/CMakeLists.txt
@@ -0,0 +1,24 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/BUILD #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_lit_test_suite(
+ NAME
+ lit
+ SRCS
+ "constant_ops.mlir"
+ "structural_ops.mlir"
+ DATA
+ iree::tools::IreeFileCheck
+ iree::tools::iree-opt
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/constant_ops.mlir b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/constant_ops.mlir
new file mode 100644
index 0000000..ef2e038
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/constant_ops.mlir
@@ -0,0 +1,10 @@
+// RUN: iree-opt -split-input-file -iree-stream-conversion %s | IreeFileCheck %s
+
+// CHECK-LABEL: @constantTensor
+func @constantTensor() {
+ // CHECK: %[[CST:.+]] = stream.tensor.constant : tensor<2xi32> in !stream.resource<constant> = dense<[1, 2]> : tensor<2xi32>
+ // CHECK: %[[SIZE:.+]] = stream.resource.size %[[CST]] : !stream.resource<constant>
+ // CHECK: %[[T:.+]] = stream.async.transfer %[[CST]] : !stream.resource<constant>{%[[SIZE]]} -> !stream.resource<*>{%[[SIZE]]}
+ %0 = arith.constant dense<[1, 2]> : tensor<2xi32>
+ return
+}
diff --git a/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/structural_ops.mlir b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/structural_ops.mlir
new file mode 100644
index 0000000..4f94561
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/StandardToStream/test/structural_ops.mlir
@@ -0,0 +1,64 @@
+// RUN: iree-opt -split-input-file -iree-stream-conversion %s | IreeFileCheck %s
+
+// CHECK-LABEL: @functionExpansion
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<*>, %[[ARG0_SIZE:.+]]: index,
+// CHECK-SAME: %[[ARG1:.+]]: i1,
+// CHECK-SAME: %[[ARG2:.+]]: !stream.resource<*>, %[[ARG2_SIZE:.+]]: index)
+// CHECK-SAME: -> (!stream.resource<*>, index, i1, !stream.resource<*>, index)
+func @functionExpansion(%arg0: tensor<4x?xf32>, %arg1: i1, %arg2: tensor<i32>)
+ -> (tensor<4x?xf32>, i1, tensor<i32>) {
+ // CHECK-NEXT: %[[RET:.+]]:5 = call @callee(%[[ARG0]], %[[ARG0_SIZE]], %[[ARG1]], %[[ARG2]], %[[ARG2_SIZE]])
+ // CHECK-SAME: : (!stream.resource<*>, index, i1, !stream.resource<*>, index) -> (!stream.resource<*>, index, i1, !stream.resource<*>, index)
+ %0:3 = call @callee(%arg0, %arg1, %arg2) : (tensor<4x?xf32>, i1, tensor<i32>) -> (tensor<4x?xf32>, i1, tensor<i32>)
+ // CHECK: return %[[RET]]#0, %[[RET]]#1, %[[RET]]#2, %[[RET]]#3, %[[RET]]#4 : !stream.resource<*>, index, i1, !stream.resource<*>, index
+ return %0#0, %0#1, %0#2 : tensor<4x?xf32>, i1, tensor<i32>
+}
+
+// CHECK: func private @callee
+func private @callee(%arg0: tensor<4x?xf32>, %arg1: i1, %arg2: tensor<i32>)
+ -> (tensor<4x?xf32>, i1, tensor<i32>)
+
+// -----
+
+// CHECK-LABEL: @brExpansion
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<*>, %[[ARG0_SIZE:.+]]: index, %arg2: i1)
+// CHECK-SAME: -> (!stream.resource<*>, index, i1)
+func @brExpansion(%arg0: tensor<1xf32>, %arg1: i1) -> (tensor<1xf32>, i1) {
+ // CHECK: br ^bb1(%[[ARG0]], %[[ARG0_SIZE]], %arg2 : !stream.resource<*>, index, i1)
+ br ^bb1(%arg0, %arg1 : tensor<1xf32>, i1)
+// CHECK: ^bb1(%[[BB_ARG0:.+]]: !stream.resource<*>, %[[BB_ARG1:.+]]: index, %[[BB_ARG2:.+]]: i1):
+^bb1(%0: tensor<1xf32>, %1: i1):
+ // CHECK: return %[[BB_ARG0]], %[[BB_ARG1]], %[[BB_ARG2]] : !stream.resource<*>, index, i1
+ return %0, %1 : tensor<1xf32>, i1
+}
+
+// -----
+
+// CHECK-LABEL: @condBrExpansion
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<*>, %[[ARG0_SIZE:.+]]: index,
+// CHECK-SAME: %[[ARG1:.+]]: !stream.resource<*>, %[[ARG1_SIZE:.+]]: index)
+// CHECK-SAME: -> (!stream.resource<*>, index)
+func @condBrExpansion(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
+ %true = arith.constant 1 : i1
+ // CHECK: cond_br %true,
+ // CHECK-SAME: ^bb1(%[[ARG0]], %[[ARG0_SIZE]] : !stream.resource<*>, index),
+ // CHECK-SAME: ^bb1(%[[ARG1]], %[[ARG1_SIZE]] : !stream.resource<*>, index)
+ cond_br %true, ^bb1(%arg0 : tensor<1xf32>), ^bb1(%arg1 : tensor<1xf32>)
+^bb1(%0: tensor<1xf32>):
+ return %0 : tensor<1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @selectExpansion
+// CHECK-SAME: (%[[ARG0:.+]]: !stream.resource<*>, %[[ARG0_SIZE:.+]]: index,
+// CHECK-SAME: %[[COND:.+]]: i1,
+// CHECK-SAME: %[[ARG1:.+]]: !stream.resource<*>, %[[ARG1_SIZE:.+]]: index)
+// CHECK-SAME: -> (!stream.resource<*>, index)
+func @selectExpansion(%arg0: tensor<1xf32>, %cond: i1, %arg1: tensor<1xf32>) -> tensor<1xf32> {
+ // CHECK-DAG: %[[RET:.+]] = select %[[COND]], %[[ARG0]], %[[ARG1]] : !stream.resource<*>
+ // CHECK-DAG: %[[RET_SIZE:.+]] = select %[[COND]], %[[ARG0_SIZE]], %[[ARG1_SIZE]] : index
+ %0 = select %cond, %arg0, %arg1 : tensor<1xf32>
+ // CHECK: return %[[RET]], %[[RET_SIZE]] : !stream.resource<*>, index
+ return %0 : tensor<1xf32>
+}
diff --git a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/BUILD b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/BUILD
new file mode 100644
index 0000000..e880f5b
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/BUILD
@@ -0,0 +1,29 @@
+# Copyright 2021 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+cc_library(
+ name = "UtilToStream",
+ srcs = [
+ "ConvertUtilToStream.cpp",
+ ],
+ hdrs = [
+ "ConvertUtilToStream.h",
+ ],
+ deps = [
+ "//iree/compiler/Dialect/Stream/Conversion",
+ "//iree/compiler/Dialect/Stream/IR",
+ "//iree/compiler/Dialect/Util/IR",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:StandardOps",
+ "@llvm-project//mlir:Transforms",
+ ],
+)
diff --git a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/CMakeLists.txt
new file mode 100644
index 0000000..274da6a
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/CMakeLists.txt
@@ -0,0 +1,30 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# iree/compiler/Dialect/Stream/Conversion/UtilToStream/BUILD #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_cc_library(
+ NAME
+ UtilToStream
+ HDRS
+ "ConvertUtilToStream.h"
+ SRCS
+ "ConvertUtilToStream.cpp"
+ DEPS
+ MLIRIR
+ MLIRStandard
+ MLIRTransforms
+ iree::compiler::Dialect::Stream::Conversion
+ iree::compiler::Dialect::Stream::IR
+ iree::compiler::Dialect::Util::IR
+ PUBLIC
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/ConvertUtilToStream.cpp b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/ConvertUtilToStream.cpp
new file mode 100644
index 0000000..d1a1c13
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/ConvertUtilToStream.cpp
@@ -0,0 +1,279 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Stream/Conversion/UtilToStream/ConvertUtilToStream.h"
+
+#include "iree/compiler/Dialect/Stream/Conversion/PatternUtils.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Globals
+//===----------------------------------------------------------------------===//
+
+struct ExpandedGlobalResource {
+ IREE::Util::GlobalOp resourceOp;
+ IREE::Util::GlobalOp resourceSizeOp;
+};
+
+struct GlobalExpansionState {
+ // A map of original symbol name to one new global for each expanded type.
+ DenseMap<StringRef, ExpandedGlobalResource> globalMap;
+};
+
+static bool isExpandedType(Type type) {
+ if (type.isa<TensorType>()) return true;
+ if (auto ptrType = type.dyn_cast<IREE::Util::PtrType>()) {
+ return isExpandedType(ptrType);
+ }
+ return false;
+}
+
+template <typename T>
+class BaseGlobalConversionPattern : public OpConversionPattern<T> {
+ public:
+ BaseGlobalConversionPattern(
+ std::shared_ptr<GlobalExpansionState> expansionState,
+ TypeConverter &typeConverter, MLIRContext *context,
+ PatternBenefit benefit = 1)
+ : OpConversionPattern<T>(typeConverter, context, benefit),
+ expansionState(std::move(expansionState)) {}
+
+ protected:
+ mutable std::shared_ptr<GlobalExpansionState> expansionState;
+};
+
+struct GlobalOpExpansion
+ : public BaseGlobalConversionPattern<IREE::Util::GlobalOp> {
+ using BaseGlobalConversionPattern::BaseGlobalConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Util::GlobalOp globalOp, llvm::ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ // Only apply to expanded types (tensors/etc).
+ if (!isExpandedType(globalOp.type())) return failure();
+
+ SmallVector<Type> newTypes;
+ if (failed(getTypeConverter()->convertType(globalOp.type(), newTypes))) {
+ return rewriter.notifyMatchFailure(globalOp,
+ "failed to convert ptr type");
+ }
+ if (newTypes.size() == 1 && newTypes.front() == globalOp.type()) {
+ return rewriter.notifyMatchFailure(globalOp, "no conversion needed");
+ }
+
+ // Start with the appropriate type. Lifetime refinement will use this as a
+ // seed. Note that what was a constant in earlier dialects becomes a mutable
+ // global holding a resource that may have constant contents.
+ bool hasConstantUsage = !globalOp.isMutable();
+ auto resourceType = IREE::Stream::ResourceType::get(
+ rewriter.getContext(), hasConstantUsage
+ ? IREE::Stream::Lifetime::Constant
+ : IREE::Stream::Lifetime::Variable);
+
+ // Special handling of the initial value: if it's a tensor then we need to
+ // materialize an initializer and initialization ops. This allows the
+ // current conversion to pick up the expanded initialization ops.
+ auto initialValue = globalOp.initial_valueAttr();
+ bool tensorInitializerRequired =
+ initialValue ? initialValue.getType().isa<TensorType>() : false;
+
+ // New global holding the initial value only if it is not a tensor type.
+ auto resourceOp = rewriter.replaceOpWithNewOp<IREE::Util::GlobalOp>(
+ globalOp, globalOp.getName(), globalOp.is_mutable(), resourceType,
+ initialValue && !tensorInitializerRequired
+ ? llvm::Optional<Attribute>{initialValue}
+ : llvm::None);
+ resourceOp.setVisibility(globalOp.getVisibility());
+
+ // NOTE: we ignore noinline here, possibly to our peril. In earlier dialects
+ // noinline indicates that the constant value should not be inlined, while
+ // here it would be indicating the reference to the constant value should
+ // not be (and that's weird).
+
+ // Also create a global for tracking the resource size. In many cases this
+ // is constant and will fold throughout the program. Global optimizations
+ // such as same-value deduplication will also take effect.
+ auto indexType = rewriter.getIndexType();
+ auto resourceSizeOp = rewriter.create<IREE::Util::GlobalOp>(
+ globalOp.getLoc(), (globalOp.getName() + "__size").str(),
+ globalOp.is_mutable(), indexType, Optional<Attribute>{});
+ resourceSizeOp.setVisibility(globalOp.getVisibility());
+
+ // Materialize the initializer if we need to setup a tensor-like constant.
+ if (tensorInitializerRequired) {
+ auto initializerOp =
+ rewriter.create<IREE::Util::InitializerOp>(globalOp.getLoc());
+ auto *entryBlock = rewriter.createBlock(&initializerOp.getBody());
+ rewriter.setInsertionPointToStart(entryBlock);
+ auto constantOp = rewriter.create<IREE::Stream::TensorConstantOp>(
+ globalOp.getLoc(), resourceOp.type(),
+ initialValue.cast<ElementsAttr>(), TypeAttr::get(globalOp.type()),
+ /*result_encoding_dims=*/ValueRange{}, /*affinity=*/nullptr);
+ auto constantSizeOp = rewriter.create<IREE::Stream::ResourceSizeOp>(
+ globalOp.getLoc(), indexType, constantOp.result());
+ rewriter.create<IREE::Util::GlobalStoreOp>(
+ globalOp.getLoc(), constantOp.result(), resourceOp.getSymbolName());
+ rewriter.create<IREE::Util::GlobalStoreOp>(
+ globalOp.getLoc(), constantSizeOp.result(),
+ resourceSizeOp.getSymbolName());
+ rewriter.create<IREE::Util::InitializerReturnOp>(globalOp.getLoc());
+ }
+
+ expansionState->globalMap[globalOp.getSymbolName()] =
+ ExpandedGlobalResource{
+ resourceOp,
+ resourceSizeOp,
+ };
+
+ return success();
+ }
+};
+
+struct GlobalLoadOpExpansion
+ : public BaseGlobalConversionPattern<IREE::Util::GlobalLoadOp> {
+ using BaseGlobalConversionPattern::BaseGlobalConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Util::GlobalLoadOp loadOp, llvm::ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Util::GlobalLoadOpAdaptor operands(newOperands,
+ loadOp->getAttrDictionary());
+
+ // Only apply to expanded types (tensors/etc).
+ if (!isExpandedType(loadOp.getType())) return failure();
+ auto &expandedGlobal =
+ expansionState->globalMap[operands.global().getValue()];
+
+ // Insert a load/transfer to the unknown resource lifetime.
+ auto unknownType = IREE::Stream::ResourceType::get(rewriter.getContext());
+ auto resource = rewriter
+ .create<IREE::Util::GlobalLoadOp>(
+ loadOp.getLoc(), expandedGlobal.resourceOp.type(),
+ expandedGlobal.resourceOp.getSymbolName())
+ .result();
+ auto resourceSize = rewriter
+ .create<IREE::Util::GlobalLoadOp>(
+ loadOp.getLoc(), rewriter.getIndexType(),
+ expandedGlobal.resourceSizeOp.getSymbolName())
+ .result();
+ rewriter.replaceOpWithNewOp<IREE::Stream::AsyncTransferOp>(
+ loadOp, unknownType, resource, resourceSize, resourceSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr);
+
+ return success();
+ }
+};
+
+struct GlobalStoreOpExpansion
+ : public BaseGlobalConversionPattern<IREE::Util::GlobalStoreOp> {
+ using BaseGlobalConversionPattern::BaseGlobalConversionPattern;
+ LogicalResult matchAndRewrite(
+ IREE::Util::GlobalStoreOp storeOp, llvm::ArrayRef<Value> newOperands,
+ ConversionPatternRewriter &rewriter) const override {
+ IREE::Util::GlobalStoreOpAdaptor operands(newOperands,
+ storeOp->getAttrDictionary());
+
+ // Only apply to expanded types (tensors/etc).
+ if (!isExpandedType(storeOp.value().getType())) return failure();
+ auto &expandedGlobal =
+ expansionState->globalMap[operands.global().getValue()];
+
+ // Insert a transfer/store to the global with unknown lifetime. Lifetime
+ // refinement will make this go away if possible.
+ auto value =
+ consumeTensorOperand(storeOp.getLoc(), operands.value(), rewriter);
+ auto transferOp = rewriter.create<IREE::Stream::AsyncTransferOp>(
+ storeOp.getLoc(), expandedGlobal.resourceOp.type(), value.resource,
+ value.resourceSize, value.resourceSize, /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr);
+ rewriter.replaceOpWithNewOp<IREE::Util::GlobalStoreOp>(
+ storeOp, transferOp.result(),
+ expandedGlobal.resourceOp.getSymbolName());
+ rewriter.create<IREE::Util::GlobalStoreOp>(
+ storeOp.getLoc(), value.resourceSize,
+ expandedGlobal.resourceSizeOp.getSymbolName());
+
+ return success();
+ }
+};
+
+} // namespace
+
+void populateUtilToStreamConversionPatterns(
+ MLIRContext *context, TypeConverter &typeConverter,
+ OwningRewritePatternList &patterns) {
+ auto expansionState = std::make_shared<GlobalExpansionState>();
+ // TODO(#7432): add indirect global expansion support to streams.
+ patterns
+ .insert<GlobalOpExpansion, GlobalLoadOpExpansion, GlobalStoreOpExpansion>(
+ expansionState, typeConverter, context);
+}
+
+void populateUtilToStreamConversionPatterns(
+ MLIRContext *context, ConversionTarget &conversionTarget,
+ TypeConverter &typeConverter, OwningRewritePatternList &patterns) {
+ typeConverter.addConversion([=](IREE::Util::PtrType type,
+ SmallVectorImpl<Type> &resultTypes) {
+ // Expand pointers to tensors to [resource, sizeof resource] pointers.
+ if (!isExpandedType(type)) return failure();
+ resultTypes.push_back(
+ IREE::Util::PtrType::get(IREE::Stream::ResourceType::get(context)));
+ resultTypes.push_back(IREE::Util::PtrType::get(IndexType::get(context)));
+ return success();
+ });
+
+ typeConverter.addConversion(
+ [=](IREE::Util::PtrType type, SmallVectorImpl<Type> &resultTypes) {
+ // Expand pointers to tensors to [ptr<resource>, ptr<sizeof resource>].
+ if (!isExpandedType(type.getTargetType())) return failure();
+ resultTypes.push_back(IREE::Stream::ResourceType::get(context));
+ resultTypes.push_back(IndexType::get(context));
+ return success();
+ });
+
+ conversionTarget
+ .addLegalOp<IREE::Util::InitializerOp, IREE::Util::InitializerReturnOp>();
+ conversionTarget.addDynamicallyLegalOp<IREE::Util::GlobalOp>(
+ [&](IREE::Util::GlobalOp op) {
+ return typeConverter.isLegal(op.type()) &&
+ (!op.initial_valueAttr() ||
+ !op.initial_valueAttr().getType().isa<TensorType>());
+ });
+ conversionTarget.addDynamicallyLegalOp<IREE::Util::GlobalAddressOp>(
+ [&](IREE::Util::GlobalAddressOp op) {
+ return typeConverter.isLegal(op.result().getType());
+ });
+ conversionTarget.addDynamicallyLegalOp<IREE::Util::GlobalLoadOp>(
+ [&](IREE::Util::GlobalLoadOp op) {
+ return typeConverter.isLegal(op.result().getType());
+ });
+ conversionTarget.addDynamicallyLegalOp<IREE::Util::GlobalLoadIndirectOp>(
+ [&](IREE::Util::GlobalLoadIndirectOp op) {
+ return typeConverter.isLegal(op.result().getType());
+ });
+ conversionTarget.addDynamicallyLegalOp<IREE::Util::GlobalStoreOp>(
+ [&](IREE::Util::GlobalStoreOp op) {
+ return typeConverter.isLegal(op.value().getType());
+ });
+ conversionTarget.addDynamicallyLegalOp<IREE::Util::GlobalStoreIndirectOp>(
+ [&](IREE::Util::GlobalStoreIndirectOp op) {
+ return typeConverter.isLegal(op.value().getType());
+ });
+
+ populateUtilToStreamConversionPatterns(context, typeConverter, patterns);
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/ConvertUtilToStream.h b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/ConvertUtilToStream.h
new file mode 100644
index 0000000..c43912d
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/ConvertUtilToStream.h
@@ -0,0 +1,31 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_DIALECT_STREAM_CONVERSION_CONVERTUTILTOSTREAM_H_
+#define IREE_COMPILER_DIALECT_STREAM_CONVERSION_CONVERTUTILTOSTREAM_H_
+
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+// Populates conversion patterns that perform util->stream conversion.
+// These patterns ensure that nested types are run through the provided
+// |typeConverter|.
+void populateUtilToStreamConversionPatterns(MLIRContext *context,
+ TypeConverter &typeConverter,
+ OwningRewritePatternList &patterns);
+void populateUtilToStreamConversionPatterns(MLIRContext *context,
+ ConversionTarget &conversionTarget,
+ TypeConverter &typeConverter,
+ OwningRewritePatternList &patterns);
+
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_DIALECT_STREAM_CONVERSION_CONVERTUTILTOSTREAM_H_
diff --git a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/BUILD b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/BUILD
new file mode 100644
index 0000000..4478764
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/BUILD
@@ -0,0 +1,26 @@
+# Copyright 2021 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+iree_lit_test_suite(
+ name = "lit",
+ srcs = enforce_glob(
+ ["global_ops.mlir"],
+ include = ["*.mlir"],
+ ),
+ data = [
+ "//iree/tools:IreeFileCheck",
+ "//iree/tools:iree-opt",
+ ],
+)
diff --git a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/CMakeLists.txt b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/CMakeLists.txt
new file mode 100644
index 0000000..c588368
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/CMakeLists.txt
@@ -0,0 +1,23 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/BUILD #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_lit_test_suite(
+ NAME
+ lit
+ SRCS
+ "global_ops.mlir"
+ DATA
+ iree::tools::IreeFileCheck
+ iree::tools::iree-opt
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/global_ops.mlir b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/global_ops.mlir
new file mode 100644
index 0000000..cc8ee40
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Conversion/UtilToStream/test/global_ops.mlir
@@ -0,0 +1,86 @@
+// RUN: iree-opt -split-input-file -iree-stream-conversion %s | IreeFileCheck %s
+
+// CHECK: util.global public mutable @var_i32 : !stream.resource<variable>
+// CHECK: util.global public mutable @var_i32__size : index
+util.global public mutable @var_i32 : tensor<i32>
+// CHECK-LABEL: @mutableGlobal
+func @mutableGlobal() {
+ // CHECK-DAG: %[[VAR:.+]] = util.global.load @var_i32 : !stream.resource<variable>
+ // CHECK-DAG: %[[SIZE:.+]] = util.global.load @var_i32__size : index
+ // CHECK: %[[LOAD_T:.+]] = stream.async.transfer %[[VAR]] : !stream.resource<variable>{%[[SIZE]]} -> !stream.resource<*>{%[[SIZE]]}
+ %0 = util.global.load @var_i32 : tensor<i32>
+ // CHECK: %[[STORE_T:.+]] = stream.async.transfer %[[LOAD_T]] : !stream.resource<*>{%[[SIZE]]} -> !stream.resource<variable>{%[[SIZE]]}
+ // CHECK-DAG: util.global.store %[[STORE_T]], @var_i32 : !stream.resource<variable>
+ // CHECK-DAG: util.global.store %[[SIZE]], @var_i32__size : index
+ util.global.store %0, @var_i32 : tensor<i32>
+ return
+}
+
+// -----
+
+// TODO(#7432): add indirect global expansion support to streams.
+// util.global public mutable @var_indirect : tensor<i32>
+// func @mutableGlobalIndirect() {
+// %0 = util.global.address @var_indirect : !util.ptr<tensor<i32>>
+// %1 = util.global.load.indirect %0 : !util.ptr<tensor<i32>> -> tensor<i32>
+// util.global.store.indirect %1, %0 : tensor<i32> -> !util.ptr<tensor<i32>>
+// return
+// }
+
+// -----
+
+// CHECK-DAG: util.global public mutable @var_with_tensor_initializer : !stream.resource<variable>
+// CHECK-DAG: util.global public mutable @var_with_tensor_initializer__size : index
+// CHECK-NEXT: util.initializer {
+// CHECK-NEXT: %[[CST:.+]] = stream.tensor.constant : tensor<f32> in !stream.resource<variable> = dense<0.000000e+00> : tensor<f32>
+// CHECK-NEXT: %[[SIZE:.+]] = stream.resource.size %[[CST]] : !stream.resource<variable>
+// CHECK-DAG: util.global.store %[[CST]], @var_with_tensor_initializer : !stream.resource<variable>
+// CHECK-DAG: util.global.store %[[SIZE]], @var_with_tensor_initializer__size : index
+util.global public mutable @var_with_tensor_initializer = dense<0.000000e+00> : tensor<f32>
+// CHECK-LABEL: @initializedGlobal
+func @initializedGlobal() {
+ // CHECK-DAG: = util.global.load @var_with_tensor_initializer : !stream.resource<variable>
+ // CHECK-DAG: = util.global.load @var_with_tensor_initializer__size : index
+ %0 = util.global.load @var_with_tensor_initializer : tensor<f32>
+ // CHECK-DAG: util.global.store %{{.+}}, @var_with_tensor_initializer : !stream.resource<variable>
+ // CHECK-DAG: util.global.store %{{.+}}, @var_with_tensor_initializer__size : index
+ util.global.store %0, @var_with_tensor_initializer : tensor<f32>
+ return
+}
+
+// -----
+
+// Checks that the implicit cast allowing a buffer_view to store into a variable
+// that maps to a buffer is permitted.
+
+// CHECK-DAG: util.global public mutable @var_with_buffer_view_store : !stream.resource<variable>
+// CHECK-DAG: util.global public mutable @var_with_buffer_view_store__size : index
+util.global public mutable @var_with_buffer_view_store : tensor<?x4xf32>
+// CHECK-LABEL: @globalStoreFromExternal
+func @globalStoreFromExternal(%arg0: !hal.buffer_view) {
+ // CHECK: %[[DIM0:.+]] = hal.buffer_view.dim
+ %dim0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
+ // CHECK: %[[SIZE:.+]] = stream.tensor.sizeof tensor<?x4xf32>{%[[DIM0]]} : index
+ // CHECK: %[[IMPORT:.+]] = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x4xf32>{%[[DIM0]]} in !stream.resource<external>{%[[SIZE]]}
+ // CHECK: %[[T:.+]] = stream.async.transfer %[[IMPORT]] : !stream.resource<external>{%[[SIZE]]} -> !stream.resource<*>{%[[SIZE]]}
+ %0 = hal.tensor.cast %arg0 : !hal.buffer_view -> tensor<?x4xf32>{%dim0}
+ // CHECK: %[[VAR:.+]] = stream.async.transfer %[[T]] : !stream.resource<*>{%[[SIZE]]} -> !stream.resource<variable>{%[[SIZE]]}
+ // CHECK: util.global.store %[[VAR]], @var_with_buffer_view_store : !stream.resource<variable>
+ // CHECK: util.global.store %[[SIZE]], @var_with_buffer_view_store__size : index
+ util.global.store %0, @var_with_buffer_view_store : tensor<?x4xf32>
+ return
+}
+
+// -----
+
+// Checks that the implicit cast allowing a buffer_view to indirect store into
+// a variable that maps to a buffer is permitted.
+
+// TODO(#7432): add indirect global expansion support to streams.
+// util.global public mutable @var_indirect_with_buffer_view_store : tensor<i32>
+// func @globalStoreFromExternalIndirect(%arg0: !hal.buffer_view) {
+// %0 = util.global.address @var_indirect_with_buffer_view_store : !util.ptr<tensor<i32>>
+// %1 = hal.tensor.cast %arg0 : !hal.buffer_view -> tensor<i32>
+// util.global.store.indirect %1, %0 : tensor<i32> -> !util.ptr<tensor<i32>>
+// return
+// }
diff --git a/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp b/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
index 84a1931..5a3906f 100644
--- a/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
+++ b/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
@@ -311,6 +311,42 @@
Block::iterator(op));
}
+namespace {
+
+// Propagates resource sizes through select ops by selecting on the sizes of the
+// select operands.
+//
+// Example:
+// %a = stream... : !stream.resource<*>{%a_sz}
+// %b = stream... : !stream.resource<*>{%b_sz}
+// %c = select %cond, %a, %b : !stream.resource<*>
+// %c_sz = stream.resource.size %c : !stream.resource<*>
+// ->
+// %c = select %cond, %a, %b : !stream.resource<*>
+// %c_sz = select %cond, %a_sz, %b_sz : index
+struct SelectResourceSizeOp : public OpRewritePattern<ResourceSizeOp> {
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(ResourceSizeOp op,
+ PatternRewriter &rewriter) const override {
+ auto selectOp = op.operand().getDefiningOp<mlir::SelectOp>();
+ if (!selectOp) return failure();
+ auto trueSize = rewriter.createOrFold<IREE::Stream::ResourceSizeOp>(
+ op.getLoc(), selectOp.true_value(), op.affinityAttr());
+ auto falseSize = rewriter.createOrFold<IREE::Stream::ResourceSizeOp>(
+ op.getLoc(), selectOp.false_value(), op.affinityAttr());
+ rewriter.replaceOpWithNewOp<mlir::SelectOp>(op, selectOp.condition(),
+ trueSize, falseSize);
+ return success();
+ }
+};
+
+} // namespace
+
+void ResourceSizeOp::getCanonicalizationPatterns(
+ OwningRewritePatternList &results, MLIRContext *context) {
+ results.insert<SelectResourceSizeOp>(context);
+}
+
//===----------------------------------------------------------------------===//
// stream.resource.map
//===----------------------------------------------------------------------===//
@@ -468,8 +504,7 @@
// Zero offsets don't do anything and can just be removed so we can avoid
// inserting a bunch of additional IR.
- if (auto constantOp = dyn_cast_or_null<arith::ConstantIndexOp>(
- baseOffset.getDefiningOp())) {
+ if (auto constantOp = baseOffset.getDefiningOp<arith::ConstantIndexOp>()) {
if (constantOp.value() == 0) {
return success();
}
@@ -916,8 +951,8 @@
LogicalResult matchAndRewrite(AsyncCloneOp cloneOp,
PatternRewriter &rewriter) const override {
if (cloneOp.use_empty()) return failure();
- auto sourceOp = dyn_cast_or_null<IREE::Stream::StreamableOpInterface>(
- cloneOp.source().getDefiningOp());
+ auto sourceOp =
+ cloneOp.source().getDefiningOp<IREE::Stream::StreamableOpInterface>();
if (!sourceOp || !sourceOp.preferCloneToConsumers()) return failure();
for (auto &use : llvm::make_early_inc_range(cloneOp.result().getUses())) {
rewriter.setInsertionPoint(use.getOwner());
@@ -968,8 +1003,7 @@
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(AsyncSliceOp sliceOp,
PatternRewriter &rewriter) const override {
- auto splatOp = dyn_cast_or_null<IREE::Stream::AsyncSplatOp>(
- sliceOp.source().getDefiningOp());
+ auto splatOp = sliceOp.source().getDefiningOp<IREE::Stream::AsyncSplatOp>();
if (!splatOp) return failure();
rewriter.replaceOpWithNewOp<IREE::Stream::AsyncSplatOp>(
sliceOp, sliceOp.result().getType(), splatOp.value(),
@@ -1053,14 +1087,13 @@
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(AsyncUpdateOp updateOp,
PatternRewriter &rewriter) const override {
- auto splatOp = dyn_cast_or_null<IREE::Stream::AsyncSplatOp>(
- updateOp.update().getDefiningOp());
+ auto splatOp =
+ updateOp.update().getDefiningOp<IREE::Stream::AsyncSplatOp>();
if (!splatOp) return failure();
rewriter.replaceOpWithNewOp<IREE::Stream::AsyncFillOp>(
updateOp, updateOp.result().getType(), updateOp.target(),
updateOp.target_size(), updateOp.target_offset(), updateOp.target_end(),
- updateOp.update_size(), splatOp.value(), updateOp.tied_operandsAttr(),
- updateOp.affinityAttr());
+ updateOp.update_size(), splatOp.value(), updateOp.affinityAttr());
return success();
}
};
@@ -1086,8 +1119,8 @@
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(AsyncUpdateOp updateOp,
PatternRewriter &rewriter) const override {
- auto sliceOp = dyn_cast_or_null<IREE::Stream::AsyncSliceOp>(
- updateOp.update().getDefiningOp());
+ auto sliceOp =
+ updateOp.update().getDefiningOp<IREE::Stream::AsyncSliceOp>();
if (!sliceOp || sliceOp->getBlock() != updateOp->getBlock()) {
// Source is not a slice or a slice from out-of-block. We don't want to
// grow memory usage by sinking the slice here (we may slice into the
@@ -1098,8 +1131,7 @@
updateOp, updateOp.result().getType(), updateOp.target(),
updateOp.target_size(), updateOp.target_offset(), updateOp.target_end(),
sliceOp.source(), sliceOp.source_size(), sliceOp.source_offset(),
- sliceOp.source_end(), sliceOp.result_size(),
- updateOp.tied_operandsAttr(), updateOp.affinityAttr());
+ sliceOp.source_end(), sliceOp.result_size(), updateOp.affinityAttr());
return success();
}
};
@@ -1140,8 +1172,7 @@
rewriter.replaceOpWithNewOp<IREE::Stream::AsyncUpdateOp>(
copyOp, copyOp.result().getType(), copyOp.target(),
copyOp.target_size(), copyOp.target_offset(), copyOp.target_end(),
- copyOp.source(), copyOp.source_size(), copyOp.tied_operandsAttr(),
- copyOp.affinityAttr());
+ copyOp.source(), copyOp.source_size(), copyOp.affinityAttr());
return success();
}
return failure();
@@ -1162,8 +1193,7 @@
//===----------------------------------------------------------------------===//
OpFoldResult AsyncTransferOp::fold(ArrayRef<Attribute> operands) {
- if (auto sourceTransferOp =
- dyn_cast_or_null<AsyncTransferOp>(source().getDefiningOp())) {
+ if (auto sourceTransferOp = source().getDefiningOp<AsyncTransferOp>()) {
if (sourceTransferOp.source().getType() == result().getType() &&
sourceTransferOp.source_affinity() == result_affinity()) {
return sourceTransferOp.source();
@@ -1201,6 +1231,29 @@
}
//===----------------------------------------------------------------------===//
+// stream.async.load
+//===----------------------------------------------------------------------===//
+
+void AsyncLoadOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+ MLIRContext *context) {
+ // TODO(benvanik): splat + load -> splat value.
+ // TODO(benvanik): clone + ex load -> slice (ranged) + load.
+ // TODO(benvanik): slice + ex load -> slice (ranged) + load.
+ // TODO(benvanik): value->transfer->load -> value->slice->transfer->load?
+ // TODO(benvanik): combine multiple loads from the same target if contiguous.
+}
+
+//===----------------------------------------------------------------------===//
+// stream.async.store
+//===----------------------------------------------------------------------===//
+
+void AsyncStoreOp::getCanonicalizationPatterns(
+ OwningRewritePatternList &results, MLIRContext *context) {
+ // TODO(benvanik): if value is a constant splat then turn into fill.
+ // TODO(benvanik): combine multiple stores to the same target if contiguous.
+}
+
+//===----------------------------------------------------------------------===//
// stream.async.dispatch
//===----------------------------------------------------------------------===//
@@ -1282,8 +1335,7 @@
SmallVector<Value> newTimepoints;
SmallVector<std::pair<unsigned, Value>> replacements;
for (auto operand : llvm::enumerate(op.operands())) {
- if (auto awaitOp = dyn_cast_or_null<TimepointAwaitOp>(
- operand.value().getDefiningOp())) {
+ if (auto awaitOp = operand.value().getDefiningOp<TimepointAwaitOp>()) {
newTimepoints.push_back(awaitOp.timepoint());
replacements.push_back(std::make_pair(
operand.index(), awaitOp.getTiedResultOperand(operand.value())));
@@ -1738,8 +1790,7 @@
SmallVector<Value> newTimepoints;
SmallVector<std::pair<unsigned, Value>> replacements;
for (auto operand : llvm::enumerate(op.operands())) {
- if (auto awaitOp = dyn_cast_or_null<TimepointAwaitOp>(
- operand.value().getDefiningOp())) {
+ if (auto awaitOp = operand.value().getDefiningOp<TimepointAwaitOp>()) {
newTimepoints.push_back(awaitOp.timepoint());
replacements.push_back(std::make_pair(
operand.index(), awaitOp.getTiedResultOperand(operand.value())));
@@ -2051,8 +2102,8 @@
rewriter.startRootUpdate(op);
bool didChange = false;
for (auto operand : llvm::enumerate(op.operands())) {
- auto subviewOp = dyn_cast_or_null<IREE::Stream::ResourceSubviewOp>(
- operand.value().getDefiningOp());
+ auto subviewOp =
+ operand.value().getDefiningOp<IREE::Stream::ResourceSubviewOp>();
if (!subviewOp) continue;
didChange = true;
unsigned operandIdx = static_cast<unsigned>(operand.index());
diff --git a/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
index cd39e77..50389ee 100644
--- a/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
+++ b/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
@@ -1091,6 +1091,41 @@
}
//===----------------------------------------------------------------------===//
+// stream.async.load
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verifyOp(AsyncLoadOp op) {
+ if (failed(verifyOpValueSizes(op, op.source(), op.source_size()))) {
+ return failure();
+ }
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// stream.async.store
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verifyOp(AsyncStoreOp op) {
+ if (failed(verifyOpValueSizes(op, op.target(), op.target_size()))) {
+ return failure();
+ }
+ return success();
+}
+
+Value AsyncStoreOp::getTiedResult(unsigned resultIndex) {
+ return IREE::Util::TiedOpInterface::findTiedBaseValue(target());
+}
+
+::llvm::Optional<unsigned> AsyncStoreOp::getTiedResultOperandIndex(
+ unsigned resultIndex) {
+ return {0}; // target
+}
+
+SmallVector<int64_t, 4> AsyncStoreOp::getTiedResultOperandIndices() {
+ return {0}; // target
+}
+
+//===----------------------------------------------------------------------===//
// stream.async.dispatch
//===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/Stream/IR/StreamOps.td b/iree/compiler/Dialect/Stream/IR/StreamOps.td
index 9d77a6e..74554ce 100644
--- a/iree/compiler/Dialect/Stream/IR/StreamOps.td
+++ b/iree/compiler/Dialect/Stream/IR/StreamOps.td
@@ -200,6 +200,7 @@
Value getResultSize(unsigned idx) { return {}; }
}];
+ let hasCanonicalizer = 1;
let hasFolder = 1;
}
@@ -949,7 +950,6 @@
Variadic<Stream_Dim>:$start_indices,
Variadic<Stream_Dim>:$lengths,
Stream_PrimitiveType:$value,
- OptionalAttr<Util_TiedOpStorageAttr>:$tied_operands,
OptionalAttr<Stream_AffinityAttr>:$affinity
);
let results = (outs
@@ -963,7 +963,7 @@
`->`
$target_encoding (`` `{` $target_encoding_dims^ `}`)?
`in`
- custom<ShapedTiedResult>(type($target), $target_size, $tied_operands)
+ custom<ShapedTiedResult>(type($target), $target_size)
attr-dict-with-keyword
}];
@@ -1009,7 +1009,6 @@
TypeAttr:$update_encoding,
Stream_ShapeDynamicDims:$update_encoding_dims,
Stream_Size:$update_size,
- OptionalAttr<Util_TiedOpStorageAttr>:$tied_operands,
OptionalAttr<Stream_AffinityAttr>:$affinity
);
let results = (outs
@@ -1025,7 +1024,7 @@
`->`
$target_encoding (`` `{` $target_encoding_dims^ `}`)?
`in`
- custom<ShapedTiedResult>(type($target), $target_size, $tied_operands)
+ custom<ShapedTiedResult>(type($target), $target_size)
attr-dict-with-keyword
}];
@@ -1090,7 +1089,7 @@
let hasCanonicalizer = 1;
}
-def Stream_TensorStoreOp : Stream_Op<"tensor.store", [
+def Stream_TensorStoreOp : Stream_PureOp<"tensor.store", [
AttrSizedOperandSegments,
AllTypesMatch<["target", "result"]>,
Stream_TensorPhaseOp,
@@ -1113,8 +1112,7 @@
Stream_ShapeDynamicDims:$target_encoding_dims,
Stream_Size:$target_size,
Variadic<Stream_Dim>:$indices,
- AnyTypeOf<[Stream_PrimitiveType, AnyVector]>:$value,
- OptionalAttr<Util_TiedOpStorageAttr>:$tied_operands
+ AnyTypeOf<[Stream_PrimitiveType, AnyVector]>:$value
);
let results = (outs
Stream_StagingResource:$result
@@ -1127,7 +1125,7 @@
`->`
$target_encoding (`` `{` $target_encoding_dims^ `}`)?
`in`
- custom<ShapedTiedResult>(type($target), $target_size, $tied_operands)
+ custom<ShapedTiedResult>(type($target), $target_size)
attr-dict-with-keyword
}];
@@ -1385,7 +1383,6 @@
Stream_Offset:$target_end,
Stream_Size:$target_length,
Stream_PrimitiveType:$value,
- OptionalAttr<Util_TiedOpStorageAttr>:$tied_operands,
OptionalAttr<Stream_AffinityAttr>:$affinity
);
let results = (outs
@@ -1397,7 +1394,7 @@
$value `,`
$target `[` $target_offset `to` $target_end `for` $target_length `]` `:`
type($value) `->`
- custom<ShapedTiedResult>(type($target), $target_size, $tied_operands)
+ custom<ShapedTiedResult>(type($target), $target_size)
attr-dict-with-keyword
}];
@@ -1439,7 +1436,6 @@
Stream_Offset:$target_end,
Stream_AnyStreamResource:$update,
Stream_Size:$update_size,
- OptionalAttr<Util_TiedOpStorageAttr>:$tied_operands,
OptionalAttr<Stream_AffinityAttr>:$affinity
);
let results = (outs
@@ -1451,7 +1447,7 @@
$update `,`
$target `[` $target_offset `to` $target_end `]` `:`
type($update) `` `{` $update_size `}` `->`
- custom<ShapedTiedResult>(type($target), $target_size, $tied_operands)
+ custom<ShapedTiedResult>(type($target), $target_size)
attr-dict-with-keyword
}];
@@ -1500,7 +1496,6 @@
Stream_Offset:$source_offset,
Stream_Offset:$source_end,
Stream_Size:$length,
- OptionalAttr<Util_TiedOpStorageAttr>:$tied_operands,
OptionalAttr<Stream_AffinityAttr>:$affinity
);
let results = (outs
@@ -1513,7 +1508,7 @@
$target `[` $target_offset `to` $target_end `]` `,`
$length `:`
type($source) `` `{` $source_size `}` `->`
- custom<ShapedTiedResult>(type($target), $target_size, $tied_operands)
+ custom<ShapedTiedResult>(type($target), $target_size)
attr-dict-with-keyword
}];
@@ -1580,6 +1575,87 @@
let hasFolder = 1;
}
+def Stream_AsyncLoadOp : Stream_PureOp<"async.load", [
+ Stream_AsyncPhaseOp,
+ Util_SizeAwareOp,
+]> {
+ let summary = [{loads a value from a resource}];
+ let description = [{
+ Returns the element at the given location from within the resource.
+ }];
+
+ let arguments = (ins
+ Stream_StagingResource:$source,
+ Stream_Size:$source_size,
+ Stream_Offset:$source_offset
+ );
+ let results = (outs
+ AnyTypeOf<[Stream_PrimitiveType, AnyVector]>:$result
+ );
+
+ let assemblyFormat = [{
+ $source `[` $source_offset `]` `:`
+ type($source) `` `{` $source_size `}`
+ `->`
+ type($result)
+ attr-dict-with-keyword
+ }];
+
+ let extraClassDeclaration = [{
+ Value getOperandSize(unsigned idx) { return source_size(); }
+ Value getResultSize(unsigned idx) { return {}; }
+ }];
+
+ let verifier = [{ return verifyOp(*this); }];
+
+ let hasCanonicalizer = 1;
+}
+
+def Stream_AsyncStoreOp : Stream_PureOp<"async.store", [
+ AllTypesMatch<["target", "result"]>,
+ Stream_AsyncPhaseOp,
+ Util_SizeAwareOp,
+ DeclareOpInterfaceMethods<Util_TiedOpInterface, [
+ "getTiedResult",
+ "getTiedResultOperandIndex",
+ "getTiedResultOperandIndices",
+ ]>,
+]> {
+ let summary = [{stores a value into a resource}];
+ let description = [{
+ Returns a resource with the element at the given offset set to the given
+ value.
+ }];
+
+ let arguments = (ins
+ Stream_StagingResource:$target,
+ Stream_Size:$target_size,
+ Stream_Offset:$target_offset,
+ AnyTypeOf<[Stream_PrimitiveType, AnyVector]>:$value
+ );
+ let results = (outs
+ Stream_StagingResource:$result
+ );
+
+ let assemblyFormat = [{
+ $value `,`
+ $target `[` $target_offset `]` `:`
+ type($value)
+ `->`
+ custom<ShapedTiedResult>(type($target), $target_size)
+ attr-dict-with-keyword
+ }];
+
+ let extraClassDeclaration = [{
+ Value getOperandSize(unsigned idx) { return target_size(); }
+ Value getResultSize(unsigned idx) { return target_size(); }
+ }];
+
+ let verifier = [{ return verifyOp(*this); }];
+
+ let hasCanonicalizer = 1;
+}
+
def Stream_AsyncDispatchOp : Stream_Op<"async.dispatch", [
AttrSizedOperandSegments,
Stream_AffinityOp,
diff --git a/iree/compiler/Dialect/Stream/IR/test/async_ops.mlir b/iree/compiler/Dialect/Stream/IR/test/async_ops.mlir
index 352a521..b02fe78 100644
--- a/iree/compiler/Dialect/Stream/IR/test/async_ops.mlir
+++ b/iree/compiler/Dialect/Stream/IR/test/async_ops.mlir
@@ -89,6 +89,26 @@
// -----
+// CHECK-LABEL: @asyncLoad
+func @asyncLoad(%arg0: !stream.resource<staging>, %arg1: index) -> f32 {
+ %c0 = arith.constant 0 : index
+ // CHECK: = stream.async.load %arg0[%c0] : !stream.resource<staging>{%arg1} -> f32
+ %0 = stream.async.load %arg0[%c0] : !stream.resource<staging>{%arg1} -> f32
+ return %0 : f32
+}
+
+// -----
+
+// CHECK-LABEL: @asyncStore
+func @asyncStore(%arg0: !stream.resource<staging>, %arg1: index, %arg2: f32) -> !stream.resource<staging> {
+ %c0 = arith.constant 0 : index
+ // CHECK: = stream.async.store %arg2, %arg0[%c0] : f32 -> %arg0 as !stream.resource<staging>{%arg1}
+ %0 = stream.async.store %arg2, %arg0[%c0] : f32 -> %arg0 as !stream.resource<staging>{%arg1}
+ return %0 : !stream.resource<staging>
+}
+
+// -----
+
// CHECK-LABEL: @asyncDispatch
func @asyncDispatch(%arg0: !stream.resource<*>, %arg1: index) -> !stream.resource<*> {
%c1 = arith.constant 1 : index
diff --git a/iree/compiler/Dialect/Stream/IR/test/resource_folding.mlir b/iree/compiler/Dialect/Stream/IR/test/resource_folding.mlir
index 0c7b333..5bd3609 100644
--- a/iree/compiler/Dialect/Stream/IR/test/resource_folding.mlir
+++ b/iree/compiler/Dialect/Stream/IR/test/resource_folding.mlir
@@ -13,6 +13,23 @@
// -----
+// CHECK-LABEL: @SelectResourceSizeOp
+func @SelectResourceSizeOp(%arg0: !stream.resource<staging>, %arg1: index, %arg2: !stream.resource<staging>, %arg3: index, %arg4: i1) -> (!stream.resource<staging>, index) {
+ // CHECK: %[[ARG0_T:.+]] = stream.async.transfer %arg0 {{.+}} -> !stream.resource<*>{%[[ARG0_SZ:.+]]}
+ %0 = stream.async.transfer %arg0 : !stream.resource<staging>{%arg1} -> !stream.resource<*>{%arg1}
+ // CHECK: %[[ARG2_T:.+]] = stream.async.transfer %arg2 {{.+}} -> !stream.resource<*>{%[[ARG2_SZ:.+]]}
+ %1 = stream.async.transfer %arg2 : !stream.resource<staging>{%arg3} -> !stream.resource<*>{%arg3}
+ // CHECK: %[[RET_T:.+]] = select %arg4, %[[ARG0_T]], %[[ARG2_T]] : !stream.resource<*>
+ %2 = select %arg4, %0, %1 : !stream.resource<*>
+ // CHECK: %[[RET_SIZE:.+]] = select %arg4, %[[ARG0_SZ]], %[[ARG2_SZ]] : index
+ %3 = stream.resource.size %2 : !stream.resource<*>
+ // CHECK: = stream.async.transfer %[[RET_T]] : !stream.resource<*>{%[[RET_SIZE]]}
+ %4 = stream.async.transfer %2 : !stream.resource<*>{%3} -> !stream.resource<staging>{%3}
+ return %4, %3 : !stream.resource<staging>, index
+}
+
+// -----
+
// CHECK-LABEL: @FoldSubviewIntoLoadOp
func @FoldSubviewIntoLoadOp(%arg0: !stream.resource<staging>, %arg1: index) -> i32 {
%c64 = arith.constant 64 : index
diff --git a/iree/compiler/Dialect/Stream/Transforms/BUILD b/iree/compiler/Dialect/Stream/Transforms/BUILD
new file mode 100644
index 0000000..1c4848c
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/BUILD
@@ -0,0 +1,71 @@
+# Copyright 2021 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+cc_library(
+ name = "Transforms",
+ srcs = [
+ "ConvertToStream.cpp",
+ "OutlineConstants.cpp",
+ "PassDetail.h",
+ "Passes.cpp",
+ "VerifyLowerings.cpp",
+ ],
+ hdrs = [
+ "Passes.h",
+ "Passes.h.inc",
+ ],
+ deps = [
+ ":PassesIncGen",
+ "//iree/compiler/Dialect/Flow/IR",
+ "//iree/compiler/Dialect/Shape/IR",
+ "//iree/compiler/Dialect/Shape/Transforms",
+ "//iree/compiler/Dialect/Shape/Utils:TypeConversion",
+ "//iree/compiler/Dialect/Stream/Conversion/FlowToStream",
+ "//iree/compiler/Dialect/Stream/Conversion/StandardToStream",
+ "//iree/compiler/Dialect/Stream/Conversion/UtilToStream",
+ "//iree/compiler/Dialect/Stream/IR",
+ "//iree/compiler/Dialect/Util/Conversion",
+ "//iree/compiler/Dialect/Util/IR",
+ "//iree/compiler/Dialect/Util/Transforms",
+ "//iree/compiler/Utils",
+ "@llvm-project//llvm:Support",
+ "@llvm-project//mlir:Affine",
+ "@llvm-project//mlir:Analysis",
+ "@llvm-project//mlir:ArithmeticDialect",
+ "@llvm-project//mlir:DialectUtils",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:Pass",
+ "@llvm-project//mlir:SCFDialect",
+ "@llvm-project//mlir:StandardOps",
+ "@llvm-project//mlir:Support",
+ "@llvm-project//mlir:TensorDialect",
+ "@llvm-project//mlir:TransformUtils",
+ "@llvm-project//mlir:Transforms",
+ ],
+)
+
+gentbl_cc_library(
+ name = "PassesIncGen",
+ tbl_outs = [
+ (
+ ["-gen-pass-decls"],
+ "Passes.h.inc",
+ ),
+ ],
+ tblgen = "@llvm-project//mlir:mlir-tblgen",
+ td_file = "Passes.td",
+ deps = [
+ "@llvm-project//mlir:PassBaseTdFiles",
+ ],
+)
diff --git a/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
new file mode 100644
index 0000000..267b6f6
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
@@ -0,0 +1,63 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# iree/compiler/Dialect/Stream/Transforms/BUILD #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_cc_library(
+ NAME
+ Transforms
+ HDRS
+ "Passes.h"
+ "Passes.h.inc"
+ SRCS
+ "ConvertToStream.cpp"
+ "OutlineConstants.cpp"
+ "PassDetail.h"
+ "Passes.cpp"
+ "VerifyLowerings.cpp"
+ DEPS
+ ::PassesIncGen
+ LLVMSupport
+ MLIRAffine
+ MLIRAnalysis
+ MLIRArithmetic
+ MLIRIR
+ MLIRPass
+ MLIRSCF
+ MLIRStandard
+ MLIRSupport
+ MLIRTensor
+ MLIRTransformUtils
+ MLIRTransforms
+ iree::compiler::Dialect::Flow::IR
+ iree::compiler::Dialect::Shape::IR
+ iree::compiler::Dialect::Shape::Transforms
+ iree::compiler::Dialect::Shape::Utils::TypeConversion
+ iree::compiler::Dialect::Stream::Conversion::FlowToStream
+ iree::compiler::Dialect::Stream::Conversion::StandardToStream
+ iree::compiler::Dialect::Stream::Conversion::UtilToStream
+ iree::compiler::Dialect::Stream::IR
+ iree::compiler::Dialect::Util::Conversion
+ iree::compiler::Dialect::Util::IR
+ iree::compiler::Dialect::Util::Transforms
+ iree::compiler::Utils
+ PUBLIC
+)
+
+iree_tablegen_library(
+ NAME
+ PassesIncGen
+ TD_FILE
+ "Passes.td"
+ OUTS
+ -gen-pass-decls Passes.h.inc
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp b/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
new file mode 100644
index 0000000..01ec65f
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
@@ -0,0 +1,257 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
+#include "iree/compiler/Dialect/Shape/IR/Builders.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeDialect.h"
+#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
+#include "iree/compiler/Dialect/Stream/Conversion/FlowToStream/ConvertFlowToStream.h"
+#include "iree/compiler/Dialect/Stream/Conversion/StandardToStream/ConvertStandardToStream.h"
+#include "iree/compiler/Dialect/Stream/Conversion/UtilToStream/ConvertUtilToStream.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
+#include "iree/compiler/Dialect/Stream/Transforms/PassDetail.h"
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
+#include "iree/compiler/Dialect/Util/Conversion/ConversionPatterns.h"
+#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
+#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
+#include "iree/compiler/Dialect/Util/Transforms/Patterns.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace Stream {
+namespace {
+
+// Builds a stream.tensor.import op that imports an external tensor value into
+// a stream resource.
+static Value buildTensorImportOp(Location loc, Value sourceTensor,
+ Type targetType, OpBuilder &builder) {
+ // Gather dynamic dimensions from the input value.
+ auto dynamicDims =
+ Shape::buildOrFindDynamicDimsForValue(loc, sourceTensor, builder);
+
+ // Compute the size of the tensor once in the stream resource.
+ // This may differ from the external encoding of the tensor as imports are
+ // a transfer operation that may need to reformat the tensor.
+ auto encodingAttr = TypeAttr::get(sourceTensor.getType());
+ auto resultSize = builder.createOrFold<IREE::Stream::TensorSizeOfOp>(
+ loc, builder.getIndexType(), encodingAttr, dynamicDims,
+ /*affinity=*/nullptr);
+
+ // Associate the external SSA value, encoding, and shape information with the
+ // stream resource. When lowering we'll then have all the metadata required
+ // even after erasing it all on the resource.
+ auto externalType = IREE::Stream::ResourceType::get(
+ builder.getContext(), IREE::Stream::Lifetime::External);
+ auto importOp = builder.create<IREE::Stream::TensorImportOp>(
+ loc, externalType, sourceTensor, encodingAttr, dynamicDims, resultSize,
+ /*affinity=*/nullptr);
+
+ // If needed insert a transfer to the target lifetime.
+ Value result = importOp.result();
+ if (targetType != externalType) {
+ result = builder
+ .create<IREE::Stream::AsyncTransferOp>(
+ loc, externalType, result, resultSize, resultSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr)
+ .result();
+ }
+
+ return result;
+}
+
+// Builds a stream.tensor.export op that exports a stream resource into an
+// external tensor value.
+static Value buildTensorExportOp(Location loc, Value sourceResource,
+ TensorType targetType, ValueRange dynamicDims,
+ OpBuilder &builder) {
+ // Query the size of the resource - which may differ from the target external
+ // value if we changed the encoding.
+ auto sourceSize = builder.createOrFold<IREE::Stream::ResourceSizeOp>(
+ loc, builder.getIndexType(), sourceResource);
+
+ // If needed insert a transfer to external resource lifetime.
+ auto externalType = IREE::Stream::ResourceType::get(
+ builder.getContext(), IREE::Stream::Lifetime::External);
+ if (sourceResource.getType() != externalType) {
+ sourceResource = builder.create<IREE::Stream::AsyncTransferOp>(
+ loc, externalType, sourceResource, sourceSize, sourceSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr);
+ }
+
+ // Associate the stream resource and external encoding and shape information.
+ auto newOp = builder.create<IREE::Stream::TensorExportOp>(
+ loc, targetType, sourceResource, TypeAttr::get(targetType), dynamicDims,
+ sourceSize,
+ /*affinity=*/nullptr);
+ return newOp.result();
+}
+
+// Returns true if |op| has tensor I/O that is not yet imported/exported using
+// the stream ops that capture encodings and shapes.
+static bool doesOperationNeedWrapping(Operation *op) {
+ return llvm::any_of(
+ op->getOperands(),
+ [&](Value operand) {
+ if (!operand.getType().isa<TensorType>()) return false;
+ return !isa_and_nonnull<TensorExportOp>(operand.getDefiningOp());
+ }) ||
+ llvm::any_of(op->getResults(), [&](Value result) {
+ if (!result.getType().isa<TensorType>()) return false;
+ return !llvm::all_of(result.getUsers(), [&](Operation *user) {
+ return isa<TensorImportOp>(user);
+ });
+ });
+}
+
+// Fallback handler for unknown ops taking/returning tensors that need to be
+// marshaled into/outof stream resource types.
+struct GenericResourcePattern : public ConversionPattern {
+ GenericResourcePattern(MLIRContext *context, TypeConverter &converter)
+ : ConversionPattern(converter, MatchAnyOpTypeTag(), 0, context) {}
+ LogicalResult matchAndRewrite(
+ Operation *op, ArrayRef<Value> operands,
+ ConversionPatternRewriter &rewriter) const override {
+ if (!doesOperationNeedWrapping(op)) return failure();
+
+ // Export resources into tensor operands for the op to consume.
+ SmallVector<Value> newOperands;
+ newOperands.reserve(op->getNumOperands());
+ rewriter.setInsertionPoint(op);
+ for (auto it : llvm::zip(op->getOperands(), operands)) {
+ auto oldOperand = std::get<0>(it);
+ auto newOperand = std::get<1>(it);
+ if (!newOperand.getType().isa<IREE::Stream::ResourceType>()) {
+ newOperands.push_back(newOperand);
+ continue;
+ }
+ auto tensorType = oldOperand.getType().dyn_cast<TensorType>();
+ assert(tensorType && "must have a tensor type to map to a resource");
+
+ auto dynamicDims = Shape::buildOrFindDynamicDimsForValue(
+ op->getLoc(), oldOperand, rewriter);
+ newOperands.push_back(buildTensorExportOp(
+ op->getLoc(), newOperand, tensorType, dynamicDims, rewriter));
+ }
+ rewriter.updateRootInPlace(op, [&]() { op->setOperands(newOperands); });
+
+ // Import into resources from tensor results produced by the op.
+ rewriter.setInsertionPointAfter(op);
+ for (auto result : op->getResults()) {
+ auto tensorType = result.getType().dyn_cast<TensorType>();
+ if (!tensorType) continue;
+
+ auto dynamicDims =
+ Shape::buildOrFindDynamicDimsForValue(op->getLoc(), result, rewriter);
+ auto importedValue = buildTensorImportOp(
+ op->getLoc(), result, IREE::Stream::ResourceType::get(getContext()),
+ rewriter);
+ result.replaceAllUsesExcept(importedValue, importedValue.getDefiningOp());
+ }
+
+ return success();
+ }
+};
+
+class ConvertToStreamPass : public ConvertToStreamBase<ConvertToStreamPass> {
+ public:
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<ShapeDialect>();
+ registry.insert<mlir::StandardOpsDialect>();
+ registry.insert<mlir::arith::ArithmeticDialect>();
+ registry.insert<IREE::Stream::StreamDialect>();
+ registry.insert<IREE::Util::UtilDialect>();
+ }
+
+ void runOnOperation() override {
+ auto *context = &getContext();
+
+ TypeConverter typeConverter;
+ ConversionTarget conversionTarget(getContext());
+ OwningRewritePatternList patterns(&getContext());
+
+ // Always allow lowerering target dialects and reasonable types.
+ conversionTarget.addLegalDialect<IREE::Stream::StreamDialect>();
+ typeConverter.addConversion(
+ [](IREE::Stream::ResourceType type) { return type; });
+
+ // Disallow tensor dialects; the goal here is to remove all tensors and
+ // turn them into stream resource ops.
+ auto indexType = IndexType::get(context);
+ conversionTarget.addIllegalDialect<tensor::TensorDialect>();
+ typeConverter.addConversion(
+ [=](TensorType type, SmallVectorImpl<Type> &resultTypes) {
+ // Expand tensors to [resource, sizeof resource].
+ resultTypes.push_back(IREE::Stream::ResourceType::get(context));
+ resultTypes.push_back(indexType);
+ return success();
+ });
+ typeConverter.addArgumentMaterialization(
+ [](OpBuilder &builder, TensorType resultType, ValueRange inputs,
+ Location loc) -> Optional<Value> {
+ assert(inputs.size() >= 2);
+ auto resourceValue = inputs[0];
+ auto resourceSize = inputs[1];
+ assert(inputs.size() == 2 &&
+ "expecting 2 operands (resource + size)");
+ return builder
+ .create<IREE::Stream::AsyncTransferOp>(
+ loc, resourceValue.getType(), resourceValue, resourceSize,
+ resourceSize,
+ /*source_affinity=*/nullptr,
+ /*result_affinity=*/nullptr)
+ .result();
+ });
+
+ populateUtilConversionPatterns(context, conversionTarget, typeConverter,
+ patterns);
+ populateUtilToStreamConversionPatterns(context, conversionTarget,
+ typeConverter, patterns);
+
+ populateStandardToStreamConversionPatterns(context, conversionTarget,
+ typeConverter, patterns);
+ populateFlowToStreamConversionPatterns(context, conversionTarget,
+ typeConverter, patterns);
+
+ conversionTarget.markUnknownOpDynamicallyLegal(
+ [&](Operation *op) -> bool { return !doesOperationNeedWrapping(op); });
+ patterns.insert<GenericResourcePattern>(context, typeConverter);
+
+ // NOTE: we allow ops that we don't know about to allow custom dialects
+ // that don't need anything Stream-specific to pass through.
+ conversionTarget.addLegalOp<UnrealizedConversionCastOp>();
+ if (failed(applyPartialConversion(getOperation(), conversionTarget,
+ std::move(patterns)))) {
+ return signalPassFailure();
+ }
+ }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createConvertToStreamPass() {
+ return std::make_unique<ConvertToStreamPass>();
+}
+
+} // namespace Stream
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Stream/Transforms/OutlineConstants.cpp b/iree/compiler/Dialect/Stream/Transforms/OutlineConstants.cpp
new file mode 100644
index 0000000..60fa6ef
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/OutlineConstants.cpp
@@ -0,0 +1,129 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <utility>
+
+#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Stream/Transforms/PassDetail.h"
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
+#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace Stream {
+
+// Returns true if |value| is worth outlining (large, etc).
+static bool isOutlinableValue(Attribute value) {
+ if (auto elementsAttr = value.dyn_cast<DenseElementsAttr>()) {
+ // Don't outline splats - we want those fused.
+ return !elementsAttr.isSplat();
+ }
+ return false;
+}
+
+struct ConstantDef {
+ Operation *op;
+ Type type;
+ ElementsAttr value;
+};
+
+// Returns a list of all constant-like shaped data ops in the module.
+static SmallVector<ConstantDef> findConstantsInModule(mlir::ModuleOp moduleOp) {
+ SmallVector<ConstantDef> results;
+ for (auto callableOp : moduleOp.getOps<CallableOpInterface>()) {
+ for (auto &block : *callableOp.getCallableRegion()) {
+ for (auto &op : block.getOperations()) {
+ if (auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
+ if (isOutlinableValue(constantOp.value())) {
+ results.push_back(ConstantDef{
+ constantOp,
+ constantOp.getType(),
+ constantOp.value().cast<ElementsAttr>(),
+ });
+ }
+ }
+ }
+ }
+ }
+ return results;
+}
+
+class OutlineConstantsPass : public OutlineConstantsBase<OutlineConstantsPass> {
+ public:
+ OutlineConstantsPass() = default;
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<mlir::StandardOpsDialect>();
+ registry.insert<mlir::arith::ArithmeticDialect>();
+ registry.insert<IREE::Util::UtilDialect>();
+ }
+
+ void runOnOperation() override {
+ auto moduleOp = getOperation();
+ if (moduleOp.getBody()->empty()) return;
+
+ SymbolTable moduleSymbols(moduleOp);
+ std::string baseName = "_constant";
+
+ // Create all top-level util.globals from constants in the module.
+ OpBuilder moduleBuilder(&moduleOp.getBody()->front());
+ std::vector<std::pair<Operation *, IREE::Util::GlobalOp>> replacements;
+ for (auto &def : findConstantsInModule(moduleOp)) {
+ // New immutable global takes the constant attribute in its specified
+ // encoding.
+ auto globalOp = moduleBuilder.create<IREE::Util::GlobalOp>(
+ def.op->getLoc(), baseName, /*isMutable=*/false, def.type, def.value);
+ globalOp.setPrivate();
+ moduleSymbols.insert(globalOp); // uniques name
+ replacements.emplace_back(def.op, globalOp);
+
+ // Prevent the variable from being re-inlined if the canonicalizer runs.
+ // By the time we've outlined things here we are sure we want them
+ // outlined even if the user runs an arbitrary number of passes between
+ // now and when we may use that information (HAL constant pooling, etc).
+ globalOp->setAttr("noinline", moduleBuilder.getUnitAttr());
+ }
+
+ // Replace all of the constants with lookups for the new variables.
+ for (auto pair : replacements) {
+ auto *originalOp = pair.first;
+ auto globalOp = pair.second;
+ OpBuilder builder(moduleOp.getContext());
+ builder.setInsertionPoint(originalOp);
+ auto loadOp = builder.create<IREE::Util::GlobalLoadOp>(
+ originalOp->getLoc(), globalOp.type(), SymbolRefAttr::get(globalOp));
+
+ Value replacement;
+ if (auto constantOp = dyn_cast<arith::ConstantOp>(originalOp)) {
+ // Directly replace constant with global constant value.
+ replacement = loadOp.result();
+ } else {
+ llvm_unreachable("unhandled constant op type");
+ }
+
+ originalOp->getResult(0).replaceAllUsesWith(replacement);
+ originalOp->erase();
+ }
+ }
+};
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createOutlineConstantsPass() {
+ return std::make_unique<OutlineConstantsPass>();
+}
+
+} // namespace Stream
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Stream/Transforms/PassDetail.h b/iree/compiler/Dialect/Stream/Transforms/PassDetail.h
new file mode 100644
index 0000000..6fab181
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/PassDetail.h
@@ -0,0 +1,25 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_DIALECT_STREAM_TRANSFORMS_PASS_DETAIL_H_
+#define IREE_COMPILER_DIALECT_STREAM_TRANSFORMS_PASS_DETAIL_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace Stream {
+
+#define GEN_PASS_CLASSES
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h.inc" // IWYU pragma: keep
+
+} // namespace Stream
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_DIALECT_STREAM_TRANSFORMS_PASS_DETAIL_H_
diff --git a/iree/compiler/Dialect/Stream/Transforms/Passes.cpp b/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
new file mode 100644
index 0000000..dd108bd
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
@@ -0,0 +1,207 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
+
+#include <memory>
+
+#include "iree/compiler/Dialect/Shape/Transforms/Passes.h"
+#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
+#include "mlir/Pass/PassOptions.h"
+#include "mlir/Pass/PassRegistry.h"
+#include "mlir/Transforms/Passes.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace Stream {
+
+//===----------------------------------------------------------------------===//
+// Utilities
+//===----------------------------------------------------------------------===//
+
+static void addCleanupPatterns(OpPassManager &passManager) {
+ // Standard MLIR cleanup.
+ passManager.addPass(mlir::createCanonicalizerPass());
+ passManager.addPass(mlir::createCSEPass());
+
+ // Cleanup and canonicalization of util.global (and other util ops).
+ passManager.addPass(IREE::Util::createApplyPatternsPass());
+ passManager.addPass(IREE::Util::createFoldGlobalsPass());
+ passManager.addPass(IREE::Util::createFuseGlobalsPass());
+
+ // Simplify util.global accesses; this can help with data flow tracking as
+ // redundant store-loads are removed.
+ passManager.addNestedPass<IREE::Util::InitializerOp>(
+ IREE::Util::createSimplifyGlobalAccessesPass());
+ passManager.addNestedPass<mlir::FuncOp>(
+ IREE::Util::createSimplifyGlobalAccessesPass());
+}
+
+//===----------------------------------------------------------------------===//
+// -iree-stream-tensor-transformation-pipeline
+//===----------------------------------------------------------------------===//
+
+void buildStreamTensorPassPipeline(OpPassManager &passManager,
+ const TransformOptions &transformOptions) {
+ //----------------------------------------------------------------------------
+ // Input cleanup and simplification
+ //----------------------------------------------------------------------------
+
+ // Verify we support the program.
+ passManager.addPass(IREE::Stream::createVerifyInputPass());
+
+ // Turn all constant ops into global variables and fix up the IR.
+ // As many locations change and constants are deduplicated we'll end up with
+ // a lot of extraneous IR (mostly global loads) and clean those up here.
+ passManager.addPass(IREE::Stream::createOutlineConstantsPass());
+
+ // Perform cleanup after constnat simplification as more canonicalizers may be
+ // able to kick in.
+ addCleanupPatterns(passManager);
+
+ //----------------------------------------------------------------------------
+ // Conversion
+ //----------------------------------------------------------------------------
+
+ // Converts from all input dialects into various levels of the stream dialect.
+ // Tensor-like things go to stream.tensor.* ops while lower level buffer-like
+ // things will go to stream.async.* ops.
+ passManager.addPass(IREE::Stream::createConvertToStreamPass());
+
+ // No more tensor.*/etc ops are allowed. This is conservative - there may be
+ // a lot of ops we convert but this will catch the majority of stragglers.
+ passManager.addPass(IREE::Stream::createVerifyLoweringToTensorsPass());
+
+ //----------------------------------------------------------------------------
+ // Constant/variable optimization
+ //----------------------------------------------------------------------------
+
+ // Cleanup globals that were created during conversion.
+ addCleanupPatterns(passManager);
+
+ // Bring all initializers together so that we can schedule them.
+ passManager.addPass(IREE::Util::createCombineInitializersPass());
+
+ //----------------------------------------------------------------------------
+ // Stream affinity/assignment
+ //----------------------------------------------------------------------------
+
+ // TODO(benvanik): pin based on target backends here.
+ // TODO(benvanik): compute affinities for executables.
+ // TODO(benvanik): annotate all dispatches with preferred executable affinity.
+ // TODO(benvanik): DFA to specify all value affinities and pin dispatches.
+}
+
+//===----------------------------------------------------------------------===//
+// -iree-stream-async-transformation-pipeline
+//===----------------------------------------------------------------------===//
+
+void buildStreamAsyncPassPipeline(OpPassManager &passManager,
+ const TransformOptions &transformOptions) {}
+
+//===----------------------------------------------------------------------===//
+// -iree-stream-cmd-transformation-pipeline
+//===----------------------------------------------------------------------===//
+
+void buildStreamCmdPassPipeline(OpPassManager &passManager,
+ const TransformOptions &transformOptions) {}
+
+//===----------------------------------------------------------------------===//
+// -iree-stream-optimization-pipeline
+//===----------------------------------------------------------------------===//
+
+void buildStreamOptimizationPassPipeline(
+ OpPassManager &passManager, const TransformOptions &transformOptions) {}
+
+//===----------------------------------------------------------------------===//
+// -iree-stream-transformation-pipeline
+//===----------------------------------------------------------------------===//
+
+void buildStreamTransformPassPipeline(
+ OpPassManager &passManager, const TransformOptions &transformOptions) {
+ //----------------------------------------------------------------------------
+ // Primary pipeline stages (required)
+ //----------------------------------------------------------------------------
+
+ buildStreamTensorPassPipeline(passManager, transformOptions);
+ buildStreamAsyncPassPipeline(passManager, transformOptions);
+ buildStreamCmdPassPipeline(passManager, transformOptions);
+
+ //----------------------------------------------------------------------------
+ // Optimizations (may be required by some targets)
+ //----------------------------------------------------------------------------
+
+ buildStreamOptimizationPassPipeline(passManager, transformOptions);
+
+ //----------------------------------------------------------------------------
+ // Post-pipeline cleanup
+ //----------------------------------------------------------------------------
+
+ // Forming streams involves a fair amount of subgraph stitching, which can
+ // cause duplication. Run CSE to collapse.
+ addCleanupPatterns(passManager);
+
+ // Symbol DCE any remaining variables/functions that are now no longer
+ // required.
+ passManager.addPass(mlir::createSymbolDCEPass());
+}
+
+//===----------------------------------------------------------------------===//
+// Registration
+//===----------------------------------------------------------------------===//
+
+void registerStreamTransformPassPipelines() {
+ PassPipelineRegistration<TransformOptions> tensorPassPipeline(
+ "iree-stream-tensor-transformation-pipeline",
+ "Lowers source dialects into stream.tensor.* IR.",
+ [](OpPassManager &passManager, const TransformOptions &transformOptions) {
+ buildStreamTensorPassPipeline(passManager, transformOptions);
+ });
+ PassPipelineRegistration<TransformOptions> asyncPassPipeline(
+ "iree-stream-async-transformation-pipeline",
+ "Lowers stream.tensor.* to stream.async.* IR.",
+ [](OpPassManager &passManager, const TransformOptions &transformOptions) {
+ buildStreamAsyncPassPipeline(passManager, transformOptions);
+ });
+ PassPipelineRegistration<TransformOptions> cmdPassPipeline(
+ "iree-stream-cmd-transformation-pipeline",
+ "Lowers stream.async.* to stream.cmd.* IR.",
+ [](OpPassManager &passManager, const TransformOptions &transformOptions) {
+ buildStreamCmdPassPipeline(passManager, transformOptions);
+ });
+ PassPipelineRegistration<TransformOptions> optimizationPassPipeline(
+ "iree-stream-optimization-pipeline",
+ "Optimizes stream commands and resources (may be required for some "
+ "targets).",
+ [](OpPassManager &passManager, const TransformOptions &transformOptions) {
+ buildStreamOptimizationPassPipeline(passManager, transformOptions);
+ });
+ PassPipelineRegistration<TransformOptions> transformPassPipeline(
+ "iree-stream-transformation-pipeline",
+ "Runs the full IREE stream dialect transformation pipeline.",
+ [](OpPassManager &passManager, const TransformOptions &transformOptions) {
+ buildStreamTransformPassPipeline(passManager, transformOptions);
+ });
+}
+
+namespace {
+#define GEN_PASS_REGISTRATION
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h.inc" // IWYU pragma: export
+} // namespace
+
+void registerStreamPasses() {
+ // Generated.
+ registerPasses();
+
+ // Pipelines.
+ registerStreamTransformPassPipelines();
+}
+
+} // namespace Stream
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Stream/Transforms/Passes.h b/iree/compiler/Dialect/Stream/Transforms/Passes.h
new file mode 100644
index 0000000..513fdd0
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/Passes.h
@@ -0,0 +1,104 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_DIALECT_STREAM_TRANSFORMS_PASSES_H_
+#define IREE_COMPILER_DIALECT_STREAM_TRANSFORMS_PASSES_H_
+
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "llvm/ADT/StringMap.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace Stream {
+
+//===----------------------------------------------------------------------===//
+// Pipelines
+//===----------------------------------------------------------------------===//
+
+struct TransformOptions : public PassPipelineOptions<TransformOptions> {
+ // TODO(benvanik): options for async/sync overrides.
+
+ Option<bool> optimizeBindings{
+ *this, "optimize-bindings",
+ llvm::cl::desc(
+ "Enables binding fusion and dispatch site specialization."),
+ llvm::cl::init(true)};
+};
+
+// Adds a set of passes to the given pass manager that run the required flow
+// transforms in the canonical order.
+//
+// Most translation code should prefer to use this instead of manually adding
+// the passes themselves to ensure that expected pass ordering is observed.
+//
+// The expected usage is:
+// Input legalization by one of:
+// - Directly passing supported flow plus core ops
+// buildStreamTransformPassPipeline
+// <run conversion from flow to sequencer/hal/vm/etc>
+//
+// This is equivalent to:
+// buildStreamTensorPassPipeline
+// buildStreamAsyncPassPipeline
+// buildStreamCmdPassPipeline
+void buildStreamTransformPassPipeline(OpPassManager &passManager,
+ const TransformOptions &transformOptions);
+
+// Lowers from source dialects into stream.tensor.* IR.
+void buildStreamTensorPassPipeline(OpPassManager &passManager,
+ const TransformOptions &transformOptions);
+// Lowers stream.tensor.* IR into stream.async.* IR.
+void buildStreamAsyncPassPipeline(OpPassManager &passManager,
+ const TransformOptions &transformOptions);
+// Lowers stream.async.* IR into stream.cmd.* IR.
+void buildStreamCmdPassPipeline(OpPassManager &passManager,
+ const TransformOptions &transformOptions);
+// Optimizes stream commands (mostly optional).
+void buildStreamOptimizationPassPipeline(
+ OpPassManager &passManager, const TransformOptions &transformOptions);
+
+void registerStreamTransformPassPipelines();
+
+//===----------------------------------------------------------------------===//
+// Optimizations
+//===----------------------------------------------------------------------===//
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createOutlineConstantsPass();
+
+//===----------------------------------------------------------------------===//
+// Conversion
+//===----------------------------------------------------------------------===//
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createConvertToStreamPass();
+
+//===----------------------------------------------------------------------===//
+// Diagnostics
+//===----------------------------------------------------------------------===//
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createVerifyInputPass();
+std::unique_ptr<OperationPass<mlir::ModuleOp>>
+createVerifyLoweringToTensorsPass();
+std::unique_ptr<OperationPass<mlir::ModuleOp>>
+createVerifyLoweringToAsyncPass();
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createVerifyLoweringToCmdPass();
+
+//===----------------------------------------------------------------------===//
+// Register all Passes
+//===----------------------------------------------------------------------===//
+
+void registerStreamPasses();
+
+} // namespace Stream
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_DIALECT_STREAM_TRANSFORMS_PASSES_H_
diff --git a/iree/compiler/Dialect/Stream/Transforms/Passes.td b/iree/compiler/Dialect/Stream/Transforms/Passes.td
new file mode 100644
index 0000000..259183b
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/Passes.td
@@ -0,0 +1,72 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_DIALECT_STREAM_PASSES
+#define IREE_DIALECT_STREAM_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+//===----------------------------------------------------------------------===//
+// Optimizations
+//===----------------------------------------------------------------------===//
+
+def OutlineConstants :
+ Pass<"iree-stream-outline-constants", "mlir::ModuleOp"> {
+ let summary = "Outlines tensor constants into util.globals at the module level.";
+ let constructor = [{
+ mlir::iree_compiler::IREE::Stream::createOutlineConstantsPass()
+ }];
+}
+
+//===----------------------------------------------------------------------===//
+// Conversion
+//===----------------------------------------------------------------------===//
+
+def ConvertToStream :
+ Pass<"iree-stream-conversion", "mlir::ModuleOp"> {
+ let summary = "Converts from flow/std/etc dialects into the stream dialect.";
+ let constructor = [{
+ mlir::iree_compiler::IREE::Stream::createConvertToStreamPass()
+ }];
+}
+
+//===----------------------------------------------------------------------===//
+// Diagnostics
+//===----------------------------------------------------------------------===//
+
+def VerifyInput :
+ Pass<"iree-stream-verify-input", "mlir::ModuleOp"> {
+ let summary = "Verifies that input dialects are supported by the streams dialect.";
+ let constructor = [{
+ mlir::iree_compiler::IREE::Stream::createVerifyInputPass()
+ }];
+}
+
+def VerifyLoweringToTensors :
+ Pass<"iree-stream-verify-lowering-to-tensors", "mlir::ModuleOp"> {
+ let summary = "Verifies that input dialects are converted to stream.tensor.* ops.";
+ let constructor = [{
+ mlir::iree_compiler::IREE::Stream::createVerifyLoweringToTensorsPass()
+ }];
+}
+
+def VerifyLoweringToAsync :
+ Pass<"iree-stream-verify-lowering-to-async", "mlir::ModuleOp"> {
+ let summary = "Verifies that all stream.tensor.* ops and types are fully lowered to stream.async.* ops.";
+ let constructor = [{
+ mlir::iree_compiler::IREE::Stream::createVerifyLoweringToAsyncPass()
+ }];
+}
+
+def VerifyLoweringToCmd :
+ Pass<"iree-stream-verify-lowering-to-cmd", "mlir::ModuleOp"> {
+ let summary = "Verifies that all stream.async.* ops and types are fully lowered to stream.cmd.* ops.";
+ let constructor = [{
+ mlir::iree_compiler::IREE::Stream::createVerifyLoweringToCmdPass()
+ }];
+}
+
+#endif // IREE_DIALECT_STREAM_PASSES
diff --git a/iree/compiler/Dialect/Stream/Transforms/VerifyLowerings.cpp b/iree/compiler/Dialect/Stream/Transforms/VerifyLowerings.cpp
new file mode 100644
index 0000000..8ecdf2c
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/VerifyLowerings.cpp
@@ -0,0 +1,420 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <utility>
+
+#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamTraits.h"
+#include "iree/compiler/Dialect/Stream/Transforms/PassDetail.h"
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
+#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace Stream {
+
+//===----------------------------------------------------------------------===//
+// Base pass utility
+//===----------------------------------------------------------------------===//
+
+class Verifier {
+ public:
+ enum class Legality {
+ LEGAL,
+ RECURSIVELY_LEGAL,
+ ILLEGAL,
+ };
+
+ using OpVerifierFn = std::function<Optional<Legality>(Operation *op)>;
+ using TypeVerifierFn = std::function<Legality(Type type)>;
+
+ void addIllegalDialect(StringRef dialectName) {
+ dialectLegality.insert({dialectName, Legality::ILLEGAL});
+ }
+ template <typename DialectT>
+ void addIllegalDialect() {
+ addIllegalDialect(DialectT::getDialectNamespace());
+ }
+
+ template <typename OpT>
+ void addLegalOp() {
+ opLegality.insert({OpT::getOperationName(), Legality::LEGAL});
+ }
+
+ template <typename OpT>
+ void addRecursivelyLegalOp() {
+ opLegality.insert({OpT::getOperationName(), Legality::RECURSIVELY_LEGAL});
+ }
+
+ template <typename OpT>
+ void addIllegalOp() {
+ opLegality.insert({OpT::getOperationName(), Legality::ILLEGAL});
+ }
+
+ void addOpVerifier(std::function<Optional<Legality>(Operation *)> fn) {
+ opVerifiers.push_back(fn);
+ }
+
+ template <typename OpT>
+ void addOpVerifier(std::function<Optional<Legality>(OpT)> fn) {
+ auto wrapperFn = [=](Operation *baseOp) -> Optional<Legality> {
+ if (auto op = dyn_cast<OpT>(baseOp)) {
+ return fn(op);
+ }
+ return llvm::None;
+ };
+ opVerifiers.push_back(wrapperFn);
+ }
+
+ template <typename TypeT>
+ void addIllegalType() {
+ typeLegality.insert({TypeID::get<TypeT>(), Legality::ILLEGAL});
+ }
+
+ template <typename TypeT>
+ void addTypeVerifier(std::function<Legality(TypeT)> fn) {
+ auto wrapperFn = [=](Type baseType) { return fn(baseType.cast<TypeT>()); };
+ if (typeVerifiers.insert({TypeID::get<TypeT>(), wrapperFn}).second ==
+ false) {
+ llvm_unreachable("already registered for this type");
+ }
+ }
+
+ LogicalResult run(Operation *rootOp) {
+ bool foundAnyIllegal = false;
+ rootOp->walk<WalkOrder::PreOrder>([&](Operation *op) {
+ auto walkResult = WalkResult::advance();
+
+ // Check for op legality - can skip the expensive work if known-illegal.
+ auto legality = getOpLegality(op);
+ switch (legality) {
+ case Legality::LEGAL:
+ // Op itself is legal but may not have valid operands/results.
+ break;
+ case Legality::RECURSIVELY_LEGAL:
+ // If the entire op w/ nested ops is legal then skip.
+ return WalkResult::skip();
+ default:
+ case Legality::ILLEGAL:
+ // Early-exit on illegal ops without recursing.
+ emitIllegalOpError(op);
+ foundAnyIllegal = true;
+ return WalkResult::skip();
+ }
+
+ // Check types for operands/results.
+ for (auto operandType : llvm::enumerate(op->getOperandTypes())) {
+ if (isTypeLegal(operandType.value())) continue;
+ emitIllegalTypeError(op, "operand", operandType.index(),
+ operandType.value());
+ foundAnyIllegal = true;
+ }
+ for (auto resultType : llvm::enumerate(op->getResultTypes())) {
+ if (isTypeLegal(resultType.value())) continue;
+ emitIllegalTypeError(op, "result", resultType.index(),
+ resultType.value());
+ foundAnyIllegal = true;
+ }
+
+ return walkResult;
+ });
+ return success(!foundAnyIllegal);
+ }
+
+ private:
+ Legality getOpLegality(Operation *op) {
+ auto opName = op->getName();
+
+ // Check specific ops first (we may override dialect settings).
+ {
+ auto legalityIt = opLegality.find(opName.getStringRef());
+ if (legalityIt != opLegality.end()) {
+ return legalityIt->second;
+ }
+ }
+
+ // Check all op verifiers (usually used for interface checks).
+ for (auto &opVerifier : opVerifiers) {
+ auto legalOr = opVerifier(op);
+ if (legalOr.hasValue()) {
+ return legalOr.getValue();
+ }
+ }
+
+ // If no op carveout is applied then check to see if the dialect is
+ // allowed at all.
+ {
+ auto legalityIt = dialectLegality.find(opName.getDialectNamespace());
+ if (legalityIt != dialectLegality.end()) {
+ return legalityIt->second;
+ }
+ }
+
+ // Assume legal by default.
+ return Legality::LEGAL;
+ }
+
+ bool isTypeLegal(Type type) {
+ // TODO(benvanik): subelements interface checks using recursive legality.
+
+ // Defer to verifiers first.
+ auto it = typeVerifiers.find(type.getTypeID());
+ if (it != typeVerifiers.end()) {
+ return it->second(type) != Legality::ILLEGAL;
+ }
+
+ // Check legality of the base type.
+ {
+ auto legalityIt = typeLegality.find(type.getTypeID());
+ if (legalityIt != typeLegality.end()) {
+ return legalityIt->second != Legality::ILLEGAL;
+ }
+ }
+
+ // Assume legal by default.
+ return true;
+ }
+
+ void emitIllegalOpError(Operation *op) {
+ op->emitOpError()
+ << "illegal for this phase of lowering in the stream dialect; "
+ "expected to have been converted or removed";
+ }
+
+ void emitIllegalTypeError(Operation *op, StringRef location, unsigned idx,
+ Type type) {
+ op->emitOpError()
+ << location << " " << idx << " type " << type
+ << " illegal for this phase of lowering in the stream dialect";
+ }
+
+ DenseMap<StringRef, Legality> dialectLegality;
+ DenseMap<StringRef, Legality> opLegality;
+ SmallVector<OpVerifierFn> opVerifiers;
+ DenseMap<TypeID, Legality> typeLegality;
+ DenseMap<TypeID, TypeVerifierFn> typeVerifiers;
+};
+
+static void markStreamTensorOpsIllegal(Verifier &verifier) {
+ verifier.addOpVerifier([](Operation *op) -> Optional<Verifier::Legality> {
+ if (op->hasTrait<OpTrait::IREE::Stream::TensorPhaseOp>()) {
+ return Verifier::Legality::ILLEGAL;
+ }
+ return llvm::None;
+ });
+}
+
+static void markStreamAsyncOpsIllegal(Verifier &verifier) {
+ verifier.addOpVerifier([](Operation *op) -> Optional<Verifier::Legality> {
+ if (op->hasTrait<OpTrait::IREE::Stream::AsyncPhaseOp>()) {
+ return Verifier::Legality::ILLEGAL;
+ }
+ return llvm::None;
+ });
+}
+
+static void markStreamCmdOpsIllegal(Verifier &verifier) {
+ verifier.addOpVerifier([](Operation *op) -> Optional<Verifier::Legality> {
+ if (op->hasTrait<OpTrait::IREE::Stream::CmdPhaseOp>()) {
+ return Verifier::Legality::ILLEGAL;
+ }
+ return llvm::None;
+ });
+}
+
+//===----------------------------------------------------------------------===//
+// -iree-stream-verify-input
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+class VerifyInputPass : public VerifyInputBase<VerifyInputPass> {
+ public:
+ VerifyInputPass() = default;
+
+ void runOnOperation() override {
+ Verifier verifier;
+
+ // TODO(#7432): add indirect global expansion support to streams.
+ verifier.addIllegalOp<IREE::Util::GlobalAddressOp>();
+ verifier.addIllegalOp<IREE::Util::GlobalLoadIndirectOp>();
+ verifier.addIllegalOp<IREE::Util::GlobalStoreIndirectOp>();
+
+ if (failed(verifier.run(getOperation()))) {
+ return signalPassFailure();
+ }
+ }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createVerifyInputPass() {
+ return std::make_unique<VerifyInputPass>();
+}
+
+//===----------------------------------------------------------------------===//
+// -iree-stream-verify-lowering-to-tensors
+//===----------------------------------------------------------------------===//
+
+static void markTensorInputsIllegal(Verifier &verifier) {
+ // Tensorish dialects should all be either converted or outlined into
+ // executables. Everything should be in resources now.
+ verifier.addIllegalDialect("tensor");
+ verifier.addIllegalDialect("linalg");
+
+ // We don't allow the flow dialect except for inside of executables for which
+ // we don't yet have a full mapping to in the stream dialect.
+ // TODO(#7277): remove this carveout once we switch over to streams fully.
+ verifier.addIllegalDialect("flow");
+ verifier.addRecursivelyLegalOp<IREE::Stream::ExecutableOp>();
+}
+
+namespace {
+
+class VerifyLoweringToTensorsPass
+ : public VerifyLoweringToTensorsBase<VerifyLoweringToTensorsPass> {
+ public:
+ VerifyLoweringToTensorsPass() = default;
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<IREE::Stream::StreamDialect>();
+ registry.insert<IREE::Util::UtilDialect>();
+ }
+
+ void runOnOperation() override {
+ // We cannot have stream.cmd.* ops mixed with stream.tensor/async.* ops
+ // as they use different memory models.
+ Verifier verifier;
+ markTensorInputsIllegal(verifier);
+ markStreamCmdOpsIllegal(verifier);
+ if (failed(verifier.run(getOperation()))) {
+ return signalPassFailure();
+ }
+ }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>>
+createVerifyLoweringToTensorsPass() {
+ return std::make_unique<VerifyLoweringToTensorsPass>();
+}
+
+//===----------------------------------------------------------------------===//
+// -iree-stream-verify-lowering-to-tensors
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+class VerifyLoweringToAsyncPass
+ : public VerifyLoweringToAsyncBase<VerifyLoweringToAsyncPass> {
+ public:
+ VerifyLoweringToAsyncPass() = default;
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<IREE::Stream::StreamDialect>();
+ registry.insert<IREE::Util::UtilDialect>();
+ }
+
+ void runOnOperation() override {
+ // We cannot have stream.cmd.* ops mixed with stream.tensor/async.* ops
+ // as they use different memory models.
+ Verifier verifier;
+ markTensorInputsIllegal(verifier);
+ markStreamTensorOpsIllegal(verifier);
+ markStreamCmdOpsIllegal(verifier);
+
+ // All resources should have had their usage assigned.
+ verifier.addTypeVerifier<IREE::Stream::ResourceType>([](auto type) {
+ if (type.getLifetime() == IREE::Stream::Lifetime::Unknown) {
+ return Verifier::Legality::ILLEGAL;
+ }
+ return Verifier::Legality::LEGAL;
+ });
+
+ // All streamable ops should be inside of execution regions.
+ verifier.addOpVerifier<IREE::Stream::StreamableOpInterface>(
+ [](auto op) -> Optional<Verifier::Legality> {
+ // Allow metadata ops outside of execution regions.
+ if (op.isMetadata()) return Verifier::Legality::LEGAL;
+
+ // TODO(benvanik): execution region interface to make this generic.
+ if (!op->template getParentOfType<IREE::Stream::AsyncExecuteOp>()) {
+ op->emitOpError()
+ << ": streamable op expected to be in an execution region";
+ return Verifier::Legality::ILLEGAL;
+ }
+ return llvm::None;
+ });
+
+ if (failed(verifier.run(getOperation()))) {
+ return signalPassFailure();
+ }
+ }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>>
+createVerifyLoweringToAsyncPass() {
+ return std::make_unique<VerifyLoweringToAsyncPass>();
+}
+
+//===----------------------------------------------------------------------===//
+// -iree-stream-verify-lowering-to-cmd
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+class VerifyLoweringToCmdPass
+ : public VerifyLoweringToCmdBase<VerifyLoweringToCmdPass> {
+ public:
+ VerifyLoweringToCmdPass() = default;
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<IREE::Stream::StreamDialect>();
+ registry.insert<IREE::Util::UtilDialect>();
+ }
+
+ void runOnOperation() override {
+ Verifier verifier;
+ markTensorInputsIllegal(verifier);
+ markStreamTensorOpsIllegal(verifier);
+ markStreamAsyncOpsIllegal(verifier);
+
+ // All resources should have had their usage assigned.
+ verifier.addTypeVerifier<IREE::Stream::ResourceType>([](auto type) {
+ if (type.getLifetime() == IREE::Stream::Lifetime::Unknown) {
+ return Verifier::Legality::ILLEGAL;
+ }
+ return Verifier::Legality::LEGAL;
+ });
+
+ if (failed(verifier.run(getOperation()))) {
+ return signalPassFailure();
+ }
+ }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createVerifyLoweringToCmdPass() {
+ return std::make_unique<VerifyLoweringToCmdPass>();
+}
+
+} // namespace Stream
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/iree/compiler/Dialect/Stream/Transforms/test/BUILD b/iree/compiler/Dialect/Stream/Transforms/test/BUILD
new file mode 100644
index 0000000..9b9e29f
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/test/BUILD
@@ -0,0 +1,29 @@
+# Copyright 2021 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+load("//iree:lit_test.bzl", "iree_lit_test_suite")
+load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+iree_lit_test_suite(
+ name = "lit",
+ srcs = enforce_glob(
+ [
+ "convert_to_stream.mlir",
+ "outline_constants.mlir",
+ ],
+ include = ["*.mlir"],
+ ),
+ data = [
+ "//iree/tools:IreeFileCheck",
+ "//iree/tools:iree-opt",
+ ],
+)
diff --git a/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt b/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
new file mode 100644
index 0000000..7401676
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
@@ -0,0 +1,24 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# iree/compiler/Dialect/Stream/Transforms/test/BUILD #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_lit_test_suite(
+ NAME
+ lit
+ SRCS
+ "convert_to_stream.mlir"
+ "outline_constants.mlir"
+ DATA
+ iree::tools::IreeFileCheck
+ iree::tools::iree-opt
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/compiler/Dialect/Stream/Transforms/test/convert_to_stream.mlir b/iree/compiler/Dialect/Stream/Transforms/test/convert_to_stream.mlir
new file mode 100644
index 0000000..eb00f10
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/test/convert_to_stream.mlir
@@ -0,0 +1,50 @@
+// RUN: iree-opt -split-input-file -iree-stream-conversion %s | IreeFileCheck %s
+
+// CHECK: stream.executable private @executable
+flow.executable private @executable {
+ // CHECK: stream.executable.export public @dispatch
+ flow.dispatch.entry public @dispatch attributes {workgroup_rank = 3 : index}
+ builtin.module {
+ // CHECK: func @dispatch(%arg0: !stream.binding, %arg1: !stream.binding, %arg2: index, %arg3: index)
+ func @dispatch(%arg0: !flow.dispatch.tensor<readonly:?x4xf32>, %arg1: !flow.dispatch.tensor<writeonly:?xf32>,
+ %arg0_dim0: index, %arg1_dim0: index) {
+ // CHECK: %[[ARG0_SHAPE:.+]] = shapex.make_ranked_shape %arg2 : (index) -> !shapex.ranked_shape<[?,4]>
+ %arg0_shape = shapex.make_ranked_shape %arg0_dim0 : (index) -> !shapex.ranked_shape<[?,4]>
+ // CHECK: %[[ARG0_DIM0:.+]] = shapex.ranked_dim %[[ARG0_SHAPE]][0] : !shapex.ranked_shape<[?,4]> -> index
+ // CHECK: %[[ARG0_SPAN:.+]] = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:?x4xf32>{%[[ARG0_DIM0]]}
+ // CHECK: = flow.dispatch.tie_shape %[[ARG0_SPAN]], %[[ARG0_SHAPE]] : (!flow.dispatch.tensor<readonly:?x4xf32>, !shapex.ranked_shape<[?,4]>) -> !flow.dispatch.tensor<readonly:?x4xf32>
+ %arg0_shaped = flow.dispatch.tie_shape %arg0, %arg0_shape : (!flow.dispatch.tensor<readonly:?x4xf32>, !shapex.ranked_shape<[?,4]>) -> !flow.dispatch.tensor<readonly:?x4xf32>
+
+ // CHECK: %[[ARG1_SHAPE:.+]] = shapex.make_ranked_shape %arg3 : (index) -> !shapex.ranked_shape<[?]>
+ %arg1_shape = shapex.make_ranked_shape %arg1_dim0 : (index) -> !shapex.ranked_shape<[?]>
+ // CHECK: %[[ARG1_DIM0:.+]] = shapex.ranked_dim %[[ARG1_SHAPE]][0] : !shapex.ranked_shape<[?]> -> index
+ // CHECK: %[[ARG1_SPAN:.+]] = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:?xf32>{%[[ARG1_DIM0]]}
+ // CHECK: = flow.dispatch.tie_shape %[[ARG1_SPAN]], %[[ARG1_SHAPE]] : (!flow.dispatch.tensor<writeonly:?xf32>, !shapex.ranked_shape<[?]>) -> !flow.dispatch.tensor<writeonly:?xf32>
+ %arg1_shaped = flow.dispatch.tie_shape %arg1, %arg1_shape : (!flow.dispatch.tensor<writeonly:?xf32>, !shapex.ranked_shape<[?]>) -> !flow.dispatch.tensor<writeonly:?xf32>
+
+ return
+ }
+ }
+}
+
+// CHECK-LABEL: @simple_mul
+func @simple_mul(%arg0: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub} {
+ // CHECK: %[[DIM0:.+]] = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
+ %dim0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
+ // CHECK: %[[ARG0_SIZE:.+]] = stream.tensor.sizeof tensor<?x4xf32>{%[[DIM0]]} : index
+ // CHECK: %[[ARG0_IMPORT:.+]] = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x4xf32>{%[[DIM0]]} in !stream.resource<external>{%[[ARG0_SIZE]]}
+ // CHECK: %[[ARG0_T:.+]] = stream.async.transfer %[[ARG0_IMPORT]] : !stream.resource<external>{%[[ARG0_SIZE]]} -> !stream.resource<*>{%[[ARG0_SIZE]]}
+ %0 = hal.tensor.cast %arg0 : !hal.buffer_view -> tensor<?x4xf32>{%dim0}
+
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ // CHECK: %[[RET0_SIZE:.+]] = stream.tensor.sizeof tensor<?xf32>{%[[DIM0]]} : index
+ // CHECK: %[[RET0:.+]] = stream.async.dispatch @executable::@dispatch[%c4, %c1, %c1](%[[ARG0_T]]) : (!stream.resource<*>{%[[ARG0_SIZE]]}) -> !stream.resource<*>{%[[RET0_SIZE]]}
+ %1 = flow.dispatch @executable::@dispatch[%c4, %c1, %c1](%0) : (tensor<?x4xf32>{%dim0}) -> tensor<?xf32>{%dim0}
+
+ // CHECK: %[[RET0_T:.+]] = stream.async.transfer %[[RET0]] : !stream.resource<*>{%[[RET0_SIZE]]} -> !stream.resource<external>{%[[RET0_SIZE]]}
+ // CHECK: %[[RET0_EXPORT:.+]] = stream.tensor.export %[[RET0_T]] : tensor<?xf32>{%[[DIM0]]} in !stream.resource<external>{%[[RET0_SIZE]]} -> !hal.buffer_view
+ %2 = hal.tensor.cast %1 : tensor<?xf32>{%dim0} -> !hal.buffer_view
+ // CHECK: return %[[RET0_EXPORT]] : !hal.buffer_view
+ return %2 : !hal.buffer_view
+}
diff --git a/iree/compiler/Dialect/Stream/Transforms/test/outline_constants.mlir b/iree/compiler/Dialect/Stream/Transforms/test/outline_constants.mlir
new file mode 100644
index 0000000..9bebe73
--- /dev/null
+++ b/iree/compiler/Dialect/Stream/Transforms/test/outline_constants.mlir
@@ -0,0 +1,30 @@
+// RUN: iree-opt -split-input-file -iree-stream-outline-constants %s | IreeFileCheck %s
+
+// CHECK-LABEL: @scalarConstant
+func @scalarConstant() {
+ // CHECK: = arith.constant 0 : i32
+ %cst = arith.constant 0 : i32
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @sparseConstant
+func @sparseConstant() {
+ // CHECK: = arith.constant dense<1.200000e+00> : tensor<512x128xf32>
+ %cst = arith.constant dense<1.2> : tensor<512x128xf32>
+ return
+}
+
+// -----
+
+// CHECK: util.global private @_constant {noinline} = dense<[0.0287729427, 0.0297581609]> : tensor<2xf32>
+// CHECK-NEXT: util.global private @_constant_0 {noinline} = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00]> : tensor<8xf32>
+// CHECK-LABEL: @denseConstants
+func @denseConstants() {
+ // CHECK: = util.global.load @_constant : tensor<2xf32>
+ %cst_0 = arith.constant dense<[0.0287729427, 0.0297581609]> : tensor<2xf32>
+ // CHECK-NEXT: = util.global.load @_constant_0 : tensor<8xf32>
+ %cst_1 = arith.constant dense<[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]> : tensor<8xf32>
+ return
+}
diff --git a/iree/compiler/Dialect/Util/IR/UtilOps.cpp b/iree/compiler/Dialect/Util/IR/UtilOps.cpp
index 792b34b..c4eb495 100644
--- a/iree/compiler/Dialect/Util/IR/UtilOps.cpp
+++ b/iree/compiler/Dialect/Util/IR/UtilOps.cpp
@@ -234,6 +234,13 @@
ParseResult parseShapedTiedResult(
OpAsmParser &parser, Type &resultType,
+ SmallVectorImpl<OpAsmParser::OperandType> &resultDims) {
+ ArrayAttr tiedOperands;
+ return parseShapedTiedResult(parser, resultType, resultDims, tiedOperands);
+}
+
+ParseResult parseShapedTiedResult(
+ OpAsmParser &parser, Type &resultType,
SmallVectorImpl<OpAsmParser::OperandType> &resultDims,
ArrayAttr &tiedOperands) {
OpAsmParser::OperandType tiedResult;
@@ -270,7 +277,7 @@
}
void printShapedTiedResult(OpAsmPrinter &p, Operation *op, Type resultType,
- ValueRange resultDims, ArrayAttr tiedOperands) {
+ ValueRange resultDims) {
auto tiedOp = cast<IREE::Util::TiedOpInterface>(op);
auto tiedOperandIndex = tiedOp.getTiedResultOperandIndex(0);
if (tiedOperandIndex.hasValue()) {
@@ -301,6 +308,11 @@
}
}
+void printShapedTiedResult(OpAsmPrinter &p, Operation *op, Type resultType,
+ ValueRange resultDims, ArrayAttr tiedOperands) {
+ printShapedTiedResult(p, op, resultType, resultDims);
+}
+
//===----------------------------------------------------------------------===//
// custom<ShapedFunctionType>
//===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/Util/IR/UtilOps.h b/iree/compiler/Dialect/Util/IR/UtilOps.h
index a00447a..02cc7ad 100644
--- a/iree/compiler/Dialect/Util/IR/UtilOps.h
+++ b/iree/compiler/Dialect/Util/IR/UtilOps.h
@@ -92,6 +92,22 @@
ParseResult parseShapedTiedResult(
OpAsmParser &parser, Type &resultType,
+ SmallVectorImpl<OpAsmParser::OperandType> &resultDims);
+inline ParseResult parseShapedTiedResult(OpAsmParser &parser, Type &resultType,
+ OpAsmParser::OperandType &resultDim) {
+ SmallVector<OpAsmParser::OperandType, 1> resultDims;
+ if (failed(parseShapedTiedResult(parser, resultType, resultDims))) {
+ return failure();
+ }
+ assert(resultDims.size() == 1 && "requires one dim");
+ resultDim = std::move(resultDims.front());
+ return success();
+}
+void printShapedTiedResult(OpAsmPrinter &p, Operation *op, Type resultType,
+ ValueRange resultDims);
+
+ParseResult parseShapedTiedResult(
+ OpAsmParser &parser, Type &resultType,
SmallVectorImpl<OpAsmParser::OperandType> &resultDims,
ArrayAttr &tiedOperands);
void printShapedTiedResult(OpAsmPrinter &p, Operation *op, Type resultType,
@@ -105,6 +121,7 @@
tiedOperands))) {
return failure();
}
+ assert(resultDims.size() == 1 && "requires one dim");
resultDim = std::move(resultDims.front());
return success();
}
diff --git a/iree/compiler/Translation/BUILD b/iree/compiler/Translation/BUILD
index 886a5a2..4b8de51 100644
--- a/iree/compiler/Translation/BUILD
+++ b/iree/compiler/Translation/BUILD
@@ -24,6 +24,7 @@
"//iree/compiler/Dialect/HAL/Conversion/HALToVM",
"//iree/compiler/Dialect/HAL/Target",
"//iree/compiler/Dialect/HAL/Transforms",
+ "//iree/compiler/Dialect/Stream/Transforms",
"//iree/compiler/Dialect/Util/Transforms",
"//iree/compiler/Dialect/VM/Conversion",
"//iree/compiler/Dialect/VM/Conversion/StandardToVM",
diff --git a/iree/compiler/Translation/CMakeLists.txt b/iree/compiler/Translation/CMakeLists.txt
index 823575e..bdaaa39 100644
--- a/iree/compiler/Translation/CMakeLists.txt
+++ b/iree/compiler/Translation/CMakeLists.txt
@@ -32,6 +32,7 @@
iree::compiler::Dialect::HAL::Conversion::HALToVM
iree::compiler::Dialect::HAL::Target
iree::compiler::Dialect::HAL::Transforms
+ iree::compiler::Dialect::Stream::Transforms
iree::compiler::Dialect::Util::Transforms
iree::compiler::Dialect::VM::Conversion
iree::compiler::Dialect::VM::Conversion::StandardToVM
diff --git a/iree/compiler/Translation/IREEVM.cpp b/iree/compiler/Translation/IREEVM.cpp
index 360a490..ed6b6c4 100644
--- a/iree/compiler/Translation/IREEVM.cpp
+++ b/iree/compiler/Translation/IREEVM.cpp
@@ -10,6 +10,7 @@
#include "iree/compiler/Bindings/TFLite/Transforms/Passes.h"
#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
#include "iree/compiler/Dialect/HAL/Transforms/Passes.h"
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
#include "iree/compiler/Dialect/VM/Target/Bytecode/TranslationFlags.h"
#include "iree/compiler/Dialect/VM/Transforms/Passes.h"
@@ -29,6 +30,16 @@
namespace mlir {
namespace iree_compiler {
+// TODO(#7277): remove flag when on by default.
+static bool getExperimentalStreamsModeFromFlags() {
+ static llvm::cl::opt<bool> *enableFlag = new llvm::cl::opt<bool>{
+ "iree-experimental-streams",
+ llvm::cl::desc("Enables experimental stream dialect pipelines."),
+ llvm::cl::init(false),
+ };
+ return *enableFlag;
+}
+
static BindingOptions getBindingOptionsFromFlags() {
static llvm::cl::OptionCategory bindingOptionsCategory(
"IREE translation binding support options");
@@ -72,59 +83,6 @@
return options;
}
-// Performs initial dialect conversion to get the canonical input lowered into
-// the IREE execution/dataflow dialect.
-//
-// This will fail if we cannot support the input yet. The hope is that any
-// error that happens after this point is either backend-specific (like
-// unsupported SPIR-V lowering) or a bug.
-static LogicalResult convertToFlowModule(ModuleOp moduleOp) {
- PassManager passManager(moduleOp.getContext());
- mlir::applyPassManagerCLOptions(passManager);
- mlir::applyDefaultTimingPassManagerCLOptions(passManager);
- passManager.addInstrumentation(std::make_unique<PassTracing>());
- IREE::Flow::TransformOptions flowOptions;
- IREE::Flow::buildFlowTransformPassPipeline(passManager, flowOptions);
- if (failed(passManager.run(moduleOp))) {
- return moduleOp.emitError()
- << "failed to run flow transformation pass pipeline";
- }
- return success();
-}
-
-// Runs the flow->HAL transform pipeline to lower a flow module and compile
-// executables for the specified target backends.
-static LogicalResult convertToHALModule(
- ModuleOp moduleOp, IREE::HAL::TargetOptions executableOptions) {
- PassManager passManager(moduleOp.getContext());
- mlir::applyPassManagerCLOptions(passManager);
- mlir::applyDefaultTimingPassManagerCLOptions(passManager);
- passManager.addInstrumentation(std::make_unique<PassTracing>());
- IREE::HAL::buildHALTransformPassPipeline(passManager, executableOptions);
- if (failed(passManager.run(moduleOp))) {
- return moduleOp.emitError()
- << "failed to run HAL transformation pass pipeline";
- }
- return success();
-}
-
-// Converts the lowered module to a canonical vm.module containing only vm ops.
-// This uses patterns to convert from standard ops and other dialects to their
-// vm ABI form.
-static LogicalResult convertToVMModule(ModuleOp moduleOp,
- IREE::VM::TargetOptions targetOptions) {
- PassManager passManager(moduleOp.getContext());
- mlir::applyPassManagerCLOptions(passManager);
- mlir::applyDefaultTimingPassManagerCLOptions(passManager);
- passManager.addInstrumentation(std::make_unique<PassTracing>());
- IREE::VM::buildVMTransformPassPipeline(passManager, targetOptions);
- if (failed(passManager.run(moduleOp))) {
- return moduleOp.emitError()
- << "failed to run VM transformation pass pipeline";
- }
- return success();
-}
-
void buildIREEVMTransformPassPipeline(
BindingOptions bindingOptions, InputDialectOptions inputOptions,
IREE::HAL::TargetOptions executableOptions,
@@ -147,11 +105,18 @@
break;
}
- IREE::Flow::TransformOptions flowOptions;
+ bool enableNewStreamsDialect = getExperimentalStreamsModeFromFlags();
buildCommonInputConversionPassPipeline(passManager);
+ IREE::Flow::TransformOptions flowOptions;
+ flowOptions.streamFormation = !enableNewStreamsDialect;
IREE::Flow::buildFlowTransformPassPipeline(passManager, flowOptions);
- IREE::HAL::buildHALTransformPassPipeline(passManager, executableOptions);
+ if (enableNewStreamsDialect) {
+ IREE::Stream::TransformOptions streamOptions;
+ IREE::Stream::buildStreamTransformPassPipeline(passManager, streamOptions);
+ } else {
+ IREE::HAL::buildHALTransformPassPipeline(passManager, executableOptions);
+ }
IREE::VM::buildVMTransformPassPipeline(passManager, targetOptions);
passManager.addPass(IREE::Util::createDropCompilerHintsPass());
}
@@ -241,6 +206,7 @@
#endif // IREE_HAVE_EMITC_DIALECT
void registerIREEVMTranslationFlags() {
+ getExperimentalStreamsModeFromFlags();
getBindingOptionsFromFlags();
getInputDialectOptionsFromFlags();
}
diff --git a/iree/tools/BUILD b/iree/tools/BUILD
index 152abc4..3f6ee17 100644
--- a/iree/tools/BUILD
+++ b/iree/tools/BUILD
@@ -113,6 +113,7 @@
"//iree/compiler/Dialect/Shape/IR",
"//iree/compiler/Dialect/Shape/Transforms",
"//iree/compiler/Dialect/Stream/IR",
+ "//iree/compiler/Dialect/Stream/Transforms",
"//iree/compiler/Dialect/Util/IR",
"//iree/compiler/Dialect/Util/Transforms",
"//iree/compiler/Dialect/VM/Analysis",
diff --git a/iree/tools/CMakeLists.txt b/iree/tools/CMakeLists.txt
index 6050277..bf2c5d1 100644
--- a/iree/tools/CMakeLists.txt
+++ b/iree/tools/CMakeLists.txt
@@ -228,6 +228,7 @@
iree::compiler::Dialect::Shape::IR
iree::compiler::Dialect::Shape::Transforms
iree::compiler::Dialect::Stream::IR
+ iree::compiler::Dialect::Stream::Transforms
iree::compiler::Dialect::Util::IR
iree::compiler::Dialect::Util::Transforms
iree::compiler::Dialect::VM::Analysis
diff --git a/iree/tools/init_iree_passes.h b/iree/tools/init_iree_passes.h
index 7c8daa3..bc5909a 100644
--- a/iree/tools/init_iree_passes.h
+++ b/iree/tools/init_iree_passes.h
@@ -21,6 +21,7 @@
#include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h"
#include "iree/compiler/Dialect/Modules/VMVX/Transforms/Passes.h"
#include "iree/compiler/Dialect/Shape/Transforms/Passes.h"
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
#include "iree/compiler/Dialect/VM/Analysis/TestPasses.h"
#include "iree/compiler/Dialect/VM/Transforms/Passes.h"
@@ -48,6 +49,7 @@
Shape::registerShapePasses();
IREE::Flow::registerFlowPasses();
IREE::HAL::registerHALPasses();
+ IREE::Stream::registerStreamPasses();
IREE::Util::registerTransformPasses();
IREE::VM::registerVMPasses();
IREE::VM::registerVMAnalysisTestPasses();
