fix iterator stability issue

Author: farzonl

llvm/lib/Target/DirectX/DXILDataScalarization.cpp

@@ -1,15 +1,15 @@
-//===- DXILDataScalarization.cpp - Prepare LLVM Module for DXIL Data
-// Legalization----===//
+//===- DXILDataScalarization.cpp - Perform DXIL Data Legalization----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
-//===--------------------------------------------------------------------------------===//
+//===----------------------------------------------------------------===//
 
 #include "DXILDataScalarization.h"
 #include "DirectX.h"
 #include "llvm/ADT/PostOrderIterator.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstVisitor.h"
@@ -20,7 +20,6 @@
 #include "llvm/IR/Type.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
-#include <utility>
 
 #define DEBUG_TYPE "dxil-data-scalarization"
 #define Max_VEC_SIZE 4
@@ -164,14 +163,16 @@ Constant *transformInitializer(Constant *Init, Type *OrigType, Type *NewType,
   if (isa<VectorType>(OrigType) && isa<ArrayType>(NewType)) {
     // Convert vector initializer to array initializer
     SmallVector<Constant *, Max_VEC_SIZE> ArrayElements;
-    if( ConstantVector *ConstVecInit = dyn_cast<ConstantVector>(Init)) {
-        for (unsigned I = 0; I < ConstVecInit->getNumOperands(); ++I) 
-            ArrayElements.push_back(ConstVecInit->getOperand(I));
-    } else if (ConstantDataVector *ConstDataVecInit = llvm::dyn_cast<llvm::ConstantDataVector>(Init)) {
-        for (unsigned I = 0; I < ConstDataVecInit->getNumElements(); ++I) 
-            ArrayElements.push_back(ConstDataVecInit->getElementAsConstant(I));
-    }  else {
-      llvm_unreachable("Expected a ConstantVector or ConstantDataVector for vector initializer!");
+    if (ConstantVector *ConstVecInit = dyn_cast<ConstantVector>(Init)) {
+      for (unsigned I = 0; I < ConstVecInit->getNumOperands(); ++I)
+        ArrayElements.push_back(ConstVecInit->getOperand(I));
+    } else if (ConstantDataVector *ConstDataVecInit =
+                   llvm::dyn_cast<llvm::ConstantDataVector>(Init)) {
+      for (unsigned I = 0; I < ConstDataVecInit->getNumElements(); ++I)
+        ArrayElements.push_back(ConstDataVecInit->getElementAsConstant(I));
+    } else {
+      llvm_unreachable("Expected a ConstantVector or ConstantDataVector for "
+                       "vector initializer!");
     }
 
     return ConstantArray::get(cast<ArrayType>(NewType), ArrayElements);
@@ -213,9 +214,10 @@ static void findAndReplaceVectors(Module &M) {
       // Create a new global variable with the updated type
       GlobalVariable *NewGlobal = new GlobalVariable(
           M, NewType, G.isConstant(), G.getLinkage(),
-          // This is set via: transformInitializer
-          nullptr, G.getName() + ".scalarized", &G, G.getThreadLocalMode(),
-          G.getAddressSpace(), G.isExternallyInitialized());
+          // Initializer is set via transformInitializer
+          /*Initializer=*/nullptr, G.getName() + ".scalarized", &G,
+          G.getThreadLocalMode(), G.getAddressSpace(),
+          G.isExternallyInitialized());
 
       // Copy relevant attributes
       NewGlobal->setUnnamedAddr(G.getUnnamedAddr());
@@ -230,31 +232,16 @@ static void findAndReplaceVectors(Module &M) {
       }
 
       // Note: we want to do G.replaceAllUsesWith(NewGlobal);, but it assumes
-      // type equality
-      //  So instead we will use the visitor pattern
+      // type equality. Instead we will use the visitor pattern.
       Impl.GlobalMap[&G] = NewGlobal;
-      for (User *U : G.users()) {
-        // Note: The GEPS are stored as constExprs
-        // This step flattens them out to instructions
+      for (User *U : make_early_inc_range(G.users())) {
         if (isa<ConstantExpr>(U) && isa<Operator>(U)) {
           ConstantExpr *CE = cast<ConstantExpr>(U);
           convertUsersOfConstantsToInstructions(CE,
                                                 /*RestrictToFunc=*/nullptr,
                                                 /*RemoveDeadConstants=*/false,
                                                 /*IncludeSelf=*/true);
         }
-      }
-      // Uses should have grown
-      std::vector<User *> UsersToProcess;
-      // Collect all users first
-      // work around so I can delete
-      // in a loop body
-      for (User *U : G.users()) {
-        UsersToProcess.push_back(U);
-      }
-
-      // Now process each user
-      for (User *U : UsersToProcess) {
         if (isa<Instruction>(U)) {
           Instruction *Inst = cast<Instruction>(U);
           Function *F = Inst->getFunction();
@@ -294,4 +281,4 @@ INITIALIZE_PASS_END(DXILDataScalarizationLegacy, DEBUG_TYPE,
 
 ModulePass *llvm::createDXILDataScalarizationLegacyPass() {
   return new DXILDataScalarizationLegacy();
-}
+}

llvm/test/CodeGen/DirectX/llc-pipeline.ll

@@ -8,6 +8,7 @@
 ; CHECK-NEXT: Target Transform Information
 ; CHECK-NEXT: ModulePass Manager
 ; CHECK-NEXT:   DXIL Intrinsic Expansion
+; CHECK-NEXT:   DXIL Data Scalarization
 ; CHECK-NEXT:   FunctionPass Manager
 ; CHECK-NEXT:     Dominator Tree Construction
 ; CHECK-NEXT:     Scalarize vector operations

llvm/test/CodeGen/DirectX/scalar-load.ll

@@ -2,17 +2,18 @@
 ; RUN: llc %s -mtriple=dxil-pc-shadermodel6.3-library --filetype=asm -o - | FileCheck %s
 @"arrayofVecData" = local_unnamed_addr addrspace(3) global [2 x <3 x float>] zeroinitializer, align 16
 @"vecData" = external addrspace(3) global <4 x i32>, align 4
-;@staticArrayOfVecData = internal global [3 x <4 x i32>] zeroinitializer, align 4
 @staticArrayOfVecData = internal global [3 x <4 x i32>] [<4 x i32> <i32 1, i32 2, i32 3, i32 4>, <4 x i32> <i32 5, i32 6, i32 7, i32 8>, <4 x i32> <i32 9, i32 10, i32 11, i32 12>], align 4
-
+@staticArray = internal global [4 x i32] [i32 1, i32 2, i32 3, i32 4], align 4
 
 ; CHECK: @arrayofVecData.scalarized = local_unnamed_addr addrspace(3) global [2 x [3 x float]] zeroinitializer, align 16
 ; CHECK: @vecData.scalarized = external addrspace(3) global [4 x i32], align 4
 ; CHECK: @staticArrayOfVecData.scalarized = internal global [3 x [4 x i32]] {{\[}}[4 x i32] [i32 1, i32 2, i32 3, i32 4], [4 x i32] [i32 5, i32 6, i32 7, i32 8], [4 x i32] [i32 9, i32 10, i32 11, i32 12]], align 4
+; Check @staticArray
 
 ; CHECK-NOT: @arrayofVecData
 ; CHECK-NOT: @vecData
 ; CHECK-NOT: @staticArrayOfVecData
+; CHECK-NOT: @staticArray.scalarized
 
 ; CHECK-LABEL: load_array_vec_test
 define <4 x i32> @load_array_vec_test() {

llvm/test/CodeGen/DirectX/scalar-store.ll

@@ -24,4 +24,4 @@ define void @store_vec_test(<4 x i32> %inputVec) {
   ; CHECK-NOT: store i32 %inputVec.{{.*}}, ptr addrspace(3)
   store <4 x i32> %inputVec, <4 x i32> addrspace(3)* @"vecData", align 4
   ret void
-}
+}