[mlir][Linalg] Disable fusion of reshape with elementwise ops for purely dynamic cases.

[mravishankar]

Summary

tensor.collapse_shape(tensor.expand_shape) op when fused with a
consumer(producer) elementwise linalg.generic operation results in
creation of tensor.expand_shape ops. In purely dynamic cases this
can end up with a dynamic dimensions being expanded to more than one
dynamic dimension. This is disallowed by the semantics of
tensor.expand_shape operation. (While the transformation is itself
correct, its a gap in the specification of tensor.expand_shape that
is the issue). So disallow fusions which result in such a pattern.

[ThomasRaoux]

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

628–632

Would the op be valid in this case? Should we have a test for this case if it is possible?

[mravishankar]

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

628–632

The reshape op that would be generated is not valid (due to shape propogation related concerns, but not actually illegal), but the transformation is still valid. The test case below checks that the fusion does not happen, and illegal (as of today) reshape op is not created.

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

 public:
   // Computes the mapping from original dimensions of the op to the dimensions
   // of the expanded op given the `indexingMap` of the fused operand/result of
   // the generic op, the `reassocationMaps` of the reshape op and the shape of
   // the expanded op.
   LogicalResult compute(LinalgOp linalgOp, OpOperand *fusableOpOperand,
                         ArrayRef<AffineMap> reassociationMaps,
                         ArrayRef<int64_t> expandedShape,
+                        ArrayRef<int64_t> collapsedShape,
                         PatternRewriter &rewriter);
   unsigned getOrigOpNumDims() const { return reassociation.size(); }
   unsigned getExpandedOpNumDims() const { return expandedOpNumDims; }
   ReassociationIndicesRef getExpandedDims(unsigned i) const {
     return reassociation[i];
   }
   ArrayRef<int64_t> getExpandedShapeOfDim(unsigned i) const {
     return expandedShapeMap[i];
   }
+  ArrayRef<int64_t> getOriginalShape() const { return originalLoopExtent; }
 
 private:
   /// Reassociation from the dimensions in the original operation to the
   /// dimension of the expanded operation.
   SmallVector<ReassociationIndices> reassociation;
   /// Mapping from extent of loops in the original operation, to the extent of
   /// loops in the expanded operation.
   SmallVector<SmallVector<int64_t>> expandedShapeMap;
