[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.

321–322

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

 /// an operand/result with the unit-extent dimension removed. This is only done
 /// if the indexing map used to access that didimensionmension has a
 /// AffineConstantExpr of value 0. Given the `type` of an result/operand of a
 /// Linalg op, and its `indexMap` the utility function returns:
 /// - the new type with dimensions of size 1 removed.
 /// - modified index map that can be used to access the replaced result/operand
 /// - the reassociation that converts from the original tensor type to the
 ///   modified tensor type.
-static UnitExtentReplacementInfo replaceUnitExtents(GenericOp genericOp,
-                                                    OpOperand *opOperand,
-                                                    MLIRContext *context) {
+static llvm::Optional<UnitExtentReplacementInfo>
+replaceUnitExtents(GenericOp genericOp, OpOperand *opOperand,
   AffineMap indexingMap = genericOp.getTiedIndexingMap(opOperand);
   ArrayRef<int64_t> shape = genericOp.getShape(opOperand);
   ArrayRef<AffineExpr> exprs = indexingMap.getResults();
   SmallVector<AffineExpr> reassociations;
   SmallVector<Attribute> reassociationMaps;
   SmallVector<AffineExpr> newIndexExprs;
   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;
   };
 
+  // Early return for memrefs with affine maps to represent that we will always
+  // leave them unchanged.
+  Type actualType = opOperand->get().getType();
+  if (auto memref = actualType.dyn_cast<MemRefType>()) {
+    if (!memref.getAffineMaps().empty())
+      return llvm::None;
+  }
+
   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)};
   return info;
LogicalResult matchAndRewrite(GenericOp genericOp,
     MLIRContext *context = rewriter.getContext();
     Location loc = genericOp.getLoc();
 
     SmallVector<AffineMap> newIndexingMaps;
     SmallVector<ArrayAttr> reassociationMaps;
     SmallVector<Type> newInputOutputTypes;
     bool doCanonicalization = false;
     for (OpOperand *opOperand : genericOp.getInputAndOutputOperands()) {
-      UnitExtentReplacementInfo replacementInfo =
-          replaceUnitExtents(genericOp, opOperand, context);
-      reassociationMaps.push_back(replacementInfo.reassociation);
-      newIndexingMaps.push_back(replacementInfo.indexMap);
-      newInputOutputTypes.push_back(replacementInfo.type);
-      doCanonicalization |= replacementInfo.type != opOperand->get().getType();
+      auto replacementInfo = replaceUnitExtents(genericOp, opOperand, context);
+      if (replacementInfo) {
+        reassociationMaps.push_back(replacementInfo->reassociation);
+        newIndexingMaps.push_back(replacementInfo->indexMap);
+        newInputOutputTypes.push_back(replacementInfo->type);
+        doCanonicalization |=
+            replacementInfo->type != opOperand->get().getType();
+      } else {
+        // If replaceUnitExtents cannot handle this case, maintain the same
+        // type, indexing map, and create a set of mappings representing an
+        // identity matrix.
+        newInputOutputTypes.push_back(opOperand->get().getType());
+        newIndexingMaps.push_back(genericOp.getTiedIndexingMap(opOperand));
+        int64_t origRank = genericOp.getRank(opOperand);
+        auto maps = 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,
+                                 getAffineDimExpr(dim, context), context));
+            }));
+        reassociationMaps.push_back(ArrayAttr::get(context, maps));
+      }
     }
 
     // If the indexing maps of the result operation are not invertible (i.e. not
     // legal), abort.
     if (!doCanonicalization ||
         !inversePermutation(concatAffineMaps(newIndexingMaps)))
       return failure();
 

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:      %[[OUT:.*]] = 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(%[[OUT]] : 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>