[Stream] Always map aliases to resolved resource ranges of the source op (#25026) This ensures that all operands in the tied operand chain for a resource will get mapped to the correct buffer range. Previously, if the resource's buffer received a new allocation, `stream.resource.subview` aliases (by nature of a byte subview, tied to a different memory range) would still receive resource ranges relatively to the old resource. Additionally, all aliases of a given source would be processed from the "nearest" subview in the alias chain. With multiple subviews in the same tied operand chain, the offsets weren't computed relative to one another. Addressing the first problem enables the most practical use case. SplitParameterEncoding in "consolidate" mode may need to place an unmodified input weight into the output buffers, which means the weight may end up in a packed allocation together with a pre-initialized output parameter. To later fetch that input weight for dynamic initialization, an explicit `subview` op is generated (*). Hoisting the constant out of the execution scope establishes the new resource range, and the subview should compute its resource range against that. Instead, the old subrange would get used that was relative to the original SSA source bound for deletion. See `Stream/Transforms/test/schedule_allocation_subviews_constant.mlir` for illustration. (*) Whether that should be happening is a different matter. Accessing the packed weight through `stream.async.slice` would've been preferable, as per a [TODO in SplitParameterEncoder](https://github.com/iree-org/iree/blob/0df50f9a8e40ab1878479124b952f45d9f8c0650/compiler/src/iree/compiler/Dialect/Stream/Transforms/SplitParameterEncoder.cpp#L1971). Generally, we should let ScheduleAllocation operate on that abstraction where a copy will implicitly be generated on writing to the range, rather than have it account for explicit aliases in that of `resource.subview`. For now I've preferred not to go into that direction, because that switch requires larger-scale testing, and currently explicit subview ops are expected to be supported in the input IR to the allocator. By implementing alias mapping as a DFS worklist and simply ignoring values that have already been resolved to a resource range, nested subviews, a.k.a. subview ops that reference each other, now correctly calculate their offset by adding up their own offset to the earlier offsets in the tied operand chain. Not sure if that does/would come up in a real case, but at least that's a sensible side effect. Assisted-by: Codex --------- Signed-off-by: Artem Gindinson <gindinson@roofline.ai>
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp index b206f27..9637a06 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp
@@ -426,24 +426,7 @@ }); // Propagate alias subranges when present. - for (auto alias : valueAliases.getValueAliases(resource)) { - ResourceRange aliasRange = resourceRange; - if (auto subviewOp = - IREE::Stream::ResourceSubviewOp::findSubviewOp(alias)) { - aliasRange.resource = subviewOp.getSubrangeResource(); - aliasRange.resourceSize = subviewOp.getSubrangeResourceSize(); - aliasRange.offset = subviewOp.getSubrangeOffset(); - aliasRange.length = subviewOp.getSubrangeLength(); - } - resourceRangeMap.insert(std::make_pair(alias, aliasRange)); - LLVM_DEBUG({ - llvm::dbgs() << " = alias "; - alias.printAsOperand(llvm::dbgs(), *asmState); - llvm::dbgs() << " = "; - aliasRange.print(llvm::dbgs(), *asmState); - llvm::dbgs() << "\n"; - }); - } + mapResourceAliases(resource, resourceRange, asmState); } // Returns a storage range backing the given stream |resource|. @@ -486,6 +469,66 @@ } private: + // Map tied operands before their results so each alias inherits its source's + // allocated buffer and offset. + void mapResourceAliases(Value resource, const ResourceRange &baseRange, + AsmState *asmState) { + SmallVector<Value> worklist; + for (auto alias : valueAliases.getValueAliases(resource)) { + worklist.push_back(alias); + } + + while (!worklist.empty()) { + Value alias = worklist.back(); + if (resourceRangeMap.contains(alias)) { + worklist.pop_back(); + continue; + } + + ResourceRange aliasRange = baseRange; + if (auto tiedOp = dyn_cast_if_present<IREE::Util::TiedOpInterface>( + alias.getDefiningOp())) { + // If the source is not mapped yet, leave this alias pending and tend + // to the source in case it aliases a subview into an already-resolved + // range. + if (Value source = tiedOp.getTiedResultOperand(alias)) { + auto it = resourceRangeMap.find(source); + if (it == resourceRangeMap.end()) { + worklist.push_back(source); + continue; + } + aliasRange = it->second; + } + } + + mapResourceAlias(alias, aliasRange, asmState); + worklist.pop_back(); + } + } + + // Map one alias using its source's allocated range. + void mapResourceAlias(Value alias, ResourceRange aliasRange, + AsmState *asmState) { + if (auto subviewOp = + alias.getDefiningOp<IREE::Stream::ResourceSubviewOp>()) { + // Keep the resolved buffer and its total size for the event that the + // original SSA value is not the final allocated buffer. Only adjust the + // view's range. + aliasRange.offset = add(subviewOp.getLoc(), aliasRange.offset, + subviewOp.getSubrangeOffset()); + aliasRange.length = subviewOp.getSubrangeLength(); + } + + resourceRangeMap.insert(std::make_pair(alias, aliasRange)); + LLVM_DEBUG({ + llvm::dbgs() << " = alias "; + alias.printAsOperand(llvm::dbgs(), *asmState); + llvm::dbgs() << " = "; + aliasRange.print(llvm::dbgs(), *asmState); + llvm::dbgs() << "\n"; + }); + } + Operation *rootOp; // Disjoint-set of values that alias each other. That two things alias does
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 6035df8..c22b60a 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
@@ -66,6 +66,9 @@ "reuse_allocations.mlir", "schedule_allocation.mlir", "schedule_allocation_parameters.mlir", + "schedule_allocation_subviews_constant.mlir", + "schedule_allocation_subviews_nested.mlir", + "schedule_allocation_subviews_transient.mlir", "schedule_concurrency.mlir", "schedule_execution.mlir", "schedule_execution_invalid.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 9db4899..8f4ba25 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
@@ -63,6 +63,9 @@ "reuse_allocations.mlir" "schedule_allocation.mlir" "schedule_allocation_parameters.mlir" + "schedule_allocation_subviews_constant.mlir" + "schedule_allocation_subviews_nested.mlir" + "schedule_allocation_subviews_transient.mlir" "schedule_concurrency.mlir" "schedule_execution.mlir" "schedule_execution_invalid.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation.mlir index ab20e41..b2194fe 100644 --- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation.mlir +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation.mlir
@@ -285,8 +285,8 @@ // CHECK-SAME: %[[SIZE:.+]]: index, %[[OFFSET0:.+]]: index, %[[OFFSET1:.+]]: index, %[[OFFSET2:.+]]: index, %[[LENGTH0:.+]]: index, %[[LENGTH1:.+]]: index, %[[LENGTH2:.+]]: index) util.func public @tiedOperandSubviews(%operand: !stream.resource<external>, %size: index, %offset0: index, %offset1: index, %offset2: index, %length0: index, %length1: index, %length2: index) { %c0 = arith.constant 0 : index - // CHECK: %[[SUBVIEW_OFFSET:.+]] = arith.addi %[[OFFSET0]], %[[OFFSET1]] - // CHECK: %[[SUBVIEW:.+]] = stream.resource.subview %[[OPERAND]][%[[SUBVIEW_OFFSET]]] {{.*}} -> !stream.resource<external>{%[[LENGTH1]]} + // CHECK-DAG: %[[SUBVIEW:.+]] = stream.resource.subview %[[OPERAND]][%[[SUBVIEW_OFFSET:.+]]] {{.*}} -> !stream.resource<external>{%[[LENGTH1]]} + // CHECK-DAG: %[[SUBVIEW_OFFSET]] = arith.addi %[[OFFSET0]], %[[OFFSET1]] // CHECK: stream.cmd.execute with(%[[OPERAND]] as %[[OPERAND_CAPTURE:.+]]: !stream.resource<external>{%[[SIZE]]}) // CHECK-NEXT: } => !stream.timepoint %result0, %result0_timepoint = stream.async.execute with(%operand as %capture: !stream.resource<external>{%size}) -> (%operand as !stream.resource<external>{%length1}) {
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation_subviews_constant.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation_subviews_constant.mlir new file mode 100644 index 0000000..a448547 --- /dev/null +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation_subviews_constant.mlir
@@ -0,0 +1,32 @@ +// RUN: iree-opt --iree-stream-schedule-allocation %s | FileCheck %s --implicit-check-not=stream.async. --implicit-check-not=stream.resource.subview + +// An extracted constant can both escape the execution region and supply a +// subview to a copy within it. The copy uses the captured constant buffer. + +// CHECK-LABEL: util.func public @constantSubview( +// CHECK-SAME: %[[TARGET:.+]]: !stream.resource<external>) +util.func public @constantSubview(%target: !stream.resource<external>) { + // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index + // CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index + // CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index + // CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index + %c0 = arith.constant 0 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c64 = arith.constant 64 : index + // CHECK: %[[CONSTANT:.+]], %[[READY:.+]] = stream.resource.constants : + // CHECK-NEXT: !stream.resource<constant>{%[[C64]]} = #stream.parameter.named<"weights"::"table"> : tensor<64xi8> + // CHECK: stream.cmd.execute await(%[[READY]]) => with(%[[TARGET]] as %[[DEST:.+]]: !stream.resource<external>{%[[C32]]}, %[[CONSTANT]] as %[[SOURCE:.+]]: !stream.resource<constant>{%[[C64]]}) { + // CHECK-NEXT: stream.cmd.copy %[[SOURCE]][%[[C16]]], %[[DEST]][%[[C0]]], %[[C32]] : !stream.resource<constant>{%[[C64]]} -> !stream.resource<external>{%[[C32]]} + %results:2, %timepoint = stream.async.execute with(%target as %target_arg: !stream.resource<external>{%c32}) -> (!stream.resource<constant>{%c64}, %target as !stream.resource<external>{%c32}) { + %constant = stream.async.constant : !stream.resource<constant>{%c64} = + #stream.parameter.named<"weights"::"table"> : tensor<64xi8> + %view = stream.resource.subview %constant[%c16] : !stream.resource<constant>{%c64} -> !stream.resource<constant>{%c32} + %updated = stream.async.update %view, %target_arg[%c0 to %c32] : !stream.resource<constant>{%c32} -> %target_arg as !stream.resource<external>{%c32} + stream.yield %constant, %updated : !stream.resource<constant>{%c64}, !stream.resource<external>{%c32} + } => !stream.timepoint + // CHECK: util.optimization_barrier %[[CONSTANT]], %[[TARGET]], + util.optimization_barrier %results#0, %results#1, %timepoint : !stream.resource<constant>, !stream.resource<external>, !stream.timepoint + // CHECK: util.return + util.return +}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation_subviews_nested.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation_subviews_nested.mlir new file mode 100644 index 0000000..56c8a16 --- /dev/null +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation_subviews_nested.mlir
@@ -0,0 +1,28 @@ +// RUN: iree-opt --iree-stream-schedule-allocation %s | FileCheck %s --implicit-check-not=stream.async. --implicit-check-not=stream.resource.subview + +// Nested views compose their offsets relative to the captured source buffer. + +// CHECK-LABEL: util.func public @nestedSubviews( +// CHECK-SAME: %[[SOURCE:.+]]: !stream.resource<external>, %[[TARGET:.+]]: !stream.resource<external>, %[[OUTER_OFFSET:.+]]: index, %[[INNER_OFFSET:.+]]: index) +util.func public @nestedSubviews(%source: !stream.resource<external>, %target: !stream.resource<external>, %outer_offset: index, %inner_offset: index) { + // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index + // CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index + // CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index + %c0 = arith.constant 0 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 64 : index + %c128 = arith.constant 128 : index + // CHECK: %[[OFFSET:.+]] = arith.addi %[[OUTER_OFFSET]], %[[INNER_OFFSET]] : index + // CHECK: stream.cmd.execute with(%[[SOURCE]] as %[[SOURCE_ARG:.+]]: !stream.resource<external>{%[[C128]]}, %[[TARGET]] as %[[TARGET_ARG:.+]]: !stream.resource<external>{%[[C16]]}) { + // CHECK-NEXT: stream.cmd.copy %[[SOURCE_ARG]][%[[OFFSET]]], %[[TARGET_ARG]][%[[C0]]], %[[C16]] : !stream.resource<external>{%[[C128]]} -> !stream.resource<external>{%[[C16]]} + %result, %timepoint = stream.async.execute with(%source as %source_arg: !stream.resource<external>{%c128}, %target as %target_arg: !stream.resource<external>{%c16}) -> %target as !stream.resource<external>{%c16} { + %outer = stream.resource.subview %source_arg[%outer_offset] : !stream.resource<external>{%c128} -> !stream.resource<external>{%c64} + %inner = stream.resource.subview %outer[%inner_offset] : !stream.resource<external>{%c64} -> !stream.resource<external>{%c16} + %updated = stream.async.update %inner, %target_arg[%c0 to %c16] : !stream.resource<external>{%c16} -> %target_arg as !stream.resource<external>{%c16} + stream.yield %updated : !stream.resource<external>{%c16} + } => !stream.timepoint + // CHECK: util.optimization_barrier %[[TARGET]], + util.optimization_barrier %result, %timepoint : !stream.resource<external>, !stream.timepoint + // CHECK: util.return + util.return +}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation_subviews_transient.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation_subviews_transient.mlir new file mode 100644 index 0000000..c1b225f --- /dev/null +++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/schedule_allocation_subviews_transient.mlir
@@ -0,0 +1,47 @@ +// RUN: iree-opt --iree-stream-schedule-allocation %s | FileCheck %s --implicit-check-not=stream.async. --implicit-check-not=stream.resource.subview + +// Overlapping temporary buffers share an allocation with distinct packed +// offsets. Views add their local offsets to their own buffer's packed offset. + +// CHECK-LABEL: util.func public @transientSubviews( +// CHECK-SAME: %[[TARGET:.+]]: !stream.resource<external>) +util.func public @transientSubviews(%target: !stream.resource<external>) { + // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index + // CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index + // CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index + // CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index + // CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index + // CHECK-DAG: %[[ONE:.+]] = arith.constant 1 : i8 + // CHECK-DAG: %[[TWO:.+]] = arith.constant 2 : i8 + %c0 = arith.constant 0 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c64 = arith.constant 64 : index + %c128 = arith.constant 128 : index + %c1_i8 = arith.constant 1 : i8 + %c2_i8 = arith.constant 2 : i8 + // CHECK: %[[PACK:.+]]:3 = stream.resource.pack slices({ + // CHECK-NEXT: [{{.+}}] = %[[C128]], + // CHECK-NEXT: [{{.+}}] = %[[C128]] + // CHECK: %[[SLAB:.+]], %[[READY:.+]] = stream.resource.alloca {{.*}}!stream.resource<transient>{%[[PACK]]#0} + // CHECK-DAG: %[[FIRST_OFFSET:.+]] = arith.addi %[[PACK]]#1, %[[C32]] : index + // CHECK-DAG: %[[SECOND_OFFSET:.+]] = arith.addi %[[PACK]]#2, %[[C16]] : index + // CHECK: stream.cmd.execute await(%[[READY]]) => with(%[[TARGET]] as %[[DEST:.+]]: !stream.resource<external>{%[[C64]]}, %[[SLAB]] as %[[SOURCE:.+]]: !stream.resource<transient>{%[[PACK]]#0}) { + // CHECK-NEXT: stream.cmd.fill %[[ONE]], %[[SOURCE]][%[[PACK]]#1 for %[[C128]]] : i8 -> !stream.resource<transient>{%[[PACK]]#0} + // CHECK-NEXT: stream.cmd.fill %[[TWO]], %[[SOURCE]][%[[PACK]]#2 for %[[C128]]] : i8 -> !stream.resource<transient>{%[[PACK]]#0} + // CHECK-NEXT: stream.cmd.copy %[[SOURCE]][%[[FIRST_OFFSET]]], %[[DEST]][%[[C0]]], %[[C32]] : !stream.resource<transient>{%[[PACK]]#0} -> !stream.resource<external>{%[[C64]]} + // CHECK-NEXT: stream.cmd.copy %[[SOURCE]][%[[SECOND_OFFSET]]], %[[DEST]][%[[C32]]], %[[C32]] : !stream.resource<transient>{%[[PACK]]#0} -> !stream.resource<external>{%[[C64]]} + %result, %timepoint = stream.async.execute with(%target as %target_arg: !stream.resource<external>{%c64}) -> %target as !stream.resource<external>{%c64} { + %first = stream.async.splat %c1_i8 : i8 -> !stream.resource<transient>{%c128} + %second = stream.async.splat %c2_i8 : i8 -> !stream.resource<transient>{%c128} + %first_view = stream.resource.subview %first[%c32] : !stream.resource<transient>{%c128} -> !stream.resource<transient>{%c32} + %second_view = stream.resource.subview %second[%c16] : !stream.resource<transient>{%c128} -> !stream.resource<transient>{%c32} + %updated = stream.async.update %first_view, %target_arg[%c0 to %c32] : !stream.resource<transient>{%c32} -> %target_arg as !stream.resource<external>{%c64} + %result = stream.async.update %second_view, %updated[%c32 to %c64] : !stream.resource<transient>{%c32} -> %updated as !stream.resource<external>{%c64} + stream.yield %result : !stream.resource<external>{%c64} + } => !stream.timepoint + // CHECK: util.optimization_barrier %[[TARGET]], + util.optimization_barrier %result, %timepoint : !stream.resource<external>, !stream.timepoint + // CHECK: util.return + util.return +}