[AMDGPU] Handle memset users in PromoteAlloca

[Pierre-vh]

Summary

Allows allocas with memset users to be promoted.

This is intended to prevent patterns such as memset(&alloca, 0, sizeof(alloca)) (which I think can be emitted by frontends) from preventing a vectorization of allocas.

Fixes SWDEV-388784

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

397

You can hide the intrinsic checks by dyn_casting to MemSetInst

405–407

I don't see why you need this restriction but if you want to lift it in a separate patch that's fine

558–561

I'd expect this to be up with the MemTransferInst handling, dyn_casting to MemSetInst

llvm/test/CodeGen/AMDGPU/promote-alloca-memset.ll

32

Not checks are fragile,, might as well generate the checks here

62

.double?

[arsenm]

llvm/test/CodeGen/AMDGPU/promote-alloca-memset.ll

2

Probably should also run SROA to make sure the end result is correct

[Pierre-vh]

Comments

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

392

I'd expect the input to be the MemSetInst and the dyn_cast would be on the caller side. Also, use TypeSize?

398

Move the size query here?

502

getTypeStoreSize?

637

store size

[jdoerfert]

Why is this not part and done by sroa?

[Pierre-vh]

Address comments

[Pierre-vh]

In D146225#4199423, @jdoerfert wrote:

Why is this not part and done by sroa?

Do you mean the PromoteAlloca pass? I'm not sure

[arsenm]

In D146225#4199423, @jdoerfert wrote:

Why is this not part and done by sroa?

@lebedev.ri had some patches to do a much less aggressive vector formation during SROA

[jdoerfert]

llvm/test/CodeGen/AMDGPU/promote-alloca-memset.ll

32

the store of val is dead and SROA drops it. What is the point?
Fixing the test and I see reasonable SROA output, what is missing? https://godbolt.org/z/Pacce89fY

[Pierre-vh]

llvm/test/CodeGen/AMDGPU/promote-alloca-memset.ll

32

What I'm trying to do here is verify our assumption that the memset is correctly interpreted by SROA & removed. (Which is verified by the presence of the zeroinitializer in the insertelement)
The interesting check line is the OPT one, the SROA check is just to ensure the alloca is removed after.

[Pierre-vh]

In D146225#4200299, @arsenm wrote:

In D146225#4199423, @jdoerfert wrote:

Why is this not part and done by sroa?

@lebedev.ri had some patches to do a much less aggressive vector formation during SROA

(Unrelated to the review) do you have a link to those patches? I was planning to do some PromoteAlloca improvements later but if we're planning to merge it with SROA later then I suppose time would be better spent moving in that direction and helping these patches land

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

397

I null check redundant with one in caller

632

Can dyn_cast and use below instead of casting again

llvm/test/CodeGen/AMDGPU/promote-alloca-memset.ll

32

Might as well another reload and store to another output

[arsenm]

llvm/test/CodeGen/AMDGPU/promote-alloca-memset.ll

4

Don't see the point of separating the run lines, it's really only the aggregate result that matters

[Pierre-vh]

Comments

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

399

!I->isVolatile() would be clearer

636

I don't see how this could differ with the current code

[Pierre-vh]

Comments

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

636

