[MPS] deformable conv2d kernel

test/test_ops.py

@@ -929,6 +929,7 @@ def test_batched_nms_implementations(self, seed):
 
 class TestDeformConv:
     dtype = torch.float64
+    mps_dtype = torch.float32
 
     def expected_fn(self, x, weight, offset, mask, bias, stride=1, padding=0, dilation=1):
         stride_h, stride_w = _pair(stride)
@@ -1050,12 +1051,11 @@ def test_is_leaf_node(self, device):
         assert len(graph_node_names[0]) == len(graph_node_names[1])
         assert len(graph_node_names[0]) == 1 + op_obj.n_inputs
 
-    @pytest.mark.parametrize("device", cpu_and_cuda())
+    @pytest.mark.parametrize("device", cpu_and_cuda_and_mps())
     @pytest.mark.parametrize("contiguous", (True, False))
     @pytest.mark.parametrize("batch_sz", (0, 33))
-    @pytest.mark.opcheck_only_one()
     def test_forward(self, device, contiguous, batch_sz, dtype=None):
-        dtype = dtype or self.dtype
+        dtype = self.mps_dtype if device == "mps" else dtype or self.dtype
         x, _, offset, mask, _, stride, padding, dilation = self.get_fn_args(device, contiguous, batch_sz, dtype)
         in_channels = 6
         out_channels = 2
@@ -1201,13 +1201,50 @@ def test_forward_scriptability(self):
         torch.jit.script(ops.DeformConv2d(in_channels=8, out_channels=8, kernel_size=3))
 
 
-optests.generate_opcheck_tests(
-    testcase=TestDeformConv,
-    namespaces=["torchvision"],
-    failures_dict_path=os.path.join(os.path.dirname(__file__), "optests_failures_dict.json"),
-    additional_decorators=[],
-    test_utils=OPTESTS,
-)
+@pytest.mark.parametrize("dtype", (torch.float16, torch.float32, torch.float64))
+@pytest.mark.parametrize("device", cpu_and_cuda())
+@pytest.mark.parametrize("requires_grad", (True, False))
+def test_deform_conv2d_opcheck(dtype, device, requires_grad):
+    batch_size, channels_in, height, width = 1, 6, 10, 10
+    kernel_size = (3, 3)
+    stride = (1, 1)
+    padding = (1, 1)
+    dilation = (1, 1)
+    groups = 2
+    out_channels = 4
+    out_h = (height + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1) // stride[0] + 1
+    out_w = (width + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1) // stride[1] + 1
+    x = torch.randn(batch_size, channels_in, height, width, dtype=dtype, device=device, requires_grad=requires_grad)
+    offset = torch.randn(batch_size, 2 * kernel_size[0] * kernel_size[1], out_h, out_w,
+                         dtype=dtype, device=device, requires_grad=requires_grad)
+    weight = torch.randn(out_channels, channels_in // groups, kernel_size[0], kernel_size[1],
+                         dtype=dtype, device=device, requires_grad=requires_grad)
+    bias = torch.randn(out_channels, dtype=dtype, device=device, requires_grad=requires_grad)
+    use_mask = True
+    mask = torch.sigmoid(torch.randn(
+        batch_size,
+        kernel_size[0] * kernel_size[1],
+        out_h,
+        out_w,
+        dtype=dtype, device=device, requires_grad=requires_grad
+    ))
+    kwargs = {
+        "offset": offset,
+        "weight": weight,
+        "bias": bias,
+        "stride_h": stride[0],
+        "stride_w": stride[1],
+        "pad_h": padding[0],
+        "pad_w": padding[1],
+        "dilation_h": dilation[0],
+        "dilation_w": dilation[1],
+        "groups": groups,
+        "offset_groups": 1,
+        "use_mask": use_mask,
+        "mask": mask,  # no modulation in this test
+    }
+    optests.opcheck(torch.ops.torchvision.deform_conv2d, args=(x,), kwargs=kwargs)
+
 
 
 class TestFrozenBNT:

torchvision/csrc/ops/mps/deform_conv2d_kernel.mm

