birth

DCNv4_op/DCNv4/functions/__init__.py

@@ -0,0 +1,11 @@
+# ------------------------------------------------------------------------------------------------
+# 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
+# ------------------------------------------------------------------------------------------------
+
+# from .ms_flash_deform_attn_func import FlashMSDeformAttnFunction
+from .flash_deform_attn_func import FlashDeformAttnFunction
+from .dcnv4_func import DCNv4Function

DCNv4_op/DCNv4/functions/dcnv4_func.py

@@ -0,0 +1,129 @@
+# --------------------------------------------------------
+# InternImage
+# Copyright (c) 2022 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import torch
+import math
+import torch.nn.functional as F
+from torch.autograd import Function
+from torch.autograd.function import once_differentiable
+from torch.cuda.amp import custom_bwd, custom_fwd
+from .table import TABLE, BWDTABLE
+
+from DCNv4 import ext
+
+def factors(N):
+    res = []
+    for i in range(1, N+1):
+        if N % i == 0:
+            res.append(i)
+    return res
+
+def findspec(B, H, W, G, C):
+    key = f"{B}x{H}x{W}x{G}x{C}"
+    if key in TABLE:
+        return TABLE[key][0], TABLE[key][1]
+
+    d_stride = 8
+    ms = factors(B*H*W)
+    multiplier = 1
+    for m in ms:
+        if m <= 64 and (m * G * C // d_stride) <= 512:
+            multiplier = m
+    n_thread = multiplier * G * C // d_stride
+    key = f"{B}x{H}x{W}x{G}x{C}"
+    TABLE[key] = (d_stride, n_thread)
+    return d_stride, n_thread
+
+def find_spec_bwd(B, H, W, G, C):
+    key = f"{B}x{H}x{W}x{G}x{C}"
+    if key in BWDTABLE:
+        return BWDTABLE[key][0], BWDTABLE[key][1]
+
+    if C >= 64:
+        d_stride = 2
+    else:
+        d_stride = 1
+
+    ms = factors(B*H*W)
+    multiplier = 1
+    for m in ms:
+        if m <= 64 and (m * G * C // d_stride) <= 256:
+            multiplier = m
+    n_thread = multiplier * G * C // d_stride
+    return d_stride, n_thread
+
+class DCNv4Function(Function):
+    @staticmethod
+    @custom_fwd
+    def forward(
+            ctx, input, offset_mask,
+            kernel_h, kernel_w, stride_h, stride_w,
+            pad_h, pad_w, dilation_h, dilation_w,
+            group, group_channels, offset_scale,
+            im2col_step, remove_center):
+
+        forward_d_stride, forward_block_thread = findspec(input.shape[0], input.shape[1], input.shape[2], group, group_channels)
+        backward_d_stride, backward_block_thread = find_spec_bwd(input.shape[0], input.shape[1], input.shape[2], group, group_channels)
+
+        ctx.kernel_h = kernel_h
+        ctx.kernel_w = kernel_w
+        ctx.stride_h = stride_h
+        ctx.stride_w = stride_w
+        ctx.pad_h = pad_h
+        ctx.pad_w = pad_w
+        ctx.dilation_h = dilation_h
+        ctx.dilation_w = dilation_w
+        ctx.group = group
+        ctx.group_channels = group_channels
+        ctx.offset_scale = offset_scale
+        ctx.im2col_step = im2col_step
+        ctx.remove_center = remove_center
+        ctx.backward_d_stride = backward_d_stride
+        ctx.backward_block_thread = backward_block_thread
+
+        args = [
+            input, offset_mask, kernel_h,
+            kernel_w, stride_h, stride_w, pad_h,
+            pad_w, dilation_h, dilation_w, group,
+            group_channels, offset_scale,
+            ctx.im2col_step,
+            remove_center,
+            forward_d_stride,
+            forward_block_thread,
+            False,
+        ]
+
+        output = ext.dcnv4_forward(*args)
+        ctx.save_for_backward(input, offset_mask)
+
+        return output
+
+    @staticmethod
+    @once_differentiable
+    @custom_bwd
+    def backward(ctx, grad_output):
+        input, offset_mask = ctx.saved_tensors
+
+        args = [
+            input, offset_mask, ctx.kernel_h,
+            ctx.kernel_w, ctx.stride_h, ctx.stride_w, ctx.pad_h,
+            ctx.pad_w, ctx.dilation_h, ctx.dilation_w, ctx.group,
+            ctx.group_channels, ctx.offset_scale, ctx.im2col_step,
+            grad_output.contiguous(), ctx.remove_center,
+            ctx.backward_d_stride, ctx.backward_block_thread,
+            False
+        ]
+
+        grad_input, grad_offset_mask = \
+            ext.dcnv4_backward(*args)
+
+        return grad_input, grad_offset_mask, \
+            None, None, None, None, None, None, None,\
+            None, None, None, None, None, None

DCNv4_op/DCNv4/functions/flash_deform_attn_func.py

@@ -0,0 +1,114 @@
+# ------------------------------------------------------------------------------------------------
+# 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
+# ------------------------------------------------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import torch
+import torch.nn.functional as F
+from torch.autograd import Function
+from torch.autograd.function import once_differentiable
+import numpy as np
+
+from DCNv4 import ext
+
+shm_size_dict = {
+    "8.0": 163000,
+    "8.6": 99000,
+    "8.7": 163000,
+    "8.9": 99000,
+    "9.0": 227000,
+    "7.5": 64000,
+    "7.0": 96000,
+}
+
+cuda_capability = f"{torch.cuda.get_device_properties(0).major}.{torch.cuda.get_device_properties(0).minor}"
+
+if cuda_capability not in shm_size_dict:
+    raise NotImplementedError
+
+shm_size_cap = shm_size_dict[cuda_capability]
+
+def factors(N):
+    res = []
+    for i in range(1, N+1):
+        if N % i == 0:
+            res.append(i)
+    return res
+
+def findspec(B, Q, G, C):
+    d_stride = 8
+    ms = factors(B*Q)
+    multiplier = 1
+    for m in ms:
+        if m <= 64 and (m * G * C // d_stride) <= 512:
+            multiplier = m
+    n_thread = multiplier * G * C // d_stride
+    return d_stride, n_thread
+
+def findspec_bwd(B, Q, G, C):
+    if C >= 64:
+        d_stride = 2
+    else:
+        d_stride = 1
+    
+    ms = factors(B*Q)
+    multiplier = 1
+    for m in ms:
+        if m <= 64 and (m * G * C // d_stride) <= 256:
+            multiplier = m
+    n_thread = multiplier * G * C // d_stride
+    return d_stride, n_thread
+
+class FlashDeformAttnFunction(Function):
+    @staticmethod
+    @torch.autocast("cuda", enabled=True, dtype=torch.float16)
+    def forward(
+        ctx, value, value_spatial_shapes, value_level_start_index, 
+        sampling_loc_attn, im2col_step, K=8
+    ):
+
+        ctx.im2col_step = im2col_step
+        ctx.K = K
+        d_stride, blockthread = findspec(value.shape[0], sampling_loc_attn.shape[1], value.shape[2], value.shape[3])
+        d_stride_backward, blockthread_backward = findspec_bwd(value.shape[0], sampling_loc_attn.shape[1], value.shape[2], value.shape[3])
+
+        ctx.d_stride_backward = d_stride_backward
+        ctx.blockthread_backward = blockthread_backward
+
+        output = ext.flash_deform_attn_forward(
+            value,
+            value_spatial_shapes,
+            value_level_start_index,
+            sampling_loc_attn,
+            ctx.im2col_step,
+            K,
+            d_stride,
+            blockthread,
+        )
+        ctx.save_for_backward(value, value_spatial_shapes, value_level_start_index, sampling_loc_attn)
+        return output
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        value, value_spatial_shapes, value_level_start_index, sampling_loc_attn = ctx.saved_tensors
+        grad_value, grad_sampling_loc_attn = ext.flash_deform_attn_backward(
+            value,
+            value_spatial_shapes,
+            value_level_start_index,
+            sampling_loc_attn,
+            grad_output.contiguous(),
+            ctx.im2col_step,
+            ctx.K,
+            ctx.d_stride_backward,
+            ctx.blockthread_backward,
+        )
+
+        return grad_value, None, None, grad_sampling_loc_attn, None, None

DCNv4_op/DCNv4/modules/__init__.py

@@ -0,0 +1,10 @@
+# ------------------------------------------------------------------------------------------------
+# 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
+# ------------------------------------------------------------------------------------------------
+
+from .flash_deform_attn import FlashDeformAttn
+from .dcnv4 import DCNv4

DCNv4_op/DCNv4/modules/dcnv4.py

@@ -0,0 +1,152 @@
+# --------------------------------------------------------
+# Deformable Convolution v4
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import math
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.init import xavier_uniform_, constant_
+from ..functions import DCNv4Function
+
+class CenterFeatureScaleModule(nn.Module):
+    def forward(self,
+                query,
+                center_feature_scale_proj_weight,
+                center_feature_scale_proj_bias):
+        center_feature_scale = F.linear(query,
+                                        weight=center_feature_scale_proj_weight,
+                                        bias=center_feature_scale_proj_bias).sigmoid()
+        return center_feature_scale
+
+class DCNv4(nn.Module):
+    def __init__(
+            self,
+            channels=64,
+            kernel_size=3,
+            stride=1,
+            pad=1,
+            dilation=1,
+            group=4,
+            offset_scale=1.0,
+            dw_kernel_size=None,
+            center_feature_scale=False,
+            remove_center=False,
+            output_bias=True,
+            without_pointwise=False,
+            **kwargs):
+        """
+        DCNv4 Module
+        :param channels
+        :param kernel_size
+        :param stride
+        :param pad
+        :param dilation
+        :param group
+        :param offset_scale
+        :param act_layer
+        :param norm_layer
+        """
+        super().__init__()
+        if channels % group != 0:
+            raise ValueError(
+                f'channels must be divisible by group, but got {channels} and {group}')
+        _d_per_group = channels // group
+
+        # you'd better set _d_per_group to a power of 2 which is more efficient in our CUDA implementation
+        assert _d_per_group % 16 == 0
+
+        self.offset_scale = offset_scale
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.dilation = dilation
+        self.pad = pad
+        self.group = group
+        self.group_channels = channels // group
+        self.offset_scale = offset_scale
+        self.dw_kernel_size = dw_kernel_size
+        self.center_feature_scale = center_feature_scale
+        self.remove_center = int(remove_center)
+        self.without_pointwise = without_pointwise
+
+        self.K =  group * (kernel_size * kernel_size - self.remove_center)
+        if dw_kernel_size is not None:
+            self.offset_mask_dw = nn.Conv2d(channels, channels, dw_kernel_size, stride=1, padding=(dw_kernel_size - 1) // 2, groups=channels)
+        self.offset_mask = nn.Linear(channels, int(math.ceil((self.K * 3)/8)*8))
+        if not without_pointwise:
+            self.value_proj = nn.Linear(channels, channels)
+            self.output_proj = nn.Linear(channels, channels, bias=output_bias)
+        self._reset_parameters()
+
+        if center_feature_scale:
+            self.center_feature_scale_proj_weight = nn.Parameter(
+                torch.zeros((group, channels), dtype=torch.float))
+            self.center_feature_scale_proj_bias = nn.Parameter(
+                torch.tensor(0.0, dtype=torch.float).view((1,)).repeat(group, ))
+            self.center_feature_scale_module = CenterFeatureScaleModule()
+
+    def _reset_parameters(self):
+        constant_(self.offset_mask.weight.data, 0.)
+        constant_(self.offset_mask.bias.data, 0.)
+        if not self.without_pointwise:
+            xavier_uniform_(self.value_proj.weight.data)
+            constant_(self.value_proj.bias.data, 0.)
+            xavier_uniform_(self.output_proj.weight.data)
+            if self.output_proj.bias is not None:
+                constant_(self.output_proj.bias.data, 0.)
+
+    def forward(self, input, shape=None):
+        """
+        :param query                       (N, H, W, C)
+        :return output                     (N, H, W, C)
+        """
+        N, L, C = input.shape
+        if shape is not None:
+            H, W = shape
+        else:
+            H, W = int(L**0.5), int(L**0.5)
+
+
+        x = input
+        if not self.without_pointwise:
+            x = self.value_proj(x)
+        x = x.reshape(N, H, W, -1)
+        if self.dw_kernel_size is not None:
+            offset_mask_input = self.offset_mask_dw(input.view(N, H, W, C).permute(0, 3, 1, 2))
+            offset_mask_input = offset_mask_input.permute(0, 2, 3, 1).view(N, L, C)
+        else:
+            offset_mask_input = input
+        offset_mask = self.offset_mask(offset_mask_input).reshape(N, H, W, -1)
+
+        x_proj = x
+
+        x = DCNv4Function.apply(
+            x, offset_mask,
+            self.kernel_size, self.kernel_size,
+            self.stride, self.stride,
+            self.pad, self.pad,
+            self.dilation, self.dilation,
+            self.group, self.group_channels,
+            self.offset_scale,
+            256,
+            self.remove_center
+            )
+
+        x = x.view(N, L, -1)
+        if self.center_feature_scale:
+            center_feature_scale = self.center_feature_scale_module(
+                x, self.center_feature_scale_proj_weight, self.center_feature_scale_proj_bias)
+            center_feature_scale = center_feature_scale[..., None].repeat(
+                1, 1, 1, 1, self.channels // self.group).flatten(-2)
+            x = x * (1 - center_feature_scale) + x_proj * center_feature_scale
+        if not self.without_pointwise:
+            x = self.output_proj(x)
+        return x
+

DCNv4_op/DCNv4/modules/flash_deform_attn.py

@@ -0,0 +1,141 @@
+# ------------------------------------------------------------------------------------------------
+# 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
+# ------------------------------------------------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import warnings
+import math
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.init import xavier_uniform_, constant_
+
+from ..functions import FlashDeformAttnFunction
+
+
+def _is_power_of_2(n):
+    if (not isinstance(n, int)) or (n < 0):
+        raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n)))
+    return (n & (n - 1) == 0) and n != 0
+
+
+class FlashDeformAttn(nn.Module):
+    def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):
+        """
+        Multi-Scale Deformable Attention Module
+        :param d_model      hidden dimension
+        :param n_levels     number of feature levels
+        :param n_heads      number of attention heads
+        :param n_points     number of sampling points per attention head per feature level
+        """
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads, but got {} and {}".format(d_model, n_heads))
+        _d_per_head = d_model // n_heads
+        # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation
+        if not _is_power_of_2(_d_per_head):
+            warnings.warn(
+                "You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 "
+                "which is more efficient in our CUDA implementation."
+            )
+
+        self.im2col_step = 64
+
+        self.d_model = d_model
+        self.n_levels = n_levels
+        self.n_heads = n_heads
+        self.n_points = n_points
+
+        self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)
+        self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
+        self.value_proj = nn.Linear(d_model, d_model)
+        self.output_proj = nn.Linear(d_model, d_model)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        constant_(self.sampling_offsets.weight.data, 0.0)
+        thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
+        grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+        grid_init = (
+            (grid_init / grid_init.abs().max(-1, keepdim=True)[0])
+            .view(self.n_heads, 1, 1, 2)
+            .repeat(1, self.n_levels, self.n_points, 1)
+        )
+        for i in range(self.n_points):
+            grid_init[:, :, i, :] *= i + 1
+        with torch.no_grad():
+            self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
+        constant_(self.attention_weights.weight.data, 0.0)
+        constant_(self.attention_weights.bias.data, 0.0)
+        xavier_uniform_(self.value_proj.weight.data)
+        constant_(self.value_proj.bias.data, 0.0)
+        xavier_uniform_(self.output_proj.weight.data)
+        constant_(self.output_proj.bias.data, 0.0)
+
+    def forward(
+        self,
+        query,
+        reference_points,
+        input_flatten,
+        input_spatial_shapes,
+        input_level_start_index,
+        input_padding_mask=None,
+    ):
+        """
+        :param query                       (N, Length_{query}, C)
+        :param reference_points            (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area
+                                        or (N, Length_{query}, n_levels, 4), add additional (w, h) to form reference boxes
+        :param input_flatten               (N, \sum_{l=0}^{L-1} H_l \cdot W_l, C)
+        :param input_spatial_shapes        (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]
+        :param input_level_start_index     (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]
+        :param input_padding_mask          (N, \sum_{l=0}^{L-1} H_l \cdot W_l), True for padding elements, False for non-padding elements
+
+        :return output                     (N, Length_{query}, C)
+        """
+        N, Len_q, _ = query.shape
+        N, Len_in, _ = input_flatten.shape
+        assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in
+
+        value = self.value_proj(input_flatten)
+        if input_padding_mask is not None:
+            value = value.masked_fill(input_padding_mask[..., None], float(0))
+        value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)
+        sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)
+        attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)
+        attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)
+        # N, Len_q, n_heads, n_levels, n_points, 2
+        if reference_points.shape[-1] == 2:
+            offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)
+            sampling_locations = (
+                reference_points[:, :, None, :, None, :]
+                + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
+            )
+        elif reference_points.shape[-1] == 4:
+            sampling_locations = (
+                reference_points[:, :, None, :, None, :2]
+                + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5
+            )
+        else:
+            raise ValueError(
+                "Last dim of reference_points must be 2 or 4, but get {} instead.".format(reference_points.shape[-1])
+            )
+
+        output = FlashDeformAttnFunction.apply(
+            value,
+            input_spatial_shapes,
+            input_level_start_index,
+            sampling_locations,
+            attention_weights,
+            self.im2col_step,
+        )
+        output = self.output_proj(output)
+        return output

DCNv4_op/MANIFEST.in

