[mlir][linalg] Support non-function ops in Detensorize

[springerm]

Summary

This fixes #58748.

[mravishankar]

I see why you made it a function pass, but I think it needs to be a pass on FunctionOpInterface. There are examples of this in IREE https://github.com/iree-org/iree/blob/55d0ffe7ea4ca5b58dd9197608151cc2f8c7521c/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td#L65 .

[springerm]

address comments

[mravishankar]

Thanks!

[mehdi_amini]

Why is this tied to the concept of a function? Why wouldn't any isolated op work equivalently here?

[springerm]

In D143655#4134041, @mehdi_amini wrote:

Why is this tied to the concept of a function? Why wouldn't any isolated op work equivalently here?

Maybe. I'm not sure. Some IsolatedFromAbove ops expose operands as block arguments. Changing the type of these block arugments could make the op invalid. There are extra checks in this pass to exclude the bbargs of the first block of a function.

[mehdi_amini]

Where is the type change happening for entry block? I missed it I think

[springerm]

In D143655#4144678, @mehdi_amini wrote:

Where is the type change happening for entry block? I missed it I think

There's a comment in line 496:

// A function is legal if all of its non-entry blocks are legal. We
// don't legalize the entry block (i.e. the function's signature)
// since detensoring can't happen along external calling convention
// boundaries, which we conservatively approximate as all function
// signatures.

It may be possible to wrap this part in AggressiveDetensoringModel::compute etc in a if (auto funcOp = dyn_cast<FunctionOpInterface>()):

for (Block &block : llvm::drop_begin(func.getFunctionBody(), 1))
  for (BlockArgument blockArgument : block.getArguments())
    blockArgsToDetensor.insert(blockArgument);

It's the first time I'm seeing this transformation so I'm a bit reluctant to make larger changes...

[springerm]

update

[mehdi_amini]

In D143655#4144708, @springerm wrote:

In D143655#4144678, @mehdi_amini wrote:

Where is the type change happening for entry block? I missed it I think

There's a comment in line 496:

// A function is legal if all of its non-entry blocks are legal. We
// don't legalize the entry block (i.e. the function's signature)
// since detensoring can't happen along external calling convention
// boundaries, which we conservatively approximate as all function
// signatures.

It may be possible to wrap this part in AggressiveDetensoringModel::compute etc in a if (auto funcOp = dyn_cast<FunctionOpInterface>()):

I don’t follow, isn’t the comment saying that this won’t touch the entry block?
I don’t quite get why any change would be needed, seems like we’re reading this the opposite way.

[springerm]

I don’t follow, isn’t the comment saying that this won’t touch the entry block?

correct, I'm wondering if it is safe to do the same for other ops (non function). or if we should not touch these at all.

[mehdi_amini]

In D143655#4144795, @springerm wrote:

I don’t follow, isn’t the comment saying that this won’t touch the entry block?

correct, I'm wondering if it is safe to do the same for other ops (non function). or if we should not touch these at all.

I don't why why it wouldn't be safe?
That was the point of my comment actually, there is nothing special to "FunctionInterface" here as far as I can tell.

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

public:
     ///   %10 = "some.op"(%9)
     ///   br ^bb2(%8 : tensor<i32>)
     /// ...
     ///
     /// if the cost model decides that the linalg.generic op should be
     /// detensored, then:
     /// - opsToDetensor should be = {linalg.generic{add}}.
     /// - blockArgsToDetensor should be = {bb1 -> {0}, bb2 -> {0}}.
