[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 ↗

(On Diff #325302)

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

155

nit: indent is weird

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

39

"to return" reads weird to me

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

182

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

193

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

44

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?

104

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

117

The .getResult() shouldn't be necessary.

[ergawy]

Handle some review comments.

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

44

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

44

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.td

let description = [{
     or potentially more primitive value(s). During this process, operations with
     such detensored operands are also converted to an equivalent form that works
     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:
+    those new primitive values as operands. Therefore, detensoring an op 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.
+
+    In addition to detensoring individual ops, this pass detensors internal
+    control flow inside a function. All blocks except for the entry block are
+    detensored by converting their arguments whenever possible.
   }];
 }
 
 #endif // MLIR_DIALECT_LINALG_PASSES

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

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_STANDARDOPS_TRANSFORMS_FUNCCONVERSIONS_H_
 #define MLIR_DIALECT_STANDARDOPS_TRANSFORMS_FUNCCONVERSIONS_H_
 
 namespace mlir {
 
 // Forward declarations.
+class ConversionTarget;
 class MLIRContext;
+class Operation;
 class OwningRewritePatternList;
 class TypeConverter;
 
 /// Add a pattern to the given pattern list to convert the operand and result
 /// types of a CallOp with the given type converter.
 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(
+/// Add a pattern to the given pattern list to rewrite branch operations 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 populateBranchOpInterfaceTypeConversionPattern(
     OwningRewritePatternList &patterns, MLIRContext *ctx,
     TypeConverter &converter);
+
+/// 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.
+///
+/// Note: Non-terminator ops are marked as legal.
+void populateBranchOpInterfaceLegality(ConversionTarget &target,
+                                       TypeConverter &converter);
+
+/// Return true if op is a BranchOpInterface op whose operands are all legal
+/// according to converter.
+bool isLegalForBranchOpInterfaceTypeConversionPattern(Operation *op,
+                                                      TypeConverter &converter);
+
+/// Add a pattern to the given pattern list to rewrite `return` ops to use
+/// operands that have been legalized by the conversion framework.
+void populateReturnOpTypeConversionPattern(OwningRewritePatternList &patterns,
+                                           MLIRContext *ctx,
+                                           TypeConverter &converter);
+
+/// For ReturnLike ops (except `return`), return True. If op is a `return` &&
+/// returnOpAlwaysLegal is false, legalize op according to converter. Otherwise,
+/// return false.
+bool isLegalForReturnOpTypeConversionPattern(Operation *op,
+                                             TypeConverter &converter,
+                                             bool returnOpAlwaysLegal = false);
+
+/// Return true if op is neither BranchOpInterface nor ReturnLike.
+bool isNotBranchOrReturnLikeOp(Operation *op);
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_STANDARDOPS_TRANSFORMS_FUNCCONVERSIONS_H_

mlir/include/mlir/Transforms/DialectConversion.h

public:
   /// replaces each block with a new block containing the updated signature. The
   /// entry block may have a special conversion if `entryConversion` is
   /// provided. On success, the new entry block to the region is returned for
   /// convenience. Otherwise, failure is returned.
   FailureOr<Block *> convertRegionTypes(
       Region *region, TypeConverter &converter,
       TypeConverter::SignatureConversion *entryConversion = nullptr);
 
+  /// Convert the types of block arguments within the given region except for
+  /// the entry region. This replaces each non-entry block with a new block
+  /// containing the updated signature.
+  LogicalResult convertNonEntryRegionTypes(Region *region,
+                                           TypeConverter &converter);
+
   /// Replace all the uses of the block argument `from` with value `to`.
   void replaceUsesOfBlockArgument(BlockArgument from, Value to);
 
   /// Return the converted value that replaces 'key'. Return 'key' if there is
   /// no such a converted value.
   Value getRemappedValue(Value key);
 
   //===--------------------------------------------------------------------===//

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;
 
+static Value sourceMaterializationCallback(OpBuilder &builder, Type type,
+                                           ValueRange inputs, Location loc) {
+  assert(inputs.size() == 1);
+  // A detensored value is converted back by creating a new tensor from its
+  // element(s).
+  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>{});
+}
+
 namespace {
 /// Defines the criteria a TensorType must follow in order to be considered
 /// "detensorable".
 ///
 /// NOTE: For now, only 0-D are supported.
 ///
 /// Returns true if tensorType can be detensored.
 bool canBeDetensored(TensorType tensorType) {
matchAndRewrite(GenericOp op, ArrayRef<Value> operands,
     rewriter.mergeBlocks(newBlock, originalBlock, {});
 
     rewriter.eraseOp(&*Block::iterator(yieldOp));
 
     return success();
   }
 };
 
+/// A conversion pattern for detensoring internal (non-entry) blocks within a
+/// function.
+struct FunctionNonEntryBlockConversion : public ConversionPattern {
+  FunctionNonEntryBlockConversion(StringRef functionLikeOpName,
+                                  MLIRContext *ctx, TypeConverter &converter)
+      : ConversionPattern(functionLikeOpName, /*benefit=*/1, converter, ctx) {}
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.startRootUpdate(op);
+
+    if (failed(rewriter.convertNonEntryRegionTypes(
+            &mlir::impl::getFunctionBody(op), *typeConverter))) {
+      rewriter.cancelRootUpdate(op);
+      return failure();
+    }
+
+    rewriter.finalizeRootUpdate(op);
+    return success();
+  }
+};
+
 class DetensorizeTypeConverter : public TypeConverter {
 public:
   DetensorizeTypeConverter() {
     addConversion([](Type type) { return type; });
 
     // A TensorType that can be detensored, is converted to the underlying
     // element type.
     addConversion([](TensorType tensorType) -> Type {
       if (canBeDetensored(tensorType))
         return tensorType.getElementType();
 
       return tensorType;
     });
 
     // A tensor value is detensoried by extracting its element(s).
     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 {
-      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>{});
-    });
+    addSourceMaterialization(sourceMaterializationCallback);
+    addArgumentMaterialization(sourceMaterializationCallback);
   }
 };
 
 /// Canonicalizes the pattern of the form
 ///
 /// %tensor = tensor.from_elements(%element) : (i32) -> tensor<1xi32>
 /// %reshaped_tensor = linalg.tensor_reshape %tensor [] : tensor<1xi32> into
 ///   tensor<i32>

 /// @see LinalgDetensorize in Linalg/Passes.td for more details.
 struct LinalgDetensorize : public LinalgDetensorizeBase<LinalgDetensorize> {
   void runOnFunction() override {
     auto *context = &getContext();
     DetensorizeTypeConverter typeConverter;
     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>());
+      // If any of the operands or results cannot be detensored (i.e. they are
+      // all legal according the DetensorizeTypeConverter), the op is considered
+      // legal and won't be detensored.
+      return llvm::any_of(op.getShapedOperandTypes(),
+                          [&](ShapedType shapedType) {
+                            return typeConverter.isLegal(shapedType);
                           });
     });
 
