Add `iree-stream-resource-alias-mutable-bindings` flag. (#13965)
This pass option is currently necessary to run some WebGPU programs
correctly, so this adds a compiler flag to set it. Without the option, I
see errors like
```
Writable storage buffer binding aliasing found between [BindGroup] set at bind group index 0, binding index 0, and [BindGroup] set at bind group index 0, binding index 2, with overlapping ranges (offset: 0, size: 256) and (offset: 0, size: 256) in [Buffer].
- While encoding [ComputePassEncoder].DispatchWorkgroups(1, 5, 1).
[Invalid CommandBuffer] is invalid.
at ValidateObject (../../third_party/dawn/src/dawn/native/Device.cpp:689)
at ValidateSubmit (../../third_party/dawn/src/dawn/native/Queue.cpp:442)
```
We can't enable the option by default yet, since it causes some severe
performance regressions on CUDA:
https://github.com/openxla/iree/pull/13323. The tests also have a few
suspicious TODOs, but I want to unblock correctness testing on WebGPU
before digging in deeper.diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamBase.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamBase.td
index 77d3427..1288878 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamBase.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamBase.td
@@ -319,7 +319,9 @@
// bindings.
"int64_t":$minBufferRangeAlignment,
// Number of bits in `index` values as passed across device boundaries.
- "int64_t":$indexBits
+ "int64_t":$indexBits,
+ // Fuses bindings that are mutable instead of leaving them split.
+ "bool":$aliasMutableBindings
);
let valueType = NoneType;
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp
index 49baca2..8a922ca 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp
@@ -63,6 +63,11 @@
"iree-stream-resource-index-bits",
llvm::cl::desc("Bit width of indices used to reference resource offsets."),
llvm::cl::init(32));
+static llvm::cl::opt<bool> clResourceAliasMutableBindings(
+ "iree-stream-resource-alias-mutable-bindings",
+ llvm::cl::desc(
+ "Fuses bindings that are mutable instead of leaving them split."),
+ llvm::cl::init(false));
//===----------------------------------------------------------------------===//
// #stream.resource_config<...>
@@ -77,6 +82,7 @@
int64_t maxBufferRange = 0;
int64_t minBufferRangeAlignment = 0;
int64_t indexBits = 32;
+ bool aliasMutableBindings = false;
while (failed(p.parseOptionalRBrace())) {
StringRef key;
int64_t value = 0;
@@ -94,14 +100,16 @@
minBufferRangeAlignment = value;
} else if (key == "index_bits") {
indexBits = value;
+ } else if (key == "alias_mutable_bindings") {
+ aliasMutableBindings = (bool)value;
}
(void)p.parseOptionalComma();
}
if (failed(p.parseGreater())) return {};
- return ResourceConfigAttr::get(p.getContext(), maxAllocationSize,
- minBufferOffsetAlignment, maxBufferRange,
- minBufferRangeAlignment, indexBits);
+ return ResourceConfigAttr::get(
+ p.getContext(), maxAllocationSize, minBufferOffsetAlignment,
+ maxBufferRange, minBufferRangeAlignment, indexBits, aliasMutableBindings);
}
void ResourceConfigAttr::print(AsmPrinter &p) const {
@@ -112,7 +120,8 @@
<< ", ";
os << "max_buffer_range = " << getMaxBufferRange() << ", ";
os << "min_buffer_range_alignment = " << getMinBufferRangeAlignment() << ", ";
- os << "index_bits = " << getIndexBits();
+ os << "index_bits = " << getIndexBits() << ", ";
+ os << "alias_mutable_bindings = " << getAliasMutableBindings();
os << "}>";
}
@@ -130,7 +139,8 @@
std::min(lhs.getMaxBufferRange(), rhs.getMaxBufferRange()),
std::max(lhs.getMinBufferRangeAlignment(),
rhs.getMinBufferRangeAlignment()),
- std::max(lhs.getIndexBits(), rhs.getIndexBits()));
+ std::max(lhs.getIndexBits(), rhs.getIndexBits()),
+ rhs.getAliasMutableBindings() && lhs.getAliasMutableBindings());
}
// static
@@ -141,7 +151,8 @@
// only use this during testing by ensuring affinities are always assigned.
return ResourceConfigAttr::get(
context, clResourceMaxAllocationSize, clResourceMinOffsetAlignment,
- clResourceMaxRange, clResourceMinOffsetAlignment, clResourceIndexBits);
+ clResourceMaxRange, clResourceMinOffsetAlignment, clResourceIndexBits,
+ clResourceAliasMutableBindings);
}
// static
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/FuseDispatchBindings.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/FuseDispatchBindings.cpp
index 017e607..4b90e35 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/FuseDispatchBindings.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/FuseDispatchBindings.cpp
@@ -315,11 +315,14 @@
IREE::Stream::ExecutableOp executableOp,
IREE::Stream::ExecutableExportOp exportOp,
ArrayRef<IREE::Stream::CmdDispatchOp> dispatchOps,
- bool aliasMutableBindings, MemoizedCmdZeros &memoizedZeros) {
+ MemoizedCmdZeros &memoizedZeros) {
if (dispatchOps.empty()) return; // no-op if no dispatches
auto anyDispatchOp = dispatchOps.front();
unsigned bindingCount = anyDispatchOp.getResources().size();
+ auto configAttr = IREE::Stream::ResourceConfigAttr::lookup(exportOp);
+ bool aliasMutableBindings = configAttr.getAliasMutableBindings();
+
LLVM_DEBUG({
AsmState asmState(executableOp->getParentOp());
llvm::dbgs() << "---- fuseDispatchBindings(@" << executableOp.getSymName()
@@ -448,7 +451,7 @@
for (auto exportOp :
executableOp.getOps<IREE::Stream::ExecutableExportOp>()) {
fuseDispatchBindings(executableOp, exportOp, entryDispatchMap[exportOp],
- aliasMutableBindings, memoizedZeros);
+ memoizedZeros);
}
}
}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
index 788771e..ed72ab0 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
@@ -190,11 +190,6 @@
let constructor = [{
mlir::iree_compiler::IREE::Stream::createFuseDispatchBindingsPass()
}];
- let options = [
- Option<"aliasMutableBindings", "alias-mutable-bindings",
- "bool", /*default=*/"false",
- "Fuses bindings that are mutable instead of leaving them split.">
- ];
}
def SpecializeDispatches :
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/fuse_dispatch_bindings.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/fuse_dispatch_bindings.mlir
index 764c788..198ad3f 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/fuse_dispatch_bindings.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/fuse_dispatch_bindings.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt --split-input-file --pass-pipeline='builtin.module(iree-stream-fuse-dispatch-bindings{alias-mutable-bindings=true})' %s | FileCheck %s
+// RUN: iree-opt --split-input-file --pass-pipeline='builtin.module(iree-stream-fuse-dispatch-bindings)' %s | FileCheck %s
// Test that bindings that are unique are rebased to the widest possible access
// and to have a 0 offset by passing in the actual offset as operands. The
@@ -8,9 +8,13 @@
// TODO(benvanik): do this in a canonicalize bindings pass. This should not be
// happening here!
+#aliasConfig = #stream.resource_config<{
+ alias_mutable_bindings = true
+}>
+
// CHECK-LABEL: @rebaseBindingsEx
stream.executable private @rebaseBindingsEx {
- stream.executable.export public @dispatch
+ stream.executable.export public @dispatch attributes {stream.resources = #aliasConfig}
builtin.module {
// CHECK: func.func @dispatch(%[[BINDING_A:.+]]: !stream.binding, %[[BINDING_B:.+]]: !stream.binding,
// CHECK-SAME: %[[OFFSET_A:.+]]: index, %[[OFFSET_B:.+]]: index, %[[OPERAND:.+]]: index)
@@ -85,9 +89,13 @@
// in the order of the deduplicated bindings instead of the original order.
// This is a bit weird and it would be nice to preserve the order in the future.
+#aliasConfig = #stream.resource_config<{
+ alias_mutable_bindings = true
+}>
+
// CHECK-LABEL: @deduplicateBindingsEx
stream.executable private @deduplicateBindingsEx {
- stream.executable.export public @dispatch
+ stream.executable.export public @dispatch attributes {stream.resources = #aliasConfig}
builtin.module {
// CHECK: func.func @dispatch(%[[BINDING_A:.+]]: !stream.binding, %[[BINDING_B:.+]]: !stream.binding,
// CHECK-SAME: %[[OFFSET_A:.+]]: index, %[[OFFSET_C:.+]]: index, %[[OFFSET_B:.+]]: index, %[[OPERAND:.+]]: index)
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/fuse_dispatch_bindings_noalias.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/fuse_dispatch_bindings_noalias.mlir
index da1abe6..686819f 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/fuse_dispatch_bindings_noalias.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/fuse_dispatch_bindings_noalias.mlir
@@ -1,13 +1,17 @@
-// RUN: iree-opt --split-input-file --pass-pipeline='builtin.module(iree-stream-fuse-dispatch-bindings{alias-mutable-bindings=false})' %s | FileCheck %s
+// RUN: iree-opt --split-input-file --pass-pipeline='builtin.module(iree-stream-fuse-dispatch-bindings)' %s | FileCheck %s
// TODO(benvanik): remove this file when aliasing mutable bindings is fixed.
// Tests that bindings that are duplicated at all dispatch sites are folded
// so long as they are not mutable.
+#noaliasConfig = #stream.resource_config<{
+ alias_mutable_bindings = false
+}>
+
// CHECK-LABEL: @deduplicateBindingsEx
stream.executable private @deduplicateBindingsEx {
- stream.executable.export public @dispatch
+ stream.executable.export public @dispatch attributes {stream.resources = #noaliasConfig}
builtin.module {
// CHECK: func.func @dispatch(%[[BINDING_A:.+]]: !stream.binding, %[[BINDING_C:.+]]: !stream.binding,
// CHECK-SAME: %[[OFFSET_A:.+]]: index, %[[OFFSET_B:.+]]: index, %[[OFFSET_C:.+]]: index, %[[OPERAND:.+]]: index)
diff --git a/experimental/web/sample_webgpu/build_sample.sh b/experimental/web/sample_webgpu/build_sample.sh
index 9360374..58dafa3 100755
--- a/experimental/web/sample_webgpu/build_sample.sh
+++ b/experimental/web/sample_webgpu/build_sample.sh
@@ -53,6 +53,7 @@
--iree-input-type=$2 \
--iree-hal-target-backends=webgpu \
--iree-codegen-gpu-native-math-precision=true \
+ --iree-stream-resource-alias-mutable-bindings=true \
--o ${BINARY_DIR}/$1_webgpu.vmfb
}
diff --git a/experimental/web/sample_webgpu/index.html b/experimental/web/sample_webgpu/index.html
index 1196e72..949d98b 100644
--- a/experimental/web/sample_webgpu/index.html
+++ b/experimental/web/sample_webgpu/index.html
@@ -177,7 +177,8 @@
<textarea type="text" readonly spellcheck="false"
class="form-control" style="width:610px; height:90px; resize:none; font-family: monospace;">
--iree-hal-target-backends=webgpu \
---iree-codegen-gpu-native-math-precision=true \</textarea>
+--iree-codegen-gpu-native-math-precision=true \
+--iree-stream-resource-alias-mutable-bindings=true \</textarea>
</div>
