[PATCH] D78464: [mlir][Linalg] Add support for fusing linalg.tensor_reshape with linalg.generic operations.

[Mahesh Ravishankar via Phabricator via llvm-commits]

mravishankar added a comment.

Submitting some responses that I had forgetten to submit



================
Comment at: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp:583
+/// `sourceMap`. Note that this implicitly assumes that the tensors dimensions
+/// are "row-major" ordered logically.
+static AffineMap linearizeCollapsedDims(AffineMap sourceMap,
----------------
nicolasvasilache wrote:
> Adding the example from your test would be informative:
> ```
> E.g.
> %0 = op ... : tensor<?x?x4x5xf32>
> with output index_map `affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>`
> 
> and reshape:
> %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
>                                  affine_map<(i, j, k, l) -> (j, k, l)>] :
>     tensor<?x?x4x5xf32> into tensor<?x?xf32>
> 
> would be rewritten into:
> %0 = op ... : tensor<?x?x4x5xf32>
> with output index_map `affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>`
> ```
THanks for the detailed comment.


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp:648
+  auto reassociationMaps = reshapeOp.getReassociationMaps();
+  unsigned dim = 0;
+  for (auto collapsedDims : reassociationMaps) {
----------------
nicolasvasilache wrote:
> I wonder if we could slightly refactor the function exposed here https://reviews.llvm.org/D75575
> and just use it instead of rewriting a slightly different version.
> 
> Basically the existing function would just need a starting dimPos = 0 by default.
> 
> You could then just use that to get:
> ```
> for ... {
>   AffineExpr stridedExpression = makeCanonicalStridedLayoutExpr(srcShape.slice(), context);
>   if (!isPureAffine(stridedExpression))
>     return AffineMap();
>   results.push_back(stridedExpression);
> }
> return AffineMap.get(...);
> ```
> 
> And check whether AffineMap is empty to know whether it is fusible.
> 
> The downside is that in your current code structure you could be recomputing it twice but IMO this would largely be beneficial to save a few dozen lines of code that is similar to something we already have in the codebase.
> 
> If the recomputing bothers you, you could refactor a bit but I wouldn't be worried at this time, these are really trivial computations.
It ended up being quite a significant change to makeCanonicalStridedLayoutExpr, but it is more general now, and it works well with the current use case as well. It does recompute twice, but thats OK by me. PTAL


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp:681
+    auto consumerIndexMaps = consumer.indexing_maps();
+    SmallVector<Attribute, 4> fusedIndexMaps;
+    fusedIndexMaps.reserve(consumerIndexMaps.size());
----------------
nicolasvasilache wrote:
> Can we simplify l680 - 695 with something along the lines of:
> ```
> SmallVector<Attribute, 4> fusedIndexMaps(consumerIndexMaps.begin(), consumerIndexMaps.end());
> fusedIndexMaps[consumerIdx] = linearize();
> ```
Nice! Done


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D78464/new/

https://reviews.llvm.org/D78464