[SROA] Prevent load atomic vector from being generated

llvm/include/llvm/IR/Instructions.h

@@ -250,6 +250,10 @@ class LoadInst : public UnaryInstruction {
            !isVolatile();
   }
 
+  /// Returns false if this type would be invalid in the
+  /// creation of a load atomic instruction.
+  static bool isValidAtomicTy(Type *Ty);
+
   Value *getPointerOperand() { return getOperand(0); }
   const Value *getPointerOperand() const { return getOperand(0); }
   static unsigned getPointerOperandIndex() { return 0U; }

llvm/lib/IR/Instructions.cpp

@@ -1247,6 +1247,12 @@ void LoadInst::AssertOK() {
          "Ptr must have pointer type.");
 }
 
+bool LoadInst::isValidAtomicTy(Type *Ty) {
+  if (!Ty->isIntOrPtrTy() && !Ty->isFloatingPointTy())
+    return false;
+  return true;
+}
+
 static Align computeLoadStoreDefaultAlign(Type *Ty, InsertPosition Pos) {
   assert(Pos.isValid() &&
          "Insertion position cannot be null when alignment not provided!");

llvm/lib/Transforms/Scalar/SROA.cpp

@@ -2853,6 +2853,11 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
 
   bool visitLoadInst(LoadInst &LI) {
     LLVM_DEBUG(dbgs() << "    original: " << LI << "\n");
+
+    // load atomic vector would be generated, which is illegal
+    if (LI.isAtomic() && !LoadInst::isValidAtomicTy(NewAI.getAllocatedType()))
+      return false;
+
     Value *OldOp = LI.getOperand(0);
     assert(OldOp == OldPtr);
 
@@ -2875,6 +2880,7 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
                (canConvertValue(DL, NewAllocaTy, TargetTy) ||
                 (IsLoadPastEnd && NewAllocaTy->isIntegerTy() &&
                  TargetTy->isIntegerTy() && !LI.isVolatile()))) {
+
       Value *NewPtr =
           getPtrToNewAI(LI.getPointerAddressSpace(), LI.isVolatile());
       LoadInst *NewLI = IRB.CreateAlignedLoad(NewAI.getAllocatedType(), NewPtr,

llvm/test/Transforms/SROA/atomic-vector.ll

@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes='sroa' -S | FileCheck %s
+
+define float @atomic_vector() {
+; CHECK-LABEL: define float @atomic_vector() {
+; CHECK-NEXT:    [[TMP1:%.*]] = alloca <1 x float>, align 4
+; CHECK-NEXT:    store <1 x float> undef, ptr [[TMP1]], align 4
+; CHECK-NEXT:    [[TMP2:%.*]] = load atomic volatile float, ptr [[TMP1]] acquire, align 4
+; CHECK-NEXT:    ret float [[TMP2]]
+;
+  %src = alloca <1 x float>
+  %val = alloca <1 x float>
+  %direct = alloca ptr
+  call void @llvm.memcpy.p0.p0.i64(ptr %val, ptr %src, i64 4, i1 false)
+  store ptr %val, ptr %direct
+  %indirect = load ptr, ptr %direct
+  %ret = load atomic volatile float, ptr %indirect acquire, align 4
+  ret float %ret
+}
+
+define i32 @atomic_vector_int() {
+; CHECK-LABEL: define i32 @atomic_vector_int() {
+; CHECK-NEXT:    [[VAL:%.*]] = alloca <1 x i32>, align 4
+; CHECK-NEXT:    store <1 x i32> undef, ptr [[VAL]], align 4
+; CHECK-NEXT:    [[RET:%.*]] = load atomic volatile i32, ptr [[VAL]] acquire, align 4
+; CHECK-NEXT:    ret i32 [[RET]]
+;
+  %src = alloca <1 x i32>
+  %val = alloca <1 x i32>
+  %direct = alloca ptr
+  call void @llvm.memcpy.p0.p0.i64(ptr %val, ptr %src, i64 4, i1 false)
+  store ptr %val, ptr %direct
+  %indirect = load ptr, ptr %direct
+  %ret = load atomic volatile i32, ptr %indirect acquire, align 4
+  ret i32 %ret
+}
+
+define ptr @atomic_vector_ptr() {
+; CHECK-LABEL: define ptr @atomic_vector_ptr() {
+; CHECK-NEXT:    [[VAL_SROA_0:%.*]] = alloca <1 x ptr>, align 8
+; CHECK-NEXT:    store <1 x ptr> undef, ptr [[VAL_SROA_0]], align 8
+; CHECK-NEXT:    [[VAL_SROA_0_0_VAL_SROA_0_0_RET:%.*]] = load atomic volatile ptr, ptr [[VAL_SROA_0]] acquire, align 4
+; CHECK-NEXT:    ret ptr [[VAL_SROA_0_0_VAL_SROA_0_0_RET]]
+;
+  %src = alloca <1 x ptr>
+  %val = alloca <1 x ptr>
+  %direct = alloca ptr
+  call void @llvm.memcpy.p0.p0.i64(ptr %val, ptr %src, i64 8, i1 false)
+  store ptr %val, ptr %direct
+  %indirect = load ptr, ptr %direct
+  %ret = load atomic volatile ptr, ptr %indirect acquire, align 4
+  ret ptr %ret
+}