+    target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
+      // 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.
+      return llvm::all_of(llvm::drop_begin(op.getBody(), 1), [&](Block &block) {
+        return typeConverter.isLegal(block.getArgumentTypes());
+      });
+    });
+
+    target.markUnknownOpDynamicallyLegal([&](Operation *op) {
+      return isNotBranchOrReturnLikeOp(op) ||
+             isLegalForBranchOpInterfaceTypeConversionPattern(op,
+                                                              typeConverter) ||
+             isLegalForReturnOpTypeConversionPattern(
+                 op, typeConverter, /*returnOpAlwaysLegal*/ true);
+    });
+
     patterns.insert<DetensorizeGenericOp>(typeConverter, context);
+    patterns.insert<FunctionNonEntryBlockConversion>(FuncOp::getOperationName(),
+                                                     context, typeConverter);
+    // Since non-entry block arguments get detensorized, we also need to update
+    // the control flow inside the function to reflect the correct types.
+    populateBranchOpInterfaceTypeConversionPattern(patterns, context,
+                                                   typeConverter);
 
-    if (failed(
-            applyPartialConversion(getFunction(), target, std::move(patterns))))
+    if (failed(applyFullConversion(getFunction(), target, std::move(patterns))))
       signalPassFailure();
 
     OwningRewritePatternList canonPatterns;
     canonPatterns.insert<ExtractFromReshapeFromElements>(context);
     if (failed(applyPatternsAndFoldGreedily(getFunction(),
                                             std::move(canonPatterns))))
       signalPassFailure();
-
-    // TODO Properly handle control flow within function boundaries.
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createLinalgDetensorizePass() {
   return std::make_unique<LinalgDetensorize>();
 }

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) {
-      // 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;
+      return isNotBranchOrReturnLikeOp(op) ||
+             isLegalForBranchOpInterfaceTypeConversionPattern(op,
+                                                              typeConverter) ||
+             isLegalForReturnOpTypeConversionPattern(op, typeConverter);
     });
 
     if (failed(applyFullConversion(module, target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createFuncBufferizePass() {
   return std::make_unique<FuncBufferizePass>();
 }

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);
+}
+
+bool mlir::isLegalForBranchOpInterfaceTypeConversionPattern(
+    Operation *op, TypeConverter &converter) {
+  // 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;
+  }
+
+  return false;
+}
+
+void mlir::populateReturnOpTypeConversionPattern(
+    OwningRewritePatternList &patterns, MLIRContext *ctx,
+    TypeConverter &typeConverter) {
+  patterns.insert<ReturnOpTypeConversion>(typeConverter, ctx);
+}
+
+bool mlir::isLegalForReturnOpTypeConversionPattern(Operation *op,
+                                                   TypeConverter &converter,
+                                                   bool returnOpAlwaysLegal) {
+  // If this is a `return` and the user pass wants to convert/transform across
+  // function boundaries, then `converter` is invoked to check whether the the
+  // `return` op is legal.
+  if (dyn_cast<ReturnOp>(op) && !returnOpAlwaysLegal)
+    return converter.isLegal(op);
+
+  // 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;
+
+  return false;
+}
+
+bool mlir::isNotBranchOrReturnLikeOp(Operation *op) {
+  // 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;
+
+  return false;
 }

