[DirectX] Bug fix for Data Scalarization crash

farzonl · farzonl · commit 7a4f55518184 · 2024-12-02T23:34:59.000-05:00
diff --git a/llvm/lib/Target/DirectX/DXILDataScalarization.cpp b/llvm/lib/Target/DirectX/DXILDataScalarization.cpp
@@ -67,28 +67,17 @@ class DataScalarizerVisitor : public InstVisitor<DataScalarizerVisitor, bool> {
 private:
   GlobalVariable *lookupReplacementGlobal(Value *CurrOperand);
   DenseMap<GlobalVariable *, GlobalVariable *> GlobalMap;
-  SmallVector<WeakTrackingVH, 32> PotentiallyDeadInstrs;
-  bool finish();
 };
 
 bool DataScalarizerVisitor::visit(Function &F) {
   assert(!GlobalMap.empty());
-  ReversePostOrderTraversal<BasicBlock *> RPOT(&F.getEntryBlock());
-  for (BasicBlock *BB : RPOT) {
-    for (BasicBlock::iterator II = BB->begin(), IE = BB->end(); II != IE;) {
-      Instruction *I = &*II;
-      bool Done = InstVisitor::visit(I);
-      ++II;
-      if (Done && I->getType()->isVoidTy())
-        I->eraseFromParent();
-    }
+  bool MadeChange = false;
+  ReversePostOrderTraversal<Function *> RPOT(&F);
+  for (BasicBlock *BB : make_early_inc_range(RPOT)) {
+    for (Instruction &I : make_early_inc_range(*BB))
+      MadeChange |= InstVisitor::visit(I);
   }
-  return finish();
-}
-
-bool DataScalarizerVisitor::finish() {
-  RecursivelyDeleteTriviallyDeadInstructionsPermissive(PotentiallyDeadInstrs);
-  return true;
+  return MadeChange;
 }
 
 GlobalVariable *
@@ -140,7 +129,7 @@ bool DataScalarizerVisitor::visitGetElementPtrInst(GetElementPtrInst &GEPI) {
         Builder.CreateGEP(NewGlobal->getValueType(), NewGlobal, Indices);
 
     GEPI.replaceAllUsesWith(NewGEP);
-    PotentiallyDeadInstrs.emplace_back(&GEPI);
+    GEPI.eraseFromParent();
   }
   return true;
 }
@@ -252,8 +241,7 @@ static bool findAndReplaceVectors(Module &M) {
                                                 /*RemoveDeadConstants=*/false,
                                                 /*IncludeSelf=*/true);
         }
