支持独立线程调度的CUDA自旋锁实现?

Question

我想回顾一下在CUDA上实现一个简单的自旋锁的情况，现在独立线程调度(ITS)已经引入了一段时间了。

我的代码是这样的:

// nvcc main.cu -arch=sm_75
#include <cstdio>
#include <iostream>
#include <vector>
#include "cuda.h"
constexpr int kN = 21;
using Ptr = uint8_t*;
struct DynamicNode {
int32_t lock = 0;
int32_t n = 0;
Ptr ptr = nullptr;
};
__global__ void func0(DynamicNode* base) {
for (int i = 0; i < kN; ++i) {
DynamicNode* dn = base + i;
atomicAdd(&(dn->n), 1);
// entering the critical section
auto* lock = &(dn->lock);
while (atomicExch(lock, 1) == 1) {
}
__threadfence();
// Use a condition to artificially boost the complexity
// of loop unrolling for the compiler
if (dn->ptr == nullptr) {
dn->ptr = reinterpret_cast<Ptr>(0xf0);
}
// leaving the critical section
atomicExch(lock, 0);
__threadfence();
}
}
int main() {
DynamicNode* dev_root = nullptr;
constexpr int kRootSize = sizeof(DynamicNode) * kN;
cudaMalloc((void**)&dev_root, kRootSize);
cudaMemset(dev_root, 0, kRootSize);
func0<<<1, kN>>>(dev_root);
cudaDeviceSynchronize();
std::vector<int32_t> host_root(kRootSize / sizeof(int32_t), 0);
cudaMemcpy(host_root.data(), dev_root, kRootSize, cudaMemcpyDeviceToHost);
cudaFree((void*)dev_root);
const auto* base = reinterpret_cast<const DynamicNode*>(host_root.data());
int sum = 0;
for (int i = 0; i < kN; ++i) {
auto& dn = base[i];
std::cout << "i=" << i << " len=" << dn.n << std::endl;
sum += dn.n;
}
std::cout << "sum=" << sum << " expected=" << kN * kN << std::endl;
return 0;
}

可以看到，在func0中实现了一个朴素自旋锁。虽然我知道这将导致旧的arch死锁(例如https://forums.developer.nvidia.com/t/atomic-locks/25522/2)，但如果我用nvcc main.cu -arch=sm_75编译代码，它实际上可以无限期地运行而不阻塞。

然而，我注意到的是每个DynamicNode中的n完全是垃圾。以下是GeForce RTX 2060 (laptop)上的输出，我可以确定地再现它:

i=0 len=21
i=1 len=230
i=2 len=19
i=3 len=18
i=4 len=17
i=5 len=16
i=6 len=15
i=7 len=14
i=8 len=13
i=9 len=12
i=10 len=11
i=11 len=10
i=12 len=9
i=13 len=8
i=14 len=7
i=15 len=6
i=16 len=5
i=17 len=4
i=18 len=3
i=19 len=2
i=20 len=1
sum=441 expected=441

理想情况下，所有DynamicNode的长度应该是kN。我也尝试过更大的kN(*)，它总是只有sum是正确的。

我是否误解了ITS?ITS真的可以保证这样的锁实现吗?如果不是，我在这里错过了什么?

(*)对于较小的kN,nvcc可能实际上展开循环，从我在PTX中看到的。当展开循环时，我从来没有发现任何问题。

更新02/02/2021

---

我应该澄清一下，我在CUDA 11.1上测试了这个。根据@robert-crovella的说法，升级到11.2可以解决这个问题。

更新02/03/2021

---

我测试了CUDA 11.2驱动程序，它仍然没有完全解决更大的kN的问题:

<表类>kN CUDA11.1 11.2tbody><<tr>21N对128NN

这似乎是编译器中的某种代码生成缺陷。解决方案似乎是更新到CUDA 11.2(或更新，可能在未来)。