[NVPTX] fix a crash bug in NVPTXFavorNonGenericAddrSpaces

[jingyue]

Summary

We used to assume V->RAUW only modifies the operand list of V's user.
However, if V and V's user are Constants, RAUW may replace and invalidate V's
user entirely.

This patch fixes the above issue by letting the caller replace the
operand instead of calling RAUW on Constants.

[jholewinski]

Looks reasonable to me. Thanks!

[broune]

Looks good to me.

[jingyue]

a minor bug

lib/Target/NVPTX/NVPTXFavorNonGenericAddrSpaces.cpp

private:
   /// outermost level. For example, this function transforms
   ///   bitcast(gep(gep(addrspacecast(X))))
   /// to
   ///   addrspacecast(bitcast(gep(gep(X)))).
   ///
   /// This reordering exposes to optimizeMemoryInstruction more
   /// optimization opportunities on loads and stores.
   ///
-  /// Returns true if this function succesfully hoists an eliminable
-  /// addrspacecast or V is already such an addrspacecast.
-  /// Transforms "gep (addrspacecast X), indices" into "addrspacecast (gep X,
-  /// indices)".
-  bool hoistAddrSpaceCastFrom(Value *V, int Depth = 0);
+  /// If this function succesfully hoists an eliminable addrspacecast or V is
+  /// already such an addrspacecast, it returns the transformed value (which is
+  /// guaranteed to be an addrspacecast); otherwise, it returns nullptr.
+  Value *hoistAddrSpaceCastFrom(Value *V, int Depth = 0);
   /// Helper function for GEPs.
-  bool hoistAddrSpaceCastFromGEP(GEPOperator *GEP, int Depth);
+  Value *hoistAddrSpaceCastFromGEP(GEPOperator *GEP, int Depth);
   /// Helper function for bitcasts.
-  bool hoistAddrSpaceCastFromBitCast(BitCastOperator *BC, int Depth);
+  Value *hoistAddrSpaceCastFromBitCast(BitCastOperator *BC, int Depth);
 };
 }
 
 char NVPTXFavorNonGenericAddrSpaces::ID = 0;
 
 namespace llvm {
 void initializeNVPTXFavorNonGenericAddrSpacesPass(PassRegistry &);
 }
if (SrcTy->getElementType() != DestTy->getElementType())
     return false;
 
   // Checks whether the addrspacecast is from a non-generic address space to the
   // generic address space.
   return (SrcTy->getAddressSpace() != AddressSpace::ADDRESS_SPACE_GENERIC &&
           DestTy->getAddressSpace() == AddressSpace::ADDRESS_SPACE_GENERIC);
 }
 