@@ -0,0 +1,134 @@
+#include <ATen/ATen.h>
+#include <ATen/mps/MPSProfiler.h>
+#include <ATen/native/mps/OperationUtils.h>
+#include "mps_kernels.h"
+
+namespace vision {
+namespace ops {
+
+namespace {
+
+at::Tensor deform_conv2d_forward_kernel(
+    const at::Tensor& input,
+    const at::Tensor& weight,
+    const at::Tensor& offset,
+    const at::Tensor& mask,
+    const at::Tensor& bias,
+    int64_t stride_h,
+    int64_t stride_w,
+    int64_t pad_h,
+    int64_t pad_w,
+    int64_t dilation_h,
+    int64_t dilation_w,
+    int64_t n_weight_grps,
+    int64_t n_offset_grps,
+    bool use_mask) {
+  using namespace at::native::mps;
+  at::Tensor input_c = input.contiguous();
+  at::Tensor weight_c = weight.contiguous();
+  at::Tensor offset_c = offset.contiguous();
+  at::Tensor mask_c = mask.contiguous();
+  at::Tensor bias_c = bias.contiguous();
+
+  TORCH_CHECK(input_c.ndimension() == 4, "Input tensor must be 4D");
+  TORCH_CHECK(weight_c.ndimension() == 4, "Weight tensor must be 4D");
+  TORCH_CHECK(offset_c.ndimension() == 4, "Offset tensor must be 4D");
+  TORCH_CHECK(!use_mask || mask_c.ndimension() == 4, "Mask tensor must be 4D if use_mask is true");
+  TORCH_CHECK(input_c.is_mps(), "input must be a MPS tensor");
+
+  at::DeviceGuard guard(input_c.device());
+
+  int batch = input_c.size(0);
+  int in_channels = input_c.size(1);
+  int in_h = input_c.size(2);
+  int in_w = input_c.size(3);
+  int weight_h = weight_c.size(2);
+  int weight_w = weight_c.size(3);
+  int out_channels = weight_c.size(0);
+  int ker_h = dilation_h * (weight_h - 1) + 1;
+  int ker_w = dilation_w * (weight_w - 1) + 1;
+  int out_h = ((in_h + 2 * pad_h - ker_h) / stride_h) + 1;
+  int out_w = ((in_w + 2 * pad_w - ker_w) / stride_w) + 1;
+
+  TORCH_CHECK(weight_c.size(1) * n_weight_grps == in_channels,
+    "Input channels (", in_channels, 
+    ") must equal weight.size(1) * n_weight_grps (", weight_c.size(1), " * ", n_weight_grps, ")");
+  TORCH_CHECK(weight_c.size(0) % n_weight_grps == 0,
+    "Weight tensor's out channels (", weight_c.size(0), 
+    ") must be divisible by n_weight_grps (", n_weight_grps, ")");
+  TORCH_CHECK(offset_c.size(1) == n_offset_grps * 2 * weight_h * weight_w,
+    "Offset tensor shape[1] is invalid: got ", offset_c.size(1), 
+    ", expected ", n_offset_grps * 2 * weight_h * weight_w);
+  TORCH_CHECK(!use_mask || mask_c.size(1) == n_offset_grps * weight_h * weight_w,
+    "Mask tensor shape[1] is invalid: got ", mask_c.size(1), 
+    ", expected ", n_offset_grps * weight_h * weight_w);
+  TORCH_CHECK(in_channels % n_offset_grps == 0,
+    "Input tensor channels (", in_channels, 
+    ") must be divisible by n_offset_grps (", n_offset_grps, ")");
+  TORCH_CHECK(offset_c.size(0) == batch,
+    "Offset tensor batch size (", offset_c.size(0),
+    ") must match input tensor batch size (", batch, ")");
+  TORCH_CHECK(offset_c.size(2) == out_h && offset_c.size(3) == out_w,
+    "Offset tensor spatial dimensions (", offset_c.size(2), ", ", offset_c.size(3), 
+    ") must match calculated output dimensions (", out_h, ", ", out_w, ")");
+  TORCH_CHECK(!use_mask || mask_c.size(0) == batch,
+    "Mask tensor batch size (", mask_c.size(0),
+    ") must match input tensor batch size (", batch, ")");
+  TORCH_CHECK(!use_mask || (mask_c.size(2) == out_h && mask_c.size(3) == out_w),
+    "Mask tensor spatial dimensions (", mask_c.size(2), ", ", mask_c.size(3),
+    ") must match calculated output dimensions (", out_h, ", ", out_w, ")");
+  TORCH_CHECK(out_h > 0 && out_w > 0,
+    "Calculated output size too small - out_h: ", out_h, " out_w: ", out_w);
+
+  auto columns = at::empty({in_channels * weight_h * weight_w, batch * out_h * out_w}, input_c.options());
+
+  id<MTLBuffer> inputBuffer  = getMTLBufferStorage(input_c);
+  id<MTLBuffer> offsetBuffer = getMTLBufferStorage(offset_c);
+  id<MTLBuffer> maskBuffer   = use_mask ? getMTLBufferStorage(mask_c) : nil;
+  id<MTLBuffer> outputBuffer = getMTLBufferStorage(columns);
+
+  id<MTLDevice> device = MPSDevice::getInstance()->device();
+  std::string kernelName = "deformable_im2col_" + scalarToMetalTypeString(input.scalar_type());
+  id<MTLComputePipelineState> pipelineState = mps::visionPipelineState(device, kernelName);
+
+  int num_kernels = in_channels * out_h * out_w * batch;
+  NSUInteger threadsPerThreadgroup = pipelineState.maxTotalThreadsPerThreadgroup;
+  NSUInteger threadgroups = (num_kernels + threadsPerThreadgroup - 1) / threadsPerThreadgroup;
+  MTLSize threadGroupSize = MTLSizeMake(threadsPerThreadgroup, 1, 1);
+  MTLSize threadgroupsPerGrid = MTLSizeMake(threadgroups, 1, 1);
+
+  MPSStream* mpsStream = getCurrentMPSStream();
+  dispatch_sync(mpsStream->queue(), ^{
+    @autoreleasepool {
+      id<MTLComputeCommandEncoder> computeEncoder = mpsStream->commandEncoder();
+      [computeEncoder setComputePipelineState:pipelineState];
+      at::native::mps::mtl_setArgs(computeEncoder, inputBuffer, offsetBuffer, maskBuffer,
+                                   in_h, in_w, weight_h, weight_w, pad_h, pad_w, stride_h, stride_w, 
+                                   dilation_h, dilation_w, batch, in_channels, n_offset_grps, out_h, out_w,
+                                   use_mask, outputBuffer);
+      [computeEncoder dispatchThreadgroups:threadgroupsPerGrid threadsPerThreadgroup:threadGroupSize];
+    }
+  });
+  int in_channels_per_grp = in_channels / n_weight_grps;
+  int out_channels_per_grp = out_channels / n_weight_grps;
+  auto weight_grouped = weight_c.view({n_weight_grps, out_channels_per_grp, in_channels_per_grp, weight_h, weight_w});
+  auto columns_grouped = columns.view({n_weight_grps,
+                                      (in_channels * weight_h * weight_w) / n_weight_grps,
+                                      batch * out_h * out_w});
+  auto weight_reshaped = weight_grouped.reshape({n_weight_grps, out_channels_per_grp, -1});
+  auto out_grouped = at::bmm(weight_reshaped, columns_grouped);
+  auto out = out_grouped.reshape({n_weight_grps * out_channels_per_grp, batch, out_h, out_w})
+              .transpose(0, 1);
+  return out + bias_c.view({1, out_channels, 1, 1});
+}
+
+} // namespace
+
+TORCH_LIBRARY_IMPL(torchvision, MPS, m) {
+  m.impl(
+      TORCH_SELECTIVE_NAME("torchvision::deform_conv2d"),
+      TORCH_FN(deform_conv2d_forward_kernel));
+}
+
+} // namespace ops
+} // namespace vision

