Adding `--iree-opt-import-parameters=` and cleaning up export. (#16828)
The new option allows for importing all or some parameters from a
parameter file into the compiler. One or more parameter files to one or
more scopes can be specified, just as with the
`--parameters=[scope=]file.*` flag in the runtime tool. If a scope is
provided only parameters with that scope will be imported. If a
parameter (optionally within a scope) is named explicitly it will be
imported. If a maximum size is specified all parameters <= that size
will be imported.
The export pass has been reworked to be the inverse of the import such
that it's possible to round-trip through export and import. The import
pass doesn't currently support all of the types we can export as
`SerializableAttrInterface` doesn't have a deserialize method (yet) and
we need special handling of packed types (i1, i4, etc). For now warnings
are emitted and the globals are left as parameter references to still
allow partial imports.
Flag renames:
* `--iree-opt-parameter-archive-export-file=` +
`--iree-opt-parameter-archive-export-scope=` ->
`--iree-opt-export-parameters=[scope=]file`
* `--iree-opt-minimum-parameter-export-size=` ->
`--iree-opt-export-parameter-minimum-size=`
* `--iree-opt-splat-parameter-archive-export-file=` ->
`--iree-opt-splat-parameters=`diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilAttrs.cpp b/compiler/src/iree/compiler/Dialect/Util/IR/UtilAttrs.cpp
index 0c88743..5135610 100644
--- a/compiler/src/iree/compiler/Dialect/Util/IR/UtilAttrs.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/IR/UtilAttrs.cpp
@@ -175,25 +175,23 @@
}
}
-// Serializes |count| copies of |splatAttr| to |os|.
-// Significantly faster than the generic ElementsAttr path that needs to perform
-// conversion of the same splat value |count| times.
-static LogicalResult serializeSplatValue(Location loc, Attribute splatAttr,
- int64_t count, llvm::endianness endian,
- llvm::raw_ostream &os) {
+// static
+LogicalResult SerializableAttrInterface::serializeSplatValue(
+ Location loc, Attribute elementAttr, int64_t count, llvm::endianness endian,
+ llvm::raw_ostream &os) {
// Get the encoded byte contents of the splat element.
SmallVector<char> elementBuffer;
- if (auto attr = llvm::dyn_cast<SerializableAttrInterface>(splatAttr)) {
+ if (auto attr = llvm::dyn_cast<SerializableAttrInterface>(elementAttr)) {
if (failed(attr.serializeToVector(loc, endian, elementBuffer))) {
return failure();
}
- } else if (auto attr = llvm::dyn_cast<IntegerAttr>(splatAttr)) {
+ } else if (auto attr = llvm::dyn_cast<IntegerAttr>(elementAttr)) {
if (failed(serializeAPIntRawData(loc, attr.getValue(),
attr.getType().getIntOrFloatBitWidth(),
endian, elementBuffer))) {
return failure();
}
- } else if (auto attr = llvm::dyn_cast<FloatAttr>(splatAttr)) {
+ } else if (auto attr = llvm::dyn_cast<FloatAttr>(elementAttr)) {
if (failed(serializeAPFloatRawData(loc, attr.getValue(),
attr.getType().getIntOrFloatBitWidth(),
endian, elementBuffer))) {
@@ -792,8 +790,9 @@
auto elementsAttr = llvm::cast<DenseElementsAttr>(baseAttr);
if (elementsAttr.isSplat()) {
// Fast-path for splat (no need to convert the value a bunch).
- return serializeSplatValue(loc, elementsAttr.getSplatValue<Attribute>(),
- elementsAttr.getNumElements(), endian, os);
+ return IREE::Util::SerializableAttrInterface::serializeSplatValue(
+ loc, elementsAttr.getSplatValue<Attribute>(),
+ elementsAttr.getNumElements(), endian, os);
}
if (canUseRawData(elementsAttr, endian)) {
diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilInterfaces.td b/compiler/src/iree/compiler/Dialect/Util/IR/UtilInterfaces.td
index aecf7a6..c2155e4 100644
--- a/compiler/src/iree/compiler/Dialect/Util/IR/UtilInterfaces.td
+++ b/compiler/src/iree/compiler/Dialect/Util/IR/UtilInterfaces.td
@@ -1532,6 +1532,13 @@
}]
>,
];
+
+ let extraClassDeclaration = [{
+ // Serializes |count| copies of |elementAttr| to |os|.
+ static LogicalResult serializeSplatValue(
+ Location loc, Attribute elementAttr, int64_t count,
+ llvm::endianness endian, llvm::raw_ostream &os);
+ }];
}
#endif // IREE_DIALECT_UTIL_IR_UTIL_INTERFACES
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
index 94dae26..ed663a4 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
+++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -70,6 +70,19 @@
void buildGlobalOptimizationPassPipeline(
OpPassManager &mainPassManager, const TransformOptions &transformOptions) {
+ // Import parameters before any global optimization passes so that the inlined
+ // parameters are available for folding.
+ if (!transformOptions.options.parameterImportPaths.empty()) {
+ IREE::IO::Parameters::ImportParametersPassOptions importParametersOptions;
+ importParametersOptions.scopePaths =
+ transformOptions.options.parameterImportPaths;
+ importParametersOptions.keys = transformOptions.options.parameterImportKeys;
+ importParametersOptions.maximumSize =
+ transformOptions.options.parameterImportMaximumSize;
+ mainPassManager.addPass(IREE::IO::Parameters::createImportParametersPass(
+ importParametersOptions));
+ }
+
// ML frontends have very uneven support for user-controlled types _and_ users
// tend to use types not well suited for the work they are doing. These
// demotions/promotions allow users to change the types after lowering out of
@@ -232,23 +245,21 @@
// constants that aren't exported and skip it for larger parameters, but this
// is a sensible place for the common case of wanting const-eval in the final
// artifact + archive.
- if (!transformOptions.options.parameterArchiveExportPath.empty()) {
+ if (!transformOptions.options.parameterExportPath.empty()) {
IREE::IO::Parameters::ExportParametersPassOptions exportParametersOptions;
- exportParametersOptions.archivePath =
- transformOptions.options.parameterArchiveExportPath;
- exportParametersOptions.parameterScope =
- transformOptions.options.parameterExportScope;
+ exportParametersOptions.scopePath =
+ transformOptions.options.parameterExportPath;
exportParametersOptions.minimumSize =
- transformOptions.options.minimumParameterExportSize;
+ transformOptions.options.parameterExportMinimumSize;
mainPassManager.addPass(IREE::IO::Parameters::createExportParametersPass(
exportParametersOptions));
}
- if (!transformOptions.options.splatParameterArchiveExportPath.empty()) {
+ if (!transformOptions.options.parameterSplatExportFile.empty()) {
IREE::IO::Parameters::GenerateSplatParameterArchivePassOptions
generateSplatOptions;
- generateSplatOptions.archivePath =
- transformOptions.options.splatParameterArchiveExportPath;
+ generateSplatOptions.filePath =
+ transformOptions.options.parameterSplatExportFile;
mainPassManager.addPass(
IREE::IO::Parameters::createGenerateSplatParameterArchivePass(
generateSplatOptions));
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.td b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
index cf5e85b..4f8dc0e 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/Passes.td
+++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
@@ -11,7 +11,7 @@
def CleanupNumericNarrowing :
Pass<"iree-global-opt-cleanup-numeric-narrowing", ""> {
- let summary = "Cleans up any numeric narrowing ops inserted by iree-global-opt-infer-numeric-narrowing";
+ let summary = "Cleans up any numeric narrowing ops inserted by iree-global-opt-infer-numeric-narrowing.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createCleanupNumericNarrowingPass()";
}
@@ -23,12 +23,12 @@
def DecomposeConcat :
Pass<"iree-global-opt-decompose-concat", ""> {
- let summary = "Decomposes concatenations into a destination and a sequence of slice inserts";
+ let summary = "Decomposes concatenations into a destination and a sequence of slice inserts.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createDecomposeConcatPass()";
let options = [
Option<"enableConcatTransposition", "enable-concat-transposition", "bool",
/*default=*/"false", "Allows transposing concatenations such that "
- "they occur on the inner most dims">,
+ "they occur on the inner most dims.">,
];
}
@@ -39,7 +39,7 @@
def DetachElementwiseFromNamedOps :
Pass<"iree-global-opt-detach-elementwise-from-named-ops", ""> {
- let summary = "Detaches elementwise ops from named Linalg ops";
+ let summary = "Detaches elementwise ops from named Linalg ops.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createDetachElementwiseFromNamedOpsPass()";
}
@@ -57,11 +57,11 @@
def FuseDequantizationMatmul:
InterfacePass<"iree-global-opt-fuse-dequantization-matmul", "mlir::FunctionOpInterface"> {
- let summary = "Fuses dequantization and matmul linalg.generic ops";
+ let summary = "Fuses dequantization and matmul linalg.generic ops.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createFuseDequantizationMatmulPass()";
let options = [
Option<"enableQuantizedMatmulReassociation", "enable-quantized-matmul-reassociation", "bool",
- /*default=*/"false", "Allow reassociation of quantized matmuls (experimental)">,
+ /*default=*/"false", "Allow reassociation of quantized matmuls (experimental).">,
];
}
@@ -73,19 +73,19 @@
def FuseSiluHorizontalMatmul:
InterfacePass<"iree-global-opt-fuse-silu-horizontal-matmul", "mlir::FunctionOpInterface"> {
- let summary = "Fuses matmul ops and silu linalg.generic op";
+ let summary = "Fuses matmul ops and silu linalg.generic op.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createFuseSiluHorizontalMatmulPass()";
}
def GeneralizeLinalgNamedOps :
InterfacePass<"iree-global-opt-generalize-linalg-named-ops", "mlir::FunctionOpInterface"> {
- let summary = "Convert some Linalg named ops into linalg.generics";
+ let summary = "Convert some Linalg named ops into linalg.generics.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createGeneralizeLinalgNamedOpsPass()";
}
def InferNumericNarrowing :
Pass<"iree-global-opt-infer-numeric-narrowing", ""> {
- let summary = "Infers and inserts util.numeric.optional_narrow ops at points that may be beneficial";
+ let summary = "Infers and inserts util.numeric.optional_narrow ops at points that may be beneficial.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createInferNumericNarrowingPass()";
}
@@ -98,34 +98,34 @@
def OptimizeNumerics :
Pass<"iree-global-opt-optimize-numerics", ""> {
- let summary = "Optimizes numerics given annotations added via iree-global-opt-infer-numeric-narrowing";
+ let summary = "Optimizes numerics given annotations added via iree-global-opt-infer-numeric-narrowing.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createOptimizeNumericsPass()";
}
def PropagateLinalgTranspose :
InterfacePass<"iree-global-opt-propagate-linalg-transpose", "mlir::FunctionOpInterface"> {
- let summary = "Propagates linalg.transpose through a restricted set of ops";
+ let summary = "Propagates linalg.transpose through a restricted set of ops.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createPropagateLinalgTransposePass()";
let options = [
Option<"enableAggressivePropagation", "enable-aggressive-propagation", "bool",
- /*default=*/"false", "Enable aggressive propagation to named ops">,
+ /*default=*/"false", "Enable aggressive propagation to named ops.">,
];
}
def RaiseSpecialOps :
Pass<"iree-global-opt-raise-special-ops", ""> {
- let summary = "Raises special ops like softmax to the high level linalg.ext representation";
+ let summary = "Raises special ops like softmax to the high level linalg.ext representation.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createRaiseSpecialOps()";
}
def RemoveZeroExtentTensors :
InterfacePass<"iree-global-opt-remove-zero-extent-tensors", "mlir::FunctionOpInterface"> {
- let summary = "Removes tensors that have 0-extents";
+ let summary = "Removes tensors that have 0-extents.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createRemoveZeroExtentTensorsPass()";
}
def SetEncoding : Pass<"iree-global-opt-set-encoding", ""> {
- let summary = "Introduces tensor encoding for compute operations";
+ let summary = "Introduces tensor encoding for compute operations.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createSetEncodingPass()";
let options = [
Option<"padFactor", "pad-factor", "int64_t", /*default=*/"0",
@@ -137,12 +137,12 @@
}
def SimplifyPackUnpack : Pass<"iree-global-opt-simplify-pack-unpack", ""> {
- let summary = "Simplifies tensor pack and unpack ops";
+ let summary = "Simplifies tensor pack and unpack ops.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createSimplifyPackUnpackPass()";
}
def GlobalLoopInvariantCodeMotion : InterfacePass<"iree-global-opt-loop-invariant-code-motion", "mlir::FunctionOpInterface"> {
- let summary = "Hoist loop invariants out of loops with zero-trip-check";
+ let summary = "Hoist loop invariants out of loops with zero-trip-check.";
let constructor = "mlir::iree_compiler::GlobalOptimization::createGlobalLoopInvariantCodeMotionPass()";
}
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ArchiveUtils.cpp b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ArchiveUtils.cpp
index 16dde66..7922add 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ArchiveUtils.cpp
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ArchiveUtils.cpp
@@ -4,11 +4,8 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-#include "iree/base/api.h"
-#include "iree/io/formats/irpa/irpa_builder.h"
-#include "iree/tooling/parameter_util.h"
-
#include "iree/compiler/Modules/IO/Parameters/Transforms/ArchiveUtils.h"
+
#include "llvm/Support/FileOutputBuffer.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Operation.h"
@@ -19,83 +16,105 @@
StringRef failureMessage) {
if (iree_status_is_ok(status))
return success();
+ iree_host_size_t buffer_length = 0;
+ if (!iree_status_format(status, /*buffer_capacity=*/0,
+ /*buffer=*/nullptr, &buffer_length))
+ return op->emitError() << failureMessage;
std::string message;
- message.resize(512);
- iree_host_size_t buffer_length;
- if (!iree_status_format(status, message.size(), &message[0],
- &buffer_length)) {
- message.resize(buffer_length + 1);
- iree_status_format(status, message.size(), &message[0], &buffer_length);
- }
- message.resize(buffer_length);
+ message.reserve(buffer_length);
+ message.resize(buffer_length - 1);
+ iree_status_format(status, message.capacity(), &message[0], &buffer_length);
iree_status_ignore(status);
- return op->emitError() << failureMessage << message;
+ return op->emitError() << failureMessage << "\n" << message;
}
-LogicalResult
-writeParameterIndex(Operation *op, iree_allocator_t allocator,
- iree_io_parameter_archive_builder_t &builder,
- std::unique_ptr<llvm::FileOutputBuffer> &fileBuffer,
- iree_io_file_handle_t **output_file_handle,
- iree_io_stream_t **output_stream,
- iree_io_parameter_index_t **output_built_index) {
+FailureOr<ArchiveBuilder> createArchiveBuilder(Operation *op) {
+ iree_allocator_t hostAllocator = iree_allocator_system();
+ iree_io_parameter_archive_builder_t *builderPtr = NULL;
+ if (failed(handleRuntimeError(op,
+ iree_allocator_malloc(hostAllocator,
+ sizeof(*builderPtr),
+ (void **)&builderPtr),
+ "allocating archive builder"))) {
+ return failure();
+ }
+ ArchiveBuilder builder(
+ builderPtr, +[](iree_io_parameter_archive_builder_t *builder) {
+ iree_allocator_t host_allocator = builder->host_allocator;
+ iree_io_parameter_archive_builder_deinitialize(builder);
+ iree_allocator_free(host_allocator, builder);
+ });
+ iree_io_parameter_archive_builder_initialize(hostAllocator, builder.get());
+ return std::move(builder);
+}
+
+FailureOr<FileStreamIndex> createParameterIndex(Operation *op,
+ ArchiveBuilder builder,
+ StringRef archivePath) {
+ // Open a file of sufficient size for writing.
+ iree_io_physical_size_t archiveLength =
+ iree_io_parameter_archive_builder_total_size(builder.get());
+ auto fileOr = llvm::FileOutputBuffer::create(archivePath, archiveLength);
+ if (!fileOr) {
+ return op->emitError()
+ << "failed to create file output buffer at " << archivePath
+ << " with error: "
+ << llvm::errorToErrorCode(fileOr.takeError()).message();
+ }
+ std::unique_ptr<llvm::FileOutputBuffer> fileBuffer = std::move(*fileOr);
// Wrap the output file for use with the parameter archive builder.
- iree_io_file_handle_t *target_file_handle = NULL;
- iree_byte_span_t file_contents = iree_make_byte_span(
- fileBuffer->getBufferStart(), fileBuffer->getBufferSize());
// Release callback is a no-op, the mapping is managed by the unique_ptr.
- iree_status_t status = iree_io_file_handle_wrap_host_allocation(
- IREE_IO_FILE_ACCESS_WRITE, file_contents,
- iree_io_file_handle_release_callback_null(), allocator,
- &target_file_handle);
- if (failed(handleRuntimeError(op, status,
- "Failed to open output parameter archive"))) {
- iree_io_file_handle_release(target_file_handle);
- return failure();
- }
-
- // Wrap the target file in a stream.
- iree_io_stream_t *target_stream = NULL;
- status = iree_io_stream_open(IREE_IO_STREAM_MODE_WRITABLE, target_file_handle,
- /*file_offset=*/0, allocator, &target_stream);
+ iree_io_file_handle_t *fileHandlePtr = nullptr;
if (failed(handleRuntimeError(
- op, status, "Failed to create I/O stream to output file"))) {
- iree_io_file_handle_release(target_file_handle);
- iree_io_stream_release(target_stream);
+ op,
+ iree_io_file_handle_wrap_host_allocation(
+ IREE_IO_FILE_ACCESS_WRITE,
+ iree_make_byte_span(fileBuffer->getBufferStart(),
+ fileBuffer->getBufferSize()),
+ iree_io_file_handle_release_callback_null(),
+ builder->host_allocator, &fileHandlePtr),
+ "failed to open output parameter archive"))) {
return failure();
}
+ auto fileHandle = FileHandle(fileHandlePtr, iree_io_file_handle_release);
+
+ // Wrap the target file in a stream. The stream will retain the file handle.
+ iree_io_stream_t *streamPtr = nullptr;
+ if (failed(handleRuntimeError(
+ op,
+ iree_io_stream_open(IREE_IO_STREAM_MODE_WRITABLE, fileHandle.get(),
+ /*file_offset=*/0, builder->host_allocator,
+ &streamPtr),
+ "failed to create I/O stream to output file"))) {
+ return failure();
+ }
+ auto stream = Stream(streamPtr, iree_io_stream_release);
// Allocate an index we'll populate during building to allow us to get the
// storage ranges of non-metadata parameters.
- iree_io_parameter_index_t *built_index = NULL;
- status = iree_io_parameter_index_create(allocator, &built_index);
- if (failed(handleRuntimeError(op, status,
- "Failed to allocate parameter index"))) {
- iree_io_file_handle_release(target_file_handle);
- iree_io_stream_release(target_stream);
- iree_io_parameter_index_release(built_index);
+ iree_io_parameter_index_t *indexPtr = nullptr;
+ if (failed(handleRuntimeError(
+ op,
+ iree_io_parameter_index_create(builder->host_allocator, &indexPtr),
+ "failed to allocate parameter index"))) {
return failure();
}
+ auto index = ParameterIndex(indexPtr, iree_io_parameter_index_release);
// Commit the archive header to the file and produce an index referencing
// it. This will allow us to know where to copy file contents.
- status = iree_io_parameter_archive_builder_write(&builder, target_file_handle,
- /*file_offset=*/0,
- target_stream, built_index);
if (failed(handleRuntimeError(
- op, status,
- "Failed to write parameter index header to output file"))) {
- iree_io_file_handle_release(target_file_handle);
- iree_io_stream_release(target_stream);
- iree_io_parameter_index_release(built_index);
+ op,
+ iree_io_parameter_archive_builder_write(
+ builder.get(), fileHandle.get(),
+ /*file_offset=*/0, stream.get(), index.get()),
+ "failed to write parameter index header to output file"))) {
return failure();
}
- *output_file_handle = target_file_handle;
- *output_stream = target_stream;
- *output_built_index = built_index;
- return success();
+ return std::make_tuple(std::move(fileBuffer), std::move(stream),
+ std::move(index));
}
} // namespace mlir::iree_compiler::IREE::IO::Parameters
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ArchiveUtils.h b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ArchiveUtils.h
index 1345e68..254bc58 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ArchiveUtils.h
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ArchiveUtils.h
@@ -7,16 +7,27 @@
#ifndef IREE_COMPILER_MODULES_IO_PARAMETERS_TRANSFORMS_ARCHIVEUTILS_H_
#define IREE_COMPILER_MODULES_IO_PARAMETERS_TRANSFORMS_ARCHIVEUTILS_H_
-#include "iree/base/api.h"
-#include "iree/io/formats/irpa/irpa_builder.h"
-#include "iree/tooling/parameter_util.h"
-
#include "llvm/Support/FileOutputBuffer.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Operation.h"
+#include "iree/base/api.h"
+#include "iree/io/file_handle.h"
+#include "iree/io/formats/irpa/irpa_builder.h"
+#include "iree/io/parameter_index.h"
+#include "iree/io/stream.h"
+
namespace mlir::iree_compiler::IREE::IO::Parameters {
+using ArchiveBuilder =
+ std::unique_ptr<iree_io_parameter_archive_builder_t,
+ void (*)(iree_io_parameter_archive_builder_t *)>;
+using FileHandle =
+ std::unique_ptr<iree_io_file_handle_t, void (*)(iree_io_file_handle_t *)>;
+using ParameterIndex = std::unique_ptr<iree_io_parameter_index_t,
+ void (*)(iree_io_parameter_index_t *)>;
+using Stream = std::unique_ptr<iree_io_stream_t, void (*)(iree_io_stream_t *)>;
+
// Wrapper around iree_io_stream for use when serializing constants.
class iree_io_stream_ostream : public llvm::raw_ostream {
public:
@@ -35,21 +46,34 @@
iree_io_stream_t *stream = NULL;
};
+using ScopePath = std::pair<StringRef, StringRef>;
+
+// Splits a `scope=path` string into two strings.
+// If no `scope=` was specified the resulting scope string will be empty.
+static inline ScopePath splitScopePath(StringRef scopePath) {
+ size_t i = scopePath.find_first_of('=');
+ if (i == StringRef::npos)
+ return ScopePath("", scopePath);
+ else
+ return ScopePath(scopePath.substr(0, i), scopePath.substr(i + 1));
+}
+
// Helper to interpret iree status messages and print the error message.
LogicalResult handleRuntimeError(Operation *op, iree_status_t status,
StringRef failureMessage);
-// Helper to write the parameter index constructed in the archive |builder|
-// to the given |fileBuffer|. Populates a file, stream, and index handle on
-// success for further writing of the data segments. The file, stream, and
-// index handled must be released by the caller if this succeeds.
-LogicalResult
-writeParameterIndex(Operation *op, iree_allocator_t allocator,
- iree_io_parameter_archive_builder_t &builder,
- std::unique_ptr<llvm::FileOutputBuffer> &fileBuffer,
- iree_io_file_handle_t **output_file_handle,
- iree_io_stream_t **output_stream,
- iree_io_parameter_index_t **output_built_index);
+// Creates an empty archive builder.
+FailureOr<ArchiveBuilder> createArchiveBuilder(Operation *op);
+
+using FileStreamIndex =
+ std::tuple<std::unique_ptr<llvm::FileOutputBuffer>, Stream, ParameterIndex>;
+
+// Creates a parameter archive from |builder| and returns a file buffer and
+// stream opened for writing the parameters at the offsets specified in the
+// returned index.
+FailureOr<FileStreamIndex> createParameterIndex(Operation *op,
+ ArchiveBuilder builder,
+ StringRef archivePath);
} // namespace mlir::iree_compiler::IREE::IO::Parameters
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/BUILD.bazel
index caef38d..4fb4155 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/BUILD.bazel
@@ -18,6 +18,7 @@
"ArchiveUtils.cpp",
"ExportParameters.cpp",
"GenerateSplatParameterArchive.cpp",
+ "ImportParameters.cpp",
"Passes.cpp",
],
hdrs = [
@@ -31,10 +32,11 @@
"//compiler/src/iree/compiler/Dialect/Util/IR",
"//runtime/src/iree/base",
"//runtime/src/iree/hal",
+ "//runtime/src/iree/io:file_handle",
"//runtime/src/iree/io:parameter_index",
- "//runtime/src/iree/io:scope_map",
+ "//runtime/src/iree/io:stream",
+ "//runtime/src/iree/io/formats:parser_registry",
"//runtime/src/iree/io/formats/irpa",
- "//runtime/src/iree/tooling:parameter_util",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:DialectUtils",
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/CMakeLists.txt
index 0a3ac29..7c27305 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/CMakeLists.txt
@@ -21,6 +21,7 @@
"ArchiveUtils.cpp"
"ExportParameters.cpp"
"GenerateSplatParameterArchive.cpp"
+ "ImportParameters.cpp"
"Passes.cpp"
DEPS
::PassesIncGen
@@ -35,10 +36,11 @@
iree::compiler::Dialect::Stream::IR
iree::compiler::Dialect::Util::IR
iree::hal
+ iree::io::file_handle
iree::io::formats::irpa
+ iree::io::formats::parser_registry
iree::io::parameter_index
- iree::io::scope_map
- iree::tooling::parameter_util
+ iree::io::stream
PUBLIC
)
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ExportParameters.cpp b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ExportParameters.cpp
index 4f7eeaa..4e67fb6 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ExportParameters.cpp
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ExportParameters.cpp
@@ -4,14 +4,10 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-#include "iree/base/api.h"
-#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
-#include "iree/io/formats/irpa/irpa_builder.h"
-#include "iree/tooling/parameter_util.h"
-
#include "iree/compiler/Dialect/Stream/IR/StreamTypes.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/Modules/IO/Parameters/Transforms/ArchiveUtils.h"
#include "iree/compiler/Modules/IO/Parameters/Transforms/Passes.h"
#include "llvm/ADT/ScopeExit.h"
@@ -30,53 +26,59 @@
namespace {
-// Hoists all serializable constants with storage size at least |minimumSize|
-// into their own globals with initial value equal to the constant value.
-static void hoistConstantsIntoGlobals(mlir::ModuleOp moduleOp,
- int64_t minimumSize) {
- SymbolTable moduleSymbols(moduleOp);
- IRRewriter rewriter(OpBuilder::atBlockBegin(moduleOp.getBody()));
- llvm::DenseMap<arith::ConstantOp, Util::GlobalOp> hoistedMap;
- moduleOp.walk([&](arith::ConstantOp constant) {
- // Constants part of a different logical program should not be hoisted.
- if (SymbolTable::getNearestSymbolTable(constant) != moduleOp) {
- return;
- }
- TypedAttr initialValueAttr = constant.getValue();
- auto serializableAttr =
- dyn_cast<IREE::Util::SerializableAttrInterface>(initialValueAttr);
- if (!serializableAttr) {
- return;
- }
-
- // Check that the serialized size of the attribute is at least as big as
- // the pass configured minimum storage size.
- iree_io_physical_size_t storageSize = serializableAttr.getStorageSize();
- if (storageSize < minimumSize) {
- return;
- }
-
- // Create a new global with initial value equal to the constant.
- Location loc = constant.getLoc();
- Util::GlobalOp globalOp = rewriter.create<Util::GlobalOp>(
- loc, "constant_hoisted", false, constant.getType());
- moduleSymbols.insert(globalOp);
- // Attributes are stored uniqued by their contents so this is not a copy.
- globalOp.setInitialValueAttr(initialValueAttr);
- SymbolTable::setSymbolVisibility(globalOp,
- SymbolTable::Visibility::Private);
- hoistedMap[constant] = globalOp;
- });
-
- // Replace all constants with their associated hoisted globals.
- for (auto it : hoistedMap) {
- arith::ConstantOp originalConstant = it.first;
- Util::GlobalOp globalOp = it.second;
- rewriter.setInsertionPointAfterValue(originalConstant);
- Value load = globalOp.createLoadOp(globalOp.getLoc(), rewriter)
- .getLoadedGlobalValue();
- rewriter.replaceOp(originalConstant, load);
+static LogicalResult
+addSplatEntry(IREE::Util::GlobalOpInterface globalOp,
+ SplatElementsAttr valueAttr, int64_t storageSize,
+ iree_io_parameter_archive_builder_t *builder) {
+ SmallVector<char, IREE_IO_PARAMETER_MAX_SPLAT_PATTERN_LENGTH> pattern;
+ llvm::raw_svector_ostream os(pattern);
+ if (failed(IREE::Util::SerializableAttrInterface::serializeSplatValue(
+ globalOp.getLoc(), valueAttr.getSplatValue<Attribute>(),
+ /*count=*/1, llvm::endianness::little, os))) {
+ return failure();
}
+
+ StringRef name = globalOp.getGlobalName();
+ return handleRuntimeError(
+ globalOp,
+ iree_io_parameter_archive_builder_add_splat_entry(
+ builder, iree_make_string_view(name.data(), name.size()),
+ /*metadata=*/iree_const_byte_span_empty(), pattern.data(),
+ static_cast<uint8_t>(pattern.size()), storageSize),
+ "failed to add splat entry for global");
+}
+
+static LogicalResult
+addDataEntry(IREE::Util::GlobalOpInterface globalOp,
+ IREE::Util::SerializableAttrInterface valueAttr,
+ int64_t storageSize,
+ iree_io_parameter_archive_builder_t *builder) {
+ StringRef name = globalOp.getGlobalName();
+ return handleRuntimeError(
+ globalOp,
+ iree_io_parameter_archive_builder_add_data_entry(
+ builder, iree_make_string_view(name.data(), name.size()),
+ /*metadata=*/iree_const_byte_span_empty(),
+ /*alignment=*/
+ IREE_IO_PARAMETER_ARCHIVE_DEFAULT_DATA_ALIGNMENT, storageSize),
+ "failed to add data entry for global");
+}
+
+// Adds an entry to the parameter archive builder for the given global.
+// If the global is mutable we allocate archive storage for the full
+// serialized parameter. This allows the parameter to be mapped for
+// read/write in the file. If the global is immutable and a splat we can
+// add a splat entry instead to save on archive size and startup time.
+static LogicalResult addEntry(IREE::Util::GlobalOpInterface globalOp,
+ IREE::Util::SerializableAttrInterface valueAttr,
+ iree_io_parameter_archive_builder_t *builder) {
+ if (!globalOp.isGlobalMutable()) {
+ if (auto elementsAttr = dyn_cast<SplatElementsAttr>(valueAttr)) {
+ return addSplatEntry(globalOp, elementsAttr, valueAttr.getStorageSize(),
+ builder);
+ }
+ }
+ return addDataEntry(globalOp, valueAttr, valueAttr.getStorageSize(), builder);
}
struct ExportParametersPass
@@ -86,159 +88,101 @@
ExportParametersPass>::ExportParametersPassBase;
void runOnOperation() override {
- // Nothing to do if no path specified.
- if (archivePath.empty()) {
- return;
- }
-
MLIRContext *context = &getContext();
+
+ // Nothing to do if no path specified.
+ if (scopePath.empty())
+ return;
+ auto [scope, path] = splitScopePath(scopePath);
+
+ // Create a builder used to accumulate the parameters.
ModuleOp moduleOp = getOperation();
+ auto builder = createArchiveBuilder(moduleOp);
+ if (failed(builder))
+ return signalPassFailure();
- // First hoist all inline constants into their own globals.
- hoistConstantsIntoGlobals(moduleOp, minimumSize);
-
- iree_allocator_t host_allocator = iree_allocator_system();
-
- // Create the parameter archive builder.
- iree_io_parameter_archive_builder_t builder;
- iree_io_parameter_archive_builder_initialize(host_allocator, &builder);
-
- auto deinitializeExit = llvm::make_scope_exit([&]() {
- return iree_io_parameter_archive_builder_deinitialize(&builder);
- });
-
- SmallVector<IREE::Util::GlobalOp> constantGlobals;
- // Walk the globals in the module.
- for (auto global : moduleOp.getOps<IREE::Util::GlobalOp>()) {
- // TODO: Support exporting mutable globals.
- if (global.getIsMutable()) {
- continue;
- }
- // Only globals initialized with initial values can be parameterized.
- auto initialValueAttr = global.getInitialValueAttr();
- if (!initialValueAttr) {
- continue;
- }
-
- // The attribute must be serializable to be turned into a parameter.
+ // Accumulate globals that match the pass options and add them to the index.
+ SmallVector<IREE::Util::GlobalOpInterface> constantGlobalOps;
+ for (auto globalOp : moduleOp.getOps<IREE::Util::GlobalOpInterface>()) {
+ // Only globals initialized with serializable initial values can be
+ // parameterized.
auto serializableAttr =
- dyn_cast<IREE::Util::SerializableAttrInterface>(initialValueAttr);
- if (!serializableAttr) {
+ dyn_cast_if_present<IREE::Util::SerializableAttrInterface>(
+ globalOp.getGlobalInitialValue());
+ if (!serializableAttr)
continue;
- }
// Check that the serialized size of the attribute is at least as big as
// the pass configured minimum storage size.
- iree_io_physical_size_t storageSize = serializableAttr.getStorageSize();
- if (storageSize < minimumSize) {
+ int64_t storageSize = serializableAttr.getStorageSize();
+ if (storageSize < minimumSize)
continue;
- }
- StringRef name = global.getSymName();
- // Add a data entry to the builder for this global.
- iree_status_t status = iree_io_parameter_archive_builder_add_data_entry(
- &builder,
- iree_string_view_t{name.data(),
- static_cast<iree_host_size_t>(name.size())},
- /*metadata=*/iree_const_byte_span_empty(),
- /*alignment=*/IREE_IO_PARAMETER_ARCHIVE_DEFAULT_DATA_ALIGNMENT,
- storageSize);
- if (failed(handleRuntimeError(moduleOp, status,
- "Failed to add data entry for global"))) {
+ // Add the entry with a type based on its contents.
+ if (failed(addEntry(globalOp, serializableAttr, builder->get())))
return signalPassFailure();
- }
- constantGlobals.push_back(global);
+ constantGlobalOps.push_back(globalOp);
}
- // Early exit if no parameterizable globals present.
- if (constantGlobals.empty()) {
+ // Early exit if no parameterizable globals are present.
+ if (constantGlobalOps.empty())
return;
- }
- // Open a file of sufficient size (now that we know it) for writing.
- iree_io_physical_size_t archive_length =
- iree_io_parameter_archive_builder_total_size(&builder);
-
- auto FileOrErr =
- llvm::FileOutputBuffer::create(archivePath, archive_length);
- if (!FileOrErr) {
- moduleOp.emitError()
- << "Failed to create file output buffer at " << archivePath
- << " with error: "
- << llvm::errorToErrorCode(FileOrErr.takeError()).message();
+ // Create the parameter archive file opened for writing.
+ auto fileStreamIndexOr =
+ createParameterIndex(moduleOp, std::move(builder.value()), path);
+ if (failed(fileStreamIndexOr))
return signalPassFailure();
- }
- std::unique_ptr<llvm::FileOutputBuffer> fileBuffer = std::move(*FileOrErr);
+ auto [file, stream, index] = *std::move(fileStreamIndexOr);
- iree_io_file_handle_t *target_file_handle = NULL;
- iree_io_stream_t *target_stream = NULL;
- iree_io_parameter_index_t *built_index = NULL;
- if (failed(writeParameterIndex(moduleOp, host_allocator, builder,
- fileBuffer, &target_file_handle,
- &target_stream, &built_index))) {
- return signalPassFailure();
- }
-
- auto releaseFileExit = llvm::make_scope_exit([&]() -> void {
- iree_io_stream_release(target_stream);
- iree_io_parameter_index_release(built_index);
- iree_io_file_handle_release(target_file_handle);
- });
-
- StringAttr scopeAttr = parameterScope.empty()
- ? StringAttr()
- : StringAttr::get(context, parameterScope);
- iree_io_stream_ostream llvm_stream(target_stream);
-
- // Write all of the global contents to the appropriate data storage
- // segments.
- for (auto constantGlobal : constantGlobals) {
- StringRef name = constantGlobal.getSymName();
-
- const iree_io_parameter_index_entry_t *target_entry = NULL;
- iree_status_t status = iree_io_parameter_index_lookup(
- built_index,
- iree_string_view_t{name.data(),
- static_cast<iree_host_size_t>(name.size())},
- &target_entry);
+ // Serialize parameters to the file.
+ for (auto globalOp : constantGlobalOps) {
+ // Lookup the entry in the index corresponding to the global.
+ const iree_io_parameter_index_entry_t *entry = nullptr;
+ StringRef name = globalOp.getGlobalName();
if (failed(handleRuntimeError(
- moduleOp, status,
- "Failed to write parameter index header to output file"))) {
- return signalPassFailure();
- }
- status = iree_io_stream_seek(target_stream, IREE_IO_STREAM_SEEK_SET,
- target_entry->storage.file.offset);
- if (failed(handleRuntimeError(
- moduleOp, status,
- "Failed to seek to location of global in index"))) {
+ globalOp,
+ iree_io_parameter_index_lookup(
+ index.get(), iree_make_string_view(name.data(), name.size()),
+ &entry),
+ "retrieve global from index"))) {
return signalPassFailure();
}
- auto initialValueAttr = constantGlobal.getInitialValueAttr();
- auto serializableAttr =
- dyn_cast<IREE::Util::SerializableAttrInterface>(initialValueAttr);
+ // Only file entries get stored; splats are in the metadata table.
+ if (entry->type == IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE) {
+ // Seek to where the serialized global begins in the file.
+ if (failed(handleRuntimeError(
+ globalOp,
+ iree_io_stream_seek(stream.get(), IREE_IO_STREAM_SEEK_SET,
+ entry->storage.file.offset),
+ "failed to seek to location of global in archive"))) {
+ return signalPassFailure();
+ }
- if (failed(serializableAttr.serializeToStream(constantGlobal.getLoc(),
- llvm::endianness::native,
- llvm_stream))) {
- moduleOp.emitError() << "Failed to serialize global " << constantGlobal;
- return signalPassFailure();
+ // Serialize the global contents to the stream.
+ iree_io_stream_ostream os(stream.get());
+ auto serializableAttr = cast<IREE::Util::SerializableAttrInterface>(
+ globalOp.getGlobalInitialValue());
+ if (failed(serializableAttr.serializeToStream(
+ globalOp.getLoc(), llvm::endianness::native, os))) {
+ globalOp.emitError() << "failed to serialize global to archive";
+ return signalPassFailure();
+ }
+ os.flush();
}
- llvm_stream.flush();
- // Now we can just replace the existing initial value with a reference to
- // the parameter.
- auto param = IREE::Stream::NamedParameterAttr::get(
- context, constantGlobal.getType(), scopeAttr,
- StringAttr::get(context, name), DictionaryAttr());
- constantGlobal.setInitialValueAttr(param);
+ // Change the global to reference the parameter.
+ globalOp.setGlobalInitialValue(IREE::Stream::NamedParameterAttr::get(
+ context, globalOp.getGlobalType(), StringAttr::get(context, scope),
+ StringAttr::get(context, name), DictionaryAttr()));
}
+
// Commit the written file.
- llvm::Error maybeCommit = fileBuffer->commit();
- if (maybeCommit) {
+ if (llvm::Error maybeCommit = file->commit()) {
InFlightDiagnostic errorStream =
- moduleOp.emitError() << "Failed to commit archive with error: ";
+ moduleOp.emitError() << "failed to commit archive with error: ";
llvm::handleAllErrors(std::move(maybeCommit),
[&](const llvm::ErrorInfoBase &PE) {
errorStream << PE.message() << "\n";
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/GenerateSplatParameterArchive.cpp b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/GenerateSplatParameterArchive.cpp
index d270372..4a6dfc2 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/GenerateSplatParameterArchive.cpp
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/GenerateSplatParameterArchive.cpp
@@ -4,14 +4,10 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-#include "iree/base/api.h"
-#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
-#include "iree/io/formats/irpa/irpa_builder.h"
-#include "iree/tooling/parameter_util.h"
-
#include "iree/compiler/Dialect/Stream/IR/StreamTypes.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/Modules/IO/Parameters/Transforms/ArchiveUtils.h"
#include "iree/compiler/Modules/IO/Parameters/Transforms/Passes.h"
#include "llvm/ADT/ScopeExit.h"
@@ -28,6 +24,15 @@
namespace {
+static Attribute getDefaultSplatAttr(Type elementType) {
+ // Today we only support basic types where 0 bits represent zeros - that lets
+ // us just splat out the right number of bits.
+ int64_t storageSize =
+ IREE::Util::getRoundedPhysicalStorageSize(1, elementType);
+ return IntegerAttr::get(
+ IntegerType::get(elementType.getContext(), storageSize * 8), 0);
+}
+
struct GenerateSplatParameterArchivePass
: public IREE::IO::Parameters::impl::GenerateSplatParameterArchivePassBase<
GenerateSplatParameterArchivePass> {
@@ -36,96 +41,72 @@
void runOnOperation() override {
// Nothing to do if no path specified.
- if (archivePath.empty()) {
+ if (filePath.empty())
return;
- }
+ // Create a builder used to accumulate the parameters.
ModuleOp moduleOp = getOperation();
+ auto builder = createArchiveBuilder(moduleOp);
+ if (failed(builder))
+ return signalPassFailure();
- iree_allocator_t host_allocator = iree_allocator_system();
-
- // Create the parameter archive builder.
- iree_io_parameter_archive_builder_t builder;
- iree_io_parameter_archive_builder_initialize(host_allocator, &builder);
-
- auto deinitializeExit = llvm::make_scope_exit([&]() {
- return iree_io_parameter_archive_builder_deinitialize(&builder);
- });
-
- bool hasParameter = false;
// Walk the globals in the module.
- for (auto global : moduleOp.getOps<IREE::Util::GlobalOp>()) {
+ for (auto globalOp : moduleOp.getOps<IREE::Util::GlobalOpInterface>()) {
+ // Only support types we can meaningfully generate splats for.
+ auto shapedType = dyn_cast<ShapedType>(globalOp.getGlobalType());
+ if (!shapedType)
+ continue;
+
// Look for globals backed by parameters.
- auto initialValueAttr = global.getInitialValueAttr();
- if (!initialValueAttr) {
- continue;
- }
auto parameterAttr =
- dyn_cast<IREE::Stream::NamedParameterAttr>(initialValueAttr);
- if (!parameterAttr) {
+ dyn_cast_if_present<IREE::Stream::NamedParameterAttr>(
+ globalOp.getGlobalInitialValue());
+ if (!parameterAttr)
continue;
- }
- // Note that the scope is not a part of the parameter archive. If the
- // module includes multiple scopes, multiple copies of the splat archive
- // would need to be passed in with all possible scopes.
- std::string parameterName = parameterAttr.getKey().str();
- iree_io_physical_size_t storageSize = parameterAttr.getStorageSize();
+ // TODO: support other patterns/generators.
+ auto elementAttr = getDefaultSplatAttr(shapedType.getElementType());
- // Add a zero-splat entry to the builder for this global.
- char c0 = 0;
- iree_status_t status = iree_io_parameter_archive_builder_add_splat_entry(
- &builder,
- iree_string_view_t{
- parameterName.data(),
- static_cast<iree_host_size_t>(parameterName.size())},
- /*metadata=*/iree_const_byte_span_empty(),
- /*pattern=*/&c0, /*pattern_length=*/1, /*data_length=*/storageSize);
- if (failed(handleRuntimeError(moduleOp, status,
- "Failed to add splate entry for global"))) {
+ // Serialize the splat pattern.
+ SmallVector<char, IREE_IO_PARAMETER_MAX_SPLAT_PATTERN_LENGTH> pattern;
+ llvm::raw_svector_ostream os(pattern);
+ if (failed(IREE::Util::SerializableAttrInterface::serializeSplatValue(
+ globalOp.getLoc(), elementAttr,
+ /*count=*/1, llvm::endianness::little, os))) {
return signalPassFailure();
}
- hasParameter = true;
+
+ // Add a zero-splat entry to the builder for this global.
+ auto parameterName = parameterAttr.getKey().getValue();
+ if (failed(handleRuntimeError(
+ moduleOp,
+ iree_io_parameter_archive_builder_add_splat_entry(
+ builder->get(),
+ iree_string_view_t{
+ parameterName.data(),
+ static_cast<iree_host_size_t>(parameterName.size())},
+ /*metadata=*/iree_const_byte_span_empty(),
+ /*pattern=*/pattern.data(),
+ /*pattern_length=*/static_cast<uint8_t>(pattern.size()),
+ /*data_length=*/parameterAttr.getStorageSize()),
+ "failed to add splat entry for global"))) {
+ return signalPassFailure();
+ }
}
// Early exit if no parameter backed globals present.
- if (!hasParameter) {
+ if (iree_io_parameter_archive_builder_is_empty(builder->get()))
return;
- }
- // Open a file of sufficient size (now that we know it) for writing.
- iree_io_physical_size_t archive_length =
- iree_io_parameter_archive_builder_total_size(&builder);
-
- auto FileOrErr =
- llvm::FileOutputBuffer::create(archivePath, archive_length);
- if (!FileOrErr) {
- moduleOp.emitError()
- << "Failed to create file output buffer at " << archivePath
- << " with error: "
- << llvm::errorToErrorCode(FileOrErr.takeError()).message();
+ // Create the parameter archive file.
+ auto fileStreamIndexOr =
+ createParameterIndex(moduleOp, std::move(builder.value()), filePath);
+ if (failed(fileStreamIndexOr))
return signalPassFailure();
- }
- std::unique_ptr<llvm::FileOutputBuffer> fileBuffer = std::move(*FileOrErr);
-
- iree_io_file_handle_t *target_file_handle = NULL;
- iree_io_stream_t *target_stream = NULL;
- iree_io_parameter_index_t *built_index = NULL;
- if (failed(writeParameterIndex(moduleOp, host_allocator, builder,
- fileBuffer, &target_file_handle,
- &target_stream, &built_index))) {
- return signalPassFailure();
- }
-
- auto releaseFileExit = llvm::make_scope_exit([&]() -> void {
- iree_io_stream_release(target_stream);
- iree_io_parameter_index_release(built_index);
- iree_io_file_handle_release(target_file_handle);
- });
+ auto [file, stream, index] = *std::move(fileStreamIndexOr);
// Commit the written file.
- llvm::Error maybeCommit = fileBuffer->commit();
- if (maybeCommit) {
+ if (llvm::Error maybeCommit = file->commit()) {
InFlightDiagnostic errorStream =
moduleOp.emitError() << "Failed to commit archive with error: ";
llvm::handleAllErrors(std::move(maybeCommit),
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ImportParameters.cpp b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ImportParameters.cpp
new file mode 100644
index 0000000..8a3888e
--- /dev/null
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/ImportParameters.cpp
@@ -0,0 +1,364 @@
+// Copyright 2024 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/IR/StreamTypes.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/Modules/IO/Parameters/Transforms/ArchiveUtils.h"
+#include "iree/compiler/Modules/IO/Parameters/Transforms/Passes.h"
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/Support/EndianStream.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/FileUtilities.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/DialectResourceBlobManager.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/FileUtilities.h"
+
+#include "iree/io/formats/parser_registry.h"
+
+namespace mlir::iree_compiler::IREE::IO::Parameters {
+
+#define GEN_PASS_DEF_IMPORTPARAMETERSPASS
+#include "iree/compiler/Modules/IO/Parameters/Transforms/Passes.h.inc"
+
+namespace {
+
+static FailureOr<FileHandle> openArchiveFile(ModuleOp moduleOp,
+ StringRef archivePath) {
+ iree_allocator_t hostAllocator = iree_allocator_system();
+
+ // Open the archive (hopefully mapped).
+ auto fileOrErr = llvm::MemoryBuffer::getFile(
+ archivePath, /*IsText=*/false, /*RequiresNullTerminator=*/false,
+ /*IsVolatile=*/false, /*Alignment=*/std::nullopt);
+ if (std::error_code error = fileOrErr.getError()) {
+ llvm::errs() << "cannot open archive input file '" + archivePath +
+ "': " + error.message();
+ return failure();
+ }
+ auto file = std::move(fileOrErr.get());
+
+ // A callback issued when a file is released to destroy the file.
+ iree_io_file_handle_release_callback_t fileReleaseCallback;
+ fileReleaseCallback.fn =
+ +[](void *user_data, iree_io_file_handle_primitive_t handle_primitive) {
+ delete reinterpret_cast<llvm::MemoryBuffer *>(user_data);
+ };
+ fileReleaseCallback.user_data = file.get();
+
+ // Wrap the archive in a file handle.
+ iree_io_file_handle_t *fileHandle = nullptr;
+ if (failed(handleRuntimeError(
+ moduleOp,
+ iree_io_file_handle_wrap_host_allocation(
+ IREE_IO_FILE_ACCESS_READ,
+ iree_make_byte_span(const_cast<char *>(file->getBufferStart()),
+ file->getBufferSize()),
+ fileReleaseCallback, hostAllocator, &fileHandle),
+ "unable to wrap archive memory buffer"))) {
+ return failure();
+ }
+ file.release(); // now owned by the fileHandle
+
+ return FileHandle(fileHandle, iree_io_file_handle_release);
+}
+
+static LogicalResult
+loadParameterIndex(ModuleOp moduleOp, StringRef path,
+ iree_io_parameter_index_t *parameterIndex) {
+ // Open the archive file (hopefully mapping it).
+ auto fileHandle = openArchiveFile(moduleOp, path);
+ if (failed(fileHandle))
+ return failure();
+
+ // Parse the archive as a particular format.
+ return handleRuntimeError(
+ moduleOp,
+ iree_io_parse_file_index(iree_make_string_view(path.data(), path.size()),
+ fileHandle->get(), parameterIndex),
+ "parsing parameter archive");
+}
+
+class ParameterIndices {
+public:
+ bool contains(StringRef scope) const {
+ return indicesByScope.contains(scope);
+ }
+
+ iree_io_parameter_index_t *lookup(StringRef scope) const {
+ auto it = indicesByScope.find(scope);
+ return it == indicesByScope.end() ? nullptr : it->second;
+ }
+
+ iree_io_parameter_index_t *lookupOrCreate(ModuleOp moduleOp,
+ StringRef scope) {
+ iree_allocator_t hostAllocator = iree_allocator_system();
+ if (iree_io_parameter_index_t *existing = lookup(scope))
+ return existing;
+ iree_io_parameter_index_t *parameterIndexPtr = nullptr;
+ if (failed(handleRuntimeError(
+ moduleOp,
+ iree_io_parameter_index_create(hostAllocator, ¶meterIndexPtr),
+ "unable to allocate empty parameter index"))) {
+ return nullptr;
+ }
+ auto parameterIndex =
+ ParameterIndex(parameterIndexPtr, iree_io_parameter_index_release);
+ auto *ptr = parameterIndex.get();
+ indices.push_back(std::move(parameterIndex));
+ indicesByScope.try_emplace(scope, ptr);
+ return ptr;
+ }
+
+private:
+ SmallVector<ParameterIndex> indices;
+ DenseMap<StringRef, iree_io_parameter_index_t *> indicesByScope;
+};
+
+// Allocates one parameter index per scope (possibly an empty string) and
+// merges in parameters from the referenced files.
+static FailureOr<ParameterIndices>
+loadParameterArchives(ModuleOp moduleOp, ArrayRef<std::string> scopePaths) {
+ ParameterIndices parameterIndices;
+ for (auto &scopePath : scopePaths) {
+ auto [scope, path] = splitScopePath(scopePath);
+ auto *parameterIndex = parameterIndices.lookupOrCreate(moduleOp, scope);
+ if (failed(loadParameterIndex(moduleOp, path, parameterIndex)))
+ return failure();
+ }
+ return parameterIndices;
+}
+
+// Today only shaped types of elements where we know we can directly access the
+// data as stored in the file.
+static bool isTypeSupported(Type type) {
+ auto shapedType = dyn_cast<ShapedType>(type);
+ if (!shapedType)
+ return false;
+ auto elementType = shapedType.getElementType();
+ // NOTE: packed types not yet supported.
+ if (!elementType.isIntOrFloat())
+ return false;
+ const unsigned logicalBitWidth = elementType.getIntOrFloatBitWidth();
+ switch (logicalBitWidth) {
+ case 8:
+ case 16:
+ case 32:
+ case 64:
+ return true;
+ default:
+ return false;
+ }
+}
+
+static FailureOr<TypedAttr>
+importParameterFromSplat(StringRef fullName, ShapedType globalType,
+ const iree_io_parameter_index_entry_t *entry) {
+ // Ensure we have the right bit count.
+ // NOTE: this will need to change for packed types.
+ auto elementType = globalType.getElementType();
+ unsigned bitWidth = elementType.getIntOrFloatBitWidth();
+ if (bitWidth != entry->storage.splat.pattern_length * 8) {
+ llvm::errs() << "splat pattern has insufficient bits for type "
+ << globalType << "\n";
+ return failure();
+ }
+
+ // Map the splat pattern into an attribute.
+ Attribute valueAttr;
+ if (auto integerType = dyn_cast<IntegerType>(elementType)) {
+ uint64_t value = 0;
+ switch (integerType.getWidth()) {
+ case 8:
+ value = entry->storage.splat.pattern[0];
+ break;
+ case 16:
+ value = llvm::support::endian::read16le(entry->storage.splat.pattern);
+ break;
+ case 32:
+ value = llvm::support::endian::read32le(entry->storage.splat.pattern);
+ break;
+ case 64:
+ value = llvm::support::endian::read64le(entry->storage.splat.pattern);
+ break;
+ default:
+ assert(false && "integer width not supported");
+ return failure();
+ }
+ valueAttr =
+ IntegerAttr::get(elementType, APInt(integerType.getWidth(), value));
+ } else if (auto floatType = dyn_cast<FloatType>(elementType)) {
+ uint64_t value = 0;
+ switch (floatType.getWidth()) {
+ case 8:
+ value = entry->storage.splat.pattern[0];
+ break;
+ case 16:
+ value = llvm::support::endian::read16le(entry->storage.splat.pattern);
+ break;
+ case 32:
+ value = llvm::support::endian::read32le(entry->storage.splat.pattern);
+ break;
+ case 64:
+ value = llvm::support::endian::read64le(entry->storage.splat.pattern);
+ break;
+ default:
+ assert(false && "integer width not supported");
+ return failure();
+ }
+ valueAttr = FloatAttr::get(elementType,
+ APFloat(floatType.getFloatSemantics(),
+ APInt(floatType.getWidth(), value)));
+ }
+ if (!valueAttr) {
+ llvm::errs() << "unsupported splat type: " << elementType << "\n";
+ return failure();
+ }
+
+ // Create a splat with the given element value.
+ return TypedAttr(SplatElementsAttr::get(globalType, valueAttr));
+}
+
+// TODO(benvanik): replace with resources, maybe? there's no FileAsmResourceBlob
+// yet but we could use that to point back to the file on disk. For now we just
+// import as a raw attr to ensure that imported parameters behave exactly as
+// constants would everywhere and can be serialized/deserialized across
+// reproducers/etc.
+static FailureOr<TypedAttr>
+importParameterFromFile(StringRef fullName, ShapedType globalType,
+ const iree_io_parameter_index_entry_t *entry) {
+ // We currently only support mapped files, but could instead handle file path
+ // references and point resource blobs directly at them.
+ iree_io_file_handle_primitive_t filePrimitive =
+ iree_io_file_handle_primitive(entry->storage.file.handle);
+ if (filePrimitive.type != IREE_IO_FILE_HANDLE_TYPE_HOST_ALLOCATION) {
+ llvm::errs() << "only host allocation file primitives are supported\n";
+ return failure();
+ }
+ const uint8_t *fileData = filePrimitive.value.host_allocation.data;
+
+ // Copy the data from the parameter file into an attribute
+ return TypedAttr(DenseElementsAttr::getFromRawBuffer(
+ globalType, ArrayRef<char>(reinterpret_cast<const char *>(
+ fileData + entry->storage.file.offset),
+ entry->length)));
+}
+
+// Import the given |parameterAttr| from |entry|.
+static FailureOr<TypedAttr>
+importParameter(StringRef fullName, ShapedType globalType,
+ IREE::Stream::NamedParameterAttr parameterAttr,
+ const iree_io_parameter_index_entry_t *entry) {
+ switch (entry->type) {
+ case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT:
+ return importParameterFromSplat(fullName, globalType, entry);
+ case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE:
+ return importParameterFromFile(fullName, globalType, entry);
+ default:
+ // Unsupported type.
+ llvm::errs() << "found parameter but type is not supported: "
+ << parameterAttr.getKey().getValue() << "\n";
+ return failure();
+ }
+}
+
+struct ImportParametersPass
+ : public IREE::IO::Parameters::impl::ImportParametersPassBase<
+ ImportParametersPass> {
+ using IREE::IO::Parameters::impl::ImportParametersPassBase<
+ ImportParametersPass>::ImportParametersPassBase;
+
+ void runOnOperation() override {
+ // Nothing to do if no path specified.
+ if (scopePaths.empty())
+ return;
+
+ // Open the archive file (hopefully mapping it) and parse the index.
+ ModuleOp moduleOp = getOperation();
+ auto parameterIndices = loadParameterArchives(moduleOp, scopePaths);
+ if (failed(parameterIndices))
+ return signalPassFailure();
+
+ // Decide whether to import a particular parameter.
+ DenseSet<StringRef> importKeys;
+ for (auto &key : keys)
+ importKeys.insert(key);
+ auto shouldImportParameter =
+ [&](IREE::Stream::NamedParameterAttr parameterAttr) -> bool {
+ // Always try to import explicitly named parameters.
+ if (importKeys.contains(parameterAttr.getKey().getValue()))
+ return true; // key match
+ // If a maximum size is specified use that to limit what we import
+ // (users may want to bring in small parameters but leave the big ones
+ // out).
+ if (maximumSize && parameterAttr.getStorageSize() <= maximumSize)
+ return true; // <= max size
+ // Default to not importing.
+ return false;
+ };
+
+ // Find all parameters and try to import them.
+ for (auto globalOp : moduleOp.getOps<IREE::Util::GlobalOpInterface>()) {
+ // Only inspect parameter globals.
+ auto parameterAttr =
+ dyn_cast_if_present<IREE::Stream::NamedParameterAttr>(
+ globalOp.getGlobalInitialValue());
+ if (!parameterAttr)
+ continue;
+
+ // Lookup the parameter index for the scope.
+ auto scope = parameterAttr.getScope().getValue();
+ auto *parameterIndex = parameterIndices->lookup(scope);
+ if (!parameterIndex)
+ continue;
+
+ // See if the parameter is present in the scope (we may have only been
+ // provided as partial index).
+ auto key = parameterAttr.getKey().getValue();
+ const iree_io_parameter_index_entry_t *entry = nullptr;
+ iree_status_t lookupStatus = iree_io_parameter_index_lookup(
+ parameterIndex, iree_make_string_view(key.data(), key.size()),
+ &entry);
+ if (!iree_status_is_ok(lookupStatus)) {
+ // Parameter not found.
+ iree_status_ignore(lookupStatus);
+ continue;
+ }
+
+ // Filter only to globals of types we serialize.
+ if (!isTypeSupported(globalOp.getGlobalType())) {
+ llvm::errs() << "WARNING: not importing parameter `"
+ << globalOp.getGlobalName().getValue() << "` of type "
+ << globalOp.getGlobalType()
+ << " as packed types are not yet supported\n";
+ continue;
+ }
+
+ // Only import if the parameter meets any filtering criteria we have
+ // from the pass options.
+ if (shouldImportParameter(parameterAttr)) {
+ std::string fullName =
+ (StringRef("__import_") + scope + "_" + key).str();
+ auto valueOr = importParameter(
+ fullName, cast<ShapedType>(globalOp.getGlobalType()), parameterAttr,
+ entry);
+ if (failed(valueOr))
+ return signalPassFailure();
+
+ // Replace the initial value with the constant.
+ globalOp.setGlobalInitialValue(*valueOr);
+ }
+ }
+ }
+};
+
+} // namespace
+} // namespace mlir::iree_compiler::IREE::IO::Parameters
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/Passes.td b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/Passes.td
index 36d374f..00d7dad 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/Passes.td
@@ -11,32 +11,49 @@
def ExportParametersPass :
Pass<"iree-io-export-parameters", "mlir::ModuleOp"> {
- let summary = "Moves all inline constants of a minimum size and constant "
- "initialized globals values to a parameter archive";
+ let summary = "Exports all global constants to an archive file when "
+ "they are larger than the specified minimum size.";
let dependentDialects = [
"IREE::Stream::StreamDialect",
"IREE::Util::UtilDialect",
];
let options = [
- Option<"parameterScope", "scope", "std::string",
+ Option<"scopePath", "path", "std::string",
/*default=*/"",
- "Optional scope to use for the exported parameters.">,
- Option<"archivePath", "archive-path", "std::string",
- /*default=*/"",
- "Path to write the parameter archive to.">,
+ "File path to an archive to export from with an optional "
+ "`scope=` prefix.">,
Option<"minimumSize", "minimum-size", "int64_t",
- /*default=*/"256",
- "Minimum size of a serialized global to export.">
+ /*default=*/"0",
+ "Minimum size of a serialized global to export.">,
];
}
def GenerateSplatParameterArchivePass :
Pass<"iree-io-generate-splat-parameter-archive", "mlir::ModuleOp"> {
- let summary = "Generates a .irpa file with splat entries for all parameters";
+ let summary = "Generates a .irpa file with splat entries for all parameters.";
let options = [
- Option<"archivePath", "archive-path", "std::string",
+ Option<"filePath", "file", "std::string",
/*default=*/"",
- "Path to write the parameter archive to.">
+ "Path to write the parameter archive to.">,
+ ];
+}
+
+def ImportParametersPass :
+ Pass<"iree-io-import-parameters", "mlir::ModuleOp"> {
+ let summary = "Imports parameters from an archive file.";
+ let dependentDialects = [
+ "IREE::Stream::StreamDialect",
+ "IREE::Util::UtilDialect",
+ ];
+ let options = [
+ ListOption<"scopePaths", "paths", "std::string",
+ "File paths to archives to import from with an optional "
+ "`scope=` prefix.">,
+ ListOption<"keys", "keys", "std::string",
+ "List of parameter keys to import.">,
+ Option<"maximumSize", "maximum-size", "int64_t",
+ /*default=*/"9223372036854775807",
+ "Maximum size of a serialized global to import.">,
];
}
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/BUILD.bazel
index 7b6620f..eb9ea3f 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/BUILD.bazel
@@ -18,6 +18,7 @@
[
"export_parameters.mlir",
"generate_splat_parameter_archive.mlir",
+ "import_parameters.mlir",
],
include = ["*.mlir"],
),
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/CMakeLists.txt
index 72297bc..1d6407b 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/CMakeLists.txt
@@ -16,6 +16,7 @@
SRCS
"export_parameters.mlir"
"generate_splat_parameter_archive.mlir"
+ "import_parameters.mlir"
TOOLS
FileCheck
iree-dump-parameters
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/export_parameters.mlir b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/export_parameters.mlir
index 5354708..81ba9da 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/export_parameters.mlir
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/export_parameters.mlir
@@ -1,26 +1,42 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(iree-io-export-parameters{archive-path="%t.irpa" scope=opt minimum-size=0})" %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline="builtin.module(iree-io-export-parameters{path="opt=%t.irpa" minimum-size=0})" %s | FileCheck %s
+// RUN: iree-dump-parameters --parameters=%t.irpa | FileCheck %s --check-prefix=DUMP
-// CHECK-LABEL: module @parameter_example
-module @parameter_example {
-// CHECK-DAG: util.global private @array_global_0 = #stream.parameter.named<"opt"::"array_global_0"> : tensor<1x2xf32>
-// CHECK-DAG: util.global private @dense_global_1 = #stream.parameter.named<"opt"::"dense_global_1"> : tensor<2x2xf32>
-// CHECK-DAG: util.global private @constant_hoisted = #stream.parameter.named<"opt"::"constant_hoisted"> : tensor<1x2xf32>
-// CHECK-DAG: util.global private @dense_global_2 = #stream.parameter.named<"opt"::"dense_global_2"> : tensor<2x2xf32>
- util.global private @array_global_0 = dense<[[11.0, 12.0]]> : tensor<1x2xf32>
- util.global private @dense_global_1 = dense<"0x0000E040000000410000104100002041"> : tensor<2x2xf32>
- util.global private @dense_global_2 = dense<"0x0000803F000000400000404000008040"> : tensor<2x2xf32>
- util.func public @parameter_example(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> {
- %cst = arith.constant 0.000000e+00 : f32
- %3 = util.global.load @array_global_0 : tensor<1x2xf32>
- %4 = util.global.load @dense_global_1 : tensor<2x2xf32>
- %5 = arith.constant dense<"0x0000A0400000C040"> : tensor<1x2xf32>
- %6 = util.global.load @dense_global_2 : tensor<2x2xf32>
- %empty = tensor.empty() : tensor<1x2xf32>
- %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<1x2xf32>) -> tensor<1x2xf32>
- %8 = linalg.matmul ins(%arg0, %6 : tensor<1x2xf32>, tensor<2x2xf32>) outs(%fill : tensor<1x2xf32>) -> tensor<1x2xf32>
- %10 = linalg.add ins(%8, %5 : tensor<1x2xf32>, tensor<1x2xf32>) outs(%empty : tensor<1x2xf32>) -> tensor<1x2xf32>
- %12 = linalg.matmul ins(%10, %4 : tensor<1x2xf32>, tensor<2x2xf32>) outs(%fill : tensor<1x2xf32>) -> tensor<1x2xf32>
- %14 = linalg.add ins(%12, %3 : tensor<1x2xf32>, tensor<1x2xf32>) outs(%empty : tensor<1x2xf32>) -> tensor<1x2xf32>
- util.return %14 : tensor<1x2xf32>
- }
-}
+// CHECK: util.global private @constant_scalar_i1 = #stream.parameter.named<"opt"::"constant_scalar_i1"> : tensor<i1>
+// DUMP: - | - | 1 | `constant_scalar_i1`
+util.global private @constant_scalar_i1 = dense<true> : tensor<i1>
+
+// CHECK-NEXT: util.global private @constant_dense_2xi1 = #stream.parameter.named<"opt"::"constant_dense_2xi1"> : tensor<2xi1>
+// DUMP-NEXT: {{[0-9]+}} | {{[0-9]+}} | 2 | `constant_dense_2xi1`
+util.global private @constant_dense_2xi1 = dense<[true, false]> : tensor<2xi1>
+
+// CHECK-NEXT: util.global private @constant_dense_3xi4 = #stream.parameter.named<"opt"::"constant_dense_3xi4"> : tensor<3xi4>
+// DUMP-NEXT: {{[0-9]+}} | {{[0-9]+}} | 2 | `constant_dense_3xi4`
+util.global private @constant_dense_3xi4 = dense<[4, 5, 6]> : tensor<3xi4>
+
+// CHECK-NEXT: util.global private @constant_dense_2xi8 = #stream.parameter.named<"opt"::"constant_dense_2xi8"> : tensor<2xi8>
+// DUMP-NEXT: {{[0-9]+}} | {{[0-9]+}} | 2 | `constant_dense_2xi8`
+util.global private @constant_dense_2xi8 = dense<[4, 5]> : tensor<2xi8>
+
+// CHECK-NEXT: util.global private @constant_dense_2xf32 = #stream.parameter.named<"opt"::"constant_dense_2xf32"> : tensor<2xf32>
+// DUMP-NEXT: {{[0-9]+}} | {{[0-9]+}} | 8 | `constant_dense_2xf32`
+util.global private @constant_dense_2xf32 = dense<[11.0, 12.0]> : tensor<2xf32>
+
+// CHECK-NEXT: util.global private @constant_splat_2xf32 = #stream.parameter.named<"opt"::"constant_splat_2xf32"> : tensor<2xf32>
+// DUMP-NEXT: - | - | 8 | `constant_splat_2xf32`
+util.global private @constant_splat_2xf32 = dense<11.0> : tensor<2xf32>
+
+// CHECK-NEXT: util.global private mutable @mutable_scalar_i1 = #stream.parameter.named<"opt"::"mutable_scalar_i1"> : tensor<i1>
+// DUMP-NEXT: {{[0-9]+}} | {{[0-9]+}} | 1 | `mutable_scalar_i1`
+util.global private mutable @mutable_scalar_i1 = dense<true> : tensor<i1>
+
+// CHECK-NEXT: util.global private mutable @mutable_dense_3xi4 = #stream.parameter.named<"opt"::"mutable_dense_3xi4"> : tensor<3xi4>
+// DUMP-NEXT: {{[0-9]+}} | {{[0-9]+}} | 2 | `mutable_dense_3xi4`
+util.global private mutable @mutable_dense_3xi4 = dense<[4, 5, 6]> : tensor<3xi4>
+
+// CHECK-NEXT: util.global private mutable @mutable_dense_2xf32 = #stream.parameter.named<"opt"::"mutable_dense_2xf32"> : tensor<2xf32>
+// DUMP-NEXT: {{[0-9]+}} | {{[0-9]+}} | 8 | `mutable_dense_2xf32`
+util.global private mutable @mutable_dense_2xf32 = dense<[11.0, 12.0]> : tensor<2xf32>
+
+// CHECK-NEXT: util.global private mutable @mutable_splat_2xf32 = #stream.parameter.named<"opt"::"mutable_splat_2xf32"> : tensor<2xf32>
+// DUMP-NEXT: {{[0-9]+}} | {{[0-9]+}} | 8 | `mutable_splat_2xf32`
+util.global private mutable @mutable_splat_2xf32 = dense<11.0> : tensor<2xf32>
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/generate_splat_parameter_archive.mlir b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/generate_splat_parameter_archive.mlir
index b77dfa8..4944c01 100644
--- a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/generate_splat_parameter_archive.mlir
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/generate_splat_parameter_archive.mlir
@@ -1,35 +1,22 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(iree-io-generate-splat-parameter-archive{archive-path="%t.irpa"})" %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline="builtin.module(iree-io-generate-splat-parameter-archive{file="%t.irpa"})" %s | FileCheck %s
// RUN: iree-dump-parameters --parameters=%t.irpa | FileCheck %s --check-prefix=DUMP
-// CHECK-LABEL: @parameter_example
-module @parameter_example {
- // CHECK: util.global private @array_global_0 = #stream.parameter.named<"model"::"global_0">
- // CHECK: util.global private @dense_global_1 = #stream.parameter.named<"model"::"global_1">
- // CHECK: util.global private @dense_global_2 = #stream.parameter.named<"model"::"global_2">
- // CHECK: util.global private @dense_global_3 = #stream.parameter.named<"model"::"global_3">
- util.global private @array_global_0 = #stream.parameter.named<"model"::"global_0"> : tensor<1x2xi32>
- util.global private @dense_global_1 = #stream.parameter.named<"model"::"global_1"> : tensor<2x2xi32>
- util.global private @dense_global_2 = #stream.parameter.named<"model"::"global_2"> : tensor<1x2xi32>
- util.global private @dense_global_3 = #stream.parameter.named<"model"::"global_3"> : tensor<2x2xi32>
- util.func public @forward(%arg0: tensor<1x2xi32>) -> tensor<1x2xi32> {
- %cst = arith.constant 0 : i32
- %3 = util.global.load @array_global_0 : tensor<1x2xi32>
- %4 = util.global.load @dense_global_1 : tensor<2x2xi32>
- %5 = util.global.load @dense_global_2 : tensor<1x2xi32>
- %6 = util.global.load @dense_global_3 : tensor<2x2xi32>
- %empty = tensor.empty() : tensor<1x2xi32>
- %fill = linalg.fill ins(%cst : i32) outs(%empty : tensor<1x2xi32>) -> tensor<1x2xi32>
- %8 = linalg.matmul ins(%arg0, %6 : tensor<1x2xi32>, tensor<2x2xi32>) outs(%fill : tensor<1x2xi32>) -> tensor<1x2xi32>
- %10 = linalg.add ins(%8, %5 : tensor<1x2xi32>, tensor<1x2xi32>) outs(%empty : tensor<1x2xi32>) -> tensor<1x2xi32>
- %12 = linalg.matmul ins(%10, %4 : tensor<1x2xi32>, tensor<2x2xi32>) outs(%fill : tensor<1x2xi32>) -> tensor<1x2xi32>
- %14 = linalg.add ins(%12, %3 : tensor<1x2xi32>, tensor<1x2xi32>) outs(%empty : tensor<1x2xi32>) -> tensor<1x2xi32>
- util.return %14 : tensor<1x2xi32>
- }
-}
+// CHECK: util.global private @tensor_i1
+// DUMP: - | - | 1 | `tensor_i1`
+util.global private @tensor_i1 = #stream.parameter.named<"opt"::"tensor_i1"> : tensor<i1>
-// Verify the generated archive is what we expect.
-// DUMP: - |{{.*}} - |{{.*}} 8 | `global_0`
-// DUMP: - |{{.*}} - |{{.*}} 16 | `global_1`
-// DUMP: - |{{.*}} - |{{.*}} 8 | `global_2`
-// DUMP: - |{{.*}} - |{{.*}} 16 | `global_3`
+// CHECK-NEXT: util.global private @tensor_i8
+// DUMP-NEXT: - | - | 1 | `tensor_i8`
+util.global private @tensor_i8 = #stream.parameter.named<"opt"::"tensor_i8"> : tensor<i8>
+// CHECK-NEXT: util.global private @tensor_1x2xi32
+// DUMP-NEXT: - | - | 8 | `tensor_1x2xi32`
+util.global private @tensor_1x2xi32 = #stream.parameter.named<"opt"::"tensor_1x2xi32"> : tensor<1x2xi32>
+
+// CHECK-NEXT: util.global private @tensor_2x2xi4
+// DUMP-NEXT: - | - | 2 | `tensor_2x2xi4`
+util.global private @tensor_2x2xi4 = #stream.parameter.named<"opt"::"tensor_2x2xi4"> : tensor<2x2xi4>
+
+// CHECK-NEXT: util.global private @tensor_3xi4
+// DUMP-NEXT: - | - | 2 | `tensor_3xi4`
+util.global private @tensor_3xi4 = #stream.parameter.named<"opt"::"tensor_3xi4"> : tensor<3xi4>
diff --git a/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/import_parameters.mlir b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/import_parameters.mlir
new file mode 100644
index 0000000..075b9d0
--- /dev/null
+++ b/compiler/src/iree/compiler/Modules/IO/Parameters/Transforms/test/import_parameters.mlir
@@ -0,0 +1,36 @@
+// RUN: iree-opt --pass-pipeline="builtin.module(iree-io-export-parameters{path="opt=%t.irpa" minimum-size=0},iree-io-import-parameters{paths="opt=%t.irpa"})" %s | FileCheck %s
+
+// NOTE: packed types not supported for import yet.
+// CHECK: util.global private @constant_scalar_i1 = #stream.parameter.named
+util.global private @constant_scalar_i1 = dense<true> : tensor<i1>
+
+// NOTE: packed types not supported for import yet.
+// CHECK: util.global private @constant_dense_2xi1 = #stream.parameter.named
+util.global private @constant_dense_2xi1 = dense<[true, false]> : tensor<2xi1>
+
+// NOTE: packed types not supported for import yet.
+// CHECK: util.global private @constant_dense_3xi4 = #stream.parameter.named
+util.global private @constant_dense_3xi4 = dense<[4, 5, 6]> : tensor<3xi4>
+
+// CHECK: util.global private @constant_dense_2xi8 = dense<[4, 5]> : tensor<2xi8>
+util.global private @constant_dense_2xi8 = dense<[4, 5]> : tensor<2xi8>
+
+// CHECK: util.global private @constant_dense_2xf32 = dense<[1.100000e+01, 1.200000e+01]> : tensor<2xf32>
+util.global private @constant_dense_2xf32 = dense<[1.100000e+01, 1.200000e+01]> : tensor<2xf32>
+
+// CHECK: util.global private @constant_splat_2xf32 = dense<1.100000e+01> : tensor<2xf32>
+util.global private @constant_splat_2xf32 = dense<1.100000e+01> : tensor<2xf32>
+
+// NOTE: packed types not supported for import yet.
+// CHECK: util.global private mutable @mutable_scalar_i1 = #stream.parameter.named
+util.global private mutable @mutable_scalar_i1 = dense<true> : tensor<i1>
+
+// NOTE: packed types not supported for import yet.
+// CHECK: util.global private mutable @mutable_dense_3xi4 = #stream.parameter.named
+util.global private mutable @mutable_dense_3xi4 = dense<[4, 5, 6]> : tensor<3xi4>
+
+// CHECK: util.global private mutable @mutable_dense_2xf32 = dense<[1.100000e+01, 1.200000e+01]> : tensor<2xf32>
+util.global private mutable @mutable_dense_2xf32 = dense<[1.100000e+01, 1.200000e+01]> : tensor<2xf32>
+
+// CHECK: util.global private mutable @mutable_splat_2xf32 = dense<1.100000e+01> : tensor<2xf32>
+util.global private mutable @mutable_splat_2xf32 = dense<1.100000e+01> : tensor<2xf32>
diff --git a/compiler/src/iree/compiler/Pipelines/Options.cpp b/compiler/src/iree/compiler/Pipelines/Options.cpp
index 53f970b..83e0bec 100644
--- a/compiler/src/iree/compiler/Pipelines/Options.cpp
+++ b/compiler/src/iree/compiler/Pipelines/Options.cpp
@@ -70,6 +70,41 @@
}
}
+void PreprocessingOptions::bindOptions(OptionsBinder &binder) {
+ static llvm::cl::OptionCategory category(
+ "IREE options for apply custom preprocessing before normal IREE "
+ "compilation flow");
+
+ binder.opt<std::string>(
+ "iree-preprocessing-pass-pipeline", preprocessingPassPipeline,
+ llvm::cl::desc("Textual description of the pass pipeline to run before "
+ "running normal IREE compilation pipelines."),
+ llvm::cl::cat(category));
+ binder.opt<std::string>(
+ "iree-preprocessing-transform-spec-filename",
+ preprocessingTransformSpecFilename,
+ llvm::cl::desc(
+ "File name of a transform dialect spec to use for preprocessing."),
+ llvm::cl::cat(category));
+ binder.opt<std::string>(
+ "iree-preprocessing-pdl-spec-filename", preprocessingPDLSpecFilename,
+ llvm::cl::desc(
+ "File name of a transform dialect spec to use for preprocessing."),
+ llvm::cl::cat(category));
+
+ // DEPRECATED: do not add explicit options for specific passes.
+ binder.opt<TransposeMatmulInput>(
+ "iree-preprocessing-transpose-matmul", preprocessingTransposeMatmulInput,
+ llvm::cl::desc("Convert Linalg matmul ops to transposed variants."),
+ llvm::cl::cat(category),
+ llvm::cl::values(clEnumValN(TransposeMatmulInput::Lhs, "lhs",
+ "Transpose LHS input matrix."),
+ clEnumValN(TransposeMatmulInput::Rhs, "rhs",
+ "Transpose RHS input matrix."),
+ clEnumValN(TransposeMatmulInput::None, "none",
+ "Transpose neither input (disable).")));
+}
+
void GlobalOptimizationOptions::bindOptions(OptionsBinder &binder) {
static llvm::cl::OptionCategory category(
"IREE options for controlling global optimizations.");
@@ -139,26 +174,35 @@
llvm::cl::desc("Strips debug assertions after any useful "
"information has been extracted."),
llvm::cl::cat(category));
- binder.opt<std::string>(
- "iree-opt-parameter-archive-export-file", parameterArchiveExportPath,
- llvm::cl::desc(
- "File path to create a parameter archive using any inline global "
- "constants."),
+
+ binder.list<std::string>(
+ "iree-opt-import-parameters", parameterImportPaths,
+ llvm::cl::desc("File paths to archives to import parameters from with an "
+ "optional `scope=` prefix."),
llvm::cl::cat(category));
+ binder.list<std::string>("iree-opt-import-parameter-keys",
+ parameterImportKeys,
+ llvm::cl::desc("List of parameter keys to import."),
+ llvm::cl::cat(category));
+ binder.opt<int64_t>("iree-opt-import-parameter-maximum-size",
+ parameterImportMaximumSize,
+ llvm::cl::desc("Maximum size of parameters to import."),
+ llvm::cl::cat(category));
+
binder.opt<std::string>(
- "iree-opt-parameter-archive-export-scope", parameterExportScope,
- llvm::cl::desc("Scope for parameters in the archive created in "
- "`iree-opt-export-parameter-archive-export-file`."),
+ "iree-opt-export-parameters", parameterExportPath,
+ llvm::cl::desc("File path to an archive to export parameters to with an "
+ "optional `scope=` prefix."),
llvm::cl::cat(category));
binder.opt<int64_t>(
- "iree-opt-minimum-parameter-export-size", minimumParameterExportSize,
+ "iree-opt-export-parameter-minimum-size", parameterExportMinimumSize,
llvm::cl::desc(
"Minimum size of constants to export to the archive created in "
"`iree-opt-export-parameter-archive-export-file`."),
llvm::cl::cat(category));
+
binder.opt<std::string>(
- "iree-opt-splat-parameter-archive-export-file",
- splatParameterArchiveExportPath,
+ "iree-opt-splat-parameters", parameterSplatExportFile,
llvm::cl::desc(
"File path to create a parameter archive of splat values out of all "
"parameter backed globals."),
@@ -216,46 +260,4 @@
llvm::cl::cat(category));
}
-void PreprocessingOptions::bindOptions(OptionsBinder &binder) {
- static llvm::cl::OptionCategory category(
- "IREE options for apply custom preprocessing before normal IREE "
- "compilation flow");
-
- binder.opt<std::string>(
- "iree-preprocessing-pass-pipeline", preprocessingPassPipeline,
- llvm::cl::desc("Textual description of the pass pipeline to run before "
- "running normal IREE compilation pipelines"),
- llvm::cl::cat(category));
- // Following ways of doing custom transformations at the pre-processing step
- // are supported.
-
- // 0. Through a preprocessing pass pipeline.
- // 1. Through a Transform dialect spec file.
- // 2. Through a PDL spec file.
- // A user may simultaneously use both. The order is transform dialect
- // transforms are applied first and then the pdl patterns.
- binder.opt<std::string>(
- "iree-preprocessing-transform-spec-filename",
- preprocessingTransformSpecFilename,
- llvm::cl::desc(
- "File name of a transform dialect spec to use for preprocessing"),
- llvm::cl::cat(category));
- binder.opt<std::string>(
- "iree-preprocessing-pdl-spec-filename", preprocessingPDLSpecFilename,
- llvm::cl::desc(
- "File name of a transform dialect spec to use for preprocessing"),
- llvm::cl::cat(category));
-
- binder.opt<TransposeMatmulInput>(
- "iree-preprocessing-transpose-matmul", preprocessingTransposeMatmulInput,
- llvm::cl::desc("Convert Linalg matmul ops to transposed variants."),
- llvm::cl::cat(category),
- llvm::cl::values(clEnumValN(TransposeMatmulInput::Lhs, "lhs",
- "Transpose LHS input matrix."),
- clEnumValN(TransposeMatmulInput::Rhs, "rhs",
- "Transpose RHS input matrix."),
- clEnumValN(TransposeMatmulInput::None, "none",
- "Transpose neither input (disable).")));
-}
-
} // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Pipelines/Options.h b/compiler/src/iree/compiler/Pipelines/Options.h
index 9734378..b704599 100644
--- a/compiler/src/iree/compiler/Pipelines/Options.h
+++ b/compiler/src/iree/compiler/Pipelines/Options.h
@@ -58,6 +58,33 @@
using FromFlags = OptionsFromFlags<InputDialectOptions>;
};
+// Allows specifying one of several ways of doing custom transformations at the
+// pre-processing phase, multiple ways may be used and they are run in order:
+// 1. Through a preprocessing pass pipeline.
+// 2. Through a Transform dialect spec file.
+// 3. Through a PDL spec file.
+struct PreprocessingOptions {
+ std::string preprocessingPassPipeline;
+ std::string preprocessingTransformSpecFilename;
+ std::string preprocessingPDLSpecFilename;
+
+ // DEPRECATED: do not put pass-specific options here and instead use the
+ // pass pipeline.
+ enum class TransposeMatmulInput {
+ /// Transpose LHS input matrix.
+ Lhs,
+ /// Transpose RHS input matrix.
+ Rhs,
+ /// Transpose neither input (disable).
+ None
+ };
+ TransposeMatmulInput preprocessingTransposeMatmulInput =
+ TransposeMatmulInput::None;
+
+ void bindOptions(OptionsBinder &binder);
+ using FromFlags = OptionsFromFlags<PreprocessingOptions>;
+};
+
// Options controlling high level optimizations.
struct GlobalOptimizationOptions {
// Gate various type based demotion passes that run before anything else.
@@ -94,18 +121,24 @@
// allow hoisting. The threshold is 1MB by default.
int64_t constExprMaxSizeIncreaseThreshold = 1024 * 1024;
- // File path to create a parameter archive out of global initial values.
- std::string parameterArchiveExportPath = "";
+ // File paths to archives to import parameters from with an optional
+ // `scope=` prefix.
+ std::vector<std::string> parameterImportPaths;
+ // List of parameter keys to import. Any matching keys from any scope will be
+ // imported.
+ std::vector<std::string> parameterImportKeys;
+ // Maximum size of parameters to import or 0 to disable automatic import.
+ int64_t parameterImportMaximumSize = 0;
- // Optional scope to use for the created parameter archive.
- std::string parameterExportScope = "";
+ // File path to an archive to export parameters to with an optional
+ // `scope=` prefix.
+ std::string parameterExportPath;
+ // Minimum size of constants to export as parameters.
+ int64_t parameterExportMinimumSize = 0;
// File path to create a splat parameter archive out of all parameters in the
// module.
- std::string splatParameterArchiveExportPath = "";
-
- // Minimum size of constants to export as parameters.
- int64_t minimumParameterExportSize = 256;
+ std::string parameterSplatExportFile = "";
void bindOptions(OptionsBinder &binder);
using FromFlags = OptionsFromFlags<GlobalOptimizationOptions>;
@@ -166,26 +199,6 @@
using FromFlags = OptionsFromFlags<SchedulingOptions>;
};
-struct PreprocessingOptions {
- /// Options for converting Linalg matmul ops to transposed variants.
- enum class TransposeMatmulInput {
- /// Transpose LHS input matrix.
- Lhs,
- /// Transpose RHS input matrix.
- Rhs,
- /// Transpose neither input (disable).
- None
- };
-
- std::string preprocessingPassPipeline;
- std::string preprocessingTransformSpecFilename;
- std::string preprocessingPDLSpecFilename;
- TransposeMatmulInput preprocessingTransposeMatmulInput =
- TransposeMatmulInput::None;
- void bindOptions(OptionsBinder &binder);
- using FromFlags = OptionsFromFlags<PreprocessingOptions>;
-};
-
} // namespace mlir::iree_compiler
#endif // IREE_COMPILER_PIPELINES_OPTIONS_H_
diff --git a/compiler/src/iree/compiler/Preprocessing/Passes.cpp b/compiler/src/iree/compiler/Preprocessing/Passes.cpp
index 7c12279..a81cac4 100644
--- a/compiler/src/iree/compiler/Preprocessing/Passes.cpp
+++ b/compiler/src/iree/compiler/Preprocessing/Passes.cpp
@@ -94,6 +94,7 @@
passManager.addPass(createCSEPass());
}
+ // DEPRECATED: do not add explicit options for specific passes.
if (preprocessingOptions.preprocessingTransposeMatmulInput !=
PreprocessingOptions::TransposeMatmulInput::None) {
Preprocessing::TransposeMatmulPassOptions transposeMatmulOptions;
diff --git a/runtime/bindings/python/CMakeLists.txt b/runtime/bindings/python/CMakeLists.txt
index 3254877..9f3151b 100644
--- a/runtime/bindings/python/CMakeLists.txt
+++ b/runtime/bindings/python/CMakeLists.txt
@@ -91,9 +91,8 @@
iree::base::internal::file_io
iree::base::internal::path
iree::io::file_handle
- iree::io::formats::gguf
iree::io::formats::irpa
- iree::io::formats::safetensors
+ iree::io::formats::parser_registry
iree::io::parameter_index
iree::io::parameter_index_provider
iree::io::parameter_provider
diff --git a/runtime/bindings/python/io.cc b/runtime/bindings/python/io.cc
index 3305698..a558cb9 100644
--- a/runtime/bindings/python/io.cc
+++ b/runtime/bindings/python/io.cc
@@ -14,10 +14,8 @@
#include "./vm.h"
#include "iree/base/internal/file_io.h"
#include "iree/base/internal/path.h"
-#include "iree/io/formats/gguf/gguf_parser.h"
#include "iree/io/formats/irpa/irpa_builder.h"
-#include "iree/io/formats/irpa/irpa_parser.h"
-#include "iree/io/formats/safetensors/safetensors_parser.h"
+#include "iree/io/formats/parser_registry.h"
#include "iree/io/parameter_index_provider.h"
#include "iree/modules/io/parameters/module.h"
#include "iree/schemas/parameter_archive.h"
@@ -100,23 +98,10 @@
void ParameterIndexParseFileHandle(ParameterIndex &self,
FileHandle &file_handle,
std::string &format) {
- if (format == "gguf") {
- CheckApiStatus(
- iree_io_parse_gguf_index(file_handle.raw_ptr(), self.raw_ptr()),
- "Could not parse gguf file into index");
- } else if (format == "irpa") {
- CheckApiStatus(
- iree_io_parse_irpa_index(file_handle.raw_ptr(), self.raw_ptr()),
- "Could not parse IREE parameter archive file into index");
- } else if (format == "safetensors") {
- CheckApiStatus(
- iree_io_parse_safetensors_index(file_handle.raw_ptr(), self.raw_ptr()),
- "Could not parse safetensors file into index");
- } else {
- throw std::invalid_argument(
- "Unrecognized file format. Expected one of: 'gguf', 'irpa', "
- "'safetensors'");
- }
+ CheckApiStatus(iree_io_parse_file_index(
+ iree_make_string_view(format.data(), format.size()),
+ file_handle.raw_ptr(), self.raw_ptr()),
+ "Could not parse parameter file index");
}
void ParameterIndexLoadFile(ParameterIndex &self, std::string &file_path,
diff --git a/runtime/src/iree/base/status.c b/runtime/src/iree/base/status.c
index 3c8a7be..2c384a9 100644
--- a/runtime/src/iree/base/status.c
+++ b/runtime/src/iree/base/status.c
@@ -712,7 +712,6 @@
} else {
buffer[buffer_length] = ';';
buffer[buffer_length + 1] = ' ';
- buffer[buffer_length + 2] = '\0';
}
}
buffer_length += 2; // '; '
@@ -772,12 +771,13 @@
iree_host_size_t message_buffer_length = 0;
bool ret = iree_status_format_message(
- status, buffer_capacity, buffer ? buffer + prefix_buffer_length : NULL,
- &message_buffer_length, /*has_prefix=*/true);
+ status, buffer ? buffer_capacity - prefix_buffer_length : 0,
+ buffer ? buffer + prefix_buffer_length : NULL, &message_buffer_length,
+ /*has_prefix=*/true);
if (!ret) {
return false;
}
- *out_buffer_length = message_buffer_length + prefix_buffer_length;
+ *out_buffer_length = prefix_buffer_length + message_buffer_length;
return true;
}
diff --git a/runtime/src/iree/io/formats/BUILD.bazel b/runtime/src/iree/io/formats/BUILD.bazel
index 522ca5d..a1bd0ba 100644
--- a/runtime/src/iree/io/formats/BUILD.bazel
+++ b/runtime/src/iree/io/formats/BUILD.bazel
@@ -4,8 +4,29 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+load("//build_tools/bazel:build_defs.oss.bzl", "iree_runtime_cc_library")
+
package(
default_visibility = ["//visibility:public"],
features = ["layering_check"],
licenses = ["notice"], # Apache 2.0
)
+
+iree_runtime_cc_library(
+ name = "parser_registry",
+ srcs = [
+ "parser_registry.c",
+ ],
+ hdrs = [
+ "parser_registry.h",
+ ],
+ deps = [
+ "//runtime/src/iree/base",
+ "//runtime/src/iree/base/internal:path",
+ "//runtime/src/iree/io:file_handle",
+ "//runtime/src/iree/io:parameter_index",
+ "//runtime/src/iree/io/formats/gguf",
+ "//runtime/src/iree/io/formats/irpa",
+ "//runtime/src/iree/io/formats/safetensors",
+ ],
+)
diff --git a/runtime/src/iree/io/formats/CMakeLists.txt b/runtime/src/iree/io/formats/CMakeLists.txt
index df50aab..72cc006 100644
--- a/runtime/src/iree/io/formats/CMakeLists.txt
+++ b/runtime/src/iree/io/formats/CMakeLists.txt
@@ -10,4 +10,22 @@
iree_add_all_subdirs()
+iree_cc_library(
+ NAME
+ parser_registry
+ HDRS
+ "parser_registry.h"
+ SRCS
+ "parser_registry.c"
+ DEPS
+ iree::base
+ iree::base::internal::path
+ iree::io::file_handle
+ iree::io::formats::gguf
+ iree::io::formats::irpa
+ iree::io::formats::safetensors
+ iree::io::parameter_index
+ PUBLIC
+)
+
### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/runtime/src/iree/io/formats/irpa/irpa_builder.c b/runtime/src/iree/io/formats/irpa/irpa_builder.c
index 91563e4..330691f 100644
--- a/runtime/src/iree/io/formats/irpa/irpa_builder.c
+++ b/runtime/src/iree/io/formats/irpa/irpa_builder.c
@@ -24,6 +24,12 @@
memset(builder, 0, sizeof(*builder));
}
+IREE_API_EXPORT bool iree_io_parameter_archive_builder_is_empty(
+ const iree_io_parameter_archive_builder_t* builder) {
+ IREE_ASSERT_ARGUMENT(builder);
+ return iree_io_parameter_index_count(builder->index) == 0;
+}
+
static iree_io_physical_size_t
iree_io_parameter_archive_builder_storage_alignment(
const iree_io_parameter_archive_builder_t* builder) {
diff --git a/runtime/src/iree/io/formats/irpa/irpa_builder.h b/runtime/src/iree/io/formats/irpa/irpa_builder.h
index 22ac538..46d9a7f 100644
--- a/runtime/src/iree/io/formats/irpa/irpa_builder.h
+++ b/runtime/src/iree/io/formats/irpa/irpa_builder.h
@@ -53,6 +53,10 @@
IREE_API_EXPORT void iree_io_parameter_archive_builder_deinitialize(
iree_io_parameter_archive_builder_t* builder);
+// Returns true if no parameters have been added to the archive.
+IREE_API_EXPORT bool iree_io_parameter_archive_builder_is_empty(
+ const iree_io_parameter_archive_builder_t* builder);
+
// Returns the total file size required to store the parameter archive header
// and contents of all added parameters. Adding new parameters will invalidate
// this value.
diff --git a/runtime/src/iree/io/formats/parser_registry.c b/runtime/src/iree/io/formats/parser_registry.c
new file mode 100644
index 0000000..716f07f
--- /dev/null
+++ b/runtime/src/iree/io/formats/parser_registry.c
@@ -0,0 +1,46 @@
+// Copyright 2024 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/io/formats/parser_registry.h"
+
+#include "iree/base/internal/path.h"
+#include "iree/io/formats/gguf/gguf_parser.h"
+#include "iree/io/formats/irpa/irpa_parser.h"
+#include "iree/io/formats/safetensors/safetensors_parser.h"
+
+IREE_API_EXPORT iree_status_t iree_io_parse_file_index(
+ iree_string_view_t path, iree_io_file_handle_t* file_handle,
+ iree_io_parameter_index_t* index) {
+ IREE_TRACE_ZONE_BEGIN(z0);
+ IREE_TRACE_ZONE_APPEND_TEXT(z0, path.data, path.size);
+
+ // Try to get the extension from the path but also support the user passing in
+ // just the raw extension as well.
+ iree_string_view_t basename = iree_string_view_empty();
+ iree_string_view_t extension = iree_string_view_empty();
+ iree_file_path_split_basename(path, &basename, &extension);
+ if (iree_string_view_is_empty(extension)) {
+ extension = basename;
+ }
+
+ iree_status_t status = iree_ok_status();
+ if (iree_string_view_equal_case(extension, IREE_SV("irpa"))) {
+ status = iree_io_parse_irpa_index(file_handle, index);
+ } else if (iree_string_view_equal_case(extension, IREE_SV("gguf"))) {
+ status = iree_io_parse_gguf_index(file_handle, index);
+ } else if (iree_string_view_equal_case(extension, IREE_SV("safetensors"))) {
+ status = iree_io_parse_safetensors_index(file_handle, index);
+ } else {
+ status = iree_make_status(
+ IREE_STATUS_UNIMPLEMENTED,
+ "unsupported file format `%.*s`; ensure the extension matches one of "
+ "the supported formats: [.irpa, .gguf, .safetensors]",
+ (int)extension.size, extension.data);
+ }
+
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
diff --git a/runtime/src/iree/io/formats/parser_registry.h b/runtime/src/iree/io/formats/parser_registry.h
new file mode 100644
index 0000000..c982d08
--- /dev/null
+++ b/runtime/src/iree/io/formats/parser_registry.h
@@ -0,0 +1,30 @@
+// Copyright 2024 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_IO_FORMATS_PARSER_REGISTRY_H_
+#define IREE_IO_FORMATS_PARSER_REGISTRY_H_
+
+#include "iree/base/api.h"
+#include "iree/io/file_handle.h"
+#include "iree/io/parameter_index.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif // __cplusplus
+
+// Parses a parameter file index in the opened |file_handle|.
+// |path| is used for logging and file format identification. It may either be
+// the original file path of |file_handle| or an extension (such as `irpa`).
+// Upon return any parameters in the file are appended to the |index|.
+IREE_API_EXPORT iree_status_t iree_io_parse_file_index(
+ iree_string_view_t path, iree_io_file_handle_t* file_handle,
+ iree_io_parameter_index_t* index);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif // __cplusplus
+
+#endif // IREE_IO_FORMATS_PARSER_REGISTRY_H_
diff --git a/runtime/src/iree/tooling/BUILD.bazel b/runtime/src/iree/tooling/BUILD.bazel
index 07ebfab..39aed05 100644
--- a/runtime/src/iree/tooling/BUILD.bazel
+++ b/runtime/src/iree/tooling/BUILD.bazel
@@ -198,15 +198,12 @@
"//runtime/src/iree/base",
"//runtime/src/iree/base/internal:file_io",
"//runtime/src/iree/base/internal:flags",
- "//runtime/src/iree/base/internal:path",
"//runtime/src/iree/hal",
"//runtime/src/iree/io:parameter_index",
"//runtime/src/iree/io:parameter_index_provider",
"//runtime/src/iree/io:parameter_provider",
"//runtime/src/iree/io:scope_map",
- "//runtime/src/iree/io/formats/gguf",
- "//runtime/src/iree/io/formats/irpa",
- "//runtime/src/iree/io/formats/safetensors",
+ "//runtime/src/iree/io/formats:parser_registry",
"//runtime/src/iree/modules/io/parameters",
"//runtime/src/iree/vm",
],
diff --git a/runtime/src/iree/tooling/CMakeLists.txt b/runtime/src/iree/tooling/CMakeLists.txt
index 2b5a1c8..451b3f8 100644
--- a/runtime/src/iree/tooling/CMakeLists.txt
+++ b/runtime/src/iree/tooling/CMakeLists.txt
@@ -228,11 +228,8 @@
iree::base
iree::base::internal::file_io
iree::base::internal::flags
- iree::base::internal::path
iree::hal
- iree::io::formats::gguf
- iree::io::formats::irpa
- iree::io::formats::safetensors
+ iree::io::formats::parser_registry
iree::io::parameter_index
iree::io::parameter_index_provider
iree::io::parameter_provider
diff --git a/runtime/src/iree/tooling/parameter_util.c b/runtime/src/iree/tooling/parameter_util.c
index fad5c72..07a5296 100644
--- a/runtime/src/iree/tooling/parameter_util.c
+++ b/runtime/src/iree/tooling/parameter_util.c
@@ -8,10 +8,7 @@
#include "iree/base/internal/file_io.h"
#include "iree/base/internal/flags.h"
-#include "iree/base/internal/path.h"
-#include "iree/io/formats/gguf/gguf_parser.h"
-#include "iree/io/formats/irpa/irpa_parser.h"
-#include "iree/io/formats/safetensors/safetensors_parser.h"
+#include "iree/io/formats/parser_registry.h"
#include "iree/io/parameter_index.h"
#include "iree/io/parameter_index_provider.h"
#include "iree/io/scope_map.h"
@@ -109,19 +106,7 @@
z0, iree_io_open_parameter_file(path, host_allocator, &file_handle));
// Index the file based on its (inferred) format.
- iree_status_t status = iree_ok_status();
- iree_string_view_t path_ext = iree_file_path_extension(path);
- if (iree_string_view_equal_case(path_ext, IREE_SV("irpa"))) {
- status = iree_io_parse_irpa_index(file_handle, index);
- } else if (iree_string_view_equal_case(path_ext, IREE_SV("gguf"))) {
- status = iree_io_parse_gguf_index(file_handle, index);
- } else if (iree_string_view_equal_case(path_ext, IREE_SV("safetensors"))) {
- status = iree_io_parse_safetensors_index(file_handle, index);
- } else {
- status = iree_make_status(IREE_STATUS_UNIMPLEMENTED,
- "unhandled parameter file format: .%.*s",
- (int)path_ext.size, path_ext.data);
- }
+ iree_status_t status = iree_io_parse_file_index(path, file_handle, index);
// Release our file reference - it's still retained by the index if it had any
// parameters in it.
diff --git a/tests/e2e/parameters/export_parameters.mlir b/tests/e2e/parameters/export_parameters.mlir
index cbf292e..9c8402e 100644
--- a/tests/e2e/parameters/export_parameters.mlir
+++ b/tests/e2e/parameters/export_parameters.mlir
@@ -1,41 +1,33 @@
-module @parameter_example {
- util.global private @array_global_0 = dense<[[11.0, 12.0]]> : tensor<1x2xf32>
- util.global private @dense_global_1 = dense<"0x0000E040000000410000104100002041"> : tensor<2x2xf32>
- util.global private @dense_global_2 = dense<"0x0000A0400000C040"> : tensor<1x2xf32>
- util.global private @dense_global_3 = dense<"0x0000803F000000400000404000008040"> : tensor<2x2xf32>
- func.func @predict(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> {
- %cst = arith.constant 0.000000e+00 : f32
- %3 = util.global.load @array_global_0 : tensor<1x2xf32>
- %4 = util.global.load @dense_global_1 : tensor<2x2xf32>
- %5 = util.global.load @dense_global_2 : tensor<1x2xf32>
- %6 = util.global.load @dense_global_3 : tensor<2x2xf32>
- %empty = tensor.empty() : tensor<1x2xf32>
- %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<1x2xf32>) -> tensor<1x2xf32>
- %8 = linalg.matmul ins(%arg0, %6 : tensor<1x2xf32>, tensor<2x2xf32>) outs(%fill : tensor<1x2xf32>) -> tensor<1x2xf32>
- %10 = linalg.add ins(%8, %5 : tensor<1x2xf32>, tensor<1x2xf32>) outs(%empty : tensor<1x2xf32>) -> tensor<1x2xf32>
- %12 = linalg.matmul ins(%10, %4 : tensor<1x2xf32>, tensor<2x2xf32>) outs(%fill : tensor<1x2xf32>) -> tensor<1x2xf32>
- %14 = linalg.add ins(%12, %3 : tensor<1x2xf32>, tensor<1x2xf32>) outs(%empty : tensor<1x2xf32>) -> tensor<1x2xf32>
- return %14 : tensor<1x2xf32>
- }
-}
-
// RUN: iree-compile %s \
// RUN: --iree-hal-target-backends=vmvx \
-// RUN: --iree-opt-parameter-archive-export-file=%t.irpa \
-// RUN: --iree-opt-parameter-archive-export-scope=compile \
-// RUN: --iree-opt-minimum-parameter-export-size=0 | \
-// RUN: iree-run-module --device=local-task --module=- \
-// RUN: --input=1x2xf32=1.0 \
-// RUN: --parameters=compile=%t.irpa \
-// RUN: --function=predict > %t.txt
-// RUN: iree-dump-parameters --parameters=compile=%t.irpa >> %t.txt
-// RUN: FileCheck %s --input-file=%t.txt
+// RUN: --iree-opt-export-parameters=scope=%t.irpa \
+// RUN: --iree-opt-export-parameter-minimum-size=0 | \
+// RUN: iree-run-module \
+// RUN: --device=local-sync \
+// RUN: --module=- \
+// RUN: --function=main \
+// RUN: --parameters=scope=%t.irpa \
+// RUN: --input=1x2xf32=1.0 | \
+// RUN: FileCheck %s
-// CHECK-LABEL: EXEC @predict
+// CHECK-LABEL: EXEC @main
// CHECK: 1x2xf32=[182 204]
-// CHECK: 512 |{{.*}} 520 |{{.*}} 8 | `constant_hoisted`
-// CHECK: 576 |{{.*}} 584 |{{.*}} 8 | `constant_hoisted_0`
-// CHECK: 640 |{{.*}} 656 |{{.*}} 16 | `constant_hoisted_1`
-// CHECK: 704 |{{.*}} 720 |{{.*}} 16 | `constant_hoisted_2`
-
+util.global private @array_global_0 = dense<[[11.0, 12.0]]> : tensor<1x2xf32>
+util.global private @dense_global_1 = dense<"0x0000E040000000410000104100002041"> : tensor<2x2xf32>
+util.global private @dense_global_2 = dense<"0x0000A0400000C040"> : tensor<1x2xf32>
+util.global private @dense_global_3 = dense<"0x0000803F000000400000404000008040"> : tensor<2x2xf32>
+func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %3 = util.global.load @array_global_0 : tensor<1x2xf32>
+ %4 = util.global.load @dense_global_1 : tensor<2x2xf32>
+ %5 = util.global.load @dense_global_2 : tensor<1x2xf32>
+ %6 = util.global.load @dense_global_3 : tensor<2x2xf32>
+ %empty = tensor.empty() : tensor<1x2xf32>
+ %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<1x2xf32>) -> tensor<1x2xf32>
+ %8 = linalg.matmul ins(%arg0, %6 : tensor<1x2xf32>, tensor<2x2xf32>) outs(%fill : tensor<1x2xf32>) -> tensor<1x2xf32>
+ %10 = linalg.add ins(%8, %5 : tensor<1x2xf32>, tensor<1x2xf32>) outs(%empty : tensor<1x2xf32>) -> tensor<1x2xf32>
+ %12 = linalg.matmul ins(%10, %4 : tensor<1x2xf32>, tensor<2x2xf32>) outs(%fill : tensor<1x2xf32>) -> tensor<1x2xf32>
+ %14 = linalg.add ins(%12, %3 : tensor<1x2xf32>, tensor<1x2xf32>) outs(%empty : tensor<1x2xf32>) -> tensor<1x2xf32>
+ return %14 : tensor<1x2xf32>
+}
diff --git a/tests/e2e/parameters/generate_splat_archive.mlir b/tests/e2e/parameters/generate_splat_archive.mlir
index 85c160d..a7b3607 100644
--- a/tests/e2e/parameters/generate_splat_archive.mlir
+++ b/tests/e2e/parameters/generate_splat_archive.mlir
@@ -1,31 +1,33 @@
-module @parameter_example {
- util.global private @array_global_0 = #stream.parameter.named<"model"::"global_0"> : tensor<1x2xi32>
- util.global private @dense_global_1 = #stream.parameter.named<"model"::"global_1"> : tensor<2x2xi32>
- util.global private @dense_global_2 = #stream.parameter.named<"model"::"global_2"> : tensor<1x2xi32>
- util.global private @dense_global_3 = #stream.parameter.named<"model"::"global_3"> : tensor<2x2xi32>
- func.func @forward(%arg0: tensor<1x2xi32>) -> tensor<1x2xi32> {
- %cst = arith.constant 0 : i32
- %3 = util.global.load @array_global_0 : tensor<1x2xi32>
- %4 = util.global.load @dense_global_1 : tensor<2x2xi32>
- %5 = util.global.load @dense_global_2 : tensor<1x2xi32>
- %6 = util.global.load @dense_global_3 : tensor<2x2xi32>
- %empty = tensor.empty() : tensor<1x2xi32>
- %fill = linalg.fill ins(%cst : i32) outs(%empty : tensor<1x2xi32>) -> tensor<1x2xi32>
- %8 = linalg.matmul ins(%arg0, %6 : tensor<1x2xi32>, tensor<2x2xi32>) outs(%fill : tensor<1x2xi32>) -> tensor<1x2xi32>
- %10 = linalg.add ins(%8, %5 : tensor<1x2xi32>, tensor<1x2xi32>) outs(%empty : tensor<1x2xi32>) -> tensor<1x2xi32>
- %12 = linalg.matmul ins(%10, %4 : tensor<1x2xi32>, tensor<2x2xi32>) outs(%fill : tensor<1x2xi32>) -> tensor<1x2xi32>
- %14 = linalg.add ins(%12, %3 : tensor<1x2xi32>, tensor<1x2xi32>) outs(%empty : tensor<1x2xi32>) -> tensor<1x2xi32>
- return %14 : tensor<1x2xi32>
- }
-}
// RUN: iree-compile %s \
// RUN: --iree-hal-target-backends=vmvx \
-// RUN: --iree-opt-splat-parameter-archive-export-file=%t.irpa | \
-// RUN: iree-run-module --device=local-task --module=- \
-// RUN: --input=1x2xi32=1 \
-// RUN: --parameters=model=%t.irpa \
-// RUN: --function=forward | FileCheck %s
+// RUN: --iree-opt-splat-parameters=%t.irpa | \
+// RUN: iree-run-module \
+// RUN: --device=local-sync \
+// RUN: --module=- \
+// RUN: --function=main \
+// RUN: --parameters=scope=%t.irpa \
+// RUN: --input=1x2xi32=1 | \
+// RUN: FileCheck %s
-// CHECK-LABEL: EXEC @forward
+// CHECK-LABEL: EXEC @main
// CHECK: 1x2xi32=[0 0]
+
+util.global private @array_global_0 = #stream.parameter.named<"scope"::"global_0"> : tensor<1x2xi32>
+util.global private @dense_global_1 = #stream.parameter.named<"scope"::"global_1"> : tensor<2x2xi32>
+util.global private @dense_global_2 = #stream.parameter.named<"scope"::"global_2"> : tensor<1x2xi32>
+util.global private @dense_global_3 = #stream.parameter.named<"scope"::"global_3"> : tensor<2x2xi32>
+func.func @main(%arg0: tensor<1x2xi32>) -> tensor<1x2xi32> {
+ %cst = arith.constant 0 : i32
+ %3 = util.global.load @array_global_0 : tensor<1x2xi32>
+ %4 = util.global.load @dense_global_1 : tensor<2x2xi32>
+ %5 = util.global.load @dense_global_2 : tensor<1x2xi32>
+ %6 = util.global.load @dense_global_3 : tensor<2x2xi32>
+ %empty = tensor.empty() : tensor<1x2xi32>
+ %fill = linalg.fill ins(%cst : i32) outs(%empty : tensor<1x2xi32>) -> tensor<1x2xi32>
+ %8 = linalg.matmul ins(%arg0, %6 : tensor<1x2xi32>, tensor<2x2xi32>) outs(%fill : tensor<1x2xi32>) -> tensor<1x2xi32>
+ %10 = linalg.add ins(%8, %5 : tensor<1x2xi32>, tensor<1x2xi32>) outs(%empty : tensor<1x2xi32>) -> tensor<1x2xi32>
+ %12 = linalg.matmul ins(%10, %4 : tensor<1x2xi32>, tensor<2x2xi32>) outs(%fill : tensor<1x2xi32>) -> tensor<1x2xi32>
+ %14 = linalg.add ins(%12, %3 : tensor<1x2xi32>, tensor<1x2xi32>) outs(%empty : tensor<1x2xi32>) -> tensor<1x2xi32>
+ return %14 : tensor<1x2xi32>
+}
