[mlir][Linalg] Add producer-consumer fusion when producer is a ConstantOp and Consumer is a GenericOp.

[nicolasvasilache]

How much of this special casing is impacted by the lack of existence (today) of the tied input/output tensor.
It seems to me once we have that, the constant op would (roughly) look like:

%1 = linalg.generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#map0], iterator_types = ["parallel"}] %0 {
  ^bb0(%arg1: f32, %arg2: f32):
     %cst = constant 42.0: f32
     linalg.yield %cst : f32
}: tensor<5xf32> -> tensor<5xf32>

And we would use a generic form of fusion.
I think it is fine for now to unblock progress but, if you agree, we should mark this as TBD.

[mravishankar]

In D79838#2037821, @nicolasvasilache wrote:

How much of this special casing is impacted by the lack of existence (today) of the tied input/output tensor.
It seems to me once we have that, the constant op would (roughly) look like:

%1 = linalg.generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#map0], iterator_types = ["parallel"}] %0 {
  ^bb0(%arg1: f32, %arg2: f32):
     %cst = constant 42.0: f32
     linalg.yield %cst : f32
}: tensor<5xf32> -> tensor<5xf32>

And we would use a generic form of fusion.
I think it is fine for now to unblock progress but, if you agree, we should mark this as TBD.

I dont think this is related to tied input/output tensors. I might be missing something here, but the tied input/output tensor is for initialization sorts of things (a symptom of such a thing is when one of the iterator types is "reduction"). There are cases I have seen in IREE where some operands to pointwise-add are splat constants (I dont know why that is the case though in the input though, havent dug into it)

[nicolasvasilache]

I dont think this is related to tied input/output tensors. I might be missing something here, but the tied input/output tensor is for initialization sorts of things (a symptom of such a thing is when one of the iterator types is "reduction"). There are cases I have seen in IREE where some operands to pointwise-add are splat constants (I dont know why that is the case though in the input though, havent dug into it)

The point is that one would translate ConstantOp to Linalg generic beforehand and not introduce op-dependent special-casing in fusion.
Special casing in transformations is what we want to avoid at all costs.

[mravishankar]

The point is that one would translate ConstantOp to Linalg generic beforehand and not introduce op-dependent special-casing in fusion.
Special casing in transformations is what we want to avoid at all costs.

I agree. I am moving this pattern from IREE into here. In IREE I was lowering a ConstantOp to a GenericOp and using the fusion between two generic Ops. But that resulted in a generic op with 0 input arguments and 1 output argument. In an offline conversation I gathered that this was unexpected and future verifier semantics might make that illegal. I also seemed to be having some really nasty segfaults which I wasnt able to root cause. This gets around that. I am happy to drop it in favor of lowering a constant op to generic op transformation and fusion. On balance I prefer this way cause both of these live in core and avoids additional hops. Totally agree that we should avoid special casing, but at the level of tensors through the special casing is very localized. So its code overhead and more patterns, but doesnt propogate past the pattern (still need to be cognizant that we shouldnt have a lot of such pairwise combinations)

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

auto fusedOp = rewriter.create<LinalgOpTy>(
         /*doc=*/nullptr,
         /*library_call=*/nullptr);
     auto &fusedRegion = fusedOp.region();
     rewriter.cloneRegionBefore(producer.region(), fusedRegion,
                                fusedRegion.begin());
     return fusedOp;
   }
 };
