[mlir] Prevent assertion failure in DropUnitDims

[tpopp]

Summary

Don't assert fail on strided memrefs when dropping unit dims.
Instead just leave them unchanged.

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

289

This seems a bit hard to follow. Can we instead make this return an Optional<...> and handle the case where nothing changes at the call site?

[tpopp]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

289

The caller takes these data structures and uses them at quite a few places and most of them need 1:1 mappings with the original input/outputs, so the alternative of returning an Optional would move this code to the caller. The caller would still need to generate this same information. There's no simple way to change that due to the caller needing the same number of data structures as the original operation had.

Given that this information needs to be generated regardless, I think it's better to just return it in here rather than have an additional if/else statement in the caller. This function can already return this NoOp like result also, so this doesn't result in a situation for the caller that didn't exist before.

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

289

I find it hard to maintain/read code if there is a whole data structure that models the no-op. It is much cleaner IMO to have to caller explicitly deal with the no-op situation than having a no-op being encoded in data structure values. I dont really see too much of an if-then-else situation on the caller side. Even if it does, it is then explicit in code that the caller is dealing with the no-op situation.

[gysit]

I added some minor comments that should be orthogonal to where the nop handling happens.

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

291–296

nit: AFAIK there is a constructor taking a single affine expression so the makeArrayRef can be dropped. It may also be an idea to start from an identity map to make this a bit simpler. However, since we need AffineMapAttr there is probably no substantially shorter solution here...

299–300

Can we use the indexingMap directly in this case?

308

nit: As long as this is the single instance I would inline the neutral result generation.

338–339

nit: I would probably delete the assertion since we handle the case above?

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

795

nit: %0 should probably be replaced by a variable matching the collapse_shape result

[tpopp]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

289

I'm doing it as you wish, but just to note, this is not making the caller handle the no-op situation. This function handles the no-op case. Then there is a separate unsupported case that results in a no-op. This is making the caller handle the unsupported situation while this function handles the no-op situation.

291–296

I did as you suggested for the single affine expression. I did not see any way to start with an identity map, as a given identity map needed to be N->N for any value N in [0,D), and the given functionalities for identity maps created a mapping for every value in a suffix of [0,D) only.

[tpopp]

Move the unsupported logic to the caller.

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

