[MLIR][LinAlg] Detensorize interal function control flow.

[ergawy]

Summary

This patch continues detensorizing implementation by detensoring
internal control flow in functions.

In order to detensorize functions, all the non-entry block's arguments
are detensored and branches between such blocks are properly updated to
reflect the detensored types as well. Function entry block (signature)
is left intact.

This continues work towards handling github/google/iree#1159.

[ergawy]

Assert on 0 inputs.

[ergawy]

mlir/include/mlir/Dialect/Linalg/Passes.h

65

Eventually, I will not implement this as a separate pass. Both passes will be merged into one. I am doing this temporarily to make understanding how dialect conversion works a bit better (i.e. to have a clear separation point between internal conversion within a BB and conversion across the BB boundary in a function).

[ergawy]

Properly handle non-entry function blocks.

[ergawy]

Update docs.

[silvas]

mlir/include/mlir/Dialect/Linalg/Passes.td

156

nit: indent is weird

mlir/include/mlir/Dialect/StandardOps/Transforms/FuncConversions.h

37

"to return" reads weird to me

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

176

more specifically: "detensoring can't happen along external calling convention boundaries, which we conservatively approximate as all function signatures"

187

seems like there is some duplication with func-bufferize. can you put a helper next to the populate*Pattern function (populate*Legality)?

mlir/test/Dialect/Linalg/detensorized_while.mlir

4

remove iree.module.export

8

can you save these attributes in a #attr_alias (see other uses in the codebase). that would make the code somewhat simpler. E.g. https://github.com/llvm/llvm-project/blob/e339bba637b941c8e78057319b7654c4babf18cb/mlir/test/Dialect/Linalg/tile.mlir#L343

28

check: tensor.extract

[ergawy]

Handle review comments.

[silvas]

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

124

I don't think that hasTrait<ReturnLike> is relevant to the branch pattern.

[rriddle]

mlir/include/mlir/Dialect/StandardOps/Transforms/FuncConversions.h

42

Just a note but, conversion target only supports one "unknown legality" callback function.

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

18

Is this really necessary?

90

Functions should not be placed inside of anonymous namespaces. They should be marked as static and at the top level.

103

The .getResult() shouldn't be necessary.

[ergawy]

Handle some review comments.

mlir/include/mlir/Dialect/StandardOps/Transforms/FuncConversions.h

42

I noticed that and and I think that might be an argument against adding this utility function. I mean it is not general enough, it's logic just happens to be needed in 2 different places. If the needs of one pass diverges from the other we might end up inlining the code again.

I guess for now we can keep the util function and see if the need arises to inline it again at call sites in the future. Let me know if you disagree.

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

124

If we remove this condition, then ops that are ReturnLike but not BranchOpInterface will be marked as illegal.

I don't disagree that this should ideally be removed. However, I guess we will have to change the expected behavior a little bit.

I think we can rewrite the entire lambda as follows:

target.markUnknownOpDynamicallyLegal([&](Operation *op) {
  // All successor operands of branch like operations must be rewritten.
  if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
    for (int p = 0, e = op->getBlock()->getNumSuccessors(); p < e; ++p) {
      auto successorOperands = branchOp.getSuccessorOperands(p);
      if (successorOperands.hasValue() &&
          !converter.isLegal(successorOperands.getValue().getTypes()))
        return false;
    }
  }

  return true;
});

With the above, all bufferization and detensorization tests pass fine except for one: https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Standard/func-bufferize.mlir#L65. In particular, test.terminator won't be marked as illegal as before.

I am not talking about this particular test of course. In general, any op that is not explicitly BranchOpInterface but might have IsTerminator will be legal. Would that be fine?

[ergawy]

mlir/include/mlir/Dialect/StandardOps/Transforms/FuncConversions.h

42

Oh by the way, is it by design that conversion target supports exactly one unknown legality callback? Or is it something that might be extended to support multiple ones?

I think multiple ones can open the door for somewhat general utils like the above and allow customization/additional behavior on top of that. But it would complicate things by quite a bit as well I guess and there might be no need for it.

If you think extension to multiple callbacks is worth doing, I would gladly give this a try.

[silvas]

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

124

