[mlir][Linalg] Fix crash in `dropUnitDims`

[springerm]

Summary

dropUnitDims is a transformation that rank-reduces the shape of operands of Linalg ops: dimensions of size 1 are removed. The transformation used to crash or produce invalid IR when dropping a dynamic unit dimensions. I.e., the dimension was statically detected to have size 1 because the dim size of another operand (mapped to the same loop) is a static 1.

This fixes #64977.

[mravishankar]

Could you post some more info on the crash? Crash is not good, so that is worth fixing, but IMO if the two dimensions are equal then one of them shouldn't be dynamic. I would expect the canonicalization to make both dimensions equal before calling these patterns. Having every pattern account for such inconsistencies is going to be very onerous. But it really depends on the crash. So some more info on the crash would help in review. Thanks!

[springerm]

There are two cases.

1. When lowering to tensor.collapse_shape:

error: 'memref.collapse_shape' op collapsed dim (0) must be dynamic if and only if reassociation group is dynamic
%1 = "memref.collapse_shape"(%arg0) <{reassociation = [[0, 1], [2]]}> : (memref<?x28x4xf16>) -> memref<28x4xf16>

We could also fix this by making the result type memref<?x4xf16>, but we would loose some static type information.

2. When lowering to tensor.extract_slice:

assert.h assertion failed at third_party/llvm/llvm-project/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp:293 in Value collapseValue(RewriterBase &, Location, Value, ArrayRef<int64_t>, ArrayRef<ReassociationIndices>, ControlDropUnitDims::RankReductionStrategy): succeeded(rankReducingExtract) && "not a unit-extent collapse"

Same thing, we know statically that the dimension is 1, but it is actually dynamic. It may be as simple as removing the assertion.

The canonicalizer does currently not promote the ? to 1, but I think it would make sense to do that (and ignore operands with ? dims in this transformation).

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

 /// whose unit dims are being dropped. Return the new indexing map
 /// to use, the shape of the operand in the replacement op
 /// and the `reassocation` to use to go from original operand shape
 /// to modified operand shape.
 struct UnitExtentReplacementInfo {
   AffineMap indexMap;
   SmallVector<ReassociationIndices> reassociation;
   SmallVector<int64_t> targetShape;
+  SmallVector<int64_t> collapsedDims;
 };
+
 static UnitExtentReplacementInfo dropUnitExtentFromOperandMetadata(
     MLIRContext *context, GenericOp genericOp, OpOperand *opOperand,
     llvm::SmallDenseMap<unsigned, unsigned> &oldDimsToNewDimsMap,
     ArrayRef<AffineExpr> dimReplacements) {
   UnitExtentReplacementInfo info;
   ReassociationIndices reassociationGroup;
   SmallVector<AffineExpr> newIndexExprs;
   AffineMap indexingMap = genericOp.getMatchingIndexingMap(opOperand);
auto isUnitDim = [&](unsigned dim) {
     if (operandShape[dim] == 1) {
       auto constAffineExpr = exprs[dim].dyn_cast<AffineConstantExpr>();
       return constAffineExpr && constAffineExpr.getValue() == 0;
     }
     return false;
   };
 
   unsigned dim = 0;
-  while (dim < operandShape.size() && isUnitDim(dim))
+  while (dim < operandShape.size() && isUnitDim(dim)) {
+    info.collapsedDims.push_back(dim);
     reassociationGroup.push_back(dim++);
+  }
   while (dim < operandShape.size()) {
     assert(!isUnitDim(dim) && "expected non unit-extent");
     reassociationGroup.push_back(dim);
     AffineExpr newExpr = exprs[dim].replaceDims(dimReplacements);
     newIndexExprs.push_back(newExpr);
     info.targetShape.push_back(operandShape[dim]);
     ++dim;
     // Fold all following dimensions that are unit-extent.
     while (dim < operandShape.size() && isUnitDim(dim)) {
+      info.collapsedDims.push_back(dim);
       reassociationGroup.push_back(dim++);
     }
     info.reassociation.push_back(reassociationGroup);
     reassociationGroup.clear();
   }
   info.indexMap =
       AffineMap::get(oldDimsToNewDimsMap.size(), indexingMap.getNumSymbols(),
                      newIndexExprs, context);
   return info;
 }
 
