[llvm-branch-commits] [mlir] [mlir][Transforms] Dialect conversion: Improve signature conversion API (PR #81997)

[Matthias Springer via llvm-branch-commits]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/81997

This commit improves the block signature conversion API of the dialect conversion.

There is the following comment in `ArgConverter::applySignatureConversion`:
```
// If no arguments are being changed or added, there is nothing to do.
```

However, the implementation actually used to replace a block with a new block even if the block argument types do not change (i.e., there is "nothing to do"). This is fixed in this commit. The documentation of the public `ConversionPatternRewriter` API is updated accordingly.

This commit also removes a check that used to *sometimes* skip a block signature conversion if the block was already converted. This is not consistent with the public `ConversionPatternRewriter` API; block should always be converted, regardless of whether they were already converted or not.

Block signature conversion also used to be silently skipped when the specified block was detached. Instead of silently skipping, an assertion is triggered. Attempting to convert a detached block (which is likely an erased block) is invalid API usage.

>From 5dc79f6af9ac00a61767062980b13eb4ae8d2571 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Fri, 16 Feb 2024 15:13:37 +0000
Subject: [PATCH] [mlir][Transforms] Dialect conversion: Improve signature
 conversion API

This commit improves the block signature conversion API of the dialect conversion.

There is the following comment in `ArgConverter::applySignatureConversion`:
```
// If no arguments are being changed or added, there is nothing to do.
```

However, the implementation actually used to replace a block with a new block even if the block argument types do not change (i.e., there is "nothing to do"). This is fixed in this commit. The documentation of the public `ConversionPatternRewriter` API is updated accordingly.

This commit also removes a check that used to *sometimes* skip a block signature conversion if the block was already converted. This is not consistent with the public `ConversionPatternRewriter` API; block should always be converted, regardless of whether they were already converted or not.

Block signature conversion also used to be silently skipped when the specified block was detached. Instead of silently skipping, an assertion is triggered. Attempting to convert a detached block (which is likely an erased block) is invalid API usage.
---
 mlir/include/mlir/Transforms/DialectConversion.h | 12 +++++++++---
 mlir/lib/Transforms/Utils/DialectConversion.cpp  | 10 +++-------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Transforms/DialectConversion.h b/mlir/include/mlir/Transforms/DialectConversion.h
index 0d7722aa07ee38..2575be4cdea1ac 100644
--- a/mlir/include/mlir/Transforms/DialectConversion.h
+++ b/mlir/include/mlir/Transforms/DialectConversion.h
@@ -663,6 +663,8 @@ class ConversionPatternRewriter final : public PatternRewriter {
   /// Apply a signature conversion to the entry block of the given region. This
   /// replaces the entry block with a new block containing the updated
   /// signature. The new entry block to the region is returned for convenience.
+  /// If no block argument types are changing, the entry original block will be
+  /// left in place and returned.
   ///
   /// If provided, `converter` will be used for any materializations.
   Block *
@@ -671,8 +673,11 @@ class ConversionPatternRewriter final : public PatternRewriter {
                            const TypeConverter *converter = nullptr);
 
   /// Convert the types of block arguments within the given region. This
-  /// replaces each block with a new block containing the updated signature. The
-  /// entry block may have a special conversion if `entryConversion` is
+  /// replaces each block with a new block containing the updated signature. If
+  /// an updated signature would match the current signature, the respective
+  /// block is left in place as is.
+  ///
+  /// 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(
@@ -681,7 +686,8 @@ class ConversionPatternRewriter final : public PatternRewriter {
 
   /// 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.
+  /// containing the updated signature. If an updated signature would match the
+  /// current signature, the respective block is left in place as is.
   ///
   /// If special conversion behavior is needed for the non-entry blocks (for
   /// example, we need to convert only a subset of a BB arguments), such
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 35028001a03dd9..c16bb144efecf5 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -544,12 +544,8 @@ FailureOr<Block *> ArgConverter::convertSignature(
     Block *block, const TypeConverter *converter,
     ConversionValueMapping &mapping,
     SmallVectorImpl<BlockArgument> &argReplacements) {
-  // Check if the block was already converted.
-  // * If the block is mapped in `conversionInfo`, it is a converted block.
-  // * If the block is detached, conservatively assume that it is going to be
-  //   deleted; it is likely the old block (before it was converted).
-  if (conversionInfo.count(block) || !block->getParent())
-    return block;
+  assert(block->getParent() && "cannot convert signature of detached block");
+
   // If a converter wasn't provided, and the block wasn't already converted,
   // there is nothing we can do.
   if (!converter)
@@ -570,7 +566,7 @@ Block *ArgConverter::applySignatureConversion(
   // If no arguments are being changed or added, there is nothing to do.
   unsigned origArgCount = block->getNumArguments();
   auto convertedTypes = signatureConversion.getConvertedTypes();
-  if (origArgCount == 0 && convertedTypes.empty())
+  if (llvm::equal(block->getArgumentTypes(), convertedTypes))
     return block;
 
   // Split the block at the beginning to get a new block to use for the updated