From 2d9c8c4aef0fad74bed24453ea65a62e3d76e368 Mon Sep 17 00:00:00 2001 From: "srishti.dutta1111" Date: Tue, 7 Jul 2026 22:53:30 +0530 Subject: [PATCH 1/4] Add multi-scale deformable attention CUDA kernel --- LICENSE.md | 8 + .../autograd/ms_deform_attn_kernel.cpp | 99 ++ .../cpu/ms_deform_attn_kernel.cpp | 44 + .../cuda/ms_deform_attn_kernel.cu | 220 +++ .../cuda/ms_deform_im2col_cuda.cuh | 1329 +++++++++++++++++ csrc/ops/ms_deform_attn/ms_deform_attn.cpp | 63 + csrc/ops/ms_deform_attn/ms_deform_attn.h | 32 + 7 files changed, 1795 insertions(+) create mode 100644 csrc/ops/ms_deform_attn/autograd/ms_deform_attn_kernel.cpp create mode 100644 csrc/ops/ms_deform_attn/cpu/ms_deform_attn_kernel.cpp create mode 100644 csrc/ops/ms_deform_attn/cuda/ms_deform_attn_kernel.cu create mode 100644 csrc/ops/ms_deform_attn/cuda/ms_deform_im2col_cuda.cuh create mode 100644 csrc/ops/ms_deform_attn/ms_deform_attn.cpp create mode 100644 csrc/ops/ms_deform_attn/ms_deform_attn.h diff --git a/LICENSE.md b/LICENSE.md index 9acb9f0..ae6bafb 100644 --- a/LICENSE.md +++ b/LICENSE.md @@ -19,3 +19,11 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +## Third-party code + +The multi-scale deformable attention kernels under +`csrc/ops/ms_deform_attn/cuda/` are vendored from Deformable-DETR +(https://github.com/fundamentalvision/Deformable-DETR), Copyright (c) 2020 +SenseTime, licensed under the Apache License, Version 2.0. The original license +headers are retained in those files. diff --git a/csrc/ops/ms_deform_attn/autograd/ms_deform_attn_kernel.cpp b/csrc/ops/ms_deform_attn/autograd/ms_deform_attn_kernel.cpp new file mode 100644 index 0000000..fdeaf0d --- /dev/null +++ b/csrc/ops/ms_deform_attn/autograd/ms_deform_attn_kernel.cpp @@ -0,0 +1,99 @@ +#include "../ms_deform_attn.h" + +#include +#include + +namespace vision { +namespace ops { + +namespace { + +class MSDeformAttnFunction + : public torch::autograd::Function { + public: + static torch::autograd::variable_list forward( + torch::autograd::AutogradContext* ctx, + const torch::autograd::Variable& value, + const torch::autograd::Variable& spatial_shapes, + const torch::autograd::Variable& level_start_index, + const torch::autograd::Variable& sampling_loc, + const torch::autograd::Variable& attn_weight, + int64_t im2col_step) { + at::AutoDispatchBelowADInplaceOrView g; + auto output = ms_deform_attn( + value, + spatial_shapes, + level_start_index, + sampling_loc, + attn_weight, + im2col_step); + + ctx->save_for_backward( + {value, spatial_shapes, level_start_index, sampling_loc, attn_weight}); + ctx->saved_data["im2col_step"] = im2col_step; + + return {output}; + } + + static torch::autograd::variable_list backward( + torch::autograd::AutogradContext* ctx, + const torch::autograd::variable_list& grad_output) { + auto saved = ctx->get_saved_variables(); + auto value = saved[0]; + auto spatial_shapes = saved[1]; + auto level_start_index = saved[2]; + auto sampling_loc = saved[3]; + auto attn_weight = saved[4]; + + auto im2col_step = ctx->saved_data["im2col_step"].toInt(); + + auto grads = detail::_ms_deform_attn_backward( + value, + spatial_shapes, + level_start_index, + sampling_loc, + attn_weight, + grad_output[0], + im2col_step); + + auto grad_value = std::get<0>(grads); + auto grad_sampling_loc = std::get<1>(grads); + auto grad_attn_weight = std::get<2>(grads); + + return { + grad_value, + torch::autograd::Variable(), // spatial_shapes + torch::autograd::Variable(), // level_start_index + grad_sampling_loc, + grad_attn_weight, + torch::autograd::Variable(), // im2col_step + }; + } +}; + +at::Tensor ms_deform_attn_autograd( + const at::Tensor& value, + const at::Tensor& spatial_shapes, + const at::Tensor& level_start_index, + const at::Tensor& sampling_loc, + const at::Tensor& attn_weight, + int64_t im2col_step) { + return MSDeformAttnFunction::apply( + value, + spatial_shapes, + level_start_index, + sampling_loc, + attn_weight, + im2col_step)[0]; +} + +} // namespace + +TORCH_LIBRARY_IMPL(torchvision, Autograd, m) { + m.impl( + TORCH_SELECTIVE_NAME("torchvision::ms_deform_attn"), + TORCH_FN(ms_deform_attn_autograd)); +} + +} // namespace ops +} // namespace vision diff --git a/csrc/ops/ms_deform_attn/cpu/ms_deform_attn_kernel.cpp b/csrc/ops/ms_deform_attn/cpu/ms_deform_attn_kernel.cpp new file mode 100644 index 0000000..b4e0e3d --- /dev/null +++ b/csrc/ops/ms_deform_attn/cpu/ms_deform_attn_kernel.cpp @@ -0,0 +1,44 @@ +#include +#include + +namespace vision { +namespace ops { + +namespace { + +// No CPU kernel: ms_deform_attn is CUDA-only. On CPU, models fall back to a +// grid_sample-based implementation. +at::Tensor ms_deform_attn_forward_kernel( + const at::Tensor& value, + const at::Tensor& spatial_shapes, + const at::Tensor& level_start_index, + const at::Tensor& sampling_loc, + const at::Tensor& attn_weight, + int64_t im2col_step) { + TORCH_CHECK(false, "ms_deform_attn is not implemented on the CPU"); +} + +std::tuple ms_deform_attn_backward_kernel( + const at::Tensor& value, + const at::Tensor& spatial_shapes, + const at::Tensor& level_start_index, + const at::Tensor& sampling_loc, + const at::Tensor& attn_weight, + const at::Tensor& grad_output, + int64_t im2col_step) { + TORCH_CHECK(false, "ms_deform_attn is not implemented on the CPU"); +} + +} // namespace + +TORCH_LIBRARY_IMPL(torchvision, CPU, m) { + m.impl( + TORCH_SELECTIVE_NAME("torchvision::ms_deform_attn"), + TORCH_FN(ms_deform_attn_forward_kernel)); + m.impl( + TORCH_SELECTIVE_NAME("torchvision::_ms_deform_attn_backward"), + TORCH_FN(ms_deform_attn_backward_kernel)); +} + +} // namespace ops +} // namespace vision diff --git a/csrc/ops/ms_deform_attn/cuda/ms_deform_attn_kernel.cu b/csrc/ops/ms_deform_attn/cuda/ms_deform_attn_kernel.cu new file mode 100644 index 0000000..ad918c0 --- /dev/null +++ b/csrc/ops/ms_deform_attn/cuda/ms_deform_attn_kernel.cu @@ -0,0 +1,220 @@ +/*! +************************************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************************************** +* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0 +************************************************************************************************** +*/ + +// Vendored from Deformable-DETR (https://github.com/fundamentalvision/Deformable-DETR) +// via LW-DETR (https://github.com/Atten4Vis/LW-DETR). Adapted to register against +// the torchvision dispatcher and to use current libtorch tensor APIs. + +#include + +#include +#include +#include +#include +#include + +#include "ms_deform_im2col_cuda.cuh" + +namespace vision { +namespace ops { + +namespace { + +at::Tensor ms_deform_attn_forward_kernel( + const at::Tensor& value, + const at::Tensor& spatial_shapes, + const at::Tensor& level_start_index, + const at::Tensor& sampling_loc, + const at::Tensor& attn_weight, + int64_t im2col_step) { + TORCH_CHECK(value.is_contiguous(), "value tensor has to be contiguous"); + TORCH_CHECK( + spatial_shapes.is_contiguous(), + "spatial_shapes tensor has to be contiguous"); + TORCH_CHECK( + level_start_index.is_contiguous(), + "level_start_index tensor has to be contiguous"); + TORCH_CHECK( + sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous"); + TORCH_CHECK( + attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous"); + + TORCH_CHECK(value.is_cuda(), "value must be a CUDA tensor"); + TORCH_CHECK(spatial_shapes.is_cuda(), "spatial_shapes must be a CUDA tensor"); + TORCH_CHECK( + level_start_index.is_cuda(), "level_start_index must be a CUDA tensor"); + TORCH_CHECK(sampling_loc.is_cuda(), "sampling_loc must be a CUDA tensor"); + TORCH_CHECK(attn_weight.is_cuda(), "attn_weight must be a CUDA tensor"); + + const int batch = value.size(0); + const int spatial_size = value.size(1); + const int num_heads = value.size(2); + const int channels = value.size(3); + + const int num_levels = spatial_shapes.size(0); + + const int num_query = sampling_loc.size(1); + const int num_point = sampling_loc.size(4); + + const int im2col_step_ = std::min(batch, static_cast(im2col_step)); + + TORCH_CHECK( + batch % im2col_step_ == 0, + "batch(", + batch, + ") must divide im2col_step(", + im2col_step_, + ")"); + + auto output = + at::zeros({batch, num_query, num_heads, channels}, value.options()); + + const int batch_n = im2col_step_; + auto output_n = output.view( + {batch / im2col_step_, batch_n, num_query, num_heads, channels}); + auto per_value_size = spatial_size * num_heads * channels; + auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; + auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; + for (int n = 0; n < batch / im2col_step_; ++n) { + auto columns = output_n.select(0, n); + AT_DISPATCH_FLOATING_TYPES( + value.scalar_type(), "ms_deform_attn_forward_cuda", ([&] { + ms_deformable_im2col_cuda( + at::cuda::getCurrentCUDAStream(), + value.data_ptr() + n * im2col_step_ * per_value_size, + spatial_shapes.data_ptr(), + level_start_index.data_ptr(), + sampling_loc.data_ptr() + + n * im2col_step_ * per_sample_loc_size, + attn_weight.data_ptr() + + n * im2col_step_ * per_attn_weight_size, + batch_n, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + columns.data_ptr()); + })); + } + + output = output.view({batch, num_query, num_heads * channels}); + + return output; +} + +std::tuple ms_deform_attn_backward_kernel( + const at::Tensor& value, + const at::Tensor& spatial_shapes, + const at::Tensor& level_start_index, + const at::Tensor& sampling_loc, + const at::Tensor& attn_weight, + const at::Tensor& grad_output, + int64_t im2col_step) { + TORCH_CHECK(value.is_contiguous(), "value tensor has to be contiguous"); + TORCH_CHECK( + spatial_shapes.is_contiguous(), + "spatial_shapes tensor has to be contiguous"); + TORCH_CHECK( + level_start_index.is_contiguous(), + "level_start_index tensor has to be contiguous"); + TORCH_CHECK( + sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous"); + TORCH_CHECK( + attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous"); + TORCH_CHECK( + grad_output.is_contiguous(), "grad_output tensor has to be contiguous"); + + TORCH_CHECK(value.is_cuda(), "value must be a CUDA tensor"); + TORCH_CHECK(spatial_shapes.is_cuda(), "spatial_shapes must be a CUDA tensor"); + TORCH_CHECK( + level_start_index.is_cuda(), "level_start_index must be a CUDA tensor"); + TORCH_CHECK(sampling_loc.is_cuda(), "sampling_loc must be a CUDA tensor"); + TORCH_CHECK(attn_weight.is_cuda(), "attn_weight must be a CUDA tensor"); + TORCH_CHECK(grad_output.is_cuda(), "grad_output must be a CUDA tensor"); + + const int batch = value.size(0); + const int spatial_size = value.size(1); + const int num_heads = value.size(2); + const int channels = value.size(3); + + const int num_levels = spatial_shapes.size(0); + + const int num_query = sampling_loc.size(1); + const int num_point = sampling_loc.size(4); + + const int im2col_step_ = std::min(batch, static_cast(im2col_step)); + + TORCH_CHECK( + batch % im2col_step_ == 0, + "batch(", + batch, + ") must divide im2col_step(", + im2col_step_, + ")"); + + auto grad_value = at::zeros_like(value); + auto grad_sampling_loc = at::zeros_like(sampling_loc); + auto grad_attn_weight = at::zeros_like(attn_weight); + + const int batch_n = im2col_step_; + auto per_value_size = spatial_size * num_heads * channels; + auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; + auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; + auto grad_output_n = grad_output.view( + {batch / im2col_step_, batch_n, num_query, num_heads, channels}); + + for (int n = 0; n < batch / im2col_step_; ++n) { + auto grad_output_g = grad_output_n.select(0, n); + AT_DISPATCH_FLOATING_TYPES( + value.scalar_type(), "ms_deform_attn_backward_cuda", ([&] { + ms_deformable_col2im_cuda( + at::cuda::getCurrentCUDAStream(), + grad_output_g.data_ptr(), + value.data_ptr() + n * im2col_step_ * per_value_size, + spatial_shapes.data_ptr(), + level_start_index.data_ptr(), + sampling_loc.data_ptr() + + n * im2col_step_ * per_sample_loc_size, + attn_weight.data_ptr() + + n * im2col_step_ * per_attn_weight_size, + batch_n, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value.data_ptr() + + n * im2col_step_ * per_value_size, + grad_sampling_loc.data_ptr() + + n * im2col_step_ * per_sample_loc_size, + grad_attn_weight.data_ptr() + + n * im2col_step_ * per_attn_weight_size); + })); + } + + return std::make_tuple(grad_value, grad_sampling_loc, grad_attn_weight); +} + +} // namespace + +TORCH_LIBRARY_IMPL(torchvision, CUDA, m) { + m.impl( + TORCH_SELECTIVE_NAME("torchvision::ms_deform_attn"), + TORCH_FN(ms_deform_attn_forward_kernel)); + m.impl( + TORCH_SELECTIVE_NAME("torchvision::_ms_deform_attn_backward"), + TORCH_FN(ms_deform_attn_backward_kernel)); +} + +} // namespace ops +} // namespace vision diff --git a/csrc/ops/ms_deform_attn/cuda/ms_deform_im2col_cuda.cuh b/csrc/ops/ms_deform_attn/cuda/ms_deform_im2col_cuda.cuh new file mode 100644 index 0000000..1d267ce --- /dev/null +++ b/csrc/ops/ms_deform_attn/cuda/ms_deform_im2col_cuda.cuh @@ -0,0 +1,1329 @@ +/*! +************************************************************************** +* Deformable DETR +* Copyright (c) 2020 SenseTime. All Rights Reserved. +* Licensed under the Apache License, Version 2.0 [see LICENSE for details] +************************************************************************** +* Modified from DCN (https://github.com/msracver/Deformable-ConvNets) +* Copyright (c) 2018 Microsoft +************************************************************************** +*/ + +#include +#include +#include + +#include +#include + +// THC/THCAtomics.cuh was removed in recent libtorch; ATen/cuda/Atomic.cuh +// provides atomicAdd for the floating point types dispatched here. +#include + +#define CUDA_KERNEL_LOOP(i, n) \ + for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ + i < (n); \ + i += blockDim.x * gridDim.x) + +const int CUDA_NUM_THREADS = 1024; +inline int GET_BLOCKS(const int N, const int num_threads) +{ + return (N + num_threads - 1) / num_threads; +} + + +template +__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + } + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + return val; +} + + +template +__device__ void ms_deform_attn_col2im_bilinear(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c, + const scalar_t &top_grad, + const scalar_t &attn_weight, + scalar_t* &grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + const scalar_t top_grad_value = top_grad * attn_weight; + scalar_t grad_h_weight = 0, grad_w_weight = 0; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + grad_h_weight -= hw * v1; + grad_w_weight -= hh * v1; + atomicAdd(grad_value+ptr1, w1*top_grad_value); + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + grad_h_weight -= lw * v2; + grad_w_weight += hh * v2; + atomicAdd(grad_value+ptr2, w2*top_grad_value); + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + grad_h_weight += hw * v3; + grad_w_weight -= lh * v3; + atomicAdd(grad_value+ptr3, w3*top_grad_value); + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + grad_h_weight += lw * v4; + grad_w_weight += lh * v4; + atomicAdd(grad_value+ptr4, w4*top_grad_value); + } + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + *grad_attn_weight = top_grad * val; + *grad_sampling_loc = width * grad_w_weight * top_grad_value; + *(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value; +} + + +template +__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t* &bottom_data, + const int &height, const int &width, const int &nheads, const int &channels, + const scalar_t &h, const scalar_t &w, const int &m, const int &c, + const scalar_t &top_grad, + const scalar_t &attn_weight, + scalar_t* &grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int h_low = floor(h); + const int w_low = floor(w); + const int h_high = h_low + 1; + const int w_high = w_low + 1; + + const scalar_t lh = h - h_low; + const scalar_t lw = w - w_low; + const scalar_t hh = 1 - lh, hw = 1 - lw; + + const int w_stride = nheads * channels; + const int h_stride = width * w_stride; + const int h_low_ptr_offset = h_low * h_stride; + const int h_high_ptr_offset = h_low_ptr_offset + h_stride; + const int w_low_ptr_offset = w_low * w_stride; + const int w_high_ptr_offset = w_low_ptr_offset + w_stride; + const int base_ptr = m * channels + c; + + const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + const scalar_t top_grad_value = top_grad * attn_weight; + scalar_t grad_h_weight = 0, grad_w_weight = 0; + + scalar_t v1 = 0; + if (h_low >= 0 && w_low >= 0) + { + const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr; + v1 = bottom_data[ptr1]; + grad_h_weight -= hw * v1; + grad_w_weight -= hh * v1; + atomicAdd(grad_value+ptr1, w1*top_grad_value); + } + scalar_t v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + { + const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr; + v2 = bottom_data[ptr2]; + grad_h_weight -= lw * v2; + grad_w_weight += hh * v2; + atomicAdd(grad_value+ptr2, w2*top_grad_value); + } + scalar_t v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + { + const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr; + v3 = bottom_data[ptr3]; + grad_h_weight += hw * v3; + grad_w_weight -= lh * v3; + atomicAdd(grad_value+ptr3, w3*top_grad_value); + } + scalar_t v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + { + const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr; + v4 = bottom_data[ptr4]; + grad_h_weight += lw * v4; + grad_w_weight += lh * v4; + atomicAdd(grad_value+ptr4, w4*top_grad_value); + } + + const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + atomicAdd(grad_attn_weight, top_grad * val); + atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value); + atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value); +} + + +template +__global__ void ms_deformable_im2col_gpu_kernel(const int n, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *data_col) +{ + CUDA_KERNEL_LOOP(index, n) + { + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + scalar_t *data_col_ptr = data_col + index; + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + scalar_t col = 0; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride); + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col) * weight; + } + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + } + } + *data_col_ptr = col; + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2]; + __shared__ scalar_t cache_grad_attn_weight[blockSize]; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + if (tid == 0) + { + scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0]; + int sid=2; + for (unsigned int tid = 1; tid < blockSize; ++tid) + { + _grad_w += cache_grad_sampling_loc[sid]; + _grad_h += cache_grad_sampling_loc[sid + 1]; + _grad_a += cache_grad_attn_weight[tid]; + sid += 2; + } + + + *grad_sampling_loc = _grad_w; + *(grad_sampling_loc + 1) = _grad_h; + *grad_attn_weight = _grad_a; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2]; + __shared__ scalar_t cache_grad_attn_weight[blockSize]; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockSize/2; s>0; s>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + } + __syncthreads(); + } + + if (tid == 0) + { + *grad_sampling_loc = cache_grad_sampling_loc[0]; + *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1]; + *grad_attn_weight = cache_grad_attn_weight[0]; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + if (tid == 0) + { + scalar_t _grad_w=cache_grad_sampling_loc[0], _grad_h=cache_grad_sampling_loc[1], _grad_a=cache_grad_attn_weight[0]; + int sid=2; + for (unsigned int tid = 1; tid < blockDim.x; ++tid) + { + _grad_w += cache_grad_sampling_loc[sid]; + _grad_h += cache_grad_sampling_loc[sid + 1]; + _grad_a += cache_grad_attn_weight[tid]; + sid += 2; + } + + + *grad_sampling_loc = _grad_w; + *(grad_sampling_loc + 1) = _grad_h; + *grad_attn_weight = _grad_a; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + if (tid + (s << 1) < spre) + { + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)]; + } + } + __syncthreads(); + } + + if (tid == 0) + { + *grad_sampling_loc = cache_grad_sampling_loc[0]; + *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1]; + *grad_attn_weight = cache_grad_attn_weight[0]; + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + +template +__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + extern __shared__ int _s[]; + scalar_t* cache_grad_sampling_loc = (scalar_t*)_s; + scalar_t* cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x; + unsigned int tid = threadIdx.x; + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + *(cache_grad_sampling_loc+(threadIdx.x << 1)) = 0; + *(cache_grad_sampling_loc+((threadIdx.x << 1) + 1)) = 0; + *(cache_grad_attn_weight+threadIdx.x)=0; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + cache_grad_sampling_loc+(threadIdx.x << 1), cache_grad_attn_weight+threadIdx.x); + } + + __syncthreads(); + + for (unsigned int s=blockDim.x/2, spre=blockDim.x; s>0; s>>=1, spre>>=1) + { + if (tid < s) { + const unsigned int xid1 = tid << 1; + const unsigned int xid2 = (tid + s) << 1; + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1]; + if (tid + (s << 1) < spre) + { + cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)]; + cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)]; + cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)]; + } + } + __syncthreads(); + } + + if (tid == 0) + { + atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]); + atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]); + atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]); + } + __syncthreads(); + + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n, + const scalar_t *grad_col, + const scalar_t *data_value, + const int64_t *data_spatial_shapes, + const int64_t *data_level_start_index, + const scalar_t *data_sampling_loc, + const scalar_t *data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t *grad_value, + scalar_t *grad_sampling_loc, + scalar_t *grad_attn_weight) +{ + CUDA_KERNEL_LOOP(index, n) + { + int _temp = index; + const int c_col = _temp % channels; + _temp /= channels; + const int sampling_index = _temp; + const int m_col = _temp % num_heads; + _temp /= num_heads; + const int q_col = _temp % num_query; + _temp /= num_query; + const int b_col = _temp; + + const scalar_t top_grad = grad_col[index]; + + int data_weight_ptr = sampling_index * num_levels * num_point; + int data_loc_w_ptr = data_weight_ptr << 1; + const int grad_sampling_ptr = data_weight_ptr; + grad_sampling_loc += grad_sampling_ptr << 1; + grad_attn_weight += grad_sampling_ptr; + const int grad_weight_stride = 1; + const int grad_loc_stride = 2; + const int qid_stride = num_heads * channels; + const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride; + + for (int l_col=0; l_col < num_levels; ++l_col) + { + const int level_start_id = data_level_start_index[l_col]; + const int spatial_h_ptr = l_col << 1; + const int spatial_h = data_spatial_shapes[spatial_h_ptr]; + const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1]; + const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride; + const scalar_t *data_value_ptr = data_value + value_ptr_offset; + scalar_t *grad_value_ptr = grad_value + value_ptr_offset; + + for (int p_col=0; p_col < num_point; ++p_col) + { + const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr]; + const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1]; + const scalar_t weight = data_attn_weight[data_weight_ptr]; + + const scalar_t h_im = loc_h * spatial_h - 0.5; + const scalar_t w_im = loc_w * spatial_w - 0.5; + if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) + { + ms_deform_attn_col2im_bilinear_gm( + data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col, + top_grad, weight, grad_value_ptr, + grad_sampling_loc, grad_attn_weight); + } + data_weight_ptr += 1; + data_loc_w_ptr += 2; + grad_attn_weight += grad_weight_stride; + grad_sampling_loc += grad_loc_stride; + } + } + } +} + + +template +void ms_deformable_im2col_cuda(cudaStream_t stream, + const scalar_t* data_value, + const int64_t* data_spatial_shapes, + const int64_t* data_level_start_index, + const scalar_t* data_sampling_loc, + const scalar_t* data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t* data_col) +{ + const int num_kernels = batch_size * num_query * num_heads * channels; + const int num_actual_kernels = batch_size * num_query * num_heads * channels; + const int num_threads = CUDA_NUM_THREADS; + ms_deformable_im2col_gpu_kernel + <<>>( + num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight, + batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ms_deformable_im2col_cuda: %s\n", cudaGetErrorString(err)); + } + +} + +template +void ms_deformable_col2im_cuda(cudaStream_t stream, + const scalar_t* grad_col, + const scalar_t* data_value, + const int64_t * data_spatial_shapes, + const int64_t * data_level_start_index, + const scalar_t * data_sampling_loc, + const scalar_t * data_attn_weight, + const int batch_size, + const int spatial_size, + const int num_heads, + const int channels, + const int num_levels, + const int num_query, + const int num_point, + scalar_t* grad_value, + scalar_t* grad_sampling_loc, + scalar_t* grad_attn_weight) +{ + const int num_threads = (channels > CUDA_NUM_THREADS)?CUDA_NUM_THREADS:channels; + const int num_kernels = batch_size * num_query * num_heads * channels; + const int num_actual_kernels = batch_size * num_query * num_heads * channels; + if (channels > 1024) + { + if ((channels & 1023) == 0) + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + else + { + ms_deformable_col2im_gpu_kernel_gm + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + } + else{ + switch(channels) + { + case 1: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 2: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 4: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 8: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 16: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 32: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 64: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 128: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 256: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 512: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + case 1024: + ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + break; + default: + if (channels < 64) + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v1 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + else + { + ms_deformable_col2im_gpu_kernel_shm_reduce_v2 + <<>>( + num_kernels, + grad_col, + data_value, + data_spatial_shapes, + data_level_start_index, + data_sampling_loc, + data_attn_weight, + batch_size, + spatial_size, + num_heads, + channels, + num_levels, + num_query, + num_point, + grad_value, + grad_sampling_loc, + grad_attn_weight); + } + } + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in ms_deformable_col2im_cuda: %s\n", cudaGetErrorString(err)); + } + +} \ No newline at end of file diff --git a/csrc/ops/ms_deform_attn/ms_deform_attn.cpp b/csrc/ops/ms_deform_attn/ms_deform_attn.cpp new file mode 100644 index 0000000..5832428 --- /dev/null +++ b/csrc/ops/ms_deform_attn/ms_deform_attn.cpp @@ -0,0 +1,63 @@ +#include "ms_deform_attn.h" + +#include +#include +#include + +namespace vision { +namespace ops { + +at::Tensor ms_deform_attn( + const at::Tensor& value, + const at::Tensor& spatial_shapes, + const at::Tensor& level_start_index, + const at::Tensor& sampling_loc, + const at::Tensor& attn_weight, + int64_t im2col_step) { + static auto op = c10::Dispatcher::singleton() + .findSchemaOrThrow("torchvision::ms_deform_attn", "") + .typed(); + return op.call( + value, + spatial_shapes, + level_start_index, + sampling_loc, + attn_weight, + im2col_step); +} + +namespace detail { + +std::tuple _ms_deform_attn_backward( + const at::Tensor& value, + const at::Tensor& spatial_shapes, + const at::Tensor& level_start_index, + const at::Tensor& sampling_loc, + const at::Tensor& attn_weight, + const at::Tensor& grad_output, + int64_t im2col_step) { + static auto op = + c10::Dispatcher::singleton() + .findSchemaOrThrow("torchvision::_ms_deform_attn_backward", "") + .typed(); + return op.call( + value, + spatial_shapes, + level_start_index, + sampling_loc, + attn_weight, + grad_output, + im2col_step); +} + +} // namespace detail + +TORCH_LIBRARY_FRAGMENT(torchvision, m) { + m.def(TORCH_SELECTIVE_SCHEMA( + "torchvision::ms_deform_attn(Tensor value, Tensor spatial_shapes, Tensor level_start_index, Tensor sampling_loc, Tensor attn_weight, int im2col_step) -> Tensor")); + m.def(TORCH_SELECTIVE_SCHEMA( + "torchvision::_ms_deform_attn_backward(Tensor value, Tensor spatial_shapes, Tensor level_start_index, Tensor sampling_loc, Tensor attn_weight, Tensor grad_output, int im2col_step) -> (Tensor, Tensor, Tensor)")); +} + +} // namespace ops +} // namespace vision diff --git a/csrc/ops/ms_deform_attn/ms_deform_attn.h b/csrc/ops/ms_deform_attn/ms_deform_attn.h new file mode 100644 index 0000000..ce76ad0 --- /dev/null +++ b/csrc/ops/ms_deform_attn/ms_deform_attn.h @@ -0,0 +1,32 @@ +#pragma once + +#include + +namespace vision { +namespace ops { + +// Multi-scale deformable attention, as used by Deformable-DETR / LW-DETR. +// Only a CUDA kernel is provided; calling on CPU tensors raises an error. +at::Tensor ms_deform_attn( + const at::Tensor& value, + const at::Tensor& spatial_shapes, + const at::Tensor& level_start_index, + const at::Tensor& sampling_loc, + const at::Tensor& attn_weight, + int64_t im2col_step); + +namespace detail { + +std::tuple _ms_deform_attn_backward( + const at::Tensor& value, + const at::Tensor& spatial_shapes, + const at::Tensor& level_start_index, + const at::Tensor& sampling_loc, + const at::Tensor& attn_weight, + const at::Tensor& grad_output, + int64_t im2col_step); + +} // namespace detail + +} // namespace ops +} // namespace vision From a1b79876ad4ea3e482ce13ef2b208d99c948422c Mon Sep 17 00:00:00 2001 From: "srishti.dutta1111" Date: Tue, 7 Jul 2026 22:53:37 +0530 Subject: [PATCH 2/4] Register ms_deform_attn op and regenerate bindings --- R/RcppExports.R | 4 ++++ csrc/CMakeLists.txt | 21 +++++++++++++++++++++ csrc/include/torchvisionlib/exports.h | 6 ++++++ csrc/src/exports.cpp | 7 +++++++ csrc/src/ops.cpp | 20 +++++++++++++++++++- csrc/src/torchvisionlib.def | 1 + inst/def/torchvisionlib.def | 1 + inst/include/torchvisionlib/exports.h | 6 ++++++ src/RcppExports.cpp | 17 +++++++++++++++++ src/exports.cpp | 4 ++++ src/exports.h | 1 + 11 files changed, 87 insertions(+), 1 deletion(-) diff --git a/R/RcppExports.R b/R/RcppExports.R index 8b5dd00..46e8c79 100644 --- a/R/RcppExports.R +++ b/R/RcppExports.R @@ -5,6 +5,10 @@ rcpp_vision_ops_nms <- function(dets, scores, iou_threshold) { .Call('_torchvisionlib_rcpp_vision_ops_nms', PACKAGE = 'torchvisionlib', dets, scores, iou_threshold) } +rcpp_vision_ops_ms_deform_attn <- function(value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step) { + .Call('_torchvisionlib_rcpp_vision_ops_ms_deform_attn', PACKAGE = 'torchvisionlib', value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step) +} + rcpp_vision_ops_deform_conv2d <- function(input, weight, offset, mask, bias, stride_h, stride_w, pad_h, pad_w, dilation_h, dilation_w, groups, offset_groups, use_mask) { .Call('_torchvisionlib_rcpp_vision_ops_deform_conv2d', PACKAGE = 'torchvisionlib', input, weight, offset, mask, bias, stride_h, stride_w, pad_h, pad_w, dilation_h, dilation_w, groups, offset_groups, use_mask) } diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index c705c75..61dbcce 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -134,10 +134,22 @@ endif() set(TORCHVISION_SRC src/torchvisionlib.cpp src/ops.cpp src/exports.cpp src/torchvisionlib_types.cpp) +# Multi-scale deformable attention op. Host/registration code always builds; +# the CUDA kernel is added only when CUDA is enabled. +list(APPEND TORCHVISION_SRC + ops/ms_deform_attn/ms_deform_attn.cpp + ops/ms_deform_attn/cpu/ms_deform_attn_kernel.cpp + ops/ms_deform_attn/autograd/ms_deform_attn_kernel.cpp +) +if (DEFINED ENV{CUDA} AND NOT '$ENV{CUDA}' STREQUAL '') + list(APPEND TORCHVISION_SRC ops/ms_deform_attn/cuda/ms_deform_attn_kernel.cu) +endif() + add_library(torchvisionlib SHARED ${TORCHVISION_SRC}) add_library(torchvisionlib::library ALIAS torchvisionlib) target_include_directories(torchvisionlib PUBLIC + ${PROJECT_SOURCE_DIR} ${PROJECT_SOURCE_DIR}/include ${TORCH_HOME}/include ${TORCHVISION_INCLUDE_DIR} @@ -163,6 +175,15 @@ add_dependencies(torchvisionlib torchvisionlib_export) set_property(TARGET torchvisionlib PROPERTY CXX_STANDARD 17) +# CUDA build settings for the ms_deform_attn kernel (CUDA-enabled builds only). +if (DEFINED ENV{CUDA} AND NOT '$ENV{CUDA}' STREQUAL '') + target_compile_definitions(torchvisionlib PRIVATE WITH_CUDA) + set_property(TARGET torchvisionlib PROPERTY CUDA_STANDARD 17) + if (NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + set_property(TARGET torchvisionlib PROPERTY CUDA_ARCHITECTURES OFF) + endif() +endif() + target_link_libraries(torchvisionlib "${TORCH_LIBRARIES}") target_link_libraries(torchvisionlib TorchVision) diff --git a/csrc/include/torchvisionlib/exports.h b/csrc/include/torchvisionlib/exports.h index a8de27c..16278a9 100644 --- a/csrc/include/torchvisionlib/exports.h +++ b/csrc/include/torchvisionlib/exports.h @@ -28,6 +28,7 @@ TORCHVISIONLIB_API void* torchvisionlib_last_error (); TORCHVISIONLIB_API void torchvisionlib_last_error_clear(); TORCHVISIONLIB_API void* _vision_ops_nms (void* dets, void* scores, double iou_threshold); +TORCHVISIONLIB_API void* _vision_ops_ms_deform_attn (void* value, void* spatial_shapes, void* level_start_index, void* sampling_loc, void* attn_weight, std::int64_t im2col_step); TORCHVISIONLIB_API void* _vision_ops_deform_conv2d (void* input, void* weight, void* offset, void* mask, void* bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask); TORCHVISIONLIB_API void* _vision_ops_ps_roi_align (void* input, void* rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t sampling_ratio); TORCHVISIONLIB_API void* _vision_ops_ps_roi_pool (void* input, void* rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width); @@ -45,6 +46,11 @@ inline void* vision_ops_nms (void* dets, void* scores, double iou_threshold) { host_exception_handler(); return ret; } +inline void* vision_ops_ms_deform_attn (void* value, void* spatial_shapes, void* level_start_index, void* sampling_loc, void* attn_weight, std::int64_t im2col_step) { + auto ret = _vision_ops_ms_deform_attn(value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step); + host_exception_handler(); + return ret; +} inline void* vision_ops_deform_conv2d (void* input, void* weight, void* offset, void* mask, void* bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask) { auto ret = _vision_ops_deform_conv2d(input, weight, offset, mask, bias, stride_h, stride_w, pad_h, pad_w, dilation_h, dilation_w, groups, offset_groups, use_mask); host_exception_handler(); diff --git a/csrc/src/exports.cpp b/csrc/src/exports.cpp index 4c1dd05..94bfc5f 100644 --- a/csrc/src/exports.cpp +++ b/csrc/src/exports.cpp @@ -21,6 +21,13 @@ TORCHVISIONLIB_API void* _vision_ops_nms (void* dets, void* scores, double iou_t } TORCHVISIONLIB_HANDLE_EXCEPTION return (void*) NULL; } +torch::Tensor vision_ops_ms_deform_attn (torch::Tensor value, torch::Tensor spatial_shapes, torch::Tensor level_start_index, torch::Tensor sampling_loc, torch::Tensor attn_weight, std::int64_t im2col_step); +TORCHVISIONLIB_API void* _vision_ops_ms_deform_attn (void* value, void* spatial_shapes, void* level_start_index, void* sampling_loc, void* attn_weight, std::int64_t im2col_step) { + try { + return make_raw::Tensor(vision_ops_ms_deform_attn(from_raw::Tensor(value), from_raw::Tensor(spatial_shapes), from_raw::Tensor(level_start_index), from_raw::Tensor(sampling_loc), from_raw::Tensor(attn_weight), im2col_step)); + } TORCHVISIONLIB_HANDLE_EXCEPTION + return (void*) NULL; +} torch::Tensor vision_ops_deform_conv2d (torch::Tensor input, torch::Tensor weight, torch::Tensor offset, torch::Tensor mask, torch::Tensor bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask); TORCHVISIONLIB_API void* _vision_ops_deform_conv2d (void* input, void* weight, void* offset, void* mask, void* bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask) { try { diff --git a/csrc/src/ops.cpp b/csrc/src/ops.cpp index 46439fa..74f3eab 100644 --- a/csrc/src/ops.cpp +++ b/csrc/src/ops.cpp @@ -7,13 +7,31 @@ #include #include #include - +#include "ops/ms_deform_attn/ms_deform_attn.h" // [[torch::export]] torch::Tensor vision_ops_nms(torch::Tensor dets, torch::Tensor scores, double iou_threshold) { return vision::ops::nms(dets, scores, iou_threshold); } +// [[torch::export]] +torch::Tensor vision_ops_ms_deform_attn( + torch::Tensor value, + torch::Tensor spatial_shapes, + torch::Tensor level_start_index, + torch::Tensor sampling_loc, + torch::Tensor attn_weight, + std::int64_t im2col_step) { + return vision::ops::ms_deform_attn( + value, + spatial_shapes, + level_start_index, + sampling_loc, + attn_weight, + im2col_step + ); +} + // [[torch::export]] torch::Tensor vision_ops_deform_conv2d( torch::Tensor input, diff --git a/csrc/src/torchvisionlib.def b/csrc/src/torchvisionlib.def index df913d4..198d070 100644 --- a/csrc/src/torchvisionlib.def +++ b/csrc/src/torchvisionlib.def @@ -4,6 +4,7 @@ EXPORTS ;------ autogenerated ------------------------- ; don't modify between the autogenerated lines _vision_ops_nms + _vision_ops_ms_deform_attn _vision_ops_deform_conv2d _vision_ops_ps_roi_align _vision_ops_ps_roi_pool diff --git a/inst/def/torchvisionlib.def b/inst/def/torchvisionlib.def index df913d4..198d070 100644 --- a/inst/def/torchvisionlib.def +++ b/inst/def/torchvisionlib.def @@ -4,6 +4,7 @@ EXPORTS ;------ autogenerated ------------------------- ; don't modify between the autogenerated lines _vision_ops_nms + _vision_ops_ms_deform_attn _vision_ops_deform_conv2d _vision_ops_ps_roi_align _vision_ops_ps_roi_pool diff --git a/inst/include/torchvisionlib/exports.h b/inst/include/torchvisionlib/exports.h index a8de27c..16278a9 100644 --- a/inst/include/torchvisionlib/exports.h +++ b/inst/include/torchvisionlib/exports.h @@ -28,6 +28,7 @@ TORCHVISIONLIB_API void* torchvisionlib_last_error (); TORCHVISIONLIB_API void torchvisionlib_last_error_clear(); TORCHVISIONLIB_API void* _vision_ops_nms (void* dets, void* scores, double iou_threshold); +TORCHVISIONLIB_API void* _vision_ops_ms_deform_attn (void* value, void* spatial_shapes, void* level_start_index, void* sampling_loc, void* attn_weight, std::int64_t im2col_step); TORCHVISIONLIB_API void* _vision_ops_deform_conv2d (void* input, void* weight, void* offset, void* mask, void* bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask); TORCHVISIONLIB_API void* _vision_ops_ps_roi_align (void* input, void* rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t sampling_ratio); TORCHVISIONLIB_API void* _vision_ops_ps_roi_pool (void* input, void* rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width); @@ -45,6 +46,11 @@ inline void* vision_ops_nms (void* dets, void* scores, double iou_threshold) { host_exception_handler(); return ret; } +inline void* vision_ops_ms_deform_attn (void* value, void* spatial_shapes, void* level_start_index, void* sampling_loc, void* attn_weight, std::int64_t im2col_step) { + auto ret = _vision_ops_ms_deform_attn(value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step); + host_exception_handler(); + return ret; +} inline void* vision_ops_deform_conv2d (void* input, void* weight, void* offset, void* mask, void* bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask) { auto ret = _vision_ops_deform_conv2d(input, weight, offset, mask, bias, stride_h, stride_w, pad_h, pad_w, dilation_h, dilation_w, groups, offset_groups, use_mask); host_exception_handler(); diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp index 4ee93d8..2deed52 100644 --- a/src/RcppExports.cpp +++ b/src/RcppExports.cpp @@ -24,6 +24,22 @@ BEGIN_RCPP return rcpp_result_gen; END_RCPP } +// rcpp_vision_ops_ms_deform_attn +torch::Tensor rcpp_vision_ops_ms_deform_attn(torch::Tensor value, torch::Tensor spatial_shapes, torch::Tensor level_start_index, torch::Tensor sampling_loc, torch::Tensor attn_weight, std::int64_t im2col_step); +RcppExport SEXP _torchvisionlib_rcpp_vision_ops_ms_deform_attn(SEXP valueSEXP, SEXP spatial_shapesSEXP, SEXP level_start_indexSEXP, SEXP sampling_locSEXP, SEXP attn_weightSEXP, SEXP im2col_stepSEXP) { +BEGIN_RCPP + Rcpp::RObject rcpp_result_gen; + Rcpp::RNGScope rcpp_rngScope_gen; + Rcpp::traits::input_parameter< torch::Tensor >::type value(valueSEXP); + Rcpp::traits::input_parameter< torch::Tensor >::type spatial_shapes(spatial_shapesSEXP); + Rcpp::traits::input_parameter< torch::Tensor >::type level_start_index(level_start_indexSEXP); + Rcpp::traits::input_parameter< torch::Tensor >::type sampling_loc(sampling_locSEXP); + Rcpp::traits::input_parameter< torch::Tensor >::type attn_weight(attn_weightSEXP); + Rcpp::traits::input_parameter< std::int64_t >::type im2col_step(im2col_stepSEXP); + rcpp_result_gen = Rcpp::wrap(rcpp_vision_ops_ms_deform_attn(value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step)); + return rcpp_result_gen; +END_RCPP +} // rcpp_vision_ops_deform_conv2d torch::Tensor rcpp_vision_ops_deform_conv2d(torch::Tensor input, torch::Tensor weight, torch::Tensor offset, torch::Tensor mask, torch::Tensor bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask); RcppExport SEXP _torchvisionlib_rcpp_vision_ops_deform_conv2d(SEXP inputSEXP, SEXP weightSEXP, SEXP offsetSEXP, SEXP maskSEXP, SEXP biasSEXP, SEXP stride_hSEXP, SEXP stride_wSEXP, SEXP pad_hSEXP, SEXP pad_wSEXP, SEXP dilation_hSEXP, SEXP dilation_wSEXP, SEXP groupsSEXP, SEXP offset_groupsSEXP, SEXP use_maskSEXP) { @@ -168,6 +184,7 @@ END_RCPP static const R_CallMethodDef CallEntries[] = { {"_torchvisionlib_rcpp_vision_ops_nms", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_nms, 3}, + {"_torchvisionlib_rcpp_vision_ops_ms_deform_attn", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_ms_deform_attn, 6}, {"_torchvisionlib_rcpp_vision_ops_deform_conv2d", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_deform_conv2d, 14}, {"_torchvisionlib_rcpp_vision_ops_ps_roi_align", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_ps_roi_align, 6}, {"_torchvisionlib_rcpp_vision_ops_ps_roi_pool", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_ps_roi_pool, 5}, diff --git a/src/exports.cpp b/src/exports.cpp index 142362a..f508f90 100644 --- a/src/exports.cpp +++ b/src/exports.cpp @@ -8,6 +8,10 @@ torch::Tensor rcpp_vision_ops_nms (torch::Tensor dets, torch::Tensor scores, dou return vision_ops_nms(dets.get(), scores.get(), iou_threshold); } // [[Rcpp::export]] +torch::Tensor rcpp_vision_ops_ms_deform_attn (torch::Tensor value, torch::Tensor spatial_shapes, torch::Tensor level_start_index, torch::Tensor sampling_loc, torch::Tensor attn_weight, std::int64_t im2col_step) { + return vision_ops_ms_deform_attn(value.get(), spatial_shapes.get(), level_start_index.get(), sampling_loc.get(), attn_weight.get(), im2col_step); +} +// [[Rcpp::export]] torch::Tensor rcpp_vision_ops_deform_conv2d (torch::Tensor input, torch::Tensor weight, torch::Tensor offset, torch::Tensor mask, torch::Tensor bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask) { return vision_ops_deform_conv2d(input.get(), weight.get(), offset.get(), mask.get(), bias.get(), stride_h, stride_w, pad_h, pad_w, dilation_h, dilation_w, groups, offset_groups, use_mask); } diff --git a/src/exports.h b/src/exports.h index 534858c..4444e58 100644 --- a/src/exports.h +++ b/src/exports.h @@ -4,6 +4,7 @@ #include "torchvisionlib_types.h" torch::Tensor rcpp_vision_ops_nms (torch::Tensor dets, torch::Tensor scores, double iou_threshold); +torch::Tensor rcpp_vision_ops_ms_deform_attn (torch::Tensor value, torch::Tensor spatial_shapes, torch::Tensor level_start_index, torch::Tensor sampling_loc, torch::Tensor attn_weight, std::int64_t im2col_step); torch::Tensor rcpp_vision_ops_deform_conv2d (torch::Tensor input, torch::Tensor weight, torch::Tensor offset, torch::Tensor mask, torch::Tensor bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask); torchvisionlib::tensor_pair rcpp_vision_ops_ps_roi_align (torch::Tensor input, torch::Tensor rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t sampling_ratio); torchvisionlib::tensor_pair rcpp_vision_ops_ps_roi_pool (torch::Tensor input, torch::Tensor rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width); From e2859a23d23bb8daeb0e5536942fd4a3ff2a12a2 Mon Sep 17 00:00:00 2001 From: "srishti.dutta1111" Date: Tue, 7 Jul 2026 22:53:44 +0530 Subject: [PATCH 3/4] Add ops_ms_deform_attn R wrapper, docs and tests --- NAMESPACE | 1 + NEWS.md | 4 + R/ops.R | 45 +++++++++ man/ops_ms_deform_attn.Rd | 54 +++++++++++ man/ops_nms.Rd | 4 + tests/testthat/test-ops-ms-deform-attn.R | 113 +++++++++++++++++++++++ 6 files changed, 221 insertions(+) create mode 100644 man/ops_ms_deform_attn.Rd create mode 100644 tests/testthat/test-ops-ms-deform-attn.R diff --git a/NAMESPACE b/NAMESPACE index 5a141c9..064fb6e 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -3,6 +3,7 @@ export(install_torchvisionlib) export(nn_ps_roi_align) export(ops_deform_conv2d) +export(ops_ms_deform_attn) export(ops_nms) export(ops_ps_roi_align) export(torchvisionlib_is_installed) diff --git a/NEWS.md b/NEWS.md index a92fc08..e105837 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,5 +1,9 @@ # torchvisionlib (development version) +- Added `ops_ms_deform_attn()`, a CUDA implementation of multi-scale deformable + attention (used by Deformable-DETR and LW-DETR). Vendored from Deformable-DETR + (Apache-2.0). (#25) + # torchvisionlib 0.5.0 - Updates to support LibTorch v2.0.1 diff --git a/R/ops.R b/R/ops.R index c83c475..c6920b5 100644 --- a/R/ops.R +++ b/R/ops.R @@ -126,6 +126,51 @@ ops_deform_conv2d <- function(input, ) } +#' Multi-scale deformable attention +#' +#' Computes multi-scale deformable attention as used by Deformable-DETR and +#' LW-DETR. Only a CUDA implementation is provided; all input tensors must be on +#' the same CUDA device. On CPU, models should fall back to a pure-R +#' implementation (e.g. based on [torch::nnf_grid_sample()]). +#' +#' @details +#' `spatial_shapes` and `level_start_index` use the kernel's 0-based indexing +#' convention: `level_start_index[l]` is the flat offset (into the `Len_in` +#' dimension of `value`) of the first element of level `l`, and the output is a +#' feature tensor, so no 1-based index adjustment is applied. +#' +#' @param value (`Tensor[batch, Len_in, n_heads, head_dim]`): flattened +#' multi-scale feature values. +#' @param spatial_shapes (`Tensor[n_levels, 2]`, integer): the `(H, W)` of each +#' feature level. The sum of `H * W` over levels must equal `Len_in`. +#' @param level_start_index (`Tensor[n_levels]`, integer): 0-based start offset +#' of each level within `Len_in`. +#' @param sampling_locations (`Tensor[batch, Len_q, n_heads, n_levels, n_points, 2]`): +#' sampling locations in `[0, 1]` (normalized `x, y`). +#' @param attention_weights (`Tensor[batch, Len_q, n_heads, n_levels, n_points]`): +#' attention weights, typically normalized over the `n_levels * n_points` axis. +#' @param im2col_step (int): batch chunk size used internally by the kernel. +#' Must divide `batch`. Default: 64. +#' +#' @returns +#' `Tensor[batch, Len_q, n_heads * head_dim]`: the attended output. +#' +#' @family ops +#' @export +ops_ms_deform_attn <- function(value, spatial_shapes, level_start_index, + sampling_locations, attention_weights, + im2col_step = 64L) { + rcpp_vision_ops_ms_deform_attn( + value, + spatial_shapes, + level_start_index, + sampling_locations, + attention_weights, + im2col_step + ) +} + + #' Performs Position-Sensitive Region of Interest (RoI) Align operator #' #' The (RoI) Align operator is mentioned in [Light-Head R-CNN](https://arxiv.org/abs/1711.07264). diff --git a/man/ops_ms_deform_attn.Rd b/man/ops_ms_deform_attn.Rd new file mode 100644 index 0000000..af34c87 --- /dev/null +++ b/man/ops_ms_deform_attn.Rd @@ -0,0 +1,54 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/ops.R +\name{ops_ms_deform_attn} +\alias{ops_ms_deform_attn} +\title{Multi-scale deformable attention} +\usage{ +ops_ms_deform_attn( + value, + spatial_shapes, + level_start_index, + sampling_locations, + attention_weights, + im2col_step = 64L +) +} +\arguments{ +\item{value}{(\code{Tensor[batch, Len_in, n_heads, head_dim]}): flattened +multi-scale feature values.} + +\item{spatial_shapes}{(\code{Tensor[n_levels, 2]}, integer): the \verb{(H, W)} of each +feature level. The sum of \code{H * W} over levels must equal \code{Len_in}.} + +\item{level_start_index}{(\code{Tensor[n_levels]}, integer): 0-based start offset +of each level within \code{Len_in}.} + +\item{sampling_locations}{(\code{Tensor[batch, Len_q, n_heads, n_levels, n_points, 2]}): +sampling locations in \verb{[0, 1]} (normalized \verb{x, y}).} + +\item{attention_weights}{(\code{Tensor[batch, Len_q, n_heads, n_levels, n_points]}): +attention weights, typically normalized over the \code{n_levels * n_points} axis.} + +\item{im2col_step}{(int): batch chunk size used internally by the kernel. +Must divide \code{batch}. Default: 64.} +} +\value{ +\code{Tensor[batch, Len_q, n_heads * head_dim]}: the attended output. +} +\description{ +Computes multi-scale deformable attention as used by Deformable-DETR and +LW-DETR. Only a CUDA implementation is provided; all input tensors must be on +the same CUDA device. On CPU, models should fall back to a pure-R +implementation (e.g. based on \code{\link[torch:nnf_grid_sample]{torch::nnf_grid_sample()}}). +} +\details{ +\code{spatial_shapes} and \code{level_start_index} use the kernel's 0-based indexing +convention: \code{level_start_index[l]} is the flat offset (into the \code{Len_in} +dimension of \code{value}) of the first element of level \code{l}, and the output is a +feature tensor, so no 1-based index adjustment is applied. +} +\seealso{ +Other ops: +\code{\link[=ops_nms]{ops_nms()}} +} +\concept{ops} diff --git a/man/ops_nms.Rd b/man/ops_nms.Rd index 3110dd9..9ced769 100644 --- a/man/ops_nms.Rd +++ b/man/ops_nms.Rd @@ -36,4 +36,8 @@ if (torchvisionlib_is_installed()) { ops_nms(torch::torch_rand(3, 4), torch::torch_rand(3), 0.5) } } +\seealso{ +Other ops: +\code{\link[=ops_ms_deform_attn]{ops_ms_deform_attn()}} +} \concept{ops} diff --git a/tests/testthat/test-ops-ms-deform-attn.R b/tests/testthat/test-ops-ms-deform-attn.R new file mode 100644 index 0000000..216a445 --- /dev/null +++ b/tests/testthat/test-ops-ms-deform-attn.R @@ -0,0 +1,113 @@ +library(torch) + +# Independent pure-R reference, equivalent to Deformable-DETR's +# `ms_deform_attn_core_pytorch`. Used to validate the CUDA kernel. +ms_deform_attn_reference <- function(value, spatial_shapes, sampling_locations, + attention_weights) { + n <- value$size(1) + n_heads <- value$size(3) + head_dim <- value$size(4) + len_q <- sampling_locations$size(2) + n_levels <- sampling_locations$size(4) + n_points <- sampling_locations$size(5) + + sizes <- as.integer(spatial_shapes[, 1] * spatial_shapes[, 2]) + value_list <- value$split(sizes, dim = 2) + sampling_grids <- 2 * sampling_locations - 1 + + sampling_value_list <- list() + for (lvl in seq_len(n_levels)) { + h_l <- as.integer(spatial_shapes[lvl, 1]) + w_l <- as.integer(spatial_shapes[lvl, 2]) + value_l <- value_list[[lvl]]$flatten(start_dim = 3)$transpose(2, 3)$ + reshape(c(n * n_heads, head_dim, h_l, w_l)) + grid_l <- sampling_grids[, , , lvl, , ]$transpose(2, 3)$ + flatten(start_dim = 1, end_dim = 2) # (n*n_heads, len_q, n_points, 2) + sampling_value_list[[lvl]] <- nnf_grid_sample( + value_l, grid_l, + mode = "bilinear", padding_mode = "zeros", align_corners = FALSE + ) # (n*n_heads, head_dim, len_q, n_points) + } + + # (n*n_heads, head_dim, len_q, n_levels, n_points) -> flatten levels*points + sampled <- torch_stack(sampling_value_list, dim = -2)$ + flatten(start_dim = 4) + attn <- attention_weights$transpose(2, 3)$ + reshape(c(n * n_heads, 1, len_q, n_levels * n_points)) + output <- (sampled * attn)$sum(-1)$ + view(c(n, n_heads * head_dim, len_q)) + output$transpose(2, 3) # (n, len_q, n_heads*head_dim) +} + +make_inputs <- function(device, requires_grad = FALSE) { + torch_manual_seed(1) + n <- 2L; n_heads <- 4L; head_dim <- 8L + n_levels <- 2L; n_points <- 3L; len_q <- 5L + shapes <- matrix(c(6L, 4L, 3L, 2L), ncol = 2, byrow = TRUE) + spatial_shapes <- torch_tensor(shapes, dtype = torch_long())$to(device = device) + level_start_index <- torch_cat(list( + torch_zeros(1, dtype = torch_long()), + (spatial_shapes[, 1] * spatial_shapes[, 2])$cumsum(1)[1:(n_levels - 1)] + ))$to(device = device) + len_in <- sum(as.integer(shapes[, 1] * shapes[, 2])) + + value <- torch_rand(n, len_in, n_heads, head_dim, device = device, + requires_grad = requires_grad) + sampling_locations <- torch_rand(n, len_q, n_heads, n_levels, n_points, 2, + device = device, requires_grad = requires_grad) + attention_weights <- torch_rand(n, len_q, n_heads, n_levels, n_points, + device = device, requires_grad = requires_grad) + + list(value = value, spatial_shapes = spatial_shapes, + level_start_index = level_start_index, + sampling_locations = sampling_locations, + attention_weights = attention_weights) +} + +test_that("ms_deform_attn matches the grid_sample reference (forward)", { + skip_if_not(cuda_is_available()) + x <- make_inputs("cuda") + + out <- ops_ms_deform_attn( + x$value, x$spatial_shapes, x$level_start_index, + x$sampling_locations, x$attention_weights, 64L + ) + ref <- ms_deform_attn_reference( + x$value, x$spatial_shapes, x$sampling_locations, x$attention_weights + ) + + expect_equal(out$shape, ref$shape) + expect_true(torch_allclose(out, ref, atol = 1e-4, rtol = 1e-4)) +}) + +test_that("ms_deform_attn gradients match the reference (backward)", { + skip_if_not(cuda_is_available()) + x <- make_inputs("cuda", requires_grad = TRUE) + xr <- make_inputs("cuda", requires_grad = TRUE) + + ops_ms_deform_attn( + x$value, x$spatial_shapes, x$level_start_index, + x$sampling_locations, x$attention_weights, 64L + )$sum()$backward() + + ms_deform_attn_reference( + xr$value, xr$spatial_shapes, xr$sampling_locations, xr$attention_weights + )$sum()$backward() + + expect_true(torch_allclose(x$value$grad, xr$value$grad, atol = 1e-3, rtol = 1e-3)) + expect_true(torch_allclose(x$sampling_locations$grad, xr$sampling_locations$grad, + atol = 1e-3, rtol = 1e-3)) + expect_true(torch_allclose(x$attention_weights$grad, xr$attention_weights$grad, + atol = 1e-3, rtol = 1e-3)) +}) + +test_that("ms_deform_attn raises on CPU tensors", { + x <- make_inputs("cpu") + expect_error( + ops_ms_deform_attn( + x$value, x$spatial_shapes, x$level_start_index, + x$sampling_locations, x$attention_weights, 64L + ), + regexp = "not implemented on the CPU" + ) +}) From ebc1aaedd57a18750ab496f4bc43600ce2bb5800 Mon Sep 17 00:00:00 2001 From: cregouby Date: Thu, 30 Jul 2026 23:28:05 +0200 Subject: [PATCH 4/4] add dependancies for fs and thus torchexport to compile on Linux (#35) * bump zlib, libpng, libjpeg to package compilation * add fs libuv1 dependancy to CI-CD * bump checkout, Jimver/cuda-toolkit, setup-pandoc , download-artifact, upload-artifact --- .github/workflows/R-CMD-check.yaml | 14 ++++++------ DESCRIPTION | 35 +++++++++++++++++------------- csrc/CMakeLists.txt | 12 +++++----- 3 files changed, 33 insertions(+), 28 deletions(-) diff --git a/.github/workflows/R-CMD-check.yaml b/.github/workflows/R-CMD-check.yaml index 12b2242..8470463 100644 --- a/.github/workflows/R-CMD-check.yaml +++ b/.github/workflows/R-CMD-check.yaml @@ -44,7 +44,7 @@ jobs: TZ: 'Etc/UTC' steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v7 - uses: mlverse/torch/.github/actions/setup-r@main with: @@ -54,7 +54,8 @@ jobs: if: matrix.container != '' run: | apt-get update - apt-get install -y curl sudo libxml2-dev wget chrpath rsync git libjpeg-dev + apt-get install -y curl sudo libxml2-dev wget chrpath rsync git libjpeg-dev libpng-dev \ + libcurl4-openssl-dev pkg-config patchelf libuv1-dev curl -fsSL https://get.docker.com -o get-docker.sh DRY_RUN=1 sh ./get-docker.sh @@ -64,7 +65,7 @@ jobs: cmake-version: '3.31' - if: ${{matrix.config.cuda != ''}} - uses: Jimver/cuda-toolkit@v0.2.23 + uses: Jimver/cuda-toolkit@v0.2.35 id: cuda-toolkit with: cuda: '${{matrix.config.cuda}}.${{matrix.config.cuda_patch}}' @@ -91,7 +92,6 @@ jobs: mkdir build && cd build cmake .. ${{ matrix.config.cmake_args }} cmake --build . --target package --config Release --parallel 4 - - id: version shell: bash run: | @@ -109,7 +109,7 @@ jobs: file_glob: true tag: v${{ steps.version.outputs.version }} - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@v7 with: name: ${{ steps.version.outputs.fname }} path: csrc/build/*.zip @@ -139,7 +139,7 @@ jobs: steps: - uses: actions/checkout@v2 - - uses: r-lib/actions/setup-pandoc@v1 + - uses: r-lib/actions/setup-pandoc@v2 if: ${{ !contains(matrix.config.os, 'm1') }} - uses: r-lib/actions/setup-r@v2 @@ -173,7 +173,7 @@ jobs: version=$(Rscript -e "cat(as.character(desc::desc_get_version()))") echo "version=$version" >> $GITHUB_OUTPUT - - uses: actions/download-artifact@v4 + - uses: actions/download-artifact@v8 with: name: 'torchvisionlib-${{ steps.version.outputs.version }}${{ matrix.config.artifact }}' path: build/ diff --git a/DESCRIPTION b/DESCRIPTION index aa0619e..2d8f39f 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -4,28 +4,33 @@ Version: 0.8.0.9000 Authors@R: c( person("Daniel", "Falbel", , "daniel@rstudio.com", role = "aut"), person("Tomasz", "Kalinowski", , "tomasz@posit.co", role = c("ctb", "cre")), - person(family = "RStudio", role = "cph") - ) -Description: Implements additional operators for computer vision models, including - operators necessary for image segmentation and object detection deep learning - models. + person(, "RStudio", role = "cph") + ) +Description: Implements additional operators for computer vision models, + including operators necessary for image segmentation and object + detection deep learning models. License: MIT + file LICENSE -Encoding: UTF-8 -Roxygen: list(markdown = TRUE) -RoxygenNote: 7.2.3 +URL: https://github.com/mlverse/torchvisionlib +BugReports: https://github.com/mlverse/torchvisionlib/issues Depends: R (>= 3.6) -LinkingTo: - Rcpp, - torch Imports: + glue, Rcpp, - torch (>= 0.17.0), rlang, - glue, + torch (>= 0.17.0), withr Suggests: testthat (>= 3.0.0) +LinkingTo: + Rcpp, + torch +Config/Needs/development: Rcpp, torch, torchexport, testthat, roxygen2 +Config/roxygen2/version: 8.0.0 Config/testthat/edition: 3 -URL: https://github.com/mlverse/torchvisionlib -BugReports: https://github.com/mlverse/torchvisionlib/issues +Config/torchvisionlib/binaries: Set TORCHVISIONLIB_URL to a URL, local + ZIP, or directory containing pre-built torchvisionlib binaries. If + unset, defaults to inst/libs/ (installed) or src/ (development). +Encoding: UTF-8 +Roxygen: list(markdown = TRUE) +SystemRequirements: C++17, CMake (>= 3.21), libtorch, libtorchvision diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index f950d6a..a6209ee 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -1,5 +1,5 @@ -project(torchvisionlib) cmake_minimum_required(VERSION 3.21) +project(torchvisionlib) # accommodate differing conventions include(GNUInstallDirs) @@ -27,7 +27,7 @@ else() SET(TORCH_INSTALL_SCRIPT "https://raw.githubusercontent.com/mlverse/torch/main/R/install.R") SET(TORCH_DESCRIPTION "https://raw.githubusercontent.com/mlverse/torch/main/DESCRIPTION") execute_process ( - COMMAND Rscript -e "cat(desc::desc(text=readLines('${TORCH_DESCRIPTION}'))$get('Version'))" + COMMAND Rscript -e "cat(desc::desc(text=readLines('${TORCH_DESCRIPTION}'))\$get('Version'))" OUTPUT_VARIABLE TORCH_R_VERSION ) message(STATUS "TORCH_R_VERSION: ${TORCH_R_VERSION}") @@ -58,20 +58,20 @@ if(WIN32) SET(COMMON_CMAKE_ARGS -DCMAKE_INSTALL_PREFIX=${INSTALL_DIR}) ExternalProject_Add(zlib GIT_REPOSITORY https://github.com/madler/zlib.git - GIT_TAG v1.2.12 + GIT_TAG v1.3.2 PREFIX "${CMAKE_CURRENT_BINARY_DIR}/zlib" CMAKE_ARGS ${COMMON_CMAKE_ARGS} -DCMAKE_ASM_COMPILER=MSVC -DCMAKE_ASM_NASM_COMPILER=MSVC ) ExternalProject_Add(libpng GIT_REPOSITORY https://github.com/glennrp/libpng - GIT_TAG v1.6.37 + GIT_TAG v1.6.55 DEPENDS zlib PREFIX "${CMAKE_CURRENT_BINARY_DIR}/libpng" CMAKE_ARGS ${COMMON_CMAKE_ARGS} -DCMAKE_IGNORE_PATH=C:/rtools44/x86_64-w64-mingw32.static.posix/include -DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH};${CURRENT_BINARY_DIR}/lib ) ExternalProject_Add(libjpeg GIT_REPOSITORY https://github.com/libjpeg-turbo/libjpeg-turbo - GIT_TAG 2.1.2 + GIT_TAG 3.1.3 PREFIX "${CMAKE_CURRENT_BINARY_DIR}/libjpeg" CMAKE_ARGS ${COMMON_CMAKE_ARGS} -DWITH_JPEG8=0 -DWITH_JPEG7=0 -DENABLE_SHARED=0 -DWITH_TURBOJPEG=0 -DCMAKE_IGNORE_PATH=C:/rtools44/x86_64-w64-mingw32.static.posix/include ) @@ -184,7 +184,7 @@ add_custom_command(TARGET torchvisionlib POST_BUILD # binaries bundles and upload them to the GitHub Releases page. set(CPACK_GENERATOR ZIP) execute_process ( - COMMAND Rscript -e "cat(desc::description$new(file = '../../DESCRIPTION')$get('Version'))" + COMMAND Rscript -e "cat(desc::description\$new(file = '../../DESCRIPTION')\$get('Version'))" OUTPUT_VARIABLE CPACK_PACKAGE_VERSION )