[PATCH] D139329: [mlir][ExpandStridedMetadata] Handle collapse_shape of dim of size 1 gracefully

[Quentin Colombet via Phabricator via llvm-commits]

qcolombet added a comment.

Thanks @jreiffers for the review!



================
Comment at: mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp:447-453
+    SmallVector<int64_t> collapsedStrides;
+    int64_t collapsedOffset;
+    bool hasKnownCollapsedStridesAndOffset = succeeded(
+        getStridesAndOffset(collapsedType, collapsedStrides, collapsedOffset));
+    (void)hasKnownCollapsedStridesAndOffset;
+    assert(hasKnownStridesAndOffset &&
+           "getStridesAndOffset must work on valid collapse_shape");
----------------
jreiffers wrote:
> This pattern is repeated multiple times in this file - maybe introduce a function returning a SmallVector<int64_t> and handling the assertion to make the call sites easier to read?
Good idea!


================
Comment at: mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp:459-460
+      for (int64_t currentDim : reassocGroup) {
+        assert(srcShape[currentDim] == 1 &&
+               "We should be dealing with 1x1x...x1");
+
----------------
jreiffers wrote:
> Does this assertion do anything useful here?
That's debatable.
Essentially I wanted to make clear that this code is not intended to be used in other cases and futur proof a potential refactoring.

Basically, comments are easy to miss, asserts have more teeth :).


================
Comment at: mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp:465
+          collapsedStride.push_back(origStrides[currentDim]);
+          break;
+        }
----------------
jreiffers wrote:
> If you return here, you can drop the condition below.
Let me play a bit with the code layout.

What I liked with the current solution was the assert below, but I can turn this into an unreachable if we early return here.


================
Comment at: mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp:473
+  } else if (groupStrides.size() == 1) {
+    // TODO: affine.min should probably fold `min(a)` => `a`.
+    collapsedStride.push_back(groupStrides[0]);
----------------
jreiffers wrote:
> Or maybe just do this now?
I feel that this should be its own PR.

In the meantime, I could drop this specialization and temporarily generate `min(a)`, that'll probably be clearer.

What do you think?


Repository:
  rG LLVM Github Monorepo

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

https://reviews.llvm.org/D139329