[mlir][linalg] Allow some fusion on mixed generics

[Hardcode84]

Summary

Relax linalg elementwise fusion check to allow mixed consumers. Producer is still required to be fully tensor to avoid potential memref aliasing.

[Hardcode84]

fix typo

[mravishankar]

Nice! I am excited for the possibilities of mixed tensor and buffer semantics.

[hw-yuexy]

Hi, I just have a question regarding this update.

I am currently having trouble lowering these generic ops that have mixed tensor and buffer into loops. I think transformation passes in linalg such as -convert-linalg-to-loops and -linalg-bufferize cannot handle such generic ops. So I'm just wondering if these passes will be updated in the future to allow lowering of generic ops with mixed tensor and buffer?

Thanks.

[Hardcode84]

In D138759#3964727, @hw-yuexy wrote:

Hi, I just have a question regarding this update.

I am currently having trouble lowering these generic ops that have mixed tensor and buffer into loops. I think transformation passes in linalg such as -convert-linalg-to-loops and -linalg-bufferize cannot handle such generic ops. So I'm just wondering if these passes will be updated in the future to allow lowering of generic ops with mixed tensor and buffer?

Thanks.

For bufferization we have our own pattern https://github.com/intel/mlir-extensions/blob/6e36adce8d211deefdf395ebcdc6c5fd47e080a5/numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/PlierToLinalg.cpp#L2657 , we are planning to upstream it eventually, but no concrete plans yet

For lowering we are using ConvertLinalgToParallelLoopsPass

[hw-yuexy]

In D138759#3965024, @Hardcode84 wrote:

In D138759#3964727, @hw-yuexy wrote:

Hi, I just have a question regarding this update.

I am currently having trouble lowering these generic ops that have mixed tensor and buffer into loops. I think transformation passes in linalg such as -convert-linalg-to-loops and -linalg-bufferize cannot handle such generic ops. So I'm just wondering if these passes will be updated in the future to allow lowering of generic ops with mixed tensor and buffer?

Thanks.

For bufferization we have our own pattern https://github.com/intel/mlir-extensions/blob/6e36adce8d211deefdf395ebcdc6c5fd47e080a5/numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/PlierToLinalg.cpp#L2657 , we are planning to upstream it eventually, but no concrete plans yet

For lowering we are using ConvertLinalgToParallelLoopsPass

I see. Thank you for the reply.

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

