Блог

В учебнике приведен пример использования packed_accessor64 для эффективного доступа к элементам тензора следующим образом. Однако я все еще не знаю, как использовать packed_accessor64. Может ли кто-нибудь привести мне более конкретный пример? Спасибо.

 __global__ void packed_accessor_kernel(
    PackedTensorAccessor64<float, 2> foo,
    float* trace) {
  int i=threadIdx.x
  gpuAtomicAdd(trace, foo[i][i])
}
 
torch::Tensor foo = torch::rand({12, 12});
 
// assert foo is 2-dimensional and holds floats.
auto foo_a = foo.packed_accessor64<float,2>();
float trace = 0;
 
packed_accessor_kernel<<<1, 12>>>(foo_a, amp;trace);