I think the core issue here is that we are conflating legality for populateReturnOpTypeConversionPattern with legality for populateBranchOpInterfaceTypeConversionPattern.

I think we should split this into two functions bool isLegalForBranchOpInterfaceTypeConversionPattern(Operation *op, TypeConverter &converter) and bool isLegalForReturnOpTypeConversionPattern(Operation *op) (it is important that these be split so that they correspond very tightly with their corresponding populate*Pattern function). It i then up to the user to call manually if they want both:

target.markUnknownOpDynamicallyLegal([&](Operation *op) {
  return isLegalForBranchOpInterfaceTypeConversionPattern(op, converter) || isLegalForReturnOpTypeConversionPattern(op);
})

to work around the fact that only one legality callback is allowed in markUnknownOpDynamicallyLegal.

[ergawy]

Handle review comments.

[ergawy]

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

124

I adopted your suggestion with a slight change. I broke up the previous lambda into 3 functions to make the purpose of each one more fine-grained. Let me know if this is not a good approach.

[ergawy]

Clean-up.

[ergawy]

Rename function.

[silvas]

One conceptual concern, but otherwise this LGTM.

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

111

The isNotBranchOpInterfaceOrReturnLikeOp is kind of annoying, and I'm sorry for getting the logic here backward in my suggestion. I realize now that for this to compose properly it basically needs to be:

target.markUnknownOpDynamicallyLegal([&](Operation *op) {
  if (needsToBeLegalizedByBranchOpInterfaceTypeConversion(op, converter))
    return false;
  if (needsToBeLegalizedByReturnOpTypeConversion(op, converter))
    return false;
  return true;
}

With a slight modification to the current code, we can write it as !isLegalForBranchOpInterfaceTypeConversionPattern(...)) return false;. The current code returns "false" for "needs to convert BranchOpInterface" and also for "i have no idea about this op". It should only return "false" for "needs to convert BranchOpInterface".

This dialect conversion / legality stuff is super tricky, sorry.

133

I don't understand why this function has to worry about ReturnLike. populateReturnOpTypeConversionPattern only handles std.return, and this function should only be concerned with that one.

[ergawy]

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

111

No problem, happy to discuss until we get it right. And apologies if I am not getting what you suggest sometimes.

With your suggestion in the above comment, we won't be able to report an error whenever a terminator neither implements BranchOpInterface nor is a ReturnOp. Such terminator would be marked legal. In order to handle this, I think we would need to:

• either centralize the legality check in one place for return-like ops and branch op interface in one place like you had before.
• or provide a catch-all utility like isNotBranchOpInterfaceOrReturnLikeOp that would mark unknown terminator ops as illegal and mark all other ops as legal.

[ergawy]

• Rebase.
• Add a TODO.

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

111

I added a TODO to revisit this and try getting rid of isNotBranchOpInterfaceOrReturnLikeOp. I will commit the patch and if you really don't like, I will handle this in a follow up patch right away.

[silvas]

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

111

ah, I see. yeah that's tricky and would have the issue you describe.

[mehdi_amini]

Nit: fix the typo in the commit title, "interal"

mlir/include/mlir/Dialect/Linalg/Passes.h

 
 /// Create a pass to conver named Linalg operations to Linalg generic
 /// operations.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgGeneralizationPass();
 
 /// Create a pass to convert Linalg operations to equivalent operations that
 /// work on primitive types, if possible.
 std::unique_ptr<Pass> createLinalgDetensorizePass();
+std::unique_ptr<Pass> createFuncDetensorizePass();
 
 /// Patterns to fold an expanding (collapsing) tensor_reshape operation with its
 /// producer (consumer) generic operation by expanding the dimensionality of the
 /// loop in the generic op.
 void populateFoldReshapeOpsByExpansionPatterns(
     MLIRContext *context, OwningRewritePatternList &patterns);
 
 /// Patterns to fold a collapsing (expanding) tensor_reshape operation with its

mlir/include/mlir/Dialect/Linalg/Passes.td

let description = [{
     on primitives.
 
     The detensoring process is driven by linalg-on-tensor ops. In particular, a
     linalg-on-tensor op is checked to see whether *all* its operands can be
     detensored. If so, those operands are converted to their primitive
     counterparts and the linalg op is replaced by an equivalent op that takes
     those new primitive values as operands. Therefore, the detensoring process
     can be divided into 2 main logical phases:
 
     1. Detect/match an op that can be detensored.
     2. Detensor the operands of the op and replace it with a primitive
        equivalent.
   }];
 }
 
