[AMDGPU] In promote-alloca, if index is dynamic, sandwich load with bitcasts to reduce excessive codegen

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

@@ -42,6 +42,7 @@
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/SSAUpdater.h"
 
@@ -644,6 +645,42 @@ static Value *promoteAllocaUserToVector(Instruction *Inst, const DataLayout &DL,
       auto *SubVecTy = FixedVectorType::get(VecEltTy, NumLoadedElts);
       assert(DL.getTypeStoreSize(SubVecTy) == DL.getTypeStoreSize(AccessTy));
 
+      // If idx is dynamic, then sandwich load with bitcasts.
+      // ie. CurValTy                 SubVecTy  AccessTy
+      //     <64 x i8> ->             <16 x i8> <8 x i16>
+      //     <64 x i8> -> <4 x i128> -> i128 -> <8 x i16>
+      // Extracting subvector with dynamic index has very large expansion in
+      // the amdgpu backend. Limit to pow2 for UDiv.
+      FixedVectorType *AccessVecTy = cast<FixedVectorType>(AccessTy);
+      FixedVectorType *VectorTy = AA.Vector.Ty;
+      uint64_t NumBits =
+          SubVecTy->getScalarSizeInBits() * SubVecTy->getNumElements();
+      uint64_t LoadAlign = cast<LoadInst>(Inst)->getAlign().value();
+      if (!isa<ConstantInt>(Index) && SubVecTy->isIntOrIntVectorTy() &&
+          llvm::isPowerOf2_32(SubVecTy->getNumElements()) &&
+          VectorTy->getNumElements() % SubVecTy->getNumElements() == 0 &&
+          llvm::isPowerOf2_32(AccessVecTy->getNumElements()) &&
+          NumBits <= (LoadAlign * 8u)) {
+        IntegerType *NewElemType = Builder.getIntNTy(NumBits);
+        const unsigned NewNumElts = VectorTy->getNumElements() *
+                                    VectorTy->getScalarSizeInBits() /
+                                    NewElemType->getScalarSizeInBits();
+        const unsigned IndexDivisor = VectorTy->getNumElements() / NewNumElts;
+        assert(VectorTy->getScalarSizeInBits() <
+                   NewElemType->getScalarSizeInBits() &&
+               "New element type should be bigger");
+        assert(IndexDivisor > 0u && "Zero index divisor");
+        FixedVectorType *BitCastType =
+            FixedVectorType::get(NewElemType, NewNumElts);
+        Value *BCVal = Builder.CreateBitCast(CurVal, BitCastType);
+        Value *NewIdx = Builder.CreateUDiv(
+            Index, ConstantInt::get(Index->getType(), IndexDivisor));
+        Value *ExtVal = Builder.CreateExtractElement(BCVal, NewIdx);
+        Value *BCOut = Builder.CreateBitCast(ExtVal, AccessTy);
+        Inst->replaceAllUsesWith(BCOut);
+        return nullptr;
+      }
+
       Value *SubVec = PoisonValue::get(SubVecTy);
       for (unsigned K = 0; K < NumLoadedElts; ++K) {
         Value *CurIdx =