[CIR][HIP] Proper Handling of address spaces in ptr-diff (#1994)

koparasy · andykaylor · web-flow · commit 630a11b5754b · 2025-11-25T10:29:33.000-08:00
Co-authored-by: Andy Kaylor &lt;akaylor@nvidia.com&gt;
diff --git a/clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp b/clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp
@@ -1524,17 +1524,59 @@ mlir::Value ScalarExprEmitter::emitSub(const BinOpInfo &Ops) {
   if (!mlir::isa<cir::PointerType>(Ops.RHS.getType()))
     return emitPointerArithmetic(CGF, Ops, /*isSubtraction=*/true);
 
-  // Otherwise, this is a pointer subtraction
-
   // Do the raw subtraction part.
-  //
-  // TODO(cir): note for LLVM lowering out of this; when expanding this into
-  // LLVM we shall take VLA's, division by element size, etc.
-  //
-  // See more in `EmitSub` in CGExprScalar.cpp.
+  mlir::Value lhs = Ops.LHS;
+  mlir::Value rhs = Ops.RHS;
+
+  cir::PointerType lhsPtrTy = mlir::dyn_cast<cir::PointerType>(lhs.getType());
+  cir::PointerType rhsPtrTy = mlir::dyn_cast<cir::PointerType>(rhs.getType());
+
+  if (lhsPtrTy && rhsPtrTy) {
+    cir::AddressSpace lhsAS = lhsPtrTy.getAddrSpace();
+    cir::AddressSpace rhsAS = rhsPtrTy.getAddrSpace();
+
+    if (lhsAS != rhsAS) {
+      // Different address spaces → use addrspacecast
+      rhs = Builder.createAddrSpaceCast(rhs, lhsPtrTy);
+    } else if (lhsPtrTy != rhsPtrTy) {
+      // Same addrspace but different pointee/type → bitcast is fine
+      rhs = Builder.createBitcast(rhs, lhsPtrTy);
+    }
+  }
+
   assert(!cir::MissingFeatures::llvmLoweringPtrDiffConsidersPointee());
-  return cir::PtrDiffOp::create(Builder, CGF.getLoc(Ops.Loc), CGF.PtrDiffTy,
-                                Ops.LHS, Ops.RHS);
+  mlir::Value diff = cir::PtrDiffOp::create(Builder, CGF.getLoc(Ops.Loc),
+                                            CGF.PtrDiffTy, lhs, rhs);
+
+  const BinaryOperator *expr = cast<BinaryOperator>(Ops.E);
+  QualType elementType = expr->getLHS()->getType()->getPointeeType();
+
+  mlir::Location loc = CGF.getLoc(Ops.Loc);
+  mlir::Value divisor;
+
+  // Check if this is a VLA pointee type.
+  if (const auto *vla = CGF.getContext().getAsVariableArrayType(elementType)) {
+    auto vlaSize = CGF.getVLASize(vla);
+    elementType = vlaSize.Type;
+    divisor = vlaSize.NumElts;
+
+    CharUnits eltSize = CGF.getContext().getTypeSizeInChars(elementType);
+    if (!eltSize.isOne()) {
+      cir::IntType cirIntTy = llvm::cast<cir::IntType>(CGF.PtrDiffTy);
+      cir::IntAttr eltSizeAttr =
+          cir::IntAttr::get(cirIntTy, eltSize.getQuantity());
+
+      if (divisor.getType() != CGF.PtrDiffTy)
+        divisor = Builder.createIntCast(divisor, CGF.PtrDiffTy);
+    }
+  } else {
+    // cir::ptrdiff correctly computes the ABI difference of 2 pointers. We
+    // do not need to compute anything else here. We just return it.
+    return diff;
+  }
+
+  return cir::BinOp::create(Builder, loc, CGF.PtrDiffTy, cir::BinOpKind::Div,
+                            diff, divisor);
 }
 
 // Helper to apply OpenCL-style shift masking. It handles both vector and scalar
diff --git a/clang/test/CIR/CodeGen/HIP/ptr-diff.cpp b/clang/test/CIR/CodeGen/HIP/ptr-diff.cpp
@@ -0,0 +1,60 @@
+#include "cuda.h"
+
+// RUN: %clang_cc1 -triple=amdgcn-amd-amdhsa -x hip -fclangir \
+// RUN:            -fcuda-is-device -fhip-new-launch-api \
+// RUN:            -I%S/../Inputs/ -emit-cir %s -o %t.ll
+// RUN: FileCheck --check-prefix=CIR-DEVICE --input-file=%t.ll %s
+
+// RUN: %clang_cc1 -triple=amdgcn-amd-amdhsa -x hip -fclangir \
+// RUN:            -fcuda-is-device -fhip-new-launch-api \
+// RUN:            -I%S/../Inputs/ -emit-llvm %s -o %t.ll
+// RUN: FileCheck --check-prefix=LLVM-DEVICE --input-file=%t.ll %s
+
+// RUN: %clang_cc1 -triple=amdgcn-amd-amdhsa -x hip  \
+// RUN:            -fcuda-is-device -fhip-new-launch-api \
+// RUN:            -I%S/../Inputs/ -emit-llvm %s -o %t.ll
+// RUN: FileCheck --check-prefix=OGCG-DEVICE --input-file=%t.ll %s
+
+__device__ int ptr_diff() {
+  const char c_str[] = "c-string"; 
+  const char* len =  c_str;  
+  return c_str - len;
+}
+
+
+// CIR-DEVICE: %[[#LenLocalAddr:]] = cir.alloca !cir.ptr<!s8i>, !cir.ptr<!cir.ptr<!s8i>>, ["len", init]
+// CIR-DEVICE: %[[#GlobalPtr:]] = cir.get_global @_ZZ8ptr_diffvE5c_str : !cir.ptr<!cir.array<!s8i x 9>, addrspace(offload_constant)>
+// CIR-DEVICE: %[[#CastDecay:]] = cir.cast array_to_ptrdecay %[[#GlobalPtr]] : !cir.ptr<!cir.array<!s8i x 9>, addrspace(offload_constant)>
+// CIR-DEVICE: %[[#LenLocalAddrCast:]] = cir.cast bitcast %[[#LenLocalAddr]] : !cir.ptr<!cir.ptr<!s8i>> -> !cir.ptr<!cir.ptr<!s8i, addrspace(offload_constant)>>
+// CIR-DEVICE: cir.store align(8) %[[#CastDecay]], %[[#LenLocalAddrCast]] : !cir.ptr<!s8i, addrspace(offload_constant)>, !cir.ptr<!cir.ptr<!s8i, addrspace(offload_constant)>>
+// CIR-DEVICE: %[[#CStr:]] = cir.cast array_to_ptrdecay %[[#GlobalPtr]] : !cir.ptr<!cir.array<!s8i x 9>, addrspace(offload_constant)> -> !cir.ptr<!s8i, addrspace(offload_constant)>
+// CIR-DEVICE: %[[#LoadedLenAddr:]] = cir.load align(8) %[[#LenLocalAddr]] : !cir.ptr<!cir.ptr<!s8i>>, !cir.ptr<!s8i> loc(#loc7)
+// CIR-DEVICE: %[[#AddrCast:]] = cir.cast address_space %[[#LoadedLenAddr]] : !cir.ptr<!s8i> -> !cir.ptr<!s8i, addrspace(offload_constant)>
+// CIR-DEVICE: %[[#DIFF:]] = cir.ptr_diff %[[#CStr]], %[[#AddrCast]] : !cir.ptr<!s8i, addrspace(offload_constant)>
+
+// LLVM-DEVICE: define dso_local i32 @_Z8ptr_diffv()
+// LLVM-DEVICE: %[[#GlobalPtrAddr:]] = alloca i32, i64 1, align 4, addrspace(5)
+// LLVM-DEVICE: %[[#GlobalPtrCast:]] = addrspacecast ptr addrspace(5) %[[#GlobalPtrAddr]] to ptr
+// LLVM-DEVICE: %[[#LenLocalAddr:]] = alloca ptr, i64 1, align 8, addrspace(5)
+// LLVM-DEVICE: %[[#LenLocalAddrCast:]] = addrspacecast ptr addrspace(5) %[[#LenLocalAddr]] to ptr
+// LLVM-DEVICE: store ptr addrspace(4) @_ZZ8ptr_diffvE5c_str, ptr %[[#LenLocalAddrCast]], align 8
+// LLVM-DEVICE: %[[#LoadedAddr:]] = load ptr, ptr %[[#LenLocalAddrCast]], align 8
+// LLVM-DEVICE: %[[#CastedVal:]] = addrspacecast ptr %[[#LoadedAddr]] to ptr addrspace(4)
+// LLVM-DEVICE: %[[#IntVal:]] = ptrtoint ptr addrspace(4) %[[#CastedVal]] to i64
+// LLVM-DEVICE: %[[#SubVal:]] = sub i64 ptrtoint (ptr addrspace(4) @_ZZ8ptr_diffvE5c_str to i64), %[[#IntVal]]
+
+// OGCG-DEVICE: define dso_local noundef i32 @_Z8ptr_diffv() #0
+// OGCG-DEVICE: %[[RETVAL:.*]] = alloca i32, align 4, addrspace(5)
+// OGCG-DEVICE: %[[C_STR:.*]] = alloca [9 x i8], align 1, addrspace(5)
+// OGCG-DEVICE: %[[LEN:.*]] = alloca ptr, align 8, addrspace(5)
+// OGCG-DEVICE: %[[RETVAL_ASCAST:.*]] = addrspacecast ptr addrspace(5) %[[RETVAL]] to ptr
+// OGCG-DEVICE: %[[C_STR_ASCAST:.*]] = addrspacecast ptr addrspace(5) %[[C_STR]] to ptr
+// OGCG-DEVICE: %[[LEN_ASCAST:.*]] = addrspacecast ptr addrspace(5) %[[LEN]] to ptr
+// OGCG-DEVICE: %[[ARRAYDECAY:.*]] = getelementptr inbounds [9 x i8], ptr %[[C_STR_ASCAST]], i64 0, i64 0
+// OGCG-DEVICE: store ptr %[[ARRAYDECAY]], ptr %[[LEN_ASCAST]], align 8
+// OGCG-DEVICE: %[[ARRAYDECAY1:.*]] = getelementptr inbounds [9 x i8], ptr %[[C_STR_ASCAST]], i64 0, i64 0
+// OGCG-DEVICE: %[[LOADED:.*]] = load ptr, ptr %[[LEN_ASCAST]], align 8
+// OGCG-DEVICE: %[[LHS:.*]] = ptrtoint ptr %[[ARRAYDECAY1]] to i64
+// OGCG-DEVICE: %[[RHS:.*]] = ptrtoint ptr %[[LOADED]] to i64
+// OGCG-DEVICE: %[[SUB:.*]] = sub i64 %[[LHS]], %[[RHS]]
+// OGCG-DEVICE: %[[CONV:.*]] = trunc i64 %[[SUB]] to i32
diff --git a/clang/test/CIR/CodeGen/vla.c b/clang/test/CIR/CodeGen/vla.c
@@ -1,4 +1,11 @@
-// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fclangir -emit-cir %s -o -  | FileCheck %s
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fclangir -emit-cir %s -o %t.cir
+// RUN: FileCheck --check-prefix=CHECK --input-file=%t.cir %s
+
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fclangir -emit-llvm %s -o %t.ll
+// RUN: FileCheck --check-prefix=LLVM --input-file=%t.ll %s
+
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -emit-llvm %s -o %t.ll
+// RUN: FileCheck --check-prefix=OGCG --input-file=%t.ll %s
 
 // CHECK:  cir.func dso_local @f0(%arg0: !s32i
 // CHECK:    [[TMP0:%.*]] = cir.alloca !s32i, !cir.ptr<!s32i>, ["len", init] {alignment = 4 : i64}
@@ -93,3 +100,92 @@ int f9(unsigned n, char (*p)[n][n+1][6]) {
 
   return p2 - p;
 }
+
+long f10(int n) {
+    int (*p)[n];
+    int (*q)[n];
+    return q - p;
+}
+// CHECK-LABEL: cir.func dso_local @f10
+// CHECK: %[[Q_VAL:[0-9]+]] = cir.load {{.*}} %{{.*}} : !cir.ptr<!cir.ptr<!s32i>>, !cir.ptr<!s32i>
+// CHECK: %[[P_VAL:[0-9]+]] = cir.load {{.*}} %{{.*}} : !cir.ptr<!cir.ptr<!s32i>>, !cir.ptr<!s32i>
+// CHECK: %[[PTRDIFF:[0-9]+]] = cir.ptr_diff %[[Q_VAL]], %[[P_VAL]] : !cir.ptr<!s32i> -> !s64i
+// CHECK: %[[N_S64:[0-9]+]] = cir.cast integral %{{.*}} : !u64i -> !s64i
+// CHECK: %[[DIV:[0-9]+]] = cir.binop(div, %[[PTRDIFF]], %[[N_S64]]) : !s64i
+
+// LLVM-LABEL: @f10(
+// LLVM: %[[QI:[0-9]+]] = ptrtoint ptr %{{.*}} to i64
+// LLVM: %[[PI:[0-9]+]] = ptrtoint ptr %{{.*}} to i64
+// LLVM: %[[DIFF_BYTES:[0-9]+]] = sub i64 %[[QI]], %[[PI]]
+// LLVM: %[[PTRDIFF_INTS:[0-9]+]] = sdiv i64 %[[DIFF_BYTES]], 4
+// LLVM: %[[RESULT:[0-9]+]] = sdiv i64 %[[PTRDIFF_INTS]], %{{.*}}
+
+// OGCG-LABEL: @f10(
+// OGCG: %{{.*}} = ptrtoint ptr %{{.*}} to i64
+// OGCG: %{{.*}} = ptrtoint ptr %{{.*}} to i64
+// OGCG: %{{.*}} = sub i64 %{{.*}}, %{{.*}}
+// OGCG: %{{.*}} = mul nuw i64 4, %{{.*}}
+// OGCG: %{{.*}} = sdiv exact i64 %{{.*}}, %{{.*}}
+
+long f11(int n, int m) {
+    int (*p)[n][m];
+    int (*q)[n][m];
+    return q - p;
+}
+// CHECK-LABEL: cir.func dso_local @f11
+
+// # allocas
+// CHECK: %[[N_ADDR:[0-9]+]] = cir.alloca !s32i, !cir.ptr<!s32i>
+// CHECK: %[[M_ADDR:[0-9]+]] = cir.alloca !s32i, !cir.ptr<!s32i>
+// CHECK: %[[RET:[0-9]+]] = cir.alloca !s64i, !cir.ptr<!s64i>
+// CHECK: %[[P:[0-9]+]] = cir.alloca !cir.ptr<!s32i>, !cir.ptr<!cir.ptr<!s32i>>
+// CHECK: %[[Q:[0-9]+]] = cir.alloca !cir.ptr<!s32i>, !cir.ptr<!cir.ptr<!s32i>>
+
+// # store n, m
+// CHECK: cir.store %arg0, %[[N_ADDR]] : !s32i, !cir.ptr<!s32i>
+// CHECK: cir.store %arg1, %[[M_ADDR]] : !s32i, !cir.ptr<!s32i>
+
+// # load n and cast to u64
+// CHECK: %[[N_LOAD:[0-9]+]] = cir.load {{.*}} %[[N_ADDR]] : !cir.ptr<!s32i>, !s32i
+// CHECK: %[[N_U64:[0-9]+]] = cir.cast integral %[[N_LOAD]] : !s32i -> !u64i
+
+// # load m and cast to u64
+// CHECK: %[[M_LOAD:[0-9]+]] = cir.load {{.*}} %[[M_ADDR]] : !cir.ptr<!s32i>, !s32i
+// CHECK: %[[M_U64:[0-9]+]] = cir.cast integral %[[M_LOAD]] : !s32i -> !u64i
+
+// # load q and p
+// CHECK: %[[Q_VAL:[0-9]+]] = cir.load {{.*}} %[[Q]] : !cir.ptr<!cir.ptr<!s32i>>, !cir.ptr<!s32i>
+// CHECK: %[[P_VAL:[0-9]+]] = cir.load {{.*}} %[[P]] : !cir.ptr<!cir.ptr<!s32i>>, !cir.ptr<!s32i>
+// 
+// # ptrdiff → (byte_diff / 4)
+// CHECK: %[[PTRDIFF:[0-9]+]] = cir.ptr_diff %[[Q_VAL]], %[[P_VAL]] : !cir.ptr<!s32i> -> !s64i
+
+// # compute n*m
+// CHECK: %[[NM_U64:[0-9]+]] = cir.binop(mul, %[[N_U64]], %[[M_U64]]) : !u64i
+// CHECK: %[[NM_S64:[0-9]+]] = cir.cast integral %[[NM_U64]] : !u64i -> !s64i
+
+// # divide ptrdiff_ints by (n*m)
+// CHECK: %[[RESULT:[0-9]+]] = cir.binop(div, %[[PTRDIFF]], %[[NM_S64]]) : !s64i
+
+// # store + return
+// CHECK: cir.store %[[RESULT]], %[[RET]] : !s64i, !cir.ptr<!s64i>
+// CHECK: %[[RETVAL:[0-9]+]] = cir.load %[[RET]] : !cir.ptr<!s64i>, !s64i
+// CHECK: cir.return %[[RETVAL]] : !s64i
+
+
+// LLVM-LABEL: @f11(
+// # load q and p
+// LLVM: %[[QI:[0-9]+]] = ptrtoint ptr %{{.*}} to i64
+// LLVM: %[[PI:[0-9]+]] = ptrtoint ptr %{{.*}} to i64
+// LLVM: %[[DIFF_BYTES:[0-9]+]] = sub i64 %[[QI]], %[[PI]]
+// LLVM: %[[PTRDIFF_INTS:[0-9]+]] = sdiv i64 %[[DIFF_BYTES]], 4
+// LLVM: %[[NM:[0-9]+]] = mul i64 %{{.*}}, %{{.*}}
+// LLVM: %[[RESULT:[0-9]+]] = sdiv i64 %[[PTRDIFF_INTS]], %[[NM]]
+
+// OGCG-LABEL: @f11(
+// OGCG: %{{.*}} = ptrtoint ptr %{{.*}} to i64
+// OGCG: %{{.*}} = ptrtoint ptr %{{.*}} to i64
+// OGCG: %{{.*}} = sub i64 %{{.*}}, %{{.*}}
+// OGCG: %{{.*}} = mul nuw i64 %{{.*}}, %{{.*}}
+// OGCG: %{{.*}} = mul nuw i64 4, %{{.*}}
+// OGCG: %{{.*}} = sdiv exact i64 %{{.*}}, %{{.*}}