+def FuncDetensorize : FunctionPass<"func-detensorize"> {
+  let summary = "TODO";
+  let constructor = "mlir::createFuncDetensorizePass()";
+}
+
 #endif // MLIR_DIALECT_LINALG_PASSES

mlir/include/mlir/Dialect/StandardOps/Transforms/FuncConversions.h

 void populateCallOpTypeConversionPattern(OwningRewritePatternList &patterns,
                                          MLIRContext *ctx,
                                          TypeConverter &converter);
 
 /// Add a pattern to the given pattern list to rewrite branch operations and
 /// `return` to use operands that have been legalized by the conversion
 /// framework. This can only be done if the branch operation implements the
 /// BranchOpInterface. Only needed for partial conversions.
-void populateBranchOpInterfaceAndReturnOpTypeConversionPattern(
+void populateBranchOpInterfaceTypeConversionPattern(
     OwningRewritePatternList &patterns, MLIRContext *ctx,
     TypeConverter &converter);
+void populateReturnOpTypeConversionPattern(OwningRewritePatternList &patterns,
+                                           MLIRContext *ctx,
+                                           TypeConverter &converter);
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_STANDARDOPS_TRANSFORMS_FUNCCONVERSIONS_H_

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

 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/StandardOps/Transforms/FuncConversions.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include <iterator>
 #include <memory>
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 namespace {
 /// Defines the criteria a TensorType must follow in order to be considered
 /// "detensorable".
DetensorizeTypeConverter() {
     addTargetMaterialization([](OpBuilder &builder, Type type,
                                 ValueRange inputs, Location loc) -> Value {
       return builder.create<tensor::ExtractOp>(loc, inputs[0], ValueRange{});
     });
 
     // A detensored value is converted back by creating a new tensor from its
     // element(s).
     addSourceMaterialization([](OpBuilder &builder, Type type,
                                 ValueRange inputs, Location loc) -> Value {
+      if (inputs.empty())
+        return builder.create<ConstantOp>(
+            loc, builder.getI32Type(),
+            builder.getZeroAttr(builder.getI32Type()));
+      for (auto input : inputs) {
+        input.dump();
+      }
+      auto createNewTensorOp = builder.create<tensor::FromElementsOp>(
+          loc, inputs[0].getType(), inputs[0]);
+
+      // FromElementsOp results in a tensor<1xdtype>, we need to reshape that to
+      // a tensor<dtype> instead.
+      return builder
+          .create<linalg::TensorReshapeOp>(loc, type, createNewTensorOp,
+                                           ArrayRef<ReassociationExprs>{})
+          .getResult();
+    });
+
+    addArgumentMaterialization([](OpBuilder &builder, Type type,
+                                  ValueRange inputs, Location loc) -> Value {
       auto createNewTensorOp = builder.create<tensor::FromElementsOp>(
           loc, inputs[0].getType(), inputs[0]);
 
       // FromElementsOp results in a tensor<1xdtype>, we need to reshape that to
       // a tensor<dtype> instead.
       return builder.create<linalg::TensorReshapeOp>(
           loc, type, createNewTensorOp, ArrayRef<ReassociationExprs>{});
     });
void runOnFunction() override {
     OwningRewritePatternList patterns;
     ConversionTarget target(*context);
 
     target.markUnknownOpDynamicallyLegal([](Operation *op) { return true; });
     target.addLegalDialect<linalg::LinalgDialect>();
     target.addDynamicallyLegalOp<GenericOp>([&](GenericOp op) {
       // If any of the operands or results cannot be detensored, the op is
       // considered legal and won't be detensored.
-      return llvm::any_of(
-          op.getShapedOperandTypes(), [](ShapedType shapedType) {
-            assert(shapedType.isa<TensorType>());
-            return !canBeDetensored(shapedType.cast<TensorType>());
+      return llvm::any_of(op.getShapedOperandTypes(),
+                          [&](ShapedType shapedType) {
+                            return typeConverter.isLegal(shapedType);
                           });
     });
 
     patterns.insert<DetensorizeGenericOp>(typeConverter, context);
 
     if (failed(
             applyPartialConversion(getFunction(), target, std::move(patterns))))
       signalPassFailure();
 
     OwningRewritePatternList canonPatterns;
     canonPatterns.insert<ExtractFromReshapeFromElements>(context);
     if (failed(applyPatternsAndFoldGreedily(getFunction(),
                                             std::move(canonPatterns))))
       signalPassFailure();
+  }
+};
+
+struct FunctionInternalBlockConversion : public ConversionPattern {
+  FunctionInternalBlockConversion(StringRef functionLikeOpName,
+                                  MLIRContext *ctx, TypeConverter &converter)
+      : ConversionPattern(functionLikeOpName, /*benefit=*/1, converter, ctx) {}
 
-    // TODO Properly handle control flow within function boundaries.
+  /// Hook to implement combined matching and rewriting for FunctionLike ops.
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    FunctionType type = mlir::impl::getFunctionType(op);
+    rewriter.startRootUpdate(op);
+
+    // Convert the original function types.
+    TypeConverter::SignatureConversion result(type.getNumInputs());
+    result.addInputs(type.getInputs());
+
+    SmallVector<Type, 1> newResults;
+    if (failed(rewriter.convertRegionTypes(&mlir::impl::getFunctionBody(op),
+                                           *typeConverter, &result))) {
+      rewriter.cancelRootUpdate(op);
+      return failure();
+    }
+
+    rewriter.finalizeRootUpdate(op);
+    return success();
+  }
+};
+
+struct FuncDetensorize : public FuncDetensorizeBase<FuncDetensorize> {
+  void runOnFunction() override {
+    auto func = getFunction();
+    auto *context = &getContext();
+
+    DetensorizeTypeConverter typeConverter;
+    OwningRewritePatternList patterns;
+    ConversionTarget target(*context);
+
+    patterns.insert<FunctionInternalBlockConversion>(FuncOp::getOperationName(),
+                                                     context, typeConverter);
+    populateBranchOpInterfaceTypeConversionPattern(patterns, context,
+                                                   typeConverter);
+
+    target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
+      return std::all_of(std::next(op.getBody().begin()), op.getBody().end(),
+                         [&](Block &block) {
+                           return typeConverter.isLegal(
+                               block.getArgumentTypes());
+                         });
+    });
+
+    target.addLegalOp<ModuleOp, ModuleTerminatorOp, ReturnOp>();
+
+    // Mark terminators as legal if they have the ReturnLike trait or
+    // implement the BranchOpInterface and have valid types. If they do not
+    // implement the trait or interface, mark them as illegal no matter what.
+    target.markUnknownOpDynamicallyLegal([&](Operation *op) {
+      // return true;
+      // If it is not a terminator, ignore it.
+      if (!op->mightHaveTrait<OpTrait::IsTerminator>())
+        return true;
+      // If it is not the last operation in the block, also ignore it. We do
+      // this to handle unknown operations, as well.
+      Block *block = op->getBlock();
+      if (!block || &block->back() != op)
+        return true;
+      // ReturnLike operations have to be legalized with their parent. For
+      // return this is handled, for other ops they remain as is.
+      if (op->hasTrait<OpTrait::ReturnLike>())
+        return true;
+      // All successor operands of branch like operations must be rewritten.
+      if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
+        for (int p = 0, e = op->getBlock()->getNumSuccessors(); p < e; ++p) {
+          auto successorOperands = branchOp.getSuccessorOperands(p);
+          if (successorOperands.hasValue() &&
+              !typeConverter.isLegal(successorOperands.getValue().getTypes()))
+            return false;
+        }
+        return true;
+      }
+      return false;
+    });
+
+    if (failed(applyPartialConversion(func, target, std::move(patterns))))
+      signalPassFailure();
+
+    OwningRewritePatternList canonPatterns;
+    canonPatterns.insert<ExtractFromReshapeFromElements>(context);
+
+    if (failed(applyPatternsAndFoldGreedily(func, std::move(canonPatterns))))
+      signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createLinalgDetensorizePass() {
   return std::make_unique<LinalgDetensorize>();
 }
+
+std::unique_ptr<Pass> mlir::createFuncDetensorizePass() {
+  return std::make_unique<FuncDetensorize>();
+}

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp

void runOnOperation() override {
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
       return typeConverter.isSignatureLegal(op.getType()) &&
              typeConverter.isLegal(&op.getBody());
     });
     populateCallOpTypeConversionPattern(patterns, context, typeConverter);
     target.addDynamicallyLegalOp<CallOp>(
         [&](CallOp op) { return typeConverter.isLegal(op); });
 
-    populateBranchOpInterfaceAndReturnOpTypeConversionPattern(patterns, context,
-                                                              typeConverter);
+    populateBranchOpInterfaceTypeConversionPattern(patterns, context,
+    populateReturnOpTypeConversionPattern(patterns, context, typeConverter);
     target.addLegalOp<ModuleOp, ModuleTerminatorOp, TensorLoadOp,
                       TensorToMemrefOp>();
     target.addDynamicallyLegalOp<ReturnOp>(
         [&](ReturnOp op) { return typeConverter.isLegal(op); });
     // Mark terminators as legal if they have the ReturnLike trait or
     // implement the BranchOpInterface and have valid types. If they do not
     // implement the trait or interface, mark them as illegal no matter what.
     target.markUnknownOpDynamicallyLegal([&](Operation *op) {

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp

matchAndRewrite(ReturnOp op, ArrayRef<Value> operands,
     Operation *operation = op.getOperation();
     rewriter.updateRootInPlace(
         op, [operands, operation]() { operation->setOperands(operands); });
     return success();
   }
 };
 } // end anonymous namespace
 
-void mlir::populateBranchOpInterfaceAndReturnOpTypeConversionPattern(
+void mlir::populateBranchOpInterfaceTypeConversionPattern(
     OwningRewritePatternList &patterns, MLIRContext *ctx,
     TypeConverter &typeConverter) {
-  patterns.insert<BranchOpInterfaceTypeConversion, ReturnOpTypeConversion>(
-      typeConverter, ctx);
+  patterns.insert<BranchOpInterfaceTypeConversion>(typeConverter, ctx);
+}
+
+void mlir::populateReturnOpTypeConversionPattern(
+    OwningRewritePatternList &patterns, MLIRContext *ctx,
+    TypeConverter &typeConverter) {
+  patterns.insert<ReturnOpTypeConversion>(typeConverter, ctx);
 }

mlir/test/Dialect/Linalg/detensorized_while.mlir

+// RUN: mlir-opt %s -linalg-detensorize -func-detensorize -canonicalize | FileCheck %s
+
+func @main(%cst: tensor<i32>, %cst_0: tensor<i32>) -> tensor<i32> attributes {iree.module.export} {
+//  %cst = constant dense<1> : tensor<i32>
+//  %cst_0 = constant dense<3> : tensor<i32>
+  br ^bb1(%cst : tensor<i32>)
+^bb1(%0: tensor<i32>):  // 2 preds: ^bb0, ^bb2
+  %1 = linalg.init_tensor [] : tensor<i1>
+  %2 = linalg.generic {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} ins(%0, %cst_0 : tensor<i32>, tensor<i32>) outs(%1 : tensor<i1>) {
+  ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):  // no predecessors
+    %8 = cmpi slt, %arg0, %arg1 : i32
+    linalg.yield %8 : i1
+  } -> tensor<i1>
+  %3 = tensor.extract %2[] : tensor<i1>
+  cond_br %3, ^bb2(%0 : tensor<i32>), ^bb3(%0 : tensor<i32>)
+^bb2(%4: tensor<i32>):  // pred: ^bb1
+  %5 = linalg.init_tensor [] : tensor<i32>
+  %6 = linalg.generic {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} ins(%4, %4 : tensor<i32>, tensor<i32>) outs(%5 : tensor<i32>) {
+  ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):  // no predecessors
+    %8 = addi %arg0, %arg1 : i32
+    linalg.yield %8 : i32
+  } -> tensor<i32>
+  br ^bb1(%6 : tensor<i32>)
+^bb3(%7: tensor<i32>):  // pred: ^bb1
+  return %7 : tensor<i32>
+}
+// CHECK-LABEL: func @main