[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/Transforms/SimplifyExtractStridedMetadata.cpp:69
+    unsigned sourceRank = sourceType.getRank();
+    SmallVector<Type> sizeStrideTypes;
+    for (unsigned i = 0; i < sourceRank; ++i)
----------------
```
 SmallVector<Type> sizeStrideTypes(sourceRank, indexType);
```


================
Comment at: mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp:88
+    bool hasKnownStridesAndOffset =
+        succeeded(getStridesAndOffset(sourceType, sourceStrides, sourceOffset));
+    // Helper function to get either the value of the newly created
----------------
qcolombet wrote:
> Is it possible for the `getStridesAndOffset` method to fail on the source of a `subview`?
> 
> Put differently, how can I write a test for that part?
If `sourceType` come from a valid subview op then you shouldn't be able to fail here.
You may want to assert success instead ans simplify code below.


================
Comment at: mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp:107
+      return origStrides[idx];
+    };
+
----------------
qcolombet wrote:
> FWIW, I'm not super happy with this helper function, because we end up materializing (and then deleting later) constants.
> 
> I've tried to use the `getStridesAndOffset` function directly with the `AffineExpr` constructs, but I was not satisfied with the code.
> 
> Essentially, I would end up in the main loop below with something like:
> ```
> AffineConstantExpr cst;
> 
> if (known  && (cst = strideExpr.dyn_cast<AffineConstantExpr>())) {
>   makeComposedFolded(s0 * cst, {subStride})
> } else {
>   makeComposedFolded(s0 * s1, {subStride, origStride})
> }
> ```
> 
> The main problem here is that the number of operands is different (`subStride` vs. `subStride, origStride`), so getting rid of the if is not easy (at least I didn't see a nice way to do it.)
> 
> What I tried for that was something like:
> ```
> Expr = s0
> Operands = {subStride}
> If ()
>   Expr *= cst
> Else {
>   Expr *= s1
>  Operands.push_back(origStride)
> }
>  makeComposedFolded(s0 * Expr, Operands)
> ```
> 
> And I didn't find it particularly readable.
I'm having trouble seeing it in this form but there must be significantly simpler ways of writing this.
Let's apply the first round of cleanups and revisit?


================
Comment at: mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp:120
+
+      offset = makeComposedFoldedAffineApply(
+          rewriter, origLoc, s0 + s1 * s2, {offset, subOffsets[i], newStride});
----------------
Iteratively composing makes it hard to control your operands and indices here.

You could do something like:
```
SmallVector<Value> newStrides;
for (unsigned i = 0; i < sourceRank; ++i) {
  newStrides.push_back(
    makeComposedFoldedAffineApply(rewriter, origLoc, s0 * s1, {getOrigStrideAtIdx(i), subStrides[i]}));
}

SmallVector<Value> values(2 * sourceRank + 1);
SmallVector<AffineExpr> symbols(2 * sourceRank + 1);

// Note: you may need to impl. bindSymbols yourself as the existing version works for variadic templates only.
bindSymbols(symbols, ctx);
AffineExpr expr = symbols.front();
for (unsigned i = 0; i < sourceRank; ++i) {
  expr = expr + symbols[1 + 2 * i] * symbols[1 + 2 * i + 1]; 
  values.push_back(getOrigStrideAtIdx(i));
  values.push_back(subStrides[i]);
}
    
Value finalOffset = makeComposedFoldedAffineApply(rewriter, origLoc, expr, values));
```



================
Comment at: mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir:222
+//       CHECK: return %[[BASE]], %[[FINAL_OFFSET]], %[[DYN_SIZE0]], %[[DYN_SIZE1]], %[[DYN_SIZE2]], %[[FINAL_STRIDE0]], %[[FINAL_STRIDE1]], %[[FINAL_STRIDE2]]
+func.func @extract_strided_metadata_of_subview_all_dynamic(%base: memref<?x?x?xf32, strided<[?,?,?], offset:?>>, %offset0: index, %offset1: index, %offset2: index, %size0: index, %size1: index, %size2: index, %stride0: index, %stride1: index, %stride2: index) -> (memref<f32>, index, index, index, index, index, index, index) {
+%subview = memref.subview %base[%offset0, %offset1, %offset2][%size0, %size1, %size2][%stride0, %stride1, %stride2] : memref<?x?x?xf32, strided<[?,?,?], offset: ?>> to memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
----------------
can you please manually reflow and align?


Repository:
  rG LLVM Github Monorepo

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

https://reviews.llvm.org/D133166