[llvm] [SROA] Single-element vectors address space convertion fix (PR #143027)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Maksim Shelegov (mshelego)

<details>
<summary>Changes</summary>

Fixes #<!-- -->139718

When a pointer value is converted into a different address space its
original or converted type can be a vector with a single element.
In such cases canConvertValue() returns true because of the types equal
bitsize, but then a crash in convertValue() is obversed.

---
Full diff: https://github.com/llvm/llvm-project/pull/143027.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/SROA.cpp (+21-3) 
- (added) llvm/test/Transforms/SROA/sev-convert-addrspace.ll (+34) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index a4e373d395b90..607c485dccf2d 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -2028,11 +2028,29 @@ static Value *convertValue(const DataLayout &DL, IRBuilderTy &IRB, Value *V,
     // cannot use `bitcast` (which has restrict on the same address space) or
     // `addrspacecast` (which is not always no-op casting). Instead, use a pair
     // of no-op `ptrtoint`/`inttoptr` casts through an integer with the same bit
-    // size.
+    // size. If the original value is a single-element vector it must be bitcast
+    // to a scalar before address space conversion. If the value must be
+    // converted to a single-element vector, a bitcast must be added after
+    // address space conversion.
     if (OldAS != NewAS) {
       assert(DL.getPointerSize(OldAS) == DL.getPointerSize(NewAS));
-      return IRB.CreateIntToPtr(IRB.CreatePtrToInt(V, DL.getIntPtrType(OldTy)),
-                                NewTy);
+      auto ConvertPtrAddrSpace = [&](Value *V, Type *NewTy) {
+        return IRB.CreateIntToPtr(
+            IRB.CreatePtrToInt(V, DL.getIntPtrType(V->getType())), NewTy);
+      };
+      if (OldTy->isVectorTy() && !NewTy->isVectorTy()) {
+        auto *VecTy = cast<FixedVectorType>(OldTy);
+        assert(VecTy->getNumElements() == 1);
+        return ConvertPtrAddrSpace(
+            IRB.CreateBitCast(V, VecTy->getElementType()), NewTy);
+      } else if (!OldTy->isVectorTy() && NewTy->isVectorTy()) {
+        auto *VecTy = cast<FixedVectorType>(NewTy);
+        assert(VecTy->getNumElements() == 1);
+        return IRB.CreateBitCast(
+            ConvertPtrAddrSpace(V, VecTy->getElementType()), NewTy);
+      } else {
+        return ConvertPtrAddrSpace(V, NewTy);
+      }
     }
   }
 
diff --git a/llvm/test/Transforms/SROA/sev-convert-addrspace.ll b/llvm/test/Transforms/SROA/sev-convert-addrspace.ll
new file mode 100644
index 0000000000000..350878a9e08ef
--- /dev/null
+++ b/llvm/test/Transforms/SROA/sev-convert-addrspace.ll
@@ -0,0 +1,34 @@
+; RUN: opt < %s -passes='sroa<preserve-cfg>' -S | FileCheck %s --check-prefixes=CHECK,CHECK-PRESERVE-CFG
+; RUN: opt < %s -passes='sroa<modify-cfg>' -S | FileCheck %s --check-prefixes=CHECK,CHECK-MODIFY-CFG
+
+define ptr addrspace(1) @sev_convert_addrspace_in(<1 x ptr> %arg) {
+; CHECK-LABEL: @sev_convert_addrspace_in(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[BITCAST:%.*]] = bitcast <1 x ptr> %arg to ptr
+; CHECK-NEXT:    [[PTRTOINT:%.*]] = ptrtoint ptr [[BITCAST]] to i64
+; CHECK-NEXT:    [[INTTOPTR:%.*]] = inttoptr i64 [[PTRTOINT]] to ptr addrspace(1)
+; CHECK-NEXT:    ret ptr addrspace(1) [[INTTOPTR]]
+entry:
+   %tmp = alloca <1 x ptr>
+   store <1 x ptr> %arg, ptr %tmp
+   %res = load ptr addrspace(1), ptr %tmp
+   ret ptr addrspace(1) %res
+ }
+
+define <1 x ptr> @sev_convert_addrspace_out(ptr addrspace(1) %arg) {
+; CHECK-LABEL: @sev_convert_addrspace_out(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[PTRTOINT:%.*]] = ptrtoint ptr addrspace(1) %arg to i64
+; CHECK-NEXT:    [[INTTOPTR:%.*]] = inttoptr i64 [[PTRTOINT]] to ptr
+; CHECK-NEXT:    [[BITCAST:%.*]] = bitcast ptr [[INTTOPTR]] to <1 x ptr>
+; CHECK-NEXT:    ret <1 x ptr> [[BITCAST]]
+entry:
+  %tmp = alloca ptr addrspace(1)
+  store ptr addrspace(1) %arg, ptr %tmp
+  %res = load <1 x ptr>, ptr %tmp
+  ret <1 x ptr> %res
+ }
+
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; CHECK-MODIFY-CFG: {{.*}}
+; CHECK-PRESERVE-CFG: {{.*}}

``````````

</details>


https://github.com/llvm/llvm-project/pull/143027