@@ -0,0 +1,2 @@
+include src/*
+include src/cuda/*

DCNv4_op/__init__.py

@@ -0,0 +1,2 @@
+from .functions import DCNv4Function, FlashDeformAttnFunction
+from .modules import DCNv4, FlashDeformAttn

DCNv4_op/make.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+# ------------------------------------------------------------------------------------------------
+# 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
+# ------------------------------------------------------------------------------------------------
+
+python setup.py build install

DCNv4_op/scripts/find_best.py

@@ -0,0 +1,61 @@
+import json
+import argparse
+class LineParser:
+    def __init__(self) -> None:
+        self.data = {}
+
+    def parse(self, line):
+        def startswith(line, lst):
+            for ele in lst:
+                if line.startswith(ele):
+                    return True
+            return False
+
+        if not startswith(line, ['1', '2', '3', '4', '5', '6', '7', '8', '9']):
+            return
+
+        eles = line.strip().split()
+        key = eles[0]
+        if key not in self.data:
+            self.data[key] = []
+        
+        self.data[key].append([eles[1], float(eles[2])])
+    
+    def sort(self):
+        for k, v in self.data.items():
+            nv = sorted(v, key=lambda x: x[1])
+            self.data[k] = nv
+    
+    def display_best(self):
+        for k, v in self.data.items():
+            print(f'{k} \t {v[0][0]} \t {v[0][1]:.4f} \t {v[1][0]} \t {v[1][1]:.4f}')
+    
+    def display_best_python(self, output):
+        res = {}
+        def parse(spec):
+            d_stride = int(spec.split('/')[0])
+            thread = int(spec.split('/')[1].split('(')[0])
+            m = int(spec.split('(')[1].split(')')[0])
+            return d_stride, thread, m
+
+        for k, v in self.data.items():
+            res[k] = parse(v[0][0])
+
+        with open(output, "w") as f:
+            json.dump(res, f, indent=4)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input', type=str)
+    parser.add_argument('--output', type=str)
+    args = parser.parse_args()
+
+    with open(args.input) as f:
+        lines = f.readlines()
+
+    lineparser = LineParser()
+    for line in lines:
+        lineparser.parse(line)
+    lineparser.sort()
+    lineparser.display_best()
+    lineparser.display_best_python(args.output)

DCNv4_op/scripts/search_bwd.sh

@@ -0,0 +1,2 @@
+python search_dcnv4_bwd_engine.py > res_bwd.txt
+python find_best.py --input res_bwd.txt --output table_bwd.py

DCNv4_op/scripts/search_dcnv4.py

@@ -0,0 +1,131 @@
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import math
+import torch
+import torch.nn as nn
+import math
+from torch.autograd import gradcheck
+import pandas as pd
+from easydict import EasyDict as edict
+import argparse
+
+from torch.cuda import Event
+
+from functions.dcnv3_func import DCNv3Function, dcnv3_core_pytorch
+from functions.dcnv4_func import DCNv4Function
+torch.set_printoptions(threshold=10000)
+
+torch.manual_seed(3)
+
+
+#@torch.no_grad()
+def speed_test(func, args, inputs, name='Unknown'):
+
+    tic = Event(enable_timing=True)
+    toc = Event(enable_timing=True)
+    # warmup
+    for i in range(args.warmup_num):
+        func(*inputs)
+
+    total_time = 0
+    tic.record()
+    for i in range(args.test_num):
+        o = func(*inputs)
+        torch.cuda.synchronize()
+    toc.record()
+
+    avg_time = tic.elapsed_time(toc) / args.test_num
+    # print(
+        # f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+@torch.no_grad()
+def test(N, H_in, W_in, M, D, spec=None):
+    Kh, Kw = 3, 3
+    remove_center = False
+    P = Kh * Kw - remove_center
+    offset_scale = 2.0
+    pad = 1
+    dilation = 1
+    stride = 1
+    H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+    W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+    input = torch.rand(N, H_in, W_in, M*D).cuda()
+    # print(input.shape)
+    offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*2
+    # offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*0
+    mask_origin = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask_origin = mask_origin.half()
+    mask = mask_origin
+    # mask = torch.nn.functional.softmax(mask_origin, dim=-1)
+    offset_mask = torch.cat([offset.unflatten(-1, (M, P * 2)), mask_origin.detach()], dim=-1).flatten(-2)
+
+    im2col_step = 128
+
+    input = input.half()
+    offset = offset.half()
+    mask = mask.half()
+    offset_mask = offset_mask.half()
+
+    dcnv3_args = [
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center,
+    ]
+    output_pytorch = DCNv3Function.apply(*dcnv3_args)
+
+    input1 = input.detach()
+
+    def pad(om):
+        padded_zero = int(math.ceil(om.shape[3]/8)*8) - om.shape[3]
+        padded = torch.zeros(om.shape[0], om.shape[1], om.shape[2], padded_zero).to(om)
+        return torch.cat([om, padded], dim=-1)
+
+    dcnv4_args = [
+        input1, pad(offset_mask),
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, 
+        spec[0], spec[1], 2, None
+        # 8, 512, 2, 256
+    ]
+    output_flash_cuda = DCNv4Function.apply(*dcnv4_args)
+
+    fwdok = torch.allclose(output_flash_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+    max_abs_err = (output_flash_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_flash_cuda - output_pytorch).abs() /
+                   (output_pytorch.abs()+ 1e-3)).max()
+    # print('>>> forward half')
+    # print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+    if not fwdok:
+        print(f"Wrong: {N}x{H_in}x{W_in}x{M}x{D} \t {spec[0]}/{spec[1]}({spec[2]})")
+        return
+    # assert(fwdok)
+
+    test_args = edict({'warmup_num': 10000, 'test_num': 10000})
+    
+    exp_time_dcnv4 = speed_test(DCNv4Function.apply, test_args, dcnv4_args, name='exp')
+    torch.cuda.synchronize()
+    print(f"{N}x{H_in}x{W_in}x{M}x{D} \t {spec[0]}/{spec[1]}({spec[2]}): {exp_time_dcnv4}")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n", type=int)
+    parser.add_argument("--h", type=int)
+    parser.add_argument("--w", type=int)
+    parser.add_argument("--g", type=int)
+    parser.add_argument("--c", type=int)
+    parser.add_argument("--dstride", type=int)
+    parser.add_argument("--blockthread", type=int)
+    parser.add_argument("--multiplier", type=int)
+    args = parser.parse_args()
+    test(args.n, args.h, args.w, args.g, args.c, (args.dstride, args.blockthread, args.multiplier))
+
+

DCNv4_op/scripts/search_dcnv4_bwd.py

@@ -0,0 +1,200 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import torch
+import torch.nn as nn
+import math
+from torch.autograd import gradcheck
+import pandas as pd
+from easydict import EasyDict as edict
+import argparse
+
+from torch.cuda import Event
+
+from functions import DCNv4Function, DCNv3Function
+torch.set_printoptions(threshold=10000)
+
+
+
+torch.manual_seed(3)
+
+def speed_test_backward(func, args, inputs, name='Unknown'):
+    # warmup
+    # for i in range(args.warmup_num):
+    #     o = func(*inputs)
+    #     o.sum().backward()
+
+    total_time = 0
+    len_input = len(inputs)
+    for i in range(args.warmup_num + args.test_num):
+        tic = Event(enable_timing=True)
+        toc = Event(enable_timing=True)
+        inputs[0] = inputs[0].detach()
+        inputs[0].requires_grad = True
+        if len_input > 1 and isinstance(inputs[1], torch.Tensor):
+            inputs[1] = inputs[1].detach()
+            inputs[1].requires_grad = True
+        if len_input > 2 and isinstance(inputs[2], torch.Tensor):
+            inputs[2] = inputs[2].detach()
+            inputs[2].requires_grad = True
+
+        o = func(*inputs)
+        torch.cuda.synchronize()
+        tic.record()
+        o.sum().backward()
+        toc.record()
+        torch.cuda.synchronize()
+        _time = tic.elapsed_time(toc)
+        if i >= args.warmup_num:
+            total_time += _time
+        o = o.detach()
+
+    # toc.record()
+    # torch.cuda.synchronize()
+
+    avg_time = total_time / args.test_num
+    #print(
+    #    f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+# @torch.no_grad()
+def test(N=64, H_in=32, W_in=32, M=4, D=16, spec=None):
+    """
+    64x56x56x128(G=4)
+    2 64: 3.66
+    - offset_mask collection write 3.4022
+    - offset_mask collection 3.1968
+    
+    """
+    Kh, Kw = 3, 3
+    remove_center = False
+    P = Kh * Kw - remove_center
+    offset_scale = 2.0
+    pad = 1
+    dilation = 1
+    stride = 1
+    H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+    W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+    additions = [None, None, spec[0], spec[1], False]
+    input = torch.rand(N, H_in, W_in, M*D).cuda() * 10
+    #offset = torch.rand(N, H_out, W_out, M*P*2).cuda() * 0
+    offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*2
+    mask_origin = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask_origin = mask_origin.half()
+    mask_origin.requires_grad = True
+    # offset_mask = torch.cat([offset.unflatten(-1, (M, P, 2)), mask_origin.detach().unsqueeze(-1)], dim=-1).flatten(-3)
+    # mask /= mask.sum(-1, keepdim=True)
+    # mask = torch.nn.functional.softmax(mask_origin, dim=-1, dtype=torch.float32)
+    mask = mask_origin
+    # mask = mask.reshape(N, H_out, W_out, M*P)
+    # offset_mask = torch.cat([offset.unflatten(-1, (M, P, 2)), mask.detach().unsqueeze(-1)], dim=-1).flatten(-3)
+    offset_mask = torch.cat([offset.detach().unflatten(-1, (M, P * 2)), mask_origin.detach()], dim=-1).flatten(-2)
+
+    im2col_step = 128
+
+    input = input.half()
+    offset = offset.half()
+    mask = mask.half()
+    input.requires_grad = True
+    offset.requires_grad = True
+    # mask.requires_grad = True
+    output_pytorch = DCNv3Function.apply(
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center)#.detach().cpu()
+    (output_pytorch.sum()/10).backward()
+
+    def pad(om):
+        padded_zero = int(math.ceil(om.shape[3]/8)*8) - om.shape[3]
+        padded = torch.zeros(om.shape[0], om.shape[1], om.shape[2], padded_zero).to(om)
+        return torch.cat([om, padded], dim=-1)
+
+    # value_offset_mask = input.detach()
+    input1 = input.detach()
+    input1.requires_grad = True
+    offset_mask = offset_mask.half()
+    offset_mask.requires_grad = True
+    # offset_mask1.requires_grad = True
+    torch.cuda.profiler.cudart().cudaProfilerStart()
+    output_flash_cuda = DCNv4Function.apply(
+        input1, offset_mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, *additions)#.detach().cpu()
+    (output_flash_cuda.sum()/10).backward()
+    torch.cuda.profiler.cudart().cudaProfilerStop()
+
+    input_grad = input.grad
+    input2_grad = input1.grad
+    bwdok = torch.allclose(input_grad.float(), input2_grad.float(), rtol=1e-2, atol=1e-3)
+    rel_err = (input_grad.abs() - input2_grad.abs())/(input_grad.abs()+1e-3)
+
+    offset_grad1 = offset.grad
+    offset_grad2 = offset_mask.grad.reshape(N, H_out, W_out, M, P*3)[..., :P*2].reshape(N, H_out, W_out, M*P*2)
+
+    bwdok2 = torch.allclose(offset_grad1.float(), offset_grad2.float(), rtol=1e-2, atol=1e-3)
+    rel_err = (offset_grad1 - offset_grad2).abs() / (offset_grad1.abs()+1e-3)
+
+    mask_grad1 = mask_origin.grad
+    mask_grad2 = offset_mask.grad.reshape(N, H_out, W_out, M, P*3)[..., P*2:].reshape(N, H_out, W_out, M, P)
+
+    bwdok3 = torch.allclose(mask_grad1, mask_grad2, rtol=1e-2, atol=1e-3)
+    rel_err = (mask_grad1 - mask_grad2).abs() / (mask_grad1.abs()+1e-3)
+
+    fwdok = torch.allclose(output_flash_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+    max_abs_err = (output_flash_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_flash_cuda - output_pytorch).abs() /
+                   (output_pytorch.abs()+ 1e-3)).max()
+    if not (bwdok and bwdok2 and bwdok3):
+        print(f"Wrong: {N}x{H_in}x{W_in}x{M}x{D} \t {spec[0]}/{spec[1]}({spec[2]})")
+        return
+    # fn_args = [
+    #     input,
+    #     offset,
+    #     mask,
+    #     Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+    #     im2col_step, remove_center
+    # ]
+
+    flash_dcn_fn_args = [
+        input1,
+        offset_mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, *additions
+    ]
+
+    test_args = edict({'warmup_num': 1000, 'test_num': 1000})
+    try:
+        exp_time = speed_test_backward(DCNv4Function.apply, test_args, flash_dcn_fn_args, name='exp')
+    except:
+        print(f"Wrong: {N}x{H_in}x{W_in}x{M}x{D} \t {spec[0]}/{spec[1]}({spec[2]})")
+        return
+
+    torch.cuda.synchronize()
+    print(f"{N}x{H_in}x{W_in}x{M}x{D} \t {spec[0]}/{spec[1]}({spec[2]}): {exp_time}")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n", type=int)
+    parser.add_argument("--h", type=int)
+    parser.add_argument("--w", type=int)
+    parser.add_argument("--g", type=int)
+    parser.add_argument("--c", type=int)
+    parser.add_argument("--dstride", type=int)
+    parser.add_argument("--blockthread", type=int)
+    parser.add_argument("--multiplier", type=int)
+    args = parser.parse_args()
+    test(args.n, args.h, args.w, args.g, args.c, (args.dstride, args.blockthread, args.multiplier))
+
+

DCNv4_op/scripts/search_dcnv4_bwd_engine.py

@@ -0,0 +1,24 @@
+import os
+
+def factors(N):
+    res = []
+    for i in range(1, N+1):
+        if N % i == 0:
+            res.append(i)
+    return res
+
+if __name__ == '__main__':
+    BATCH=64
+    for N, Hin, Win in [(BATCH, 56, 56), (BATCH, 28, 28), (BATCH, 14, 14), (BATCH, 7, 7), 
+                        (1, 200, 320), (1, 100, 160), (1, 50, 80), (1, 25, 40), (1, 64, 64)]:
+        for group_channel in [16, 32, 64]:
+            for group in [4, 5, 6, 7, 8]:
+                for d_stride in [1, 2, 4]:
+                    for m in factors(N*Hin*Win):
+                        if m > 64:
+                            break
+                        block_thread = group * (group_channel//d_stride) * m
+                        if block_thread > 1024:
+                            break
+                        cmd = f"python search_dcnv4_bwd.py --n {N} --h {Hin} --w {Win} --g {group} --c {group_channel} --dstride {d_stride} --blockthread {block_thread} --multiplier {m}"
+                        os.system(cmd)

DCNv4_op/scripts/search_dcnv4_engine.py

@@ -0,0 +1,24 @@
+import os
+
+def factors(N):
+    res = []
+    for i in range(1, N+1):
+        if N % i == 0:
+            res.append(i)
+    return res
+
+if __name__ == '__main__':
+    BATCH=64
+    for group_channel in [16, 32, 64]:
+        for group in [4, 5, 6, 7, 8]:
+            for N, Hin, Win in [(BATCH, 56, 56), (BATCH, 28, 28), (BATCH, 14, 14), (BATCH, 7, 7), 
+                                (1, 200, 320), (1, 100, 160), (1, 50, 80), (1, 25, 40), (1, 64, 64)]:
+                for d_stride in [2, 4, 8, 16]:
+                    for m in factors(N*Hin*Win):
+                        if m > 64:
+                            break
+                        block_thread = group * (group_channel//d_stride) * m
+                        if block_thread > 1024:
+                            break
+                        cmd = f"python search_dcnv4.py --n {N} --h {Hin} --w {Win} --g {group} --c {group_channel} --dstride {d_stride} --blockthread {block_thread} --multiplier {m}"
+                        os.system(cmd)

DCNv4_op/scripts/search_fwd.sh

@@ -0,0 +1,2 @@
+python search_dcnv4_engine.py > res.txt
+python find_best.py --input res.txt --output table.py

DCNv4_op/scripts/test_dcnv4.py

@@ -0,0 +1,144 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import torch
+import torch.nn as nn
+import math
+from torch.autograd import gradcheck
+import pandas as pd
+from easydict import EasyDict as edict
+
+from torch.cuda import Event
+
+from functions.dcnv3_func import DCNv3Function, dcnv3_core_pytorch
+from functions.dcnv4_func import DCNv4Function
+torch.set_printoptions(threshold=10000)
+
+H_in, W_in = 56, 56
+N, M, D = 64, 4, 32
+
+# H_in, W_in = 28, 28
+# N, M, D = 64, 8, 32
+
+# H_in, W_in = 14, 14
+# N, M, D = 64, 16, 32
+
+# H_in, W_in = 7, 7
+# N, M, D = 64, 32, 32
+
+# H_in, W_in = 8, 8
+# N, M, D = 128, 4, 16
+
+
+Kh, Kw = 3, 3
+remove_center = False
+P = Kh * Kw - remove_center
+offset_scale = 2.0
+pad = 1
+dilation = 1
+stride = 1
+H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+torch.manual_seed(3)
+
+#@torch.no_grad()
+def speed_test(func, args, inputs, name='Unknown'):
+
+    tic = Event(enable_timing=True)
+    toc = Event(enable_timing=True)
+    # warmup
+    for i in range(args.warmup_num):
+        func(*inputs)
+
+    total_time = 0
+    tic.record()
+    for i in range(args.test_num):
+        o = func(*inputs)
+        torch.cuda.synchronize()
+    toc.record()
+
+    avg_time = tic.elapsed_time(toc) / args.test_num
+    print(
+        f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+@torch.no_grad()
+def check_forward_equal_with_pytorch_half():
+    input = torch.rand(N, H_in, W_in, M*D).cuda()
+    print(input.shape)
+    offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*10
+    # offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*0
+    mask_origin = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask_origin = mask_origin.half()
+    mask = mask_origin
+    # mask = torch.nn.functional.softmax(mask_origin, dim=-1)
+    offset_mask = torch.cat([offset.unflatten(-1, (M, P * 2)), mask_origin.detach()], dim=-1).flatten(-2)
+
+    im2col_step = 128
+
+    input = input.half()
+    offset = offset.half()
+    mask = mask.half()
+    offset_mask = offset_mask.half()
+
+    dcnv3_args = [
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center,
+    ]
+    output_pytorch = DCNv3Function.apply(*dcnv3_args)
+
+    input1 = input.detach()
+
+    def pad(om):
+        padded_zero = int(math.ceil(om.shape[3]/8)*8) - om.shape[3]
+        padded = torch.zeros(om.shape[0], om.shape[1], om.shape[2], padded_zero).to(om)
+        return torch.cat([om, padded], dim=-1)
+
+    dcnv4_args = [
+        input1, pad(offset_mask),
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, 8, 512, 2, 256, True, True,
+    ]
+    output_flash_cuda = DCNv4Function.apply(*dcnv4_args)
+
+    fwdok = torch.allclose(output_flash_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+    max_abs_err = (output_flash_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_flash_cuda - output_pytorch).abs() /
+                   (output_pytorch.abs()+ 1e-3)).max()
+    print('>>> forward half')
+    print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+    assert(fwdok)
+
+    test_args = edict({'warmup_num': 1000, 'test_num': 1000})
+    
+    exp_time_dcnv4 = speed_test(DCNv4Function.apply, test_args, dcnv4_args, name='exp')
+    exp_time_dcnv3 = speed_test(DCNv3Function.apply, test_args, dcnv3_args, name='exp')
+    torch.cuda.synchronize()
+
+    results = [{}]
+    results[0]['dcnv3_time'] = exp_time_dcnv3
+    results[0]['dcnv4_time'] = exp_time_dcnv4
+    columns = list(results[0].keys())
+
+    outputs = pd.DataFrame(results, columns=columns)
+    with pd.option_context(
+        'display.max_rows', None, 'display.max_columns', None,
+        'display.max_colwidth', None, 'display.width', None,
+        'display.precision', 4, ):
+        print(outputs)
+
+
+if __name__ == '__main__':
+    check_forward_equal_with_pytorch_half()

DCNv4_op/scripts/test_dcnv4_bwd.py

@@ -0,0 +1,221 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import torch
+import torch.nn as nn
+import math
+from torch.autograd import gradcheck
+import pandas as pd
+from easydict import EasyDict as edict
+
+from torch.cuda import Event
+
+from functions import DCNv4Function, DCNv3Function
+torch.set_printoptions(threshold=10000)
+
+H_in, W_in = 56, 56
+N, M, D = 64, 4, 32
+
+# H_in, W_in = 28, 28
+# N, M, D = 64, 16, 16
+
+# H_in, W_in = 14, 14
+# N, M, D = 64, 32, 16
+
+
+# H_in, W_in = 7, 7
+# N, M, D = 64, 64, 16
+
+# H_in, W_in = 8, 8
+# N, M, D = 128, 4, 16
+
+
+Kh, Kw = 3, 3
+remove_center = False
+P = Kh * Kw - remove_center
+offset_scale = 2.0
+pad = 1
+dilation = 1
+stride = 1
+H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+torch.manual_seed(3)
+
+def speed_test_backward(func, args, inputs, name='Unknown'):
+    # warmup
+    # for i in range(args.warmup_num):
+    #     o = func(*inputs)
+    #     o.sum().backward()
+
+    total_time = 0
+    len_input = len(inputs)
+    for i in range(args.warmup_num + args.test_num):
+        tic = Event(enable_timing=True)
+        toc = Event(enable_timing=True)
+        inputs[0] = inputs[0].detach()
+        inputs[0].requires_grad = True
+        if len_input > 1 and isinstance(inputs[1], torch.Tensor):
+            inputs[1] = inputs[1].detach()
+            inputs[1].requires_grad = True
+        if len_input > 2 and isinstance(inputs[2], torch.Tensor):
+            inputs[2] = inputs[2].detach()
+            inputs[2].requires_grad = True
+
+        o = func(*inputs)
+        torch.cuda.synchronize()
+        tic.record()
+        o.sum().backward()
+        toc.record()
+        torch.cuda.synchronize()
+        _time = tic.elapsed_time(toc)
+        if i >= args.warmup_num:
+            total_time += _time
+        o = o.detach()
+
+    # toc.record()
+    # torch.cuda.synchronize()
+
+    avg_time = total_time / args.test_num
+    #print(
+    #    f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+# @torch.no_grad()
+def check_forward_equal_with_pytorch_half():
+    """
+    64x56x56x128(G=4)
+    2 64: 3.66
+    - offset_mask collection write 3.4022
+    - offset_mask collection 3.1968
+    
+    """
+    additions = [8, 128, 2, 256, False]
+    input = torch.rand(N, H_in, W_in, M*D).cuda() * 10
+    print(input.shape)
+    #offset = torch.rand(N, H_out, W_out, M*P*2).cuda() * 0
+    offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*2
+    mask_origin = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask_origin = mask_origin.half()
+    mask_origin.requires_grad = True
+    # offset_mask = torch.cat([offset.unflatten(-1, (M, P, 2)), mask_origin.detach().unsqueeze(-1)], dim=-1).flatten(-3)
+    # mask /= mask.sum(-1, keepdim=True)
+    # mask = torch.nn.functional.softmax(mask_origin, dim=-1, dtype=torch.float32)
+    mask = mask_origin
+    # mask = mask.reshape(N, H_out, W_out, M*P)
+    # offset_mask = torch.cat([offset.unflatten(-1, (M, P, 2)), mask.detach().unsqueeze(-1)], dim=-1).flatten(-3)
+    offset_mask = torch.cat([offset.detach().unflatten(-1, (M, P * 2)), mask_origin.detach()], dim=-1).flatten(-2)
+
+    im2col_step = 128
+
+    input = input.half()
+    offset = offset.half()
+    mask = mask.half()
+    input.requires_grad = True
+    offset.requires_grad = True
+    # mask.requires_grad = True
+    output_pytorch = DCNv3Function.apply(
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center)#.detach().cpu()
+    (output_pytorch.sum()/10).backward()
+
+    def pad(om):
+        padded_zero = int(math.ceil(om.shape[3]/8)*8) - om.shape[3]
+        padded = torch.zeros(om.shape[0], om.shape[1], om.shape[2], padded_zero).to(om)
+        return torch.cat([om, padded], dim=-1)
+
+    # value_offset_mask = input.detach()
+    input1 = input.detach()
+    input1.requires_grad = True
+    offset_mask = offset_mask.half()
+    offset_mask.requires_grad = True
+    # offset_mask1.requires_grad = True
+    torch.cuda.profiler.cudart().cudaProfilerStart()
+    output_flash_cuda = DCNv4Function.apply(
+        input1, offset_mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, *additions)#.detach().cpu()
+    (output_flash_cuda.sum()/10).backward()
+    torch.cuda.profiler.cudart().cudaProfilerStop()
+
+    input_grad = input.grad
+    input2_grad = input1.grad
+    bwdok = torch.allclose(input_grad.float(), input2_grad.float(), rtol=1e-2, atol=1e-3)
+    print("bwdok")
+    print(bwdok)
+    rel_err = (input_grad.abs() - input2_grad.abs())/(input_grad.abs()+1e-3)
+    print(rel_err.max())
+
+    offset_grad1 = offset.grad
+    offset_grad2 = offset_mask.grad.reshape(N, H_out, W_out, M, P*3)[..., :P*2].reshape(N, H_out, W_out, M*P*2)
+    # print(offset_grad1)
+    # print("====================")
+    # print(offset_grad2)
+    bwdok2 = torch.allclose(offset_grad1.float(), offset_grad2.float(), rtol=1e-2, atol=1e-3)
+    print("bwdok2")
+    print(bwdok2)
+    rel_err = (offset_grad1 - offset_grad2).abs() / (offset_grad1.abs()+1e-3)
+    print(rel_err.max())
+
+    mask_grad1 = mask_origin.grad
+    mask_grad2 = offset_mask.grad.reshape(N, H_out, W_out, M, P*3)[..., P*2:].reshape(N, H_out, W_out, M, P)
+
+    bwdok3 = torch.allclose(mask_grad1, mask_grad2, rtol=1e-2, atol=1e-3)
+    print("bwdok3")
+    print(bwdok3)
+    rel_err = (mask_grad1 - mask_grad2).abs() / (mask_grad1.abs()+1e-3)
+    print(rel_err.max())
+
+    fwdok = torch.allclose(output_flash_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+    max_abs_err = (output_flash_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_flash_cuda - output_pytorch).abs() /
+                   (output_pytorch.abs()+ 1e-3)).max()
+    print('>>> forward half')
+    print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+    fn_args = [
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center
+    ]
+
+    flash_dcn_fn_args = [
+        input1,
+        offset_mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, *additions
+    ]
+
+    test_args = edict({'warmup_num': 1000, 'test_num': 1000})
+    exp_time = speed_test_backward(
+            DCNv4Function.apply, test_args, flash_dcn_fn_args, name='exp')
+    exp_time_base = speed_test_backward(
+            DCNv3Function.apply, test_args, fn_args, name='exp')
+
+    results = [{}]
+    results[0]['time'] = exp_time
+    results[0]['time_base'] = exp_time_base
+    columns = list(results[0].keys())
+
+    outputs = pd.DataFrame(results, columns=columns)
+    with pd.option_context(
+        'display.max_rows', None, 'display.max_columns', None,
+        'display.max_colwidth', None, 'display.width', None,
+        'display.precision', 4, ):
+        print(outputs)
+
+
+if __name__ == '__main__':
+    check_forward_equal_with_pytorch_half()

DCNv4_op/scripts/test_flash_deform_attn.py

@@ -0,0 +1,174 @@
+# ------------------------------------------------------------------------------------------------
+# 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
+# ------------------------------------------------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+from easydict import EasyDict as edict
+from torch.cuda import Event
+import pandas as pd
+
+import time
+import torch
+import torch.nn as nn
+from torch.autograd import gradcheck
+
+from functions import MSDeformAttnFunction, FlashDeformAttnFunction, ms_deform_attn_core_pytorch
+
+
+# N, M, D = 1, 4, 8
+# # Lq, L, P = 2, 2, 2
+# # shapes = torch.as_tensor([(6, 4), (3, 2)], dtype=torch.long).cuda()
+# Lq, L, P = 1, 2, 8
+# shapes = torch.as_tensor([(8, 16), (4, 8)], dtype=torch.long).cuda()
+
+# N, M, D = 1, 8, 32
+# # Lq, L, P = 2, 2, 2
+# # shapes = torch.as_tensor([(6, 4), (3, 2)], dtype=torch.long).cuda()
+# Lq, L, P = 300, 4, 4
+# # shapes = torch.as_tensor([(134, 151), (67, 76), (34, 38), (17, 19)], dtype=torch.long).cuda()
+# # shapes = torch.as_tensor([(134, 151), (67, 76), (34, 38), (16, 16)], dtype=torch.long).cuda()
+# # shapes = torch.as_tensor([(134, 151), (67, 76), (34, 38), (17, 19)], dtype=torch.long).cuda()
+# # shapes = torch.as_tensor([(17, 19), (4, 4)], dtype=torch.long).cuda()
+# shapes = torch.as_tensor([(100, 151), (50, 76), (25, 38), (13, 19)], dtype=torch.long).cuda()
+# # shapes = torch.as_tensor([(110, 151)], dtype=torch.long).cuda()
+
+# B:6
+# H:232
+# W:400
+# G:5
+# D: 16
+# channels: 80
+# kernel: 3 points = 3 * 3
+# num_split = 45 = kernel *kernel * G
+
+H = 256
+W = 256
+N, M, D = 1, 8, 32
+Lq, L, P = 100*152, 4, 8
+
+shapes = torch.Tensor([[100, 152], [ 50,  76], [ 25,  38], [ 13,  19]]).long().cuda()
+
+# x = x.reshape([B, H*W, G, D + self.num_split * 3])
+# shapes = torch.as_tensor([(H, W)], dtype=torch.long).cuda()
+# shapes = torch.as_tensor([(H, W), (H // 2, W // 2)], dtype=torch.long).cuda()
+# shapes = torch.as_tensor([(H, W), (H // 2, W // 2), (H // 4, W // 4), (H // 8, W // 8)], dtype=torch.long).cuda()
+
+level_start_index = torch.cat((shapes.new_zeros((1,)), shapes.prod(1).cumsum(0)[:-1]))
+S = sum([(H * W).item() for H, W in shapes])
+print(S)
+
+def get_reference_points(spatial_shapes, device):
+    reference_points_list = []
+    for lvl, (H_, W_) in enumerate(spatial_shapes):
+
+        ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
+                                        torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
+        ref_y = ref_y.reshape(-1)[None] / (H_)
+        ref_x = ref_x.reshape(-1)[None] / (W_)
+        ref = torch.stack((ref_x, ref_y), -1)
+        reference_points_list.append(ref)
+    reference_points = torch.cat(reference_points_list, 1)
+    # reference_points = reference_points[:, :, None] * valid_ratios[:, None]
+    return reference_points
+
+
+torch.manual_seed(3)
+
+@torch.no_grad()
+def speed_test(func, args, inputs, name='Unknown'):
+
+    tic = Event(enable_timing=True)
+    toc = Event(enable_timing=True)
+    # warmup
+    for i in range(args.warmup_num):
+        func(*inputs)
+
+    tic.record()
+    for i in range(args.test_num):
+        func(*inputs)
+    toc.record()
+    torch.cuda.synchronize()
+
+    avg_time = tic.elapsed_time(toc) / args.test_num
+    print(
+        f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+
+@torch.no_grad()
+def check_forward_equal_with_pytorch_half():
+    value = torch.rand(N, S, M, D).cuda() * 0.01
+    # offset = (torch.rand(N, Lq, M, L, P, 2).cuda() * 2 - 1) / 10
+    sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
+    attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
+    sampling_loc_attn = torch.cat([sampling_locations.reshape(N, Lq, M, L*P*2), attention_weights.reshape(N, Lq, M, L*P)], dim=-1)
+    attention_weights = torch.nn.functional.softmax(attention_weights.flatten(-2, -1), dim=-1).unflatten(-1, (L, P))
+
+
+    im2col_step = 128
+
+    flash_fn_args = (
+        value.half(),
+        shapes,
+        level_start_index,
+        sampling_loc_attn.half(),
+        im2col_step,
+        P, 16
+    )
+    output_cuda = (
+        FlashDeformAttnFunction.apply(*flash_fn_args)
+        .detach()
+        .cpu()
+    ).double()
+    
+    fn_args = (
+        value,
+        shapes,
+        level_start_index,
+        sampling_locations,
+        attention_weights,
+        im2col_step,
+    )
+
+    output_pytorch = (
+        MSDeformAttnFunction.apply(*fn_args)
+        .detach().double()
+        .cpu()
+    )
+
+    max_abs_err = (output_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()
+    fwdok = torch.allclose(output_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+
+    print(
+        f"* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}"
+    )
+
+    test_args = edict({'warmup_num': 1000, 'test_num': 1000})
+    exp_time_base = speed_test(
+            MSDeformAttnFunction.apply, test_args, fn_args, name='exp')
+    exp_time = speed_test(
+            FlashDeformAttnFunction.apply, test_args, flash_fn_args, name='exp')
+
+    results = [{}]
+    results[0]['time'] = exp_time
+    results[0]['time_base'] = exp_time_base
+    columns = list(results[0].keys())
+
+    outputs = pd.DataFrame(results, columns=columns)
+    with pd.option_context(
+        'display.max_rows', None, 'display.max_columns', None,
+        'display.max_colwidth', None, 'display.width', None,
+        'display.precision', 4, ):
+        print(outputs)
+
+
+if __name__ == "__main__":
+    check_forward_equal_with_pytorch_half()
+

DCNv4_op/scripts/test_flash_deform_attn_backward.py

@@ -0,0 +1,194 @@
+# ------------------------------------------------------------------------------------------------
+# 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
+# ------------------------------------------------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+from easydict import EasyDict as edict
+from torch.cuda import Event
+import pandas as pd
+
+import time
+import torch
+import torch.nn as nn
+from torch.autograd import gradcheck
+
+from functions import MSDeformAttnFunction, ms_deform_attn_core_pytorch, FlashDeformAttnFunction
+
+
+H = 256
+W = 256
+N, M, D = 1, 8, 16
+Lq, L, P = H * W, 1, 8
+
+# x = x.reshape([B, H*W, G, D + self.num_split * 3])
+shapes = torch.as_tensor([(H, W)], dtype=torch.long).cuda()
+# shapes = torch.as_tensor([(H, W), (H // 2, W // 2)], dtype=torch.long).cuda()
+# shapes = torch.as_tensor([(H, W), (H // 2, W // 2), (H // 4, W // 4), (H // 8, W // 8)], dtype=torch.long).cuda()
+
+H = 256
+W = 256
+N, M, D = 1, 8, 32
+Lq, L, P = 100*152, 4, 8
+
+shapes = torch.Tensor([[100, 152], [ 50,  76], [ 25,  38], [ 13,  19]]).long().cuda()
+
+level_start_index = torch.cat((shapes.new_zeros((1,)), shapes.prod(1).cumsum(0)[:-1]))
+S = sum([(H * W).item() for H, W in shapes])
+
+def get_reference_points(spatial_shapes, device):
+    reference_points_list = []
+    for lvl, (H_, W_) in enumerate(spatial_shapes):
+
+        ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
+                                        torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
+        ref_y = ref_y.reshape(-1)[None] / (H_)
+        ref_x = ref_x.reshape(-1)[None] / (W_)
+        ref = torch.stack((ref_x, ref_y), -1)
+        reference_points_list.append(ref)
+    reference_points = torch.cat(reference_points_list, 1)
+    # reference_points = reference_points[:, :, None] * valid_ratios[:, None]
+    return reference_points
+
+
+torch.manual_seed(3)
+
+@torch.no_grad()
+def speed_test(func, args, inputs, name='Unknown'):
+
+    tic = Event(enable_timing=True)
+    toc = Event(enable_timing=True)
+    # warmup
+    for i in range(args.warmup_num):
+        func(*inputs)
+
+    tic.record()
+    for i in range(args.test_num):
+        func(*inputs)
+    toc.record()
+    torch.cuda.synchronize()
+
+    avg_time = tic.elapsed_time(toc) / args.test_num
+    print(
+        f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+
+def check_forward_equal_with_pytorch_half():
+    value = torch.rand(N, S, M, D).cuda() * 0.01
+    offset = (torch.rand(N, Lq, M, L, P, 2).cuda() * 2 - 1) / 10
+    sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
+    attention_weights_origin = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
+    attention_weights_origin.requires_grad = True
+    sampling_loc_attn = torch.cat([sampling_locations.detach().reshape(N, Lq, M, L*P*2), attention_weights_origin.detach().reshape(N, Lq, M, L*P)], dim=-1)
+
+    attention_weights = torch.nn.functional.softmax(attention_weights_origin.flatten(-2, -1), dim=-1).unflatten(-1, (L, P))
+
+
+    im2col_step = 128
+
+    value.requires_grad = True
+    sampling_loc_attn.requires_grad = True
+    output_cuda = (
+        FlashDeformAttnFunction.apply(
+            value.float(),
+            shapes,
+            level_start_index,
+            sampling_loc_attn.float(),
+            im2col_step,
+        )
+    )
+    (output_cuda.float().sum()/10).backward()
+
+
+    value1 = value.detach()
+    value1.requires_grad = True
+    sampling_locations.requires_grad = True
+    #attention_weights.requires_grad = True
+    output_pytorch = (
+        ms_deform_attn_core_pytorch(value1, shapes, sampling_locations, attention_weights)
+    )
+    (output_pytorch.sum()/10).backward()
+
+    max_abs_err = (output_cuda.float() - output_pytorch).abs().max()
+    max_rel_err = ((output_cuda.float() - output_pytorch).abs() / output_pytorch.abs()).max()
+    fwdok = torch.allclose(output_cuda.float(), output_pytorch, rtol=1e-2, atol=1e-3)
+    print(fwdok)
+    print(max_abs_err, max_rel_err)
+    #exit()
+
+    bwdok1 = torch.allclose(value.grad, value1.grad, rtol=1e-2, atol=1e-3)
+    print(bwdok1)
+    # rel_err = (sampling_locations.grad - sampling_loc_attn.grad[..., :L*P*2].reshape(*sampling_locations.shape)).abs()/(sampling_locations.grad.abs()+1e-3)
+    # print(rel_err.max())
+
+    locgrad1 = sampling_locations.grad
+    locgrad2 = sampling_loc_attn.grad[..., :L*P*2].reshape(*sampling_locations.shape)
+    bwdok2 = torch.allclose(locgrad1, locgrad2, rtol=1e-2, atol=1e-3)
+    print(bwdok2)
+    rel_err = (locgrad1 - locgrad2).abs()/(locgrad1.abs()+1e-3)
+    print(rel_err.max())
+
+    attngrad1 = attention_weights_origin.grad
+    attngrad2 = sampling_loc_attn.grad[..., L*P*2:].reshape(*attention_weights_origin.shape)
+    bwdok3 = torch.allclose(locgrad1, locgrad2, rtol=1e-2, atol=1e-3)
+    print(bwdok3)
+    rel_err = (attngrad1 - attngrad2).abs()/(attngrad1.abs()+1e-3)
+    print(rel_err.max())
+    exit()
+    #exit()
+
+    # pdb.set_trace()
+    max_abs_err = (output_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()
+    fwdok = torch.allclose(output_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+
+    print(
+        f"* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}"
+    )
+
+
+    fn_args = (
+        value,
+        shapes,
+        level_start_index,
+        sampling_locations,
+        attention_weights,
+        im2col_step,
+    )
+
+    flash_dcn_fn_args = (
+        value.half(),
+        shapes,
+        level_start_index,
+        sampling_loc_attn.half(),
+        im2col_step,
+    )
+
+
+    test_args = edict({'warmup_num': 50, 'test_num': 100})
+    exp_time = speed_test(
+            FlashMSDeformAttnFunction.apply, test_args, flash_dcn_fn_args, name='exp')
+    exp_time_base = speed_test(
+            MSDeformAttnFunction.apply, test_args, fn_args, name='exp')
+
+    results = [{}]
+    results[0]['time'] = exp_time
+    results[0]['time_base'] = exp_time_base
+    columns = list(results[0].keys())
+
+    outputs = pd.DataFrame(results, columns=columns)
+    with pd.option_context(
+        'display.max_rows', None, 'display.max_columns', None,
+        'display.max_colwidth', None, 'display.width', None,
+        'display.precision', 4, ):
+        print(outputs)
+
+
+if __name__ == "__main__":
+    check_forward_equal_with_pytorch_half()

DCNv4_op/setup.py

@@ -0,0 +1,72 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable Convolution v4
+# Copyright (c) 2024 OpenGVLab
+# 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
+# ------------------------------------------------------------------------------------------------
+
+import os
+import glob
+
+import torch
+
+from torch.utils.cpp_extension import CUDA_HOME
+from torch.utils.cpp_extension import CppExtension
+from torch.utils.cpp_extension import CUDAExtension
+
+from setuptools import find_packages
+from setuptools import setup
+
+requirements = ["torch", "torchvision"]
+
+def get_extensions():
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    extensions_dir = os.path.join(this_dir, "src")
+
+    main_file = glob.glob(os.path.join(extensions_dir, "*.cpp"))
+    source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp"))
+    source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu"))
+
+    sources = main_file + source_cpu
+    extension = CppExtension
+    extra_compile_args = {"cxx": []}
+    define_macros = []
+
+    if torch.cuda.is_available() and CUDA_HOME is not None:
+        extension = CUDAExtension
+        sources += source_cuda
+        define_macros += [("WITH_CUDA", None)]
+        extra_compile_args["nvcc"] = [
+            "-DCUDA_HAS_FP16=1",
+            "-D__CUDA_NO_HALF_OPERATORS__",
+            "-D__CUDA_NO_HALF_CONVERSIONS__",
+            "-D__CUDA_NO_HALF2_OPERATORS__",
+            "-O3",
+        ]
+    else:
+        raise NotImplementedError('Cuda is not available')
+
+    sources = [os.path.join(extensions_dir, s) for s in sources]
+    include_dirs = [extensions_dir]
+    ext_modules = [
+        extension(
+            "DCNv4.ext",
+            sources,
+            include_dirs=include_dirs,
+            define_macros=define_macros,
+            extra_compile_args=extra_compile_args,
+        )
+    ]
+    return ext_modules
+
+setup(
+    name="DCNv4",
+    version="1.0.0",
+    author="Yuwen Xiong, Feng Wang",
+    url="",
+    description="PyTorch Wrapper for CUDA Functions of DCNv4",
+    packages=['DCNv4', 'DCNv4/functions', 'DCNv4/modules'],
+    ext_modules=get_extensions(),
+    cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension},
+)

DCNv4_op/src/cuda/common.h

@@ -0,0 +1,212 @@
+#ifndef FMSDACOMMON
+#define FMSDACOMMON
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+constexpr int kWarpSize = 32;
+#define opmath_t at::opmath_type<scalar_t>
+
+inline int GET_BLOCKS(const int N, const int num_threads) {
+  return (N + num_threads - 1) / num_threads;
+}
+
+#define CUDA_KERNEL_LOOP(i, n)                                                 \
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);                 \
+       i += blockDim.x * gridDim.x)
+
+inline bool check_backward_warpp(int d_stride, int D){
+  int n_group_threads = D / d_stride;
+  return (n_group_threads <= kWarpSize) && (kWarpSize % n_group_threads == 0);
+}
+
+template <typename scalar_t, typename transfer_t, int c_per_thread>
+__device__ void ms_deform_attn_im2col_bilinear(
+    opmath_t out_reg_array[], const scalar_t *&p_value, const int &height,
+    const int &width, const opmath_t &h_px, const opmath_t &w_px,
+    const opmath_t &attn, const int &w_stride, const int &base_ptr) {
+
+  const int h_low = floor(h_px);
+  const int w_low = floor(w_px);
+  const int h_high = h_low + 1;
+  const int w_high = w_low + 1;
+  const opmath_t lh = h_px - h_low;
+  const opmath_t lw = w_px - w_low;
+  const opmath_t hh = 1 - lh;
+  const opmath_t hw = 1 - lw;
+
+  const opmath_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+
+  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;
+
+  int idx1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+  int idx2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+  int idx3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+  int idx4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+
+  scalar_t v1_array[c_per_thread] = {0.};
+  scalar_t v2_array[c_per_thread] = {0.};
+  scalar_t v3_array[c_per_thread] = {0.};
+  scalar_t v4_array[c_per_thread] = {0.};
+
+  if (h_low >= 0 && w_low >= 0) {
+    auto p1 = p_value + idx1;
+    *(transfer_t *)(v1_array) = *(transfer_t *)(p1);
+  }
+
+  if (h_low >= 0 && w_high < width) {
+    auto p2 = p_value + idx2;
+    *(transfer_t *)(v2_array) = *(transfer_t *)(p2);
+  }
+  if (h_high < height && w_low >= 0) {
+    auto p3 = p_value + idx3;
+    *(transfer_t *)(v3_array) = *(transfer_t *)(p3);
+  }
+  if (h_high < height && w_high < width) {
+    auto p4 = p_value + idx4;
+    *(transfer_t *)(v4_array) = *(transfer_t *)(p4);
+  }
+#pragma unroll
+  for (int i = 0; i < c_per_thread; i++) {
+    out_reg_array[i] +=
+        (opmath_t)attn *
+        (w1 * (opmath_t)v1_array[i] + w2 * (opmath_t)v2_array[i] +
+         w3 * (opmath_t)v3_array[i] + w4 * (opmath_t)v4_array[i]);
+  }
+}
+
+template <typename scalar_t, typename transfer_t, int c_per_thread>
+__device__ void ms_deform_attn_col2im_bilinear(
+    const scalar_t *&p_value, const int &height, const int &width,
+    const opmath_t &h_px, const opmath_t &w_px, const opmath_t &attn,
+    const int &w_stride, const int &base_ptr, const opmath_t offset_scale_h,
+    const opmath_t offset_scale_w, const scalar_t *&top_grad,
+    opmath_t *&grad_im, opmath_t *grad_offset) {
+
+  const int h_low = floor(h_px);
+  const int w_low = floor(w_px);
+  const int h_high = h_low + 1;
+  const int w_high = w_low + 1;
+  const opmath_t lh = h_px - h_low;
+  const opmath_t lw = w_px - w_low;
+  const opmath_t hh = 1 - lh;
+  const opmath_t hw = 1 - lw;
+
+  const opmath_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+
+  scalar_t _top_grad_array[c_per_thread] = {0.};
+  *(transfer_t *)(_top_grad_array) = *(transfer_t *)(top_grad);
+
+  opmath_t top_grad_array[c_per_thread] = {0.};
+  for (int i = 0; i < c_per_thread; ++i) {
+    top_grad_array[i] = (opmath_t)(_top_grad_array[i]);
+  }
+
+  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;
+
+  int idx1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+  int idx2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+  int idx3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+  int idx4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+
+  scalar_t v1_array[c_per_thread] = {0.};
+  scalar_t v2_array[c_per_thread] = {0.};
+  scalar_t v3_array[c_per_thread] = {0.};
+  scalar_t v4_array[c_per_thread] = {0.};
+
+  opmath_t grad_h_weight[c_per_thread] = {0.};
+  opmath_t grad_w_weight[c_per_thread] = {0.};
+
+  if (h_low >= 0 && w_low >= 0) {
+    auto p1 = p_value + idx1;
+    *(transfer_t *)(v1_array) = *(transfer_t *)(p1);
+#pragma unroll
+    for (int i = 0; i < c_per_thread; ++i) {
+      grad_h_weight[i] -= hw * v1_array[i];
+      grad_w_weight[i] -= hh * v1_array[i];
+      atomicAdd(grad_im + idx1 + i, top_grad_array[i] * attn * w1);
+    }
+  }
+
+  if (h_low >= 0 && w_high < width) {
+    auto p2 = p_value + idx2;
+    *(transfer_t *)(v2_array) = *(transfer_t *)(p2);
+#pragma unroll
+    for (int i = 0; i < c_per_thread; ++i) {
+      grad_h_weight[i] -= lw * v2_array[i];
+      grad_w_weight[i] += hh * v2_array[i];
+      atomicAdd(grad_im + idx2 + i, top_grad_array[i] * attn * w2);
+    }
+  }
+  if (h_high < height && w_low >= 0) {
+    auto p3 = p_value + idx3;
+    *(transfer_t *)(v3_array) = *(transfer_t *)(p3);
+#pragma unroll
+    for (int i = 0; i < c_per_thread; ++i) {
+      grad_h_weight[i] += hw * v3_array[i];
+      grad_w_weight[i] -= lh * v3_array[i];
+      atomicAdd(grad_im + idx3 + i, top_grad_array[i] * attn * w3);
+    }
+  }
+  if (h_high < height && w_high < width) {
+    auto p4 = p_value + idx4;
+    *(transfer_t *)(v4_array) = *(transfer_t *)(p4);
+#pragma unroll
+    for (int i = 0; i < c_per_thread; ++i) {
+      grad_h_weight[i] += lw * v4_array[i];
+      grad_w_weight[i] += lh * v4_array[i];
+      atomicAdd(grad_im + idx4 + i, top_grad_array[i] * attn * w4);
+    }
+  }
+
+  opmath_t _grad_offset_x = 0;
+  opmath_t _grad_offset_y = 0;
+#pragma unroll
+  for (int i = 0; i < c_per_thread; ++i) {
+    _grad_offset_x +=
+        grad_w_weight[i] * top_grad_array[i]; // channel aware term
+    _grad_offset_y +=
+        grad_h_weight[i] * top_grad_array[i]; // channel aware term
+  }
+  _grad_offset_x *= (offset_scale_w * attn); // channel shared term
+  _grad_offset_y *= (offset_scale_h * attn); // channel shared term
+
+  *grad_offset = _grad_offset_x;
+  *(grad_offset + 1) = _grad_offset_y;
+
+  opmath_t current_val;
+  opmath_t _grad_offset_z = 0;
+#pragma unroll
+  for (int i = 0; i < c_per_thread; i++) {
+    current_val = (opmath_t)(w1 * v1_array[i] + w2 * v2_array[i] +
+                             w3 * v3_array[i] + w4 * v4_array[i]);
+    _grad_offset_z += current_val * top_grad_array[i];
+  }
+  *(grad_offset + 2) = _grad_offset_z;
+}
+
+
+
+#endif

DCNv4_op/src/cuda/dcnv4_col2im_cuda.cuh

@@ -0,0 +1,561 @@
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include "common.h"
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K,
+          bool softmax>
+__global__ void backward_kernel_dcn(
+    const scalar_t *p_value, const scalar_t *p_offset,
+    const scalar_t *grad_output, const int G, const int D, const int Q,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int height_in, const int width_in,
+    const int height_out, const int width_out, const opmath_t offset_scale,
+    const int remove_center, const int block_multiplier, opmath_t *grad_im,
+    opmath_t *grad_offset, const int padded_offset_dim) {
+
+  extern __shared__ char _s[];
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+  constexpr int li = 0;
+
+  opmath_t *const cache_g_mask_before_softmax = (opmath_t *)(_s); // mG x K
+  opmath_t *const cache_grad_offset =
+      (opmath_t *)(cache_g_mask_before_softmax +
+                   block_multiplier * G * K); // mG x blockDim.x x 3
+  opmath_t *const p_mask_shm =
+      (opmath_t *)(cache_grad_offset + block_multiplier * G * blockDim.x * 3) +
+      (threadIdx.z * G + gi) * K;
+
+  const scalar_t *p_offset_ptr = p_offset + (bi*Q + qi)*padded_offset_dim + gi*K*3;
+
+  const int mask_length = K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+
+  const scalar_t *top_grad = grad_output + ((bi * Q + qi) * G + gi) * D + di_s;
+
+  __syncthreads();
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) = *(
+        scalar_t *)(p_offset_ptr + K * 2 + num_thread * num_iter + threadIdx.x);
+  }
+
+  if (softmax) {
+    __syncthreads();
+    // transfer offset from global memory to shared memory >
+
+    // Calculate softmax over L and K
+    if (threadIdx.x == 0) { // gi != 0, di = 0, li = 0
+      opmath_t softmax_max = -1e100;
+      opmath_t softmax_sum = 0.0;
+
+      // get max
+      for (int j = 0; j < K; j++) {
+        softmax_max = max(softmax_max, p_mask_shm[j]);
+      }
+
+      // get sumexp
+      for (int j = 0; j < K; j++) {
+        opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+        p_mask_shm[j] = exp_results;
+        softmax_sum += exp_results;
+      }
+
+      // normalize
+      for (int j = 0; j < K; j++) {
+        p_mask_shm[j] /= softmax_sum;
+      }
+    }
+
+    __syncthreads();
+  }
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+  const scalar_t *p_value_ptr =
+      p_value + (bi * (height_in * width_in)) * (G * D);
+  opmath_t *grad_im_ptr = grad_im + (bi * (height_in * width_in)) * (G * D);
+
+  const int p0_w = ((dilation_w * (kernel_w - 1)) >> 1) - pad_w +
+                   (qi % width_out) * stride_w;
+  const int p0_h = ((dilation_h * (kernel_h - 1)) >> 1) - pad_h +
+                   (qi / width_out) * stride_h;
+  const opmath_t p0_w_ =
+      p0_w - ((dilation_w * (kernel_w - 1)) >> 1) * offset_scale;
+  const opmath_t p0_h_ =
+      p0_h - ((dilation_h * (kernel_h - 1)) >> 1) * offset_scale;
+  const int center_h = kernel_h / 2;
+  const int center_w = kernel_w / 2;
+
+  grad_offset += (bi*Q + qi)*padded_offset_dim + gi*K*3;
+  opmath_t *grad_offset_softmax = grad_offset + K * 2;
+
+  int cache_grad_off_idx =
+      ((threadIdx.z * G + threadIdx.y) * blockDim.x + threadIdx.x) * 3;
+  for (int i = 0; i < kernel_w; ++i) {
+    for (int j = 0; j < kernel_h; ++j) {
+      if (i != center_w || j != center_h || !remove_center) {
+        const opmath_t w_im =
+            p0_w_ + (i * dilation_w + (opmath_t)p_offset_ptr[offset_idx]) *
+                        offset_scale;
+        const opmath_t h_im =
+            p0_h_ + (j * dilation_h + (opmath_t)p_offset_ptr[offset_idx + 1]) *
+                        offset_scale;
+        const opmath_t attn = p_mask_shm[mask_idx];
+        cache_grad_offset[cache_grad_off_idx] = 0;
+        cache_grad_offset[cache_grad_off_idx + 1] = 0;
+        cache_grad_offset[cache_grad_off_idx + 2] = 0;
+
+        if (h_im > -1 && w_im > -1 && h_im < height_in && w_im < width_in) {
+          ms_deform_attn_col2im_bilinear<scalar_t, transfer_t, d_stride>(
+              p_value_ptr, height_in, width_in, h_im, w_im, attn, w_stride,
+              base_ptr, offset_scale, offset_scale, top_grad, grad_im_ptr,
+              cache_grad_offset + cache_grad_off_idx);
+        }
+
+        // aggregated across different channel for offset
+        
+        __syncthreads();
+        if (threadIdx.x == 0) { //
+          int _didx = (threadIdx.z * G + threadIdx.y) * blockDim.x * 3;
+          opmath_t _grad_w = cache_grad_offset[_didx],
+                   _grad_h = cache_grad_offset[_didx + 1],
+                   _grad_a = cache_grad_offset[_didx + 2];
+
+          for (int c_id = 1; c_id < blockDim.x; ++c_id) {
+            _grad_w += cache_grad_offset[_didx + 3 * c_id];
+            _grad_h += cache_grad_offset[_didx + 3 * c_id + 1];
+            _grad_a += cache_grad_offset[_didx + 3 * c_id + 2];
+          }
+
+          *(grad_offset) = _grad_w;     // B x H x W x G x L x K x 3
+          *(grad_offset + 1) = _grad_h; // B x H x W x G x L x K x 3
+          if (softmax) {
+            cache_g_mask_before_softmax[(threadIdx.z * G + threadIdx.y) * K +
+                                        mask_idx] = _grad_a * attn;
+          }
+          else{
+            grad_offset_softmax[mask_idx] = _grad_a;
+          }
+        }
+        __syncthreads();
+
+        offset_idx += 2;
+        mask_idx += 1;
+        grad_offset += 2;
+      }
+    }
+  }
+  // backward for softmax
+  if(softmax){
+    if (threadIdx.x == 0) {
+      const opmath_t* group_g_mask = cache_g_mask_before_softmax + (threadIdx.z*G + threadIdx.y)*K;
+      #pragma unroll
+      for (int i = 0; i < K; ++i) {
+        opmath_t sum = 0.;
+        for (int j = 0; j < K; ++j) {
+          sum += group_g_mask[j]; // dL/di * di/dj
+        }
+        *(grad_offset_softmax) = group_g_mask[i] - p_mask_shm[i] * sum;
+
+        grad_offset_softmax += 1;
+      }
+    }
+    __syncthreads();
+  }
+}
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K,
+          bool softmax>
+__global__ void backward_kernel_dcn_warp_primitive(
+    const scalar_t *p_value, const scalar_t *p_offset,
+    const scalar_t *grad_output, const int G, const int D, const int Q,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int height_in, const int width_in,
+    const int height_out, const int width_out, const opmath_t offset_scale,
+    const int remove_center, const int block_multiplier, opmath_t *grad_im,
+    opmath_t *grad_offset, const int padded_offset_dim) {
+
+  extern __shared__ char _s[];
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+
+  constexpr int li = 0;
+  const int tid = (threadIdx.z * blockDim.y + threadIdx.y)*blockDim.x + threadIdx.x;
+
+  const int lane_id = tid % kWarpSize;
+
+  // find the position of current group in the current warp
+  const int group_per_warp = kWarpSize / blockDim.x;
+  const int group_in_warp_id = (threadIdx.z * G + threadIdx.y) % group_per_warp;
+  const unsigned lane_mask = ((1 << blockDim.x) - 1) << (group_in_warp_id * blockDim.x);
+
+  opmath_t *const p_mask_shm = (opmath_t *)(_s) + (threadIdx.z * G + gi) * K;
+  opmath_t *cache_g_mask_before_softmax = (opmath_t *)((opmath_t *)(_s) + block_multiplier * G * K) +
+                                          (threadIdx.z*G+gi)*K; // only used by threadIdx.x = 0
+
+  const scalar_t *p_offset_ptr = p_offset + (bi*Q + qi)*padded_offset_dim + gi*K*3;
+
+  const int mask_length = K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+
+  const scalar_t *top_grad = grad_output + ((bi * Q + qi) * G + gi) * D + di_s;
+
+  __syncthreads();
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) = *(
+        scalar_t *)(p_offset_ptr + K * 2 + num_thread * num_iter + threadIdx.x);
+  }
+
+  if (softmax) {
+    __syncthreads();
+    // transfer offset from global memory to shared memory >
+
+    // Calculate softmax over L and K
+    if (threadIdx.x == 0) { // gi != 0, di = 0, li = 0
+      opmath_t softmax_max = -1e100;
+      opmath_t softmax_sum = 0.0;
+
+      // get max
+      for (int j = 0; j < K; j++) {
+        softmax_max = max(softmax_max, p_mask_shm[j]);
+      }
+
+      // get sumexp
+      for (int j = 0; j < K; j++) {
+        opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+        p_mask_shm[j] = exp_results;
+        softmax_sum += exp_results;
+      }
+
+      // normalize
+      for (int j = 0; j < K; j++) {
+        p_mask_shm[j] /= softmax_sum;
+      }
+    }
+
+    __syncthreads();
+  }
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+  const scalar_t *p_value_ptr =
+      p_value + (bi * (height_in * width_in)) * (G * D);
+  opmath_t *grad_im_ptr = grad_im + (bi * (height_in * width_in)) * (G * D);
+
+  const int p0_w = ((dilation_w * (kernel_w - 1)) >> 1) - pad_w +
+                   (qi % width_out) * stride_w;
+  const int p0_h = ((dilation_h * (kernel_h - 1)) >> 1) - pad_h +
+                   (qi / width_out) * stride_h;
+  const opmath_t p0_w_ =
+      p0_w - ((dilation_w * (kernel_w - 1)) >> 1) * offset_scale;
+  const opmath_t p0_h_ =
+      p0_h - ((dilation_h * (kernel_h - 1)) >> 1) * offset_scale;
+  const int center_h = kernel_h / 2;
+  const int center_w = kernel_w / 2;
+
+  grad_offset += (bi * Q + qi)*padded_offset_dim + gi*K*3;
+  opmath_t *grad_offset_softmax = grad_offset + K * 2;
+
+  int cache_grad_off_idx =
+      ((threadIdx.z * G + threadIdx.y) * blockDim.x + threadIdx.x) * 3;
+
+  opmath_t reg_grad_offset[3] = {0.};
+  for (int i = 0; i < kernel_w; ++i) {
+    for (int j = 0; j < kernel_h; ++j) {
+      if (i != center_w || j != center_h || !remove_center) {
+        const opmath_t w_im =
+            p0_w_ + (i * dilation_w + (opmath_t)p_offset_ptr[offset_idx]) *
+                        offset_scale;
+        const opmath_t h_im =
+            p0_h_ + (j * dilation_h + (opmath_t)p_offset_ptr[offset_idx + 1]) *
+                        offset_scale;
+        const opmath_t attn = p_mask_shm[mask_idx];
+        reg_grad_offset[0] = 0;
+        reg_grad_offset[1] = 0;
+        reg_grad_offset[2] = 0;
+
+        if (h_im > -1 && w_im > -1 && h_im < height_in && w_im < width_in) {
+          ms_deform_attn_col2im_bilinear<scalar_t, transfer_t, d_stride>(
+              p_value_ptr, height_in, width_in, h_im, w_im, attn, w_stride,
+              base_ptr, offset_scale, offset_scale, top_grad, grad_im_ptr,
+              reg_grad_offset);
+        }
+
+        // aggregated across different channel for offset
+        for (uint32_t offset = blockDim.x>>1; offset > 0; offset >>= 1){
+          reg_grad_offset[0] += __shfl_down_sync(lane_mask, reg_grad_offset[0], offset);
+          reg_grad_offset[1] += __shfl_down_sync(lane_mask, reg_grad_offset[1], offset);
+          reg_grad_offset[2] += __shfl_down_sync(lane_mask, reg_grad_offset[2], offset);
+        }
+
+        if (threadIdx.x == 0) { //
+          *(grad_offset) = reg_grad_offset[0];     // B x H x W x G x L x K x 3
+          *(grad_offset + 1) = reg_grad_offset[1]; // B x H x W x G x L x K x 3
+          if (softmax) {
+            cache_g_mask_before_softmax[mask_idx] = reg_grad_offset[2] * attn;
+          }
+          else{
+            grad_offset_softmax[mask_idx] = reg_grad_offset[2];
+          }
+        }
+        offset_idx += 2;
+        mask_idx += 1;
+        grad_offset += 2;
+      }
+    }
+  }
+  // backward for softmax
+  if(softmax){
+    if (threadIdx.x == 0) {
+      opmath_t sum = 0.;
+      #pragma unroll
+      for (int i=0; i < K; ++i){
+        sum += cache_g_mask_before_softmax[i];
+      }
+      #pragma unroll
+      for (int i = 0; i < K; ++i) {
+        *(grad_offset_softmax) = cache_g_mask_before_softmax[i] - p_mask_shm[i] * sum;
+        grad_offset_softmax += 1;
+      }
+    }
+  }
+}
+
+template <typename scalar_t, typename stride_type, int d_stride>
+void _dcnv4_col2im_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,       // B, H * W, (G * D)
+    const scalar_t *p_offset,    // B, H * W, (G*K*3)
+    const scalar_t *grad_output, // B, H_out*W_out, G * D
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int G, const int D, const int B,
+    const int height_in, const int width_in, const int height_out,
+    const int width_out, const opmath_t offset_scale, const int remove_center,
+    opmath_t *grad_im, opmath_t *grad_offset, const int block_thread, 
+    const bool softmax, const int padded_offset_dim) {
+
+  constexpr int L = 1;
+
+  auto kernel =
+      backward_kernel_dcn_warp_primitive<scalar_t, d_stride, stride_type, 1, 9, false>;
+
+  int N = height_in * width_in;
+  int Q = height_out * width_out;
+  int K = kernel_h * kernel_w;
+
+  if (remove_center) {
+    K -= 1;
+  }
+
+  if (softmax) {
+    switch (K) {
+    case 9:
+      if(check_backward_warpp(d_stride, D)){
+        kernel = backward_kernel_dcn_warp_primitive<scalar_t, d_stride, stride_type, 1, 9, true>;
+      }
+      else{
+        kernel = backward_kernel_dcn<scalar_t, d_stride, stride_type, 1, 9, true>;
+      }
+      break;
+    case 8:
+      if(check_backward_warpp(d_stride, D)){
+        kernel = backward_kernel_dcn_warp_primitive<scalar_t, d_stride, stride_type, 1, 8, true>;
+      }
+      else {
+        kernel = backward_kernel_dcn<scalar_t, d_stride, stride_type, 1, 8, true>;
+      }
+      break;
+    default:
+      printf("K=%ld\n", K);
+      throw std::invalid_argument("invalid kernel shape");
+    }
+  } else {
+    switch (K) {
+    case 9:
+      if(check_backward_warpp(d_stride, D)){
+        kernel = backward_kernel_dcn_warp_primitive<scalar_t, d_stride, stride_type, 1, 9, false>;
+      }
+      else{
+        kernel = backward_kernel_dcn<scalar_t, d_stride, stride_type, 1, 9, false>;
+      }
+      break;
+    case 8:
+      if(check_backward_warpp(d_stride, D)){
+        kernel = backward_kernel_dcn_warp_primitive<scalar_t, d_stride, stride_type, 1, 8, false>;
+      }
+      else {
+        kernel = backward_kernel_dcn<scalar_t, d_stride, stride_type, 1, 8, false>;
+      }
+      break;
+    default:
+      printf("K=%ld\n", K);
+      throw std::invalid_argument("invalid kernel shape");
+    }
+  }
+
+  const int block_multiplier = block_thread / (D / d_stride) / G;
+  assert((B*Q) % block_multiplier == 0);
+
+  dim3 num_blocks(B*Q / block_multiplier);
+  dim3 num_threads(D / d_stride, G, block_multiplier);
+
+  const int blockdimX = D / d_stride;
+
+  int shm_size = sizeof(opmath_t) * (G * block_multiplier * K) * 2;
+  if(!check_backward_warpp(d_stride, D)){
+    shm_size = sizeof(opmath_t) * ((G * block_multiplier * K) * 2 + G * block_multiplier * blockdimX * 3);
+  }
+
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       shm_size);
+
+  kernel<<<num_blocks, num_threads, shm_size, stream>>>(
+      value, p_offset, grad_output, G, D, Q, kernel_h, kernel_w, stride_h,
+      stride_w, pad_h, pad_w, dilation_h, dilation_w, height_in, width_in,
+      height_out, width_out, offset_scale, remove_center, block_multiplier,
+      grad_im, grad_offset, padded_offset_dim);
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf("error in dcnv4_im2col_cuda: %s\n", cudaGetErrorString(err));
+    printf("launch arguments: gridDim=(%d, %d, %d), blockDim=(%d, %d, %d), "
+           "shm_size=%d\n\n",
+           num_blocks.x, num_blocks.y, num_blocks.z, num_threads.x,
+           num_threads.y, num_threads.z, shm_size);
+    AT_ASSERTM(false, "kernel launch error");
+  }
+}
+
+template <typename scalar_t>
+void dcnv4_col2im_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,       // B, H * W, (G * D)
+    const scalar_t *p_offset,    // B, H * W, (G*K*3)
+    const scalar_t *grad_output, // B, H_out*W_out, G * D
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int G, const int D, const int B,
+    const int height_in, const int width_in, const int height_out,
+    const int width_out, const opmath_t offset_scale, const int remove_center,
+    opmath_t *grad_im, opmath_t *grad_offset, const int d_stride,
+    const int block_thread, const bool softmax, const int padded_offset_dim) {
+
+  assert(D % d_stride == 0);
+  const int size_scalar = sizeof(scalar_t);
+  if (size_scalar == 2) {
+    switch (d_stride) {
+    case 1:
+      _dcnv4_col2im_cuda<scalar_t, scalar_t, 1>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 2:
+      _dcnv4_col2im_cuda<scalar_t, uint, 2>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 4:
+      _dcnv4_col2im_cuda<scalar_t, uint2, 4>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 8:
+      _dcnv4_col2im_cuda<scalar_t, uint4, 8>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 16:
+      _dcnv4_col2im_cuda<scalar_t, ulonglong4, 16>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    }
+  } else {
+    assert(size_scalar == 4);
+    switch (d_stride) {
+    case 1:
+      _dcnv4_col2im_cuda<scalar_t, uint, 1>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 2:
+      _dcnv4_col2im_cuda<scalar_t, uint2, 2>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 4:
+      _dcnv4_col2im_cuda<scalar_t, uint4, 4>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 8:
+      _dcnv4_col2im_cuda<scalar_t, ulonglong4, 8>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    }
+  }
+}

DCNv4_op/src/cuda/dcnv4_cuda.cu

@@ -0,0 +1,175 @@
+/*!
+**************************************************************************************************
+* 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
+**************************************************************************************************
+*/
+
+#include "cuda/dcnv4_im2col_cuda.cuh"
+#include "cuda/dcnv4_col2im_cuda.cuh"
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <torch/torch.h>
+
+at::Tensor dcnv4_cuda_forward(
+    const at::Tensor &value,
+    const at::Tensor &p_offset,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const int remove_center,
+    const int d_stride, const int block_thread, const bool softmax) {
+  AT_ASSERTM(value.is_contiguous(), "input tensor has to be contiguous");
+  AT_ASSERTM(value.type().is_cuda(), "input must be a CUDA tensor");
+  AT_ASSERTM(p_offset.is_contiguous(), "input tensor has to be contiguous");
+  AT_ASSERTM(p_offset.type().is_cuda(), "input must be a CUDA tensor");
+
+  const int batch = value.size(0);
+  const int height_in = value.size(1);
+  const int width_in = value.size(2);
+  const int channels = value.size(3);
+  const int padded_offset_dim = p_offset.size(3);
+
+  // tensor core requirement
+  assert(padded_offset_dim % 8 == 0);
+
+  const int height_out =
+      (height_in + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h +
+      1;
+  const int width_out =
+      (width_in + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
+
+  const int im2col_step_ = std::min(batch, im2col_step);
+  AT_ASSERTM(batch % im2col_step_ == 0, "batch(", batch,
+             ") must divide im2col_step(", im2col_step_, ")");
+  AT_ASSERTM(
+      channels == (group * group_channels),
+      "Input channels and group times group channels wont match: (%d vs %d).",
+      channels, group * group_channels);
+
+  auto output = at::zeros(
+      {batch, height_out, width_out, group * group_channels}, value.options());
+
+  const int batch_n = im2col_step_;
+  auto output_n = output.view({batch / batch_n, batch_n, height_out, width_out,
+                               group * group_channels});
+  auto per_value_size = height_in * width_in * channels;
+  auto per_offset_size = height_out * width_out * padded_offset_dim;
+
+  for (int n = 0; n < batch / im2col_step_; ++n) {
+    auto columns = output_n.select(0, n);
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half, at::ScalarType::BFloat16, value.scalar_type(),
+        "dcnv4_forward_cuda", ([&] {
+          dcnv4_im2col_cuda(
+              at::cuda::getCurrentCUDAStream(),
+              value.data_ptr<scalar_t>() + n * im2col_step_ * per_value_size,
+              p_offset.data_ptr<scalar_t>() +
+                  n * im2col_step_ * per_offset_size,
+              columns.data_ptr<scalar_t>(), kernel_h, kernel_w, stride_h,
+              stride_w, pad_h, pad_w, dilation_h, dilation_w, group,
+              group_channels, batch_n, height_in, width_in, height_out,
+              width_out, offset_scale, remove_center, d_stride, block_thread,
+              softmax, padded_offset_dim);
+        }));
+  }
+
+  return output;
+}
+
+std::vector<at::Tensor>
+dcnv4_cuda_backward(
+    const at::Tensor &value, 
+    const at::Tensor &p_offset, 
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const at::Tensor &grad_output, 
+    const int remove_center, const int d_stride, const int block_thread,
+    const bool softmax) {
+  AT_ASSERTM(value.is_contiguous(), "input tensor has to be contiguous");
+  AT_ASSERTM(p_offset.is_contiguous(), "offset tensor has to be contiguous");
+  AT_ASSERTM(grad_output.is_contiguous(), 
+             "grad_output tensor has to be contiguous");
+
+  AT_ASSERTM(value.type().is_cuda(), "input must be a CUDA tensor");
+  AT_ASSERTM(p_offset.type().is_cuda(), "offset must be a CUDA tensor");
+  AT_ASSERTM(grad_output.type().is_cuda(),
+             "grad_output must be a CUDA tensor");
+
+  const int batch = value.size(0);
+  const int height_in = value.size(1);
+  const int width_in = value.size(2);
+  const int channels = value.size(3);
+  const int padded_offset_dim = p_offset.size(3);
+  assert(padded_offset_dim % 8 == 0);
+
+  const int height_out =
+      (height_in + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h +
+      1;
+  const int width_out =
+      (width_in + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
+
+  const int im2col_step_ = std::min(batch, im2col_step);
+  AT_ASSERTM(batch % im2col_step_ == 0, "batch(", batch,
+             ") must divide im2col_step(", im2col_step_, ")");
+  AT_ASSERTM(
+      channels == (group * group_channels),
+      "Input channels and group times group channels wont match: (%d vs %d).",
+      channels, group * group_channels);
+  
+  auto dtype = value.dtype();
+  if (dtype == at::kHalf){
+    dtype = at::kFloat;
+  }
+
+  auto grad_input = at::zeros_like(value, dtype);
+  auto grad_offset = at::zeros_like(p_offset, dtype);
+
+  const int batch_n = im2col_step_;
+  auto grad_output_n = grad_output.view({batch / batch_n, batch_n, height_out, width_out,
+                               group, group_channels});
+  auto per_value_size = height_in * width_in * channels;
+  auto per_offset_size = height_out * width_out * padded_offset_dim;
+
+  for (int n = 0; n < batch / im2col_step_; ++n) {
+    auto columns = grad_output_n.select(0, n);
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(value.type(),
+        "dcnv4_backward_cuda", ([&] {
+          dcnv4_col2im_cuda(
+              at::cuda::getCurrentCUDAStream(),
+              value.data_ptr<scalar_t>() + n * im2col_step_ * per_value_size,
+              p_offset.data_ptr<scalar_t>() +
+                  n * im2col_step_ * per_offset_size,
+              columns.data_ptr<scalar_t>(), kernel_h, kernel_w, stride_h,
+              stride_w, pad_h, pad_w, dilation_h, dilation_w, group,
+              group_channels, batch_n, height_in, width_in, height_out,
+              width_out, offset_scale, remove_center,
+              grad_input.data<opmath_t>() + n * im2col_step_ * per_value_size,
+              grad_offset.data<opmath_t>() +
+                  n * im2col_step_ * per_offset_size,
+              d_stride, block_thread, softmax, padded_offset_dim
+              );
+        }));
+  }
+
+  if(value.dtype() == torch::kHalf){
+    return {grad_input.to(torch::kHalf), grad_offset.to(torch::kHalf)};
+  }
+  else{
+    return {grad_input, grad_offset};
+  }
+}
+

DCNv4_op/src/cuda/dcnv4_cuda.h

@@ -0,0 +1,33 @@
+/*!
+**************************************************************************************************
+* 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
+**************************************************************************************************
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor dcnv4_cuda_forward(
+    const at::Tensor &value,
+    const at::Tensor &p_offset,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const int remove_center,
+    const int d_stride, const int block_thread, const bool softmax);
+
+std::vector<at::Tensor>
+dcnv4_cuda_backward(
+    const at::Tensor &value,
+    const at::Tensor &p_offset,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const at::Tensor &grad_output, 
+    const int remove_center, const int d_stride, const int block_thread, 
+    const bool softmax);

DCNv4_op/src/cuda/dcnv4_im2col_cuda.cuh

@@ -0,0 +1,416 @@
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include "common.h"
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K,
+          bool softmax>
+__global__ void forward_kernel_dcn(
+    const scalar_t *p_value, const scalar_t *p_offset, scalar_t *p_output,
+    const int G, const int D, const int Q, const int kernel_h,
+    const int kernel_w, const int stride_h, const int stride_w, const int pad_h,
+    const int pad_w, const int dilation_h, const int dilation_w,
+    const int height_in, const int width_in, const int height_out,
+    const int width_out, const opmath_t offset_scale, const int remove_center,
+    const int block_multiplier, const int padded_offset_dim) {
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+  constexpr int li = 0;
+
+  extern __shared__ char _s[];
+  opmath_t *const p_mask_shm =
+      (opmath_t *)(_s) + ((threadIdx.z * G + gi) * L + li) * K;
+
+  opmath_t p_out_shm[d_stride] = {0.};
+
+  const scalar_t *p_offset_ptr = p_offset + (bi*Q + qi)*padded_offset_dim + gi*K*3;
+
+  const int mask_length = K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) = *(
+        scalar_t *)(p_offset_ptr + K * 2 + num_thread * num_iter + threadIdx.x);
+  }
+
+  int mask_idx;
+  if (softmax) {
+    __syncthreads();
+
+    // Calculate softmax over L and K
+    if (threadIdx.x == 0) { // gi != 0, di = 0, li = 0
+      opmath_t softmax_max = -1e100;
+      opmath_t softmax_sum = 0.0;
+
+      // get max
+      // #pragma unroll
+      for (int j = 0; j < K; j++) {
+        softmax_max = max(softmax_max, p_mask_shm[j]);
+      }
+
+      // get sumexp
+      // #pragma unroll
+      for (int j = 0; j < K; j++) {
+        opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+        p_mask_shm[j] = exp_results;
+        softmax_sum += exp_results;
+      }
+
+      // normalize
+      // #pragma unroll
+      for (int j = 0; j < K; j++) {
+        p_mask_shm[j] /= softmax_sum;
+      }
+    }
+
+    __syncthreads();
+  }
+  int offset_idx = 0;
+  mask_idx = 0;
+
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+  const scalar_t *p_value_ptr =
+      p_value + (bi * (height_in * width_in)) * (G * D);
+
+  const int p0_w = ((dilation_w * (kernel_w - 1)) >> 1) - pad_w +
+                   (qi % width_out) * stride_w;
+  const int p0_h = ((dilation_h * (kernel_h - 1)) >> 1) - pad_h +
+                   (qi / width_out) * stride_h;
+  const opmath_t p0_w_ =
+      p0_w - ((dilation_w * (kernel_w - 1)) >> 1) * offset_scale;
+  const opmath_t p0_h_ =
+      p0_h - ((dilation_h * (kernel_h - 1)) >> 1) * offset_scale;
+  const int center_h = kernel_h / 2;
+  const int center_w = kernel_w / 2;
+
+  int out_idx = ((bi * Q + qi) * G + gi) * D + di_s;
+
+  for (int i = 0; i < kernel_w; ++i) {
+    for (int j = 0; j < kernel_h; ++j) {
+      if (i != center_w || j != center_h || !remove_center) {
+        const opmath_t w_im =
+            p0_w_ + (i * dilation_w + (opmath_t)p_offset_ptr[offset_idx]) *
+                        offset_scale;
+        const opmath_t h_im =
+            p0_h_ + (j * dilation_h + (opmath_t)p_offset_ptr[offset_idx + 1]) *
+                        offset_scale;
+        const opmath_t attn = p_mask_shm[mask_idx];
+
+        if (h_im > -1 && w_im > -1 && h_im < height_in && w_im < width_in) {
+          ms_deform_attn_im2col_bilinear<scalar_t, transfer_t, d_stride>(
+              p_out_shm, p_value_ptr, height_in, width_in, h_im, w_im, attn,
+              w_stride, base_ptr);
+        }
+        offset_idx += 2;
+        mask_idx += 1;
+      }
+    }
+  }
+  scalar_t *fp16_regs = (scalar_t *)(p_out_shm);
+#pragma unroll
+  for (int ds = 0; ds < d_stride; ds++) {
+    fp16_regs[ds] = p_out_shm[ds];
+  }
+  *(transfer_t *)(p_output + out_idx) = *(transfer_t *)(p_out_shm);
+}
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K,
+          bool softmax>
+__global__ void forward_kernel_dcn_reg(
+    const scalar_t *p_value, const scalar_t *p_offset, scalar_t *p_output,
+    const int G, const int D, const int Q, const int kernel_h,
+    const int kernel_w, const int stride_h, const int stride_w, const int pad_h,
+    const int pad_w, const int dilation_h, const int dilation_w,
+    const int height_in, const int width_in, const int height_out,
+    const int width_out, const opmath_t offset_scale, const int remove_center,
+    const int block_multiplier, const int padded_offset_dim) {
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+  constexpr int li = 0;
+
+  opmath_t p_mask_shm[K] = {0.};
+  opmath_t p_out_shm[d_stride] = {0.};
+
+  const scalar_t *p_offset_ptr = p_offset + (bi*Q + qi)*padded_offset_dim + gi*K*3;
+  const int mask_length = K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+
+  for (int i=0; i < K; i++){
+    p_mask_shm[i] = *(p_offset_ptr + K*2 + i);
+  }
+
+  if (softmax) {
+    // Calculate softmax over L and K
+      opmath_t softmax_max = -1e100;
+      opmath_t softmax_sum = 0.0;
+      // get max
+      for (int j = 0; j < K; j++) {
+        softmax_max = max(softmax_max, p_mask_shm[j]);
+      }
+
+      // get sumexp
+      for (int j = 0; j < K; j++) {
+        opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+        p_mask_shm[j] = exp_results;
+        softmax_sum += exp_results;
+      }
+
+      // normalize
+      for (int j = 0; j < K; j++) {
+        p_mask_shm[j] /= softmax_sum;
+      }
+  }
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+  const scalar_t *p_value_ptr =
+      p_value + (bi * (height_in * width_in)) * (G * D);
+
+  const int p0_w = ((dilation_w * (kernel_w - 1)) >> 1) - pad_w +
+                   (qi % width_out) * stride_w;
+  const int p0_h = ((dilation_h * (kernel_h - 1)) >> 1) - pad_h +
+                   (qi / width_out) * stride_h;
+  const opmath_t p0_w_ =
+      p0_w - ((dilation_w * (kernel_w - 1)) >> 1) * offset_scale;
+  const opmath_t p0_h_ =
+      p0_h - ((dilation_h * (kernel_h - 1)) >> 1) * offset_scale;
+  const int center_h = kernel_h / 2;
+  const int center_w = kernel_w / 2;
+
+  int out_idx = ((bi * Q + qi) * G + gi) * D + di_s;
+
+  for (int i = 0; i < kernel_w; ++i) {
+    for (int j = 0; j < kernel_h; ++j) {
+      if (i != center_w || j != center_h || !remove_center) {
+        const opmath_t w_im =
+            p0_w_ + (i * dilation_w + (opmath_t)p_offset_ptr[offset_idx]) *
+                        offset_scale;
+        const opmath_t h_im =
+            p0_h_ + (j * dilation_h + (opmath_t)p_offset_ptr[offset_idx + 1]) *
+                        offset_scale;
+        const opmath_t attn = p_mask_shm[mask_idx];
+
+        if (h_im > -1 && w_im > -1 && h_im < height_in && w_im < width_in) {
+          ms_deform_attn_im2col_bilinear<scalar_t, transfer_t, d_stride>(
+              p_out_shm, p_value_ptr, height_in, width_in, h_im, w_im, attn,
+              w_stride, base_ptr);
+        }
+        offset_idx += 2;
+        mask_idx += 1;
+      }
+    }
+  }
+  scalar_t *fp16_regs = (scalar_t *)(p_out_shm);
+#pragma unroll
+  for (int ds = 0; ds < d_stride; ds++) {
+    fp16_regs[ds] = p_out_shm[ds];
+  }
+
+  *(transfer_t *)(p_output + out_idx) = *(transfer_t *)(p_out_shm);
+}
+
+template <typename scalar_t, typename stride_type, int d_stride>
+void _dcnv4_im2col_cuda(cudaStream_t stream,
+                              const scalar_t *value,    // B, H * W, (G * D)
+                              const scalar_t *p_offset, // B, H * W, G * K * 3)
+                              scalar_t *output,         // B, H_out*W_out, G * D
+                              const int kernel_h, const int kernel_w,
+                              const int stride_h, const int stride_w,
+                              const int pad_h, const int pad_w,
+                              const int dilation_h, const int dilation_w,
+                              const int G, const int D, const int B,
+                              const int height_in, const int width_in,
+                              const int height_out, const int width_out,
+                              const opmath_t offset_scale,
+                              const int remove_center, const int block_thread,
+                              const int softmax,
+                              const int padded_offset_dim) {
+
+  constexpr int L = 1;
+
+  auto kernel = forward_kernel_dcn_reg<scalar_t, d_stride, stride_type, 1, 9, true>;
+
+  int N = height_in * width_in;
+  int Q = height_out * width_out;
+  int K = kernel_h * kernel_w;
+
+  if (remove_center) {
+    K -= 1;
+  }
+  if (softmax) {
+    switch (K) {
+    case 9:
+      kernel = forward_kernel_dcn_reg<scalar_t, d_stride, stride_type, 1, 9, true>;
+      break;
+    case 8:
+      kernel = forward_kernel_dcn_reg<scalar_t, d_stride, stride_type, 1, 8, true>;
+      break;
+    default:
+      printf("K=%ld\n", K);
+      throw std::invalid_argument("invalid kernel shape");
+    }
+  } else {
+    switch (K) {
+    case 9:
+      kernel = forward_kernel_dcn_reg<scalar_t, d_stride, stride_type, 1, 9, false>;
+      break;
+    case 8:
+      kernel = forward_kernel_dcn_reg<scalar_t, d_stride, stride_type, 1, 8, false>;
+    break;
+    default:
+      printf("K=%ld\n", K);
+      throw std::invalid_argument("invalid kernel shape");
+    }
+  }
+
+  const int block_multiplier = block_thread / (D / d_stride) / G;
+  assert((B*Q) % block_multiplier == 0);
+
+  dim3 num_blocks(B*Q / block_multiplier);
+  dim3 num_threads(D / d_stride, G, block_multiplier);
+
+  int shm_size = 0;
+
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       shm_size);
+
+  kernel<<<num_blocks, num_threads, shm_size, stream>>>(
+      value, p_offset, output, G, D, Q, kernel_h, kernel_w, stride_h, stride_w,
+      pad_h, pad_w, dilation_h, dilation_w, height_in, width_in, height_out,
+      width_out, offset_scale, remove_center, block_multiplier, padded_offset_dim);
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf("error in dcnv4_im2col_cuda: %s\n", cudaGetErrorString(err));
+    printf("launch arguments: gridDim=(%d, %d, %d), blockDim=(%d, %d, %d), "
+           "shm_size=%d\n\n",
+           num_blocks.x, num_blocks.y, num_blocks.z, num_threads.x,
+           num_threads.y, num_threads.z, shm_size);
+    AT_ASSERTM(false, "kernel launch error");
+  }
+}
+
+template <typename scalar_t>
+void dcnv4_im2col_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,    // B, H * W, (G * D)
+    const scalar_t *p_offset, // B, H * W, G * K * 3)
+    scalar_t *output,         // B, H_out*W_out, G * D
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int G, const int D, const int B,
+    const int height_in, const int width_in, const int height_out,
+    const int width_out, const opmath_t offset_scale, const int remove_center,
+    const int d_stride, const int block_thread, const bool softmax,
+    const int padded_offset_dim) {
+
+  assert(D % d_stride == 0);
+  if (sizeof(scalar_t) == 2) {
+    switch (d_stride) {
+    case 1:
+      _dcnv4_im2col_cuda<scalar_t, scalar_t, 1>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 2:
+      _dcnv4_im2col_cuda<scalar_t, uint, 2>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 4:
+      _dcnv4_im2col_cuda<scalar_t, uint2, 4>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 8:
+      _dcnv4_im2col_cuda<scalar_t, uint4, 8>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 16:
+      _dcnv4_im2col_cuda<scalar_t, ulonglong4, 16>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    }
+  } else {
+    assert(sizeof(scalar_t) == 4);
+    switch (d_stride) {
+    case 1:
+      _dcnv4_im2col_cuda<scalar_t, uint, 1>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 2:
+      _dcnv4_im2col_cuda<scalar_t, uint2, 2>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 4:
+      _dcnv4_im2col_cuda<scalar_t, uint4, 4>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 8:
+      _dcnv4_im2col_cuda<scalar_t, ulonglong4, 8>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    default:
+      printf("not supported for d_stride > 8 for fp32");
+      throw std::invalid_argument("invalid d_stride");
+    }
+  }
+}

DCNv4_op/src/cuda/flash_deform_attn_cuda.cu

@@ -0,0 +1,162 @@
+/*!
+**************************************************************************************************
+* 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
+**************************************************************************************************
+*/
+
+#include "cuda/flash_deform_im2col_cuda.cuh"
+#include "cuda/flash_deform_col2im_cuda.cuh"
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <torch/torch.h>
+
+at::Tensor flash_deform_attn_cuda_forward(
+    const at::Tensor &value, const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index, const at::Tensor &sampling_loc_attn,
+    const int im2col_step = 64, const int K=8, const int d_stride=8, 
+    const int block_thread=0) {
+  AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+  AT_ASSERTM(spatial_shapes.is_contiguous(),
+             "spatial_shapes tensor has to be contiguous");
+  AT_ASSERTM(level_start_index.is_contiguous(),
+             "level_start_index tensor has to be contiguous");
+  AT_ASSERTM(sampling_loc_attn.is_contiguous(),
+             "sampling_loc_attn tensor has to be contiguous");
+
+  AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor");
+  AT_ASSERTM(spatial_shapes.type().is_cuda(),
+             "spatial_shapes must be a CUDA tensor");
+  AT_ASSERTM(level_start_index.type().is_cuda(),
+             "level_start_index must be a CUDA tensor");
+  AT_ASSERTM(sampling_loc_attn.type().is_cuda(),
+             "sampling_loc_attn 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 num_channels = value.size(3);
+
+  const int num_levels = spatial_shapes.size(0);
+  const int num_query = sampling_loc_attn.size(1);
+  const int num_point = K;
+
+  const int im2col_step_ = std::min(batch, im2col_step);
+  AT_ASSERTM(batch % im2col_step_ == 0, "batch(", batch,
+             ") must divide im2col_step(", im2col_step_, ")");
+
+  auto output =
+      at::zeros({batch, num_query, num_heads, num_channels}, value.options());
+
+  auto per_value_size = spatial_size * num_heads * num_channels;
+  auto per_offset_size = num_query * num_heads * num_levels * num_point * 3;
+  auto per_out_size = num_query * num_heads * num_channels;
+
+  for (int n = 0; n < batch / im2col_step_; ++n) {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half, at::ScalarType::BFloat16, value.scalar_type(),
+        "flash_deform_attn_forward_cuda", ([&] {
+          flash_deformable_im2col_cuda(
+              at::cuda::getCurrentCUDAStream(),
+              value.data_ptr<scalar_t>() + n * im2col_step_ * per_value_size,
+              spatial_shapes.data<int64_t>(), level_start_index.data<int64_t>(),
+              sampling_loc_attn.data_ptr<scalar_t>() +
+                  n * im2col_step_ * per_offset_size,
+              output.data_ptr<scalar_t>() + n * im2col_step_ * per_out_size,
+              im2col_step_, spatial_size, num_heads, num_channels, num_levels,
+              num_query, num_point, d_stride, block_thread, true);
+        }));
+  }
+  output = output.view({batch, num_query, num_heads * num_channels});
+  return output;
+}
+
+std::vector<at::Tensor>
+flash_deform_attn_cuda_backward(
+    const at::Tensor &value, const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index, const at::Tensor &sampling_loc_attn,
+    const at::Tensor &grad_output, const int im2col_step = 64, const int K=8,
+    const int d_stride=2, const int block_thread=0) {
+  AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+  AT_ASSERTM(spatial_shapes.is_contiguous(),
+             "spatial_shapes tensor has to be contiguous");
+  AT_ASSERTM(level_start_index.is_contiguous(),
+             "level_start_index tensor has to be contiguous");
+  AT_ASSERTM(sampling_loc_attn.is_contiguous(),
+             "sampling_loc_attn tensor has to be contiguous");
+  AT_ASSERTM(grad_output.is_contiguous(),
+             "grad_output tensor has to be contiguous");
+
+  AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor");
+  AT_ASSERTM(spatial_shapes.type().is_cuda(),
+             "spatial_shapes must be a CUDA tensor");
+  AT_ASSERTM(level_start_index.type().is_cuda(),
+             "level_start_index must be a CUDA tensor");
+  AT_ASSERTM(sampling_loc_attn.type().is_cuda(),
+             "sampling_loc_attn must be a CUDA tensor");
+  AT_ASSERTM(grad_output.type().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 num_channels = value.size(3);
+
+  const int num_levels = spatial_shapes.size(0);
+  const int num_query = sampling_loc_attn.size(1);
+  const int num_point = K;
+
+  const int im2col_step_ = std::min(batch, im2col_step);
+  AT_ASSERTM(batch % im2col_step_ == 0, "batch(", batch,
+             ") must divide im2col_step(", im2col_step_, ")");
+
+  auto dtype = value.dtype();
+  if (dtype == at::kHalf){
+    dtype = at::kFloat;
+  }
+
+  auto grad_input = at::zeros_like(value, dtype);
+  auto grad_offset = at::zeros_like(sampling_loc_attn, dtype);
+
+  auto per_value_size = spatial_size * num_heads * num_channels;
+  auto per_offset_size = num_query * num_heads * num_levels * num_point * 3;
+  auto per_out_size = num_query * num_heads * num_channels;
+
+  for (int n = 0; n < batch / im2col_step_; ++n) {
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(value.type(),
+        "flash_deform_attn_backward_cuda", ([&] {
+          flash_deformable_col2im_cuda(
+              at::cuda::getCurrentCUDAStream(),
+              value.data_ptr<scalar_t>() + n * im2col_step_ * per_value_size,
+              spatial_shapes.data<int64_t>(), level_start_index.data<int64_t>(),
+              sampling_loc_attn.data_ptr<scalar_t>() +
+                  n * im2col_step_ * per_offset_size,
+              grad_output.data_ptr<scalar_t>() + n * im2col_step_ * per_out_size,
+              im2col_step_, spatial_size, num_heads, num_channels, num_levels,
+              num_query, num_point, 
+              grad_input.data<opmath_t>() + n * im2col_step_ * per_value_size,
+              grad_offset.data<opmath_t>() + n * im2col_step_ * per_offset_size,
+              d_stride, block_thread
+              );
+        }));
+  }
+
+  if(value.dtype() == torch::kHalf){
+    return {grad_input.to(torch::kHalf), grad_offset.to(torch::kHalf)};
+  }
+  else{
+    return {grad_input, grad_offset};
+  }
+}

DCNv4_op/src/cuda/flash_deform_attn_cuda.h

@@ -0,0 +1,25 @@
+/*!
+**************************************************************************************************
+* 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
+**************************************************************************************************
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor flash_deform_attn_cuda_forward(
+    const at::Tensor &value, const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index, const at::Tensor &sampling_loc_attn,
+    const int im2col_step, const int K, const int d_stride, const int block_thread);
+
+std::vector<at::Tensor>
+flash_deform_attn_cuda_backward(
+    const at::Tensor &value, const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index, const at::Tensor &sampling_loc_attn,
+    const at::Tensor &grad_output, const int im2col_step, const int K,
+    const int d_stride, const int block_thread);

DCNv4_op/src/cuda/flash_deform_col2im_cuda.cuh

@@ -0,0 +1,580 @@
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include "common.h"
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K>
+__global__ void
+backward_kernel(const scalar_t *p_value, const int64_t *data_spatial_shapes,
+                const int64_t *data_level_start_index, const scalar_t *p_offset,
+                const scalar_t *grad_output, const int N, const int G,
+                const int D, const int Q, 
+                const int block_multiplier, opmath_t *grad_im,
+                opmath_t *grad_offset) {
+
+  extern __shared__ char _s[];
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+
+  opmath_t *cache_g_mask_before_softmax =
+      (opmath_t *)(_s); // (block_multiplier*G) * (L * K)
+  opmath_t *cache_grad_offset =
+      (opmath_t *)(cache_g_mask_before_softmax +
+                   block_multiplier * G * L *
+                       K); // (block_multiplier*G*D/d_stride*3)
+  opmath_t *const p_mask_shm =
+      ((opmath_t *)(cache_grad_offset +
+                    block_multiplier * G * D / d_stride * 3)) +
+      (threadIdx.z * G + gi) * L * K; // G*block_multiplier * L * K
+
+  const scalar_t *p_offset_ptr =
+      p_offset + (((bi * Q + qi) * G + gi) * L) * K * 3;
+  const int mask_length = L * K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+  const scalar_t *top_grad = grad_output + ((bi * Q + qi) * G + gi) * D + di_s;
+
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * num_iter +
+                      threadIdx.x);
+  }
+
+  __syncthreads();
+  // Calculate softmax over L and K
+  if (threadIdx.x == 0) { // gi != 0, di = 0, li = 0
+    opmath_t softmax_max = -1e100;
+    opmath_t softmax_sum = 0.0;
+
+    // get max
+    for (int j = 0; j < L * K; j++) {
+      softmax_max = max(softmax_max, p_mask_shm[j]);
+    }
+
+    // get sumexp
+    for (int j = 0; j < L * K; j++) {
+      opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+      p_mask_shm[j] = exp_results;
+      softmax_sum += exp_results;
+    }
+
+    // normalize
+    for (int j = 0; j < L * K; j++) {
+      p_mask_shm[j] /= softmax_sum;
+    }
+  }
+
+  __syncthreads();
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+
+  for (int li = 0; li < L; li++) {
+
+    const int spatial_h = data_spatial_shapes[li * 2];
+    const int spatial_w = data_spatial_shapes[li * 2 + 1];
+    const int level_start_id = data_level_start_index[li];
+    const scalar_t *p_value_ptr = p_value + (bi * N + level_start_id) * G * D;
+    opmath_t *grad_im_ptr = grad_im + (bi * N + level_start_id) * G * D;
+
+    int cache_grad_off_idx =
+        ((threadIdx.z * G + threadIdx.y) * blockDim.x + threadIdx.x) * 3;
+    for (int ki = 0; ki < K; ki++) {
+      const opmath_t loc_w = p_offset_ptr[offset_idx];
+      const opmath_t loc_h = p_offset_ptr[offset_idx + 1];
+      const opmath_t attn = p_mask_shm[mask_idx];
+      const opmath_t h_im = loc_h * spatial_h - 0.5;
+      const opmath_t w_im = loc_w * spatial_w - 0.5;
+      // for cache_grad_offset (mG) x D/d x 3
+      cache_grad_offset[cache_grad_off_idx] = 0;
+      cache_grad_offset[cache_grad_off_idx + 1] = 0;
+      cache_grad_offset[cache_grad_off_idx + 2] = 0;
+
+      if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+        ms_deform_attn_col2im_bilinear<scalar_t, transfer_t, d_stride>(
+            p_value_ptr, spatial_h, spatial_w, h_im, w_im, attn, w_stride,
+            base_ptr, spatial_h, spatial_w, top_grad, grad_im_ptr,
+            cache_grad_offset + cache_grad_off_idx);
+
+        // aggregate across different channel for offset
+        __syncthreads();
+        if (threadIdx.x == 0) {
+          int _didx = (threadIdx.z * G + threadIdx.y) * blockDim.x * 3;
+          opmath_t _grad_w = cache_grad_offset[_didx];
+          opmath_t _grad_h = cache_grad_offset[_didx + 1];
+          opmath_t _grad_a = cache_grad_offset[_didx + 2];
+          for (int c_id = 1; c_id < blockDim.x; ++c_id) {
+            _grad_w += cache_grad_offset[_didx + 3 * c_id];
+            _grad_h += cache_grad_offset[_didx + 3 * c_id + 1];
+            _grad_a += cache_grad_offset[_didx + 3 * c_id + 2];
+          }
+
+          grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + li * K * 2 +
+                      ki * 2] = _grad_w;
+          grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + li * K * 2 +
+                      ki * 2 + 1] = _grad_h;
+          cache_g_mask_before_softmax
+              [((threadIdx.y + threadIdx.z * G) * L + li) * K + ki] = _grad_a;
+        }
+      }
+      __syncthreads();
+
+      offset_idx += 2;
+      mask_idx += 1;
+    }
+  }
+  // backward for softmax
+  if (threadIdx.x == 0) {
+    for (int i = 0; i < L * K; ++i) {
+      opmath_t grad_i = 0.;
+      const opmath_t *group_g_mask = cache_g_mask_before_softmax +
+                                      (threadIdx.y + threadIdx.z * G) * L * K;
+      for (int j = 0; j < L * K; ++j) {
+        if (i != j) {
+          grad_i -= group_g_mask[j] * p_mask_shm[i] * p_mask_shm[j];
+        } else {
+          grad_i += group_g_mask[i] * p_mask_shm[i] * (1 - p_mask_shm[i]);
+        }
+      }
+      grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + L * K * 2 + i] =
+          grad_i;
+    }
+  }
+  __syncthreads();
+}
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K>
+__global__ void
+backward_kernel_warp_primitive(const scalar_t *p_value, const int64_t *data_spatial_shapes,
+                const int64_t *data_level_start_index, const scalar_t *p_offset,
+                const scalar_t *grad_output, const int N, const int G,
+                const int D, const int Q, 
+                const int block_multiplier, opmath_t *grad_im,
+                opmath_t *grad_offset) {
+
+  extern __shared__ char _s[];
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+
+  const int tid = (threadIdx.z * blockDim.y + threadIdx.y)*blockDim.x + threadIdx.x;
+  const int lane_id = tid % kWarpSize;
+  const int group_per_warp = kWarpSize / blockDim.x;
+  const int group_in_warp_id = (threadIdx.z * G + threadIdx.y) % group_per_warp;
+  const unsigned lane_mask = ((1 << blockDim.x) - 1) << (group_in_warp_id * blockDim.x);
+
+  opmath_t *cache_g_mask_before_softmax =
+      (opmath_t *)(_s); // (block_multiplier*G) * (L * K)
+
+  opmath_t *const p_mask_shm =
+      ((opmath_t *)(cache_g_mask_before_softmax + block_multiplier * G * L * K)) +
+        (threadIdx.z * G + gi) * L * K; // G*block_multiplier * L * K
+
+  const scalar_t *p_offset_ptr =
+      p_offset + (((bi * Q + qi) * G + gi) * L) * K * 3;
+  const int mask_length = L * K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+  const scalar_t *top_grad = grad_output + ((bi * Q + qi) * G + gi) * D + di_s;
+
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * num_iter +
+                      threadIdx.x);
+  }
+
+  __syncthreads();
+  // Calculate softmax over L and K
+  if (threadIdx.x == 0) { // gi != 0, di = 0, li = 0
+    opmath_t softmax_max = -1e100;
+    opmath_t softmax_sum = 0.0;
+
+    // get max
+    for (int j = 0; j < L * K; j++) {
+      softmax_max = max(softmax_max, p_mask_shm[j]);
+    }
+
+    // get sumexp
+    for (int j = 0; j < L * K; j++) {
+      opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+      p_mask_shm[j] = exp_results;
+      softmax_sum += exp_results;
+    }
+
+    // normalize
+    for (int j = 0; j < L * K; j++) {
+      p_mask_shm[j] /= softmax_sum;
+    }
+  }
+
+  __syncthreads();
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+
+  for (int li = 0; li < L; li++) {
+
+    const int spatial_h = data_spatial_shapes[li * 2];
+    const int spatial_w = data_spatial_shapes[li * 2 + 1];
+    const int level_start_id = data_level_start_index[li];
+    const scalar_t *p_value_ptr = p_value + (bi * N + level_start_id) * G * D;
+    opmath_t *grad_im_ptr = grad_im + (bi * N + level_start_id) * G * D;
+
+    int cache_grad_off_idx =
+        ((threadIdx.z * G + threadIdx.y) * blockDim.x + threadIdx.x) * 3;
+
+    opmath_t reg_grad_offset[3] = {0.};
+    for (int ki = 0; ki < K; ki++) {
+      const opmath_t loc_w = p_offset_ptr[offset_idx];
+      const opmath_t loc_h = p_offset_ptr[offset_idx + 1];
+      const opmath_t attn = p_mask_shm[mask_idx];
+      const opmath_t h_im = loc_h * spatial_h - 0.5;
+      const opmath_t w_im = loc_w * spatial_w - 0.5;
+      reg_grad_offset[0] = 0;
+      reg_grad_offset[1] = 0;
+      reg_grad_offset[2] = 0;
+
+      if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+        ms_deform_attn_col2im_bilinear<scalar_t, transfer_t, d_stride>(
+            p_value_ptr, spatial_h, spatial_w, h_im, w_im, attn, w_stride,
+            base_ptr, spatial_h, spatial_w, top_grad, grad_im_ptr,
+            reg_grad_offset);
+
+        // aggregate across different channel for offset
+        for (uint32_t offset = blockDim.x>>1; offset > 0; offset >>= 1){
+          reg_grad_offset[0] += __shfl_down_sync(lane_mask, reg_grad_offset[0], offset);
+          reg_grad_offset[1] += __shfl_down_sync(lane_mask, reg_grad_offset[1], offset);
+          reg_grad_offset[2] += __shfl_down_sync(lane_mask, reg_grad_offset[2], offset);
+        }
+
+        if (threadIdx.x == 0) {
+          grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + li * K * 2 +
+                      ki * 2] = reg_grad_offset[0];
+          grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + li * K * 2 +
+                      ki * 2 + 1] = reg_grad_offset[1];
+          cache_g_mask_before_softmax
+              [((threadIdx.y + threadIdx.z * G) * L + li) * K + ki] = reg_grad_offset[2];
+        }
+      }
+      __syncthreads();
+
+      offset_idx += 2;
+      mask_idx += 1;
+    }
+  }
+  // backward for softmax
+  if (threadIdx.x == 0) {
+    for (int i = 0; i < L * K; ++i) {
+      opmath_t grad_i = 0.;
+      const opmath_t *group_g_mask = cache_g_mask_before_softmax +
+                                      (threadIdx.y + threadIdx.z * G) * L * K;
+      for (int j = 0; j < L * K; ++j) {
+        if (i != j) {
+          grad_i -= group_g_mask[j] * p_mask_shm[i] * p_mask_shm[j];
+        } else {
+          grad_i += group_g_mask[i] * p_mask_shm[i] * (1 - p_mask_shm[i]);
+        }
+      }
+      grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + L * K * 2 + i] =
+          grad_i;
+    }
+  }
+  __syncthreads();
+}
+
+template <typename scalar_t, typename stride_type, int K, int d_stride>
+void _flash_deformable_col2im_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    const scalar_t *grad_output,           // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, opmath_t *grad_im, opmath_t *grad_offset,
+    const int block_thread) {
+
+  assert(D % d_stride == 0);
+
+  const int block_multiplier = block_thread / (D / d_stride) / G;
+  assert((B*Q) % block_multiplier == 0);
+  dim3 num_blocks(B*Q / block_multiplier);
+  dim3 num_threads(D / d_stride, G, block_multiplier);
+
+  int shm_size;
+  if(check_backward_warpp(d_stride, D)){
+    shm_size =
+      sizeof(opmath_t) * (block_multiplier * G * L * K) +
+      sizeof(opmath_t) * (G * block_multiplier * L * K);
+  }
+  else{
+    shm_size =
+      sizeof(opmath_t) * (block_multiplier * G * L * K) +
+      sizeof(opmath_t) * (G * block_multiplier * L * K) + 
+      sizeof(opmath_t) * (G * block_multiplier * D / d_stride * 3);
+  }
+
+  auto kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 1, K>;
+
+  switch (L) {
+  case 1:
+    if(check_backward_warpp(d_stride, D)){
+      kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 1, K>;
+    } else {
+      kernel = backward_kernel<scalar_t, d_stride, stride_type, 1, K>;
+    }
+    break;
+  case 2:
+    if(check_backward_warpp(d_stride, D)){
+      kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 2, K>;
+    } else {
+      kernel = backward_kernel<scalar_t, d_stride, stride_type, 2, K>;
+    }
+    break;
+  case 3:
+    if(check_backward_warpp(d_stride, D)){
+      kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 3, K>;
+    } else {
+      kernel = backward_kernel<scalar_t, d_stride, stride_type, 3, K>;
+    }
+    break;
+  case 4:
+    if(check_backward_warpp(d_stride, D)){
+      kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 4, K>;
+    } else {
+      kernel = backward_kernel<scalar_t, d_stride, stride_type, 4, K>;
+    }
+    break;
+  case 5:
+    if(check_backward_warpp(d_stride, D)){
+      kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 5, K>;
+    } else {
+      kernel = backward_kernel<scalar_t, d_stride, stride_type, 5, K>;
+    }
+    break;
+  default:
+    printf("L=%ld\n", L);
+    throw std::invalid_argument("invalid number of scales");
+  }
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       shm_size);
+
+  kernel<<<num_blocks, num_threads, shm_size, stream>>>(
+      value, data_spatial_shapes, data_level_start_index, offset, grad_output,
+      N, G, D, Q, block_multiplier, grad_im, grad_offset);
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf("error in flash_deformable_im2col_cuda: %s\n",
+           cudaGetErrorString(err));
+    printf("launch arguments: gridDim=(%d, %d, %d), blockDim=(%d, %d, %d), "
+           "shm_size=%d, Q=%d\n\n",
+           num_blocks.x, num_blocks.y, num_blocks.z, num_threads.x,
+           num_threads.y, num_threads.z, shm_size, Q);
+    AT_ASSERTM(false, "kernel launch error");
+  }
+}
+
+template <typename scalar_t, int K>
+void flash_deformable_col2im_cuda_inner(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    const scalar_t *grad_output,           // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, opmath_t *grad_im, opmath_t *grad_offset, 
+    const int d_stride, const int block_thread) {
+
+  assert(D % d_stride == 0);
+  if(sizeof(scalar_t) == 2) {
+    switch(d_stride) {
+    case 1:
+      _flash_deformable_col2im_cuda<scalar_t, scalar_t, K, 1>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          grad_output,            // B, N, G, D
+          B, N, G, D, L, Q, grad_im, grad_offset,
+          block_thread);
+      break;
+    case 2:
+      _flash_deformable_col2im_cuda<scalar_t, uint, K, 2>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          grad_output,            // B, N, G, D
+          B, N, G, D, L, Q, grad_im, grad_offset,
+          block_thread);
+      break;
+    case 4:
+      _flash_deformable_col2im_cuda<scalar_t, uint2, K, 4>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          grad_output,            // B, N, G, D
+          B, N, G, D, L, Q, grad_im, grad_offset,
+          block_thread);
+      break;
+    case 8:
+      _flash_deformable_col2im_cuda<scalar_t, uint4, K, 8>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          grad_output,            // B, N, G, D
+          B, N, G, D, L, Q, grad_im, grad_offset,
+          block_thread);
+      break;
+    case 16:
+      _flash_deformable_col2im_cuda<scalar_t, ulonglong4, K, 16>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          grad_output,            // B, N, G, D
+          B, N, G, D, L, Q, grad_im, grad_offset,
+          block_thread);
+      break;
+    default:
+      printf("not supported for d_stride > 16 for fp16");
+      throw std::invalid_argument("invalid d_stride");
+    }
+  } else {
+    assert(sizeof(scalar_t) == 4);
+    switch(d_stride) {
+    case 1:  
+      _flash_deformable_col2im_cuda<scalar_t, scalar_t, K, 1>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        block_thread);
+      break;
+    case 2:  
+      _flash_deformable_col2im_cuda<scalar_t, uint2, K, 2>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        block_thread);
+      break;
+    case 4:  
+      _flash_deformable_col2im_cuda<scalar_t, uint4, K, 4>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        block_thread);
+      break;
+    case 8:  
+      _flash_deformable_col2im_cuda<scalar_t, ulonglong4, K, 8>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        block_thread);
+      break;
+    default:
+      printf("not supported for d_stride > 8 for fp32");
+      throw std::invalid_argument("invalid d_stride");
+    }
+  }
+}
+
+template <typename scalar_t>
+void flash_deformable_col2im_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    const scalar_t *grad_output,           // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, const int K, opmath_t *grad_im, opmath_t *grad_offset,
+    const int d_stride, const int block_thread) {
+
+  switch (K) {
+  case 4:
+    flash_deformable_col2im_cuda_inner<scalar_t, 4>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        d_stride, block_thread);
+    break;
+  case 8:
+    flash_deformable_col2im_cuda_inner<scalar_t, 8>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        d_stride, block_thread);
+    break;
+  default:
+    printf("not supported for K not in [4, 8]");
+    throw std::invalid_argument("invalid K");
+  }
+}

DCNv4_op/src/cuda/flash_deform_im2col_cuda.cuh

@@ -0,0 +1,451 @@
+/*!
+**************************************************************************
+* 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 <algorithm>
+#include <cstdio>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include "common.h"
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K>
+__global__ void
+forward_kernel(const scalar_t *p_value, const int64_t *data_spatial_shapes,
+               const int64_t *data_level_start_index, const scalar_t *p_offset,
+               scalar_t *p_output, const int N, const int G, const int D,
+               const int Q, const int block_multiplier) {
+
+  extern __shared__ char _s[];
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+
+  opmath_t p_out_shm[d_stride] = {0.};
+  opmath_t *const p_mask_shm =
+      (opmath_t *)(_s) + (threadIdx.z * G + gi) * L * K;
+
+  const scalar_t *p_offset_ptr =
+      p_offset + (((bi * Q + qi) * G + gi) * L) * K * 3;
+  const int mask_length = L * K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * num_iter +
+                      threadIdx.x);
+  }
+
+  __syncthreads();
+  // Calculate softmax over L and K
+  if (threadIdx.x == 0) { // di = 0
+    opmath_t softmax_max = -1e100;
+    opmath_t softmax_sum = 0.0;
+
+    // get max
+    for (int j = 0; j < L * K; j++) {
+      softmax_max = max(softmax_max, p_mask_shm[j]);
+    }
+
+    // get sumexp
+    for (int j = 0; j < L * K; j++) {
+      opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+      p_mask_shm[j] = exp_results;
+      softmax_sum += exp_results;
+    }
+
+    // normalize
+    for (int j = 0; j < L * K; j++) {
+      p_mask_shm[j] /= softmax_sum;
+    }
+  }
+
+  __syncthreads();
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+
+  for (int li = 0; li < L; li++) {
+    const int spatial_h = data_spatial_shapes[li * 2];
+    const int spatial_w = data_spatial_shapes[li * 2 + 1];
+    const int level_start_id = data_level_start_index[li];
+    const scalar_t *p_value_ptr = p_value + (bi * N + level_start_id) * G * D;
+
+    for (int ki = 0; ki < K; ki++) {
+      const opmath_t loc_w = p_offset_ptr[offset_idx];
+      const opmath_t loc_h = p_offset_ptr[offset_idx + 1];
+      const opmath_t attn = p_mask_shm[mask_idx];
+      const opmath_t h_im = loc_h * spatial_h - 0.5;
+      const opmath_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_im2col_bilinear<scalar_t, transfer_t, d_stride>(
+            p_out_shm, p_value_ptr, spatial_h, spatial_w, h_im, w_im, attn,
+            w_stride, base_ptr);
+      }
+      offset_idx += 2;
+      mask_idx += 1;
+    }
+  }
+
+  int out_idx = ((bi * Q + qi) * G + gi) * D + di_s;
+
+  scalar_t *fp16_regs = (scalar_t *)(p_out_shm);
+#pragma unroll
+  for (int ds = 0; ds < d_stride; ds++) {
+    fp16_regs[ds] = p_out_shm[ds];
+  }
+
+  *(transfer_t *)(p_output + out_idx) = *(transfer_t *)(p_out_shm);
+}
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K>
+__global__ void
+forward_kernel_reg(const scalar_t *p_value, const int64_t *data_spatial_shapes,
+               const int64_t *data_level_start_index, const scalar_t *p_offset,
+               scalar_t *p_output, const int N, const int G, const int D,
+               const int Q, const int block_multiplier) {
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+
+  opmath_t p_out_shm[d_stride] = {0.};
+  opmath_t p_mask_shm[L*K] = {0.};
+
+  const scalar_t *p_offset_ptr =
+      p_offset + (((bi * Q + qi) * G + gi) * L) * K * 3;
+
+  for (int i=0; i < L*K; i++){
+    p_mask_shm[i] = *(p_offset_ptr + L * K * 2 + i);
+  }
+
+  // Calculate softmax over L and K
+  opmath_t softmax_max = -1e100;
+  opmath_t softmax_sum = 0.0;
+
+  // get max
+  for (int j = 0; j < L * K; j++) {
+    softmax_max = max(softmax_max, p_mask_shm[j]);
+  }
+
+  // get sumexp
+  for (int j = 0; j < L * K; j++) {
+    opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+    p_mask_shm[j] = exp_results;
+    softmax_sum += exp_results;
+  }
+
+  // normalize
+  for (int j = 0; j < L * K; j++) {
+    p_mask_shm[j] /= softmax_sum;
+  }
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+
+  for (int li = 0; li < L; li++) {
+    const int spatial_h = data_spatial_shapes[li * 2];
+    const int spatial_w = data_spatial_shapes[li * 2 + 1];
+    const int level_start_id = data_level_start_index[li];
+    const scalar_t *p_value_ptr = p_value + (bi * N + level_start_id) * G * D;
+
+    for (int ki = 0; ki < K; ki++) {
+      const opmath_t loc_w = p_offset_ptr[offset_idx];
+      const opmath_t loc_h = p_offset_ptr[offset_idx + 1];
+      const opmath_t attn = p_mask_shm[mask_idx];
+      const opmath_t h_im = loc_h * spatial_h - 0.5;
+      const opmath_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_im2col_bilinear<scalar_t, transfer_t, d_stride>(
+            p_out_shm, p_value_ptr, spatial_h, spatial_w, h_im, w_im, attn,
+            w_stride, base_ptr);
+      }
+      offset_idx += 2;
+      mask_idx += 1;
+    }
+  }
+
+  int out_idx = ((bi * Q + qi) * G + gi) * D + di_s;
+
+  scalar_t *fp16_regs = (scalar_t *)(p_out_shm);
+#pragma unroll
+  for (int ds = 0; ds < d_stride; ds++) {
+    fp16_regs[ds] = p_out_shm[ds];
+  }
+
+  *(transfer_t *)(p_output + out_idx) = *(transfer_t *)(p_out_shm);
+}
+
+template <typename scalar_t, typename stride_type, int K, int d_stride>
+void _flash_deformable_im2col_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    scalar_t *output,                      // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, const int block_thread, 
+    const bool _use_reg) {
+
+  assert(D % d_stride == 0);
+
+  const int block_multiplier = block_thread / (D / d_stride) / G;;
+  assert((B*Q) % block_multiplier == 0);
+  dim3 num_blocks(B*Q / block_multiplier);
+  dim3 num_threads(D / d_stride, G, block_multiplier);
+
+  const int shm_size = 0;
+
+  auto kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 1, K>;
+
+  switch (L) {
+  case 1:
+    kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 1, K>;
+    break;
+  case 2:
+    kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 2, K>;
+    break;
+  case 3:
+    kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 3, K>;
+    break;
+  case 4:
+    kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 4, K>;
+    break;
+  case 5:
+    kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 5, K>;
+    break;
+  default:
+    printf("L=%ld\n", L);
+    throw std::invalid_argument("invalid number of scales");
+  }
+
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       shm_size);
+
+  kernel<<<num_blocks, num_threads, shm_size, stream>>>(
+      value, data_spatial_shapes, data_level_start_index, offset, output, N, G,
+      D, Q, block_multiplier);
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf("error in flash_deformable_im2col_cuda: %s\n",
+           cudaGetErrorString(err));
+    printf("launch arguments: gridDim=(%d, %d, %d), blockDim=(%d, %d, %d), "
+           "shm_size=%d, Q=%d\n\n",
+           num_blocks.x, num_blocks.y, num_blocks.z, num_threads.x,
+           num_threads.y, num_threads.z, shm_size, Q);
+    AT_ASSERTM(false, "kernel launch error");
+  }
+}
+
+template <typename scalar_t, int K>
+void flash_deformable_im2col_cuda_inner(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    scalar_t *output,                      // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, const int d_stride, 
+    const int block_thread,
+    const bool _use_reg) {
+  
+  assert(D % d_stride == 0);
+  if(sizeof(scalar_t) == 2) {
+    switch(d_stride) {
+    case 1:
+      _flash_deformable_im2col_cuda<scalar_t, scalar_t, K, 1>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 2:
+      _flash_deformable_im2col_cuda<scalar_t, uint, K, 2>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 4:
+      _flash_deformable_im2col_cuda<scalar_t, uint2, K, 4>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 8:
+      _flash_deformable_im2col_cuda<scalar_t, uint4, K, 8>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 16:
+      _flash_deformable_im2col_cuda<scalar_t, ulonglong4, K, 16>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    default:
+      printf("not supported for d_stride > 16 for fp16");
+      throw std::invalid_argument("invalid d_stride");
+    }
+  } else {
+    assert(sizeof(scalar_t) == 4);
+    switch(d_stride) {
+    case 1:
+      _flash_deformable_im2col_cuda<scalar_t, scalar_t, K, 1>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 2:
+      _flash_deformable_im2col_cuda<scalar_t, uint2, K, 2>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 4:
+      _flash_deformable_im2col_cuda<scalar_t, uint4, K, 4>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 8:
+      _flash_deformable_im2col_cuda<scalar_t, ulonglong4, K, 8>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    default:
+      printf("not supported for d_stride > 8 for fp32");
+      throw std::invalid_argument("invalid d_stride");
+    }
+  }
+}
+
+template <typename scalar_t>
+void flash_deformable_im2col_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    scalar_t *output,                      // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, const int K, const int d_stride, 
+    const int block_thread,
+    const bool _use_reg) {
+  switch (K) {
+  case 4:
+    flash_deformable_im2col_cuda_inner<scalar_t, 4>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        output,                 // B, N, G, D
+        B, N, G, D, L, Q, d_stride, 
+        block_thread, _use_reg);
+    break;
+  case 8:
+    flash_deformable_im2col_cuda_inner<scalar_t, 8>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        output,                 // B, N, G, D
+        B, N, G, D, L, Q, d_stride,
+        block_thread, _use_reg);
+    break;
+  default:
+    printf("not supported for K not in [4, 8]");
+    throw std::invalid_argument("invalid K");
+  }
+}

DCNv4_op/src/dcnv4.h

@@ -0,0 +1,107 @@
+/*!
+**************************************************************************************************
+* 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
+**************************************************************************************************
+*/
+
+#pragma once
+
+#ifdef WITH_CUDA
+#include "cuda/dcnv4_cuda.h"
+#include "cuda/flash_deform_attn_cuda.h"
+#endif
+
+at::Tensor flash_deform_attn_forward(const at::Tensor &value,
+                                        const at::Tensor &spatial_shapes,
+                                        const at::Tensor &level_start_index,
+                                        const at::Tensor &sampling_loc_attn,
+                                        const int im2col_step, const int K,
+                                        const int d_stride, const int block_thread) {
+  if (value.device().is_cuda()) {
+#ifdef WITH_CUDA
+    return flash_deform_attn_cuda_forward(value, spatial_shapes,
+                                             level_start_index,
+                                             sampling_loc_attn, im2col_step, 
+                                             K, d_stride, block_thread);
+#else
+    AT_ERROR("Not compiled with GPU support");
+#endif
+  }
+  AT_ERROR("Not implemented on the CPU");
+}
+
+std::vector<at::Tensor>
+flash_deform_attn_backward(const at::Tensor &value, 
+                              const at::Tensor &spatial_shapes,
+                              const at::Tensor &level_start_index, 
+                              const at::Tensor &sampling_loc_attn,
+                              const at::Tensor &grad_output, 
+                              const int im2col_step, 
+                              const int K, 
+                              const int d_stride, const int block_thread){
+  if (value.device().is_cuda()) {
+#ifdef WITH_CUDA
+    return flash_deform_attn_cuda_backward(value, 
+                                              spatial_shapes,
+                                              level_start_index,
+                                              sampling_loc_attn, 
+                                              grad_output,
+                                              im2col_step, 
+                                              K, d_stride, 
+                                              block_thread);
+#else
+    AT_ERROR("Not compiled with GPU support");
+#endif
+  }
+  AT_ERROR("Not implemented on the CPU");
+}
+
+at::Tensor dcnv4_forward(
+    const at::Tensor &value,
+    const at::Tensor &p_offset,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const int remove_center,
+    const int d_stride, const int block_thread, const bool softmax) {
+  if (value.device().is_cuda()) {
+#ifdef WITH_CUDA
+    return dcnv4_cuda_forward(
+        value, p_offset, kernel_h, kernel_w, stride_h, stride_w, pad_h,
+        pad_w, dilation_h, dilation_w, group, group_channels, offset_scale,
+        im2col_step, remove_center, d_stride, block_thread, softmax);
+#else
+    AT_ERROR("Not compiled with GPU support");
+#endif
+  }
+  AT_ERROR("Not implemented on the CPU");
+}
+
+std::vector<at::Tensor>
+dcnv4_backward(
+    const at::Tensor &value, 
+    const at::Tensor &p_offset, 
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const at::Tensor &grad_output, 
+    const int remove_center, const int d_stride, const int block_thread,
+    const bool softmax){
+  if (value.device().is_cuda()) {
+#ifdef WITH_CUDA
+    return dcnv4_cuda_backward(
+        value, p_offset, kernel_h, kernel_w, stride_h, stride_w, pad_h,
+        pad_w, dilation_h, dilation_w, group, group_channels, offset_scale,
+        im2col_step, grad_output, remove_center, d_stride, block_thread,
+        softmax);
+#else
+    AT_ERROR("Not compiled with GPU support");
+#endif
+  }
+  AT_ERROR("Not implemented on the CPU");
+}

DCNv4_op/src/vision.cpp

@@ -0,0 +1,21 @@
+/*!
+**************************************************************************************************
+* 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
+**************************************************************************************************
+*/
+
+#include "dcnv4.h"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("flash_deform_attn_forward", &flash_deform_attn_forward,
+        "flash_deform_attn_forward");
+  m.def("flash_deform_attn_backward", &flash_deform_attn_backward,
+        "flash_deform_attn_backward");
+  m.def("dcnv4_forward", &dcnv4_forward, "dcnv4_forward");
+  m.def("dcnv4_backward", &dcnv4_backward, "dcnv4_backward");
+}