[llvm] [AMDGPU] Handle natively unsupported types in addrspace(7) lowering (PR #110572)
Krzysztof Drewniak via llvm-commits
llvm-commits at lists.llvm.org
Fri Nov 1 09:53:45 PDT 2024
================
@@ -576,6 +595,434 @@ bool StoreFatPtrsAsIntsVisitor::visitStoreInst(StoreInst &SI) {
return true;
}
+namespace {
+/// Convert loads/stores of types that the buffer intrinsics can't handle into
+/// one ore more such loads/stores that consist of legal types.
+///
+/// Do this by
+/// 1. Converting arrays of non-aggregate, byte-sized types into their
+/// correspondinng vectors
+/// 2. Bitcasting unsupported types, namely overly-long scalars and byte
+/// vectors, into vectors of supported types.
+/// 3. Splitting up excessively long reads/writes into multiple operations.
+///
+/// Note that this doesn't handle complex data strucures, but, in the future,
+/// the aggregate load splitter from SROA could be refactored to allow for that
+/// case.
+class LegalizeBufferContentTypesVisitor
+ : public InstVisitor<LegalizeBufferContentTypesVisitor, bool> {
+ friend class InstVisitor<LegalizeBufferContentTypesVisitor, bool>;
+
+ IRBuilder<> IRB;
+
+ const DataLayout &DL;
+
+ /// If T is [N x U], where U is a scalar type, return the vector type
+ /// <N x U>, otherwise, return T.
+ Type *scalarArrayTypeAsVector(Type *MaybeArrayType);
+ Value *arrayToVector(Value *V, Type *TargetType, StringRef Name);
+ Value *vectorToArray(Value *V, Type *OrigType, StringRef Name);
+
+ /// Break up the loads of a struct into the loads of its components
+
+ /// Convert a vector or scalar type that can't be operated on by buffer
+ /// intrinsics to one that would be legal through bitcasts and/or truncation.
+ /// Uses the wider of i32, i16, or i8 where possible.
+ Type *legalNonAggregateFor(Type *T);
+ Value *makeLegalNonAggregate(Value *V, Type *TargetType, StringRef Name);
+ Value *makeIllegalNonAggregate(Value *V, Type *OrigType, StringRef Name);
+
+ struct Slice {
+ unsigned Offset;
+ unsigned Length;
+ Slice(unsigned Offset, unsigned Length) : Offset(Offset), Length(Length) {}
+ };
+ // Return the [offset, length] pairs into which `T` needs to be cut to form
+ // legal buffer load or store operations. Clears `Slices`. Creates an empty
+ // `Slices` for non-vector inputs and creates one slice if no slicing will be
+ // needed.
+ void getSlices(Type *T, SmallVectorImpl<Slice> &Slices);
+
+ Value *extractSlice(Value *Vec, Slice S, StringRef Name);
+ Value *insertSlice(Value *Whole, Value *Part, Slice S, StringRef Name);
+
+ // In most cases, return `LegalType`. However, when given an input that would
+ // normally be a legal type for the buffer intrinsics to return but that isn't
+ // hooked up through SelectionDAG, return a type of the same width that can be
+ // used with the relevant intrinsics. Specifically, handle the cases:
+ // - <1 x T> => T for all T
+ // - <N x i8> <=> i16, i32, 2xi32, 4xi32 (as needed)
+ // - <N x T> where T is under 32 bits and the total size is 96 bits <=> <3 x
+ // i32>
+ Type *intrinsicTypeFor(Type *LegalType);
+
+ bool visitInstruction(Instruction &I) { return false; }
+ bool visitLoadInst(LoadInst &LI);
+ bool visitStoreInst(StoreInst &SI);
+
+public:
+ LegalizeBufferContentTypesVisitor(const DataLayout &DL, LLVMContext &Ctx)
+ : IRB(Ctx), DL(DL) {}
+ bool processFunction(Function &F);
+};
+} // namespace
+
+Type *LegalizeBufferContentTypesVisitor::scalarArrayTypeAsVector(Type *T) {
+ ArrayType *AT = dyn_cast<ArrayType>(T);
+ if (!AT)
+ return T;
+ Type *ET = AT->getElementType();
+ if (!ET->isSingleValueType() || isa<VectorType>(ET))
+ report_fatal_error(
+ "loading non-scalar arrays from buffer fat pointers is unimplemented");
+ if (!DL.typeSizeEqualsStoreSize(AT))
+ report_fatal_error(
+ "loading padded arrays from buffer fat pinters is unimplemented");
+ return FixedVectorType::get(ET, AT->getNumElements());
+}
+
+Value *LegalizeBufferContentTypesVisitor::arrayToVector(Value *V,
+ Type *TargetType,
+ StringRef Name) {
+ Value *VectorRes = PoisonValue::get(TargetType);
+ auto *VT = cast<FixedVectorType>(TargetType);
+ unsigned EC = VT->getNumElements();
+ for (auto I : iota_range<unsigned>(0, EC, /*Inclusive=*/false)) {
+ Value *Elem = IRB.CreateExtractValue(V, I, Name + ".elem." + Twine(I));
+ VectorRes = IRB.CreateInsertElement(VectorRes, Elem, I,
+ Name + ".as.vec." + Twine(I));
+ }
+ return VectorRes;
+}
+
+Value *LegalizeBufferContentTypesVisitor::vectorToArray(Value *V,
+ Type *OrigType,
+ StringRef Name) {
+ Value *ArrayRes = PoisonValue::get(OrigType);
+ ArrayType *AT = cast<ArrayType>(OrigType);
+ unsigned EC = AT->getNumElements();
+ for (auto I : iota_range<unsigned>(0, EC, /*Inclusive=*/false)) {
+ Value *Elem = IRB.CreateExtractElement(V, I, Name + ".elem." + Twine(I));
+ ArrayRes = IRB.CreateInsertValue(ArrayRes, Elem, I,
+ Name + ".as.array." + Twine(I));
+ }
+ return ArrayRes;
+}
+
+Type *LegalizeBufferContentTypesVisitor::legalNonAggregateFor(Type *T) {
+ TypeSize Size = DL.getTypeStoreSizeInBits(T);
+ // Implicitly zero-extend to the next byte if needed
+ if (!DL.typeSizeEqualsStoreSize(T))
+ T = IRB.getIntNTy(Size.getFixedValue());
+ auto *VT = dyn_cast<VectorType>(T);
+ Type *ElemTy = T;
+ if (VT) {
+ ElemTy = VT->getElementType();
+ }
+ if (isa<PointerType>(ElemTy))
+ return T; // Pointers are always big enough
+ unsigned ElemSize = DL.getTypeSizeInBits(ElemTy).getFixedValue();
+ if (isPowerOf2_32(ElemSize) && ElemSize >= 16 && ElemSize <= 128) {
+ // [vectors of] anything that's 16/32/64/128 bits can be cast and split into
+ // legal buffer operations.
+ return T;
+ }
+ Type *BestVectorElemType = nullptr;
+ if (Size.isKnownMultipleOf(32))
+ BestVectorElemType = IRB.getInt32Ty();
+ else if (Size.isKnownMultipleOf(16))
+ BestVectorElemType = IRB.getInt16Ty();
+ else
+ BestVectorElemType = IRB.getInt8Ty();
+ unsigned NumCastElems =
+ Size.getFixedValue() / BestVectorElemType->getIntegerBitWidth();
+ if (NumCastElems == 1)
+ return BestVectorElemType;
+ return FixedVectorType::get(BestVectorElemType, NumCastElems);
+}
+
+Value *LegalizeBufferContentTypesVisitor::makeLegalNonAggregate(
+ Value *V, Type *TargetType, StringRef Name) {
+ Type *SourceType = V->getType();
+ if (DL.getTypeSizeInBits(SourceType) != DL.getTypeSizeInBits(TargetType)) {
+ Type *ShortScalarTy =
+ IRB.getIntNTy(DL.getTypeSizeInBits(SourceType).getFixedValue());
+ Type *ByteScalarTy =
+ IRB.getIntNTy(DL.getTypeSizeInBits(TargetType).getFixedValue());
+ Value *AsScalar = IRB.CreateBitCast(V, ShortScalarTy, Name + ".as.scalar");
+ Value *Zext = IRB.CreateZExt(AsScalar, ByteScalarTy, Name + ".zext");
+ V = Zext;
+ SourceType = ByteScalarTy;
+ }
+ if (SourceType == TargetType)
+ return V;
+ return IRB.CreateBitCast(V, TargetType, Name + ".legal");
+}
+
+Value *LegalizeBufferContentTypesVisitor::makeIllegalNonAggregate(
+ Value *V, Type *OrigType, StringRef Name) {
+ Type *LegalType = V->getType();
+ if (DL.getTypeSizeInBits(LegalType) != DL.getTypeSizeInBits(OrigType)) {
+ Type *ShortScalarTy =
+ IRB.getIntNTy(DL.getTypeSizeInBits(OrigType).getFixedValue());
+ Type *ByteScalarTy =
+ IRB.getIntNTy(DL.getTypeSizeInBits(LegalType).getFixedValue());
+ Value *AsScalar = IRB.CreateBitCast(V, ByteScalarTy, Name + ".bytes.cast");
+ Value *Trunc = IRB.CreateTrunc(AsScalar, ShortScalarTy, Name + ".trunc");
+ if (OrigType != ShortScalarTy)
+ return IRB.CreateBitCast(Trunc, OrigType, Name + ".orig");
+ return Trunc;
+ }
+ if (LegalType == OrigType)
+ return V;
+ return IRB.CreateBitCast(V, OrigType, Name + ".real.ty");
+}
+
+Type *LegalizeBufferContentTypesVisitor::intrinsicTypeFor(Type *LegalType) {
+ auto *VT = dyn_cast<FixedVectorType>(LegalType);
+ if (!VT)
+ return LegalType;
+ Type *ET = VT->getElementType();
+ if (VT->getNumElements() == 1)
+ return ET;
+ if (DL.getTypeSizeInBits(LegalType) == 96 && DL.getTypeSizeInBits(ET) < 32)
+ return FixedVectorType::get(IRB.getInt32Ty(), 3);
+ if (ET->isIntegerTy(8)) {
+ switch (VT->getNumElements()) {
+ default:
+ return LegalType; // Let it crash later
+ case 1:
+ return IRB.getInt8Ty();
+ case 2:
+ return IRB.getInt16Ty();
+ case 4:
+ return IRB.getInt32Ty();
+ case 8:
+ return FixedVectorType::get(IRB.getInt32Ty(), 2);
+ case 16:
+ return FixedVectorType::get(IRB.getInt32Ty(), 4);
+ }
+ }
+ return LegalType;
+}
+
+void LegalizeBufferContentTypesVisitor::getSlices(
+ Type *T, SmallVectorImpl<Slice> &Slices) {
+ Slices.clear();
+ auto *VT = dyn_cast<FixedVectorType>(T);
+ if (!VT)
+ return;
+
+ unsigned ElemBitWidth =
+ DL.getTypeSizeInBits(VT->getElementType()).getFixedValue();
+
+ unsigned ElemsPer4Words = 128 / ElemBitWidth;
+ unsigned ElemsPer2Words = ElemsPer4Words / 2;
+ unsigned ElemsPerWord = ElemsPer2Words / 2;
+ unsigned ElemsPerShort = ElemsPerWord / 2;
+ unsigned ElemsPerByte = ElemsPerShort / 2;
+ // If the elements evenly pack into 32-bit words, we can use 3-word stores,
+ // such as for <6 x bfloat> or <3 x i32>, but we can't dot his for, for
+ // example, <3 x i64>, since that's not slicing.
+ unsigned ElemsPer3Words = ElemsPerWord * 3;
+
+ unsigned TotalElems = VT->getNumElements();
+ unsigned Off = 0;
+ auto TrySlice = [&](unsigned MaybeLen) {
+ if (MaybeLen > 0 && Off + MaybeLen <= TotalElems) {
+ Slices.emplace_back(/*Offset=*/Off, /*Length=*/MaybeLen);
+ Off += MaybeLen;
+ return true;
+ }
+ return false;
+ };
+ while (Off < TotalElems) {
+ TrySlice(ElemsPer4Words) || TrySlice(ElemsPer3Words) ||
+ TrySlice(ElemsPer2Words) || TrySlice(ElemsPerWord) ||
+ TrySlice(ElemsPerShort) || TrySlice(ElemsPerByte);
+ }
+}
+
+Value *LegalizeBufferContentTypesVisitor::extractSlice(Value *Vec, Slice S,
+ StringRef Name) {
+ if (S.Length == 1)
+ return IRB.CreateExtractElement(Vec, S.Offset,
+ Name + ".slice." + Twine(S.Offset));
+ SmallVector<int> Mask = llvm::to_vector(llvm::iota_range<int>(
+ S.Offset, S.Offset + S.Length, /*Inclusive=*/false));
+ return IRB.CreateShuffleVector(Vec, Mask, Name + ".slice." + Twine(S.Offset));
+}
+
+Value *LegalizeBufferContentTypesVisitor::insertSlice(Value *Whole, Value *Part,
+ Slice S, StringRef Name) {
+ if (S.Length == 1) {
+ return IRB.CreateInsertElement(Whole, Part, S.Offset,
+ Name + ".slice." + Twine(S.Offset));
+ }
+ int NumElems = cast<FixedVectorType>(Whole->getType())->getNumElements();
+
+ // Extend the slice with poisons to make the main shufflevector happy.
+ SmallVector<int> ExtPartMask(NumElems, -1);
+ for (auto [I, E] : llvm::enumerate(
+ MutableArrayRef<int>(ExtPartMask).take_front(S.Length))) {
+ E = I;
+ }
+ Value *ExtPart = IRB.CreateShuffleVector(Part, ExtPartMask,
+ Name + ".ext." + Twine(S.Offset));
+
+ SmallVector<int> Mask =
+ llvm::to_vector(llvm::iota_range<int>(0, NumElems, /*Inclusive=*/false));
+ for (auto [I, E] :
+ llvm::enumerate(MutableArrayRef<int>(Mask).slice(S.Offset, S.Length)))
+ E = I + NumElems;
+ return IRB.CreateShuffleVector(Whole, ExtPart, Mask,
+ Name + ".parts." + Twine(S.Offset));
+}
+
+bool LegalizeBufferContentTypesVisitor::visitLoadInst(LoadInst &LI) {
+ if (LI.getPointerAddressSpace() != AMDGPUAS::BUFFER_FAT_POINTER)
+ return false;
+ Type *OrigType = LI.getType();
+ Type *ArrayAsVecType = scalarArrayTypeAsVector(OrigType);
+ Type *LegalType = legalNonAggregateFor(ArrayAsVecType);
+
+ SmallVector<Slice> Slices;
+ getSlices(LegalType, Slices);
+ bool NeedToSplit = Slices.size() > 1;
+ Value *LoadsRes;
+ StringRef Name = LI.getName();
+ if (!NeedToSplit) {
+ Type *LoadableType = intrinsicTypeFor(LegalType);
+ if (LoadableType == OrigType)
+ return false;
+
+ IRB.SetInsertPoint(&LI);
+ auto *NLI = cast<LoadInst>(LI.clone());
+ NLI->mutateType(LoadableType);
+ NLI = IRB.Insert(NLI);
+ NLI->setName(Name + ".loadable");
+
+ LoadsRes = NLI;
+ if (LoadableType != LegalType) {
+ LoadsRes =
+ IRB.CreateBitCast(LoadsRes, LegalType, Name + ".from.loadable");
+ }
+ } else {
+ IRB.SetInsertPoint(&LI);
+ LoadsRes = PoisonValue::get(LegalType);
+ Value *OrigPtr = LI.getPointerOperand();
+ // If we're needing to spill something into more than one load, its legal
+ // type will be a vector (ex. an i256 load will have LegalType = <8 x i32>).
+ Type *ElemType = cast<VectorType>(LegalType)->getElementType();
+ unsigned ElemBytes = DL.getTypeStoreSize(ElemType);
+ AAMDNodes AANodes = LI.getAAMetadata();
+ for (Slice S : Slices) {
+ Type *SliceType =
+ S.Length != 1 ? FixedVectorType::get(ElemType, S.Length) : ElemType;
+ unsigned ByteOffset = S.Offset * ElemBytes;
+ // You can't reasonably expect loads to wrap around the edge of memory.
+ Value *NewPtr = IRB.CreateGEP(
+ IRB.getInt8Ty(), LI.getPointerOperand(), IRB.getInt32(ByteOffset),
+ OrigPtr->getName() + ".part.ptr." + Twine(S.Offset),
+ GEPNoWrapFlags::noUnsignedWrap());
+ Type *LoadableType = intrinsicTypeFor(SliceType);
+ LoadInst *NewLI = IRB.CreateAlignedLoad(
+ LoadableType, NewPtr, commonAlignment(LI.getAlign(), ByteOffset),
+ Name + ".part." + Twine(S.Offset));
+ copyMetadataForLoad(*NewLI, LI);
+ NewLI->setAAMetadata(
+ AANodes.adjustForAccess(ByteOffset, LoadableType, DL));
----------------
krzysz00 wrote:
Maybe? Not sure this comes up all that often
https://github.com/llvm/llvm-project/pull/110572
More information about the llvm-commits
mailing list