+  /// Extent of the loop in the original operation.
+  SmallVector<int64_t> originalLoopExtent;
   unsigned expandedOpNumDims;
 };
 } // namespace
 
 LogicalResult ExpansionInfo::compute(LinalgOp linalgOp,
                                      OpOperand *fusableOpOperand,
                                      ArrayRef<AffineMap> reassociationMaps,
                                      ArrayRef<int64_t> expandedShape,
+                                     ArrayRef<int64_t> collapsedShape,
                                      PatternRewriter &rewriter) {
   if (reassociationMaps.empty())
     return failure();
   AffineMap fusedIndexMap = linalgOp.getTiedIndexingMap(fusableOpOperand);
 
   Optional<SmallVector<int64_t, 4>> originalLoopRange =
       linalgOp.getStaticLoopRanges();
   if (!originalLoopRange)
     return rewriter.notifyMatchFailure(linalgOp, "unable to find loop range");
+  originalLoopExtent.assign(originalLoopRange->begin(),
+                            originalLoopRange->end());
 
   reassociation.clear();
   expandedShapeMap.clear();
   // Compute the number of dimension in the expanded op that correspond to each
   // dimension of the original op.
   SmallVector<unsigned> numExpandedDims(fusedIndexMap.getNumDims(), 1);
   expandedShapeMap.resize(fusedIndexMap.getNumDims());
   for (const auto &resultExpr : llvm::enumerate(fusedIndexMap.getResults())) {
     unsigned pos = resultExpr.value().cast<AffineDimExpr>().getPosition();
     AffineMap foldedDims = reassociationMaps[resultExpr.index()];
     numExpandedDims[pos] = foldedDims.getNumResults();
     ArrayRef<int64_t> shape =
         expandedShape.slice(foldedDims.getDimPosition(0), numExpandedDims[pos]);
     expandedShapeMap[pos].assign(shape.begin(), shape.end());
   }
   // The remaining dimensions remain the same.
   for (unsigned i : llvm::seq<unsigned>(0, fusedIndexMap.getNumDims()))
     if (expandedShapeMap[i].empty())
-      expandedShapeMap[i] = {(*originalLoopRange)[i]};
+      expandedShapeMap[i] = {originalLoopExtent[i]};
 
   // Compute reassociation map from the original op to the expanded op.
   unsigned sum = 0;
   reassociation.reserve(fusedIndexMap.getNumDims());
   for (const auto &numFoldedDim : llvm::enumerate(numExpandedDims)) {
     auto seq = llvm::seq<int64_t>(sum, sum + numFoldedDim.value());
     reassociation.emplace_back(seq.begin(), seq.end());
     sum += numFoldedDim.value();
   }
   expandedOpNumDims = sum;
   return success();
 }
 
 /// Epanding the body of a linalg operation requires adaptations of the accessed
 /// loop indices. Specifically, access of indices in the original operation need
 /// to be replaced with linearizations of indices in the expanded op. That
 /// requires the shape of the expanded dimensions to be static (at least all but
 /// the most significant). For now check that these are all statically sized.
 /// Note that this could be extended to handle dynamic case, but the
 /// implementation below uses `affine.apply` which seems to have issues when the
 /// shapes are not static.
 LogicalResult isGenericOpExpandable(GenericOp genericOp,
                                     const ExpansionInfo &expansionInfo,
                                     PatternRewriter &rewriter) {
+  // Current reshape only supports expansion of a dynamic dim when only one of
+  // the expanded dims are dynamic.
+  for (auto originalShape : llvm::enumerate(expansionInfo.getOriginalShape()))
+    if (ShapedType::isDynamic(originalShape.value())) {
+      // All but one of the expanded dims must be static.
+      bool foundDynamicExpandedDim = false;
+      for (auto expandedShape :
+           expansionInfo.getExpandedShapeOfDim(originalShape.index())) {
+        if (ShapedType::isDynamic(expandedShape)) {
+          if (foundDynamicExpandedDim) {
+            return rewriter.notifyMatchFailure(
+                genericOp,
+                "cannot expanded dynamic dims into multiple dynamic dims");
+          }
+          foundDynamicExpandedDim = true;
+        }
+      }
+      if (!foundDynamicExpandedDim) {
+        return rewriter.notifyMatchFailure(
+            genericOp, "dynamic dim expansion needs at least one dynamic dim "
+                       "in result shape");
+      }
+    }
+
   if (!genericOp.hasIndexSemantics())
     return success();
   for (unsigned i : llvm::seq<unsigned>(0, expansionInfo.getOrigOpNumDims())) {
     ArrayRef<int64_t> expandedShape = expansionInfo.getExpandedShapeOfDim(i);
     if (expandedShape.size() == 1)
       continue;
     for (int64_t shape : expandedShape.drop_front()) {
       if (ShapedType::isDynamic(shape)) {
assert(isFusableWithReshapeByDimExpansion(genericOp, fusableOpOperand) &&
          "preconditions for fuse operation failed");
   // Check if reshape is expanding or collapsing.
   auto expandingReshapeOp = dyn_cast<tensor::ExpandShapeOp>(*reshapeOp);
   auto collapsingReshapeOp = dyn_cast<tensor::CollapseShapeOp>(*reshapeOp);
   bool isExpanding = (expandingReshapeOp != nullptr);
   RankedTensorType expandedType = isExpanding
                                       ? expandingReshapeOp.getResultType()
                                       : collapsingReshapeOp.getSrcType();
+  RankedTensorType collapsedType = isExpanding
+                                       ? expandingReshapeOp.getSrcType()
+                                       : collapsingReshapeOp.getResultType();
 
   ExpansionInfo expansionInfo;
   if (failed(expansionInfo.compute(
           genericOp, fusableOpOperand,
           isExpanding ? expandingReshapeOp.getReassociationMaps()
                       : collapsingReshapeOp.getReassociationMaps(),
-          expandedType.getShape(), rewriter)))
+          expandedType.getShape(), collapsedType.getShape(), rewriter)))
     return llvm::None;
 
   if (failed(isGenericOpExpandable(genericOp, expansionInfo, rewriter)))
     return llvm::None;
 
   SmallVector<AffineMap, 4> expandedOpIndexingMaps = llvm::to_vector<4>(
       llvm::map_range(genericOp.getIndexingMaps(), [&](AffineMap m) {
         return getIndexingMapInExpandedOp(rewriter, m, expansionInfo);

mlir/test/Dialect/Linalg/reshape_fusion.mlir

 //         FOLDUNITDIM: func @unit_dim_reshape_expansion_full
 //    FOLDUNITDIM-SAME:   %[[ARG0:.+]]: tensor<1x?x1x2x1x4xf32>
 //    FOLDUNITDIM-SAME:   %[[ARG1:.+]]: tensor<?x2x4xf32>
 //     FOLDUNITDIM-DAG:   %[[RESHAPE:.+]] = tensor.expand_shape %[[ARG1]]
 //         FOLDUNITDIM:   linalg.generic
 //    FOLDUNITDIM-SAME:     ins(%[[ARG0]], %[[RESHAPE]] : tensor<1x?x1x2x1x4xf32>, tensor<1x?x1x2x1x4xf32>)
 //    FOLDUNITDIM-SAME:     outs(%{{.+}} : tensor<1x?x1x2x1x4xf32>)
 
+// -----
+
+func @no_fuse_dynamic_dims(%arg0: tensor<?x?xf32>) -> tensor<?xf32> {
+  %c0 = arith.constant 0 : index
+  %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<?x?xf32> into tensor<?xf32>
+  %1 = tensor.dim %0, %c0 : tensor<?xf32>
+  %2 = linalg.init_tensor [%1] : tensor<?xf32>
+  %3 = linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]}
+    ins(%0 : tensor<?xf32>) outs(%2 : tensor<?xf32>) {
+      ^bb0(%arg1 : f32, %arg2: f32):
+        %4 = arith.addf %arg1, %arg1 : f32
+        linalg.yield %4 : f32
+    } -> tensor<?xf32>
+  return %3 : tensor<?xf32>
+}
+//      CHECK: func @no_fuse_dynamic_dims
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<?x?xf32>
+//      CHECK:   %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]]
+//      CHECK:   %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME:       ins(%[[RESHAPE]] : tensor<?xf32>)
+//      CHECK:   return %[[GENERIC]]