torchvision/csrc/ops/mps/mps_kernels.h

@@ -91,6 +91,52 @@ inline T bilinear_interpolate(
   return val;
 }
 
+template <typename T, typename integer_t>
+inline T bilinear_interpolate_deformable_conv2d(
+    constant T* input,
+    integer_t height,
+    integer_t width,
+    T y,
+    T x,
+    uint index /* index for debug only*/) {
+  if (y <= -1.0 || y >= height || x <= -1.0 || x >= width) {
+    return 0;
+  }
+  integer_t y_low = static_cast<integer_t>(floor(y));
+  integer_t x_low = static_cast<integer_t>(floor(x));
+  integer_t y_high = y_low + 1;
+  integer_t x_high = x_low + 1;
+
+  T ly = y - static_cast<T>(y_low);
+  T lx = x - static_cast<T>(x_low);
+  T hh = 1.0 - ly;
+  T hw = 1.0 - lx;
+
+  T v1 = 0;
+  if (y_low >= 0 && x_low >= 0)
+    v1 = input[y_low * width + x_low];
+
+  T v2 = 0;
+  if (y_low >= 0 && x_high <= width - 1)
+    v2 = input[y_low * width + x_high];
+
+  T v3 = 0;
+  if (y_high <= height - 1 && x_low >= 0)
+    v3 = input[y_high * width + x_low];
+
+  T v4 = 0;
+  if (y_high <= height - 1 && x_high <= width - 1)
+    v4 = input[y_high * width + x_high];
+
+  T w1 = hh * hw;
+  T w2 = hh * lx;
+  T w3 = ly * hw;
+  T w4 = ly * lx;
+
+  T val = w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4;
+  return val;
+}
+
 template <typename T, typename integer_t>
 inline void bilinear_interpolate_gradient(
     integer_t height,
@@ -225,6 +271,117 @@ kernel void nms<DTYPE ## 4, DTYPE>(                        \
   uint2    tgid   [[threadgroup_position_in_grid]],        \
   uint2    tid2   [[thread_position_in_threadgroup]]);
 