mlir/lib/Transforms/Utils/DialectConversion.cpp

struct ConversionPatternRewriterImpl {
   applySignatureConversion(Region *region,
                            TypeConverter::SignatureConversion &conversion);
 
   /// Convert the types of block arguments within the given region.
   FailureOr<Block *>
   convertRegionTypes(Region *region, TypeConverter &converter,
                      TypeConverter::SignatureConversion *entryConversion);
 
+  /// Convert the types of non-entry block arguments within the given region.
+  LogicalResult convertNonEntryRegionTypes(Region *region,
+                                           TypeConverter &converter);
+
   //===--------------------------------------------------------------------===//
   // Rewriter Notification Hooks
   //===--------------------------------------------------------------------===//
 
   /// PatternRewriter hook for replacing the results of an operation.
   void notifyOpReplaced(Operation *op, ValueRange newValues);
 
   /// Notifies that a block is about to be erased.

 
 FailureOr<Block *> ConversionPatternRewriterImpl::convertRegionTypes(
     Region *region, TypeConverter &converter,
     TypeConverter::SignatureConversion *entryConversion) {
   argConverter.setConverter(region, &converter);
   if (region->empty())
     return nullptr;
 
-  // Convert the arguments of each block within the region.
+  if (failed(convertNonEntryRegionTypes(region, converter)))
+    return failure();
+
   FailureOr<Block *> newEntry =
       convertBlockSignature(&region->front(), converter, entryConversion);
+  return newEntry;
+}
+
+LogicalResult ConversionPatternRewriterImpl::convertNonEntryRegionTypes(
+    Region *region, TypeConverter &converter) {
+  argConverter.setConverter(region, &converter);
+  if (region->empty())
+    return success();
+
+  // Convert the arguments of each block within the region.
   for (Block &block : llvm::make_early_inc_range(llvm::drop_begin(*region, 1)))
     if (failed(convertBlockSignature(&block, converter)))
       return failure();
-  return newEntry;
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Rewriter Notification Hooks
 
 void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
                                                      ValueRange newValues) {
   assert(newValues.size() == op->getNumResults());

 }
 
 FailureOr<Block *> ConversionPatternRewriter::convertRegionTypes(
     Region *region, TypeConverter &converter,
     TypeConverter::SignatureConversion *entryConversion) {
   return impl->convertRegionTypes(region, converter, entryConversion);
 }
 
+LogicalResult ConversionPatternRewriter::convertNonEntryRegionTypes(
+    Region *region, TypeConverter &converter) {
+  return impl->convertNonEntryRegionTypes(region, converter);
+}
+
 void ConversionPatternRewriter::replaceUsesOfBlockArgument(BlockArgument from,
                                                            Value to) {
   LLVM_DEBUG({
     Operation *parentOp = from.getOwner()->getParentOp();
     impl->logger.startLine() << "** Replace Argument : '" << from
                              << "'(in region of '" << parentOp->getName()
                              << "'(" << from.getOwner()->getParentOp() << ")\n";
   });

mlir/test/Dialect/Linalg/detensorized_while.mlir

+// RUN: mlir-opt %s -linalg-detensorize | FileCheck %s
+
+#map0 = affine_map<() -> ()>
+
+#attrs = {
+  indexing_maps = [#map0, #map0, #map0],
+  iterator_types = []
+}
+
+func @main(%farg0: tensor<i32>, %farg1: tensor<i32>) -> tensor<i32> attributes {} {
+  br ^bb1(%farg0 : tensor<i32>)
+
+^bb1(%0: tensor<i32>):  // 2 preds: ^bb0, ^bb2
+  %1 = linalg.init_tensor [] : tensor<i1>
+  %2 = linalg.generic #attrs
+    ins(%0, %farg1 : 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 #attrs
+    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
+// CHECK-SAME:    (%{{.*}}: tensor<i32>, %{{.*}}: tensor<i32>)
+// CHECK:         tensor.extract {{.*}}
+// CHECK:         br ^[[bb1:.*]](%{{.*}} : i32)
+// CHECK:       ^[[bb1]](%{{.*}}: i32)
+// CHECK:         cmpi slt, {{.*}}
+// CHECK:         cond_br {{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^[[bb3:.*]](%{{.*}} : i32)
+// CHECK:       ^[[bb2]](%{{.*}}: i32)
+// CHECK:         addi {{.*}}
+// CHECK:         br ^[[bb1]](%{{.*}} : i32)
+// CHECK:       ^[[bb3]](%{{.*}}: i32)
+// CHECK:         tensor.from_elements {{.*}}
+// CHECK:         linalg.tensor_reshape {{.*}}
+// CHECK:         return %{{.*}} : tensor<i32>