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

[Johannes Reifferscheid via Phabricator via llvm-commits]

jreiffers added inline comments.


================
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");
----------------
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?


================
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");
+
----------------
Does this assertion do anything useful here?


================
Comment at: mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp:465
+          collapsedStride.push_back(origStrides[currentDim]);
+          break;
+        }
----------------
If you return here, you can drop the condition below.


================
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]);
----------------
Or maybe just do this now?


Repository:
  rG LLVM Github Monorepo

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

https://reviews.llvm.org/D139329