[mlir][linalg] Fix bug in the fusion on tensors index op handling.

[gysit]

Summary

The old index op handling created the index operations used to replace the producer indices directly in the fused block and let them point back to the producer block. As a result, after fusion some index operations in the fuse block had back references to the old producer block resulting in illegal IR. The patch performs the index operation replacement in the producer block and then clones all operations to the fused block avoiding the aforementioned inconsistencies.

[pifon2a]

Thank you, Tobias!

[gysit]

Avoid modifying the producer block since it may be inlined into another consumer. The updated revision thus introduces the fused index operations directly in the fused block and uses the mapper to updated the index op references during cloning.

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

for (auto it :
        llvm::zip(producerBlock.getArguments().take_front(numProducerIndices),
                  fusedBlock->getArguments().take_front(numProducerIndices))) {
     auto newIndex = rewriter.create<mlir::AffineApplyOp>(
         producer.getLoc(),
         consumerToProducerLoopsMap.getSubMap(std::get<0>(it).getArgNumber()),
         fusedBlock->getArguments().take_front(numFusedOpIndices));
     mapper.map(std::get<0>(it), newIndex);
   }
-  // 2b. Replace the producer index operations by index operations placed in the
-  // fused block using the `consumerToProducerLoopsMap` to map the index spaces.
+  // 2b. Add an index operation for every fused loop dimension and use the
+  // `consumerToProducerLoopsMap` to map the producer indices.
+  if (producer.hasIndexSemantics()) {
+    // Add an index operation for every fused loop dimension.
     unsigned numFusedOpLoops =
         std::max(producer.getNumLoops(), consumer.getNumLoops());
-  if (producer.hasIndexSemantics()) {
     SmallVector<Value> fusedIndices;
     fusedIndices.reserve(numFusedOpLoops);
     llvm::transform(llvm::seq<int64_t>(0, numFusedOpLoops),
                     std::back_inserter(fusedIndices), [&](int64_t dim) {
                       return rewriter.create<IndexOp>(producer.getLoc(), dim);
                     });
     for (IndexOp indexOp :
          llvm::make_early_inc_range(producerBlock.getOps<IndexOp>())) {
       Value newIndex = rewriter.create<mlir::AffineApplyOp>(
           producer.getLoc(),
           consumerToProducerLoopsMap.getSubMap(indexOp.dim()), fusedIndices);
-      // Replace the producer index operation by the index value computed in the
-      // fused block. All remaining operations in the producer block are later
-      // on cloned to the fused block.
-      rewriter.replaceOp(indexOp, newIndex);
+      mapper.map(indexOp.getResult(), newIndex);
     }
   }
   // TODO: allow fusing the producer of an output operand.
   assert(consumerIdx < consumer.getNumInputs() &&
          "expected producer of input operand");
   // 3. Consumer input operands up to consumerIdx (exclusive).
   for (BlockArgument bbArg : consumerBlock.getArguments()
                                  .drop_front(numConsumerIndices)
generateFusedElementwiseOpRegion(PatternRewriter &rewriter, Operation *fusedOp,
   // 6. All of consumer's output operands.
   for (BlockArgument bbArg :
        consumerBlock.getArguments().take_back(consumer.getNumOutputs()))
     mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType()));
   // 7. All of producer's output operands except the one fused.
   // TODO: allow fusion of multi-result producers.
   assert(producer->getNumResults() == 1 && "expected single result producer");
 
-  // 8. Clone operations from producer (except the yield operation) to the fused
-  // op.
-  for (auto &op : producerBlock.without_terminator())
+  // 8. Clone all producer operations except for the yield and index operations
+  // to the fused operation.
+  for (auto &op : producerBlock.without_terminator()) {
+    if (!isa<IndexOp>(op))
       rewriter.clone(op, mapper);
+  }
   // 9. Now we can map the consumerBlock's `consumerIdx` block argument. Just
   // forward the yield operand.
   auto yieldOp = cast<linalg::YieldOp>(producerBlock.getTerminator());
   // TODO: allow fusion of multi-result producers.
   assert(producer->getNumResults() == 1 && "expected single result producer");
   unsigned producerResultNumber = 0;
   Value replacement =
       mapper.lookupOrDefault(yieldOp.getOperand(producerResultNumber));

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

 //      CHECK:   %[[VAL2:.+]] = subi %[[VAL1]], %[[SUB_OPERAND]] : i32
 //      CHECK:   %[[VAL3:.+]] = addi %[[VAL2]], %[[ARG3]] : i32
 //      CHECK:   linalg.yield %[[VAL3]] : i32
 //   CHECK-NOT: linalg.generic
 
 // -----
 
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
-func @indexed_producer_consumer_fusion(%arg0: tensor<?x?xi32>,
-                                       %arg1: tensor<?x?xi32>) -> tensor<?x?xi32> {
+func @indexed_producer_consumer_fusion(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %0 = memref.dim %arg0, %c0 : tensor<?x?xi32>
   %1 = memref.dim %arg0, %c1 : tensor<?x?xi32>
   %2 = linalg.init_tensor [%0, %1] : tensor<?x?xi32>
   %3 = linalg.generic {
     indexing_maps = [#map0, #map0],
     iterator_types = ["parallel", "parallel"] }
     ins(%arg0 : tensor<?x?xi32>)
     outs(%2 : tensor<?x?xi32>) {
     ^bb0(%arg4: i32, %arg5: i32):       // no predecessors
       %idx0 = linalg.index 0 : index
       %idx1 = linalg.index 1 : index
       %4 = index_cast %idx0 : index to i32
       %5 = index_cast %idx1 : index to i32
       %6 = addi %arg4, %4 : i32
       %7 = subi %6, %5 : i32
       linalg.yield %7 : i32
     } -> tensor<?x?xi32>
   %4 = linalg.generic {
     indexing_maps = [#map0, #map0, #map0],
     iterator_types = ["parallel", "parallel"] }
-    ins(%3, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>)
+    ins(%3, %arg0 : tensor<?x?xi32>, tensor<?x?xi32>)
     outs(%2 : tensor<?x?xi32>) {
     ^bb0(%arg2: i32, %arg3: i32, %arg4: i32):       // no predecessors
       %10 = addi %arg2, %arg3 : i32
       linalg.yield %10 : i32
     } -> tensor<?x?xi32>
   return %4 : tensor<?x?xi32>
 }
 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-LABEL: func @indexed_producer_consumer_fusion
 //       CHECK: linalg.generic
-// CHECK-SAME:    indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]]]
+// CHECK-SAME:    indexing_maps = [#[[$MAP0]], #[[$MAP0]]]
 //      CHECK: ^{{[a-zA-Z0-9_]*}}
 // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: i32
 // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: i32
 //      CHECK:   %[[IDX0:.+]] = linalg.index 0 : index
 //      CHECK:   %[[IDX1:.+]] = linalg.index 1 : index
 //      CHECK:   %[[ADD_OPERAND:.+]] = index_cast %[[IDX0]] : index to i32
 //      CHECK:   %[[SUB_OPERAND:.+]] = index_cast %[[IDX1]] : index to i32
 //      CHECK:   %[[VAL1:.+]] = addi %[[ARG0]], %[[ADD_OPERAND]] : i32
 //      CHECK:   %[[VAL2:.+]] = subi %[[VAL1]], %[[SUB_OPERAND]] : i32
-//      CHECK:   %[[VAL3:.+]] = addi %[[VAL2]], %[[ARG1]] : i32
+//      CHECK:   %[[VAL3:.+]] = addi %[[VAL2]], %[[ARG0]] : 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)>