Don't scalarize groupshared vector arrays

- Also fix some bugs handling addrspacecast exposed by making this change

lib/Transforms/Scalar/ScalarReplAggregatesHLSL.cpp

@@ -4619,6 +4619,12 @@ void SROA_Parameter_HLSL::flattenGlobal(GlobalVariable *GV) {
 
     // Flat Global vector if no dynamic vector indexing.
     bool bFlatVector = !hasDynamicVectorIndexing(EltGV);
+
+    // Disable scalarization of groupshared vector arrays
+    if (GV->getType()->getAddressSpace() == DXIL::kTGSMAddrSpace &&
+        Ty->isArrayTy())
+      bFlatVector = false;
+
     std::vector<Value *> Elts;
     bool SROAed = SROA_Helper::DoScalarReplacement(
         EltGV, Elts, Builder, bFlatVector,
@@ -6636,6 +6642,8 @@ private:
   void ReplaceVectorWithArray(Value *Vec, Value *Array);
   void ReplaceVectorArrayWithArray(Value *VecArray, Value *Array);
   void ReplaceStaticIndexingOnVector(Value *V);
+  void ReplaceAddrSpaceCast(ConstantExpr *CE,
+                            Value *A, IRBuilder<> &Builder);
 };
 
 void DynamicIndexingVectorToArray::applyOptions(PassOptions O) {
@@ -6738,13 +6746,23 @@ void DynamicIndexingVectorToArray::ReplaceVecGEP(Value *GEP, ArrayRef<Value *> i
   }
 }
 
+void DynamicIndexingVectorToArray::ReplaceAddrSpaceCast(ConstantExpr *CE,
+                                              Value *A, IRBuilder<> &Builder) {
+  // create new AddrSpaceCast.
+  Value *NewAddrSpaceCast = Builder.CreateAddrSpaceCast(
+    A,
+    PointerType::get(A->getType()->getPointerElementType(),
+                      CE->getType()->getPointerAddressSpace()));
+  ReplaceVectorWithArray(CE, NewAddrSpaceCast);
+}
+
 void DynamicIndexingVectorToArray::ReplaceVectorWithArray(Value *Vec, Value *A) {
   unsigned size = Vec->getType()->getPointerElementType()->getVectorNumElements();
   for (auto U = Vec->user_begin(); U != Vec->user_end();) {
     User *User = (*U++);
 
     // GlobalVariable user.
-    if (isa<ConstantExpr>(User)) {
+    if (ConstantExpr * CE = dyn_cast<ConstantExpr>(User)) {
       if (User->user_empty())
         continue;
       if (GEPOperator *GEP = dyn_cast<GEPOperator>(User)) {
@@ -6752,7 +6770,12 @@ void DynamicIndexingVectorToArray::ReplaceVectorWithArray(Value *Vec, Value *A)
         SmallVector<Value *, 4> idxList(GEP->idx_begin(), GEP->idx_end());
         ReplaceVecGEP(GEP, idxList, A, Builder);
         continue;
+      } else if (CE->getOpcode() == Instruction::AddrSpaceCast) {
+        IRBuilder<> Builder(Vec->getContext());
+        ReplaceAddrSpaceCast(CE, A, Builder);
+        continue;
       }
+      DXASSERT(0, "not implemented yet");
     }
     // Instrution user.
     Instruction *UserInst = cast<Instruction>(User);
@@ -6961,7 +6984,7 @@ void ReplaceMultiDimGEP(User *GEP, Value *OneDim, IRBuilder<> &Builder) {
 void MultiDimArrayToOneDimArray::lowerUseWithNewValue(Value *MultiDim, Value *OneDim) {
   LLVMContext &Context = MultiDim->getContext();
   // All users should be element type.
-  // Replace users of AI.
+  // Replace users of AI or GV.
   for (auto it = MultiDim->user_begin(); it != MultiDim->user_end();) {
     User *U = *(it++);
     if (U->user_empty())
@@ -6970,21 +6993,29 @@ void MultiDimArrayToOneDimArray::lowerUseWithNewValue(Value *MultiDim, Value *On
       BCI->setOperand(0, OneDim);
       continue;
     }
-    // Must be GEP.
-    GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(U);
 
-    if (!GEP) {
-      DXASSERT_NOMSG(isa<GEPOperator>(U));
-      // NewGEP must be GEPOperator too.
-      // No instruction will be build.
+    if (ConstantExpr *CE = dyn_cast<ConstantExpr>(U)) {
       IRBuilder<> Builder(Context);
-      ReplaceMultiDimGEP(U, OneDim, Builder);
-    } else {
+      if (GEPOperator *GEP = dyn_cast<GEPOperator>(U)) {
+        // NewGEP must be GEPOperator too.
+        // No instruction will be build.
+        ReplaceMultiDimGEP(U, OneDim, Builder);
+      } else if (CE->getOpcode() == Instruction::AddrSpaceCast) {
+        Value *NewAddrSpaceCast = Builder.CreateAddrSpaceCast(
+          OneDim,
+          PointerType::get(OneDim->getType()->getPointerElementType(),
+                           CE->getType()->getPointerAddressSpace()));
+        lowerUseWithNewValue(CE, NewAddrSpaceCast);
+      } else {
+        DXASSERT(0, "not implemented");
+      }
+    } else if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(U)) {
       IRBuilder<> Builder(GEP);
       ReplaceMultiDimGEP(U, OneDim, Builder);
-    }
-    if (GEP)
       GEP->eraseFromParent();
+    } else {
+      DXASSERT(0, "not implemented");
+    }
   }
 }
 

tools/clang/test/CodeGenHLSL/quick-test/addrspacecast.hlsl

@@ -0,0 +1,28 @@
+// RUN: %dxc -E main -T cs_6_0 %s | FileCheck %s
+
+// Make sure generate addrspacecast.
+// CHECK: addrspacecast ([6 x float] addrspace(3)*
+
+struct ST
+{
+  float3 a; // center
+  float3 b; // half extents
+
+  void func(float3 x, float3 y)
+  {
+    a = x + y;
+    b = x * y;
+  }
+};
+
+groupshared ST myST[2];
+StructuredBuffer<ST> buf0;
+float3 a;
+float3 b;
+RWBuffer<float3> buf1;
+[numthreads(8,8,1)]
+void main() {
+  myST[0] = buf0[0];
+  myST[0].func(a, b);
+  buf1[0] = myST[0].b;
+}

tools/clang/test/CodeGenHLSL/share_mem_dbg.hlsl

@@ -12,12 +12,9 @@
 // CHECK: !dx.source.mainFileName
 // CHECK: !dx.source.args
 
-// CHECK: DIGlobalVariable(name: "dataC.1.0"
-// CHECK: DIDerivedType(tag: DW_TAG_member, name: ".1.0"
-// CHECK: DIGlobalVariable(name: "dataC.1.1"
-// CHECK: DIDerivedType(tag: DW_TAG_member, name: ".1.1"
-// CHECK: DIGlobalVariable(name: "dataC.0
-// CHECK: DIDerivedType(tag: DW_TAG_member, name: ".0"
+// CHECK: DIGlobalVariable(name: "dataC"
+// CHECK: DIDerivedType(tag: DW_TAG_member, name: "d"
+// CHECK: DIDerivedType(tag: DW_TAG_member, name: "b"
 
 // Make sure source info contents exist.
 // CHECK: share_mem_dbg.hlsl", !"// RUN: %dxc