I think t's better to be safe especially since this is (I think) cheap to do, but I could replace it with an assertion if you prefer.

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

 #include "Utils/AMDGPUBaseInfo.h"
 #include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/IntrinsicsR600.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/Pass.h"
 #include "llvm/Target/TargetMachine.h"
 
 #define DEBUG_TYPE "amdgpu-promote-alloca"
 
 using namespace llvm;
 
 namespace {
static Value *GEPToVectorIndex(GetElementPtrInst *GEP, AllocaInst *Alloca,
   return ConstantInt::get(GEP->getContext(), Quot);
 }
 
 struct MemTransferInfo {
   ConstantInt *SrcIndex = nullptr;
   ConstantInt *DestIndex = nullptr;
 };
 
+// Checks if the instruction I is a memset user of the alloca AI that we can
+// deal with. Currently, only non-volatile memsets that affect the whole alloca
+// are handled.
+static bool isSupportedMemset(MemSetInst *I, AllocaInst *AI,
+                              const DataLayout &DL) {
+  using namespace PatternMatch;
+  // For now we only care about non-volatile memsets that affect the whole type
+  // (start at index 0 and fill the whole alloca).
+  const unsigned Size = DL.getTypeStoreSize(AI->getAllocatedType());
+  return I->getOperand(0) == AI &&
+         match(I->getOperand(2), m_SpecificInt(Size)) && !I->isVolatile();
+}
+
 static bool tryPromoteAllocaToVector(AllocaInst *Alloca, const DataLayout &DL,
                                      unsigned MaxVGPRs) {
 
   if (DisablePromoteAllocaToVector) {
     LLVM_DEBUG(dbgs() << "  Promotion alloca to vector is disabled\n");
     return false;
   }
 
   Type *AllocaTy = Alloca->getAllocatedType();
   auto *VectorTy = dyn_cast<FixedVectorType>(AllocaTy);
   if (auto *ArrayTy = dyn_cast<ArrayType>(AllocaTy)) {
     if (VectorType::isValidElementType(ArrayTy->getElementType()) &&
         ArrayTy->getNumElements() > 0)
       VectorTy = arrayTypeToVecType(ArrayTy);
   }
if (auto *GEP = dyn_cast<GetElementPtrInst>(Inst)) {
       }
 
       GEPVectorIdx[GEP] = Index;
       for (Use &U : Inst->uses())
         Uses.push_back(&U);
       continue;
     }
 
+    if (MemSetInst *MSI = dyn_cast<MemSetInst>(Inst);
+        MSI && isSupportedMemset(MSI, Alloca, DL)) {
+      WorkList.push_back(Inst);
+      continue;
+    }
+
     if (MemTransferInst *TransferInst = dyn_cast<MemTransferInst>(Inst)) {
       if (TransferInst->isVolatile())
         return false;
 
       ConstantInt *Len = dyn_cast<ConstantInt>(TransferInst->getLength());
       if (!Len || !!(Len->getZExtValue() % ElementSize))
         return false;
 
while (!Uses.empty()) {
     if (isAssumeLikeIntrinsic(Inst))
       continue;
 
     if (isa<ICmpInst>(Inst) && all_of(Inst->users(), [](User *U) {
           return isAssumeLikeIntrinsic(cast<Instruction>(U));
         }))
       continue;
 
     // Unknown user.
     return false;
   }
 
   while (!DeferredInsts.empty()) {
     Instruction *Inst = DeferredInsts.pop_back_val();
     MemTransferInst *TransferInst = cast<MemTransferInst>(Inst);
     // TODO: Support the case if the pointers are from different alloca or
     // from different address spaces.
     MemTransferInfo &Info = TransferInfo[TransferInst];
     if (!Info.SrcIndex || !Info.DestIndex)
       return false;
case Instruction::Call: {
         Type *VecPtrTy = VectorTy->getPointerTo(Alloca->getAddressSpace());
         Value *BitCast = Builder.CreateBitCast(Alloca, VecPtrTy);
         Value *VecValue =
             Builder.CreateAlignedLoad(VectorTy, BitCast, Alloca->getAlign());
         Value *NewVecValue = Builder.CreateShuffleVector(VecValue, Mask);
         Builder.CreateAlignedStore(NewVecValue, BitCast, Alloca->getAlign());
 
         Inst->eraseFromParent();
+      } else if (MemSetInst *MSI = dyn_cast<MemSetInst>(Inst)) {
+        // Ensure the length parameter of the memsets matches the new vector
+        // type's. In general, the type size shouldn't change so this is a
+        // no-op, but it's better to be safe.
+        MSI->setOperand(2, Builder.getInt64(DL.getTypeStoreSize(VectorTy)));
       } else {
         llvm_unreachable("Unsupported call when promoting alloca to vector");
       }
       break;
     }
 
     default:
       llvm_unreachable("Inconsistency in instructions promotable to vector");
     }

llvm/test/CodeGen/AMDGPU/promote-alloca-memset.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -mtriple=amdgcn-unknown-amdhsa -passes=amdgpu-promote-alloca,sroa < %s | FileCheck %s
+
+; Checks that memsets don't block PromoteAlloca.
+
+; Note: memsets are just updated with the new type size. They are not eliminated which means
+; the original alloca also stay. This puts a bit more load on SROA.
+; If PromoteAlloca is moved to SSAUpdater, we could just entirely replace the memsets with
+; e.g. ConstantAggregate.
+
+define amdgpu_kernel void @memset_all_zero(i64 %val) {
+; CHECK-LABEL: @memset_all_zero(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = insertelement <6 x i64> zeroinitializer, i64 [[VAL:%.*]], i32 0
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <6 x i64> [[TMP0]], i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <6 x i64> [[TMP0]], i64 [[VAL]], i64 1
+; CHECK-NEXT:    ret void
+;
+entry:
+  %stack = alloca [6 x i64], align 4, addrspace(5)
+  call void @llvm.memset.p5.i64(ptr addrspace(5) %stack, i8 0, i64 48, i1 false)
+  store i64 %val, ptr addrspace(5) %stack
+  %reload = load i64, ptr addrspace(5) %stack
+  %stack.1 = getelementptr [6 x i64], ptr addrspace(5) %stack, i64 0, i64 1
+  store i64 %val, ptr addrspace(5) %stack.1
+  ret void
+}
+
+define amdgpu_kernel void @memset_all_5(i64 %val) {
+; CHECK-LABEL: @memset_all_5(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = insertelement <4 x i64> <i64 361700864190383365, i64 361700864190383365, i64 361700864190383365, i64 361700864190383365>, i64 [[VAL:%.*]], i32 0
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <4 x i64> [[TMP0]], i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <4 x i64> [[TMP0]], i64 [[VAL]], i64 1
+; CHECK-NEXT:    ret void
+;
+entry:
+  %stack = alloca [4 x i64], align 4, addrspace(5)
+  call void @llvm.memset.p5.i64(ptr addrspace(5) %stack, i8 5, i64 32, i1 false)
+  store i64 %val, ptr addrspace(5) %stack
+  %reload = load i64, ptr addrspace(5) %stack
+  %stack.1 = getelementptr [6 x i64], ptr addrspace(5) %stack, i64 0, i64 1
+  store i64 %val, ptr addrspace(5) %stack.1
+  ret void
+}
+
+define amdgpu_kernel void @memset_volatile_nopromote(i64 %val) {
+; CHECK-LABEL: @memset_volatile_nopromote(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[STACK_SROA_0:%.*]] = alloca i64, align 8, addrspace(5)
+; CHECK-NEXT:    [[STACK_SROA_2:%.*]] = alloca [3 x i64], align 8, addrspace(5)
+; CHECK-NEXT:    call void @llvm.memset.p5.i64(ptr addrspace(5) align 8 [[STACK_SROA_0]], i8 0, i64 8, i1 true)
+; CHECK-NEXT:    call void @llvm.memset.p5.i64(ptr addrspace(5) align 8 [[STACK_SROA_2]], i8 0, i64 24, i1 true)
+; CHECK-NEXT:    store i64 [[VAL:%.*]], ptr addrspace(5) [[STACK_SROA_0]], align 8
+; CHECK-NEXT:    ret void
+;
+entry:
+  %stack = alloca [4 x i64], align 4, addrspace(5)
+  call void @llvm.memset.p5.i64(ptr addrspace(5) %stack, i8 0, i64 32, i1 true)
+  store i64 %val, ptr addrspace(5) %stack
+  ret void
+}
+
+define amdgpu_kernel void @memset_badsize_nopromote(i64 %val) {
+; CHECK-LABEL: @memset_badsize_nopromote(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[STACK_SROA_0:%.*]] = alloca i64, align 8, addrspace(5)
+; CHECK-NEXT:    [[STACK_SROA_2:%.*]] = alloca [23 x i8], align 4, addrspace(5)
+; CHECK-NEXT:    call void @llvm.memset.p5.i64(ptr addrspace(5) align 8 [[STACK_SROA_0]], i8 0, i64 8, i1 true)
+; CHECK-NEXT:    call void @llvm.memset.p5.i64(ptr addrspace(5) align 4 [[STACK_SROA_2]], i8 0, i64 23, i1 true)
+; CHECK-NEXT:    store i64 [[VAL:%.*]], ptr addrspace(5) [[STACK_SROA_0]], align 8
+; CHECK-NEXT:    ret void
+;
+entry:
+  %stack = alloca [4 x i64], align 4, addrspace(5)
+  call void @llvm.memset.p5.i64(ptr addrspace(5) %stack, i8 0, i64 31, i1 true)
+  store i64 %val, ptr addrspace(5) %stack
+  ret void
+}
+
+define amdgpu_kernel void @memset_offset_ptr_nopromote(i64 %val) {
+; CHECK-LABEL: @memset_offset_ptr_nopromote(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[STACK_SROA_1:%.*]] = alloca [3 x i64], align 8, addrspace(5)
+; CHECK-NEXT:    call void @llvm.memset.p5.i64(ptr addrspace(5) align 8 [[STACK_SROA_1]], i8 0, i64 24, i1 true)
+; CHECK-NEXT:    ret void
+;
+entry:
+  %stack = alloca [4 x i64], align 4, addrspace(5)
+  %gep = getelementptr [4 x i64], ptr addrspace(5) %stack, i64 0, i64 1
+  call void @llvm.memset.p5.i64(ptr addrspace(5) %gep, i8 0, i64 24, i1 true)
+  store i64 %val, ptr addrspace(5) %stack
+  ret void
+}
+
+declare void @llvm.memset.p5.i64(ptr addrspace(5) nocapture writeonly, i8, i64, i1 immarg)