Skip to content

roll kernel as grid stride loop #2852

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jerrymannil merged 1 commit into from
Dec 4, 2025
Merged

roll kernel as grid stride loop #2852

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 17 additions & 15 deletions aten/src/ATen/native/cuda/TensorTransformations.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,20 +90,20 @@ __global__ void roll_cuda_kernel(
int64_t size,
int64_t stride,
int64_t total_dims) {
int64_t linear_index = ((int64_t) blockIdx.x) * blockDim.x + threadIdx.x;
if (linear_index >= N) {
return;
}
// roll dim idx is the index of linear_index along the rolling dimension.
int64_t roll_dim_idx = linear_index % (stride * size) / stride;
// index into the source data to find appropriate value.
int64_t source_idx = 0;
if( roll_dim_idx >= (size - start) ) {
source_idx = linear_index - ((size - start) * stride);
} else {
source_idx = linear_index + (start * stride);
for (int64_t linear_index = ((int64_t) blockIdx.x) * blockDim.x + threadIdx.x;
linear_index < N; linear_index += blockDim.x*gridDim.x)
{
// roll dim idx is the index of linear_index along the rolling dimension.
int64_t roll_dim_idx = linear_index % (stride * size) / stride;
// index into the source data to find appropriate value.
int64_t source_idx = 0;
if( roll_dim_idx >= (size - start) ) {
source_idx = linear_index - ((size - start) * stride);
} else {
source_idx = linear_index + (start * stride);
}
out_tensor[linear_index] = in_tensor[source_idx];
}
out_tensor[linear_index] = in_tensor[source_idx];
}

// Roll a tensor along a dimension
Expand All @@ -129,8 +129,10 @@ Tensor roll_cuda(const Tensor& self, IntArrayRef shifts, IntArrayRef dims) {
if( start < 0 ) start = start + size;

dim3 dim_block = cuda::getApplyBlock();
dim3 dim_grid;
TORCH_CHECK(cuda::getApplyGrid(N, dim_grid, in_tensor.get_device()), "unable to get dim grid");

const int num_mp = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
// Given a thread block size of 512, we launch with 4 blocks per SM/CU
dim3 dim_grid(num_mp * 4);

auto total_dims = in_tensor.dim();

Expand Down