add elide transfers pass - #1
Draft
ziereis wants to merge 92 commits into
Draft
Conversation
ziereis
marked this pull request as draft
May 22, 2025 14:22
ziereis
force-pushed
the
stream_topology_attr
branch
from
May 23, 2025 07:38
fb33a0f to
e231073
Compare
ziereis
force-pushed
the
stream_elide_transfer_ops
branch
2 times, most recently
from
May 23, 2025 13:03
c2a3ea0 to
501f6a7
Compare
Signed-off-by: Thomas Ziereis <ziereis@roofline.ai>
ziereis
force-pushed
the
stream_elide_transfer_ops
branch
from
May 26, 2025 14:10
501f6a7 to
b48d0b1
Compare
…-org#20873) New Revert: llvm/llvm-project#139521 will be fixed in a followup as it is a large change Currently, it carries 4 reverts: - iree-org/llvm-project@13631cf - iree-org/llvm-project@0357fd9 - iree-org/llvm-project@20a9a98 - iree-org/llvm-project@435bb50
This PR adds support for dynamic unit trip (0 or 1 trip) scf.for using scf.if. --------- Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>
This should almost always be 'auto' but when performing executable compilation (`--compile-mode=hal-executable`) it can be explicitly set to `hsaco` to have the output of the compile command be the raw HSACO ELF image. Example: ``` iree-compile \ --compile-mode=hal-executable \ --iree-hal-target-device=hip --iree-hip-target=gfx942 \ --iree-rocm-container-type=hsaco \ tools/test/iree-benchmark-executable.mlir \ -o=hsaco.elf ```
…ree-org#20897) Renames `iree_codegen.load_from_memref` to `iree_codegen.load_from_buffer`, and `iree_codegen.store_to_memref` to `iree_codegen.store_to_buffer` for consistency with the op definition. Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
This allows for kernel specialization by way of annotating tilable ops
with the preferred iteration ranges and divisibilities to specialize
for. For example,
```
%35 = linalg.matmul ins(%30, %31 : tensor<?x?xf16>, tensor<?x?xf16>) outs(%34 : tensor<?x?xf32>)
{iree_codegen.specialization_ranges = #util<int.assumption.multi_array[
[<umin = 128, umax = 4096, udiv = 128>, <umin = 128, umax = 4096, udiv = 128>, <umin = 64, udiv = 64>],
[<umin = 4096, udiv = 256>, <umin = 4096, udiv = 256>, <udiv = 64>]
]>} -> tensor<?x?xf32>
```
This matmul will be specialized into up to 3 variants, one default, and
two for each list of ranges. More details can be found in the code
documentation.
To specialize a dispatch, the pass queries the iteration domain of the
tilable op and looks for a producer `util.int.assume` op. The function
is then cloned and the assume is updated to match the requested ranges.
This is just one way to implement kernel specialization, in the future
we can implement many more options (e.g. based on estimated TOPs or
smallest individual dimension) and that's without considering what
device we're targeting. The implementation here supports many common
cases of minimum/maximum sizes and alignments. Learnings here will
translate to future improvements for other strategies.
This pass runs first thing in the configuration pipelines because it
needs to apply to the variant, and should run before passes like
BlockDynamicDims which takes advantage of `util.assume.int` information
that this pass refines.
## Summary
This PR sets the foundation for using `global_load_lds` instruction to
load values from global to LDS memory. The pipeline is as follows:
* Only convert `linalg.copy` emitted in `PromoteGPUMatMulOperands`. When
it sees fit, insert a different attribute
(`#iree_gpu.use_global_load_dma`) to `linalg.copy` to tag it along the
pipeline.
* Tagged `linalg.copy` will not be decomposed/tiled until bufferization.
* after distributed to threads and bufferization, the tagged
`linalg.copy` will then be lowered to a sequence of code responsible for
subgroup-coalesced loading op `iree_gpu.global_load_dma`.
* `iree_gpu.global_load_dma` will be mapped to `amdgpu.gather_to_lds`
op, which will mapped to corresponding rocdl op.
* Disable padding to reduce bank conflict pass because the destination
workgroup memory has to be contiguous.
## Lowering `linalg.copy`
After bufferization and distribute to threads, tagged `linalg.copy`
still exists in the IR:
```
linalg.copy {lowering_config = #iree_gpu.use_global_load_dma}
ins(%subview_12 : memref<64x128xi8, strided<[256, 1], offset: ?>, #amdgpu.address_space<fat_raw_buffer>>)
outs(%alloc_4 : memref<64x128xi8, #gpu.address_space<workgroup>>)
```
Note that this `linalg.copy` is kept in the thread's code. The op itself
is then converted into a `for loop`, in which subgroup of threads loads
coalesced chunk of values. For example, assume there are N subgroups
loading from `tensor<a x b x c>`:
* then `i`-th subgruop will load a sub tensor of size `[a/N, b, c]`, so
each slice is consecutive.
* At this moment, assume row-major, and only tile the outermost dim.
* The reason right now we are only dealing with `linalg.copy` emitted by
`GPUPromoteMatmulOperands` is that we know the destination is allocated
contiguously.
* TODO: expand to any memref slices.
* given `gpu.subgroup_id` and `gpu.lane_id`, each thread calculates the
consecutive data chunk the subgroup the thread belongs to is responsible
to load:
* the chunk indices is the delinearized indices of the input tensor,
from:
* `affine.delinearize_index[gpu.subgroup_id * (num_elems_of(tensor) /
num_subgroups)]`, to
* `affine.delinearize_index[(gpu.subgroup_id + 1) *
(num_elems_of(tensor) / num_subgroups) - 1]`
* Assume each subgroup will load `n` values from linearized index `[N_f,
N_b]`, then thread with lane id `i` will try to load: `iter = 0 to n :
N_f + subgroup_size * iter + (i - 1)` .
Then it will be converted to something like the following (in the
example, assume `workgroup size = 256`, `subgroup_size = 64`, loading
`64x128xi8`):
```miler
scf.for %indvar = %c0 to %c32 step %c1 {
;; thread-specific gathering address from global address
%17 = affine.apply affine_map<()[s0, s1, s2] -> (s0 + s1 * 2048 + s2 * 64)>()[%lane_id, %subgroup_id, %indvar]
%18:2 = affine.delinearize_index %17 into (128, 64) : index, index
;; this iteration's base storing index
%19 = affine.apply affine_map<()[s0, s1] -> (s0 * 2048 + s1 * 64)>()[%subgroup_id, %indvar]
%20:2 = affine.delinearize_index %19 into (128, 64) : index, index
iree_gpu.global_load_dma %subview_13[%18#0, %18#1] -> %alloc_5[%20#0, %20#1] : memref<128x64xi8, strided<[256, 1], offset: ?>, #amdgpu.address_space<fat_raw_buffer>> -> memref<128x64xi8, #gpu.address_space<workgroup>>
}
;; if there are residual elements (subgroup_copy_region_size % subgroup_size != 0), copy residual elements here
gpu.barrier
```
## Dependent PRs:
* design doc: https://hackmd.io/N0RitxPzT9GPhM0jEPtOCg?view
* upstream changes required:
* llvm/llvm-project#133498
* llvm/llvm-project#136405
* llvm/llvm-project#137671
* llvm/llvm-project#137425
* iree-org#20800 (review)
---------
Signed-off-by: Alan Li <me@alanli.org>
…#20911) It closes iree-org#20907. It closes iree-org#20093. Newly supported APIs: - PJRT_Client_AddressableMemories - PJRT_Device_AddressableMemories - PJRT_Device_DefaultMemory - PJRT_Buffer_Memory - PJRT_Memory_Id - PJRT_Memory_Kind - PJRT_Memory_ToString - PJRT_Memory_DebugString - PJRT_Memory_AddressableByDevices Currently the memory is one-to-one mapped from its device, just for JAX to successfully initialize the PJRT plugin. After this patch, the PJRT plugin can support JAX from 0.5.1 up to the latest version (0.6.1) (see iree-org#20907 for the current status). ci-exactly: build_packages, test_pjrt Signed-off-by: PragmaTwice <twice@apache.org>
It was rarely doing anything but deleting unused ops after cloning but running on the full module (including dispatches). This does not change the number of iterations required but does make each iteration significantly faster and makes the code easier to refactor in subsequent changes.
InlineConstantGlobalInitializer was not using a symbol table and is redundant with what FoldGlobalsPass does anyway. The global info queried in resource usage analysis was just used to get the type - instead, the type can be accessed locally on the load/store op.
Previously the pass would always clone an op if analysis resulted in multiple producer affinities. This leads to bad behavior where the pass will keep trying to clone ops in order to fix the situation but do nothing meaningful before the next try (and eventually hit the limit). This should help situations where CloneToConsumersPass was hitting its iteration limit. It does not fix affinity analysis iteration limits - those are caused by ambiguous IR and will still fail to converge with these changes.
Mostly NFC besides early-exiting on assigning affinities to function-like ops, as they can never have affinities. Removes a lot of debug spam (entire function/initializer contents).
… semantic change. (iree-org#20913) This drops a revert: iree-org/llvm-project@435bb50 To make IREE compatible with the new `memref.assuem_assignment` semantics, the revision moves the op creation from OneShotBufferization to post bufferization passes. It is mainly for reducing the burdens of bufferization analysis. Otherwise, it'd easily generate redundant copies between the original memref and the result of AssumeAlignment op. Furthermore, there are several missing support when AssumeAlignment op is present in SSA chain. - The `FlattenMemRefSubspanPass` needs to be improved to take the op into account. - We are missing a pattern that propagates the `memref.dim` ops for `memref.dim(memref.assume_alignment)` cases. It reduces the deps from AssumeAlignment op, so we don't need to add more patterns for it. - The chained memref.assume_alignment ops are not folded away. This can be done as a canonicalization pattern. The revision adds a workaround for early-bufferization dispatch, and it will be removed once we implement the pattern. - Need to teach SPIRVVectorizeLoadStore pass to take AssumeAlignment ops into account. - Need to teach VMVX ResolveBufferDescriptors pass to take AssumeAlignment ops into account. The new pattern is copied from the other one, and a refactoring can be done in the future. For now, we mainly want to drop the revert, and we can easily refactor the code out. So it is a TODO for now. - It triggers failures in hoisting vector transfer ops when AssumeAlignment ops are present. We don't know the root cause yet, and it is too complicated and hard to fix it within the revision. Because the implementation is in upstream. See iree-org#20912 for reproducer. Closes iree-org#20908 --------- Signed-off-by: hanhanW <hanhan0912@gmail.com>
…ils. (iree-org#20915) Signed-off-by: hanhanW <hanhan0912@gmail.com>
…0925) Few tests are updated to use the result of `memref.assume_alignment` because they care about the ops. Many `memref.assume_alignment` ops are removed from many lit tests because they do not need them. They don't use them at their pass scope. There are two tests not changed because of recent regression. I'll revisit it if I come back to the issue. For now, leave them as what they are: - SPIRV/test/tile_and_promote_cooperative_matrix.mlir - SPIRV/test/tile_and_vectorize_to_cooperative_ops.mlir This is a follow-up to iree-org@cbb7536 --------- Signed-off-by: hanhanW <hanhan0912@gmail.com>
Currently we just build and test the IREE PJRT CPU plugin in pkgci. Since we don't have Nvidia GPUs for CI runners refer to iree-org#18814 now, we can at least confirm that the CUDA PJRT plugin can be successfully built in CI. This PR tries to enable it. ci-exactly: build_packages, test_pjrt Signed-off-by: PragmaTwice <twice@apache.org>
This change extends `ReshapeFusion.cpp` to support propagating reshapes through `iree_linalg_ext.gather`. The logic is very similar to `scatter` but with the 0th and 2nd operands flipped. Additionally, the `LinalgExt` reshape propagation tests were moved out of `DispatchCreation`. A test pass was created in `LinalgExt` to run the reshape propagation patterns. --------- Signed-off-by: Ian Wood <ianwood2024@u.northwestern.edu>
…org#20923) `--iree-stream-affinity-solver-max-iterations=N` will allow an easy way to isolate analysis failures to the solver iteration count. Now that CloneToConsumersPass isn't running the solver itself in a loop (much) it's less scary to up the default, though it's still risky in the failure cases as it could lead to _very_ long times to failure. The proper solution is to reduce the iteration count by somehow dealing with very long tied op chains that pass through multiple affinities.
Dropped reverts, which are fixed by the head: - iree-org/llvm-project@0357fd9 - iree-org/llvm-project@20a9a98 Carried revert: - (old) iree-org/llvm-project@13631cf - (new) iree-org/llvm-project@5f144b5 - (new) iree-org/llvm-project@cddcc13 - (new) llvm/llvm-project@989aadf Additional cherry-picks, which can be dropped after we bump ahead of it: - iree-org/llvm-project@92b9689 The fixes in IREE are for - llvm/llvm-project@05494f3 - llvm/llvm-project@61d5fdf The second and the third reverts because they break torch-mlir. I have a fix for IREE side and can fix torch-mlir together in our repo. I prefer fixing them later because we are already ~1 week behind the LLVM head. The four revert is because of an assertion in convert-to-llvm lowering in ARM SME pipeline tests. Fixes iree-org#20737 --------- Signed-off-by: hanhanW <hanhan0912@gmail.com>
There are methods on the MMA operation and the MMAKind attribute that I found no uses of, and so I'm going to go ahead and delete them.
Currently default tuning specs are registered to a global embedded directory when the plugin for a backend is registered. Then the MaterializeTuningSpecs pass would try to interpret the target attribute as a gpu target attr to get a target arch to use. This is action at a distance that only works for gpu targets. This patch replaces the global directory with a simple lookup for "iree_codegen.default_tuning_spec" in the `hal.executable.target` configuration. The attribute stored there maintains a reference to a (typically static) module that can be retrieved by attribute interface, giving target backends freedom over the mechanism used to deliver the tuning spec + caching implementations. Today only the ROCMTarget implements a default tuning spec (for a single arch at that), so this PR adds a ROCMDialect to the ROCM plugin to implement the `rocm.builtin.tuning_module` attribute. Any builtin accessed through this attribute is cached on the dialect (same as the iree_codegen dialect). The same move to attribute based reference will be done for ukernels as a follow up (but with a different interface given that parsing is not needed).
…-org#20988) If we successfully created the buffer view, it holds a ref to the buffer, so we want to release our own. If we failed then we want to release the imported buffer anyway so we don't leak. Signed-off-by: Andrew Woloszyn <andrew.woloszyn@gmail.com>
This test has been reported as flaky on certain builds: https://discord.com/channels/689900678990135345/1062405112292712499/1379561856460652706 and currently has no active maintainer. Disable the test until the non-determinism is fixed or the code is dropped.
Improve documentation related to sharktuner. See Issue: nod-ai/amd-shark-ai#1248 --------- Signed-off-by: Muzammiluddin Syed <muzasyed@amd.com>
Signed-off-by: hanhanW <hanhan0912@gmail.com>
…iree-org#20885) Adds `IREE::Stream::AffinityTopologyAttrInterface` so that analysis and passes in the stream dialect can query for topology information from target device layers. `#hal.device.topology<...>` is the first implementation. Fixes iree-org#20854. --------- Signed-off-by: Thomas Ziereis <ziereis@roofline.ai> Co-authored-by: Ben Vanik <ben.vanik@gmail.com>
iree-org#20468) When performing cross-lane reductions using subgroup_reduce ops across contiguous lanes on AMD GPUs, lower to Data Parallel Primitives (DPP) ops when possible. This reduces latency on applicable devices. See related iree-org#20007 --------- Signed-off-by: Muzammiluddin Syed <muzasyed@amd.com>
This PR is part of a larger change moving `LLVMCPUTileRootAndFuseProducerConsumer`'s tiling of parallel dimensions to use `scf.forall` rather than `scf.for`; allowing us to take advantage of some canonicalization patterns available on scf.forall. The patterns available prevent redundant stack allocations seen in iree-org#20792. Tiling parallel dimensions with `scf.forall` rather than `scf.for` could also be said to be semantically cleaner. As part of that change, this PR intrudes shared tiling utilities pulled out of `TileDispatchUsingForall` that will later be used in `LLVMCPUTileRootAndFuseProducerConsumer`. PR 1/4 addressing iree-org#20792
This PR is part of a larger change moving `LLVMCPUTileRootAndFuseProducerConsumer`'s tiling of parallel dimensions to use `scf.forall` rather than `scf.for`; allowing us to take advantage of some canonicalization patterns available on scf.forall. The patterns available prevent redundant stack allocations seen in iree-org#20792. Tiling parallel dimensions with `scf.forall` rather than `scf.for` could also be said to be semantically cleaner. As part of that change, the pipeline will need to moved the tiled `scf.forall` back to a `scf.for` as later stages expect `scf.for`. This PR introduces a pass to do the conversion. PR 2/4 addressing iree-org#20792 --------- Co-authored-by: Prashant Kumar <pk5561@gmail.com> Co-authored-by: Han-Chung Wang <hanhan0912@gmail.com>
std::string is not a round-trippable attr parameter. Change to a stringref parameter.
) This change combines `BubbleUpExpandShapes` and `BubbleUpExtractSlices` into a single pass. In cases like `linalg.generic -> extract -> expand -> extract -> expand` the expand & extract bubbling patterns must be done together and iteratively to be able to bubble all ops through the `linalg.generic`. This fixes missed paged attention fusion with the mask producer. --------- Signed-off-by: Ian Wood <ianwood2024@u.northwestern.edu>
Integrate to iree-org/llvm-project@d96447b Drops aall cherry-picks/reverts.
…rg#21011) The IGEMM works out of the box for the dilated convolutions by removing the previous guard condition. The output numerics are correct for forward convs. The performance can be improved from 475us to 195us with IGEMM for the following conv kernel. ``` convbfp16 -n 16 -c 144 -H 24 -W 16 -k 144 -y 3 -x 3 -p 2 -q 2 -u 1 -v 1 -l 2 -j 2 -m conv -g 1 -F 1 -t 1 --in_layout NHWC --out_layout NHWC --fil_layout NHWC --iter 100 ``` --------- Signed-off-by: yzhang93 <zhyuhang88@gmail.com>
In experiments with shortfin, this prevents the kernel launches from appearing before they are enqueued. Signed-off-by: Eric Eaton <erieaton@amd.com>
iree-org#21020) …(iree-org#20805)" This reverts commit d7bb36b. Issue Tracker: iree-org#21019 Signed-off-by: Praveen G <praveen.g2@amd.com>
…ree-org#21027) Fixes iree-org#21024 Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>
…ee-org#21025) This commonly arises in dynamic resource offsets on stream bindings: a bunch of `umin=0, umax=0` (or whatever offset) values that ugly up the IR.
Signed-off-by: default <ziereis@roofline.ai>
Signed-off-by: default <ziereis@roofline.ai>
Signed-off-by: default <ziereis@roofline.ai>
Signed-off-by: default <ziereis@roofline.ai>
Signed-off-by: default <ziereis@roofline.ai>
Signed-off-by: default <ziereis@roofline.ai>
Signed-off-by: default <ziereis@roofline.ai>
ziereis
pushed a commit
that referenced
this pull request
Aug 27, 2025
This PR adds a pattern to swap the `collapse_shape` with the `extract_slice` op to enable more loop fusion opportunity. Note this pattern is adapted from the upstream[ `BubbleUpCollapseShapeThroughExtractSlice`](https://github.com/llvm/llvm-project/blob/ffb453989b0e95d85b6cfa543b65fec23b65649d/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp#L322) by allowing some special cases we've seen through the IGEMM path. The special case is processed under `CASE #1` session with the corresponding tests mainly focusing on the changes. Note that some strict conditions are added to match the special case for convolutions which maybe not general and robust enough to add to upstream. In addition, we break the [upstream single pattern](https://github.com/llvm/llvm-project/blob/ffb453989b0e95d85b6cfa543b65fec23b65649d/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp#L767) into two separate ones to be able to apply them separately in different passes. --------- Signed-off-by: yzhang93 <zhyuhang88@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
20 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
draft implementation for iree-org#20853