CUDA共享内存封装在模板类中,指向相同的内存

CUDA shared memory wrapped in templated class, points to same memory

本文关键字:内存 封装 共享 CUDA      更新时间:2023-10-16

我正试图在模板类中的CUDA内核中分配共享内存

template<typename T, int Size>
struct SharedArray {
    __device__ T* operator()(){
        __shared__ T x[Size];
        return x;
    }
};

只要不检索两次具有相同类型和大小的共享内存,这就可以工作。但当我试图获得两次相同类型和大小的共享内存时,第二次共享内存指向第一次:

__global__
void test() {
    // Shared array
    SharedArray<int, 5> sharedArray;
    int* x0 = sharedArray();
    int* y0 = sharedArray();
    x0[0] = 1;
    y0[0] = 0;
    printf("%i %inn", x0[0], y0[0]);
    // Prints:
    // 0 0
}

一种解决方案是在每次调用共享内存类时添加一个id,如:

template<int ID, typename T, int Size>
struct StaticSharedArrayWithID {
    __device__ static T* shared(){
        __shared__ T x[Size];
        return x;
    }   
};

但我必须提供一些计数器,它提供了一个非常丑陋的用户界面:

__global__
void test() {
    int& x1 = StaticSharedArrayWithID<__COUNTER__, int, 5>::shared();
    int& y1 = StaticSharedArrayWithID<__COUNTER__, int, 5>::shared();
    x1[0] = 1;
    y1[0] = 0;
    printf("%i %inn", x1[0], y1[0]);
    // Prints:
    // 1 0
}

有没有人想去掉用户界面中的__COUNTER__宏?当它被隐藏的时候是可以的。

原因是默认情况下__shared__变量为static。同一函数的同一实例引用同一变量。这种行为的最初原因是编译器无法推断何时可以回收内存。拥有一个变量static使它和内核一样长。

副作用是,如果同一个函数在程序中的两个位置被调用两次,则会得到相同的结果。事实上,当多个CUDA线程在同一位置调用您的函数时,这正是您所期望的,不是吗?

没有一种干净的方法可以动态分配共享内存。在我的项目中,我通过自己的共享内存内存管理器(前面是难看的指针运算,小心!):

typedef unsigned char byte;
/*
  Simple shared memory manager.
    With any luck if invoked with constant parameters this will not take up any register whatsoever
    Must be called uniformly by whole block which is going to use these
    sSize - amount of preallocated memory
*/
template <size_t sSize>
class SharedMemoryManager {
private:
    byte* shArray;
    byte* head;
public:
    __device__ SharedMemoryManager() {
        __shared__ byte arr[sSize];
        shArray=arr;
        head=arr;
    }
    __device__  void reset() {
        head=shArray;
    }
    __device__  byte* getHead() {return head;}
    __device__  void setHead(byte* newHead) {head=newHead;}
    template <typename T>
    __device__  T* alloc(size_t count) {
      size_t addr = head;
      size_t alignment = __alignof(T); //assuming alignment is power of 2
      addr = ((addr-1) | (alignment-1)) +1; //round up to match the alignment requirement
      head = (byte*)(addr);
      T* var = (T*)(head);
      head+=sizeof(T)*size;
      return allocAt<T>(head,count);
    }
template <typename T>
    __device__  T& alloc() {
      return *alloc<T>(1);
    }
};

当您知道可以回收共享内存时,可以使用getHead/setHead来回收共享内存,但只能以堆栈方式回收。

当CUDA不是您的目标时,这种方法应该很容易在非共享内存上抽象。

然后你应该能够写:

__global__
 void test() {
   SharedMemoryManager shMem<1024>();
   int& xValue = shMem.alloc<int>();
   int& yValue = shMem.alloc<int>();
   int* xArray = shMem.alloc<int>(5);
   int* yArray = shMem.alloc<int>(5);
   xArray[0] = 1;
   yArray[0] = 0;
   printf("%i %inn", xArray[0], yArray[0]);
   __syncthreads();
   shMem.reset(); //memory reclaimed
   ... 
}