From 4fab22e16e90c81976b689a155837fcb6f6938ac Mon Sep 17 00:00:00 2001
From: Diego Caballero <dcaballero@nvidia.com>
Date: Thu, 29 May 2025 03:47:30 +0000
Subject: [PATCH] [mlir][LLVMIR] Add folders for `llvm.inttoptr` and
 `llvm.ptrtoint`

This PR adds folders for `inttoptr(ptrtoint(x))` and `ptrtoint(inttoptr(x))`.
---
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td |  8 ++++--
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp  | 29 +++++++++++++++++++
 mlir/test/Dialect/LLVMIR/canonicalize.mlir  | 32 +++++++++++++++++++--
 3 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 6fde45ce5c556..3eb48964bd378 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -622,10 +622,14 @@ def LLVM_AddrSpaceCastOp : LLVM_CastOp<"addrspacecast", "AddrSpaceCast",
 }
 def LLVM_IntToPtrOp : LLVM_DereferenceableCastOp<"inttoptr", "IntToPtr",
                                   LLVM_ScalarOrVectorOf<AnySignlessInteger>,
-                                  LLVM_ScalarOrVectorOf<LLVM_AnyPointer>>;
+                                  LLVM_ScalarOrVectorOf<LLVM_AnyPointer>> {
+  let hasFolder = 1;
+}
 def LLVM_PtrToIntOp : LLVM_CastOp<"ptrtoint", "PtrToInt",
                                   LLVM_ScalarOrVectorOf<LLVM_AnyPointer>,
-                                  LLVM_ScalarOrVectorOf<AnySignlessInteger>>;
+                                  LLVM_ScalarOrVectorOf<AnySignlessInteger>> {
+  let hasFolder = 1;
+}
 def LLVM_SExtOp : LLVM_CastOp<"sext", "SExt",
                               LLVM_ScalarOrVectorOf<AnySignlessInteger>,
                               LLVM_ScalarOrVectorOf<AnySignlessInteger>> {
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index c7528c970a4ba..74745752fd3dd 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -3562,6 +3562,35 @@ LogicalResult LLVM::BitcastOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// IntToPtrOp
+//===----------------------------------------------------------------------===//
+
+template <typename CastOp, typename ReverseCastOp>
+static OpFoldResult foldReversibleCast(CastOp castOp) {
+  auto reverseCastOp = castOp.getArg().template getDefiningOp<ReverseCastOp>();
+  if (!reverseCastOp)
+    return {};
+
+  // cast(reverse_cast(x)) -> x
+  if (reverseCastOp.getArg().getType() == castOp.getType())
+    return reverseCastOp.getArg();
+
+  return {};
+}
+
+OpFoldResult LLVM::IntToPtrOp::fold(FoldAdaptor adaptor) {
+  return foldReversibleCast<LLVM::IntToPtrOp, LLVM::PtrToIntOp>(*this);
+}
+
+//===----------------------------------------------------------------------===//
+// PtrToIntOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult LLVM::PtrToIntOp::fold(FoldAdaptor adaptor) {
+  return foldReversibleCast<LLVM::PtrToIntOp, LLVM::IntToPtrOp>(*this);
+}
+
 //===----------------------------------------------------------------------===//
 // Folder for LLVM::AddrSpaceCastOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
index 8accf6e263863..51f583868aee2 100644
--- a/mlir/test/Dialect/LLVMIR/canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
@@ -103,10 +103,10 @@ llvm.func @fold_unrelated_extractvalue(%arr: !llvm.array<4 x f32>) -> f32 {
 // -----
 // CHECK-LABEL: fold_extract_extractvalue
 llvm.func @fold_extract_extractvalue(%arr: !llvm.struct<(i64, array<1 x ptr<1>>)>) -> !llvm.ptr<1> {
-  // CHECK: llvm.extractvalue %{{.*}}[1, 0] 
+  // CHECK: llvm.extractvalue %{{.*}}[1, 0]
   // CHECK-NOT: extractvalue
-  %a = llvm.extractvalue %arr[1] : !llvm.struct<(i64, array<1 x ptr<1>>)> 
-  %b = llvm.extractvalue %a[0] : !llvm.array<1 x ptr<1>> 
+  %a = llvm.extractvalue %arr[1] : !llvm.struct<(i64, array<1 x ptr<1>>)>
+  %b = llvm.extractvalue %a[0] : !llvm.array<1 x ptr<1>>
   llvm.return %b : !llvm.ptr<1>
 }
 
@@ -134,6 +134,32 @@ llvm.func @fold_extract_splat() -> f64 {
 
 // -----
 
+// CHECK-LABEL: fold_inttoptr_ptrtoint
+//  CHECK-SAME: %[[ARG:.+]]: i32) -> i32
+//   CHECK-NOT: inttoptr
+//   CHECK-NOT: ptrtoint
+//  CHECK-NEXT: llvm.return %[[ARG]]
+llvm.func @fold_inttoptr_ptrtoint(%x : i32) -> i32 {
+  %c = llvm.inttoptr %x : i32 to !llvm.ptr
+  %d = llvm.ptrtoint %c : !llvm.ptr to i32
+  llvm.return %d : i32
+}
+
+// -----
+
+// CHECK-LABEL: fold_ptrtoint_inttoptr
+//  CHECK-SAME: %[[ARG:.+]]: !llvm.ptr<3>) -> !llvm.ptr<3>
+//   CHECK-NOT: inttoptr
+//   CHECK-NOT: ptrtoint
+//  CHECK-NEXT: llvm.return %[[ARG]]
+llvm.func @fold_ptrtoint_inttoptr(%x : !llvm.ptr<3>) -> !llvm.ptr<3> {
+  %c = llvm.ptrtoint %x : !llvm.ptr<3> to i32
+  %d = llvm.inttoptr %c : i32 to !llvm.ptr<3>
+  llvm.return %d : !llvm.ptr<3>
+}
+
+// -----
+
 // CHECK-LABEL: fold_bitcast
 // CHECK-SAME: %[[ARG:[[:alnum:]]+]]
 // CHECK-NEXT: llvm.return %[[ARG]]