[SROA] Single-element vectors address space convertion fix

[mshelego]

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.

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Maksim Shelegov (mshelego)

Changes

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 --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: {{.*}}

[mshelego]

The issue seems to be already fixed in ea90466