int managed_counter = 0;
#[inline(always)] int atomic_increment(int* addr) {
return atomicAdd(addr, 1);
}
void increment_kernel(int* counter, int n) {
int i = block_idx * block_dim + thread_idx;
if (i < n) {
atomic_increment(counter);
}
}
void scale_kernel(float* data, float scale, int n) {
int i = block_idx * block_dim + thread_idx;
if (i < n) {
data[i] *= scale;
}
}
int main(void) {
const int N = 1 << 16;
float* managed_data = nullptr;
cudaMallocManaged((void**)&managed_data, N * sizeof(float));
for (int i = 0; i < N; i++) {
managed_data[i] = (float)i;
}
float* pinned_data = nullptr;
cudaMallocHost((void**)&pinned_data, N * sizeof(float));
dim3 grid(N / 256);
dim3 block(256);
{ let _kernel = modules.get_function("scale_kernel"); unsafe { let _ = launch!( _kernel<<<grid as grid_size, block as block_size, 0 as usize, default>>>(managed_data, 2.0f, N) ); } };
{ let _kernel = modules.get_function("increment_kernel"); unsafe { let _ = launch!( _kernel<<<grid as grid_size, block as block_size, 0 as usize, default>>>(&managed_counter, N) ); } };
cudaDeviceSynchronize();
printf("counter = %d, data[0] = %f\n", managed_counter, managed_data[0]);
cudaMemcpy(pinned_data, managed_data, N * sizeof(float), cudaMemcpyDeviceToHost);
cudaFree(managed_data);
cudaFreeHost(pinned_data);
return 0;
}