static UnitExtentReplacementInfo replaceUnitExtents(GenericOp genericOp,
   SmallVector<int64_t> newShape;
 
   int64_t origRank = genericOp.getRank(opOperand);
   AffineExpr zeroExpr = getAffineConstantExpr(0, context);
   auto isUnitExtent = [&](int64_t dim) -> bool {
     return shape[dim] == 1 && exprs[dim] == zeroExpr;
   };
 
+  // Provide a struct that represents an unchanged replacement result. This is
+  // used in place of error conditions.
+  Type actualType = opOperand->get().getType();
+  auto generateDoNotChange = [&]() -> UnitExtentReplacementInfo {
+    auto reassociationMaps = llvm::to_vector<8>(llvm::map_range(
+        llvm::seq<int64_t>(0, origRank), [&](int64_t dim) -> Attribute {
+          return AffineMapAttr::get(AffineMap::get(
+              origRank, /*symbolCount = */ 0,
+              llvm::makeArrayRef(getAffineDimExpr(dim, context)), context));
+        }));
+
+    return {actualType,
+            AffineMap::get(indexingMap.getNumDims(),
+                           indexingMap.getNumSymbols(), exprs, context),
+            ArrayAttr::get(context, reassociationMaps)};
+  };
+
+  // Early return for memrefs with affine maps to represent that we will always
+  // leave them unchanged.
+  if (auto memref = actualType.dyn_cast<MemRefType>()) {
+    if (!memref.getAffineMaps().empty())
+      return generateDoNotChange();
+  }
+
   int64_t dim = 0;
   // Fold dimensions that are unit-extent at the beginning of the tensor.
   while (dim < origRank && isUnitExtent(dim))
     reassociations.push_back(getAffineDimExpr(dim++, context));
   while (dim < origRank) {
     reassociations.push_back(getAffineDimExpr(dim, context));
     newIndexExprs.push_back(exprs[dim]);
     newShape.push_back(shape[dim]);
     // Fold all following dimensions that are unit-extent.
     while (dim + 1 < origRank && isUnitExtent(dim + 1)) {
       ++dim;
       reassociations.push_back(getAffineDimExpr(dim, context));
     }
     reassociationMaps.push_back(AffineMapAttr::get(AffineMap::get(
         origRank, /*symbolCount = */ 0, reassociations, context)));
     reassociations.clear();
     ++dim;
   }
+
   // Compute the tensor or scalar replacement type.
-  Type actualType = opOperand->get().getType();
   Type elementType = getElementTypeOrSelf(opOperand->get());
   Type replacementType;
   if (elementType == opOperand->get().getType()) {
     replacementType = elementType;
   } else if (actualType.isa<RankedTensorType>()) {
     replacementType = RankedTensorType::get(newShape, elementType);
   } else if (actualType.isa<MemRefType>()) {
     assert(actualType.cast<MemRefType>().getAffineMaps().empty() &&
            "unsupported strided memrefs");
     replacementType = MemRefType::get(newShape, elementType);
   }
   assert(replacementType && "unsupported shaped type");
   UnitExtentReplacementInfo info = {replacementType,
                                     AffineMap::get(indexingMap.getNumDims(),
                                                    indexingMap.getNumSymbols(),
                                                    newIndexExprs, context),
                                     ArrayAttr::get(context, reassociationMaps)};

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

 //       CHECK: linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP1]], #[[MAP2]]]
 //  CHECK-SAME:     iterator_types = ["reduction"]
 //  CHECK-SAME:   ins(%[[INPUT_RESHAPE]] : memref<1000xf32>)
 //  CHECK-SAME:   outs(%[[INIT_RESHAPE]] : memref<f32>)
 //       CHECK: return %[[INIT:.+]] : memref<1xf32>
 
 
+// -----
+// Test that nothing changes and no assertions are fired for memrefs with affine
+// maps while still changing the other operations.
+
+#map0 = affine_map<(d0, d1, d2)[s0] -> (d0 * s0 + d1 + d2)>
+
+#accesses = [
+  affine_map<(i, j, k, l, m) -> (i, k, m)>,
+  affine_map<(i, j, k, l, m) -> ()>,
+  affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>
+]
+
+#trait = {
+  iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],
+  indexing_maps = #accesses,
+  library_call = "some_external_func"
+}
+
+func @input_stays_same(%arg0 : memref<?x1x?xf32, #map0>, %arg1 : f32, %shape: memref<?x1x?x1x?xf32>) -> memref<?x1x?x1x?xf32> {
+  linalg.generic #trait
+     ins(%arg0, %arg1 : memref<?x1x?xf32, #map0>, f32)
+    outs(%shape : memref<?x1x?x1x?xf32>) {
+       ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) :
+         linalg.yield %arg3 : f32
+       }
+  return %shape : memref<?x1x?x1x?xf32>
+}
 
+// CHECK:     #[[MAP0:.*]] = affine_map<(d0, d1, d2)[s0] -> (d0 * s0 + d1 + d2)>
+// CHECK:     #[[MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)>
+// CHECK:     #[[MAP2:.*]] = affine_map<(d0, d1, d2) -> ()>
+// CHECK:     #[[MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK:     builtin.func @input_stays_same(
+// CHECK-SAME:  %[[ARG0:.*]]: memref<?x1x?xf32, #[[MAP0]]>,
+// CHECK-SAME:  %[[ARG1:.*]]: f32, %[[ARG2:.*]]: memref<?x1x?x1x?xf32>)
+// CHECK-SAME   -> memref<?x1x?x1x?xf32> {
+// CHECK:      memref.collapse_shape %[[ARG2]] {{\[}}[0, 1], [2, 3], [4]]
+// CHECK-SAME:   : memref<?x1x?x1x?xf32> into memref<?x?x?xf32>
+// CHECK:      linalg.generic
+// CHECK-SAME:   {indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]],
+// CHECK-SAME:   iterator_types = ["parallel", "parallel", "parallel"]}
+// CHECK-SAME:   ins(%[[ARG0]], %[[ARG1]] : memref<?x1x?xf32, #[[MAP0]]>, f32)
+// CHECK-SAME:   outs(%0 : memref<?x?x?xf32>) {
+// CHECK:      ^bb0(%{{.*}}: f32, %[[ARG:.*]]: f32, %{{.*}}: f32): // no predecessors
+// CHECK:       linalg.yield %[[ARG]] : f32
+// CHECK:      }
+// CHECK:      return %[[ARG2]] : memref<?x1x?x1x?xf32>