-    virtual void compute(FunctionOpInterface func,
-                         DetensorizeTypeConverter typeConverter,
+    virtual void compute(Operation *op, DetensorizeTypeConverter typeConverter,
                          DenseSet<Operation *> &opsToDetensor,
                          DenseSet<BlockArgument> &blockArgsToDetensor) = 0;
 
     /// From the blockArgsToDetensor set computed by a CostModel
     /// implementation, this method computes the corresponding branch op
     /// detensoring. The result is a map from a branch op to a subset of indices
     /// of its operands. The indices specify which of the branch op's operands
     /// should be detensored.
struct LinalgDetensorize
   /// comes from. If the condition value is (indirectly) computed by a linalg op
   /// that can be detensored, the model then continues walking the use-def chain
   /// in order to understand where the linalg op's operands come from. This
   /// leads to discovering a "detensoring component". A detensoring component is
   /// the set of operations + block arguments that are involved in control-flow
   /// AND can be detensored.
   class ControlFlowDetectionModel : public CostModel {
   public:
-    void compute(FunctionOpInterface func,
-                 DetensorizeTypeConverter typeConverter,
+    void compute(Operation *op, DetensorizeTypeConverter typeConverter,
                  DenseSet<Operation *> &opsToDetensor,
                  DenseSet<BlockArgument> &blockArgsToDetensor) override {
       SmallVector<Value> workList;
 
-      func->walk([&](cf::CondBranchOp condBr) {
+      op->walk([&](cf::CondBranchOp condBr) {
         llvm::append_range(workList, condBr.getOperands());
       });
 
-      func->walk([&](cf::BranchOp br) {
+      op->walk([&](cf::BranchOp br) {
         llvm::append_range(workList, br.getOperands());
       });
 
       DenseSet<Value> visitedValues;
       DenseSet<Operation *> visitedOps;
 
       // For a (to-be-detesored) value, check if it "escapes" the block by being
       // passed to terminator. If it does, then workList is updated with the
void compute(Operation *op, DetensorizeTypeConverter typeConverter,
         blockArgsToDetensor.erase(blockArg);
       }
     }
   };
 
   /// Detensorize everything that can detensored.
   class AggressiveDetensoringModel : public CostModel {
   public:
-    void compute(FunctionOpInterface func,
-                 DetensorizeTypeConverter typeConverter,
+    void compute(Operation *op, DetensorizeTypeConverter typeConverter,
                  DenseSet<Operation *> &opsToDetensor,
                  DenseSet<BlockArgument> &blockArgsToDetensor) override {
-      func->walk([&](GenericOp genericOp) {
+      op->walk([&](GenericOp genericOp) {
         if (shouldBeDetensored(genericOp, typeConverter))
           opsToDetensor.insert(genericOp);
       });
 
+      if (auto func = dyn_cast<FunctionOpInterface>(op)) {
         for (Block &block : llvm::drop_begin(func.getFunctionBody(), 1))
           for (BlockArgument blockArgument : block.getArguments())
             blockArgsToDetensor.insert(blockArgument);
       }
+    }
   };
 
   void runOnOperation() override {
     MLIRContext *context = &getContext();
     DetensorizeTypeConverter typeConverter;
     RewritePatternSet patterns(context);
     ConversionTarget target(*context);
     DenseSet<Operation *> opsToDetensor;
     DenseMap<Operation *, DenseSet<int>> detensorableBranchOps;
     DenseSet<BlockArgument> blockArgsToDetensor;
-    FunctionOpInterface funcOp = cast<FunctionOpInterface>(getOperation());
+    Operation *op = getOperation();
 
     if (aggressiveMode.getValue()) {
       AggressiveDetensoringModel costModel;
-      costModel.compute(funcOp, typeConverter, opsToDetensor,
-                        blockArgsToDetensor);
+      costModel.compute(op, typeConverter, opsToDetensor, blockArgsToDetensor);
     } else {
       ControlFlowDetectionModel costModel;
-      costModel.compute(funcOp, typeConverter, opsToDetensor,
-                        blockArgsToDetensor);
+      costModel.compute(op, typeConverter, opsToDetensor, blockArgsToDetensor);
     }
 
     detensorableBranchOps =
         CostModel::computeBranchOpDetensoring(blockArgsToDetensor);
 
     target.addDynamicallyLegalOp<GenericOp>(
         [&](GenericOp op) { return !opsToDetensor.count(op); });
 

mlir/test/Dialect/Linalg/detensorize_0d.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline="builtin.module(func.func(linalg-detensorize{aggressive-mode}))" | FileCheck %s
+// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-detensorize="aggressive-mode" | FileCheck %s
 
 #map = affine_map<() -> ()>
 
 func.func @detensor_simple(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} {
   %0 = tensor.empty() : tensor<f32>
   %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []}
     ins(%arg1, %arg2 : tensor<f32>, tensor<f32>)
     outs(%0 : tensor<f32>) {

 }
 // CHECK-LABEL: func @detensor_foreign_op
 // CHECK-SAME:    (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>)
 // CHECK-DAG:     %[[arg1_val:.*]] = tensor.extract %[[arg1]]
 // CHECK-DAG:     %[[arg2_val:.*]] = tensor.extract %[[arg2]]
 // CHECK:         %[[detensored_res:.*]] = "foreign.do_something"(%[[arg1_val]], %[[arg2_val]])
 // CHECK:         %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res]]
 // CHECK:         return %[[new_tensor_res]]
+
+// CHECK-LABEL: func @regression()
+func.func @regression() {
+  %0 = tensor.empty() : tensor<4xf32>
+  return
+}