diff --git a/setup.py b/setup.py index ad728ef7f18..d6c365ff052 100644 --- a/setup.py +++ b/setup.py @@ -163,6 +163,7 @@ def get_macros_and_flags(): CSRS_DIR / "ops/nms.cpp", CSRS_DIR / "ops/cpu/nms_kernel.cpp", CSRS_DIR / "ops/mps/nms_kernel.mm", + CSRS_DIR / "ops/mps/deform_conv2d_kernel.mm", CSRS_DIR / "ops/quantized/cpu/qnms_kernel.cpp", CSRS_DIR / "io/image/common_stable.cpp", CSRS_DIR / "io/image/cpu/encode_png.cpp", diff --git a/torchvision/csrc/ops/mps/deform_conv2d_kernel.mm b/torchvision/csrc/ops/mps/deform_conv2d_kernel.mm index 63371365655..bc0d52ac6b3 100644 --- a/torchvision/csrc/ops/mps/deform_conv2d_kernel.mm +++ b/torchvision/csrc/ops/mps/deform_conv2d_kernel.mm @@ -1,19 +1,100 @@ -#include -#include -#include -#include "mps_kernels.h" +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "../StableABICompat.h" +#include "deform_conv2d_metal_shader.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, +using torch::stable::Tensor; + +AOTIMetalShaderLibraryHandle deform_conv2d_shader_library() { + static AOTIMetalShaderLibraryHandle library = []() { + AOTIMetalShaderLibraryHandle handle = nullptr; + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_create_shader_library( + deform_conv2d_metal_shader, &handle)); + return handle; + }(); + return library; +} + +const char* metal_type_string(torch::headeronly::ScalarType scalar_type) { + if (scalar_type == torch::headeronly::ScalarType::Float) { + return "float"; + } + if (scalar_type == torch::headeronly::ScalarType::Half) { + return "half"; + } + return ""; +} + +// The im2col kernel takes its (h, w)-style pairs as Metal `int2&` params (8 +// bytes). The shim only sets int64 args, so pack each pair into one int64 +// (little-endian: x in the low 32 bits, y in the high 32 bits). Single `int&` +// params read the low 32 bits of the int64, and `bool&` reads the low byte. +int64_t pack_int2(int64_t x, int64_t y) { + return static_cast( + static_cast(x) | (static_cast(static_cast(y)) << 32)); +} + +struct DeformIm2colArgs { + AtenTensorHandle input; + AtenTensorHandle offset; + AtenTensorHandle mask; + AtenTensorHandle columns; + int64_t input_size; + int64_t weight_size; + int64_t pad; + int64_t stride; + int64_t dilation; + int64_t batch; + int64_t in_channels; + int64_t n_offset_grps; + int64_t out_size; + int64_t use_mask; + uint64_t num_kernels; +}; + +void deform_im2col_encode( + AOTIMetalKernelFunctionHandle func, + void* user_data) { + const auto* a = static_cast(user_data); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_start_encoding(func)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_tensor(func, 0, a->input)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_tensor(func, 1, a->offset)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_tensor(func, 2, a->mask)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 3, a->input_size)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 4, a->weight_size)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 5, a->pad)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 6, a->stride)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 7, a->dilation)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 8, a->batch)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 9, a->in_channels)); + TORCH_ERROR_CODE_CHECK( + aoti_torch_mps_set_arg_int(func, 10, a->n_offset_grps)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 11, a->out_size)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 12, a->use_mask)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_tensor(func, 13, a->columns)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_dispatch_single(func, a->num_kernels)); +} + +Tensor deform_conv2d_forward_kernel( + const Tensor& input, + const Tensor& weight, + const Tensor& offset, + const Tensor& mask, + const Tensor& bias, int64_t stride_h, int64_t stride_w, int64_t pad_h, @@ -23,126 +104,106 @@ 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"); - TORCH_CHECK(weight.is_mps(), "weight must be a MPS tensor"); - TORCH_CHECK(offset.is_mps(), "offset must be a MPS tensor"); - TORCH_CHECK(mask.is_mps(), "mask must be a MPS tensor"); - TORCH_CHECK(bias.is_mps(), "bias must be a MPS tensor"); - - at::DeviceGuard guard(input_c.device()); - - uint32_t batch = input_c.size(0); - uint32_t in_channels = input_c.size(1); - uint32_t in_h = input_c.size(2); - uint32_t in_w = input_c.size(3); - uint32_t weight_h = weight_c.size(2); - uint32_t weight_w = weight_c.size(3); - uint32_t out_channels = weight_c.size(0); - uint32_t ker_h = dilation_h * (weight_h - 1) + 1; - uint32_t ker_w = dilation_w * (weight_w - 1) + 1; - uint32_t out_h = ((in_h + 2 * pad_h - ker_h) / stride_h) + 1; - uint32_t out_w = ((in_w + 2 * pad_w - ker_w) / stride_w) + 1; - uint32_t pad_h_u = static_cast(pad_h); - uint32_t pad_w_u = static_cast(pad_w); - uint32_t stride_h_u = static_cast(stride_h); - uint32_t stride_w_u = static_cast(stride_w); - uint32_t dilation_h_u = static_cast(dilation_h); - uint32_t dilation_w_u = static_cast(dilation_w); - - 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 inputBuffer = getMTLBufferStorage(input_c); - id offsetBuffer = getMTLBufferStorage(offset_c); - id maskBuffer = use_mask ? getMTLBufferStorage(mask_c) : nil; - id outputBuffer = getMTLBufferStorage(columns); - - id device = MPSDevice::getInstance()->device(); - std::string kernelName = "deformable_im2col_" + scalarToMetalTypeString(input.scalar_type()); - id 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 computeEncoder = mpsStream->commandEncoder(); - [computeEncoder setComputePipelineState:pipelineState]; - at::native::mps::mtl_setArgs(computeEncoder, inputBuffer, offsetBuffer, maskBuffer, - std::array{in_h, in_w}, - std::array{weight_h, weight_w}, - std::array{pad_h_u, pad_w_u}, - std::array{stride_h_u, stride_w_u}, - std::array{dilation_h_u, dilation_w_u}, - batch, in_channels, n_offset_grps, - std::array{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}); + Tensor input_c = torch::stable::contiguous(input); + Tensor weight_c = torch::stable::contiguous(weight); + Tensor offset_c = torch::stable::contiguous(offset); + Tensor mask_c = torch::stable::contiguous(mask); + Tensor bias_c = torch::stable::contiguous(bias); + + STD_TORCH_CHECK(input_c.dim() == 4, "Input tensor must be 4D"); + STD_TORCH_CHECK(weight_c.dim() == 4, "Weight tensor must be 4D"); + STD_TORCH_CHECK(offset_c.dim() == 4, "Offset tensor must be 4D"); + STD_TORCH_CHECK( + !use_mask || mask_c.dim() == 4, + "Mask tensor must be 4D if use_mask is true"); + STD_TORCH_CHECK( + input_c.device().type() == torch::headeronly::DeviceType::MPS, + "input must be a MPS tensor"); + + int64_t batch = input_c.size(0); + int64_t in_channels = input_c.size(1); + int64_t in_h = input_c.size(2); + int64_t in_w = input_c.size(3); + int64_t weight_h = weight_c.size(2); + int64_t weight_w = weight_c.size(3); + int64_t out_channels = weight_c.size(0); + int64_t ker_h = dilation_h * (weight_h - 1) + 1; + int64_t ker_w = dilation_w * (weight_w - 1) + 1; + int64_t out_h = ((in_h + 2 * pad_h - ker_h) / stride_h) + 1; + int64_t out_w = ((in_w + 2 * pad_w - ker_w) / stride_w) + 1; + + STD_TORCH_CHECK( + weight_c.size(1) * n_weight_grps == in_channels, + "Input channels must equal weight.size(1) * n_weight_grps"); + STD_TORCH_CHECK( + out_channels % n_weight_grps == 0, + "Weight tensor's out channels must be divisible by n_weight_grps"); + STD_TORCH_CHECK(out_h > 0 && out_w > 0, "Calculated output size too small"); + + Tensor columns = torch::stable::new_empty( + input_c, {in_channels * weight_h * weight_w, batch * out_h * out_w}); + + const std::string kernel = "deformable_im2col_" + + std::string(metal_type_string(input_c.scalar_type())); + AOTIMetalKernelFunctionHandle func = nullptr; + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_get_kernel_function( + deform_conv2d_shader_library(), kernel.c_str(), &func)); + + DeformIm2colArgs args{ + input_c.get(), + offset_c.get(), + // When use_mask is false the kernel never reads the mask buffer; bind the + // input handle as a harmless placeholder so a valid buffer is always set. + use_mask ? mask_c.get() : input_c.get(), + columns.get(), + pack_int2(in_h, in_w), + pack_int2(weight_h, weight_w), + pack_int2(pad_h, pad_w), + pack_int2(stride_h, stride_w), + pack_int2(dilation_h, dilation_w), + batch, + in_channels, + n_offset_grps, + pack_int2(out_h, out_w), + use_mask ? 1 : 0, + static_cast(in_channels * out_h * out_w * batch)}; + TORCH_ERROR_CODE_CHECK( + aoti_torch_mps_run_command_block(func, &deform_im2col_encode, &args)); + + int64_t in_channels_per_grp = in_channels / n_weight_grps; + int64_t out_channels_per_grp = out_channels / n_weight_grps; + Tensor weight_grouped = torch::stable::view( + weight_c, + {n_weight_grps, + out_channels_per_grp, + in_channels_per_grp, + weight_h, + weight_w}); + Tensor columns_grouped = torch::stable::view( + columns, + {n_weight_grps, + (in_channels * weight_h * weight_w) / n_weight_grps, + batch * out_h * out_w}); + Tensor weight_reshaped = torch::stable::reshape( + weight_grouped, {n_weight_grps, out_channels_per_grp, -1}); + Tensor out_grouped = torch::stable::matmul(weight_reshaped, columns_grouped); + Tensor out = torch::stable::transpose( + torch::stable::reshape( + out_grouped, + {n_weight_grps * out_channels_per_grp, batch, out_h, out_w}), + 0, + 1); + Tensor bias_view = + torch::stable::view(bias_c, {1, out_channels, 1, 1}); + // subtract(out, bias_view, alpha=-1) computes out - (-1) * bias_view, i.e. + // out + bias_view; stable ops.h ships subtract but not add. + return torch::stable::subtract(out, bias_view, /*alpha=*/-1.0); } } // namespace -TORCH_LIBRARY_IMPL(torchvision, MPS, m) { - m.impl( - TORCH_SELECTIVE_NAME("torchvision::deform_conv2d"), - TORCH_FN(deform_conv2d_forward_kernel)); +STABLE_TORCH_LIBRARY_IMPL(torchvision, MPS, m) { + m.impl("deform_conv2d", TORCH_BOX(&deform_conv2d_forward_kernel)); } } // namespace ops diff --git a/torchvision/csrc/ops/mps/deform_conv2d_metal_shader.h b/torchvision/csrc/ops/mps/deform_conv2d_metal_shader.h new file mode 100644 index 00000000000..74f0d3e6e10 --- /dev/null +++ b/torchvision/csrc/ops/mps/deform_conv2d_metal_shader.h @@ -0,0 +1,174 @@ +#pragma once + +// Metal shader source for the deform_conv2d MPS forward kernel (deformable +// im2col). Carved out of ops/mps/mps_kernels.h so it can compile into the +// stable-ABI _C_stable extension without pulling in +// at::native::mps::MetalShaderLibrary (unavailable under +// -DTORCH_TARGET_VERSION). + +namespace vision { +namespace ops { + +static const char* deform_conv2d_metal_shader = R"VISION_METAL( + +#include +using namespace metal; + +template +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(floor(y)); + integer_t x_low = static_cast(floor(x)); + integer_t y_high = y_low + 1; + integer_t x_high = x_low + 1; + + T ly = y - static_cast(y_low); + T lx = x - static_cast(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 +kernel void deformable_im2col_kernel( + constant T* input_ptr [[ buffer(0) ]], + constant T* offset_ptr [[ buffer(1) ]], + constant T* mask_ptr [[ buffer(2) ]], + constant int2& input_size [[ buffer(3) ]], // (height, width) + constant int2& weight_size [[ buffer(4) ]], // (weight_h, weight_w) + constant int2& pad [[ buffer(5) ]], // (pad_h, pad_w) + constant int2& stride [[ buffer(6) ]], // (stride_h, stride_w) + constant int2& dilation [[ buffer(7) ]], // (dilation_h, dilation_w) + constant int& batch_size [[ buffer(8) ]], + constant int& n_in_channels [[ buffer(9) ]], + constant int& n_offset_grps [[ buffer(10)]], + constant int2& out_size [[ buffer(11)]], // (out_h, out_w) + constant bool& use_mask [[ buffer(12)]], + device T* columns_ptr [[ buffer(13)]], + uint tid [[ thread_position_in_grid ]], + uint tpg [[ threads_per_grid ]] +) +{ + int height = input_size.x, width = input_size.y; + int weight_h = weight_size.x, weight_w = weight_size.y; + int pad_h = pad.x, pad_w = pad.y; + int stride_h = stride.x, stride_w = stride.y; + int dilation_h = dilation.x, dilation_w = dilation.y; + int out_h = out_size.x, out_w = out_size.y; + + int total = out_w * out_h * batch_size * n_in_channels; + 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( \ + constant DTYPE* input_ptr [[ buffer(0) ]], \ + constant DTYPE* offset_ptr [[ buffer(1) ]], \ + constant DTYPE* mask_ptr [[ buffer(2) ]], \ + constant int2& input_size [[ buffer(3) ]], /* (h, w) */ \ + constant int2& weight_size [[ buffer(4) ]], /* (h, w) */ \ + constant int2& pad [[ buffer(5) ]], /* (h, w) */ \ + constant int2& stride [[ buffer(6) ]], /* (h, w) */ \ + constant int2& dilation [[ buffer(7) ]], /* (h, w) */ \ + constant int& batch_size [[ buffer(8) ]], \ + constant int& n_in_channels [[ buffer(9) ]], \ + constant int& n_offset_grps [[ buffer(10)]], \ + constant int2& out_size [[ buffer(11)]], /* (h, w) */ \ + constant bool& use_mask [[ buffer(12)]], \ + device DTYPE* columns_ptr [[ buffer(13)]], \ + uint tid [[ thread_position_in_grid ]], \ + uint tpg [[ threads_per_grid ]]); + +REGISTER_DEFORMABLE_IM2COL_OP(float); +REGISTER_DEFORMABLE_IM2COL_OP(half); + +)VISION_METAL"; + +} // namespace ops +} // namespace vision