+
+template<typename T>
+kernel void deformable_im2col_kernel(
+    constant T*           input_ptr     [[ buffer(0) ]],
+    constant T*           offset_ptr    [[ buffer(1) ]],
+    constant T*           mask_ptr      [[ buffer(2) ]],
+    constant int&         height        [[ buffer(3) ]],
+    constant int&         width         [[ buffer(4) ]],
+    constant int&         weight_h      [[ buffer(5) ]],
+    constant int&         weight_w      [[ buffer(6) ]],
+    constant int&         pad_h         [[ buffer(7) ]],
+    constant int&         pad_w         [[ buffer(8) ]],
+    constant int&         stride_h      [[ buffer(9) ]],
+    constant int&         stride_w      [[ buffer(10)]],
+    constant int&         dilation_h    [[ buffer(11)]],
+    constant int&         dilation_w    [[ buffer(12)]],
+    constant int&         batch_size      [[ buffer(13)]],
+    constant int&         n_in_channels [[ buffer(14)]],
+    constant int&         n_offset_grps [[ buffer(15)]],
+    constant int&         out_h         [[ buffer(16)]],
+    constant int&         out_w         [[ buffer(17)]],
+    constant bool&        use_mask      [[ buffer(18)]],
+    device T*             columns_ptr   [[ buffer(19)]],
+    uint                  tid           [[ thread_position_in_grid ]],
+    uint                  tpg           [[ threads_per_grid ]])
+{
+    int total = out_w * out_h * batch_size * n_in_channels;
+    int gridSize = tpg;
+    if (tid >= total) {
+        return;
+    }
+
+    int out_x = tid % out_w;
+    int out_y = (tid / out_w) % out_h;
+    int out_b = (tid / (out_w * out_h)) % batch_size;
+    int in_c  = tid / (out_w * out_h * batch_size);
+    int out_c = in_c * weight_h * weight_w;
+
+    int c_per_offset_grp = n_in_channels / n_offset_grps;
+    int grp_idx = in_c / c_per_offset_grp;
+
+    int col_offset = out_c * (batch_size * out_h * out_w)
+                      + out_b * (out_h * out_w)
+                      + out_y * out_w + out_x;
+    device T* local_columns_ptr = columns_ptr + col_offset;
+
+    int input_offset = out_b * (n_in_channels * height * width)
+                        + in_c * (height * width);
+    constant T* local_input_ptr = input_ptr + input_offset;
+
+    int offset_offset = (out_b * n_offset_grps + grp_idx) * 2 * weight_h * weight_w * out_h * out_w;
+    constant T* local_offset_ptr = offset_ptr + offset_offset;
+
+    constant T* local_mask_ptr = nullptr;
+    if (use_mask) {
+        int mask_offset = (out_b * n_offset_grps + grp_idx) * weight_h * weight_w * out_h * out_w;
+        local_mask_ptr = mask_ptr + mask_offset;
+    }
+
+    for (int i = 0; i < weight_h; ++i) {
+        for (int j = 0; j < weight_w; ++j) {
+            int mask_index = i * weight_w + j;
+            int offset_index = 2 * mask_index;
+
+            T mask_value = 1;
+            if (use_mask) {
+                mask_value = local_mask_ptr[mask_index * (out_h * out_w) + out_y * out_w + out_x];
+            }
+
+            T offset_h_val = local_offset_ptr[offset_index * (out_h * out_w) + out_y * out_w + out_x];
+            T offset_w_val = local_offset_ptr[(offset_index + 1) * (out_h * out_w) + out_y * out_w + out_x];
+
+            T y = (out_y * stride_h - pad_h) + i * dilation_h + offset_h_val;
+            T x = (out_x * stride_w - pad_w) + j * dilation_w + offset_w_val;
+
+            T interp = bilinear_interpolate_deformable_conv2d(local_input_ptr, height, width, y, x, tid);
+
+            *local_columns_ptr = mask_value * interp;
+
+            local_columns_ptr += batch_size * out_h * out_w;
+        }
+    }
+}
+
+#define REGISTER_DEFORMABLE_IM2COL_OP(DTYPE)                                  \
+template                                                                      \
+[[host_name("deformable_im2col_" #DTYPE)]]                                    \
+kernel void deformable_im2col_kernel<DTYPE>(                                  \
+    constant DTYPE*           input_ptr    [[ buffer(0) ]],                   \
+    constant DTYPE*           offset_ptr   [[ buffer(1) ]],                   \
+    constant DTYPE*           mask_ptr     [[ buffer(2) ]],                   \
+    constant int&                 height       [[ buffer(3) ]],               \
+    constant int&                 width        [[ buffer(4) ]],               \
+    constant int&                 weight_h     [[ buffer(5) ]],               \
+    constant int&                 weight_w     [[ buffer(6) ]],               \
+    constant int&                 pad_h        [[ buffer(7) ]],               \
+    constant int&                 pad_w        [[ buffer(8) ]],               \
+    constant int&                 stride_h     [[ buffer(9) ]],               \
+    constant int&                 stride_w     [[ buffer(10)]],               \
+    constant int&                 dilation_h   [[ buffer(11)]],               \
+    constant int&                 dilation_w   [[ buffer(12)]],               \
+    constant int&                 batch_sz     [[ buffer(13)]],               \
+    constant int&                 n_in_channels[[ buffer(14)]],               \
+    constant int&                 n_offset_grps[[ buffer(15)]],               \
+    constant int&                 out_h        [[ buffer(16)]],               \
+    constant int&                 out_w        [[ buffer(17)]],               \
+    constant bool&                use_mask     [[ buffer(18)]],               \
+    device DTYPE*                 columns_ptr  [[ buffer(19)]],               \
+    uint                          tid          [[ thread_position_in_grid ]], \
+    uint                          tpg           [[ threads_per_grid ]]);
+
 template<typename T, typename integer_t>
 kernel void roi_align(
     constant T       * input          [[buffer(0)]],
@@ -1013,6 +1170,8 @@ kernel void ps_roi_pool_backward<DTYPE, INT_DTYPE>(          \
 
 REGISTER_NMS_OP(float);
 REGISTER_NMS_OP(half);
+REGISTER_DEFORMABLE_IM2COL_OP(float);
+REGISTER_DEFORMABLE_IM2COL_OP(half);
 REGISTER_ROI_ALIGN_OP(float, int64_t);
 REGISTER_ROI_ALIGN_OP(half, int64_t);
 REGISTER_ROI_ALIGN_BACKWARD_OP(float, int64_t);