Support buffers in LinalgFoldUnitExtentDims

[tpopp]

Summary

This doesn't add any canonicalizations, but executes the same
simplification on bufferSemantic linalg.generic ops by using
linalg::ReshapeOp instead of linalg::TensorReshapeOp.

[tpopp]

Update a comment

[mravishankar]

Interesting that this makes it work for memref types. A quick glance at it, it looks mostly OK. One of the things though is that I am not sure how valid this will be when we have memref with strides. Maybe to be conservative, only do this when the operands dont have memrefs with affine maps?

[mravishankar]

Have a comment above (clicked enter early there). Will review more, but if we are going down this path, do we just replicate all the tests here for tensors with memrefs as well?

[tpopp]

In D103513#2794948, @mravishankar wrote:

Have a comment above (clicked enter early there). Will review more, but if we are going down this path, do we just replicate all the tests here for tensors with memrefs as well?

If it's agreed upon as an acceptable idea, I can create a clone of all relevant tests (with an additional mixed mode that would be supported in the latest patch). Primarily I don't want to do the additional work while there is risk of it being discarded.

In D103513#2794945, @mravishankar wrote:

Interesting that this makes it work for memref types. A quick glance at it, it looks mostly OK. One of the things though is that I am not sure how valid this will be when we have memref with strides. Maybe to be conservative, only do this when the operands dont have memrefs with affine maps?

Good catch. This is added. To add the additional check no longer creates the simple templating, so I needed to setup 2 classes with mostly shared functionality or to have one class do everything. I went with a single class because this also supports linalg.generics with mixed types.

Also, please feel free to add any other reviewers if they are appropriate. I wasn't sure who was appropriate as refactors have created many authors of this file.

[tpopp]

Create a single method to support both Tensor and MemRef operands.

[mravishankar]

LGTM-ing this. But needs to rebased on ToT before landing.

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

370

Nit: please use { } when the body spans multiple lines (even though its one statement)

371

This might not work on ToT now since reshape and tensor_reshape are split into collapse_* and expand_* variants. Suggest moving the reshape generation logic (to compute the reshapedValue ) into a separate method.

[tpopp]

Rebase, add tests, and extract reshape logic into separate methods.

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

LogicalResult matchAndRewrite(GenericOp genericOp,
     genericOp.iterator_typesAttr(ArrayAttr::get(context, newIteratorTypes));
     replaceUnitDimIndexOps(genericOp, unitDims, rewriter);
     rewriter.finalizeRootUpdate(genericOp);
     return success();
   }
 };
 
 struct UnitExtentReplacementInfo {
-  RankedTensorType type;
+  ShapedType type;
   AffineMap indexMap;
   ArrayAttr reassociation;
 };
 } // namespace
 
 /// Utility function for replacing operands/results to a linalg generic
 /// operation on tensors with unit-extent dimensions. These can be replaced with
 /// 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(AffineMap indexMap,
-                                                    RankedTensorType type,
+template <typename T>
+static UnitExtentReplacementInfo replaceUnitExtents(AffineMap indexMap, T type,
                                                     MLIRContext *context) {
   ArrayRef<int64_t> shape = type.getShape();
   ArrayRef<AffineExpr> exprs = indexMap.getResults();
   SmallVector<AffineExpr, 2> reassociations;
   SmallVector<Attribute, 4> reassociationMaps;
   SmallVector<AffineExpr, 4> newIndexExprs;
   SmallVector<int64_t, 4> newShape;
 
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;
   }
