[MLIR] Simplify affine maps + operands exploiting IV info
authorUday Bondhugula <uday@polymagelabs.com>
Mon, 26 Sep 2022 05:23:16 +0000 (10:53 +0530)
committerUday Bondhugula <uday@polymagelabs.com>
Tue, 4 Oct 2022 12:48:34 +0000 (18:18 +0530)
commitddff3766b78d1b1d2c003639ab44c8afd1e458c1
tree6f96b388891ed060d40ec5ce02f153a3cb8f186f
parent82cac65dd286e2bef6b0e482c9562a37ff859409
[MLIR] Simplify affine maps + operands exploiting IV info

Simplify affine expressions and maps while exploiting simple range and
step info of any IVs that are operands. This simplification is local,
O(1) and practically useful in several scenarios. Accesses with
floordiv's and mod's where the LHS is non-negative and bounded or is a
known multiple of a constant can often be simplified. This is
implemented as a canonicalization for all affine ops in a generic way:
all affine.load/store, vector_load/store, affine.apply, affine.min/max,
etc. ops.

Eg: For tiled loop nests accessing buffers this way:

affine.for %i = 0 to 1024 step 32 {
  affine.for %ii = 0 to 32 {
    affine.load [(%i + %ii) floordiv 32, (%i + %ii) mod 32]
  }
}

// Note that %i is a multiple of 32 and %ii < 32, hence:

(%i + %ii) floordiv 32 is the same as %i floordiv 32
(%i + %ii) mod 32 is the same as %ii mod 32.

The simplification leads to simpler index/subscript arithmetic for
multi-dimensional arrays and also in turn enables detection of spatial
locality (for vectorization for eg.), temporal locality or loop
invariance for hoisting or scalar replacement.

Differential Revision: https://reviews.llvm.org/D135085
mlir/include/mlir/IR/AffineMap.h
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/lib/IR/AffineMap.cpp
mlir/test/Dialect/Affine/canonicalize.mlir