Fix non-const AddrSpaceCast, and constant cleanup in HLMatrixLower

lib/HLSL/HLMatrixLowerPass.cpp

@@ -510,18 +510,29 @@ void HLMatrixLowerPass::replaceAllVariableUses(
       
       // Discard the GEP
       DXASSERT_NOMSG(GEP->use_empty());
-      Use.set(UndefValue::get(Use->getType()));
-      if (GetElementPtrInst *GEPInst = dyn_cast<GetElementPtrInst>(GEP))
+      if (GetElementPtrInst *GEPInst = dyn_cast<GetElementPtrInst>(GEP)) {
+        Use.set(UndefValue::get(Use->getType()));
         addToDeadInsts(GEPInst);
+      } else {
+        // constant GEP
+        cast<Constant>(GEP)->destroyConstant();
+      }
       continue;
     }
 
     if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Use.getUser())) {
       DXASSERT(CE->getOpcode() == Instruction::AddrSpaceCast,
                "Unexpected constant user");
-      if (!CE->user_empty())
-        replaceAllVariableUses(GEPIdxStack, CE, LoweredPtr);
+      replaceAllVariableUses(GEPIdxStack, CE, LoweredPtr);
+      DXASSERT_NOMSG(CE->use_empty());
+      CE->destroyConstant();
+      continue;
+    }
+
+    if (AddrSpaceCastInst *CI = dyn_cast<AddrSpaceCastInst>(Use.getUser())) {
+      replaceAllVariableUses(GEPIdxStack, CI, LoweredPtr);
       Use.set(UndefValue::get(Use->getType()));
+      addToDeadInsts(CI);
       continue;
     }
 

lib/Transforms/Scalar/ScalarReplAggregatesHLSL.cpp

@@ -113,7 +113,7 @@ private:
 
   void RewriteForConstExpr(ConstantExpr *user, IRBuilder<> &Builder);
   void RewriteForGEP(GEPOperator *GEP, IRBuilder<> &Builder);
-  void RewriteForAddrSpaceCast(ConstantExpr *user, IRBuilder<> &Builder);
+  void RewriteForAddrSpaceCast(Value *user, IRBuilder<> &Builder);
   void RewriteForLoad(LoadInst *loadInst);
   void RewriteForStore(StoreInst *storeInst);
   void RewriteMemIntrin(MemIntrinsic *MI, Value *OldV);
@@ -3303,17 +3303,17 @@ void SROA_Helper::RewriteCall(CallInst *CI) {
   }
 }
 
-/// RewriteForConstExpr - Rewrite the GEP which is ConstantExpr.
-void SROA_Helper::RewriteForAddrSpaceCast(ConstantExpr *CE,
+/// RewriteForAddrSpaceCast - Rewrite the AddrSpaceCast, either ConstExpr or Inst.
+void SROA_Helper::RewriteForAddrSpaceCast(Value *CE,
                                           IRBuilder<> &Builder) {
   SmallVector<Value *, 8> NewCasts;
   // create new AddrSpaceCast.
   for (unsigned i = 0, e = NewElts.size(); i != e; ++i) {
-    Value *NewGEP = Builder.CreateAddrSpaceCast(
+    Value *NewCast = Builder.CreateAddrSpaceCast(
         NewElts[i],
         PointerType::get(NewElts[i]->getType()->getPointerElementType(),
                          CE->getType()->getPointerAddressSpace()));
-    NewCasts.emplace_back(NewGEP);
+    NewCasts.emplace_back(NewCast);
   }
   SROA_Helper helper(CE, NewCasts, DeadInsts, typeSys, DL);
   helper.RewriteForScalarRepl(CE, Builder);
@@ -3379,7 +3379,9 @@ void SROA_Helper::RewriteForScalarRepl(Value *V, IRBuilder<> &Builder) {
       RewriteCall(CI);
     else if (BitCastInst *BCI = dyn_cast<BitCastInst>(User))
       RewriteBitCast(BCI);
-    else {
+    else if (AddrSpaceCastInst *CI = dyn_cast<AddrSpaceCastInst>(User)) {
+      RewriteForAddrSpaceCast(CI, Builder);
+    } else {
       assert(0 && "not support.");
     }
   }

tools/clang/test/CodeGenHLSL/quick-test/groupshared-member-matrix-subscript-col.hlsl

@@ -35,15 +35,18 @@ class Obj {
 };
 
 RWByteAddressBuffer RWBuf;
-groupshared Obj obj;
+groupshared Obj obj[2];
+int objIndex;
 
 [numthreads(4, 1, 1)]
 void main(uint3 groupThreadID: SV_GroupThreadID) {
   if (groupThreadID.x == 0) {
-    obj.set();
+    obj[0].set();
+    obj[objIndex].set();
   }
   GroupMemoryBarrierWithGroupSync();
-  float4 row = obj.mat[groupThreadID.x];
+  float4 row = obj[0].mat[groupThreadID.x];
+  row *= obj[objIndex].mat[groupThreadID.x];
   uint addr = groupThreadID.x * 4;
   RWBuf.Store4(addr, uint4(asuint(row.x), asuint(row.y), asuint(row.z), asuint(row.w)));
 }