-        if (isa<Instruction>(U)) {
-          Instruction *Inst = cast<Instruction>(U);
+        if (Instruction *Inst = dyn_cast<Instruction>(U)) {
           Function *F = Inst->getFunction();
           if (F)
             Impl.visit(*F);
diff --git a/llvm/lib/Target/DirectX/DXILFlattenArrays.cpp b/llvm/lib/Target/DirectX/DXILFlattenArrays.cpp
@@ -321,6 +321,18 @@ static void collectElements(Constant *Init,
     Elements.push_back(Init);
     return;
   }
+  auto *LLVMArrayType = dyn_cast<ArrayType>(Init->getType());
+  if (isa<ConstantAggregateZero>(Init)) {
+    for (unsigned I = 0; I < LLVMArrayType->getNumElements(); ++I)
+      Elements.push_back(
+          Constant::getNullValue(LLVMArrayType->getElementType()));
+    return;
+  }
+  if (isa<UndefValue>(Init)) {
+    for (unsigned I = 0; I < LLVMArrayType->getNumElements(); ++I)
+      Elements.push_back(UndefValue::get(LLVMArrayType->getElementType()));
+    return;
+  }
 
   // Recursive case: Process each element in the array.
   if (auto *ArrayConstant = dyn_cast<ConstantArray>(Init)) {
diff --git a/llvm/test/CodeGen/DirectX/flatten-bug-117273.ll b/llvm/test/CodeGen/DirectX/flatten-bug-117273.ll
@@ -0,0 +1,25 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes='dxil-flatten-arrays,dxil-op-lower' -mtriple=dxil-pc-shadermodel6.3-library %s | FileCheck %s
+
+
+@_ZL7Palette = internal constant [2 x [3 x float]] [[3 x float] zeroinitializer, [3 x float] undef], align 16
+
+; CHECK: @_ZL7Palette.1dim = internal constant [6 x float] [float 0.000000e+00, float 0.000000e+00, float 0.000000e+00, float undef, float undef, float undef], align 16
+
+define internal void @_Z4mainDv3_j(<3 x i32> noundef %DID) {
+; CHECK-LABEL: define internal void @_Z4mainDv3_j(
+; CHECK-SAME: <3 x i32> noundef [[DID:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr [24 x float], ptr @_ZL7Palette.1dim, i32 1
+; CHECK-NEXT:    [[DOTI0:%.*]] = load float, ptr [[TMP0]], align 16
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [24 x float], ptr @_ZL7Palette.1dim, i32 2
+; CHECK-NEXT:    [[DOTI03:%.*]] = load float, ptr [[TMP1]], align 16
+; CHECK-NEXT:    ret void
+;
+entry:
+  %0 = getelementptr [8 x [3 x float]], ptr @_ZL7Palette, i32 0, i32 1
+  %.i0 = load float, ptr %0, align 16
+  %1 = getelementptr [8 x [3 x float]], ptr @_ZL7Palette, i32 0, i32 2
+  %.i03 = load float, ptr %1, align 16
+  ret void
+}
diff --git a/llvm/test/CodeGen/DirectX/scalar-bug-117273.ll b/llvm/test/CodeGen/DirectX/scalar-bug-117273.ll
@@ -0,0 +1,32 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes='dxil-data-scalarization,dxil-flatten-arrays,function(scalarizer<load-store>),dxil-op-lower' -mtriple=dxil-pc-shadermodel6.3-library %s | FileCheck %s
+
+
+@_ZL7Palette = internal constant [8 x <3 x float>] [<3 x float> zeroinitializer, <3 x float> splat (float 5.000000e-01), <3 x float> <float 1.000000e+00, float 5.000000e-01, float 5.000000e-01>, <3 x float> <float 5.000000e-01, float 1.000000e+00, float 5.000000e-01>, <3 x float> <float 5.000000e-01, float 5.000000e-01, float 1.000000e+00>, <3 x float> <float 5.000000e-01, float 1.000000e+00, float 1.000000e+00>, <3 x float> <float 1.000000e+00, float 5.000000e-01, float 1.000000e+00>, <3 x float> <float 1.000000e+00, float 1.000000e+00, float 5.000000e-01>], align 16
+
+; Function Attrs: alwaysinline convergent mustprogress norecurse nounwind
+define internal void @_Z4mainDv3_j(<3 x i32> noundef %DID) #1 {
+; CHECK-LABEL: define internal void @_Z4mainDv3_j(
+; CHECK-SAME: <3 x i32> noundef [[DID:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr [24 x float], ptr @_ZL7Palette.scalarized.1dim, i32 1
+; CHECK-NEXT:    [[DOTI0:%.*]] = load float, ptr [[TMP0]], align 16
+; CHECK-NEXT:    [[DOTI1:%.*]] = getelementptr float, ptr [[TMP0]], i32 1
+; CHECK-NEXT:    [[DOTI11:%.*]] = load float, ptr [[DOTI1]], align 4
+; CHECK-NEXT:    [[DOTI2:%.*]] = getelementptr float, ptr [[TMP0]], i32 2
+; CHECK-NEXT:    [[DOTI22:%.*]] = load float, ptr [[DOTI2]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [24 x float], ptr @_ZL7Palette.scalarized.1dim, i32 2
+; CHECK-NEXT:    [[DOTI03:%.*]] = load float, ptr [[TMP1]], align 16
+; CHECK-NEXT:    [[DOTI14:%.*]] = getelementptr float, ptr [[TMP1]], i32 1
+; CHECK-NEXT:    [[DOTI15:%.*]] = load float, ptr [[DOTI14]], align 4
+; CHECK-NEXT:    [[DOTI26:%.*]] = getelementptr float, ptr [[TMP1]], i32 2
+; CHECK-NEXT:    [[DOTI27:%.*]] = load float, ptr [[DOTI26]], align 8
+; CHECK-NEXT:    ret void
+;
+entry:
+  %arrayidx = getelementptr inbounds [8 x <3 x float>], ptr @_ZL7Palette, i32 0, i32 1
+  %2 = load <3 x float>, ptr %arrayidx, align 16
+  %arrayidx2 = getelementptr inbounds [8 x <3 x float>], ptr @_ZL7Palette, i32 0, i32 2
+  %3 = load <3 x float>, ptr %arrayidx2, align 16
+  ret void
+}
