[Stream] Add UnifyEncodingForGlobals pass. 1/n (#22767)
This pass identifies immutable globals that are encoded with multiple
different encodings and unifies them to a single encoding. This reduces
memory footprint when the same source data (e.g., model weights) is used
in multiple dispatch operations with different encoding requirements.
The pass:
- Traces from stream.tensor.encode ops back to source globals
- Requires both source and encoded globals to be immutable
- Uses identity encoding as a conservative fallback
- Updates both stream.tensor.encode and stream.tensor.sizeof
consistently
This is the minimal functioning implementation. The pass will be built
incrementally with resolver integration and subgraph equivalence
support.
Note that it does not work completely as dispatch sites and executables
are not updated.
It is a step towards #22485.
---------
Signed-off-by: hanhanW <hanhan0912@gmail.com>
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
index e4aa5ea..f84c181 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
@@ -53,6 +53,7 @@
"SpecializeDispatches.cpp",
"SpecializeEncodings.cpp",
"SyncInitializers.cpp",
+ "UnifyEncodingForGlobals.cpp",
"VerifyAffinities.cpp",
"VerifyAsyncAccessRanges.cpp",
"VerifyLowerings.cpp",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
index 187675e..ca92382 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
@@ -49,6 +49,7 @@
"SpecializeDispatches.cpp"
"SpecializeEncodings.cpp"
"SyncInitializers.cpp"
+ "UnifyEncodingForGlobals.cpp"
"VerifyAffinities.cpp"
"VerifyAsyncAccessRanges.cpp"
"VerifyLowerings.cpp"
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
index b2846c6..8a17539 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
@@ -562,6 +562,42 @@
];
}
+def UnifyEncodingForGlobalsPass :
+ Pass<"iree-stream-unify-encoding-for-globals", "mlir::ModuleOp"> {
+ let summary = "Unifies multiple encodings of the same immutable global to reduce memory.";
+ let description = [{
+ When a single source global (e.g., a model weight) is encoded multiple
+ times with different encodings for different uses, this pass selects a
+ unified encoding and updates all references to reduce memory footprint.
+
+ The pass analyzes initializers to identify:
+ 1. Source globals (e.g., stream.tensor.constant with parameters)
+ 2. Encoded globals derived from the same source
+ 3. Dispatch sites using these encoded globals
+
+ When multiple encoded globals share the same source, the pass queries
+ layout resolvers to determine a unified encoding. If a common encoding
+ is found, the pass updates:
+ 1. The encoding operations in initializers
+ 2. The corresponding stream.tensor.sizeof operations
+ 3. The dispatch operand encodings
+ 4. The executable binding types
+
+ Constraints:
+ - Only handles immutable globals with !stream.resource<constant> type
+ - Bails out if an encoded global can have multiple sources due to CFG
+ - Requires exactly one AffinityAnalysisDialectInterface implementation
+ - Skips hand-authored executables to avoid incorrect modifications
+
+ This pass should run after CombineInitializersPass and before
+ SpecializeEncodingsPass in the Stream pipeline.
+ }];
+ let dependentDialects = [
+ "IREE::Encoding::IREEEncodingDialect",
+ "IREE::Stream::StreamDialect",
+ ];
+}
+
def SpecializeEncodingsPass :
Pass<"iree-stream-specialize-encodings", "mlir::ModuleOp"> {
let summary = "Specializes serializable encodings based on layout analysis.";
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/UnifyEncodingForGlobals.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/UnifyEncodingForGlobals.cpp
new file mode 100644
index 0000000..a756717
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/UnifyEncodingForGlobals.cpp
@@ -0,0 +1,333 @@
+// Copyright 2025 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/Encoding/IR/EncodingDialect.h"
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
+#include "iree/compiler/Dialect/Util/Analysis/GlobalTable.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/DebugLog.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir::iree_compiler::IREE::Stream {
+
+#define DEBUG_TYPE "iree-stream-unify-encoding-for-globals"
+
+#define GEN_PASS_DEF_UNIFYENCODINGFORGLOBALSPASS
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h.inc"
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Analysis.
+//===----------------------------------------------------------------------===//
+
+// Information about an encoded global and its relationship to its source.
+struct EncodedGlobalInfo {
+ // The destination global for the final encoded data.
+ IREE::Util::GlobalOpInterface encodedGlobal;
+ Attribute encodingAttr;
+ IREE::Stream::TensorSizeOfOp sizeofOp;
+ IREE::Stream::TensorEncodeOp encodeOp;
+};
+
+// Information about a source global and all its encoded versions.
+struct SourceGlobalInfo {
+ IREE::Util::GlobalOpInterface sourceGlobal;
+ SmallVector<EncodedGlobalInfo> encodedVersions;
+};
+
+// Returns the global store op that uses the given value, traversing through
+// passthrough ops like clone. Returns nullptr if not found or multiple users.
+static IREE::Util::GlobalStoreOpInterface findStoreOp(Value value) {
+ if (!llvm::hasSingleElement(value.getUsers())) {
+ return nullptr;
+ }
+ Operation *user = *value.getUsers().begin();
+ if (auto store = dyn_cast<IREE::Util::GlobalStoreOpInterface>(user)) {
+ return store;
+ }
+ if (auto cloneOp = dyn_cast<IREE::Stream::AsyncCloneOp>(user)) {
+ return findStoreOp(cloneOp.getResult());
+ }
+ return nullptr;
+}
+
+// Analyzes a module to find immutable globals that have multiple encoded
+// versions. Use run() to perform analysis, then query results with
+// getSourcesWithMultipleEncodings() or getSourceGlobals().
+class GlobalEncodingAnalyzer {
+public:
+ explicit GlobalEncodingAnalyzer(ModuleOp moduleOp)
+ : moduleOp(moduleOp), symbolTable(moduleOp), globalTable(moduleOp) {}
+
+ LogicalResult run();
+
+ // Returns all source globals that have multiple distinct encodings.
+ SmallVector<StringRef> getSourcesWithMultipleEncodings() const {
+ SmallVector<StringRef> result;
+ for (const auto &[name, info] : sourceGlobals) {
+ llvm::SmallDenseSet<Attribute, 4> uniqueEncodings;
+ for (const EncodedGlobalInfo &encoded : info.encodedVersions) {
+ uniqueEncodings.insert(encoded.encodingAttr);
+ }
+ if (uniqueEncodings.size() > 1) {
+ result.push_back(name);
+ }
+ }
+ return result;
+ }
+
+ // Returns the SourceGlobalInfo for the given source global name, or
+ // std::nullopt if not found.
+ std::optional<SourceGlobalInfo> getSourceGlobals(StringRef name) const {
+ if (sourceGlobals.contains(name)) {
+ return sourceGlobals.find(name)->second;
+ }
+ return std::nullopt;
+ }
+
+private:
+ // Walks all initializers to find encoding patterns and populates
+ // sourceGlobals map. Looks for patterns like:
+ // %source = util.global.load @source_global
+ // %encoded = stream.tensor.encode %source ...
+ // util.global.store %encoded, @encoded_global
+ // Only considers immutable source and encoded globals.
+ LogicalResult collectGlobalEncodings();
+
+ // Traces from encode op's source operand back to a source global.
+ // Returns nullptr if tracing fails or source is mutable.
+ IREE::Util::GlobalOpInterface traceToSourceGlobal(Value value);
+
+ // Returns true if the type has a recognized encoding attribute.
+ bool hasRecognizedEncoding(Type type) const;
+
+ ModuleOp moduleOp;
+ SymbolTable symbolTable;
+ IREE::Util::GlobalTable globalTable;
+
+ // Maps source global name to its info. Populated by run(). The global name
+ // must match the name of `sourceGlobal` inside SourceGlobalInfo. StringRef is
+ // used for easier lookup, which works better with SymbolTable, etc.
+ llvm::MapVector<StringRef, SourceGlobalInfo> sourceGlobals;
+};
+
+LogicalResult GlobalEncodingAnalyzer::run() {
+ LDBG() << "=== GlobalEncodingAnalyzer::run() ===";
+ globalTable.rebuild();
+ if (failed(collectGlobalEncodings())) {
+ return failure();
+ }
+ LDBG_OS([&](llvm::raw_ostream &os) {
+ os << "Analysis complete:\n";
+ os << " Source globals: " << sourceGlobals.size() << "\n";
+ for (const auto &[name, info] : sourceGlobals) {
+ os << " " << name << ": " << info.encodedVersions.size()
+ << " encoded versions\n";
+ }
+ });
+ return success();
+}
+
+bool GlobalEncodingAnalyzer::hasRecognizedEncoding(Type type) const {
+ auto rankedTensorType = dyn_cast<RankedTensorType>(type);
+ if (!rankedTensorType) {
+ return false;
+ }
+ auto encoding = rankedTensorType.getEncoding();
+ if (!encoding) {
+ return false;
+ }
+ return isa<IREE::Encoding::SerializableAttr>(encoding);
+}
+
+LogicalResult GlobalEncodingAnalyzer::collectGlobalEncodings() {
+ LDBG() << "--- collectGlobalEncodings ---";
+ // Walk all initializers to find encoding patterns:
+ // %source = util.global.load @source_global
+ // %encoded = stream.tensor.encode %source ...
+ // util.global.store %encoded, @encoded_global
+ for (auto initOp : moduleOp.getOps<IREE::Util::InitializerOp>()) {
+ initOp.walk([&](IREE::Stream::TensorEncodeOp encodeOp) {
+ LDBG() << " Found TensorEncodeOp: " << encodeOp;
+ Type resultType = encodeOp.getResultEncoding();
+ if (!hasRecognizedEncoding(resultType)) {
+ LDBG() << " Skipping: no recognized encoding";
+ return;
+ }
+
+ auto storeOp = findStoreOp(encodeOp.getResult());
+ if (!storeOp) {
+ LDBG() << " Skipping: no store found";
+ return;
+ }
+ StringRef encodedGlobalName = storeOp.getGlobalName();
+ auto encodedGlobalOp =
+ symbolTable.lookup<IREE::Util::GlobalOpInterface>(encodedGlobalName);
+ if (!encodedGlobalOp) {
+ LDBG() << " Skipping: encoded global not found";
+ return;
+ }
+ if (encodedGlobalOp.isGlobalMutable()) {
+ LDBG() << " Skipping: mutable encoded global";
+ return;
+ }
+ auto sourceGlobal = traceToSourceGlobal(encodeOp.getSource());
+ if (!sourceGlobal) {
+ LDBG() << " Skipping: failed to trace to source global";
+ return;
+ }
+
+ LDBG() << " Source: " << sourceGlobal.getGlobalName();
+ EncodedGlobalInfo encodedInfo;
+ encodedInfo.encodedGlobal = encodedGlobalOp;
+ encodedInfo.encodeOp = encodeOp;
+ encodedInfo.encodingAttr =
+ cast<RankedTensorType>(resultType).getEncoding();
+ encodedInfo.sizeofOp = encodeOp.getResultSize()
+ .getDefiningOp<IREE::Stream::TensorSizeOfOp>();
+ if (!encodedInfo.sizeofOp) {
+ LDBG() << " Skipping: no sizeof op found for result size, can't "
+ "update corresponding size computations later";
+ return;
+ }
+ SourceGlobalInfo &sourceInfo =
+ sourceGlobals[sourceGlobal.getGlobalName()];
+ if (!sourceInfo.sourceGlobal) {
+ sourceInfo.sourceGlobal = sourceGlobal;
+ }
+ sourceInfo.encodedVersions.push_back(encodedInfo);
+ });
+ }
+ return success();
+}
+
+IREE::Util::GlobalOpInterface
+GlobalEncodingAnalyzer::traceToSourceGlobal(Value value) {
+ IREE::Util::GlobalOpInterface foundSourceGlobal;
+ bool shouldContinue = true;
+ while (shouldContinue) {
+ Operation *defOp = value.getDefiningOp();
+ if (!defOp) {
+ LDBG() << " Bail: block argument";
+ return nullptr;
+ }
+ shouldContinue = false;
+ bool isValid =
+ llvm::TypeSwitch<Operation *, bool>(defOp)
+ .Case([&](IREE::Util::GlobalLoadOpInterface loadOp) {
+ StringRef globalName = loadOp.getGlobalName();
+ auto globalOp =
+ symbolTable.lookup<IREE::Util::GlobalOpInterface>(globalName);
+ if (!globalOp) {
+ LDBG() << " Bail: global not found: " << globalName;
+ return false;
+ }
+ if (globalOp.isGlobalMutable()) {
+ LDBG() << " Bail: mutable source global";
+ return false;
+ }
+ assert(!foundSourceGlobal &&
+ "should only find one source global");
+ foundSourceGlobal = globalOp;
+ // Has inline initial value, done tracing.
+ if (globalOp.getGlobalInitialValue()) {
+ return true;
+ }
+ auto &global = globalTable.lookup(globalName);
+ if (global.storeOps.size() != 1) {
+ LDBG() << " Bail: expected immutable global is only "
+ "initialized once";
+ return false;
+ }
+ // Continue tracing the source of the store op.
+ value = global.storeOps[0].getStoredGlobalValue();
+ shouldContinue = true;
+ return true;
+ })
+ .Case([&](IREE::Stream::TensorConstantOp) {
+ // Constant is a valid leaf source.
+ return true;
+ })
+ .Default([&](Operation *op) {
+ LDBG() << " Bail: unknown op " << op->getName();
+ return false;
+ });
+ if (!isValid) {
+ return nullptr;
+ }
+ }
+ return foundSourceGlobal;
+}
+
+//===----------------------------------------------------------------------===//
+// Pass Implementation
+//===----------------------------------------------------------------------===//
+
+struct UnifyEncodingForGlobalsPass
+ : public impl::UnifyEncodingForGlobalsPassBase<
+ UnifyEncodingForGlobalsPass> {
+ void runOnOperation() override {
+ ModuleOp moduleOp = getOperation();
+ GlobalEncodingAnalyzer analyzer(moduleOp);
+ if (failed(analyzer.run())) {
+ LDBG() << "Analysis failed, skipping.";
+ return;
+ }
+ auto candidates = analyzer.getSourcesWithMultipleEncodings();
+ if (candidates.empty()) {
+ LDBG() << "No source globals with multiple encodings found.";
+ return;
+ }
+ LDBG() << "Found " << candidates.size()
+ << " source globals with multiple encodings:";
+ for (auto name : candidates) {
+ LDBG() << " - " << name;
+ }
+
+ // Unify encodings for each source global with multiple encodings.
+ for (auto sourceName : candidates) {
+ std::optional<SourceGlobalInfo> sourceInfo =
+ analyzer.getSourceGlobals(sourceName);
+ if (!sourceInfo) {
+ LDBG() << " ERROR: source global info not found for " << sourceName;
+ continue;
+ }
+
+ // TODO(#22485): Select unified encoding via resolver. For now, use
+ // identity encoding.
+ auto unifiedEncoding =
+ IREE::Encoding::IdentityAttr::get(moduleOp.getContext());
+
+ // Update each encode op to use the unified encoding.
+ for (EncodedGlobalInfo &encodedInfo : sourceInfo->encodedVersions) {
+ auto encodeOp = encodedInfo.encodeOp;
+ auto oldResultType =
+ cast<RankedTensorType>(encodeOp.getResultEncoding());
+ RankedTensorType newResultType =
+ oldResultType.cloneWithEncoding(unifiedEncoding);
+ encodeOp.setResultEncodingAttr(TypeAttr::get(newResultType));
+ LDBG() << " Updated encode op: " << encodeOp;
+
+ encodedInfo.sizeofOp.setEncodingAttr(TypeAttr::get(newResultType));
+ LDBG() << " Updated sizeof op: " << encodedInfo.sizeofOp;
+ }
+ }
+
+ // TODO(#22485): Update tensor ops, dispatch sites and executables.
+ }
+};
+
+} // namespace
+
+} // namespace mlir::iree_compiler::IREE::Stream
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
index 1a7d36d..40f34e3 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
@@ -67,6 +67,7 @@
"specialize_dispatches.mlir",
"specialize_encodings.mlir",
"sync_initializers.mlir",
+ "unify_encoding_for_globals.mlir",
"verify_affinities.mlir",
"verify_async_access_ranges.mlir",
],
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
index cf32fcb..d6f192d 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
@@ -65,6 +65,7 @@
"specialize_dispatches.mlir"
"specialize_encodings.mlir"
"sync_initializers.mlir"
+ "unify_encoding_for_globals.mlir"
"verify_affinities.mlir"
"verify_async_access_ranges.mlir"
TOOLS
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/unify_encoding_for_globals.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/unify_encoding_for_globals.mlir
new file mode 100644
index 0000000..55316c6
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/unify_encoding_for_globals.mlir
@@ -0,0 +1,209 @@
+// RUN: iree-opt --split-input-file --pass-pipeline='builtin.module(iree-stream-unify-encoding-for-globals)' %s | FileCheck %s
+
+// Test: immutable source global (with initial value) with two encodings -
+// should unify to identity encoding.
+
+#encoding1 = #iree_encoding.testing<layouts = [#iree_encoding.specialized<123>]>
+#encoding2 = #iree_encoding.testing<layouts = [#iree_encoding.specialized<456>]>
+
+// CHECK-LABEL: module @immutable_source_with_initial_value
+module @immutable_source_with_initial_value {
+ util.global private @source = #stream.parameter.named<"model"::"weight"> : !stream.resource<constant>
+ util.global private @encoded_v1 : !stream.resource<constant>
+ util.global private @encoded_v2 : !stream.resource<constant>
+
+ // CHECK: util.initializer
+ util.initializer {
+ // CHECK: %[[SOURCE:.+]] = util.global.load @source
+ %source = util.global.load @source : !stream.resource<constant>
+ %source_size = stream.resource.size %source : !stream.resource<constant>
+
+ // CHECK: stream.tensor.sizeof tensor<4096x4096xf32, #iree_encoding.identity>
+ // CHECK: stream.tensor.encode %[[SOURCE]] : {{.*}} -> tensor<4096x4096xf32, #iree_encoding.identity>
+ %size1 = stream.tensor.sizeof tensor<4096x4096xf32, #encoding1> : index
+ %enc1 = stream.tensor.encode %source : tensor<4096x4096xf32> in !stream.resource<constant>{%source_size} -> tensor<4096x4096xf32, #encoding1> in !stream.resource<*>{%size1}
+ %const1 = stream.async.clone %enc1 : !stream.resource<*>{%size1} -> !stream.resource<constant>{%size1}
+ util.global.store %const1, @encoded_v1 : !stream.resource<constant>
+
+ // CHECK: stream.tensor.sizeof tensor<4096x4096xf32, #iree_encoding.identity>
+ // CHECK: stream.tensor.encode %[[SOURCE]] : {{.*}} -> tensor<4096x4096xf32, #iree_encoding.identity>
+ %size2 = stream.tensor.sizeof tensor<4096x4096xf32, #encoding2> : index
+ %enc2 = stream.tensor.encode %source : tensor<4096x4096xf32> in !stream.resource<constant>{%source_size} -> tensor<4096x4096xf32, #encoding2> in !stream.resource<*>{%size2}
+ %const2 = stream.async.clone %enc2 : !stream.resource<*>{%size2} -> !stream.resource<constant>{%size2}
+ util.global.store %const2, @encoded_v2 : !stream.resource<constant>
+
+ util.return
+ }
+}
+
+// -----
+
+// Test: immutable source global (initialized from parameter in initializer) with
+// two encodings - should unify to identity encoding.
+
+#encoding1 = #iree_encoding.testing<layouts = [#iree_encoding.specialized<123>]>
+#encoding2 = #iree_encoding.testing<layouts = [#iree_encoding.specialized<456>]>
+
+// CHECK-LABEL: module @immutable_source_initialized_from_parameter
+module @immutable_source_initialized_from_parameter {
+ util.global private @weight : !stream.resource<constant>
+ util.global private @weight_size : index
+ util.global private @encoded_v1 : !stream.resource<constant>
+ util.global private @encoded_v2 : !stream.resource<constant>
+
+ // CHECK: util.initializer
+ util.initializer {
+ %cst = stream.tensor.constant : tensor<4096x4096xf8E4M3FNUZ> in !stream.resource<constant> = #stream.parameter.named<"model"::"weight"> : tensor<4096x4096xf32>
+ %0 = stream.resource.size %cst : !stream.resource<constant>
+ util.global.store %cst, @weight : !stream.resource<constant>
+ util.global.store %0, @weight_size : index
+ // CHECK: %[[SOURCE:.+]] = util.global.load @weight
+ %source = util.global.load @weight : !stream.resource<constant>
+ %source_size = util.global.load @weight_size : index
+
+ // CHECK: stream.tensor.sizeof tensor<4096x4096xf32, #iree_encoding.identity>
+ // CHECK: stream.tensor.encode %[[SOURCE]] : {{.*}} -> tensor<4096x4096xf32, #iree_encoding.identity>
+ %size1 = stream.tensor.sizeof tensor<4096x4096xf32, #encoding1> : index
+ %enc1 = stream.tensor.encode %source : tensor<4096x4096xf32> in !stream.resource<constant>{%source_size} -> tensor<4096x4096xf32, #encoding1> in !stream.resource<*>{%size1}
+ %const1 = stream.async.clone %enc1 : !stream.resource<*>{%size1} -> !stream.resource<constant>{%size1}
+ util.global.store %const1, @encoded_v1 : !stream.resource<constant>
+
+ // CHECK: stream.tensor.sizeof tensor<4096x4096xf32, #iree_encoding.identity>
+ // CHECK: stream.tensor.encode %[[SOURCE]] : {{.*}} -> tensor<4096x4096xf32, #iree_encoding.identity>
+ %size2 = stream.tensor.sizeof tensor<4096x4096xf32, #encoding2> : index
+ %enc2 = stream.tensor.encode %source : tensor<4096x4096xf32> in !stream.resource<constant>{%source_size} -> tensor<4096x4096xf32, #encoding2> in !stream.resource<*>{%size2}
+ %const2 = stream.async.clone %enc2 : !stream.resource<*>{%size2} -> !stream.resource<constant>{%size2}
+ util.global.store %const2, @encoded_v2 : !stream.resource<constant>
+
+ util.return
+ }
+}
+
+// -----
+
+// Test: mutable source global - should be skipped, encoding unchanged.
+
+#encoding1 = #iree_encoding.testing<layouts = [#iree_encoding.specialized<123>]>
+
+// CHECK: #[[$ENC:.+]] = #iree_encoding.testing<layouts = [#iree_encoding.specialized<123>]>
+// CHECK-LABEL: module @mutable_source_skipped
+module @mutable_source_skipped {
+ util.global private mutable @mutable_source : !stream.resource<constant>
+ util.global private @encoded : !stream.resource<constant>
+
+ util.initializer {
+ %cst = stream.tensor.constant : tensor<4096x4096xf32> in !stream.resource<constant> = dense<0.0> : tensor<4096x4096xf32>
+ %cst_size = stream.resource.size %cst : !stream.resource<constant>
+ util.global.store %cst, @mutable_source : !stream.resource<constant>
+
+ %source = util.global.load @mutable_source : !stream.resource<constant>
+ %source_size = stream.resource.size %source : !stream.resource<constant>
+
+ // CHECK: stream.tensor.sizeof tensor<4096x4096xf32, #[[$ENC]]>
+ // CHECK: stream.tensor.encode {{.*}} -> tensor<4096x4096xf32, #[[$ENC]]>
+ %size1 = stream.tensor.sizeof tensor<4096x4096xf32, #encoding1> : index
+ %enc1 = stream.tensor.encode %source : tensor<4096x4096xf32> in !stream.resource<constant>{%source_size} -> tensor<4096x4096xf32, #encoding1> in !stream.resource<*>{%size1}
+ %const1 = stream.async.clone %enc1 : !stream.resource<*>{%size1} -> !stream.resource<constant>{%size1}
+ util.global.store %const1, @encoded : !stream.resource<constant>
+
+ util.return
+ }
+}
+
+// -----
+
+// Test: mutable encoded global - should be skipped, encoding unchanged.
+
+#encoding1 = #iree_encoding.testing<layouts = [#iree_encoding.specialized<123>]>
+#encoding2 = #iree_encoding.testing<layouts = [#iree_encoding.specialized<456>]>
+
+// CHECK: #[[$ENC1:.+]] = #iree_encoding.testing<layouts = [#iree_encoding.specialized<123>]>
+// CHECK: #[[$ENC2:.+]] = #iree_encoding.testing<layouts = [#iree_encoding.specialized<456>]>
+// CHECK-LABEL: module @mutable_encoded_global_skipped
+module @mutable_encoded_global_skipped {
+ util.global private @source = #stream.parameter.named<"model"::"weight"> : !stream.resource<constant>
+ util.global private mutable @encoded_mutable_v1 : !stream.resource<constant>
+ util.global private mutable @encoded_mutable_v2 : !stream.resource<constant>
+
+ util.initializer {
+ %source = util.global.load @source : !stream.resource<constant>
+ %source_size = stream.resource.size %source : !stream.resource<constant>
+
+ // CHECK: stream.tensor.sizeof tensor<4096x4096xf32, #[[$ENC1]]>
+ // CHECK: stream.tensor.encode {{.*}} -> tensor<4096x4096xf32, #[[$ENC1]]>
+ %size1 = stream.tensor.sizeof tensor<4096x4096xf32, #encoding1> : index
+ %enc1 = stream.tensor.encode %source : tensor<4096x4096xf32> in !stream.resource<constant>{%source_size} -> tensor<4096x4096xf32, #encoding1> in !stream.resource<*>{%size1}
+ %const1 = stream.async.clone %enc1 : !stream.resource<*>{%size1} -> !stream.resource<constant>{%size1}
+ util.global.store %const1, @encoded_mutable_v1 : !stream.resource<constant>
+
+ // CHECK: stream.tensor.sizeof tensor<4096x4096xf32, #[[$ENC2]]>
+ // CHECK: stream.tensor.encode {{.*}} -> tensor<4096x4096xf32, #[[$ENC2]]>
+ %size2 = stream.tensor.sizeof tensor<4096x4096xf32, #encoding2> : index
+ %enc2 = stream.tensor.encode %source : tensor<4096x4096xf32> in !stream.resource<constant>{%source_size} -> tensor<4096x4096xf32, #encoding2> in !stream.resource<*>{%size2}
+ %const2 = stream.async.clone %enc2 : !stream.resource<*>{%size2} -> !stream.resource<constant>{%size2}
+ util.global.store %const2, @encoded_mutable_v2 : !stream.resource<constant>
+ util.return
+ }
+}
+
+// -----
+
+// Test: single encoding - not a candidate for unification, encoding unchanged.
+
+#encoding1 = #iree_encoding.testing<layouts = [#iree_encoding.specialized<123>]>
+
+// CHECK: #[[$ENC:.+]] = #iree_encoding.testing<layouts = [#iree_encoding.specialized<123>]>
+// CHECK-LABEL: module @single_encoding_no_unification
+module @single_encoding_no_unification {
+ util.global private @source = #stream.parameter.named<"model"::"weight"> : !stream.resource<constant>
+ util.global private @encoded : !stream.resource<constant>
+
+ util.initializer {
+ %source = util.global.load @source : !stream.resource<constant>
+ %source_size = stream.resource.size %source : !stream.resource<constant>
+
+ // CHECK: stream.tensor.sizeof tensor<4096x4096xf32, #[[$ENC]]>
+ // CHECK: stream.tensor.encode {{.*}} -> tensor<4096x4096xf32, #[[$ENC]]>
+ %size1 = stream.tensor.sizeof tensor<4096x4096xf32, #encoding1> : index
+ %enc1 = stream.tensor.encode %source : tensor<4096x4096xf32> in !stream.resource<constant>{%source_size} -> tensor<4096x4096xf32, #encoding1> in !stream.resource<*>{%size1}
+ %const1 = stream.async.clone %enc1 : !stream.resource<*>{%size1} -> !stream.resource<constant>{%size1}
+ util.global.store %const1, @encoded : !stream.resource<constant>
+
+ util.return
+ }
+}
+
+// -----
+
+// Test: same encoding used twice - not a candidate (only one unique encoding).
+
+#encoding1 = #iree_encoding.testing<layouts = [#iree_encoding.specialized<123>]>
+
+// CHECK: #[[$ENC:.+]] = #iree_encoding.testing<layouts = [#iree_encoding.specialized<123>]>
+// CHECK-LABEL: module @same_encoding_twice_no_unification
+module @same_encoding_twice_no_unification {
+ util.global private @source = #stream.parameter.named<"model"::"weight"> : !stream.resource<constant>
+ util.global private @encoded_v1 : !stream.resource<constant>
+ util.global private @encoded_v2 : !stream.resource<constant>
+
+ util.initializer {
+ %source = util.global.load @source : !stream.resource<constant>
+ %source_size = stream.resource.size %source : !stream.resource<constant>
+
+ // CHECK: stream.tensor.sizeof tensor<4096x4096xf32, #[[$ENC]]>
+ // CHECK: stream.tensor.encode {{.*}} -> tensor<4096x4096xf32, #[[$ENC]]>
+ %size1 = stream.tensor.sizeof tensor<4096x4096xf32, #encoding1> : index
+ %enc1 = stream.tensor.encode %source : tensor<4096x4096xf32> in !stream.resource<constant>{%source_size} -> tensor<4096x4096xf32, #encoding1> in !stream.resource<*>{%size1}
+ %const1 = stream.async.clone %enc1 : !stream.resource<*>{%size1} -> !stream.resource<constant>{%size1}
+ util.global.store %const1, @encoded_v1 : !stream.resource<constant>
+
+ // CHECK: stream.tensor.sizeof tensor<4096x4096xf32, #[[$ENC]]>
+ // CHECK: stream.tensor.encode {{.*}} -> tensor<4096x4096xf32, #[[$ENC]]>
+ %size2 = stream.tensor.sizeof tensor<4096x4096xf32, #encoding1> : index
+ %enc2 = stream.tensor.encode %source : tensor<4096x4096xf32> in !stream.resource<constant>{%source_size} -> tensor<4096x4096xf32, #encoding1> in !stream.resource<*>{%size2}
+ %const2 = stream.async.clone %enc2 : !stream.resource<*>{%size2} -> !stream.resource<constant>{%size2}
+ util.global.store %const2, @encoded_v2 : !stream.resource<constant>
+
+ util.return
+ }
+}
