[llvm] [DirectX] Scalarize `extractelement` with dynamic index (PR #141676)
Deric C. via llvm-commits
llvm-commits at lists.llvm.org
Tue May 27 14:47:16 PDT 2025
https://github.com/Icohedron created https://github.com/llvm/llvm-project/pull/141676
Fixes #141136
- Implement `visitExtractElementInst` in `DXILDataScalarizerVisitor` to scalarize `extractelement` instructions whose index operand is not a `ConstantInt` by converting the vector to an array and then loading from the array
- Rename the `replaceVectorWithArray` helper function to `equivalentArrayTypeFromVector`, relocate the function toward the top of the file, and remove the unused `Ctx` parameter
>From d3f1a51cef21e74b281ba1dbf38bff16410b20e1 Mon Sep 17 00:00:00 2001
From: Icohedron <cheung.deric at gmail.com>
Date: Tue, 27 May 2025 21:28:12 +0000
Subject: [PATCH] Scalarize extractelement with dynamic index
---
.../Target/DirectX/DXILDataScalarization.cpp | 66 ++++++++++++++-----
.../DirectX/scalarize-dynamic-vector-index.ll | 38 +++++++++++
2 files changed, 86 insertions(+), 18 deletions(-)
create mode 100644 llvm/test/CodeGen/DirectX/scalarize-dynamic-vector-index.ll
diff --git a/llvm/lib/Target/DirectX/DXILDataScalarization.cpp b/llvm/lib/Target/DirectX/DXILDataScalarization.cpp
index 06708cec00cec..7bd0539c6bfe0 100644
--- a/llvm/lib/Target/DirectX/DXILDataScalarization.cpp
+++ b/llvm/lib/Target/DirectX/DXILDataScalarization.cpp
@@ -27,6 +27,19 @@ static const int MaxVecSize = 4;
using namespace llvm;
+// Recursively creates an array-like version of a given vector type.
+static Type *equivalentArrayTypeFromVector(Type *T) {
+ if (auto *VecTy = dyn_cast<VectorType>(T))
+ return ArrayType::get(VecTy->getElementType(),
+ dyn_cast<FixedVectorType>(VecTy)->getNumElements());
+ if (auto *ArrayTy = dyn_cast<ArrayType>(T)) {
+ Type *NewElementType = equivalentArrayTypeFromVector(ArrayTy->getElementType());
+ return ArrayType::get(NewElementType, ArrayTy->getNumElements());
+ }
+ // If it's not a vector or array, return the original type.
+ return T;
+}
+
class DXILDataScalarizationLegacy : public ModulePass {
public:
@@ -55,7 +68,7 @@ class DataScalarizerVisitor : public InstVisitor<DataScalarizerVisitor, bool> {
bool visitCastInst(CastInst &CI) { return false; }
bool visitBitCastInst(BitCastInst &BCI) { return false; }
bool visitInsertElementInst(InsertElementInst &IEI) { return false; }
- bool visitExtractElementInst(ExtractElementInst &EEI) { return false; }
+ bool visitExtractElementInst(ExtractElementInst &EEI);
bool visitShuffleVectorInst(ShuffleVectorInst &SVI) { return false; }
bool visitPHINode(PHINode &PHI) { return false; }
bool visitLoadInst(LoadInst &LI);
@@ -90,20 +103,6 @@ DataScalarizerVisitor::lookupReplacementGlobal(Value *CurrOperand) {
return nullptr; // Not found
}
-// Recursively creates an array version of the given vector type.
-static Type *replaceVectorWithArray(Type *T, LLVMContext &Ctx) {
- if (auto *VecTy = dyn_cast<VectorType>(T))
- return ArrayType::get(VecTy->getElementType(),
- dyn_cast<FixedVectorType>(VecTy)->getNumElements());
- if (auto *ArrayTy = dyn_cast<ArrayType>(T)) {
- Type *NewElementType =
- replaceVectorWithArray(ArrayTy->getElementType(), Ctx);
- return ArrayType::get(NewElementType, ArrayTy->getNumElements());
- }
- // If it's not a vector or array, return the original type.
- return T;
-}
-
static bool isArrayOfVectors(Type *T) {
if (ArrayType *ArrType = dyn_cast<ArrayType>(T))
return isa<VectorType>(ArrType->getElementType());
@@ -116,8 +115,7 @@ bool DataScalarizerVisitor::visitAllocaInst(AllocaInst &AI) {
ArrayType *ArrType = cast<ArrayType>(AI.getAllocatedType());
IRBuilder<> Builder(&AI);
- LLVMContext &Ctx = AI.getContext();
- Type *NewType = replaceVectorWithArray(ArrType, Ctx);
+ Type *NewType = equivalentArrayTypeFromVector(ArrType);
AllocaInst *ArrAlloca =
Builder.CreateAlloca(NewType, nullptr, AI.getName() + ".scalarize");
ArrAlloca->setAlignment(AI.getAlign());
@@ -173,6 +171,38 @@ bool DataScalarizerVisitor::visitStoreInst(StoreInst &SI) {
return false;
}
+bool DataScalarizerVisitor::visitExtractElementInst(ExtractElementInst &EEI) {
+ // If the index is a constant then we don't need to scalarize it
+ Value *Index = EEI.getIndexOperand();
+ Type *IndexTy = Index->getType();
+ if (isa<ConstantInt>(Index))
+ return false;
+
+ IRBuilder<> Builder(&EEI);
+ VectorType *VecTy = EEI.getVectorOperandType();
+ assert(VecTy->getElementCount().isFixed() &&
+ "Vector operand of ExtractElement must have a fixed size");
+
+ Type *ArrTy = equivalentArrayTypeFromVector(VecTy);
+ Value *ArrAlloca = Builder.CreateAlloca(ArrTy);
+
+ for (unsigned I = 0; I < ArrTy->getArrayNumElements(); ++I) {
+ Value *EE = Builder.CreateExtractElement(EEI.getVectorOperand(), I);
+ Value *GEP = Builder.CreateInBoundsGEP(
+ ArrTy, ArrAlloca,
+ {ConstantInt::get(IndexTy, 0), ConstantInt::get(IndexTy, I)});
+ Builder.CreateStore(EE, GEP);
+ }
+
+ Value *GEP = Builder.CreateInBoundsGEP(ArrTy, ArrAlloca,
+ {ConstantInt::get(IndexTy, 0), Index});
+ Value *Load = Builder.CreateLoad(ArrTy->getArrayElementType(), GEP);
+
+ EEI.replaceAllUsesWith(Load);
+ EEI.eraseFromParent();
+ return true;
+}
+
bool DataScalarizerVisitor::visitGetElementPtrInst(GetElementPtrInst &GEPI) {
unsigned NumOperands = GEPI.getNumOperands();
@@ -257,7 +287,7 @@ static bool findAndReplaceVectors(Module &M) {
for (GlobalVariable &G : M.globals()) {
Type *OrigType = G.getValueType();
- Type *NewType = replaceVectorWithArray(OrigType, Ctx);
+ Type *NewType = equivalentArrayTypeFromVector(OrigType);
if (OrigType != NewType) {
// Create a new global variable with the updated type
// Note: Initializer is set via transformInitializer
diff --git a/llvm/test/CodeGen/DirectX/scalarize-dynamic-vector-index.ll b/llvm/test/CodeGen/DirectX/scalarize-dynamic-vector-index.ll
new file mode 100644
index 0000000000000..74e9202b540c1
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/scalarize-dynamic-vector-index.ll
@@ -0,0 +1,38 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes='dxil-data-scalarization' -mtriple=dxil-pc-shadermodel6.3-library %s | FileCheck %s
+
+define float @extract_float_vec_dynamic(<4 x float> %0, i32 %1) {
+; CHECK-LABEL: define float @extract_float_vec_dynamic(
+; CHECK-SAME: <4 x float> [[TMP0:%.*]], i32 [[TMP1:%.*]]) {
+; CHECK-NEXT: [[TMP3:%.*]] = alloca [4 x float], align 4
+; CHECK-NEXT: [[TMP4:%.*]] = extractelement <4 x float> [[TMP0]], i64 0
+; CHECK-NEXT: [[TMP5:%.*]] = getelementptr inbounds [4 x float], ptr [[TMP3]], i32 0, i32 0
+; CHECK-NEXT: store float [[TMP4]], ptr [[TMP5]], align 4
+; CHECK-NEXT: [[TMP6:%.*]] = extractelement <4 x float> [[TMP0]], i64 1
+; CHECK-NEXT: [[TMP7:%.*]] = getelementptr inbounds [4 x float], ptr [[TMP3]], i32 0, i32 1
+; CHECK-NEXT: store float [[TMP6]], ptr [[TMP7]], align 4
+; CHECK-NEXT: [[TMP8:%.*]] = extractelement <4 x float> [[TMP0]], i64 2
+; CHECK-NEXT: [[TMP9:%.*]] = getelementptr inbounds [4 x float], ptr [[TMP3]], i32 0, i32 2
+; CHECK-NEXT: store float [[TMP8]], ptr [[TMP9]], align 4
+; CHECK-NEXT: [[TMP10:%.*]] = extractelement <4 x float> [[TMP0]], i64 3
+; CHECK-NEXT: [[TMP11:%.*]] = getelementptr inbounds [4 x float], ptr [[TMP3]], i32 0, i32 3
+; CHECK-NEXT: store float [[TMP10]], ptr [[TMP11]], align 4
+; CHECK-NEXT: [[TMP12:%.*]] = getelementptr inbounds [4 x float], ptr [[TMP3]], i32 0, i32 [[TMP1]]
+; CHECK-NEXT: [[TMP13:%.*]] = load float, ptr [[TMP12]], align 4
+; CHECK-NEXT: ret float [[TMP13]]
+;
+ %e = extractelement <4 x float> %0, i32 %1
+ ret float %e
+}
+
+; An extractelement with a constant index should not be converted to array form
+define i16 @extract_i16_vec_constant(<4 x i16> %0) {
+; CHECK-LABEL: define i16 @extract_i16_vec_constant(
+; CHECK-SAME: <4 x i16> [[TMP0:%.*]]) {
+; CHECK-NEXT: [[E:%.*]] = extractelement <4 x i16> [[TMP0]], i32 1
+; CHECK-NEXT: ret i16 [[E]]
+;
+ %e = extractelement <4 x i16> %0, i32 1
+ ret i16 %e
+}
+
More information about the llvm-commits
mailing list