summaryrefslogtreecommitdiff
path: root/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
AgeCommit message (Collapse)Author
2025-10-22[mlir][vector][xegpu] Accept uniform values in `getDistributedType` (#163887)Charitha Saumya
Uniform values should not be distributed during vector distribution. Example would be a reduction result where reduction happens across lanes. However, current `getDistributedType` does not accept a zero result affine map (i.e. no distributed dims) when describing the distributed dimensions. This result in null type being returned and crashing the vector distribution in some cases. An example case would be a `scf.for` op (about to be distributed) in which one of the for result is a uniform value and it does not have a user outside the warp op. This necessitates querying the `getDistributedType` to figure our the distributed type of this value.
2025-10-17[mlir][vector] Add support for yielding loop bounds in `scf.for` ↵Charitha Saumya
distribution. (#163443) In some cases, loop bounds (lower, upper and step) of `scf.for` can come locally from the parent warp op the `scf.for`. Current logic will not yield the loop bounds in the new warp op generated during lowering causing sinked `scf.for` to have non dominating use. In this PR, we have added logic to yield loop bounds by default (treat them as other operands of `scf.for`) which fixes this bug.
2025-10-15[MLIR][Vector] Improve warp distribution robustness (#161647)Artem Kroviakov
2025-09-15[mlir][vector] Use `source` as the source argument name (#158258)Andrzej Warzyński
This patch updates the following ops to use `source` (instead of `vector`) as the name for their source argument: * `vector.extract` * `vector.scalable.extract` * `vector.extract_strided_slice` This change ensures naming consistency with the "builders" for these Ops that already use the name `source` rather than `vector`. It also addresses part of: * https://github.com/llvm/llvm-project/issues/131602 Specifically, it ensures that we use `source` and `dest` for read and write operations, respectively (as opposed to `vector` and `dest`).
2025-09-10[MLIR][Vector] Add warp distribution for `scf.if` (#157119)Artem Kroviakov
This PR adds `scf.if` op distribution to the existing `VectorDistribute` patterns. The logic mostly follows that of `scf.for`: move op outside, wrap each branch with `gpu.warp_execute_on_lane_0`. A notable difference to `scf.for` is that each branch has its own set of escaping values, and `scf.if` itself does not have block arguments.
2025-08-28[MLIR][Vector] Add warp distribution for `vector.step` op (#155425)Artem Kroviakov
This PR adds a distribution pattern for [`vector.step`](https://mlir.llvm.org/docs/Dialects/Vector/#vectorstep-vectorstepop) op. The result of the step op is a vector containing a sequence `[0,1,...,N-1]`. For the warp distribution, we consider a vector with `N == warp_size` (think SIMD). Distributing it to SIMT, means that each lane is represented by a thread/lane id scalar. More complex cases with the support for warp size multiples (e.g., `[0,1,...,2*N-1]`) require additional layout information to be handled properly. Such support may be added later. The lane id scalar is wrapped into a `vector<1xindex>` to emulate the sequence distribution result. Other than that, the distribution is similar to that of `arith.constant`.
2025-08-22[mlir][gpu] Warp execute terminator getter (#154729)Adam Siemieniuk
Adds a utility getter to `warp_execute_on_lane_0` which simplifies access to the op's terminator. Uses are refactored to utilize the new terminator getter.
2025-08-15[mlir][SCF] `scf.for`: Add support for unsigned integer comparison (#153379)Matthias Springer
Add a new unit attribute to allow for unsigned integer comparison. Example: ```mlir scf.for unsigned %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 { // body } ``` Discussion: https://discourse.llvm.org/t/scf-should-scf-for-support-unsigned-comparison/84655
2025-07-25[mlir][NFC] update `mlir/Dialect` create APIs (32/n) (#150657)Maksim Levental
See https://github.com/llvm/llvm-project/pull/147168 for more info.
2025-07-22[mlir][NFC] update `mlir/Dialect` create APIs (24/n) (#149931)Maksim Levental
See https://github.com/llvm/llvm-project/pull/147168 for more info.
2025-07-14Reapply [mlir][vector] Refactor WarpOpScfForOp to support unused or swapped ↵Charitha Saumya
forOp results. (#148313) Reapply attempt for : https://github.com/llvm/llvm-project/pull/148291 Fix for the build failure reported in : https://lab.llvm.org/buildbot/#/builders/116/builds/15477 ----- This crash is caused by mismatch of distributed type returned by `getDistributedType` and intended distributed type for forOp results. Solution diff: https://github.com/llvm/llvm-project/commit/20c2cf67662c3b3fdecf95a0e280809f98d8db50 Example: ``` func.func @warp_scf_for_broadcasted_result(%arg0: index) -> vector<1xf32> { %c128 = arith.constant 128 : index %c1 = arith.constant 1 : index %c0 = arith.constant 0 : index %2 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<1xf32>) { %ini = "some_def"() : () -> (vector<1xf32>) %0 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini) -> (vector<1xf32>) { %1 = "some_op"(%arg4) : (vector<1xf32>) -> (vector<1xf32>) scf.yield %1 : vector<1xf32> } gpu.yield %0 : vector<1xf32> } return %2 : vector<1xf32> } ``` In this case the distributed type for forOp result is `vector<1xf32>` (result is not distributed and broadcasted to all lanes instead). However, in this case `getDistributedType` will return NULL type. Therefore, if the distributed type can be recovered from warpOp, we should always do that first before using `getDistributedType`
2025-07-11Revert "[mlir][vector] Refactor WarpOpScfForOp to support unused or swapped ↵Charitha Saumya
forOp results." (#148291) Reverts llvm/llvm-project#147620 Reverting due to build failure: https://lab.llvm.org/buildbot/#/builders/116/builds/15477
2025-07-11[mlir][vector] Refactor WarpOpScfForOp to support unused or swapped forOp ↵Charitha Saumya
results. (#147620) Current implementation generates incorrect code or crashes in the following valid cases. 1. At least one of the for op results are not yielded by the warpOp. Example: ``` %0 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>) { .... %3:2 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini, %arg5 = %ini1) -> (vector<128xf32>, vector<128xf32>) { %1 = ... %acc = .... scf.yield %acc, %1 : vector<128xf32>, vector<128xf32> } gpu.yield %3#0 : vector<128xf32> // %3#1 is not used but can not be removed as dead code (loop carried). } "some_use"(%0) : (vector<4xf32>) -> () return ``` 2. Enclosing warpOp yields the forOp results in different order compared to the forOp results. Example: ``` %0:3 = gpu.warp_execute_on_lane_0(%arg0)[32] -> (vector<4xf32>, vector<4xf32>, vector<8xf32>) { .... %3:3 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini1, %arg5 = %ini2, %arg6 = %ini3) -> (vector<256xf32>, vector<128xf32>, vector<128xf32>) { ..... scf.yield %acc1, %acc2, %acc3 : vector<256xf32>, vector<128xf32>, vector<128xf32> } gpu.yield %3#2, %3#1, %3#0 : vector<128xf32>, vector<128xf32>, vector<256xf32> // swapped order } "some_use_1"(%0#0) : (vector<4xf32>) -> () "some_use_2"(%0#1) : (vector<4xf32>) -> () "some_use_3"(%0#2) : (vector<8xf32>) -> () ```
2025-07-09[mlir][Vector] Remove usage of `vector.insertelement/extractelement` from ↵Diego Caballero
Vector (#144413) This PR is part of the last step to remove `vector.extractelement` and `vector.insertelement` ops. RFC: https://discourse.llvm.org/t/rfc-psa-remove-vector-extractelement-and-vector-insertelement-ops-in-favor-of-vector-extract-and-vector-insert-ops It removes instances of `vector.extractelement` and `vector.insertelement` from the Vector dialect layer.
2025-06-30[mlir][vector] Avoid setting padding by default to `0` in ↵Fabian Mora
`vector.transfer_read` prefer `ub.poison` (#146088) Context: `vector.transfer_read` always requires a padding value. Most of its builders take no `padding` value and assume the safe value of `0`. However, this should be a conscious choice by the API user, as it makes it easy to introduce bugs. For example, I found several occasions while making this patch that the padding value was not getting propagated (`vector.transfer_read` was transformed into another `vector.transfer_read`). These bugs, were always caused because of constructors that don't require specifying padding. Additionally, using `ub.poison` as a possible default value is better, as it indicates the user "doesn't care" about the actual padding value, forcing users to specify the actual padding semantics they want. With that in mind, this patch changes the builders in `vector.transfer_read` to always having a `std::optional<Value> padding` argument. This argument is never optional, but for convenience users can pass `std::nullopt`, padding the transfer read with `ub.poison`. --------- Signed-off-by: Fabian Mora <fabian.mora-cordero@amd.com>
2025-06-25[mlir][vector] Add support for vector extract/insert_strided_slice in vector ↵Charitha Saumya
distribution. (#145421) This PR adds initial support for `vector.extract_strided_slice` and `vector.insert_strided_slice` ops in vector distribution.
2025-06-13Revert "[mlir][vector] Fix for WarpOpScfForOp failure when scf.for has ↵Charitha Saumya
results that are unused." (#144124) Reverts llvm/llvm-project#141853 Reverting the bug fix because it does not handle all cases correctly.
2025-06-09[mlir][vector] Fix for WarpOpScfForOp failure when scf.for has results that ↵Charitha Saumya
are unused. (#141853) Currently, only the values defined outside ForOp but inside the original WarpOp are considered "escaping values". However this is not true if the ForOp has some unused results. In this case, corresponding IterArgs must also be yielded by the original WarpOp. This PR adds the required code changes to achieve this.
2025-05-12[mlir][vector] Standardize `base` Naming Across Vector Ops (NFC) (#137859)Andrzej Warzyński
[mlir][vector] Standardize base Naming Across Vector Ops (NFC) This change standardizes the naming convention for the argument representing the value to read from or write to in Vector ops that interface with Tensors or MemRefs. Specifically, it ensures that all such ops use the name `base` (i.e., the base address or location to which offsets are applied). Updated operations: * `vector.transfer_read`, * `vector.transfer_write`. For reference, these ops already use `base`: * `vector.load`, `vector.store`, `vector.scatter`, `vector.gather`, `vector.expandload`, `vector.compressstore`, `vector.maskedstore`, `vector.maskedload`. This is a non-functional change (NFC) and does not alter the semantics of these operations. However, it does require users of the XFer ops to switch from `op.getSource()` to `op.getBase()`. To ease the transition, this PR temporarily adds a `getSource()` interface method for compatibility. This is intended for downstream use only and should not be relied on upstream. The method will be removed prior to the LLVM 21 release. Implements #131602
2025-05-06[mlir] Remove unused local variables (NFC) (#138642)Kazu Hirata
2025-04-07[mlir][vector] Standardise `valueToStore` Naming Across Vector Ops (NFC) ↵Andrzej Warzyński
(#134206) This change standardises the naming convention for the argument representing the value to store in various vector operations. Specifically, it ensures that all vector ops storing a value—whether into memory, a tensor, or another vector — use `valueToStore` for the corresponding argument name. Updated operations: * `vector.transfer_write`, `vector.insert`, `vector.scalable_insert`, `vector.insert_strided_slice`. For reference, here are operations that currently use `valueToStore`: * `vector.store` `vector.scatter`, `vector.compressstore`, `vector.maskedstore`. This change is non-functional (NFC) and does not affect the functionality of these operations. Implements #131602
2024-12-13[MLIR] Create GPU utils library & move distribution utils (#119264)Petr Kurapov
Continue the move of `warp_execute_on_lane_0` op to the gpu dialect (#116994). This patch creates a utils library in GPU and moves generic helper functions there.
2024-11-22[MLIR] Move warp_execute_on_lane_0 from vector to gpu (#116994)Petr Kurapov
Please see the related RFC here: https://discourse.llvm.org/t/rfc-move-execute-on-lane-0-from-vector-to-gpu-dialect/82989. This patch does exactly one thing - moves the op to gpu.
2024-11-18[mlir][Vector] Move insert/extractelement distribution patterns to ↵Kunwar Grover
insert/extract (#116425) This is a NFC-ish change that moves vector.extractelement/vector.insertelement vector distribution patterns to vector.insert/vector.extract. Before: 0-d/1-d vector.extract -> vector.extractelement -> distributed vector.extractelement 2-d+ vector.extract -> distributed vector.extract After: scalar input vector.extract -> distributed vector.extract vector.extractelement -> distributed vector.extract 2d+ vector.extract -> distributed vector.extract The same changes are done for insertelement/insert. The change allows us to remove reliance on vector.extractelement/vector.insertelement, which are soon to be depreciated: https://discourse.llvm.org/t/rfc-psa-remove-vector-extractelement-and-vector-insertelement-ops-in-favor-of-vector-extract-and-vector-insert-ops/71116/8 No extra tests are included because this patch doesn't introduce / remove any functionality. It only changes the chain of lowerings. This change can be completly NFC if we make the distributed operation vector.extractelement/vector.insertelement, but that is slightly weird, because you are going from extractelement -> extract -> extractelement.
2024-11-08[mlir][IR][NFC] Cleanup insertion point API usage (#115415)Matthias Springer
Use `setInsertionPointToStart` / `setInsertionPointToEnd` when possible.
2024-08-14[mlir][vector] Add extra check on distribute types to avoid crashes (#102952)Bangtian Liu
This PR addresses the issue detailed in https://github.com/iree-org/iree/issues/17948. The problem occurs when distributed types are set to NULL, leading to compilation crashes. --------- Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
2024-08-04[mlir] Construct SmallVector with ArrayRef (NFC) (#101896)Kazu Hirata
2024-04-01[mlir][NFC] Simplify type checks with isa predicates (#87183)Jakub Kuderski
For more context on isa predicates, see: https://github.com/llvm/llvm-project/pull/83753.
2024-02-28[mlir][vector] Fix n-d transfer write distribution (#83215)Quinn Dawkins
Currently n-d transfer write distribution can be inconsistent with distribution of reductions if a value has multiple users, one of which is a transfer_write with a non-standard distribution map, and the other of which is a vector.reduction. We may want to consider removing the distribution map functionality in the future for this reason.
2024-01-17[mlir][IR] Rename "update root" to "modify op" in rewriter API (#78260)Matthias Springer
This commit renames 4 pattern rewriter API functions: * `updateRootInPlace` -> `modifyOpInPlace` * `startRootUpdate` -> `startOpModification` * `finalizeRootUpdate` -> `finalizeOpModification` * `cancelRootUpdate` -> `cancelOpModification` The term "root" is a misnomer. The root is the op that a rewrite pattern matches against (https://mlir.llvm.org/docs/PatternRewriter/#root-operation-name-optional). A rewriter must be notified of all in-place op modifications, not just in-place modifications of the root (https://mlir.llvm.org/docs/PatternRewriter/#pattern-rewriter). The old function names were confusing and have contributed to various broken rewrite patterns. Note: The new function names use the term "modify" instead of "update" for consistency with the `RewriterBase::Listener` terminology (`notifyOperationModified`).
2024-01-12[mlir][vector] Fix dominance error in warp vector distribution (#77771)Matthias Springer
This commit fixes a test in `vector-warp-distribute.mlir` when `MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS` is enabled. ``` within split at /usr/local/google/home/springerm/mlir_public/llvm-project/mlir/test/Dialect/Vector/vector-warp-distribute.mlir:1 offset :18:10: error: operand #0 does not dominate this use %1 = vector.extract %0[9] : f32 from vector<64xf32> ^ within split at /usr/local/google/home/springerm/mlir_public/llvm-project/mlir/test/Dialect/Vector/vector-warp-distribute.mlir:1 offset :18:10: note: see current operation: %1 = "affine.apply"(%8) <{map = affine_map<()[s0] -> (s0 ceildiv 2)>}> : (index) -> index within split at /usr/local/google/home/springerm/mlir_public/llvm-project/mlir/test/Dialect/Vector/vector-warp-distribute.mlir:1 offset :18:10: note: operand defined here (op in a child region) "func.func"() <{function_type = (index) -> f32, sym_name = "vector_extract_1d"}> ({ ^bb0(%arg0: index): %0:2 = "vector.warp_execute_on_lane_0"(%arg0) <{warp_size = 32 : i64}> ({ %7 = "some_def"() : () -> vector<64xf32> %8 = "arith.constant"() <{value = 9 : index}> : () -> index %9 = "vector.extractelement"(%7, %8) : (vector<64xf32>, index) -> f32 "vector.yield"(%9, %7) : (f32, vector<64xf32>) -> () }) : (index) -> (f32, vector<2xf32>) %1 = "affine.apply"(%8) <{map = affine_map<()[s0] -> (s0 ceildiv 2)>}> : (index) -> index %2 = "affine.apply"(%8) <{map = affine_map<()[s0] -> (s0 mod 2)>}> : (index) -> index %3 = "vector.extractelement"(%0#1, %2) : (vector<2xf32>, index) -> f32 %4 = "arith.index_cast"(%1) : (index) -> i32 %5 = "arith.constant"() <{value = 32 : i32}> : () -> i32 %6:2 = "gpu.shuffle"(%3, %4, %5) <{mode = #gpu<shuffle_mode idx>}> : (f32, i32, i32) -> (f32, i1) "func.return"(%6#0) : (f32) -> () }) : () -> () LLVM ERROR: IR failed to verify after pattern application ``` The position at which `vector.extractelement` extracts must also be distributed. The fix in `WarpOpExtractElement` is similar to `WarpOpInsertElement`.
2024-01-12[mlir][vector] Support warp distribution of `transfer_read` with ↵Matthias Springer
dependencies (#77779) Support distribution of `vector.transfer_read` ops when operands are defined inside of the region of `warp_execute_on_lane_0` (except for the buffer from which the op is reading). Such IR was previously not supported. This commit changes the implementation such that indices and the padding value are also distributed. This commit simplifies the implementation considerably: the original implementation created a new `transfer_read` op and then checked if this new op is valid. If not, the rewrite pattern failed. This was a bit hacky. It was also a violation of the rewrite pattern API (detected by `MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS`) because the IR was modified, but the pattern returned "failure".
2023-12-12[mlir][vector] Allow vector distribution with multiple written elements (#75122)Jakub Kuderski
Add a configuration option to allow vector distribution with multiple elements written by a single lane. This is so that we can perform vector multi-reduction with multiple results per workgroup.
2023-11-27[mlir][vector] Distribute all non-permutation or broadcasted masked transfer ↵Quinn Dawkins
reads (#73539) The primary difficulty with distribution of masked transfers is when the permutation map permutes the vector, in which case the distribution logic needs to make sure the correct mask elements end up with the distributed transfer. This is only tricky when the permutation map has a permutation in it, so we can relax the condition for distribution.
2023-11-14[mlir][affine][nfc] cleanup deprecated T.cast style functions (#71269)long.chen
detail see the docment: https://mlir.llvm.org/deprecation/ Not all changes are made manually, most of them are made through a clang tool I wrote https://github.com/lipracer/cpp-refactor.
2023-11-10[mlir][vector] Drop incorrect startRootUpdate calls in vector distribution ↵Quinn Dawkins
(#71988) Fixes asan failures in https://lab.llvm.org/buildbot/#/builders/5/builds/38191 introduced by #71964.
2023-11-10[mlir][vector] Notify the rewriter when sinking out of warp ops (#71964)Quinn Dawkins
A number of the warp distribution patterns work by rewriting a warp op in place by moving a contained op outside. This notifies the rewriter that the warp op is changing in this case.
2023-11-10[mlir][vector] Add distribution pattern for vector.create_mask (#71619)Quinn Dawkins
This is the last step needed for basic support for distributing masked vector code. The lane id gets delinearized based on the distributed mask shape and then compared against the original mask sizes to compute the bounds for the distributed mask. Note that the distribution of masks is implicit on the shape specified by the warp op. As a result, it is the responsibility of the consumer of the mask to ensure the distributed mask will match its own distribution semantics.
2023-11-10[mlir][vector] Root the transfer write distribution pattern on the warp op ↵Quinn Dawkins
(#71868) Currently when there is a mix of transfer read ops and transfer write ops that need to be distributed, because the pattern for write distribution is rooted on the transfer write, it is hard to guarantee that the write gets distributed after the read when the two aren't directly connected by SSA. This is likely still relatively unsafe when there are undistributable ops, but structurally these patterns are a bit difficult to work with. For now pattern benefits give fairly good guarantees for happy paths.
2023-11-09[mlir][vector] Fix cases with multiple yielded transfer_read ops (#71625)Quinn Dawkins
This fixes two bugs: 1) When deciding whether a transfer read could be propagated out of a warp op, it looked for the first yield operand that was produced by a transfer read. If this transfer read wasn't ready to be distributed, the pattern would not re-check for any other transfer reads that could have been propagated. 2) When dropping dead warp results, we do so by updating the warp op signature and splicing in the old region. This does not add the ops in the body of the warp op back to the pattern applicator's worklist, and thus those operations won't be DCE'd. This is a problem for patterns like the one for transfer reads that will still see the dead operation as a user.
2023-11-09[mlir][vector] Add pattern to distribute masked reads (#71610)Quinn Dawkins
Because the distribution is based on types, supporting general masked reads requires first materializing the permutation map in IR to align the elements of the mask with the elements read by the transfer op. For now just support cases with the trivial permutation map.
2023-11-07[mlir][vector] Add support for distributing masked writes (#71482)Quinn Dawkins
General distribution of masked writes requires materializing the permutation on the vector of the write in IR to ensure the vector lines up with the mask. For now just support cases with trivial permutation maps.
2023-11-06[mlir][vector] Hoist uniform scalar loop code after scf.for distribution ↵Quinn Dawkins
(#71422) After propagation of `vector.warp_execute_on_lane_0` through `scf.for`, uniform operations like those on the loop iterators can now be hoisted out of the inner warp op.
2023-09-28[mlir][vector] add result type to vector.extract assembly format (#66499)Cullen Rhodes
The vector.extract assembly format currently only contains the source type, for example: %1 = vector.extract %0[1] : vector<3x7x8xf32> it's not immediately obvious if this is the source or result type. This patch improves the assembly format to make this clearer, so the above becomes: %1 = vector.extract %0[1] : vector<7x8xf32> from vector<3x7x8xf32>
2023-09-22[mlir][Vector] Add support for Value indices to vector.extract/insertDiego Caballero
`vector.extract/insert` ops only support constant indices. This PR is extending them so that arbitrary values can be used instead. This work is part of the RFC: https://discourse.llvm.org/t/rfc-psa-remove-vector-extractelement-and-vector-insertelement-ops-in-favor-of-vector-extract-and-vector-insert-ops Differential Revision: https://reviews.llvm.org/D155034
2023-09-18[mlir][SCF] scf.for: Consistent API around `initArgs` (#66512)Matthias Springer
* Always use the auto-generated `getInitArgs` function. Remove the hand-written `getInitOperands` duplicate. * Remove `hasIterOperands` and `getNumIterOperands`. The names were inconsistent because the "arg" is called `initArgs` in TableGen. Use `getInitArgs().size()` instead. * Fix verification around ops with no results.
2023-08-22[mlir][vector] Use dyn_cast in if conditionsLei Zhang
Reviewed By: dcaballe Differential Revision: https://reviews.llvm.org/D158336
2023-08-17[mlir][vector] Fix uniform transfer_read distributionLei Zhang
If the original shape and the distributed shape is the same, we don't distribute at all--every thread is handling the whole. Reviewed By: hanchung Differential Revision: https://reviews.llvm.org/D158235
2023-08-16[mlir][vector] Enable distribution over multiple dimensionsLei Zhang
This commit starts enabling vector distruction over multiple dimensions. It requires delinearize the lane ID to match the expected rank. shape_cast and transfer_read now can properly handle multiple dimensions. Reviewed By: hanchung Differential Revision: https://reviews.llvm.org/D157931
2023-07-31[mlir][vector] Use DenseI64ArrayAttr for ExtractOp/InsertOp positionsMatthias Springer
`DenseI64ArrayAttr` provides a better API than `I64ArrayAttr`. E.g., accessors returning `ArrayRef<int64_t>` (instead of `ArrayAttr`) are generated. Differential Revision: https://reviews.llvm.org/D156684