+/// Cast the given shaped value to the specified shape.
+static Value castToShape(OpBuilder &b, Value shapedValue,
+                         ArrayRef<int64_t> shape) {
+  if (auto tensorType = dyn_cast<RankedTensorType>(shapedValue.getType())) {
+    auto targetType = RankedTensorType::get(shape, tensorType.getElementType(),
+                                            tensorType.getEncoding());
+    return b.create<tensor::CastOp>(shapedValue.getLoc(), targetType,
+                                    shapedValue);
+  }
+
+  auto memrefType = cast<MemRefType>(shapedValue.getType());
+  MemRefLayoutAttrInterface layout;
+  auto targetType = MemRefType::get(shape, memrefType.getElementType(), layout,
+                                    memrefType.getMemorySpace());
+  return b.create<memref::CastOp>(shapedValue.getLoc(), targetType,
+                                  shapedValue);
+}
+
 LogicalResult linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
                                    const ControlDropUnitDims &options) {
   SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
   if (indexingMaps.empty())
     return failure();
 
   // 1. Check if any of the iteration dimensions are unit-trip count. They will
   //    end up being unit-trip count if they are used to index into a unit-dim
LogicalResult linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
   // Note that the indexing maps might need changing even if there are no
   // unit dimensions that are dropped to handle cases where `0` is used to
   // access a unit-extent tensor. Consider moving this out of this specific
   // transformation as a stand-alone transformation. Kept here right now due
   // to legacy.
   SmallVector<AffineMap> newIndexingMaps;
   SmallVector<SmallVector<ReassociationIndices>> reassociations;
   SmallVector<SmallVector<int64_t>> targetShapes;
-  SmallVector<bool> collapsed;
+  SmallVector<SmallVector<int64_t>> collapsedDims;
   auto hasCollapsibleType = [](OpOperand &operand) {
     Type operandType = operand.get().getType();
     if (auto memrefOperandType = dyn_cast_or_null<MemRefType>(operandType)) {
       return memrefOperandType.getLayout().isIdentity();
     } else if (auto tensorOperandType =
                    dyn_cast<RankedTensorType>(operandType)) {
       return tensorOperandType.getEncoding() == nullptr;
     }
     return false;
   };
   for (OpOperand &opOperand : genericOp->getOpOperands()) {
     auto indexingMap = genericOp.getMatchingIndexingMap(&opOperand);
     ArrayRef<int64_t> shape = genericOp.getShape(&opOperand);
     if (!hasCollapsibleType(opOperand)) {
       AffineMap newIndexingMap = indexingMap.replaceDimsAndSymbols(
           dimReplacements, ArrayRef<AffineExpr>{}, oldDimToNewDimMap.size(), 0);
       newIndexingMaps.push_back(newIndexingMap);
       targetShapes.push_back(llvm::to_vector(shape));
-      collapsed.push_back(false);
+      collapsedDims.push_back({});
       reassociations.push_back({});
       continue;
     }
     auto replacementInfo = dropUnitExtentFromOperandMetadata(
         rewriter.getContext(), genericOp, &opOperand, oldDimToNewDimMap,
         dimReplacements);
     reassociations.push_back(replacementInfo.reassociation);
     newIndexingMaps.push_back(replacementInfo.indexMap);
     targetShapes.push_back(replacementInfo.targetShape);
-    collapsed.push_back(!(replacementInfo.indexMap.getNumResults() ==
-                          indexingMap.getNumResults()));
+    collapsedDims.push_back(replacementInfo.collapsedDims);
   }
 
   // Abort if the indexing maps of the result operation are not invertible
   // (i.e. not legal) or if no dimension was reduced.
   if (newIndexingMaps == indexingMaps ||
       !inversePermutation(concatAffineMaps(newIndexingMaps)))
     return failure();
 
-  Location loc = genericOp.getLoc();
   // 4. For each of the operands, collapse the operand to convert
   //    from original shape to shape in the modified operation if needed,
   //    either through use of reshapes or rank-reducing slices as
   //    specified in `options`.
+  Location loc = genericOp.getLoc();
   SmallVector<Value> newOperands;
   for (OpOperand &opOperand : genericOp->getOpOperands()) {
     int64_t idx = opOperand.getOperandNumber();
-    if (!collapsed[idx]) {
+
+    // Check if there are any dims to collapse.
+    if (collapsedDims[idx].empty()) {
       newOperands.push_back(opOperand.get());
       continue;
     }
-    newOperands.push_back(collapseValue(rewriter, loc, opOperand.get(),
+
+    // Look for operand dims that should be collapsed but have a dynamic size.
+    // Cast such dims to static `1`.
+    SmallVector<int64_t> operandShape(
+        cast<ShapedType>(opOperand.get().getType()).getShape());
+    bool castNeeded = false;
+    for (int64_t i = 0; i < collapsedDims[idx].size(); ++i) {
+      auto dim = collapsedDims[idx][i];
+      if (operandShape[dim] != 1) {
+        assert(ShapedType::isDynamic(operandShape[dim]) &&
+               "expected dynamic dim size");
+        operandShape[dim] = 1;
+        castNeeded = true;
+      }
+    }
+
+    Value source = castNeeded
+                       ? castToShape(rewriter, opOperand.get(), operandShape)
+                       : opOperand.get();
+    newOperands.push_back(collapseValue(rewriter, loc, source,
                                         targetShapes[idx], reassociations[idx],
                                         options.rankReductionStrategy));
   }
 
   // 5. Create the `linalg.generic` operation with the new operands,
   //    indexing maps, iterator types and result types.
   ArrayRef<Value> newInputs =
       ArrayRef<Value>(newOperands).take_front(genericOp.getNumDpsInputs());
   ArrayRef<Value> newOutputs =
       ArrayRef<Value>(newOperands).take_back(genericOp.getNumDpsInits());
   SmallVector<Type> resultTypes;
   resultTypes.reserve(genericOp.getNumResults());
   for (unsigned i : llvm::seq<unsigned>(0, genericOp.getNumResults()))
     resultTypes.push_back(newOutputs[i].getType());
   GenericOp replacementOp =
       rewriter.create<GenericOp>(loc, resultTypes, newInputs, newOutputs,
                                  newIndexingMaps, newIteratorTypes);
   rewriter.inlineRegionBefore(genericOp.getRegion(), replacementOp.getRegion(),
                               replacementOp.getRegion().begin());
   // 5a. Replace `linalg.index` operations that refer to the dropped unit
   // dimensions.
   replaceUnitDimIndexOps(replacementOp, unitDims, rewriter);
 
   // 6. If any result type changes, insert a reshape/slice to convert from the
-  // original
-  //    type to the new type.
+  // original type to the new type.
   SmallVector<Value> resultReplacements;
   for (auto [index, result] : llvm::enumerate(replacementOp.getResults())) {
     unsigned opOperandIndex = index + replacementOp.getNumDpsInputs();
     Value origDest = genericOp.getDpsInitOperand(index)->get();
-    if (!collapsed[opOperandIndex]) {
+    if (collapsedDims[opOperandIndex].empty()) {
       resultReplacements.push_back(result);
       continue;
     }
     resultReplacements.push_back(expandValue(rewriter, loc, result, origDest,
                                              reassociations[opOperandIndex],
                                              options.rankReductionStrategy));
   }
 

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

 // CHECK-LABEL: func @drop_all_loops
 //       CHECK:   memref.collapse_shape
 //  CHECK-SAME:     [] : memref<1x1xf32, 3> into memref<f32, 3>
 //       CHECK:   linalg.generic{{.*}}memref<f32, 3>
 
 // CHECK-SLICES-LABEL: func @drop_all_loops
 //       CHECK-SLICES:   memref.subview %{{.*}}[0, 0] [1, 1] [1, 1] : memref<1x1xf32, 3> to memref<f32, strided<[]>, 3>
 //       CHECK-SLICES:   linalg.generic{{.*}}memref<f32, strided<[]>, 3>
+
+// -----
+
+#map4 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#map5 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+func.func @dynamic_dim_not_collapsible(%arg0: memref<?x28x4xf16>, %arg1: tensor<1x28xf16>) {
+  linalg.generic {indexing_maps = [#map4, #map5], iterator_types = ["parallel", "parallel", "parallel"]}
+      ins(%arg1 : tensor<1x28xf16>) outs(%arg0 : memref<?x28x4xf16>) {
+  ^bb0(%in: f16, %out: f16):
+    linalg.yield %in : f16
+  }
+  return
+}
+
+// CHECK-LABEL: func @dynamic_dim_not_collapsible(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?x28x4xf16>, %[[arg1:.*]]: tensor<1x28xf16>
+//       CHECK:   %[[c1:.*]] = tensor.collapse_shape %[[arg1]] {{\[}}[0, 1]] : tensor<1x28xf16> into tensor<28xf16>
+//       CHECK:   %[[cast:.*]] = memref.cast %[[arg0]] : memref<?x28x4xf16> to memref<1x28x4xf16>
+//       CHECK:   %[[c2:.*]] = memref.collapse_shape %[[cast]] {{\[}}[0, 1], [2]] : memref<1x28x4xf16> into memref<28x4xf16>
+//       CHECK:   linalg.generic {{.*}} ins(%[[c1]] : tensor<28xf16>) outs(%[[c2]] : memref<28x4xf16>)
+
+// CHECK-SLICES-LABEL: func @dynamic_dim_not_collapsible(
+//  CHECK-SLICES-SAME:     %[[arg0:.*]]: memref<?x28x4xf16>, %[[arg1:.*]]: tensor<1x28xf16>
+//       CHECK-SLICES:   %[[c1:.*]] = tensor.extract_slice %[[arg1]][0, 0] [1, 28] [1, 1] : tensor<1x28xf16> to tensor<28xf16>
+//       CHECK-SLICES:   %[[cast:.*]] = memref.cast %[[arg0]] : memref<?x28x4xf16> to memref<1x28x4xf16>
+//       CHECK-SLICES:   %[[c2:.*]] = memref.subview %[[cast]][0, 0, 0] [1, 28, 4] [1, 1, 1] : memref<1x28x4xf16> to memref<28x4xf16, strided<[4, 1]>>
+//       CHECK-SLICES:   linalg.generic {{.*}} ins(%[[c1]] : tensor<28xf16>) outs(%[[c2]] : memref<28x4xf16, strided<[4, 1]>>)