[all-commits] [llvm/llvm-project] d57163: [mlir][Linalg] SplitReduction implementation witho...

[Nicolas Vasilache via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: d5716395792696f2b56a0d4debadd040ee385143
      https://github.com/llvm/llvm-project/commit/d5716395792696f2b56a0d4debadd040ee385143
  Author: Nicolas Vasilache <nicolas.vasilache at gmail.com>
  Date:   2022-06-22 (Wed, 22 Jun 2022)

  Changed paths:
    M mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
    M mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
    M mlir/include/mlir/IR/AffineMap.h
    M mlir/include/mlir/IR/BuiltinTypes.h
    M mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
    M mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp
    A mlir/test/Dialect/Linalg/transform-op-split-reduction-by-scaling.mlir

  Log Message:
  -----------
  [mlir][Linalg] SplitReduction implementation without tensor::ExpandShapeOp

This revision proposes a different implementation of the SplitReductoin transformation that does
not rely on tensor::ExpandShapeOp.

Previously, a dimension `[k]` would be split into `[k][kk]` via an ExpandShapeOp.
Instead, this revision proposes to rewrite `[k]` into `[factor * k + kk]`.

There are different tradeoffs involved  but the proposed implementation is more general because
the affine rewrite is well-defined. In particular, it works naturally with `?` parallel dimensions and
non-trivial indexing maps.

A further rewrite of `[factor * k + kk]` + ExpandShapeOp is possible as a followup.

Differential Revision: https://reviews.llvm.org/D128266