LW-DETR multi-scale deformable attention CUDA op

.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/

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

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.

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)

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.8.0
 
 - Updates to support LibTorch v2.8

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)
 }

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).

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
   )
@@ -138,10 +138,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}
@@ -167,6 +179,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)
 
@@ -184,7 +205,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
 )
 

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();