[mlir] detensorize: don't accidentally convert function entry blocks

[ftynse]

Summary

In the Linalg detensorize pass, dialect conversion could accidentally
trigger signature conversion of the function entry block after inlining
the body of a Linalg generic into it. Such a conversion is not desirable
because it would break the internal validity of the function op, that is
futhermore not supposed to be detensorized at the boundary. Mitigate
this by creating a dummy (empty) entry block so Linalg operations are
never inlined into it and the conversion is never triggered.

Closes #62249.

[silvas]

Thanks! This makes sense to me :)

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

 bool shouldBeDetensored(Operation *op, TypeConverter typeConverter) {
   GenericOp genericOp = dyn_cast_or_null<GenericOp>(op);
   return genericOp &&
          llvm::all_of(genericOp->getOpOperands(), [&](OpOperand &opOperand) {
            return !typeConverter.isLegal(opOperand.get().getType());
          });
 }
 
-/// A conversion patttern for detensoring `linalg.generic` ops.
+/// A conversion pattern for detensoring `linalg.generic` ops.
 class DetensorizeGenericOp : public OpConversionPattern<GenericOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(GenericOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Block *originalBlock = op->getBlock();
 
-    // Gather some information about the op before inling its region.
+    // Gather some information about the op before inlining its region.
     Block *opEntryBlock = &*op.getRegion().begin();
     YieldOp yieldOp = dyn_cast<YieldOp>(op.getRegion().back().getTerminator());
 
     // Split the op's region before the op. This way, we have a clear insertion
     // point in which the op can be inlined.
     Block *newBlock = rewriter.splitBlock(originalBlock, Block::iterator(op));
     rewriter.inlineRegionBefore(op.getRegion(), newBlock);
     // Now that op's region is inlined, the operands of its YieldOp are mapped
void runOnOperation() override {
     DetensorizeTypeConverter typeConverter;
     RewritePatternSet patterns(context);
     ConversionTarget target(*context);
     DenseSet<Operation *> opsToDetensor;
     DenseMap<Operation *, DenseSet<int>> detensorableBranchOps;
     DenseSet<BlockArgument> blockArgsToDetensor;
     FunctionOpInterface funcOp = getOperation();
 
+    // Make sure the entry block of the function doesn't contain any Linalg ops.
+    // Otherwise, it may lead to the signature of the block being changed by the
+    // dialect conversion below, which would make the function op invalid
+    // because its type shouldn't change.
+    IRRewriter rewriter(funcOp->getContext());
+    Block *entryBlock = &funcOp.getFunctionBody().front();
+    Block *postEntryBlock =
+        rewriter.splitBlock(entryBlock, entryBlock->begin());
+    rewriter.setInsertionPointToStart(entryBlock);
+    auto branch =
+        rewriter.create<cf::BranchOp>(rewriter.getUnknownLoc(), postEntryBlock);
+
     if (aggressiveMode.getValue()) {
       AggressiveDetensoringModel costModel;
       costModel.compute(funcOp, typeConverter, opsToDetensor,
                         blockArgsToDetensor);
     } else {
       ControlFlowDetectionModel costModel;
       costModel.compute(funcOp, typeConverter, opsToDetensor,
                         blockArgsToDetensor);
if (failed(
             applyFullConversion(getOperation(), target, std::move(patterns))))
       signalPassFailure();
 
     RewritePatternSet canonPatterns(context);
     tensor::FromElementsOp::getCanonicalizationPatterns(canonPatterns, context);
     if (failed(applyPatternsAndFoldGreedily(getOperation(),
                                             std::move(canonPatterns))))
       signalPassFailure();
+
+    // Get rid of the dummy entry block we created in the beginning to work
+    // around dialect conversion signature rewriting.
+    rewriter.eraseOp(branch);
+    rewriter.mergeBlocks(postEntryBlock, entryBlock);
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createLinalgDetensorizePass() {
   return std::make_unique<LinalgDetensorize>();
 }

mlir/test/Dialect/Linalg/detensorize_entry_block.mlir

+// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline="builtin.module(func.func(linalg-detensorize))" | FileCheck %s
+
+#map = affine_map<() -> ()>
+func.func @main(%arg0: tensor<f32>) -> tensor<f32> {
+  %0 = tensor.empty() : tensor<f32>
+  %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = []} ins(%arg0 : tensor<f32>) outs(%0 : tensor<f32>) {
+  ^bb0(%in: f32, %out: f32):
+    linalg.yield %in : f32
+  } -> tensor<f32>
+  cf.br ^bb1(%1 : tensor<f32>)
+^bb1(%2: tensor<f32>):  // pred: ^bb0
+  return %2 : tensor<f32>
+}
+
+// CHECK-LABEL: @main
+// CHECK-SAME:       (%[[ARG0:.+]]: tensor<f32>) -> tensor<f32>
+// CHECK:   %[[EXTRACTED:.+]] = tensor.extract %[[ARG0]][] : tensor<f32>
+// CHECK: cf.br ^{{.*}}(%[[EXTRACTED]] : f32)
+// CHECK: ^{{.*}}(%[[ARG1:.+]]: f32):
+// CHECK:   %[[ELEMENTS:.+]] = tensor.from_elements %[[ARG1]] : tensor<f32>
+// CHECK:   return %[[ELEMENTS]] : tensor<f32>

mlir/test/Dialect/Linalg/detensorize_while_pure_cf.mlir

^bb2(%6: tensor<i32>):  // pred: ^bb1
   } -> tensor<i32>
   cf.br ^bb1(%8 : tensor<i32>)
 
 ^bb3:  // pred: ^bb1
   return
 }
 
 // CHECK-LABEL: func @main
-// CHECK-NEXT:    arith.constant 0 : i32
-// CHECK-NEXT:    arith.constant 10
+// CHECK-DAG:    arith.constant 0 : i32
+// CHECK-DAG:    arith.constant 10
 // CHECK-NEXT:    cf.br ^[[bb1:.*]](%{{.*}} : i32)
 // CHECK-NEXT:  ^[[bb1]](%{{.*}}: i32)
 // CHECK-NEXT:    %{{.*}} = arith.cmpi slt, %{{.*}}, %{{.*}}
 // CHECK-NEXT:    cf.cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^[[bb3:.*]]
 // CHECK-NEXT:  ^[[bb2]](%{{.*}}: i32)
 // CHECK-NEXT:    %{{.*}} = arith.addi %{{.*}}, %{{.*}}
 // CHECK-NEXT:    cf.br ^[[bb1]](%{{.*}} : i32)
 // CHECK-NEXT:  ^[[bb3]]:
 // CHECK-NEXT:    return
 // CHECK-NEXT:  }