using RegionComputationFn =
       std::function<APIntOrFloat(const APIntOrFloatArray &)>;
 
   FoldConstantBase(MLIRContext *context, const ControlFusionFn &controlFn,
                    PatternBenefit benefit = 1)
       : OpRewritePattern<GenericOp>(context, benefit), controlFn(controlFn) {}
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
-    if (genericOp.hasBufferSemantics())
+    // Mixed and buffer sematics aren't supported.
+    if (!genericOp.hasTensorSemantics())
       return failure();
 
     // Only support ops generating one output for now.
     if (genericOp.getNumDpsInits() != 1)
       return failure();
 
     auto outputType = genericOp.getResultTypes().front().dyn_cast<ShapedType>();
     // Require the output types to be static given that we are generating

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

 bool mlir::linalg::areElementwiseOpsFusable(OpOperand *fusedOperand) {
   auto producer = fusedOperand->get().getDefiningOp<GenericOp>();
   auto consumer = dyn_cast<GenericOp>(fusedOperand->getOwner());
 
   // Check producer and consumer are generic ops.
   if (!producer || !consumer)
     return false;
 
-  // Producer and consumer must have tensor semantics.
-  if (!producer.hasTensorSemantics() || !consumer.hasTensorSemantics())
+  // Consumer can have mixed semantics, just check operand itself has tensor
+  // type. Producer must have full tensor semantics to avoid potential
+  // aliasing between producer and consumer memrefs.
+  if (!producer.hasTensorSemantics() ||
+      !fusedOperand->get().getType().isa<RankedTensorType>())
     return false;
 
   // Verify that
   // - the producer has all "parallel" iterator type.
   if (producer.getNumParallelLoops() != producer.getNumLoops())
     return false;
 
   // Only allow fusing the producer of an input operand for now.
for (auto opOperand : llvm::enumerate(producer.getDpsInitOperands())) {
     fusedIndexMaps.push_back(map);
     fusedResultTypes.push_back(opOperand.value()->get().getType());
   }
 
   // 7. All of consumer's output operands (skip operands: added by the builder).
   for (OpOperand *opOperand : consumer.getDpsInitOperands()) {
     fusedOutputOperands.push_back(opOperand->get());
     fusedIndexMaps.push_back(consumer.getMatchingIndexingMap(opOperand));
-    fusedResultTypes.push_back(opOperand->get().getType());
+    Type resultType = opOperand->get().getType();
+    if (!resultType.isa<MemRefType>())
+      fusedResultTypes.push_back(resultType);
   }
 
   // Generate the fused op.
   auto fusedOp = rewriter.create<GenericOp>(
       consumer.getLoc(), fusedResultTypes, fusedInputOperands,
       fusedOutputOperands, rewriter.getAffineMapArrayAttr(fusedIndexMaps),
       consumer.getIteratorTypes(),
       /*doc=*/nullptr,

mlir/lib/Interfaces/DestinationStyleOpInterface.cpp

if (failed(OpTrait::impl::verifyNOperands(op, numInputs + numInits)))
     return failure();
   // Verify the number of results matches the number of output tensors.
   if (op->getNumResults() != outputTensorOperands.size())
     return op->emitOpError("expected the number of results (")
            << op->getNumResults()
            << ") to be equal to the number of output tensors ("
            << outputTensorOperands.size() << ")";
 
-  // Simplifying assumption: either full tensor or full buffer mode.
-  // This allows simpler verification of output operands vs result types
-  // without premature tracking of which operand is what in mixed-mode.
-  // TODO: relax when mixed-mode needs to pass verification.
-  if (!outputBufferOperands.empty() && !outputTensorOperands.empty())
-    return op->emitOpError(
-        "expected output operands to all have tensor type or "
-        "all have buffer type");
-
   for (OpOperand *opOperand : outputTensorOperands) {
     OpResult result = dstStyleOp.getTiedOpResult(opOperand);
     if (result.getType() != opOperand->get().getType())
       return op->emitOpError("expected type of operand #")
              << opOperand->getOperandNumber() << " ("
              << opOperand->get().getType() << ")"
              << " to match type of corresponding result (" << result.getType()
              << ")";
   }
   return success();
 }

mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir

 //  CHECK-SAME:         %[[B0:[a-zA-Z0-9_]+]]: f32
 //  CHECK-SAME:         %[[B1:[a-zA-Z0-9_]+]]: f32
 //   CHECK-DAG:     %[[T0:.+]] = arith.addf %[[B0]], %[[B1]]
 //   CHECK-DAG:     %[[T1:.+]] = arith.addf %[[T0]], %[[B1]]
 //   CHECK-DAG:     %[[T2:.+]] = arith.addf %[[T1]], %[[B1]]
 //   CHECK-DAG:     %[[T3:.+]] = arith.addf %[[T2]], %[[B1]]
 //       CHECK:     linalg.yield %[[T3]] : f32
 //       CHECK:   return %[[GENERIC]]
+
+// -----
+
+// CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @mixed_fusion
+func.func @mixed_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>, %arg8 : memref<?x?xf32>)
+{
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
+  %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
+  %3 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}
+      ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%2 : tensor<?x?xf32>) {
+    ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
+      %4 = arith.addf %arg3, %arg4 : f32
+      linalg.yield %4 : f32
+  } -> tensor<?x?xf32>
+  // CHECK: linalg.generic {
+  // CHECK-SAME: indexing_maps = {{\[}}[[$MAP0]], [[$MAP0]], [[$MAP0]], [[$MAP0]]{{\]}}
+  linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}
+      ins(%3, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg8 : memref<?x?xf32>) {
+    // CHECK: ^{{[a-zA-Z0-9_]*}}
+    // CHECK-SAME: [[ARG0:%[a-zA-Z0-9_]*]]
+    // CHECK-SAME: [[ARG1:%[a-zA-Z0-9_]*]]
+    // CHECK-SAME: [[ARG2:%[a-zA-Z0-9_]*]]
+    ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):
+      // CHECK: [[T1:%[a-zA-Z0-9_]*]] = arith.addf [[ARG0]], [[ARG1]]
+      // CHECK-NOT: linalg.yield
+      // CHECK: arith.mulf [[T1]], [[ARG2]]
+      // CHECK: linalg.yield
+      %5 = arith.mulf %arg5, %arg6 : f32
+      linalg.yield %5 : f32
+    }
+  return
+}