diff --git a/setup.py b/setup.py index ad728ef7f18..7bc4bfab7e6 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/ps_roi_pool_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/ps_roi_pool_kernel.mm b/torchvision/csrc/ops/mps/ps_roi_pool_kernel.mm index 75d0ff4845f..fc23fa46d3e 100644 --- a/torchvision/csrc/ops/mps/ps_roi_pool_kernel.mm +++ b/torchvision/csrc/ops/mps/ps_roi_pool_kernel.mm @@ -1,198 +1,260 @@ -#include -#include -#include "mps_helpers.h" -#include "mps_kernels.h" +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "../StableABICompat.h" +#include "ps_roi_pool_metal_shader.h" namespace vision { namespace ops { namespace { -std::tuple ps_roi_pool_forward_kernel(const at::Tensor& input, - const at::Tensor& rois, - double spatial_scale, - int64_t pooled_height, - int64_t pooled_width) { - using namespace at::native::mps; - TORCH_CHECK(input.is_mps(), "input must be a MPS tensor"); - TORCH_CHECK(rois.is_mps(), "rois must be a MPS tensor"); - TORCH_CHECK(rois.size(1) == 5, "rois must have shape as Tensor[K, 5]"); +using torch::stable::Tensor; - at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2}; +AOTIMetalShaderLibraryHandle ps_roi_pool_shader_library() { + static AOTIMetalShaderLibraryHandle library = []() { + AOTIMetalShaderLibraryHandle handle = nullptr; + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_create_shader_library( + ps_roi_pool_metal_shader, &handle)); + return handle; + }(); + return library; +} - at::CheckedFrom c = "ps_roi_pool_forward_kernel"; - at::checkAllSameGPU(c, {input_t, rois_t}); - at::checkAllSameType(c, {input_t, rois_t}); +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 ""; +} + +// spatial_scale rides in as a 1-element float32 tensor: the MPS shim has no +// scalar-float arg setter yet (same workaround as nms's iou_threshold). +Tensor make_scalar_tensor(const Tensor& ref, double value) { + Tensor t = torch::stable::new_empty( + ref, {1}, torch::headeronly::ScalarType::Float); + torch::stable::fill_(t, value); + return t; +} + +struct PsRoiPoolForwardArgs { + AtenTensorHandle input; + AtenTensorHandle rois; + AtenTensorHandle output; + AtenTensorHandle channel_mapping; + int64_t output_size; + int64_t channels; + int64_t height; + int64_t width; + int64_t pooled_height; + int64_t pooled_width; + int64_t channels_out; + AtenTensorHandle spatial_scale; +}; + +void ps_roi_pool_forward_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->rois)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_tensor(func, 2, a->output)); + TORCH_ERROR_CODE_CHECK( + aoti_torch_mps_set_arg_tensor(func, 3, a->channel_mapping)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 4, a->output_size)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 5, a->channels)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 6, a->height)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 7, a->width)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 8, a->pooled_height)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 9, a->pooled_width)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 10, a->channels_out)); + TORCH_ERROR_CODE_CHECK( + aoti_torch_mps_set_arg_tensor(func, 11, a->spatial_scale)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_dispatch_single( + func, static_cast(a->output_size))); +} + +struct PsRoiPoolBackwardArgs { + AtenTensorHandle grad; + AtenTensorHandle rois; + AtenTensorHandle channel_mapping; + AtenTensorHandle grad_input; + int64_t output_size; + int64_t channels; + int64_t height; + int64_t width; + int64_t pooled_height; + int64_t pooled_width; + int64_t channels_out; + AtenTensorHandle spatial_scale; +}; + +void ps_roi_pool_backward_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->grad)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_tensor(func, 1, a->rois)); + TORCH_ERROR_CODE_CHECK( + aoti_torch_mps_set_arg_tensor(func, 2, a->channel_mapping)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_tensor(func, 3, a->grad_input)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 4, a->output_size)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 5, a->channels)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 6, a->height)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 7, a->width)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 8, a->pooled_height)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 9, a->pooled_width)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_set_arg_int(func, 10, a->channels_out)); + TORCH_ERROR_CODE_CHECK( + aoti_torch_mps_set_arg_tensor(func, 11, a->spatial_scale)); + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_dispatch_single( + func, static_cast(a->output_size))); +} + +std::tuple ps_roi_pool_forward_kernel( + const Tensor& input, + const Tensor& rois, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width) { + STD_TORCH_CHECK( + input.device().type() == torch::headeronly::DeviceType::MPS, + "input must be a MPS tensor"); + STD_TORCH_CHECK( + rois.device().type() == torch::headeronly::DeviceType::MPS, + "rois must be a MPS tensor"); + STD_TORCH_CHECK(rois.size(1) == 5, "rois must have shape as Tensor[K, 5]"); int64_t num_rois = rois.size(0); int64_t channels = input.size(1); int64_t height = input.size(2); int64_t width = input.size(3); - float spatial_scale_f = static_cast(spatial_scale); - - TORCH_CHECK(channels % (pooled_height * pooled_width) == 0, - "input channels must be a multiple of pooling height * pooling width"); + STD_TORCH_CHECK( + channels % (pooled_height * pooled_width) == 0, + "input channels must be a multiple of pooling height * pooling width"); int64_t channels_out = channels / (pooled_height * pooled_width); - auto output = at::zeros({num_rois, channels_out, pooled_height, pooled_width}, input.options()); - auto channel_mapping = at::zeros(output.sizes(), input.options().dtype(at::kLong)); - auto output_size = output.numel(); - + Tensor output = torch::stable::new_empty( + input, {num_rois, channels_out, pooled_height, pooled_width}); + Tensor channel_mapping = torch::stable::new_empty( + input, + {num_rois, channels_out, pooled_height, pooled_width}, + torch::headeronly::ScalarType::Long); + int64_t output_size = output.numel(); if (output_size == 0) { return std::make_tuple(output, channel_mapping); } - auto input_ = input.contiguous(); - auto rois_ = rois.contiguous(); - - id inputBuffer = getMTLBufferStorage(input_); - id roisBuffer = getMTLBufferStorage(rois_); - id outputBuffer = getMTLBufferStorage(output); - id channelMappingBuffer = getMTLBufferStorage(channel_mapping); - id device = MPSDevice::getInstance()->device(); - MPSStream* mpsStream = getCurrentMPSStream(); - dispatch_sync(mpsStream->queue(), ^() { - @autoreleasepool { - id computeEncoder = mpsStream->commandEncoder(); - MTLSize threadgroupsPerGrid = MTLSizeMake( - std::min(ceil_div(static_cast(output_size), static_cast(512)), static_cast(4096)), - 1, - 1); - - const std::string kernel = "ps_roi_pool_" + scalarToMetalTypeString(input.scalar_type()); - id visionPSO = mps::visionPipelineState(device, kernel); - - // this function call is a no-op if MPS Profiler is not enabled - getMPSProfiler().beginProfileKernel(visionPSO, kernel, {input_, rois_}); - - [computeEncoder setComputePipelineState:visionPSO]; - // [N, C, H, W] - [computeEncoder setBuffer:inputBuffer offset:input_.storage_offset() * input_.element_size() atIndex:0]; - [computeEncoder setBuffer:roisBuffer offset:rois_.storage_offset() * rois_.element_size() atIndex:1]; - [computeEncoder setBuffer:outputBuffer offset:output.storage_offset() * output.element_size() atIndex:2]; - [computeEncoder setBuffer:channelMappingBuffer - offset:channel_mapping.storage_offset() * channel_mapping.element_size() - atIndex:3]; - - [computeEncoder setBytes:&output_size length:sizeof(int64_t) atIndex:4]; - [computeEncoder setBytes:&channels length:sizeof(int64_t) atIndex:5]; - [computeEncoder setBytes:&height length:sizeof(int64_t) atIndex:6]; - [computeEncoder setBytes:&width length:sizeof(int64_t) atIndex:7]; - [computeEncoder setBytes:&pooled_height length:sizeof(int64_t) atIndex:8]; - [computeEncoder setBytes:&pooled_width length:sizeof(int64_t) atIndex:9]; - [computeEncoder setBytes:&channels_out length:sizeof(int64_t) atIndex:10]; - [computeEncoder setBytes:&spatial_scale_f length:sizeof(float) atIndex:11]; - - // A threadGroup is equivalent to a cuda's block. - NSUInteger tgSize = visionPSO.maxTotalThreadsPerThreadgroup; - if (tgSize > threadsPerBlock) { - tgSize = threadsPerBlock; - } - - MTLSize threadGroupSize = MTLSizeMake(tgSize, 1, 1); - [computeEncoder dispatchThreadgroups:threadgroupsPerGrid threadsPerThreadgroup:threadGroupSize]; - - getMPSProfiler().endProfileKernel(visionPSO); - } - }); + Tensor input_ = torch::stable::contiguous(input); + Tensor rois_ = torch::stable::contiguous(rois); + Tensor spatial_scale_t = make_scalar_tensor(input_, spatial_scale); + + const std::string kernel = + "ps_roi_pool_" + std::string(metal_type_string(input_.scalar_type())); + AOTIMetalKernelFunctionHandle func = nullptr; + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_get_kernel_function( + ps_roi_pool_shader_library(), kernel.c_str(), &func)); + + PsRoiPoolForwardArgs args{ + input_.get(), + rois_.get(), + output.get(), + channel_mapping.get(), + output_size, + channels, + height, + width, + pooled_height, + pooled_width, + channels_out, + spatial_scale_t.get()}; + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_run_command_block( + func, &ps_roi_pool_forward_encode, &args)); return std::make_tuple(output, channel_mapping); } -at::Tensor ps_roi_pool_backward_kernel(const at::Tensor& grad, - const at::Tensor& rois, - const at::Tensor& channel_mapping, - double spatial_scale, - int64_t pooled_height, - int64_t pooled_width, - int64_t batch_size, - int64_t channels, - int64_t height, - int64_t width) { - using namespace at::native::mps; - TORCH_CHECK(grad.is_mps(), "grad must be a MPS tensor"); - TORCH_CHECK(rois.is_mps(), "rois must be a MPS tensor"); - TORCH_CHECK(grad.scalar_type() != at::kHalf, "MPS does not support ps_roi_pool backward with float16 inputs."); - TORCH_CHECK(channel_mapping.is_mps(), "channel_mapping must be a MPS tensor"); - - at::TensorArg grad_t{grad, "grad", 1}, rois_t{rois, "rois", 2}, - channel_mapping_t{channel_mapping, "channel_mapping", 3}; - - at::CheckedFrom c = "ps_roi_pool_backward_kernel"; - at::checkAllSameGPU(c, {grad_t, rois_t, channel_mapping_t}); - at::checkAllSameType(c, {grad_t, rois_t}); - - float spatial_scale_f = static_cast(spatial_scale); - - auto grad_input = at::zeros({batch_size, channels, height, width}, grad.options()); - +Tensor ps_roi_pool_backward_kernel( + const Tensor& grad, + const Tensor& rois, + const Tensor& channel_mapping, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width, + int64_t batch_size, + int64_t channels, + int64_t height, + int64_t width) { + STD_TORCH_CHECK( + grad.device().type() == torch::headeronly::DeviceType::MPS, + "grad must be a MPS tensor"); + STD_TORCH_CHECK( + rois.device().type() == torch::headeronly::DeviceType::MPS, + "rois must be a MPS tensor"); + STD_TORCH_CHECK( + channel_mapping.device().type() == torch::headeronly::DeviceType::MPS, + "channel_mapping must be a MPS tensor"); + STD_TORCH_CHECK( + grad.scalar_type() != torch::headeronly::ScalarType::Half, + "MPS does not support ps_roi_pool backward with float16 inputs."); + + Tensor grad_input = torch::stable::new_zeros( + grad, {batch_size, channels, height, width}); if (grad.numel() == 0) { return grad_input; } int64_t channels_out = channels / (pooled_height * pooled_width); - int64_t output_size = grad.numel(); - - at::globalContext().alertNotDeterministic("ps_roi_pool_backward_kernel"); - auto grad_ = grad.contiguous(), rois_ = rois.contiguous(); - - id inputBuffer = getMTLBufferStorage(grad_); - id roisBuffer = getMTLBufferStorage(rois_); - id channelMappingBuffer = getMTLBufferStorage(channel_mapping); - id outputBuffer = getMTLBufferStorage(grad_input); - id device = MPSDevice::getInstance()->device(); - MPSStream* mpsStream = getCurrentMPSStream(); - dispatch_sync(mpsStream->queue(), ^() { - @autoreleasepool { - id computeEncoder = mpsStream->commandEncoder(); - MTLSize threadgroupsPerGrid = MTLSizeMake( - std::min(ceil_div(static_cast(grad.numel()), static_cast(512)), static_cast(4096)), - 1, - 1); - - const std::string kernel = "ps_roi_pool_backward_" + scalarToMetalTypeString(grad.scalar_type()); - id visionPSO = mps::visionPipelineState(device, kernel); - - // this function call is a no-op if MPS Profiler is not enabled - getMPSProfiler().beginProfileKernel(visionPSO, kernel, {grad_, rois_, channel_mapping}); - - [computeEncoder setComputePipelineState:visionPSO]; - // [N, C, H, W] - [computeEncoder setBuffer:inputBuffer offset:grad_.storage_offset() * grad_.element_size() atIndex:0]; - [computeEncoder setBuffer:roisBuffer offset:rois_.storage_offset() * rois_.element_size() atIndex:1]; - [computeEncoder setBuffer:channelMappingBuffer - offset:channel_mapping.storage_offset() * channel_mapping.element_size() - atIndex:2]; - [computeEncoder setBuffer:outputBuffer offset:grad_input.storage_offset() * grad_input.element_size() atIndex:3]; - - [computeEncoder setBytes:&output_size length:sizeof(int64_t) atIndex:4]; - [computeEncoder setBytes:&channels length:sizeof(int64_t) atIndex:5]; - [computeEncoder setBytes:&height length:sizeof(int64_t) atIndex:6]; - [computeEncoder setBytes:&width length:sizeof(int64_t) atIndex:7]; - [computeEncoder setBytes:&pooled_height length:sizeof(int64_t) atIndex:8]; - [computeEncoder setBytes:&pooled_width length:sizeof(int64_t) atIndex:9]; - [computeEncoder setBytes:&channels_out length:sizeof(int64_t) atIndex:10]; - [computeEncoder setBytes:&spatial_scale_f length:sizeof(float) atIndex:11]; - - // A threadGroup is equivalent to a cuda's block. - NSUInteger tgSize = visionPSO.maxTotalThreadsPerThreadgroup; - if (tgSize > threadsPerBlock) { - tgSize = threadsPerBlock; - } - - MTLSize threadGroupSize = MTLSizeMake(tgSize, 1, 1); - [computeEncoder dispatchThreadgroups:threadgroupsPerGrid threadsPerThreadgroup:threadGroupSize]; - - getMPSProfiler().endProfileKernel(visionPSO); - } - }); + + Tensor grad_ = torch::stable::contiguous(grad); + Tensor rois_ = torch::stable::contiguous(rois); + Tensor channel_mapping_ = torch::stable::contiguous(channel_mapping); + Tensor spatial_scale_t = make_scalar_tensor(grad_, spatial_scale); + + int64_t output_size = grad_.numel(); + + const std::string kernel = "ps_roi_pool_backward_" + + std::string(metal_type_string(grad_.scalar_type())); + AOTIMetalKernelFunctionHandle func = nullptr; + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_get_kernel_function( + ps_roi_pool_shader_library(), kernel.c_str(), &func)); + + PsRoiPoolBackwardArgs args{ + grad_.get(), + rois_.get(), + channel_mapping_.get(), + grad_input.get(), + output_size, + channels, + height, + width, + pooled_height, + pooled_width, + channels_out, + spatial_scale_t.get()}; + TORCH_ERROR_CODE_CHECK(aoti_torch_mps_run_command_block( + func, &ps_roi_pool_backward_encode, &args)); return grad_input; } } // namespace -TORCH_LIBRARY_IMPL(torchvision, MPS, m) { - m.impl(TORCH_SELECTIVE_NAME("torchvision::ps_roi_pool"), TORCH_FN(ps_roi_pool_forward_kernel)); - m.impl(TORCH_SELECTIVE_NAME("torchvision::_ps_roi_pool_backward"), TORCH_FN(ps_roi_pool_backward_kernel)); +STABLE_TORCH_LIBRARY_IMPL(torchvision, MPS, m) { + m.impl("ps_roi_pool", TORCH_BOX(&ps_roi_pool_forward_kernel)); + m.impl("_ps_roi_pool_backward", TORCH_BOX(&ps_roi_pool_backward_kernel)); } } // namespace ops diff --git a/torchvision/csrc/ops/mps/ps_roi_pool_metal_shader.h b/torchvision/csrc/ops/mps/ps_roi_pool_metal_shader.h new file mode 100644 index 00000000000..971a9fadd54 --- /dev/null +++ b/torchvision/csrc/ops/mps/ps_roi_pool_metal_shader.h @@ -0,0 +1,339 @@ +#pragma once + +// Metal shader source for the ps_roi_pool MPS kernels (forward + backward). +// +// Carved out of the shared ops/mps/mps_kernels.h so it can be compiled into the +// stable-ABI _C_stable extension: mps_kernels.h opens with +// #include and instantiates +// at::native::mps::MetalShaderLibrary, both unavailable under +// -DTORCH_TARGET_VERSION. This header carries only the Metal source as a plain +// string, handed to aoti_torch_mps_create_shader_library at runtime. + +namespace vision { +namespace ops { + +static const char* ps_roi_pool_metal_shader = R"VISION_METAL( + +#include +#include +using namespace metal; + +#define MPS_1D_KERNEL_LOOP_T(i, n, n_tgs, index_t) \ + for (index_t i = (tgid.x * tptg.x) + tid2.x; i < (n); \ + i += (tptg.x * n_tgs)) + +#define MPS_1D_KERNEL_LOOP(i, n, n_tgs) MPS_1D_KERNEL_LOOP_T(i, n, n_tgs, uint) + +template +inline T ceil_div(T n, T m) { + return (n + m - 1) / m; +} + +inline void atomic_add_float(device float* data_ptr, const float val) +{ + atomic_fetch_add_explicit((device atomic_float*) data_ptr, val, memory_order_relaxed); +} + + +inline void atomic_add_float(device half* data_ptr, const half val) +{ + atomic_fetch_add_explicit((device atomic_float*) data_ptr, static_cast(val), memory_order_relaxed); +} + +template +inline T bilinear_interpolate( + constant T* input, + integer_t height, + integer_t width, + T y, + T x, + uint index /* index for debug only*/) { + // deal with cases that inverse elements are out of feature map boundary + if (y < -1.0 || y > height || x < -1.0 || x > width) { + // empty + return 0; + } + + if (y <= 0) + y = 0; + if (x <= 0) + x = 0; + + integer_t y_low = (integer_t)y; + integer_t x_low = (integer_t)x; + integer_t y_high; + integer_t x_high; + + if (y_low >= height - 1) { + y_high = y_low = height - 1; + y = (T)y_low; + } else { + y_high = y_low + 1; + } + + if (x_low >= width - 1) { + x_high = x_low = width - 1; + x = (T)x_low; + } else { + x_high = x_low + 1; + } + + T ly = y - y_low; + T lx = x - x_low; + T hy = 1. - ly, hx = 1. - lx; + + // do bilinear interpolation + T v1 = input[y_low * width + x_low]; + T v2 = input[y_low * width + x_high]; + T v3 = input[y_high * width + x_low]; + T v4 = input[y_high * width + x_high]; + T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; + + T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + + return val; +} + +template +inline void bilinear_interpolate_gradient( + integer_t height, + integer_t width, + T y, + T x, + thread T& w1, + thread T& w2, + thread T& w3, + thread T& w4, + thread integer_t& x_low, + thread integer_t& x_high, + thread integer_t& y_low, + thread integer_t& y_high, + uint index /* index for debug only*/) { + // deal with cases that inverse elements are out of feature map boundary + if (y < -1.0 || y > height || x < -1.0 || x > width) { + // empty + w1 = w2 = w3 = w4 = 0.; + x_low = x_high = y_low = y_high = -1; + return; + } + + if (y <= 0) + y = 0; + if (x <= 0) + x = 0; + + y_low = (integer_t)y; + x_low = (integer_t)x; + + if (y_low >= height - 1) { + y_high = y_low = height - 1; + y = (T)y_low; + } else { + y_high = y_low + 1; + } + + if (x_low >= width - 1) { + x_high = x_low = width - 1; + x = (T)x_low; + } else { + x_high = x_low + 1; + } + + T ly = y - y_low; + T lx = x - x_low; + T hy = 1. - ly, hx = 1. - lx; + + // reference in forward + // T v1 = input[y_low * width + x_low]; + // T v2 = input[y_low * width + x_high]; + // T v3 = input[y_high * width + x_low]; + // T v4 = input[y_high * width + x_high]; + // T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + + w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; +} + +template +kernel void ps_roi_pool( + constant T * input [[buffer(0)]], + constant T * rois [[buffer(1)]], + device T * output [[buffer(2)]], + device int64_t * channel_mapping [[buffer(3)]], + constant int64_t & output_size [[buffer(4)]], + constant int64_t & channels [[buffer(5)]], + constant int64_t & height [[buffer(6)]], + constant int64_t & width [[buffer(7)]], + constant int64_t & pooled_height [[buffer(8)]], + constant int64_t & pooled_width [[buffer(9)]], + constant int64_t & channels_out [[buffer(10)]], + constant float & spatial_scale [[buffer(11)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint2 tptg [[threads_per_threadgroup]], + uint2 tid2 [[thread_position_in_threadgroup]]){ + MPS_1D_KERNEL_LOOP(index, output_size, 1) { + // (n, c_out, ph, pw) is an element in the pooled output + integer_t pw = index % pooled_width; + integer_t ph = (index / pooled_width) % pooled_height; + integer_t c_out = (index / (pooled_width * pooled_height)) % channels_out; + integer_t n = index / pooled_width / pooled_height / channels_out; + + // (n, c_in, ph, pw) is the associated element in the input + integer_t c_in = (c_out * pooled_height + ph) * pooled_width + pw; + + // [start, end) interval for spatial sampling + constant T* offset_rois = rois + n * 5; + integer_t roi_batch_ind = offset_rois[0]; + integer_t roi_start_w = round(offset_rois[1] * spatial_scale); + integer_t roi_start_h = round(offset_rois[2] * spatial_scale); + integer_t roi_end_w = round(offset_rois[3] * spatial_scale); + integer_t roi_end_h = round(offset_rois[4] * spatial_scale); + + // Force too small ROIs to be 1x1 + integer_t roi_width = max(roi_end_w - roi_start_w, static_cast(1)); + integer_t roi_height = max(roi_end_h - roi_start_h, static_cast(1)); + T bin_size_h = static_cast(roi_height) / static_cast(pooled_height); + T bin_size_w = static_cast(roi_width) / static_cast(pooled_width); + + integer_t hstart = static_cast(floor(static_cast(ph) * bin_size_h)); + integer_t wstart = static_cast(floor(static_cast(pw) * bin_size_w)); + integer_t hend = static_cast(ceil(static_cast(ph + 1) * bin_size_h)); + integer_t wend = static_cast(ceil(static_cast(pw + 1) * bin_size_w)); + + // Add roi offsets and clip to input boundaries + hstart = min(max(hstart + roi_start_h, static_cast(0)), static_cast(height - 1)); + hend = min(max(hend + roi_start_h, static_cast(0)), static_cast(height - 1)); + wstart = min(max(wstart + roi_start_w, static_cast(0)), static_cast(width - 1)); + wend = min(max(wend + roi_start_w, static_cast(0)), static_cast(width - 1)); + bool is_empty = (hend <= hstart) || (wend <= wstart); + + constant T* offset_input = + input + (roi_batch_ind * channels + c_in) * height * width; + T out_sum = 0; + for (integer_t h = hstart; h < hend; ++h) { + for (integer_t w = wstart; w < wend; ++w) { + integer_t input_index = h * width + w; + out_sum += offset_input[input_index]; + } + } + + T bin_area = (hend - hstart) * (wend - wstart); + output[index] = is_empty ? static_cast(0) : out_sum / bin_area; + channel_mapping[index] = c_in; + } +} + +#define REGISTER_PS_ROI_POOL_OP(DTYPE, INT_DTYPE) \ +template \ +[[host_name("ps_roi_pool_" #DTYPE)]] \ +kernel void ps_roi_pool( \ + constant DTYPE * input [[buffer(0)]], \ + constant DTYPE * rois [[buffer(1)]], \ + device DTYPE * output [[buffer(2)]], \ + device int64_t * channel_mapping [[buffer(3)]], \ + constant int64_t & output_size [[buffer(4)]], \ + constant int64_t & channels [[buffer(5)]], \ + constant int64_t & height [[buffer(6)]], \ + constant int64_t & width [[buffer(7)]], \ + constant int64_t & pooled_height [[buffer(8)]], \ + constant int64_t & pooled_width [[buffer(9)]], \ + constant int64_t & channels_out [[buffer(10)]], \ + constant float & spatial_scale [[buffer(11)]], \ + uint2 tgid [[threadgroup_position_in_grid]], \ + uint2 tptg [[threads_per_threadgroup]], \ + uint2 tid2 [[thread_position_in_threadgroup]]); + +template +kernel void ps_roi_pool_backward( + constant T * grad_output [[buffer(0)]], + constant T * rois [[buffer(1)]], + constant int64_t * channel_mapping [[buffer(2)]], + device T * grad_input [[buffer(3)]], + constant int64_t & output_size [[buffer(4)]], + constant int64_t & channels [[buffer(5)]], + constant int64_t & height [[buffer(6)]], + constant int64_t & width [[buffer(7)]], + constant int64_t & pooled_height [[buffer(8)]], + constant int64_t & pooled_width [[buffer(9)]], + constant int64_t & channels_out [[buffer(10)]], + constant float & spatial_scale [[buffer(11)]], + uint2 tgid [[threadgroup_position_in_grid]], + uint2 tptg [[threads_per_threadgroup]], + uint2 tid2 [[thread_position_in_threadgroup]]){ + + MPS_1D_KERNEL_LOOP(index, output_size, 1) { + // (n, *, ph, pw) is an element in the pooled output + integer_t pw = index % pooled_width; + integer_t ph = (index / pooled_width) % pooled_height; + integer_t n = index / pooled_width / pooled_height / channels_out; + + constant T* offset_rois = rois + n * 5; + integer_t roi_batch_ind = offset_rois[0]; + integer_t roi_start_w = round(offset_rois[1] * spatial_scale); + integer_t roi_start_h = round(offset_rois[2] * spatial_scale); + integer_t roi_end_w = round(offset_rois[3] * spatial_scale); + integer_t roi_end_h = round(offset_rois[4] * spatial_scale); + + // Force too small ROIs to be 1x1 + integer_t roi_width = max(roi_end_w - roi_start_w, static_cast(1)); + integer_t roi_height = max(roi_end_h - roi_start_h, static_cast(1)); + T bin_size_h = static_cast(roi_height) / static_cast(pooled_height); + T bin_size_w = static_cast(roi_width) / static_cast(pooled_width); + + integer_t hstart = static_cast(floor(static_cast(ph) * bin_size_h)); + integer_t wstart = static_cast(floor(static_cast(pw) * bin_size_w)); + integer_t hend = static_cast(ceil(static_cast(ph + 1) * bin_size_h)); + integer_t wend = static_cast(ceil(static_cast(pw + 1) * bin_size_w)); + + // Add roi offsets and clip to input boundaries + hstart = min(max(hstart + roi_start_h, static_cast(0)), static_cast(height)); + hend = min(max(hend + roi_start_h, static_cast(0)), static_cast(height)); + wstart = min(max(wstart + roi_start_w, static_cast(0)), static_cast(width)); + wend = min(max(wend + roi_start_w, static_cast(0)), static_cast(width)); + bool is_empty = (hend <= hstart) || (wend <= wstart); + + integer_t c_in = channel_mapping[index]; + T bin_area = (hend - hstart) * (wend - wstart); + T diff_val = is_empty ? static_cast(0) : grad_output[index] / bin_area; + + const integer_t offset = (roi_batch_ind * channels + c_in) * height * width; + + for (integer_t h = hstart; h < hend; ++h) { + for (integer_t w = wstart; w < wend; ++w) { + integer_t grad_input_index = h * width + w; + atomic_add_float(grad_input + offset + grad_input_index, diff_val); + } + } + + } // MPS_1D_KERNEL_LOOP +} + +#define REGISTER_PS_ROI_POOL_BACKWARD_OP(DTYPE, INT_DTYPE) \ +template \ +[[host_name("ps_roi_pool_backward_" #DTYPE)]] \ +kernel void ps_roi_pool_backward( \ + constant DTYPE * grad_output [[buffer(0)]], \ + constant DTYPE * rois [[buffer(1)]], \ + constant int64_t * channel_mapping [[buffer(2)]], \ + device DTYPE * grad_input [[buffer(3)]], \ + constant int64_t & output_size [[buffer(4)]], \ + constant int64_t & channels [[buffer(5)]], \ + constant int64_t & height [[buffer(6)]], \ + constant int64_t & width [[buffer(7)]], \ + constant int64_t & pooled_height [[buffer(8)]], \ + constant int64_t & pooled_width [[buffer(9)]], \ + constant int64_t & channels_out [[buffer(10)]], \ + constant float & spatial_scale [[buffer(11)]], \ + uint2 tgid [[threadgroup_position_in_grid]], \ + uint2 tptg [[threads_per_threadgroup]], \ + uint2 tid2 [[thread_position_in_threadgroup]]); + + +REGISTER_PS_ROI_POOL_OP(float, int64_t); +REGISTER_PS_ROI_POOL_OP(half, int64_t); +REGISTER_PS_ROI_POOL_BACKWARD_OP(float, int64_t); +REGISTER_PS_ROI_POOL_BACKWARD_OP(half, int64_t); + +)VISION_METAL"; + +} // namespace ops +} // namespace vision