Fix a specialization bug and shMemPool refactor

Fix a bug in specialization where we should run peephole optimization
after one iteration of the specialization, otherwise something like
sizeof(Type) won't be folded into constant, which will block the
specialization of some generics.

Refactor the shMemPool type, now we can use the Type as generic
parameter to declare the AccelerateVectorCoopMat type.

Author: kaizhangNV

source/slang/slang-ir-specialize.cpp

@@ -1237,6 +1237,7 @@ struct SpecializationContext
             {
                 this->changed = true;
                 eliminateDeadCode(module->getModuleInst());
+                peepholeOptimizeGlobalScope(targetProgram, this->module);
                 applySparseConditionalConstantPropagationForGlobalScope(this->module, this->sink);
             }
 

source/standard-modules/neural/accelerate-vector-coopmat.slang

@@ -1011,8 +1011,6 @@ public struct AccelerateVectorCoopMat<T, ShMemPool : ISharedMemoryPool, int N, i
     public typealias Differential = AccelerateVectorCoopMat<T.Differential, ShMemPool, N, SubgroupSize>;
 
     public static const int Size = N;
-
-    public no_diff SPtr<uint4> sharedMemoryPtr;
     public no_diff ShMemPool shMemPool;
 
     private typealias DTypeMatC = half;

source/standard-modules/neural/shared-memory-pool.slang

@@ -14,17 +14,18 @@ internal typealias SPtr<T> = Ptr<T, Access::ReadWrite, AddressSpace::GroupShared
 
 internal interface ISharedMemoryPool
 {
-    static SPtr<uint4> getPointer();
+    internal static SPtr<uint4> getPointer();
 }
 
 public interface ISharedMemorySize
 {
     public static const uint Bytes;
 }
 
-public struct SharedMemoryPool<int Bytes> : ISharedMemoryPool
+public struct SharedMemoryPool<ShMemSize: ISharedMemorySize> : ISharedMemoryPool
 {
-    internal static groupshared uint4 data[Bytes / sizeof(uint4)];
+    public static const uint sizeInBytes = ShMemSize.Bytes;
+    internal static groupshared uint4 data[sizeInBytes / sizeof(uint4)];
     VISIBILITY_LEVEL static SPtr<uint4> getPointer()
     {
         return __getAddress(data[0]);

tests/neural/basic-coopmat-vector-test.slang

@@ -39,8 +39,8 @@ static const int BatchSize = 32;
 static const int SubgroupSize = 32;
 static const int workgroupCount = BatchSize / SubgroupSize;
 
-typealias ShMemSize = SharedMemorySize<ElementType, TargetEnum.CUDA, ExecutionMode.Training, SubgroupSize, BatchSize/SubgroupSize>;
-static const int ShMemSizeInBytes = ShMemSize.OfLayer1<InputSize, OutputSize>.Bytes;
+typealias ShMemSize = SharedMemorySize< ElementType, TargetEnum.CUDA, ExecutionMode.Training, SubgroupSize, BatchSize / SubgroupSize>;
+typealias ShMemSizeLayer1 = ShMemSize.OfLayer1<InputSize, OutputSize>;
 
 typealias SPtr<T> = Ptr<T, Access::ReadWrite, AddressSpace::GroupShared>;
 
@@ -73,21 +73,20 @@ OutputVector TestInlineVectorMatMul<InputVector, OutputVector>(
 // Basic test on MatMul without bias, this test covers both forward and backward pass
 void BasicTestWithoutBias(int tid, int resIndex)
 {
-    typealias ShMemPool = SharedMemoryPool<ShMemSizeInBytes>;
+    typealias ShMemPool = SharedMemoryPool<ShMemSizeLayer1>;
     typealias InVectorType = AccelerateVectorCoopMat<ElementType, ShMemPool, InputSize, SubgroupSize>;
     typealias OutVectorType = AccelerateVectorCoopMat<ElementType, ShMemPool, OutputSize, SubgroupSize>;
 
     ElementType[InputSize] inputData = { ElementType(1.0), ElementType(2.0), ElementType(3.0), ElementType(4.0) };
     InVectorType input = InVectorType(inputData);
 
-    input.sharedMemoryPtr = ShMemPool.getPointer();
-
     BufferStorage weightStorage = BufferStorage(parameters);
     BufferStorage dweightStorage = BufferStorage(dParameters);
     BufferStorage.Address weightAddress = 0;
 
     // Run the forward pass
     let outputVec = TestInlineVectorMatMul<InVectorType, OutVectorType>(input, weightStorage, weightAddress);
+
     // serialRead<16, half>(tid, __getAddress(shMem[0]));
     // serialRead<16, half>(tid, __getAddress(shMem[0]) + 32);
 
@@ -96,7 +95,6 @@ void BasicTestWithoutBias(int tid, int resIndex)
     bool isPassed = true;
     isPassed = isPassed && (outputVec[0] == 30.0 && outputVec[1] == 70.0);
 
-
     var weightDiffPair = DifferentialPtrPair<BufferStorage>(weightStorage, dweightStorage);
     let dRes = OutVectorType(1.0f);
     var dPair = diffPair(input);