[mlir][Linalg] Implement fusion of linalg.generic operation on tensors.

[mravishankar]

Summary

The initial implementation of the fusion operation exposes a method to
fuse a consumer with its producer, when

• both the producer and consumer operate on tensors
• the producer has only a single result value
• the producer has only "parallel" iterator types

A new interface method hasTensorSemantics is added to verify that an
operation has all operands and results of type RankedTensorType.

[mravishankar]

Updating with changes to fix build failures

[mravishankar]

Actually updating the file to fix build errors.

[mravishankar]

Addressing some clang format issues

[rriddle]

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

228

nit: Can you separate the functor to avoid duplication?

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

328

Please remove all trivial braces.

374

Start sentences with a capital.

403

nit: you could use getOperands()[producerOp.getNumInputs()]

427

This seems like an unnecessary usage of for_each. Why not just write a for loop? You don't even need to do that really. You could just use getResultTypes() directly as it is just an ArrayRef<Type>

[nicolasvasilache]

This looks really great, thanks a lot for pushing this Mahesh!
I think it is an elegant solution that shows how proper usage of the StructuredOps abstraction avoid creating problems that should never have existed in the first place (e.g. carrying broadcast and all the phase ordering issues related to buffer allocation, fusion and erasing unnecessary buffers).

Approving conditioned on addressing all comments.

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

337

trivial braces here and below plz

346

re. AffineMap, note that the current "inversion" procedure is best effort atm and will fail in many cases.
You want to check that the inverse computation succeeeds and return false if not.
Otherwise I suspect you'll crash later.

Example that I think should fail today and may crash your transformation down the line: (i, j) -> (i, i+j)

390

This can fail in a bunch of even simple cases.
Plz see my comment above re. checking it in the precondition.

411

Haven't checked deeply that you don't have an off by 1 here somewhere, please be extra careful but you know better than me :)

520

Thanks for writing this as a pattern that will compose with other patterns!

524

trivial braces here and everywhere please

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

92

Very cool, thanks Mahesh!

[mravishankar]

Updating after addressing comments

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

346

Thanks. Forgot to check for permutation. Added that check now. I think that is enough for now. If we need to support check of invertability, we can extend that later.

403

I am not sure about that. I am inserting a range of elements. I think the insert method used needs an iterator for the begin/end.

411

I did hit the off by one. Dont need to worry at erase time. THe first insertion will be fine too. Only need to increment the insert position after each insertion to make sure things are inserted at the right place.

427

Thanks for getResultTypes() (cant keep all the utility methods in my head, but should have checked)

[rriddle]

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

426

Is there any other location you can use besides unknown? We should try to propagate location as much as possible.

437

fusedBlock = b.createBlock(&fusedOpRegion);

?

528

Are all of these .getOperation calls really necessary?

[phosek]

The test seems to be failing on aarch64, this is the failure we're seeing on aarch64-linux-gnu:

FAIL: Clang :: CodeGen/stack-clash-protection.c (3099 of 16811)
******************** TEST 'Clang :: CodeGen/stack-clash-protection.c' FAILED ********************
Script:
--
: 'RUN: at line 2';   /b/s/w/ir/k/recipe_cleanup/clanggb_2RS/llvm_build_dir/bin/clang -target x86_64 -O0 -o /b/s/w/ir/k/recipe_cleanup/clanggb_2RS/llvm_build_dir/tools/clang/test/CodeGen/Output/stack-clash-protection.c.tmp.out /b/s/w/ir/k/llvm-project/clang/test/CodeGen/stack-clash-protection.c -fstack-clash-protection && /b/s/w/ir/k/recipe_cleanup/clanggb_2RS/llvm_build_dir/tools/clang/test/CodeGen/Output/stack-clash-protection.c.tmp.out
--
Exit Code: 1

Command Output (stderr):
--
clang-11: warning: argument unused during compilation: '-fstack-clash-protection' [-Wunused-command-line-argument]
gcc: error: unrecognized command line option '-m64'
clang-11: error: linker (via gcc) command failed with exit code 1 (use -v to see invocation)

--

I think this test needs REQUIRES: x86.

[phosek]

I just noticed it's failing on macOS as well with a different error:

FAIL: Clang :: CodeGen/stack-clash-protection.c (3400 of 16811)
******************** TEST 'Clang :: CodeGen/stack-clash-protection.c' FAILED ********************
Script:
--
: 'RUN: at line 2';   /b/s/w/ir/k/recipe_cleanup/clanghb5vUd/llvm_build_dir/bin/clang -target x86_64 -O0 -o /b/s/w/ir/k/recipe_cleanup/clanghb5vUd/llvm_build_dir/tools/clang/test/CodeGen/Output/stack-clash-protection.c.tmp.out /b/s/w/ir/k/llvm-project/clang/test/CodeGen/stack-clash-protection.c -fstack-clash-protection && /b/s/w/ir/k/recipe_cleanup/clanghb5vUd/llvm_build_dir/tools/clang/test/CodeGen/Output/stack-clash-protection.c.tmp.out
--
Exit Code: 1

Command Output (stderr):
--
clang-11: warning: argument unused during compilation: '-fstack-clash-protection' [-Wunused-command-line-argument]
ld: warning: ignoring file /b/s/w/ir/tmp/t/stack-clash-protection-ccc8da.o, file was built for unsupported file format ( 0x7F 0x45 0x4C 0x46 0x02 0x01 0x01 0x00 0x00 0x00 0x00 0x00 0x00 0x00 0x00 0x00 ) which is not the architecture being linked (x86_64): /b/s/w/ir/tmp/t/stack-clash-protection-ccc8da.o
Undefined symbols for architecture x86_64:
  "_main", referenced from:
     implicit entry/start for main executable
ld: symbol(s) not found for architecture x86_64
clang: error: linker command failed with exit code 1 (use -v to see invocation)
clang-11: error: linker (via gcc) command failed with exit code 1 (use -v to see invocation)

---

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

let methods = [
     >,
     InterfaceMethod<
       "Query the indexing maps attribute within the current operation.",
       "ArrayAttr", "indexing_maps"
     >,
     InterfaceMethod<[{
         Query whether the op has only MemRef input and outputs.
       }], "bool", "hasBufferSemantics">,
+    InterfaceMethod<[{
+        Query whether the op has only RankedTensor input and outputs.
+      }], "bool", "hasTensorSemantics">,
 
     //========================================================================//
     // Other static interface methods.
     //========================================================================//
     StaticInterfaceMethod<[{
         Create an operation of the current type with the given location,
         operands, and attributes.
       }],

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

public:
   //==========================================================================//
   /// Query whether the op has only buffer inputs and no returns.
   bool hasBufferSemantics() {
     return this->getOperation()->getNumResults() == 0 &&
            llvm::all_of(getInputs(),
                         [](Value v) { return v.getType().isa<MemRefType>(); });
   }
 
+  /// Query whether the op has only tensor inputs and outputs.
+  bool hasTensorSemantics() {
+    return llvm::all_of(
+               getInputs(),
+               [](Value v) { return v.getType().isa<RankedTensorType>(); }) &&
+           llvm::all_of(this->getOperation()->getResults(), [](Value v) {
+             return v.getType().isa<RankedTensorType>();
+           });
+  }
+
   //==========================================================================//
   // Other static interface methods.
   //==========================================================================//
   static LogicalResult verifyTrait(Operation *op) {
     auto nOperands = cast<ConcreteType>(op).getNumInputsAndOutputBuffers();
     if (failed(OpTrait::impl::verifyAtLeastNOperands(op, nOperands)))
       return failure();
     return success();
   }
 };
 
 } // namespace linalg
 } // namespace OpTrait
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_LINALGTRAITS_H_

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 /// When non-null, the optional pointer `folder` is used to call into the
 /// `createAndFold` builder method. If `folder` is null, the regular `create`
 /// method is called.
 Optional<FusionInfo> fuseProducerOf(OpBuilder &b, LinalgOp consumer,
                                     unsigned consumerIdx,
                                     const LinalgDependenceGraph &graph,
                                     OperationFolder *folder = nullptr);
 
+/// Fuse linalg operation on tensors, where the result of the producer is used
+/// as the operand of the consumer at position `consumerIdx`.
+Optional<LinalgOp> fuseTensorOps(OpBuilder &b, LinalgOp producer,
+                                 LinalgOp consumer, unsigned consumerIdx,
+                                 OperationFolder *folder = nullptr);
+
 /// Returns the linearized list of all view dimensions in a linalgOp. Applying
 /// the inverse, concatenated loopToOperandRangeMaps to this list allows the
 /// derivation of loop ranges for any linalgOp.
 template <typename ConcreteOp>
 SmallVector<Value, 8> getViewSizes(OpBuilder &builder, ConcreteOp linalgOp) {
   auto loc = linalgOp.getLoc();
   SmallVector<Value, 8> res;
   for (auto v : linalgOp.getInputsAndOutputBuffers()) {

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

 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/EDSC/Helpers.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/STLExtras.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
for (auto dependence : graph.getDependencesInto(
     auto fusedProducer = fuse(producedView, producer, consumer, consumerIdx,
                               producerIdx, folder);
 
     return FusionInfo{producer, fusedProducer};
   }
   return llvm::None;
 }
 
+/// Checks if two Generic ops are fusible, when one is a producer and another is
+/// a consumer (with the result of the producer being the `consumerIdx` operand
+/// of the consumer).
+static bool areTensorOpsFusible(LinalgOp producer, LinalgOp consumer,
+                                unsigned consumerIdx) {
+  // Verify that the producer and consumer are ops on tensors.
+  if (!producer.hasTensorSemantics() || !consumer.hasTensorSemantics()) {
+    return false;
+  }
+  auto producerOp = dyn_cast<linalg::GenericOp>(producer.getOperation());
+  auto consumerOp = dyn_cast<linalg::GenericOp>(consumer.getOperation());
+  // Verify that the producer and consumers are generic ops, and only handle
+  // cases where the producer has a single return value that should be the same
+  // as argument at `consumerIdx` of the consumer.
+  if (!producerOp || !consumerOp || producerOp.getNumOutputs() != 1 ||
+      producerOp.getResult(0) != consumerOp.getOperand(consumerIdx)) {
+    return false;
+  }
+  // Check that the producer has all "parallel" iterator types.
+  if (producerOp.getNumParallelLoops() != producerOp.getNumLoops()) {
+    return false;
+  }
+  // Get the consumer index map. The number of results of the consumer index map
+  // must match the number of loops of the producer.
+  AffineMap consumerIndexMap = consumerOp.getIndexingMap(consumerIdx);
+  if (consumerIndexMap.getNumResults() != producerOp.getNumLoops()) {
+    return false;
+  }
+
+  // For now, only handle ops which have regions and does not use functions to
+  // specify the region.
+  if (producerOp.fun() || consumerOp.fun()) {
+    return false;
+  }
+  return true;
+}
+
+/// Computes the indexing maps for arguments of a producer generic op when the
+/// result of the producer is fused with the consumer.
+/// - consumerIndexMap is the indexing_map for the argument in the consumer op
+///   that is the result of the producer op.
+/// - invProducerResultIndexMap is the inverse of the indexing_map for the
+///   result in the producer op.
+/// - producerArgIndexMap is the indexing_map of the argument of the producer
+///   op.
+/// The result is the indexing_map to use for the producer argument when the
+/// producer and consumer ops are fused.
+static AffineMap computeProducerArgMap(AffineMap consumerIndexMap,
+                                       AffineMap invProducerResultIndexMap,
+                                       AffineMap producerArgIndexMap) {
+  // t1 is map from producer result tensor index -> producer arg tensor index.
+  auto t1 = producerArgIndexMap.compose(invProducerResultIndexMap);
+  // return is consumer loop -> producer arg tensor index, i.e. indexing_map for
+  // the producer argument in the fused operation.
+  return t1.compose(consumerIndexMap);
+}
+
+Optional<LinalgOp> fuseTensorOps(OpBuilder &b, LinalgOp producer,
+                                 LinalgOp consumer, unsigned consumerIdx,
+                                 OperationFolder *folder = nullptr) {
+  if (!areTensorOpsFusible(producer, consumer, consumerIdx)) {
+    return {};
+  }
+  MLIRContext *context = b.getContext();
+  auto producerOp = cast<linalg::GenericOp>(producer.getOperation());
+  auto consumerOp = cast<linalg::GenericOp>(consumer.getOperation());
+  AffineMap consumerIndexMap = consumerOp.getIndexingMap(consumerIdx);
+  AffineMap invProducerResultIndexMap =
+      inversePermutation(producerOp.getOutputIndexingMap(0));
+
+  // Compute the fused op operandslist by replacing the operand corresponding to
+  // the result of the producer, with the operands of the producer.
+  unsigned fusedArgsIn =
+      producerOp.getNumInputs() + consumerOp.getNumInputs() - 1;
+  auto fusedArgsOut = consumerOp.getNumOutputs();
+  SmallVector<Value, 2> fusedOperandsList(consumerOp.getOperands());
+  fusedOperandsList.erase(std::next(fusedOperandsList.begin(), consumerIdx));
+  fusedOperandsList.reserve(fusedArgsIn + fusedArgsOut);
+  fusedOperandsList.insert(
+      std::next(fusedOperandsList.begin(), consumerIdx),
+      producerOp.operand_begin(),
+      std::next(producerOp.operand_begin(), producerOp.getNumInputs()));
+
+  // Compute the fused indexing_maps of the operands/results of the fused op.
+  SmallVector<Attribute, 2> fusedIndexingMapAttrs;
+  fusedIndexingMapAttrs.reserve(fusedArgsIn + fusedArgsOut);
+  fusedIndexingMapAttrs.append(consumerOp.indexing_maps().begin(),
+                               consumerOp.indexing_maps().end());
+  fusedIndexingMapAttrs.erase(
+      std::next(fusedIndexingMapAttrs.begin(), consumerIdx));
+  auto insertPos = std::next(fusedIndexingMapAttrs.begin(), consumerIdx);
+  for (auto producerArgIndexAttr :
+       llvm::enumerate(producerOp.indexing_maps())) {
+    if (producerArgIndexAttr.index() == producerOp.getNumInputs()) {
+      break;
+    }
+    auto composedIndexMap = computeProducerArgMap(
+        consumerIndexMap, invProducerResultIndexMap,
+        producerArgIndexAttr.value().cast<AffineMapAttr>().getValue());
+    insertPos = std::next(fusedIndexingMapAttrs.insert(
+        insertPos, AffineMapAttr::get(composedIndexMap)));
+  }
+
+  // Generate the fused op.
+  SmallVector<Type, 2> fusedOpResultTypes;
+  llvm::for_each(consumerOp.getResults(), [&fusedOpResultTypes](Value v) {
+    fusedOpResultTypes.push_back(v.getType());
+  });
+  auto fusedLinalgOp = b.create<GenericOp>(
+      UnknownLoc::get(context), fusedOpResultTypes, fusedOperandsList,
+      b.getI64IntegerAttr(fusedArgsIn), b.getI64IntegerAttr(fusedArgsOut),
+      b.getArrayAttr(fusedIndexingMapAttrs), consumerOp.iterator_types(),
+      /*doc=*/nullptr,
+      /*fun=*/nullptr,
+      /*library_call=*/nullptr);
+
+  // Build the region of the fused op.
+  auto &fusedOpRegion = fusedLinalgOp.region();
+  Block &producerOpBlock = producerOp.region().front();
+  Block &consumerOpBlock = consumerOp.region().front();
+  Block *fusedBlock = new Block();
+  fusedOpRegion.push_back(fusedBlock);
+  BlockAndValueMapping mapper;
+  // Map the arguments for the unmodified args from the consumer.
+  for (auto consumerOpArg : llvm::enumerate(consumerOpBlock.getArguments())) {
+    if (consumerOpArg.index() == consumerIdx) {
+      // Map the arguments for the args from the producer.
+      for (auto producerOpArg : producerOpBlock.getArguments()) {
+        mapper.map(producerOpArg,
+                   fusedBlock->addArgument(producerOpArg.getType()));
+      }
+      continue;
+    }
+    mapper.map(consumerOpArg.value(),
+               fusedBlock->addArgument(consumerOpArg.value().getType()));
+  }
+
+  // Add operations from producer (except the yield operation) to the fused op.
+  for (auto &op : producerOpBlock.getOperations()) {
+    if (auto yieldOp = dyn_cast<YieldOp>(op)) {
+      // Lookup the value the yield operation is mapped to.
+      Value yieldVal = yieldOp.getOperand(0);
+      auto clonedVal = mapper.lookup(yieldVal);
+      mapper.map(consumerOpBlock.getArgument(consumerIdx), clonedVal);
+      continue;
+    }
+    fusedBlock->push_back(op.clone(mapper));
+  }
+  for (auto &op : consumerOpBlock.getOperations()) {
+    fusedBlock->push_back(op.clone(mapper));
+  }
+  return cast<LinalgOp>(fusedLinalgOp.getOperation());
+}
+
 static void fuseLinalgOpsGreedily(FuncOp f) {
   LLVM_DEBUG(f.print(dbgs() << "\nBefore linalg-fusion: \n"));
 
   OpBuilder b(f);
   OperationFolder folder(f.getContext());
   DenseSet<Operation *> eraseSet;
 
   // Save original Linalg ops, we only want to make a pass over those.
static void fuseLinalgOpsGreedily(FuncOp f) {
   // dependences and are guaranteed by construction to produce the whole view.
   // We may thus erase the producer once it is fused.
   for (auto *e : eraseSet)
     e->erase();
   LLVM_DEBUG(f.print(dbgs() << "\nAfter linalg-fusion: \n"));
 }
 
 namespace {
+
+/// Patterns to fuse a generic op, with the producer of its operands.
+struct FuseGenericTensorOps : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(GenericOp op,
+                                     PatternRewriter &rewriter) const override {
+    if (!op.hasTensorSemantics()) {
+      return matchFailure();
+    }
+    // Find the first operand that is defined by another generic op on tensors.
+    for (auto operand : llvm::enumerate(op.getOperation()->getOperands())) {
+      auto definingOp =
+          dyn_cast_or_null<GenericOp>(operand.value().getDefiningOp());
+      if (!definingOp || !definingOp.hasTensorSemantics()) {
+        return matchFailure();
+      }
+      auto fusedOp =
+          fuseTensorOps(rewriter, cast<LinalgOp>(definingOp.getOperation()),
+                        cast<LinalgOp>(op.getOperation()), operand.index());
+      if (!fusedOp) {
+        return matchFailure();
+      }
+      rewriter.replaceOp(op, fusedOp.getValue().getOperation()->getResults());
+      return matchSuccess();
+    }
+    return matchFailure();
+  }
+};
+
+/// Pass that fuses generic ops on tensors. Used only for testing.
+struct FusionOfTensorOpsPass : public OperationPass<FusionOfTensorOpsPass> {
+  void runOnOperation() override {
+    OwningRewritePatternList patterns;
+    Operation *op = getOperation();
+    patterns.insert<FuseGenericTensorOps>(op->getContext());
+    applyPatternsGreedily(op->getRegions(), patterns);
+  };
+};
+
 struct LinalgFusionPass : public FunctionPass<LinalgFusionPass> {
   void runOnFunction() override { fuseLinalgOpsGreedily(getFunction()); }
 };
 } // namespace
 
 std::unique_ptr<OpPassBase<FuncOp>> mlir::createLinalgFusionPass() {
   return std::make_unique<LinalgFusionPass>();
 }
 
 static PassRegistration<LinalgFusionPass>
     pass("linalg-fusion", "Fuse operations in the linalg dialect");
+
+static PassRegistration<FusionOfTensorOpsPass>
+    tensorOpsPass("linalg-fusion-for-tensor-ops",
+                  "Fuse operations on RankedTensorType in linalg dialect");

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

+// RUN: mlir-opt %s -linalg-fusion-for-tensor-ops -split-input-file | FileCheck %s --dump-input-on-failure
+
+// CHECK-DAG: [[MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @add_mul_fusion
+func @add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+{
+  %0 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]} %arg0, %arg1 {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = addf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  }: tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
+  // CHECK: linalg.generic {args_in = 3 : i64, args_out = 1 : i64
+  // CHECK-SAME: indexing_maps = {{\[}}[[MAP0]], [[MAP0]], [[MAP0]], [[MAP0]]{{\]}}
+  %2 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]} %0, %arg2 {
+    // 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):       // no predecessors
+      // CHECK: [[T1:%[a-zA-Z0-9_]*]] = addf [[ARG0]], [[ARG1]]
+      // CHECK-NOT: linalg.yield
+      // CHECK: mulf [[T1]], [[ARG2]]
+      // CHECK: linalg.yield
+      %3 = mulf %arg5, %arg6 : f32
+      linalg.yield %3 : f32
+    }: tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
+  return %2 : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK-DAG: [[MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: [[MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d1, d0)>
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1, d0)>
+
+// CHECK-LABEL: @transpose_add_mul_fusion
+func @transpose_add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+{
+  %0 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]} %arg0, %arg1 {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = addf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  }: tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
+  // CHECK: linalg.generic {args_in = 3 : i64, args_out = 1 : i64
+  // CHECK-SAME: indexing_maps = {{\[}}[[MAP0]], [[MAP1]], [[MAP0]], [[MAP0]]{{\]}}
+  %2 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]} %0, %arg2 {
+    ^bb0(%arg5: f32, %arg6: f32):       // no predecessors
+      %3 = mulf %arg5, %arg6 : f32
+      linalg.yield %3 : f32
+    }: tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
+  return %2 : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK-DAG: [[MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: [[MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d1, d0)>
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1, d0)>
+
+// CHECK-LABEL: @add_transpose_mul_fusion
+func @add_transpose_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+{
+  %0 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]} %arg0, %arg1 {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = addf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  }: tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
+  // CHECK: linalg.generic {args_in = 3 : i64, args_out = 1 : i64
+  // CHECK-SAME: indexing_maps = {{\[}}[[MAP1]], [[MAP0]], [[MAP0]], [[MAP0]]{{\]}}
+  %2 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]} %0, %arg2 {
+    ^bb0(%arg5: f32, %arg6: f32):       // no predecessors
+      %3 = mulf %arg5, %arg6 : f32
+      linalg.yield %3 : f32
+    }: tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
+  return %2 : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK-DAG: [[MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: [[MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0)>
+// CHECK-DAG: [[MAP2:#[a-zA-Z0-9_]*]] = affine_map<(d0) -> (d0)>
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d0)>
+#map2 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL: @add_broadcast_mul_fusion
+func @add_broadcast_mul_fusion(%arg0: tensor<?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+{
+  %0 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map2, #map2, #map2], iterator_types = ["parallel"]} %arg0, %arg1 {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = addf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  }: tensor<?xf32>, tensor<?xf32> -> tensor<?xf32>
+  // CHECK: linalg.generic {args_in = 3 : i64, args_out = 1 : i64
+  // CHECK-SAME: indexing_maps = {{\[}}[[MAP1]], [[MAP1]], [[MAP0]], [[MAP0]]
+  %2 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]} %0, %arg2 {
+    ^bb0(%arg5: f32, %arg6: f32):       // no predecessors
+      %3 = mulf %arg5, %arg6 : f32
+      linalg.yield %3 : f32
+    }: tensor<?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
+  return %2 : tensor<?x?xf32>
+}