[Mlir-commits] [mlir] [mlir][IndexToLLVM] Lower UnrealizedConversionCastOp related to IndexType in IndexToLLVM (PR #75255)

[Xiang Li]

https://github.com/python3kgae created https://github.com/llvm/llvm-project/pull/75255


Lower pattern like
```
    %5 = builtin.unrealized_conversion_cast %a : i32 to index
    %6 = builtin.unrealized_conversion_cast %5 : index to i64
```
into
```
    llvm.zext %a : i32 to i64
```
This is for https://github.com/llvm/llvm-project/issues/74939

>From 13ae1a02c8bf534f66e0ba51b54687ebec84dddf Mon Sep 17 00:00:00 2001
From: Xiang Li <python3kgae at outlook.com>
Date: Tue, 12 Dec 2023 18:53:11 -0500
Subject: [PATCH] [mlir][IndexToLLVM] Lower UnrealizedConversionCastOp related
 to IndexType in IndexToLLVM.

Lower pattern like
    %5 = builtin.unrealized_conversion_cast %a : i32 to index
    %6 = builtin.unrealized_conversion_cast %5 : index to i64
into
    llvm.zext %a : i32 to i64

This is for https://github.com/llvm/llvm-project/issues/74939
---
 .../Conversion/IndexToLLVM/IndexToLLVM.cpp    | 42 ++++++++++++++++++-
 .../Conversion/IndexToLLVM/index-to-llvm.mlir | 36 ++++++++++++++++
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp b/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
index 9d8a5d8a0e1c06..db073f542e131e 100644
--- a/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
+++ b/mlir/lib/Conversion/IndexToLLVM/IndexToLLVM.cpp
@@ -285,6 +285,45 @@ using ConvertIndexXor = mlir::OneToOneConvertToLLVMPattern<XOrOp, LLVM::XOrOp>;
 using ConvertIndexBoolConstant =
     mlir::OneToOneConvertToLLVMPattern<BoolConstantOp, LLVM::ConstantOp>;
 
+/// Define lowering patterns for UnrealizedConversionCastOp ops
+struct UnrealizedConversionCastOpLowering
+    : public ConvertOpToLLVMPattern<UnrealizedConversionCastOp> {
+  UnrealizedConversionCastOpLowering(const LLVMTypeConverter &converter)
+      : ConvertOpToLLVMPattern<UnrealizedConversionCastOp>(converter) {}
+
+  LogicalResult
+  matchAndRewrite(UnrealizedConversionCastOp castOp,
+                  typename UnrealizedConversionCastOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type srcType = castOp.getOperand(0).getType();
+    Type dstType = castOp.getType(0);
+    // Skip cast not related to IndexType.
+    if (!isa<IndexType>(srcType) && !isa<IndexType>(dstType))
+      return failure();
+
+    srcType = getTypeConverter()->convertType(srcType);
+    IntegerType srcIntType = dyn_cast<IntegerType>(srcType);
+    if (!srcIntType)
+      return failure();
+    dstType = getTypeConverter()->convertType(dstType);
+    IntegerType dstIntType = dyn_cast<IntegerType>(dstType);
+    if (!dstIntType)
+      return failure();
+
+    Value src = adaptor.getInputs()[0];
+    uint32_t srcWidth = srcIntType.getWidth();
+    uint32_t dstWidth = dstIntType.getWidth();
+    if (srcWidth == dstWidth)
+      rewriter.replaceOp(castOp, src);
+    else if (srcWidth > dstWidth)
+      rewriter.replaceOpWithNewOp<LLVM::TruncOp>(castOp, dstType, src);
+    else
+      rewriter.replaceOpWithNewOp<LLVM::ZExtOp>(castOp, dstType, src);
+
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
@@ -320,7 +359,8 @@ void index::populateIndexToLLVMConversionPatterns(
       ConvertIndexCmp,
       ConvertIndexSizeOf,
       ConvertIndexConstant,
-      ConvertIndexBoolConstant
+      ConvertIndexBoolConstant,
+      UnrealizedConversionCastOpLowering
       // clang-format on
       >(typeConverter);
 }
diff --git a/mlir/test/Conversion/IndexToLLVM/index-to-llvm.mlir b/mlir/test/Conversion/IndexToLLVM/index-to-llvm.mlir
index 1b13ebb38dc9e9..8930265a6a1020 100644
--- a/mlir/test/Conversion/IndexToLLVM/index-to-llvm.mlir
+++ b/mlir/test/Conversion/IndexToLLVM/index-to-llvm.mlir
@@ -194,3 +194,39 @@ func.func @index_constant() {
   %3 = index.constant 3000000000
   return
 }
+
+// CHECK-LABEL: @unrealized_conversion_cast_i32_to_i64
+// CHECK-SAME: %[[A:.*]]: i32
+// INDEX32-LABEL: @unrealized_conversion_cast_i32_to_i64
+// INDEX32-SAME: %[[A:.*]]: i32
+// INDEX64-LABEL: @unrealized_conversion_cast_i32_to_i64
+// INDEX64-SAME: %[[A:.*]]: i32
+func.func @unrealized_conversion_cast_i32_to_i64(%a: i32) -> i64 {
+  // CHECK-NEXT: %[[RET:.*]] = llvm.zext %[[A]] : i32 to i64
+  // CHECK-NEXT: return %[[RET]] : i64
+  // INDEX32-NEXT: %[[RET:.*]] = llvm.zext %[[A]] : i32 to i64
+  // INDEX32-NEXT: return %[[RET]] : i64
+  // INDEX64-NEXT: %[[RET:.*]] = llvm.zext %[[A]] : i32 to i64
+  // INDEX64-NEXT: return %[[RET]] : i64
+  %0 = builtin.unrealized_conversion_cast %a : i32 to index
+  %1 = builtin.unrealized_conversion_cast %0 : index to i64
+  return %1 : i64
+}
+
+// CHECK-LABEL: @unrealized_conversion_cast_i64_to_i32
+// CHECK-SAME: %[[A:.*]]: i64
+// INDEX32-LABEL: @unrealized_conversion_cast_i64_to_i32
+// INDEX32-SAME: %[[A:.*]]: i64
+// INDEX64-LABEL: @unrealized_conversion_cast_i64_to_i32
+// INDEX64-SAME: %[[A:.*]]: i64
+func.func @unrealized_conversion_cast_i64_to_i32(%a: i64) -> i32 {
+  // CHECK-NEXT: %[[RET:.*]] = llvm.trunc %[[A]] : i64 to i32
+  // CHECK-NEXT: return %[[RET]] : i32
+  // INDEX32-NEXT: %[[RET:.*]] = llvm.trunc %[[A]] : i64 to i32
+  // INDEX32-NEXT: return %[[RET]] : i32
+  // INDEX64-NEXT: %[[RET:.*]] = llvm.trunc %[[A]] : i64 to i32
+  // INDEX64-NEXT: return %[[RET]] : i32
+  %0 = builtin.unrealized_conversion_cast %a : i64 to index
+  %1 = builtin.unrealized_conversion_cast %0 : index to i32
+  return %1 : i32
+}
\ No newline at end of file