+
+/// Implementation of fusion on tensor ops when producer is a splat constant.
+template <typename LinalgOpTy> struct FuseConstantOpAsProducer {
+  static bool isFusible(ConstantOp producer, LinalgOpTy consumer,
+                        unsigned consumerIdx) {
+    return producer.getResult().getType().isa<RankedTensorType>() &&
+           producer.value().template cast<DenseElementsAttr>().isSplat();
+  }
+
+  static Operation *fuse(ConstantOp producer, LinalgOpTy consumer,
+                         unsigned consumerIdx, PatternRewriter &rewriter,
+                         OperationFolder *folder = nullptr) {
+    if (!isFusible(producer, consumer, consumerIdx))
+      return nullptr;
+
+    // The indexing_maps for the operands of the fused operation are same as
+    // those for the operands of the consumer without the indexing map at
+    // consumerIdx
+    SmallVector<AffineMap, 4> fusedIndexMaps =
+        llvm::to_vector<4>(llvm::map_range(
+            consumer.indexing_maps(), [](Attribute attr) -> AffineMap {
+              return attr.cast<AffineMapAttr>().getValue();
+            }));
+    fusedIndexMaps.erase(std::next(fusedIndexMaps.begin(), consumerIdx));
+
+    // The operands list is same as the consumer with the argument for constant
+    // index dropped.
+    SmallVector<Value, 4> fusedOperands(consumer.operand_begin(),
+                                        consumer.operand_end());
+    fusedOperands.erase(std::next(fusedOperands.begin(), consumerIdx));
+
+    // Create a constant scalar value from the splat constant.
+    Value scalarConstant = rewriter.create<ConstantOp>(
+        producer.getLoc(),
+        producer.value().template cast<DenseElementsAttr>().getSplatValue());
+
+    auto fusedOp = rewriter.create<LinalgOpTy>(
+        rewriter.getUnknownLoc(), consumer.getResultTypes(), fusedOperands,
+        rewriter.getI64IntegerAttr(consumer.getNumOperands() - 1),
+        rewriter.getI64IntegerAttr(consumer.getNumResults()),
+        rewriter.getAffineMapArrayAttr(fusedIndexMaps),
+        consumer.iterator_types(),
+        /*doc=*/nullptr,
+        /*library_call=*/nullptr);
+
+    // Map the block argument corresponding to the replaced argument with the
+    // scalar constant.
+    Region &consumerRegion = consumer.region();
+    Block &entryBlock = *consumerRegion.begin();
+    unsigned argIndex =
+        entryBlock.getNumArguments() - consumer.getNumOperands() + consumerIdx;
+    BlockAndValueMapping mapping;
+    mapping.map(entryBlock.getArgument(argIndex), scalarConstant);
+    Region &fusedRegion = fusedOp.region();
+    rewriter.cloneRegionBefore(consumerRegion, fusedRegion, fusedRegion.begin(),
+                               mapping);
+    return fusedOp;
+  }
+};
+
 } // namespace
 
 Operation *mlir::linalg::fuseTensorOps(PatternRewriter &rewriter,
                                        Operation *consumer,
                                        unsigned consumerIdx,
                                        OperationFolder *folder) {
   if (consumerIdx >= consumer->getNumOperands())
     return nullptr;
   Operation *producer = consumer->getOperand(consumerIdx).getDefiningOp();
   if (!producer || producer->getNumResults() != 1)
     return nullptr;
 
   // Fuse when consumer is GenericOp.
   if (GenericOp genericOp = dyn_cast<GenericOp>(consumer)) {
     if (!genericOp.hasTensorSemantics())
       return nullptr;
     if (auto genericOpProducer = dyn_cast<GenericOp>(producer)) {
       if (genericOpProducer.hasTensorSemantics())
         return FuseGenericOpsOnTensors::fuse(genericOpProducer, genericOp,
                                              consumerIdx, rewriter, folder);
     } else if (auto reshapeOpProducer = dyn_cast<TensorReshapeOp>(producer)) {
       return FuseTensorReshapeOpAsProducer<GenericOp>::fuse(
           reshapeOpProducer, genericOp, consumerIdx, rewriter, folder);
+    } else if (auto constantOpProducer = dyn_cast<ConstantOp>(producer)) {
+      return FuseConstantOpAsProducer<GenericOp>::fuse(
+          constantOpProducer, genericOp, consumerIdx, rewriter, folder);
     }
     return nullptr;
   }
 
   // Fuse when consumer is a TensorReshapeOp.
   if (TensorReshapeOp reshapeOp = dyn_cast<TensorReshapeOp>(consumer)) {
     if (auto genericOpProducer = dyn_cast<GenericOp>(producer)) {
       if (genericOpProducer.hasTensorSemantics())

mlir/test/Dialect/Linalg/fusion-tensor.mlir

func @generic_op_reshape_consumer_nofusion(%arg0 : tensor<?x?x?x5xf32>,
   %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
                                  affine_map<(i, j, k, l) -> (j, k, l)>] :
     tensor<?x?x?x5xf32> into tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
 
 // CHECK-LABEL: func @generic_op_reshape_consumer_nofusion
 //       CHECK: linalg.tensor_reshape
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2) -> (d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func @generic_op_constant_fusion(%arg0 : tensor<5x?x?xf32>) -> tensor<5x?x?xf32>
+{
+  %0 = constant dense<42.0> : tensor<5xf32>
+  %1 = linalg.generic
+       {args_in = 2 : i64, args_out = 1 : i64,
+         indexing_maps = [#map0, #map1, #map1],
+         iterator_types = ["parallel", "parallel", "parallel"]}
+       %0, %arg0 {
+       ^bb0(%arg1: f32, %arg2: f32):
+         %2 = mulf %arg1, %arg2 : f32
+         linalg.yield %2 : f32
+       }: tensor<5xf32>, tensor<5x?x?xf32> -> tensor<5x?x?xf32>
+  return %1 : tensor<5x?x?xf32>
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func @generic_op_constant_fusion
+//       CHECK:   %[[CST:.*]] = constant {{.*}} : f32
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     args_in = 1 : i64
+//  CHECK-SAME:     args_out = 1 : i64
+//       CHECK:   ^{{.*}}(%[[ARG1:.*]]: f32)
+//       CHECK:     mulf %[[CST]], %[[ARG1]]
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2) -> ()>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func @generic_op_zero_dim_constant_fusion(%arg0 : tensor<5x?x?xf32>)
+  -> tensor<5x?x?xf32>
+{
+  %0 = constant dense<42.0> : tensor<f32>
+  %1 = linalg.generic
+       {args_in = 2 : i64, args_out = 1 : i64,
+         indexing_maps = [#map0, #map1, #map1],
+         iterator_types = ["parallel", "parallel", "parallel"]}
+       %0, %arg0 {
+       ^bb0(%arg1: f32, %arg2: f32):
+         %2 = mulf %arg1, %arg2 : f32
+         linalg.yield %2 : f32
+       }: tensor<f32>, tensor<5x?x?xf32> -> tensor<5x?x?xf32>
+  return %1 : tensor<5x?x?xf32>
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func @generic_op_zero_dim_constant_fusion
+//       CHECK:   %[[CST:.*]] = constant {{.*}} : f32
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     args_in = 1 : i64
+//  CHECK-SAME:     args_out = 1 : i64
+//       CHECK:   ^{{.*}}(%[[ARG1:.*]]: f32)
+//       CHECK:     mulf %[[CST]], %[[ARG1]]