[PATCH] D79766: [mlir][Linalg] Add pass to remove unit-extent dims from tensor operands of Generic ops.

[Nicolas Vasilache via Phabricator via llvm-commits]

nicolasvasilache accepted this revision.
nicolasvasilache added a comment.
This revision is now accepted and ready to land.

Thanks!

Note the last minor points.



================
Comment at: mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp:100
+/// canonicalized op with these dims removed, given the original `indexingMaps`.
+static ArrayAttr replaceUnitDims(DenseSet<unsigned> &unitDims,
+                                 ArrayRef<AffineMap> indexingMaps,
----------------
It seems like this whole function could be replaced by:

1. create an affine map that either sets unitDims to 0 or drops them (depending on whether you want a projection or explictly set to 0).
2. compose with the `indexingMaps`

?


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp:142
+    return nullptr;
+  return ArrayAttr::get(
+      llvm::to_vector<4>(llvm::map_range(
----------------
We have `ArrayAttr::getAsRange`, seems a similar `ArrayAttr::getFromRange` would be useful?
No need to do it in this CL if this is too intrusive.


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp:175
+        if (dims[dimExpr.getPosition()] == 1 &&
+            iteratorTypes[expr.index()].dyn_cast<StringAttr>().getValue() ==
+                getParallelIteratorTypeName())
----------------
I am curious why does the iterator type matter here?


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp:192
+    SmallVector<Attribute, 4> newIteratorTypes;
+    for (auto attr : llvm::enumerate(iteratorTypes)) {
+      if (unitDims.count(attr.index()))
----------------
nit:
```
    for (auto attr : llvm::enumerate(iteratorTypes))
      if (unitDims.count(attr.index()) == 0)
        newIteratorTypes.push_back(attr.value());
```
?


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp:215
+/// operation on tensors with unit-extent dimensions. These can be replaced with
+/// an operand/result with the unit-extent dimenion removed. This is only done
+/// if the indexing map used to access that dimension has a AffineConstantExpr
----------------
typo: dimension


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp:243
+    reassociations.push_back(getAffineDimExpr(dim++, context));
+  while (dim < origRank) {
+    reassociations.push_back(getAffineDimExpr(dim, context));
----------------
Thanks, much simpler to follow AFAIR !


================
Comment at: mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir:166
+//  CHECK-SAME:     %[[A]]
\ No newline at end of file

----------------
nit:newline 


Repository:
  rG LLVM Github Monorepo

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

https://reviews.llvm.org/D79766