-bool NVPTXFavorNonGenericAddrSpaces::hoistAddrSpaceCastFromGEP(GEPOperator *GEP,
-                                                               int Depth) {
-  if (!hoistAddrSpaceCastFrom(GEP->getPointerOperand(), Depth + 1))
-    return false;
-
-  // That hoistAddrSpaceCastFrom succeeds implies GEP's pointer operand is now
-  // an eliminable addrspacecast.
-  assert(isEliminableAddrSpaceCast(GEP->getPointerOperand()));
-  Operator *Cast = cast<Operator>(GEP->getPointerOperand());
+Value *NVPTXFavorNonGenericAddrSpaces::hoistAddrSpaceCastFromGEP(
+    GEPOperator *GEP, int Depth) {
+  Value *NewOperand =
+      hoistAddrSpaceCastFrom(GEP->getPointerOperand(), Depth + 1);
+  if (NewOperand == nullptr)
+    return nullptr;
+
+  // hoistAddrSpaceCastFrom returns an eliminable addrspacecast or nullptr.
+  assert(isEliminableAddrSpaceCast(NewOperand));
+  Operator *Cast = cast<Operator>(NewOperand);
 
   SmallVector<Value *, 8> Indices(GEP->idx_begin(), GEP->idx_end());
+  Value *NewASC;
   if (Instruction *GEPI = dyn_cast<Instruction>(GEP)) {
     // GEP = gep (addrspacecast X), indices
     // =>
     // NewGEP = gep X, indices
     // NewASC = addrspacecast NewGEP
     GetElementPtrInst *NewGEP = GetElementPtrInst::Create(
         GEP->getSourceElementType(), Cast->getOperand(0), Indices,
         "", GEPI);
     NewGEP->setIsInBounds(GEP->isInBounds());
-    Value *NewASC = new AddrSpaceCastInst(NewGEP, GEP->getType(), "", GEPI);
+    NewASC = new AddrSpaceCastInst(NewGEP, GEP->getType(), "", GEPI);
     NewASC->takeName(GEP);
+    // Without RAUWing GEP, the compiler would visit GEP again and emit
+    // redundant instructions. This is exercised in test @rauw in
+    // access-non-generic.ll.
     GEP->replaceAllUsesWith(NewASC);
   } else {
     // GEP is a constant expression.
     Constant *NewGEP = ConstantExpr::getGetElementPtr(
         GEP->getSourceElementType(), cast<Constant>(Cast->getOperand(0)),
         Indices, GEP->isInBounds());
-    GEP->replaceAllUsesWith(
-        ConstantExpr::getAddrSpaceCast(NewGEP, GEP->getType()));
+    NewASC = ConstantExpr::getAddrSpaceCast(NewGEP, GEP->getType());
   }
-
-  return true;
+  return NewASC;
 }
 
-bool NVPTXFavorNonGenericAddrSpaces::hoistAddrSpaceCastFromBitCast(
+Value *NVPTXFavorNonGenericAddrSpaces::hoistAddrSpaceCastFromBitCast(
     BitCastOperator *BC, int Depth) {
-  if (!hoistAddrSpaceCastFrom(BC->getOperand(0), Depth + 1))
-    return false;
-
-  // That hoistAddrSpaceCastFrom succeeds implies BC's source operand is now
-  // an eliminable addrspacecast.
-  assert(isEliminableAddrSpaceCast(BC->getOperand(0)));
-  Operator *Cast = cast<Operator>(BC->getOperand(0));
+  Value *NewOperand = hoistAddrSpaceCastFrom(BC->getOperand(0), Depth + 1);
+  if (NewOperand == nullptr)
+    return nullptr;
+
+  // hoistAddrSpaceCastFrom returns an eliminable addrspacecast or nullptr.
+  assert(isEliminableAddrSpaceCast(NewOperand));
+  Operator *Cast = cast<Operator>(NewOperand);
 
   // Cast  = addrspacecast Src
   // BC    = bitcast Cast
   //   =>
   // Cast' = bitcast Src
   // BC'   = addrspacecast Cast'
   Value *Src = Cast->getOperand(0);
   Type *TypeOfNewCast =
       PointerType::get(BC->getType()->getPointerElementType(),
                        Src->getType()->getPointerAddressSpace());
+  Value *NewBC;
   if (BitCastInst *BCI = dyn_cast<BitCastInst>(BC)) {
     Value *NewCast = new BitCastInst(Src, TypeOfNewCast, "", BCI);
-    Value *NewBC = new AddrSpaceCastInst(NewCast, BC->getType(), "", BCI);
+    NewBC = new AddrSpaceCastInst(NewCast, BC->getType(), "", BCI);
     NewBC->takeName(BC);
+    // Without RAUWing BC, the compiler would visit BC again and emit
+    // redundant instructions. This is exercised in test @rauw in
+    // access-non-generic.ll.
     BC->replaceAllUsesWith(NewBC);
   } else {
     // BC is a constant expression.
     Constant *NewCast =
         ConstantExpr::getBitCast(cast<Constant>(Src), TypeOfNewCast);
-    Constant *NewBC = ConstantExpr::getAddrSpaceCast(NewCast, BC->getType());
-    BC->replaceAllUsesWith(NewBC);
+    NewBC = ConstantExpr::getAddrSpaceCast(NewCast, BC->getType());
   }
-  return true;
+  return NewBC;
 }
 
-bool NVPTXFavorNonGenericAddrSpaces::hoistAddrSpaceCastFrom(Value *V,
+Value *NVPTXFavorNonGenericAddrSpaces::hoistAddrSpaceCastFrom(Value *V,
                                                               int Depth) {
-  // Returns true if V is already an eliminable addrspacecast.
+  // Returns V if V is already an eliminable addrspacecast.
   if (isEliminableAddrSpaceCast(V))
-    return true;
+    return V;
 
   // Limit the depth to prevent this recursive function from running too long.
   const int MaxDepth = 20;
   if (Depth >= MaxDepth)
-    return false;
+    return nullptr;
 
   // If V is a GEP or bitcast, hoist the addrspacecast if any from its pointer
   // operand. This enables optimizeMemoryInstruction to shortcut addrspacecasts
   // that are not directly used by the load/store.
   if (GEPOperator *GEP = dyn_cast<GEPOperator>(V))
     return hoistAddrSpaceCastFromGEP(GEP, Depth);
 
   if (BitCastOperator *BC = dyn_cast<BitCastOperator>(V))
     return hoistAddrSpaceCastFromBitCast(BC, Depth);
 
-  return false;
+  return nullptr;
 }
 
 bool NVPTXFavorNonGenericAddrSpaces::optimizeMemoryInstruction(Instruction *MI,
                                                                unsigned Idx) {
-  if (hoistAddrSpaceCastFrom(MI->getOperand(Idx))) {
-    // load/store (addrspacecast X) => load/store X if shortcutting the
-    // addrspacecast is valid and can improve performance.
-    //
-    // e.g.,
-    // %1 = addrspacecast float addrspace(3)* %0 to float*
-    // %2 = load float* %1
-    // ->
-    // %2 = load float addrspace(3)* %0
-    //
-    // Note: the addrspacecast can also be a constant expression.
-    assert(isEliminableAddrSpaceCast(MI->getOperand(Idx)));
-    Operator *ASC = dyn_cast<Operator>(MI->getOperand(Idx));
-    MI->setOperand(Idx, ASC->getOperand(0));
-    return true;
-  }
-  return false;
-}
+  Value *NewOperand = hoistAddrSpaceCastFrom(MI->getOperand(Idx));
+  if (NewOperand == nullptr)
+    return false;
+
+  assert(isEliminableAddrSpaceCast(NewOperand));
+  Operator *ASC = dyn_cast<Operator>(NewOperand);
 
 bool NVPTXFavorNonGenericAddrSpaces::runOnFunction(Function &F) {
   if (DisableFavorNonGeneric)
     return false;
 
   bool Changed = false;
   for (Function::iterator B = F.begin(), BE = F.end(); B != BE; ++B) {
     for (BasicBlock::iterator I = B->begin(), IE = B->end(); I != IE; ++I) {

test/CodeGen/NVPTX/access-non-generic.ll

; IR-NEXT: getelementptr float, float addrspace(1)* {{%[^,]+}}, i32 %j
   %4 = bitcast float* %3 to i32*
 ; IR-NEXT: bitcast float addrspace(1)* {{%[^ ]+}} to i32 addrspace(1)*
   %5 = load i32, i32* %4
 ; IR-NEXT: load i32, i32 addrspace(1)* {{%.+}}
 ; PTX-LABEL: ld.global
   ret i32 %5
 }
 
+define void @nested_const_expr() {
+; PTX-LABEL: nested_const_expr(
+  ; store 1 to bitcast(gep(addrspacecast(array), 0, 1))
+  store i32 1, i32* bitcast (float* getelementptr ([10 x float], [10 x float]* addrspacecast ([10 x float] addrspace(3)* @array to [10 x float]*), i64 0, i64 1) to i32*), align 4
+; PTX: mov.u32 %r1, 1;
+; PTX-NEXT: st.shared.u32 [array+4], %r1;
+  ret void
+}
+
+define void @rauw(float addrspace(1)* %input) {
+  %generic_input = addrspacecast float addrspace(1)* %input to float*
+  %addr = getelementptr float, float* %generic_input, i64 10
+  %v = load float, float* %addr
+  store float %v, float* %addr
+  ret void
+; IR-LABEL: @rauw(
+; IR-NEXT: %1 = getelementptr float, float addrspace(1)* %input, i64 10
+; IR-NEXT: %v = load float, float addrspace(1)* %1
+; IR-NEXT: store float %v, float addrspace(1)* %1
+; IR-NEXT: ret void
+}
+
 declare void @llvm.cuda.syncthreads() #3
 
 attributes #3 = { noduplicate nounwind }