[all-commits] [llvm/llvm-project] 31aa7f: [mlir][Affine] Let affine.[de]linearize_index omit...

[Krzysztof Drewniak via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 31aa7f34e07c901773993dac0f33568307f96da6
      https://github.com/llvm/llvm-project/commit/31aa7f34e07c901773993dac0f33568307f96da6
  Author: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
  Date:   2024-11-18 (Mon, 18 Nov 2024)

  Changed paths:
    M mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
    M mlir/include/mlir/Dialect/Affine/Utils.h
    M mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    M mlir/lib/Dialect/Affine/Transforms/AffineExpandIndexOps.cpp
    M mlir/lib/Dialect/Affine/Utils/Utils.cpp
    M mlir/test/Dialect/Affine/affine-expand-index-ops.mlir
    M mlir/test/Dialect/Affine/canonicalize.mlir
    M mlir/test/Dialect/Affine/invalid.mlir
    M mlir/test/python/dialects/affine.py

  Log Message:
  -----------
  [mlir][Affine] Let affine.[de]linearize_index omit outer bounds (#116103)

The affine.delinearize_index and affine.linearize_index operations, as
currently defined, require providing a length N basis to [de]linearize N
values. The first value in this basis is never used during lowering and
is unused during lowering. (Note that, even though it isn't used during
lowering it can still be used to, for example, remove length-1 outputs
from a delinearize).

This dead value makes sense in the original context of these operations,
which is linearizing or de-linearizing indexes to memref<>s, vector<>s,
and other shaped types, where that outer bound is avaliable and may be
useful for analysis.

However, other usecases exist where the outer bound is not known. For
example:

    %thread_id_x = gpu.thread_id x : index
%0:3 = affine.delinearize_index %thread_id_x into (4, 16) : index,index,
index

In this code, we don't know the upper bound of the thread ID, but we do
want to construct the ?x4x16 grid of delinearized values in order to
further partition the GPU threads.

In order to support such usecases, we broaden the definition of
affine.delinearize_index and affine.linearize_index to make the outer
bound optional.

In the case of affine.delinearize_index, where the number of results is
a function of the size of the passed-in basis, we augment all existing
builders with a `hasOuterBound` argument, which, for backwards
compatibilty and to preserve the natural usage of the op, defaults to
`true`. If this flag is true, the op returns one result per basis
element, if it is false, it returns one extra result in position 0.

We also update existing canonicalization patterns (and move one of them
into the folder) to handle these cases. Note that disagreements about
the outer bound now no longer prevent delinearize/linearize
cancelations.



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications