[mlir][Linalg] Add support for fusion between indexed_generic ops and generic ops on tensors.

[hanchung]

Summary

Different from the fusion between generic ops, indices are involved. In this
context, we need to re-map the indices for producer since the fused op is built
on consumer's perspective. This patch supports all combination of the fusion
between indexed_generic ops and generic ops, which includes tests case:

1. generic op as producer and indexed_generic op as consumer.
2. indexed_generic op as producer and generic op as consumer.
3. indexed_generic op as producer and indexed_generic op as consumer.

[mravishankar]

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

764

It seems like the logic used for generating the fused region for the indexed generic op would work for generic op as well. The numProducerIndices and numConsumerIndices would just remain the same. Can we combine these two methods?

[hanchung]

Merge methods together for generic ops and indexed_generic ops.

[hanchung]

Update parents.

[mravishankar]

Really cool! Just one comment. Approving assuming that it will be addressed. Thanks a lot!

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

441

Oh I really like this! It was something I wanted to get to anyway. Thanks for doing it right here!

531–539

I think this should be outside of the else?

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

226

Lets just make all shapes dynamic in the tests. THere shouldnt be any reason for static shape right?

[hanchung]

Address comments.

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

531–539

Good catch, in the beginning I thought we can't compose the maps sometimes. After review the comment of compose method, I found that the last two maps should swap.

See comments of AffineMap::compose():

/// Example:
///   map1: `(d0, d1)[s0, s1] -> (d0 + 1 + s1, d1 - 1 - s0)`
///   map2: `(d0)[s0] -> (d0 + s0, d0 - s0)`
///   map1.compose(map2):
///     `(d0)[s0, s1, s2] -> (d0 + s1 + s2 + 1, d0 - s0 - s2 - 1)`

Add some comment for it.

[hanchung]

rebase

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

 //===- Fusion.cpp - Implementation of linalg Fusion -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the linalg dialect Fusion pass.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
 #include "mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h"
 #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/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"

 }
 
 //====---------------------------------------------------------------------===//
 // Fusion on Tensor operation.
 //====---------------------------------------------------------------------===//
 
 namespace {
 
-/// Implementation of fusion of generic ops.
+/// Implementation of fusion of generic ops and indexed_generic ops.
 struct FuseGenericOpsOnTensors {
-  static bool isFusible(GenericOp producer, GenericOp consumer,
+  static bool isFusible(LinalgOp producer, LinalgOp consumer,
                         unsigned consumerIdx) {
     // Verify that
     // - the producer has all "parallel" iterator type.
     if (producer.getNumParallelLoops() != producer.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 = consumer.getIndexingMap(consumerIdx);
     if (consumerIndexMap.getNumResults() != producer.getNumLoops())
       return false;
 
     // Finally the index_map for the result must be invertible. For now just
     // verify it is a permutation.
     AffineMap producerResultIndexMap = producer.getOutputIndexingMap(0);
     return producerResultIndexMap.isPermutation();
   }
 
-  static Operation *fuse(GenericOp producer, GenericOp consumer,
+  static Operation *fuse(LinalgOp producer, LinalgOp consumer,
                          unsigned consumerIdx, PatternRewriter &rewriter,
                          OperationFolder *folder = nullptr) {
     if (!isFusible(producer, consumer, consumerIdx))
       return nullptr;
 
     unsigned numFusedOperands = producer.getOperation()->getNumOperands() +
                                 consumer.getOperation()->getNumOperands() - 1;
 
fusedOperands.append(std::next(consumerOperands.begin(), consumerIdx + 1),
                          consumerOperands.end());
 
     // Compute indexing_maps for the fused operation. The indexing_maps for the
     // operands of the consumers that arent fused are the same. The
     // indexing_maps for the producers need to be computed based on the
     // indexing_map of the operand at consumerIdx in the consumer.
     SmallVector<Attribute, 4> fusedIndexMaps;
     auto consumerIndexMaps = consumer.indexing_maps();
-    fusedIndexMaps.reserve(fusedOperands.size() + consumer.getNumResults());
+    fusedIndexMaps.reserve(fusedOperands.size() +
+                           consumer.getOperation()->getNumResults());
     fusedIndexMaps.assign(consumerIndexMaps.begin(),
                           std::next(consumerIndexMaps.begin(), consumerIdx));
     // Compute indexing maps for the producer args in the fused operation.
     computeProducerOperandIndex(
         producer, consumer.getInputIndexingMap(consumerIdx), fusedIndexMaps);
 
     // Append the indexing maps for the remaining consumer operands.
     fusedIndexMaps.append(std::next(consumerIndexMaps.begin(), consumerIdx + 1),
                           consumerIndexMaps.end());
 
     // Generate the fused op.
-    auto fusedOp = rewriter.create<GenericOp>(
-        rewriter.getUnknownLoc(), consumer.getResultTypes(), fusedOperands,
+    LinalgOp fusedOp;
+    if (isa<GenericOp>(producer.getOperation()) &&
+        isa<GenericOp>(consumer.getOperation())) {
+      fusedOp =
+          rewriter
+              .create<GenericOp>(
+                  rewriter.getUnknownLoc(),
+                  consumer.getOperation()->getResultTypes(), fusedOperands,
                   rewriter.getI64IntegerAttr(fusedOperands.size()),
-        rewriter.getI64IntegerAttr(consumer.getNumResults()),
-        rewriter.getArrayAttr(fusedIndexMaps), consumer.iterator_types(),
+                  rewriter.getI64IntegerAttr(
+                      consumer.getOperation()->getNumResults()),
+                  rewriter.getArrayAttr(fusedIndexMaps),
+                  consumer.iterator_types(),
                   /*doc=*/nullptr,
-        /*library_call=*/nullptr);
-    generateFusedRegion(rewriter, fusedOp.region(), producer.region(),
-                        consumer.region(), consumerIdx);
+                  /*library_call=*/nullptr)
+              .getOperation();
+    } else {
+      fusedOp =
+          rewriter
+              .create<IndexedGenericOp>(
+                  rewriter.getUnknownLoc(),
+                  consumer.getOperation()->getResultTypes(), fusedOperands,
+                  rewriter.getI64IntegerAttr(fusedOperands.size()),
+                  rewriter.getI64IntegerAttr(
+                      consumer.getOperation()->getNumResults()),
+                  rewriter.getArrayAttr(fusedIndexMaps),
+                  consumer.iterator_types(),
+                  /*doc=*/nullptr,
+                  /*library_call=*/nullptr)
+              .getOperation();
+    }
+
+    // Construct an AffineMap from consumer loops to producer loops.
+    // consumer loop -> tensor index
+    AffineMap consumerResultIndexMap =
+        consumer.getInputIndexingMap(consumerIdx);
+    // producer loop -> tensor index
+    AffineMap producerResultIndexMap = producer.getOutputIndexingMap(0);
+    // tensor index -> producer loop
+    AffineMap invProducerResultIndexMap =
+        inversePermutation(producerResultIndexMap);
+    assert(invProducerResultIndexMap &&
+           "expected producer result indexig map to be invertible");
+    // consumer loop -> producer loop
+    AffineMap consumerToProducerLoopsMap =
+        invProducerResultIndexMap.compose(consumerResultIndexMap);
+
+    generateFusedRegion(rewriter, fusedOp, producer, consumer,
+                        consumerToProducerLoopsMap, consumerIdx,
+                        consumer.getNumLoops());
     return fusedOp;
   }
 
 private:
   /// Append to `fusedOpIndexingMapAttrs` the indexing maps for the operands of
   /// the `producer` to use in the fused operation given the indexing map of the
   /// result of the producer in the consumer.
   static void computeProducerOperandIndex(
-      GenericOp producer, AffineMap fusedConsumerArgIndexMap,
+      LinalgOp producer, AffineMap fusedConsumerArgIndexMap,
       SmallVectorImpl<Attribute> &fusedOpIndexingMapAttrs) {
     // The indexing map in the consumer op (fusedConsumerArgIndexMap) is a map
     // from consumer loop -> consumer arg tensor index/producer result tensor
     // index. The fused loop is same as the consumer loop. For each producer arg
     // the indexing map to be computed is a map from consumer loop -> producer
     // arg tensor index.
 
     AffineMap producerResultIndexMap = producer.getOutputIndexingMap(0);
for (unsigned argNum : llvm::seq<unsigned>(0, producer.getNumInputs())) {
       // consumer loop/ fused loop -> producer arg tensor index.
       AffineMap indexingMap = t1.compose(fusedConsumerArgIndexMap);
       fusedOpIndexingMapAttrs.push_back(AffineMapAttr::get(indexingMap));
     }
   }
 
   /// Generate the region of the fused operation. The region of the fused op
   /// must be empty.
-  static void generateFusedRegion(PatternRewriter &rewriter,
-                                  Region &fusedRegion, Region &producerRegion,
-                                  Region &consumerRegion,
-                                  unsigned consumerIdx) {
+  static void generateFusedRegion(PatternRewriter &rewriter, Operation *fusedOp,
+                                  LinalgOp producer, LinalgOp consumer,
+                                  AffineMap consumerToProducerLoopsMap,
+                                  unsigned consumerIdx, unsigned nloops) {
     // Build the region of the fused op.
-    Block &producerBlock = producerRegion.front();
-    Block &consumerBlock = consumerRegion.front();
+    Block &producerBlock = producer.getOperation()->getRegion(0).front();
+    Block &consumerBlock = consumer.getOperation()->getRegion(0).front();
     Block *fusedBlock = new Block();
-    fusedRegion.push_back(fusedBlock);
+    fusedOp->getRegion(0).push_back(fusedBlock);
     BlockAndValueMapping mapper;
     OpBuilder::InsertionGuard guard(rewriter);
     rewriter.setInsertionPointToStart(fusedBlock);
+
+    // The block arguments are
+    // [index_0, index_1, ... ,
+    //   consumer_operand_0, ... , consumer_operand_(`consumerIdx`-1),
+    //   producer_operand_0, ... , producer_operand_(n-1)],
+    //   consumer_operand_(`consumerIdx`), .. consumer_operand_(m-1)]
+    // , where n is the number of producer's operand and m is the number
+    // consumer's operand.
+    // If both `numProducerIndices` and `numConsumerIndices` are zero, this is a
+    // generic op. In this case, there are no indices in block arguments.
+    unsigned numProducerIndices =
+        isa<IndexedGenericOp>(producer.getOperation()) ? nloops : 0;
+    unsigned numConsumerIndices =
+        isa<IndexedGenericOp>(consumer.getOperation()) ? nloops : 0;
+    // Firstly, add all the indices to the block arguments.
+    for (unsigned i = 0, e = std::max(numProducerIndices, numConsumerIndices);
+         i < e; ++i)
+      fusedBlock->addArgument(rewriter.getIndexType());
     // Map the arguments for the unmodified args from the consumer.
     for (auto consumerArg : llvm::enumerate(consumerBlock.getArguments())) {
-      if (consumerArg.index() == consumerIdx) {
+      if (consumerArg.index() == consumerIdx + numConsumerIndices) {
         // Map the arguments for the args from the producer.
-        for (auto producerArg : producerBlock.getArguments())
-          mapper.map(producerArg,
-                     fusedBlock->addArgument(producerArg.getType()));
+        for (auto producerArg : llvm::enumerate(producerBlock.getArguments())) {
+          // If producer is an indexed_generic op, map the indices from consumer
+          // loop to producer loop (because the fusedOp is built based on
+          // consumer's perspective).
+          if (producerArg.index() < numProducerIndices) {
+            auto newIndex = rewriter.create<mlir::AffineApplyOp>(
+                producer.getLoc(),
+                consumerToProducerLoopsMap.getSubMap(producerArg.index()),
+                fusedBlock->getArguments().take_front(nloops));
+            mapper.map(producerArg.value(), newIndex);
+          } else {
+            mapper.map(producerArg.value(),
+                       fusedBlock->addArgument(producerArg.value().getType()));
+          }
+        }
         continue;
       }
+
+      // If consumer is an indexed_generic op, map the indices to the block
+      // arguments directly. Otherwise, add the same type of arugment and map to
+      // it.
+      if (consumerArg.index() < numConsumerIndices) {
+        mapper.map(consumerArg.value(),
+                   fusedBlock->getArgument(consumerArg.index()));
+      } else {
         mapper.map(consumerArg.value(),
                    fusedBlock->addArgument(consumerArg.value().getType()));
       }
+    }
 
     // Add operations from producer (except the yield operation) to the fused
     // op.
     for (auto &op : producerBlock.getOperations()) {
       if (auto yieldOp = dyn_cast<YieldOp>(op)) {
         // Lookup the value the yield operation is mapped to.
         Value yieldVal = yieldOp.getOperand(0);
         if (Value clonedVal = mapper.lookupOrNull(yieldVal))
-          mapper.map(consumerBlock.getArgument(consumerIdx), clonedVal);
+          mapper.map(
+              consumerBlock.getArgument(consumerIdx + numConsumerIndices),
+              clonedVal);
         continue;
       }
       rewriter.clone(op, mapper);
     }
     for (auto &op : consumerBlock.getOperations())
       rewriter.clone(op, mapper);
   }
 };
static Operation *fuse(TensorReshapeOp producer, LinalgOpTy consumer,
 
     // Compute indexing_maps for the fused operation. The indexing_maps for the
     // operands of the consumers that arent fused are the same.
     SmallVector<AffineMap, 4> fusedIndexMaps =
         llvm::to_vector<4>(llvm::map_range(
             consumer.indexing_maps(), [](Attribute attr) -> AffineMap {
               return attr.cast<AffineMapAttr>().getValue();
             }));
 
     // Compute the indexing map to use for the operand of the producer.
     AffineMap modifiedMap = linearizeCollapsedDims(
         fusedIndexMaps[consumerIdx], producer.getResultType().getShape(),
         producer.getReassociationMaps());
     for (AffineExpr expr : modifiedMap.getResults()) {
       if (!expr.isPureAffine())
         return nullptr;
     }
Operation *mlir::linalg::fuseTensorOps(PatternRewriter &rewriter,
                                        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,
+  // Fuse when consumer is GenericOp or IndexedGenericOp.
+  if (isa<GenericOp>(consumer) || isa<IndexedGenericOp>(consumer)) {
+    auto linalgOpConsumer = cast<LinalgOp>(consumer);
+    if (!linalgOpConsumer.hasTensorSemantics())
+      return nullptr;
+    if (isa<GenericOp>(producer) || isa<IndexedGenericOp>(producer)) {
+      auto linalgOpProducer = cast<LinalgOp>(producer);
+      if (linalgOpProducer.hasTensorSemantics())
+        return FuseGenericOpsOnTensors::fuse(linalgOpProducer, linalgOpConsumer,
                                              consumerIdx, rewriter, folder);
     } else if (auto reshapeOpProducer = dyn_cast<TensorReshapeOp>(producer)) {
+      if (auto genericOpConsumer = dyn_cast<GenericOp>(consumer)) {
         return FuseTensorReshapeOpAsProducer<GenericOp>::fuse(
-          reshapeOpProducer, genericOp, consumerIdx, rewriter, folder);
+            reshapeOpProducer, genericOpConsumer, consumerIdx, rewriter,
+            folder);
+      }
     } else if (auto constantOpProducer = dyn_cast<ConstantOp>(producer)) {
+      if (auto genericOpConsumer = dyn_cast<GenericOp>(consumer)) {
         return FuseConstantOpAsProducer<GenericOp>::fuse(
-          constantOpProducer, genericOp, consumerIdx, rewriter, folder);
+            constantOpProducer, genericOpConsumer, 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())
         return FuseTensorReshapeOpAsConsumer<GenericOp>::fuse(
             genericOpProducer, reshapeOp, consumerIdx, rewriter, folder);
     }
     return nullptr;
   }
+
   return nullptr;
 }
 
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 template <typename LinalgOpTy>
 struct FuseTensorOps : public OpRewritePattern<LinalgOpTy> {
   using OpRewritePattern<LinalgOpTy>::OpRewritePattern;

 
 struct LinalgFusionPass : public LinalgFusionBase<LinalgFusionPass> {
   void runOnFunction() override { fuseLinalgOpsGreedily(getFunction()); }
 };
 } // namespace
 
 void mlir::populateLinalgTensorOpsFusionPatterns(
     MLIRContext *context, OwningRewritePatternList &patterns) {
-  patterns.insert<FuseTensorOps<GenericOp>, FuseTensorOps<TensorReshapeOp>>(
-      context);
+  patterns.insert<FuseTensorOps<GenericOp>, FuseTensorOps<IndexedGenericOp>,
+                  FuseTensorOps<TensorReshapeOp>>(context);
 }
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgFusionPass() {
   return std::make_unique<LinalgFusionPass>();
 }
 
 std::unique_ptr<Pass> mlir::createLinalgFusionOfTensorOpsPass() {
   return std::make_unique<FusionOfTensorOpsPass>();
 }

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

 }
 
 // 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 {

 //   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]]
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+func @generic_op_indexed_generic_op_fusion(%arg0: tensor<?x?xi32>,
+                                           %arg1: tensor<?x?xi32>) {
+    %0 = linalg.generic {
+      args_in = 2 : i64,
+      args_out = 1 : i64,
+      indexing_maps = [#map0, #map0, #map0],
+      iterator_types = ["parallel", "parallel"] } %arg0, %arg1 {
+    ^bb0(%arg2: i32, %arg3: i32):       // no predecessors
+      %10 = addi %arg2, %arg3 : i32
+      linalg.yield %10 : i32
+    } : tensor<?x?xi32>, tensor<?x?xi32> -> tensor<?x?xi32>
+    %1 = linalg.indexed_generic {
+      args_in = 1 : i64,
+      args_out = 1 : i64,
+      indexing_maps = [#map0, #map0],
+      iterator_types = ["parallel", "parallel"] } %0 {
+    ^bb0(%arg2: index, %arg3: index, %arg4: i32):       // no predecessors
+      %2 = index_cast %arg2 : index to i32
+      %3 = index_cast %arg3 : index to i32
+      %4 = addi %arg4, %2 : i32
+      %5 = subi %4, %3 : i32
+      linalg.yield %5 : i32
+    }: tensor<?x?xi32> -> tensor<?x?xi32>
+  return
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @generic_op_indexed_generic_op_fusion
+//   CHECK-NOT: linalg.generic
+//       CHECK: linalg.indexed_generic
+// CHECK-SAME:    args_in = 2
+// CHECK-SAME:    args_out = 1
+// CHECK-SAME:    indexing_maps = [#[[MAP0]], #[[MAP0]], #[[MAP0]]]
+//      CHECK: ^{{[a-zA-Z0-9_]*}}
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: index
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: index
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: i32
+// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]*]]: i32
+//      CHECK:   %[[VAL1:.+]] = addi %[[ARG2]], %[[ARG3]] : i32
+//      CHECK:   %[[ADD_OPERAND:.+]] = index_cast %[[ARG0]] : index to i32
+//      CHECK:   %[[SUB_OPERAND:.+]] = index_cast %[[ARG1]] : index to i32
+//      CHECK:   %[[VAL2:.+]] = addi %[[VAL1]], %[[ADD_OPERAND]] : i32
+//      CHECK:   %[[VAL3:.+]] = subi %[[VAL2]], %[[SUB_OPERAND]] : i32
+//      CHECK:   linalg.yield %[[VAL3]] : i32
+
+// -----
+
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+func @indexed_generic_op_generic_op_fusion(%arg0: tensor<?x?xi32>,
+                                           %arg1: tensor<?x?xi32>) {
+  %0 = linalg.indexed_generic {
+    args_in = 1 : i64,
+    args_out = 1 : i64,
+    indexing_maps = [#map0, #map0],
+    iterator_types = ["parallel", "parallel"] } %arg0 {
+  ^bb0(%arg2: index, %arg3: index, %arg4: i32):       // no predecessors
+    %2 = index_cast %arg2 : index to i32
+    %3 = index_cast %arg3 : index to i32
+    %4 = addi %arg4, %2 : i32
+    %5 = subi %4, %3 : i32
+    linalg.yield %5 : i32
+  }: tensor<?x?xi32> -> tensor<?x?xi32>
+  %1 = linalg.generic {
+    args_in = 2 : i64,
+    args_out = 1 : i64,
+    indexing_maps = [#map0, #map0, #map0],
+    iterator_types = ["parallel", "parallel"] } %0, %arg1 {
+  ^bb0(%arg2: i32, %arg3: i32):       // no predecessors
+    %10 = addi %arg2, %arg3 : i32
+    linalg.yield %10 : i32
+  } : tensor<?x?xi32>, tensor<?x?xi32> -> tensor<?x?xi32>
+  return
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @indexed_generic_op_generic_op_fusion
+//       CHECK: linalg.indexed_generic
+// CHECK-SAME:    args_in = 2
+// CHECK-SAME:    args_out = 1
+// CHECK-SAME:    indexing_maps = [#[[MAP0]], #[[MAP0]], #[[MAP0]]]
+//      CHECK: ^{{[a-zA-Z0-9_]*}}
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: index
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: index
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: i32
+// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]*]]: i32
+//      CHECK:   %[[ADD_OPERAND:.+]] = index_cast %[[ARG0]] : index to i32
+//      CHECK:   %[[SUB_OPERAND:.+]] = index_cast %[[ARG1]] : index to i32
+//      CHECK:   %[[VAL1:.+]] = addi %[[ARG2]], %[[ADD_OPERAND]] : i32
+//      CHECK:   %[[VAL2:.+]] = subi %[[VAL1]], %[[SUB_OPERAND]] : i32
+//      CHECK:   %[[VAL3:.+]] = addi %[[VAL2]], %[[ARG3]] : i32
+//      CHECK:   linalg.yield %[[VAL3]] : i32
+//   CHECK-NOT: linalg.generic
+
+// -----
+
+// The indices of the first indexed_generic op are swapped after fusion.
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+func @indexed_generic_op_fusion(%arg0: tensor<?x?xi32>) {
+    %0 = linalg.indexed_generic {
+      args_in = 1 : i64,
+      args_out = 1 : i64,
+      indexing_maps = [#map0, #map0],
+      iterator_types = ["parallel", "parallel"] } %arg0 {
+    ^bb0(%arg2: index, %arg3: index, %arg4: i32):       // no predecessors
+      %2 = index_cast %arg2 : index to i32
+      %3 = index_cast %arg3 : index to i32
+      %4 = addi %arg4, %2 : i32
+      %5 = subi %4, %3 : i32
+      linalg.yield %5 : i32
+    }: tensor<?x?xi32> -> tensor<?x?xi32>
+    %1 = linalg.indexed_generic {
+      args_in = 1 : i64,
+      args_out = 1 : i64,
+      indexing_maps = [#map1, #map1],
+      iterator_types = ["parallel", "parallel"] } %0 {
+    ^bb0(%arg2: index, %arg3: index, %arg4: i32):       // no predecessors
+      %2 = index_cast %arg2 : index to i32
+      %3 = index_cast %arg3 : index to i32
+      %4 = addi %arg4, %2 : i32
+      %5 = subi %4, %3 : i32
+      linalg.yield %5 : i32
+    }: tensor<?x?xi32> -> tensor<?x?xi32>
+  return
+}
+//   CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @indexed_generic_op_fusion
+//       CHECK: linalg.indexed_generic
+// CHECK-SAME:    args_in = 1
+// CHECK-SAME:    args_out = 1
+// CHECK-SAME:    indexing_maps = [#[[MAP0]], #[[MAP0]]]
+//      CHECK: ^{{[a-zA-Z0-9_]*}}
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: index
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: index
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: i32
+//      CHECK:   %[[ADD_OPERAND1:.+]] = index_cast %[[ARG1]] : index to i32
+//      CHECK:   %[[SUB_OPERAND1:.+]] = index_cast %[[ARG0]] : index to i32
+//      CHECK:   %[[VAL1:.+]] = addi %[[ARG2]], %[[ADD_OPERAND1]] : i32
+//      CHECK:   %[[VAL2:.+]] = subi %[[VAL1]], %[[SUB_OPERAND1]] : i32
+//      CHECK:   %[[ADD_OPERAND2:.+]] = index_cast %[[ARG0]] : index to i32
+//      CHECK:   %[[SUB_OPERAND2:.+]] = index_cast %[[ARG1]] : index to i32
+//      CHECK:   %[[VAL3:.+]] = addi %[[VAL2]], %[[ADD_OPERAND2]] : i32
+//      CHECK:   %[[VAL4:.+]] = subi %[[VAL3]], %[[SUB_OPERAND2]] : i32
+//      CHECK:   linalg.yield %[[VAL4]] : i32
+//   CHECK-NOT: linalg.indexed_generic