-  UnitExtentReplacementInfo info = {
-      RankedTensorType::get(newShape, type.getElementType()),
-      AffineMap::get(indexMap.getNumDims(), indexMap.getNumSymbols(),
+  UnitExtentReplacementInfo info = {T::get(newShape, type.getElementType()),
+                                    AffineMap::get(indexMap.getNumDims(),
+                                                   indexMap.getNumSymbols(),
                                                    newIndexExprs, context),
                                     ArrayAttr::get(context, reassociationMaps)};
   return info;
 }
 
 namespace {
 
 SmallVector<ReassociationExprs, 2>
 convertAffineMapArrayToExprs(ArrayAttr affineMapArrayAttr) {
   SmallVector<ReassociationExprs, 2> reassociationExprs;
   for (auto attr : affineMapArrayAttr)
     reassociationExprs.push_back(
         llvm::to_vector<4>(attr.cast<AffineMapAttr>().getValue().getResults()));
   return reassociationExprs;
 }
 
-/// Pattern to replace tensors operands/results that are unit extents.
-struct ReplaceUnitExtentTensors : public OpRewritePattern<GenericOp> {
+/// Pattern to replace tensor/buffer operands/results that are unit extents.
+template <typename OperandType, typename ReshapeOp>
+struct ReplaceUnitExtents : public OpRewritePattern<GenericOp> {
   using OpRewritePattern<GenericOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
-    if (!genericOp.hasTensorSemantics())
-      return failure();
-
     MLIRContext *context = rewriter.getContext();
     Location loc = genericOp.getLoc();
 
     SmallVector<AffineMap, 4> newIndexingMaps;
     SmallVector<ArrayAttr, 4> reassociationMaps;
     SmallVector<ShapedType, 4> newInputOutputTypes;
     bool doCanonicalization = false;
     for (auto it : llvm::zip(genericOp.getIndexingMaps(),
                              genericOp.getShapedOperandTypes())) {
-      auto replacementInfo = replaceUnitExtents(
-          std::get<0>(it), std::get<1>(it).template cast<RankedTensorType>(),
-          context);
+      auto operand = std::get<1>(it).template dyn_cast<OperandType>();
+      if (!operand)
+        return failure();
+
+      auto replacementInfo =
+          replaceUnitExtents(std::get<0>(it), operand, context);
       reassociationMaps.push_back(replacementInfo.reassociation);
       newIndexingMaps.push_back(replacementInfo.indexMap);
       newInputOutputTypes.push_back(replacementInfo.type);
       doCanonicalization |= replacementInfo.type != std::get<1>(it);
     }
 
     // If the indexing maps of the result operation are not invertible (i.e. not
     // legal), abort.
     if (!doCanonicalization ||
         !inversePermutation(concatAffineMaps(newIndexingMaps)))
       return failure();
 
     // If any operand type change, insert a reshape to convert from the original
     // type to the new type.
     // TODO: get rid of flattenedIdx which assumes operand order and contiguity.
     unsigned flattenedIdx = 0;
     auto insertReshapes = [&](ValueRange values) {
       SmallVector<Value, 4> res;
       res.reserve(values.size());
       for (auto operand : llvm::enumerate(values)) {
         if (operand.value().getType() == newInputOutputTypes[flattenedIdx])
           res.push_back(operand.value());
         else
-          res.push_back(rewriter.create<linalg::TensorReshapeOp>(
+          res.push_back(rewriter.create<ReshapeOp>(
               loc, newInputOutputTypes[flattenedIdx], operand.value(),
               convertAffineMapArrayToExprs(reassociationMaps[flattenedIdx])));
         ++flattenedIdx;
       }
       return res;
     };
 
     SmallVector<Value, 4> newInputs = insertReshapes(genericOp.inputs());
     SmallVector<Value, 4> newOutputs = insertReshapes(genericOp.outputs());
 
     // If any result type changes, insert a reshape to convert from the original
     // type to the new type.
     SmallVector<Type, 4> resultTypes;
     resultTypes.reserve(genericOp.getNumResults());
     for (unsigned i : llvm::seq<unsigned>(0, genericOp.getNumResults()))
       resultTypes.push_back(newInputOutputTypes[i + genericOp.getNumInputs()]);
     GenericOp replacementOp = rewriter.create<GenericOp>(
         loc, resultTypes, newInputs, newOutputs, newIndexingMaps,
         llvm::to_vector<4>(
             genericOp.iterator_types().template getAsValueRange<StringAttr>()));
     rewriter.inlineRegionBefore(genericOp.region(), replacementOp.region(),
                                 replacementOp.region().begin());
 
     // If any result tensor has a modified shape, then add reshape to recover
     // the original shape.
     SmallVector<Value, 4> resultReplacements;
     for (auto result : llvm::enumerate(replacementOp.getResults())) {
       unsigned index = result.index() + replacementOp.getNumInputs();
-      RankedTensorType origResultType = genericOp.getResult(result.index())
-                                            .getType()
-                                            .template cast<RankedTensorType>();
+      auto origResultType = genericOp.getResult(result.index())
+                                .template cast<OperandType>();
       if (origResultType != result.value().getType())
-        resultReplacements.push_back(rewriter.create<linalg::TensorReshapeOp>(
+        resultReplacements.push_back(rewriter.create<ReshapeOp>(
             loc, origResultType, result.value(),
             convertAffineMapArrayToExprs(reassociationMaps[index])));
       else
         resultReplacements.push_back(result.value());
     }
     rewriter.replaceOp(genericOp, resultReplacements);
     return success();
   }

 };
 } // namespace
 
 /// Patterns that are used to canonicalize the use of unit-extent dims for
 /// broadcasting.
 void mlir::linalg::populateFoldUnitExtentDimsPatterns(
     RewritePatternSet &patterns) {
   auto *context = patterns.getContext();
-  patterns.add<FoldUnitDimLoops, ReplaceUnitExtentTensors,
-               UseRankReducedSubTensorOp, UseRankReducedSubTensorInsertOp>(
-      context);
+  patterns
+      .add<FoldUnitDimLoops, ReplaceUnitExtents<MemRefType, linalg::ReshapeOp>,
+           ReplaceUnitExtents<RankedTensorType, linalg::TensorReshapeOp>,
+           UseRankReducedSubTensorOp, UseRankReducedSubTensorInsertOp>(context);
   TensorReshapeOp::getCanonicalizationPatterns(patterns, context);
 }
 
 namespace {
 /// Pass that removes unit-extent dims within generic ops.
 struct LinalgFoldUnitExtentDimsPass
     : public LinalgFoldUnitExtentDimsBase<LinalgFoldUnitExtentDimsPass> {
   void runOnFunction() override {

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

func @subtensor_insert_unit_dims(%arg0: tensor<1x3xf32>, %arg1: tensor<1x1xf32>) -> tensor<1x3xf32> {
   %0 = subtensor_insert %arg1 into %arg0[0, 2] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x3xf32>
   return %0 : tensor<1x3xf32>
 }
 // CHECK-LABEL: func @subtensor_insert_unit_dims
 //       CHECK:   %[[RESHAPE:.+]] = linalg.tensor_reshape %{{.+}} []
 //       CHECK:   %[[RESULT:.+]] = subtensor_insert %[[RESHAPE]]
 //  CHECK-SAME:     tensor<f32> into tensor<1x3xf32>
 //       CHECK:   return %[[RESULT]]
+
+// -----
+
+#accesses = [
+  affine_map<(d0, d1) -> (0, 0)>,
+  affine_map<(d0, d1) -> (d0, d1)>
+]
+
+#trait = {
+  indexing_maps = #accesses,
+  iterator_types = ["parallel", "parallel"],
+  library_call = "some_external_fn"
+}
+
+// Confirm the linalg.generic patterns also work on memrefs
+func @broadcast_scalar_memref(%arg0 : memref<1x1xf32>, %shape : memref<?x?xf32>) -> memref<?x?xf32>
+{
+   linalg.generic #trait
+     ins(%arg0 : memref<1x1xf32>)
+    outs(%shape : memref<?x?xf32>) {
+      ^bb0(%arg2 : f32, %arg3 : f32):
+        linalg.yield %arg2 : f32
+   }
+   return %shape : memref<?x?xf32>
+}
+//   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> ()>
+//   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @broadcast_scalar_memref
+//  CHECK-SAME:   %[[ARG0:.*]]: memref<1x1xf32>
+//       CHECK:   %[[A:.*]] = linalg.reshape %[[ARG0]] []
+//  CHECK-SAME:     memref<1x1xf32> into memref<f32>
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+//  CHECK-SAME:     iterator_types = ["parallel", "parallel"]
+//  CHECK-SAME:     %[[A]]
+
+