[PATCH] D133166: [mlir][MemRef] Canonicalize extract_strided_metadata(subview)

[Nicolas Vasilache via Phabricator via llvm-commits]

nicolasvasilache added inline comments.


================
Comment at: mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp:1222
+      Value subStride;
+      if (subview.isDynamicStride(i))
+        subStride = subview.getDynamicStride(i);
----------------
qcolombet wrote:
> qcolombet wrote:
> > nicolasvasilache wrote:
> > > We should have something in StaticValueUtils.h|cpp to hide away the if/else complexity and return an OpFoldResult.
> > > Could you reuse or extend ?
> > Thanks for the pointer!
> > I dug a little bit more and I was struggling to find a generic way to describe the ValueRange we want to iterate on.
> > 
> > The high level interface is OffsetSizeAndStrideOpInterface and I thought of using that and then I stumbled on `OffsetSizeAndStrideOpInterface::getMixedStrides()` which covers that and returns directly `OpFoldResult`.
> > 
> > So looks like I may avoid the StaticValueUtils all together.
> I ended up giving up on OpFoldResult here because unlike the AffineOps, the ArithmeticOps don't play nicely with that type.
> 
> Well more precisely, you need to convert them to actual values to be able to construct the arith operations, unless I missed something.
Makes sense, thanks for trying!

Yes, the difference here is that AffineApply + AffineMap can represent constants as attribute / part of the attribute, but this is not the case for individual `arith` ops.
Whatever you do you have to materialize the constant, even when it folds perfectly.

Side note, a lower-level indexing dialect with an `affine_apply` op in it would allow breaking unnecessary dependences but the discussion churn trying to get progress on that has been too high, this is as good as we can do right now.


================
Comment at: mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp:1228
+
+      Value newStride = rewriter.createOrFold<arith::MulIOp>(
+          origLoc, newExtractStridedMetadata.getStrides()[i], subStride);
----------------
qcolombet wrote:
> qcolombet wrote:
> > nicolasvasilache wrote:
> > > Can we use a single AffineApplyOp here and below?
> > > At the higher levels of abstraction this is a better way of representing such indexing computations (i.e. 1 op instead of N ops).
> > > 
> > > See e.g. https://sourcegraph.com/github.com/llvm/llvm-project/-/blob/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp?L154
> > > You can also make the variadic process of defining symbols bound to a position better by using something like: https://sourcegraph.com/github.com/llvm/llvm-project/-/blob/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp?L281
> > Thanks for that pointer as well.
> > Makes a lot of sense.
> > 
> > However regarding the 1 op vs. N ops. That's going to be true for the offset, but not for the strides.
> > Indeed, we still need to produce one value per stride since these appear in the final result.
> > 
> > I need to look closer how `AffineApplyOp` are built to avoid creating useless one (when one of the stride is 1 or both strides are constant).
> Alright reporting back on this.
> 
> If we do this, the `MemRef` dialect will depend on the `Affine` dialect, which I am not sure is what the community wants in general.
> 
> On top of that, the `Affine` dialect already depends on the `MemRef` dialect, so we would create a circular dependency. That may not be a deal breaker, but I don't have enough MLIR experience at this point to know how it could be broken.
Fair enough.
Given that strides indeed require multiple ops anyway, this is fine.

For reference, one way to break cycles (but not applicable here) is to realize that canonicalization patterns don't necessarily need to be rooted on the consumer operation.
Here is one example: https://discourse.llvm.org/t/cross-dialect-folding-and-canonicalization/2740

The key insight (that does not apply here) is: "In this case, no dialect dependency (in the milr registration/loading sense) is needed because your pattern does not produce ops from a different dialect."
Here we would produce an `affine` dialect op.



================
Comment at: mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp:1271
+    // the subview.
+    for (int i = 0; i < sourceRank; ++i) {
+      if (droppedDims.test(i))
----------------
qcolombet wrote:
> qcolombet wrote:
> > nicolasvasilache wrote:
> > > Can this be factored out into a helper function that would live next to e.g. `SubViewOp::inferResultType` ?
> > > It could return a vector of OpFoldResult to avoid passing a builder and then StaticValueUtils can be used or extended to materialize the constants here.
> > > 
> > > All this should then become much more idiomatic (and shorter).
> > Thanks for the suggestion, looking!
> I ended up populating the resulting strides and sizes in the same loop and I didn't see a good refactoring at this point.
> 
> I decided to go this way, because we already have to loop through the dimensions to populate the strides.
SG!

Not using AffineApplyOp triggers a different set of tradeoffs and this looks fine to me.



================
Comment at: mlir/test/Dialect/MemRef/canonicalize.mlir:910
+//
+//       CHECK: return %[[BASE]], %[[FINAL_OFFSET]], %[[C6]], %[[C3]], %[[STRIDES]]#1, %[[STRIDES]]#2
+func.func @extract_strided_metadata_of_rank_reduced_subview(%base: memref<8x16x4xf32>) -> (memref<f32>, index, index, index, index, index) {
----------------
In this example and below I believe some foldings and tests are missing.
The strides of `%ARG` are known, they are `[64, 4, 1]`, so the return should be `%[[C4]]`, `%[[C1]]`.

It would be nice to also have some mixed tests where the function argument type has e.g. `strides<?x2>` etc.

You could have 1 fully dynamic test that spells out all the IR (adds, muls etc) and a few other tests that are just 1-line cheks and look that the return have the right constant in the proper place.


Repository:
  rG LLVM Github Monorepo

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

https://reviews.llvm.org/D133166