From d312aeac2cffc07fae6ae029e86732b056a3d2b6 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Mon, 12 Oct 2026 12:45:47 -0700 Subject: [PATCH] [ET-VK][convolution] Nest convolution sources by route Pull Request resolved: https://github.com/pytorch/executorch/pull/23492 ## Ulterior Motive Organize Vulkan operators by implementation family so operators with multiple routes can use clear, scalable directory layouts. ## Rationale **What**: Move convolution C++ implementations, GLSL templates, YAML variant definitions, and convolution-only shader headers under `convolution/` in both `impl/` and `glsl/`, split into `conv1d/` and `conv2d/`. Split the monolithic `Convolution.cpp` and `Q8taConv2d.{h,cpp}` so each file holds either routing logic or a single implementation route. **Why**: Convolution is implemented through several routes (direct, depthwise, pointwise, im2col/GEMM, transposed), and the q8ta quantized family repeats most of them. In the flat directories their sources were interleaved by name, and single files mixed routing with several route implementations. This prepares for adding further routes such as Winograd. ## Details - Layout (`impl/` and `glsl/` mirror each other) - `convolution/`: `Convolution.{h,cpp}` registers `aten.convolution` and dispatches to conv1d or conv2d by input rank; it also holds `prepack_biases` and `check_conv_args`, used by both. - `convolution/conv1d/`: `Conv1d.{h,cpp}` (`conv1d_impl` routing), `Conv1dDirect.{h,cpp}`, `Conv1dDW.cpp`, `Conv1dPW.cpp`, and the conv1d shaders. - `convolution/conv2d/`: `Conv2d.{h,cpp}` (`conv2d_impl` routing, `Conv2dMethod`, `resize_conv2d_node`), `Conv2dDirect.{h,cpp}` (direct and transposed shader dispatch), `Conv2dDW.cpp`, `Conv2dPW.cpp`, `ConvolutionUtils.h`, `conv2d_common.glslh`, and the direct/depthwise/pointwise shaders. - `convolution/conv2d/{im2col,transposed}/`: im2col/GEMM sources and shaders; transposed shaders. - `convolution/conv2d/q8ta/`: `Q8taConv2d.{h,cpp}` (`q8ta_conv2d` routing and the im2col heuristic, formerly `Q8taConv2dRoute`), `Q8taConv2dCommon.{h,cpp}` (`ActivationType`, packed-dim checks), `Q8taConv2dDirect.{h,cpp}` (general route), `Q8taConv2dDW.{h,cpp}`, `Q8taConv2dPW.{h,cpp}`, plus `im2col/` and `transposed/`. Each route has its own header; consumers include only what they use. - `convolution/conv2d/quantized_legacy/`: `QuantizedConvolution.{h,cpp}` and its shaders and headers. - Weight packing shaders sit in `pack_weights/` subdirectories of the route that dispatches them (`conv2d/`, `conv2d/im2col/`, `conv2d/q8ta/`, `conv2d/transposed/`). - Naming - Functions that add a `DynamicDispatchNode` are named `add_*_node()`; functions that call `add_*_node()` to implement an operator are named `*_impl()`. Registered operator names are unchanged. - `add_conv2d_node` is folded into `conv2d_impl`. Its test-only `force_direct` flag is removed; the conv2d test's `direct` selector calls `conv2d_direct_impl` directly, which is equivalent for its sliding-window cases. - Includes - Route shaders include `conv2d_common.glslh` unchanged through ancestor include resolution. - Root shaders `quantize_and_pack_4w4c` and `unpack_4w4c_and_dequantize` include `convolution/conv2d/conv2d_common.glslh`. - Headers shared with other operator families (`linear_*`, `conv2d_fp_input_tile_load.glslh`, indexing headers) stay at the shader root. - Moved files are recorded as renames; files split from or merged into another file are recorded as copies of their source, so reviews show moved code as unchanged. - Moving code into new files exposes it to CLANGTIDY, so its findings there are fixed: `utils::safe_downcast` for narrowing, `(void)` for unused parameters, an anonymous namespace for the file-local `Conv2dParams`, and removal of an unused `ShaderNameUtils.h` include. Generated shader names and runtime dispatch behavior are unchanged. ghstack-source-id: 446951717 @exported-using-ghexport Differential Revision: [D123450342](https://our.internmc.facebook.com/intern/diff/D123450342/) --- backends/vulkan/op_registry.py | 4 +- .../glsl/{ => convolution/conv1d}/conv1d.glsl | 0 .../glsl/{ => convolution/conv1d}/conv1d.yaml | 0 .../{ => convolution/conv1d}/conv1d_dw.glsl | 0 .../{ => convolution/conv1d}/conv1d_dw.yaml | 0 .../{ => convolution/conv1d}/conv1d_pw.glsl | 0 .../{ => convolution/conv1d}/conv1d_pw.yaml | 0 .../glsl/{ => convolution/conv2d}/conv2d.glsl | 0 .../glsl/{ => convolution/conv2d}/conv2d.yaml | 0 .../conv2d}/conv2d_common.glslh | 0 .../{ => convolution/conv2d}/conv2d_dw.glsl | 0 .../{ => convolution/conv2d}/conv2d_dw.yaml | 0 .../conv2d}/conv2d_dw_output_tile.glsl | 0 .../conv2d}/conv2d_dw_output_tile.yaml | 0 .../conv2d}/conv2d_dw_sned_output_tile.glsl | 0 .../conv2d}/conv2d_dw_sned_output_tile.yaml | 0 .../{ => convolution/conv2d}/conv2d_pw.glsl | 0 .../{ => convolution/conv2d}/conv2d_pw.yaml | 0 .../conv2d}/conv2d_pw_s1p0.glsl | 0 .../conv2d}/conv2d_pw_s1p0.yaml | 0 .../conv2d}/conv2d_pw_tiled.glsl | 0 .../conv2d}/conv2d_pw_tiled.yaml | 0 .../conv2d/im2col}/conv2d_gemm.glsl | 0 .../conv2d/im2col}/conv2d_gemm.yaml | 0 .../conv2d/im2col}/conv2d_im2col.glsl | 0 .../conv2d/im2col}/conv2d_im2col.yaml | 0 .../pack_conv2d_gemm_weight.glsl | 0 .../pack_conv2d_gemm_weight.yaml | 0 .../conv2d_dw_prepack_weights.glsl | 0 .../conv2d_dw_prepack_weights.yaml | 0 .../pack_weights}/conv2d_prepack_weights.glsl | 0 .../pack_weights}/conv2d_prepack_weights.yaml | 0 .../conv2d/q8ta}/conv2d_dw_q8_utils.glslh | 0 .../q8ta}/conv2d_dw_q8ta_q8csw_q8to.glsl | 0 .../q8ta}/conv2d_dw_q8ta_q8csw_q8to.yaml | 0 .../conv2d/q8ta/im2col}/q8ta_im2col.glsl | 0 .../conv2d/q8ta/im2col}/q8ta_im2col.yaml | 0 .../pack_q8_conv2d_dw_weights.glsl | 0 .../pack_q8_conv2d_dw_weights.yaml | 0 .../pack_weights}/pack_q8_conv2d_weights.glsl | 0 .../pack_weights}/pack_q8_conv2d_weights.yaml | 0 .../conv2d/q8ta}/q8ta_conv2d.glsl | 0 .../conv2d/q8ta}/q8ta_conv2d.yaml | 0 .../conv2d/q8ta}/q8ta_conv2d_dw.glsl | 0 .../conv2d/q8ta}/q8ta_conv2d_dw.yaml | 0 .../conv2d/q8ta}/q8ta_conv2d_pw.glsl | 0 .../conv2d/q8ta}/q8ta_conv2d_pw.yaml | 0 .../transposed}/q8ta_conv2d_transposed.glsl | 0 .../transposed}/q8ta_conv2d_transposed.yaml | 0 .../conv2d/quantized_legacy}/col2im.glsl | 0 .../conv2d/quantized_legacy}/col2im.yaml | 0 .../conv2d_fp_im2col_block.glslh | 0 .../conv2d_fp_im2col_block_load.glslh | 0 .../conv2d_fp_im2col_block_store.glslh | 0 .../conv2d_int8_input_block_load.glslh | 0 .../conv2d_int8_input_tile_load.glslh | 0 .../conv2d_int8_output_tile_store.glslh | 0 .../conv2d_pw_q8ta_q8csw_q8to_tiled.glsl | 0 .../conv2d_pw_q8ta_q8csw_q8to_tiled.yaml | 0 .../quantized_legacy}/conv2d_q8_utils.glslh | 0 .../conv2d_q8csw_linear_tiled.glsl | 0 .../conv2d_q8csw_linear_tiled.yaml | 0 .../conv2d_q8ta_q8csw_linear_tiled.glsl | 0 .../conv2d_q8ta_q8csw_linear_tiled.yaml | 0 .../conv2d_q8ta_q8csw_q8to.glsl | 0 .../conv2d_q8ta_q8csw_q8to.yaml | 0 .../conv2d_q8ta_q8csw_q8to_linear_tiled.glsl | 0 .../conv2d_q8ta_q8csw_q8to_linear_tiled.yaml | 0 .../conv2d/quantized_legacy}/im2col.glsl | 0 .../conv2d/quantized_legacy}/im2col.yaml | 0 .../quantized_legacy}/im2col_packed_int8.glsl | 0 .../quantized_legacy}/im2col_packed_int8.yaml | 0 .../im2col_packed_int8_utils.glslh | 0 .../quantize_and_pack_im2col.glsl | 0 .../quantize_and_pack_im2col.yaml | 0 .../conv2d/transposed}/conv_transpose2d.glsl | 0 .../conv2d/transposed}/conv_transpose2d.yaml | 0 .../conv_transpose2d_prepack_weights.glsl | 0 .../conv_transpose2d_prepack_weights.yaml | 0 .../ops/glsl/quantize_and_pack_4w4c.glsl | 2 +- .../ops/glsl/unpack_4w4c_and_dequantize.glsl | 2 +- .../runtime/graph/ops/impl/Convolution.cpp | 871 ------------------ .../runtime/graph/ops/impl/Q8taConv2d.h | 216 ----- .../runtime/graph/ops/impl/Q8taLinear.cpp | 2 +- .../runtime/graph/ops/impl/Q8taLinearGemv.cpp | 2 +- .../ops/impl/convolution/Convolution.cpp | 84 ++ .../graph/ops/impl/convolution/Convolution.h | 28 + .../ops/impl/convolution/conv1d/Conv1d.cpp | 100 ++ .../ops/impl/convolution/conv1d/Conv1d.h | 17 + .../{ => convolution/conv1d}/Conv1dDW.cpp | 4 +- .../impl/convolution/conv1d/Conv1dDirect.cpp | 179 ++++ .../impl/convolution/conv1d/Conv1dDirect.h | 29 + .../{ => convolution/conv1d}/Conv1dPW.cpp | 4 +- .../ops/impl/convolution/conv2d/Conv2d.cpp | 234 +++++ .../conv2d/Conv2d.h} | 40 +- .../{ => convolution/conv2d}/Conv2dDW.cpp | 3 +- .../impl/convolution/conv2d/Conv2dDirect.cpp | 402 ++++++++ .../impl/convolution/conv2d/Conv2dDirect.h | 33 + .../{ => convolution/conv2d}/Conv2dPW.cpp | 3 +- .../conv2d}/ConvolutionUtils.h | 0 .../conv2d/im2col}/Conv2dGemm.cpp | 6 +- .../conv2d/im2col}/Conv2dGemm.h | 0 .../conv2d/im2col}/Conv2dIm2Col.cpp | 2 +- .../conv2d/im2col}/Conv2dIm2Col.h | 0 .../conv2d/q8ta/Q8taConv2d.cpp} | 59 +- .../conv2d/q8ta/Q8taConv2d.h} | 4 + .../conv2d/q8ta/Q8taConv2dCommon.cpp | 37 + .../conv2d/q8ta/Q8taConv2dCommon.h | 28 + .../conv2d/q8ta}/Q8taConv2dDW.cpp | 12 +- .../convolution/conv2d/q8ta/Q8taConv2dDW.h | 61 ++ .../conv2d/q8ta/Q8taConv2dDirect.cpp} | 85 +- .../conv2d/q8ta/Q8taConv2dDirect.h | 50 + .../conv2d/q8ta}/Q8taConv2dPW.cpp | 5 +- .../convolution/conv2d/q8ta/Q8taConv2dPW.h | 53 ++ .../conv2d/q8ta/im2col}/Q8taConv2dIm2Col.cpp | 11 +- .../conv2d/q8ta/im2col/Q8taConv2dIm2Col.h | 71 ++ .../q8ta/transposed}/Q8taConv2dTransposed.cpp | 12 +- .../q8ta/transposed/Q8taConv2dTransposed.h | 19 + .../QuantizedConvolution.cpp | 13 +- .../quantized_legacy}/QuantizedConvolution.h | 2 +- .../test/custom_ops/impl/TestConv2d.cpp | 20 +- .../test/custom_ops/impl/TestConv2dDw.cpp | 3 +- .../test/custom_ops/impl/TestQ8taConv2d.cpp | 3 +- .../custom_ops/q8ta_conv2d_route_test.cpp | 2 +- .../q8ta_conv2d_stream_plan_test.cpp | 2 +- .../vulkan/test/custom_ops/test_conv2d.cpp | 2 +- .../test/custom_ops/test_q8ta_conv2d.cpp | 3 +- .../test/custom_ops/test_q8ta_conv2d_pw.cpp | 2 +- .../op_tests/conv2d_gemm_dynamic_test.cpp | 2 +- 129 files changed, 1574 insertions(+), 1254 deletions(-) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv1d}/conv1d.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv1d}/conv1d.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv1d}/conv1d_dw.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv1d}/conv1d_dw.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv1d}/conv1d_pw.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv1d}/conv1d_pw.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_common.glslh (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_dw.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_dw.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_dw_output_tile.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_dw_output_tile.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_dw_sned_output_tile.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_dw_sned_output_tile.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_pw.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_pw.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_pw_s1p0.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_pw_s1p0.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_pw_tiled.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d}/conv2d_pw_tiled.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/im2col}/conv2d_gemm.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/im2col}/conv2d_gemm.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/im2col}/conv2d_im2col.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/im2col}/conv2d_im2col.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/im2col/pack_weights}/pack_conv2d_gemm_weight.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/im2col/pack_weights}/pack_conv2d_gemm_weight.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/pack_weights}/conv2d_dw_prepack_weights.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/pack_weights}/conv2d_dw_prepack_weights.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/pack_weights}/conv2d_prepack_weights.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/pack_weights}/conv2d_prepack_weights.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta}/conv2d_dw_q8_utils.glslh (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta}/conv2d_dw_q8ta_q8csw_q8to.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta}/conv2d_dw_q8ta_q8csw_q8to.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta/im2col}/q8ta_im2col.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta/im2col}/q8ta_im2col.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta/pack_weights}/pack_q8_conv2d_dw_weights.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta/pack_weights}/pack_q8_conv2d_dw_weights.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta/pack_weights}/pack_q8_conv2d_weights.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta/pack_weights}/pack_q8_conv2d_weights.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta}/q8ta_conv2d.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta}/q8ta_conv2d.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta}/q8ta_conv2d_dw.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta}/q8ta_conv2d_dw.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta}/q8ta_conv2d_pw.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta}/q8ta_conv2d_pw.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta/transposed}/q8ta_conv2d_transposed.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/q8ta/transposed}/q8ta_conv2d_transposed.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/col2im.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/col2im.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_fp_im2col_block.glslh (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_fp_im2col_block_load.glslh (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_fp_im2col_block_store.glslh (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_int8_input_block_load.glslh (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_int8_input_tile_load.glslh (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_int8_output_tile_store.glslh (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_pw_q8ta_q8csw_q8to_tiled.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_pw_q8ta_q8csw_q8to_tiled.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_q8_utils.glslh (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_q8csw_linear_tiled.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_q8csw_linear_tiled.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_q8ta_q8csw_linear_tiled.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_q8ta_q8csw_linear_tiled.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_q8ta_q8csw_q8to.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_q8ta_q8csw_q8to.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_q8ta_q8csw_q8to_linear_tiled.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/conv2d_q8ta_q8csw_q8to_linear_tiled.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/im2col.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/im2col.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/im2col_packed_int8.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/im2col_packed_int8.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/im2col_packed_int8_utils.glslh (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/quantize_and_pack_im2col.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/quantized_legacy}/quantize_and_pack_im2col.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/transposed}/conv_transpose2d.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/transposed}/conv_transpose2d.yaml (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/transposed/pack_weights}/conv_transpose2d_prepack_weights.glsl (100%) rename backends/vulkan/runtime/graph/ops/glsl/{ => convolution/conv2d/transposed/pack_weights}/conv_transpose2d_prepack_weights.yaml (100%) delete mode 100644 backends/vulkan/runtime/graph/ops/impl/Convolution.cpp delete mode 100644 backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/Convolution.cpp create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/Convolution.h create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1d.cpp create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1d.h rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv1d}/Conv1dDW.cpp (97%) create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDirect.cpp create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDirect.h rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv1d}/Conv1dPW.cpp (98%) create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2d.cpp rename backends/vulkan/runtime/graph/ops/impl/{Convolution.h => convolution/conv2d/Conv2d.h} (50%) rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d}/Conv2dDW.cpp (98%) create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDirect.cpp create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDirect.h rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d}/Conv2dPW.cpp (97%) rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d}/ConvolutionUtils.h (100%) rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d/im2col}/Conv2dGemm.cpp (98%) rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d/im2col}/Conv2dGemm.h (100%) rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d/im2col}/Conv2dIm2Col.cpp (98%) rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d/im2col}/Conv2dIm2Col.h (100%) rename backends/vulkan/runtime/graph/ops/impl/{Q8taConv2dRoute.cpp => convolution/conv2d/q8ta/Q8taConv2d.cpp} (70%) rename backends/vulkan/runtime/graph/ops/impl/{Q8taConv2dRoute.h => convolution/conv2d/q8ta/Q8taConv2d.h} (84%) create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dCommon.cpp create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dCommon.h rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d/q8ta}/Q8taConv2dDW.cpp (97%) create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDW.h rename backends/vulkan/runtime/graph/ops/impl/{Q8taConv2d.cpp => convolution/conv2d/q8ta/Q8taConv2dDirect.cpp} (84%) create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDirect.h rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d/q8ta}/Q8taConv2dPW.cpp (98%) create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dPW.h rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d/q8ta/im2col}/Q8taConv2dIm2Col.cpp (96%) create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.h rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d/q8ta/transposed}/Q8taConv2dTransposed.cpp (95%) create mode 100644 backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/transposed/Q8taConv2dTransposed.h rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d/quantized_legacy}/QuantizedConvolution.cpp (98%) rename backends/vulkan/runtime/graph/ops/impl/{ => convolution/conv2d/quantized_legacy}/QuantizedConvolution.h (78%) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 461aeeaafcc..22c9da4aab1 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -896,7 +896,7 @@ def register_max_pool2d_with_indices(): # ============================================================================= -# Convolution.cpp +# convolution/Convolution.cpp # ============================================================================= @@ -989,7 +989,7 @@ def pick_conv_storage( # ============================================================================= -# Q8taConv2d*.cpp +# convolution/conv2d/q8ta/**/Q8taConv2d*.cpp # ============================================================================= diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv1d.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv1d.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv1d.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv1d.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv1d_dw.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_dw.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv1d_dw.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_dw.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv1d_dw.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_dw.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv1d_dw.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_dw.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv1d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_pw.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv1d_pw.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_pw.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv1d_pw.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_pw.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv1d_pw.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv1d/conv1d_pw.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_common.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_common.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_common.glslh rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_common.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw_output_tile.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw_output_tile.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw_output_tile.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw_output_tile.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_sned_output_tile.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw_sned_output_tile.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_sned_output_tile.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw_sned_output_tile.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_sned_output_tile.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw_sned_output_tile.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_sned_output_tile.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_dw_sned_output_tile.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_s1p0.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_s1p0.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_s1p0.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_s1p0.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_tiled.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_tiled.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_tiled.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_tiled.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/conv2d_pw_tiled.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_gemm.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_gemm.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_gemm.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_gemm.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_gemm.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_gemm.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_gemm.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_gemm.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_im2col.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_im2col.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_im2col.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_im2col.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_im2col.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_im2col.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/conv2d_im2col.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_conv2d_gemm_weight.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/pack_weights/pack_conv2d_gemm_weight.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_conv2d_gemm_weight.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/pack_weights/pack_conv2d_gemm_weight.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_conv2d_gemm_weight.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/pack_weights/pack_conv2d_gemm_weight.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_conv2d_gemm_weight.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/im2col/pack_weights/pack_conv2d_gemm_weight.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_prepack_weights.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_prepack_weights.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_prepack_weights.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_prepack_weights.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_prepack_weights.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_prepack_weights.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_q8_utils.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/conv2d_dw_q8_utils.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_q8_utils.glslh rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/conv2d_dw_q8_utils.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_q8ta_q8csw_q8to.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/conv2d_dw_q8ta_q8csw_q8to.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_q8ta_q8csw_q8to.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/conv2d_dw_q8ta_q8csw_q8to.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_q8ta_q8csw_q8to.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/conv2d_dw_q8ta_q8csw_q8to.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_q8ta_q8csw_q8to.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/conv2d_dw_q8ta_q8csw_q8to.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/im2col/q8ta_im2col.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/im2col/q8ta_im2col.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/im2col/q8ta_im2col.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/im2col/q8ta_im2col.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_dw_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/pack_weights/pack_q8_conv2d_dw_weights.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_dw_weights.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/pack_weights/pack_q8_conv2d_dw_weights.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_dw_weights.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/pack_weights/pack_q8_conv2d_dw_weights.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_dw_weights.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/pack_weights/pack_q8_conv2d_dw_weights.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/pack_weights/pack_q8_conv2d_weights.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/pack_weights/pack_q8_conv2d_weights.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/pack_weights/pack_q8_conv2d_weights.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/pack_weights/pack_q8_conv2d_weights.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d_dw.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d_dw.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d_dw.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_dw.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d_dw.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d_pw.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d_pw.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d_pw.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/q8ta_conv2d_pw.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_transposed.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/transposed/q8ta_conv2d_transposed.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_transposed.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/transposed/q8ta_conv2d_transposed.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_transposed.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/transposed/q8ta_conv2d_transposed.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_transposed.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/q8ta/transposed/q8ta_conv2d_transposed.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/col2im.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/col2im.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/col2im.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/col2im.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/col2im.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/col2im.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/col2im.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/col2im.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block.glslh rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_store.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block_store.glslh rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_fp_im2col_block_store.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_int8_input_block_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_input_block_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_int8_input_block_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_input_block_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_int8_input_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_input_tile_load.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_int8_input_tile_load.glslh rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_input_tile_load.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_int8_output_tile_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_output_tile_store.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_int8_output_tile_store.glslh rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_int8_output_tile_store.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_q8ta_q8csw_q8to_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_pw_q8ta_q8csw_q8to_tiled.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_q8ta_q8csw_q8to_tiled.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_pw_q8ta_q8csw_q8to_tiled.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_q8ta_q8csw_q8to_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_pw_q8ta_q8csw_q8to_tiled.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_q8ta_q8csw_q8to_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_pw_q8ta_q8csw_q8to_tiled.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8_utils.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8_utils.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_q8_utils.glslh rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8_utils.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8csw_linear_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8csw_linear_tiled.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_q8csw_linear_tiled.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8csw_linear_tiled.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8csw_linear_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8csw_linear_tiled.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_q8csw_linear_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8csw_linear_tiled.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_linear_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_linear_tiled.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_linear_tiled.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_linear_tiled.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_linear_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_linear_tiled.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_linear_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_linear_tiled.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_q8to.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_q8to.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_q8to.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_q8to.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_q8to_linear_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to_linear_tiled.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_q8to_linear_tiled.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to_linear_tiled.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_q8to_linear_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to_linear_tiled.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_q8to_linear_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/conv2d_q8ta_q8csw_q8to_linear_tiled.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/im2col.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/im2col.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/im2col.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/im2col.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/im2col.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/im2col.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/im2col.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/im2col_packed_int8.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/im2col_packed_int8.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/im2col_packed_int8.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/im2col_packed_int8.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/im2col_packed_int8.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/im2col_packed_int8.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/im2col_packed_int8.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/im2col_packed_int8.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/im2col_packed_int8_utils.glslh b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/im2col_packed_int8_utils.glslh similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/im2col_packed_int8_utils.glslh rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/im2col_packed_int8_utils.glslh diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/quantize_and_pack_im2col.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_im2col.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/quantize_and_pack_im2col.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_im2col.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/quantize_and_pack_im2col.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_im2col.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/quantized_legacy/quantize_and_pack_im2col.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/conv_transpose2d.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/conv_transpose2d.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/conv_transpose2d.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/conv_transpose2d.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.glsl similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d_prepack_weights.glsl rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.glsl diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d_prepack_weights.yaml b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.yaml similarity index 100% rename from backends/vulkan/runtime/graph/ops/glsl/conv_transpose2d_prepack_weights.yaml rename to backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w4c.glsl b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w4c.glsl index 205929d7456..6f58021be7f 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w4c.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w4c.glsl @@ -28,7 +28,7 @@ $if INPUT_STORAGE == "buffer": layout(std430) buffer; -#include "conv2d_common.glslh" +#include "convolution/conv2d/conv2d_common.glslh" ${layout_declare_tensor(B, "w", "t_packed_int8_input", "int", OUTPUT_STORAGE, is_scalar_array=False)} ${layout_declare_tensor(B, "r", "t_fp_input", DTYPE, INPUT_STORAGE)} diff --git a/backends/vulkan/runtime/graph/ops/glsl/unpack_4w4c_and_dequantize.glsl b/backends/vulkan/runtime/graph/ops/glsl/unpack_4w4c_and_dequantize.glsl index 3baa21f54bf..fd7fa0aa7f1 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/unpack_4w4c_and_dequantize.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/unpack_4w4c_and_dequantize.glsl @@ -28,7 +28,7 @@ $if INPUT_STORAGE == "buffer": layout(std430) buffer; -#include "conv2d_common.glslh" +#include "convolution/conv2d/conv2d_common.glslh" ${layout_declare_tensor(B, "w", "t_fp_output", DTYPE, OUTPUT_STORAGE)} ${layout_declare_tensor(B, "r", "t_packed_int8_output", "int", INPUT_STORAGE, is_scalar_array=False)} diff --git a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp deleted file mode 100644 index 272e886b7b1..00000000000 --- a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp +++ /dev/null @@ -1,871 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#include - -#include - -#include -#include -#include - -#include - -#include -#include - -#include - -namespace vkcompute { - -enum class Conv2dMethod : uint8_t { - Depthwise, - Pointwise, - SlidingWindow, - Transposed, -}; - -void resize_conv2d_node( - ComputeGraph* graph, - const std::vector& args, - const std::vector& extra_args) { - const ValueRef out = args.at(0).refs.at(0); - const ValueRef self = args.at(1).refs.at(0); - - size_t ndim = graph->dim_of(self); - std::vector new_out_sizes(ndim); - const bool transposed = graph->get_bool(extra_args.at(4)); - - std::vector self_sizes = graph->sizes_of(self); - // Batch, Channel - if (ndim == 4) { - new_out_sizes.at(ndim - 4) = self_sizes.at(ndim - 4); - } - - TensorRefPtr weight_ref = graph->get_tref(extra_args.at(0)); - const auto& weight_sizes = weight_ref->sizes; - new_out_sizes.at(ndim - 3) = - transposed ? weight_sizes.at(ndim - 3) : weight_sizes.at(ndim - 4); - - // Height, Width - const auto& new_out_sizes_hw = calc_out_sizes_hw( - *graph, - self_sizes, - extra_args.at(0), - /*kernel_size_only = */ false, - {extra_args.at(1), extra_args.at(2), extra_args.at(3), extra_args.at(5)}, - transposed); - new_out_sizes.at(ndim - 2) = new_out_sizes_hw.at(0); - new_out_sizes.at(ndim - 1) = new_out_sizes_hw.at(1); - - graph->virtual_resize(out, new_out_sizes); -} - -void resize_conv1d_node( - ComputeGraph* graph, - const std::vector& args, - const std::vector& extra_args) { - const ValueRef out = args.at(0).refs.at(0); - const ValueRef self = args.at(1).refs.at(0); - TensorRefPtr weight_ref = graph->get_tref(extra_args.at(0)); - - const int64_t stride_size = graph->get_int_list(extra_args.at(1))->at(0); - const int64_t padding_size = graph->get_int_list(extra_args.at(2))->at(0); - const int64_t dilation_size = graph->get_int_list(extra_args.at(3))->at(0); - - const std::vector& weight_sizes = weight_ref->sizes; - - const std::vector in_sizes = graph->sizes_of(self); - const size_t ndim = in_sizes.size(); - std::vector new_out_sizes(ndim); - - const int64_t kernel_size = weight_sizes.at(2); - const int64_t in_length = in_sizes.at(2); - - new_out_sizes.at(0) = in_sizes.at(0); - new_out_sizes.at(1) = weight_sizes.at(0); - new_out_sizes.at(2) = calc_out_size( - in_length, kernel_size, stride_size, padding_size, dilation_size, false); - - graph->virtual_resize(out, new_out_sizes); -} - -ValueRef prepack_biases( - ComputeGraph& graph, - const ValueRef vref, - const ValueRef weight, - const bool transposed, - const utils::StorageType storage_type, - const utils::GPUMemoryLayout memory_layout) { - auto sizes = graph.sizes_of(weight); - const int64_t out_channels = transposed ? sizes.at(1) : sizes.at(0); - - ValueRef v = graph.add_tensor( - {out_channels}, graph.dtype_of(weight), storage_type, memory_layout); - - vkapi::ShaderInfo shader = - get_nchw_to_tensor_shader(graph, v, graph.get_staging_dtype_for(weight)); - - vkapi::ParamsBindList param_buffers = {}; - if (graph.is_buffer_storage(v)) { - param_buffers.append(graph.buffer_meta_ubo(v)); - } else { - param_buffers.append(graph.texture_meta_ubo(v)); - } - - std::vector pcs; - if (graph.is_buffer_storage(v)) { - pcs = {graph.sizes_pc_of(v), graph.strides_pc_of(v), graph.numel_pc_of(v)}; - } - - graph.prepack_nodes().emplace_back(new PrepackNode( - graph, - shader, - graph.create_gwg(v), - graph.create_lwg(v), - vref, - v, - param_buffers, - // Specialization constants - {graph.hashed_layout_of(v)}, - pcs)); - - return v; -} - -vkapi::ShaderInfo get_conv2d_shader( - ComputeGraph& graph, - const ValueRef out, - const bool prepack_weights, - const Conv2dMethod method, - const ValueRef weight, - const bool clamp_out = false, - const bool stride_equals_dilation = false, - const bool stride_1_padding_0 = false) { - std::string kernel_name; - kernel_name.reserve(kShaderNameReserve); - switch (method) { - case Conv2dMethod::Depthwise: - kernel_name = "conv2d_dw"; - break; - case Conv2dMethod::Pointwise: - if (prepack_weights) { - kernel_name = "conv2d"; - } else { - kernel_name = stride_1_padding_0 ? "conv2d_pw_s1p0" : "conv2d_pw"; - } - break; - case Conv2dMethod::SlidingWindow: - kernel_name = "conv2d"; - break; - case Conv2dMethod::Transposed: - kernel_name = "conv_transpose2d"; - break; - } - if (prepack_weights) { - kernel_name += "_prepack_weights"; - } else if (clamp_out) { - kernel_name += "_clamp"; - } - add_dtype_suffix(kernel_name, graph.dtype_of(out)); - - if (prepack_weights) { - add_dtype_suffix(kernel_name, graph.get_staging_dtype_for(weight)); - } - - return VK_KERNEL_FROM_STR(kernel_name); -} - -std::vector get_final_sizes( - const std::vector& original_sizes, - const Conv2dMethod method) { - int64_t batch_padded = utils::align_up_4(utils::val_at(-4, original_sizes)); - int64_t channels_padded = - utils::align_up_4(utils::val_at(-3, original_sizes)); - int64_t height = utils::val_at(-2, original_sizes); - int64_t width = utils::val_at(-1, original_sizes); - - switch (method) { - case Conv2dMethod::Depthwise: - return std::vector{4, batch_padded / 4, height * width}; - case Conv2dMethod::Pointwise: - case Conv2dMethod::SlidingWindow: - return std::vector{ - 4, batch_padded * height / 4, channels_padded * width}; - case Conv2dMethod::Transposed: - return std::vector{ - 4, channels_padded * height / 4, batch_padded * width}; - } -} - -ValueRef prepack_weights( - ComputeGraph& graph, - const ValueRef vref, - const Conv2dMethod method) { - const auto original_sizes = graph.sizes_of(vref); - const auto final_sizes = get_final_sizes(original_sizes, method); - - ValueRef v = graph.add_tensor( - final_sizes, - graph.dtype_of(vref), - utils::kTexture2D, - utils::kChannelsPacked); - - vkapi::ShaderInfo shader = - get_conv2d_shader(graph, v, /*prepack_weights = */ true, method, vref); - - const auto original_sizes_pc = - utils::make_ivec4(original_sizes, /*reverse = */ true); - graph.prepack_nodes().emplace_back(new PrepackNode( - graph, - shader, - graph.create_gwg(v), - graph.create_lwg(v), - vref, - v, - {}, - // Specialization constants - {graph.packed_dim_of(v)}, - {graph.sizes_pc_of(v), - PushConstantDataInfo(&original_sizes_pc, sizeof(original_sizes_pc))})); - - return v; -} - -void check_conv_args( - ComputeGraph& graph, - const ValueRef in, - const ValueRef out) { - VK_CHECK_COND(graph.packed_dim_of(in) == WHCN::kChannelsDim); - VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kChannelsDim); -} - -struct Conv2dParams final { - utils::ivec2 overlay_region; - int in_group_size; -}; - -struct OutputParams final { - float out_min; - float out_max; -}; - -Conv2dParams create_conv2d_params( - ComputeGraph& graph, - const ValueRef weight, - const Kernel2dParams& p, - const bool transposed) { - const auto& overlay_region = utils::make_ivec2({ - p.kernel_size[0] + (p.kernel_size[0] - 1) * (p.dilation[0] - 1), - p.kernel_size[1] + (p.kernel_size[1] - 1) * (p.dilation[1] - 1), - }); - const auto weight_sizes = graph.sizes_of(weight); - const int32_t in_group_size = utils::safe_downcast( - utils::align_up_4(transposed ? weight_sizes.at(0) : weight_sizes.at(1))); - return {overlay_region, in_group_size}; -} - -void check_conv2d_params(const Kernel2dParams& p, const bool transposed) { - if (transposed) { - if (p.dilation[0] > 1 || p.dilation[1] > 1) { - VK_THROW( - "aten.convolution.default: transposed = true, dilation > 1 is not supported yet!"); - } - } -} - -Conv2dMethod get_conv2d_method( - ComputeGraph& graph, - const ValueRef weight, - const int64_t groups, - const bool transposed) { - const auto weight_sizes = graph.sizes_of(weight); - if (!transposed && weight_sizes.at(0) == groups && weight_sizes.at(1) == 1) { - return Conv2dMethod::Depthwise; - } - if (transposed) { - return Conv2dMethod::Transposed; - } - if (weight_sizes.at(2) == 1 && weight_sizes.at(3) == 1) { - return Conv2dMethod::Pointwise; - } - return Conv2dMethod::SlidingWindow; -} - -// Decide whether a SlidingWindow conv2d should be computed via the -// im2col + GEMM path (conv2d_gemm_impl) instead of the direct convolution -// shader. Across 26 configs on Mali-G715 (buffer path) and Adreno SM8650 -// (texture path): FP32 cases were numerically verified against the reference; -// FP16 cases were routing/dispatch-validated only (the reference is float-only -// for the large shapes, so FP16 outputs were not numerically checked). -// -// Only called for SlidingWindow conv2d (1x1 is routed to conv2d_pw and -// Depthwise/Transposed are handled before the call site). -// -// Preconditions (fall back to direct conv if any fail — the im2col path is -// either not applicable or not beneficial): -// - groups == 1 -// - dilation == 1 (all dims) -// -// Selection rule: use im2col once the output channel count is large enough to -// amortize the fixed ~N*K_total im2col gather cost. -constexpr int64_t kIm2colMinCOut = 128; - -// A cheap gather is a second, independent reason to take im2col. The gather -// materializes an N x K_total matrix before the GEMM runs, and a large gather -// is fine when c_out is large because the GEMM reads it back c_out times. When -// c_out is small it is only worth paying if the matrix is small outright, which -// is where the direct shader on Mali loses badly: an 80x80 3x3 conv wants ~4M -// elements, a 640x640 9x9 conv wants ~100M. -constexpr int64_t kIm2colMaxCheapGatherElements = 32 * 1024 * 1024; - -bool should_use_conv2d_im2col( - ComputeGraph& graph, - const ValueRef weight_data, - const int64_t groups_val, - const Kernel2dParams& kernel_params, - const ValueRef out) { - if (groups_val != 1) { - return false; - } - if (kernel_params.dilation[0] != 1 || kernel_params.dilation[1] != 1) { - return false; - } - const auto weight_sizes = graph.sizes_of(weight_data); - const int64_t c_out = weight_sizes.at(0); - const int64_t k_total = - weight_sizes.at(1) * weight_sizes.at(2) * weight_sizes.at(3); - - const auto out_sizes = graph.sizes_of(out); - const size_t ndim = out_sizes.size(); - const int64_t n = out_sizes.at(ndim - 1) * out_sizes.at(ndim - 2); - return c_out >= kIm2colMinCOut || - (graph.device_is_mali() && n * k_total <= kIm2colMaxCheapGatherElements); -} - -GlobalWorkGrid create_conv2d_gwg( - ComputeGraph& graph, - const Conv2dMethod method, - const ValueRef out, - const ValueRef weight_data, - const bool stride_equals_dilation) { - if (method == Conv2dMethod::Pointwise) { - const utils::uvec3 image_extents = graph.logical_limits_of(out); - return GlobalWorkGrid( - {utils::div_up(image_extents[0u], 1u), - utils::div_up(image_extents[1u], 4u), - image_extents[2u]}, - kTiledWorkGrid); - } else { - return graph.create_gwg(out); - } -} - -// Determines which convolution method a dispatch uses. -// -// Depthwise and transposed convolutions have shader names of their own, but -// the name alone cannot separate pointwise from sliding window: the sliding -// window shader is itself named "conv2d", and a pointwise convolution also -// takes that name when its weights are prepacked. Those two are therefore -// separated by the weight's spatial extent. Shared by the global and local -// workgroup size functions below so that the two cannot disagree about the -// same dispatch. -Conv2dMethod infer_conv2d_method_from_shader( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const ValueRef weight_data) { - const std::string& kernel_name = shader.kernel_name; - // Checked before the plain "conv2d" test below, which "conv2d_dw" and - // "conv2d_pw" would otherwise match too. - if (kernel_name.find("conv2d_dw") != std::string::npos) { - return Conv2dMethod::Depthwise; - } - if (kernel_name.find("conv2d_pw") != std::string::npos) { - return Conv2dMethod::Pointwise; - } - if (kernel_name.find("conv_transpose2d") != std::string::npos) { - return Conv2dMethod::Transposed; - } - const auto& weight_sizes = graph->get_tref(weight_data)->sizes; - if (weight_sizes.at(2) == 1 && weight_sizes.at(3) == 1) { - return Conv2dMethod::Pointwise; - } - return Conv2dMethod::SlidingWindow; -} - -// Custom global workgroup size function for conv2d -GlobalWorkGrid conv2d_gwg( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const std::vector& args, - const std::vector& resize_args) { - const ValueRef out = args.at(0).refs.at(0); - const ValueRef weight_data = resize_args.at(0); - - const Conv2dMethod method = - infer_conv2d_method_from_shader(graph, shader, weight_data); - - // Determine stride_equals_dilation from shader name - bool stride_equals_dilation = - shader.kernel_name.find("_sned") == std::string::npos; - - const GlobalWorkGrid wg_size = create_conv2d_gwg( - *graph, method, out, weight_data, stride_equals_dilation); - - if (method == Conv2dMethod::Pointwise) { - utils::uvec3 pointwise_wg_size = {wg_size[0] * wg_size[1], wg_size[2], 1u}; - - if (shader.kernel_name.find("s1p0") != std::string::npos) { - pointwise_wg_size[0] *= 4; - } - return GlobalWorkGrid(pointwise_wg_size, kTiledWorkGrid); - } - - return wg_size; -} - -// Custom local workgroup size function for conv2d -LocalWorkGroup conv2d_lwg( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const GlobalWorkGrid& gwg, - const std::vector& args, - const std::vector& resize_args) { - (void)args; - - const ValueRef weight_data = resize_args.at(0); - const Conv2dMethod method = - infer_conv2d_method_from_shader(graph, shader, weight_data); - - if (method == Conv2dMethod::Pointwise) { - uint32_t lwg_y = 1; - if (gwg[1] % 8 == 0) { - lwg_y = 8; - } else if (gwg[1] % 4 == 0) { - lwg_y = 4; - } else if (gwg[1] % 2 == 0) { - lwg_y = 2; - } - return LocalWorkGroup(64u / lwg_y, lwg_y, 1u); - } else { - return graph->create_lwg(gwg); - } -} - -// Custom global workgroup size function for conv1d -GlobalWorkGrid conv1d_gwg( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)resize_args; - const ValueRef out = args.at(0).refs.at(0); - - return GlobalWorkGrid( - {// out length - graph->size_at(-1, out), - // out channels - static_cast(graph->size_at(-2, out)), - // out batches - utils::div_up_4(graph->size_at(-3, out))}, - kTiledWorkGrid); -} - -void add_conv2d_node( - ComputeGraph& graph, - const ValueRef in, - const ValueRef weight_data, - const ValueRef bias, - const ValueRef stride, - const ValueRef padding, - const ValueRef dilation, - const ValueRef transposed, - const ValueRef output_padding, - const ValueRef groups, - const ValueRef out_min, - const ValueRef out_max, - const ValueRef out, - const bool clamp_out, - const bool force_direct) { - const bool transposed_val = graph.get_bool(transposed); - - float out_min_val = 0.0f; - float out_max_val = 0.0f; - if (out_min != kDummyValueRef) { - out_min_val = graph.extract_scalar(out_min); - } - if (out_max != kDummyValueRef) { - out_max_val = graph.extract_scalar(out_max); - } - - const int64_t groups_val = graph.get_int(groups); - - const Conv2dMethod method = - get_conv2d_method(graph, weight_data, groups_val, transposed_val); - - // Use tiled path for all pointwise conv2d - if (method == Conv2dMethod::Pointwise) { - return conv2d_pw_impl( - graph, - in, - weight_data, - bias, - stride, - padding, - out, - transposed_val, - clamp_out, - out_min_val, - out_max_val); - } - - if (method == Conv2dMethod::Depthwise) { - return conv2d_dw_impl( - graph, - in, - weight_data, - bias, - stride, - padding, - dilation, - out, - clamp_out, - out_min_val, - out_max_val); - } - - const Kernel2dParams kernel_params = create_kernel2d_params( - graph, - weight_data, - /*kernel_size_only = */ false, - stride, - padding, - dilation); - - // SlidingWindow conv2d: route to the im2col + GEMM path when the heuristic - // indicates it is beneficial, falling back to the direct convolution shader - // otherwise. `force_direct` bypasses the heuristic entirely and forces the - // direct path (used by tests to exercise the direct shader regardless of - // device); the default (false) reproduces the production routing exactly. - const bool use_im2col = !force_direct && - method == Conv2dMethod::SlidingWindow && - should_use_conv2d_im2col( - graph, weight_data, groups_val, kernel_params, out); - if (use_im2col) { - return conv2d_gemm_impl( - graph, - in, - weight_data, - bias, - stride, - padding, - dilation, - out, - clamp_out, - out_min_val, - out_max_val); - } - - ValueRef arg_weight = prepack_weights(graph, weight_data, method); - ValueRef arg_bias = prepack_biases( - graph, - bias, - weight_data, - transposed_val, - /* storage_type = */ utils::kTexture2D, - /* memory_layout = */ utils::kWidthPacked); - - const std::vector in_sizes = graph.sizes_of(in); - if (in_sizes.at(0) > 1) { - VK_THROW("conv2d: input batch size > 1 is not supported yet!"); - } - - check_conv_args(graph, in, out); - - Conv2dParams extra_params = - create_conv2d_params(graph, weight_data, kernel_params, transposed_val); - - const bool stride_equals_dilation = - (kernel_params.stride[0] == kernel_params.dilation[0] && - kernel_params.stride[1] == kernel_params.dilation[1]); - - const bool stride_1_padding_0 = - (kernel_params.stride[0] == 1 && kernel_params.stride[1] == 1 && - kernel_params.padding[0] == 0 && kernel_params.padding[1] == 0); - - OutputParams out_params = {out_min_val, out_max_val}; - - check_conv2d_params(kernel_params, transposed_val); - - vkapi::ShaderInfo shader = get_conv2d_shader( - graph, - out, - /*prepack_weights = */ false, - method, - weight_data, - clamp_out, - stride_equals_dilation, - stride_1_padding_0); - - vkapi::ParamsBindList param_buffers = { - graph.logical_limits_ubo(out), - graph.sizes_ubo(in), - graph.create_params_buffer(kernel_params), - graph.create_params_buffer(extra_params), - graph.create_params_buffer(out_params), - }; - - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - shader, - conv2d_gwg, - conv2d_lwg, - // Inputs and Outputs - {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, - // Shader params buffers - param_buffers, - // Push Constants - {}, - // Specialization Constants - {utils::safe_downcast(groups_val)}, - // Resize Args - {weight_data, stride, padding, dilation, transposed, output_padding}, - // Resizing Logic - resize_conv2d_node)); -} - -void add_conv1d_node( - ComputeGraph& graph, - const ValueRef in, - const ValueRef weight, - const ValueRef bias, - const ValueRef stride, - const ValueRef padding, - const ValueRef dilation, - const ValueRef groups, - const ValueRef out_min, - const ValueRef out_max, - const ValueRef out, - const bool clamp_out) { - ValueRef arg_weight = prepack_standard( - graph, - weight, - graph.storage_type_of(out), - utils::kChannelsPacked, - /* passthrough = */ false); - ValueRef arg_bias = prepack_biases( - graph, - bias, - weight, - /*transposed = */ false, - /*storage_type = */ utils::kTexture3D, - /*memory_layout = */ utils::kWidthPacked); - - float out_min_val = 0.0f; - float out_max_val = 0.0f; - if (out_min != kDummyValueRef) { - out_min_val = graph.extract_scalar(out_min); - } - if (out_max != kDummyValueRef) { - out_max_val = graph.extract_scalar(out_max); - } - - const int64_t groups_val = graph.get_int(groups); - - const std::vector in_sizes = graph.sizes_of(in); - const std::vector weight_sizes = graph.sizes_of(arg_weight); - const std::vector out_sizes = graph.sizes_of(out); - - check_conv_args(graph, in, out); - - const int32_t in_channels = in_sizes.at(1); - const int32_t out_channels = weight_sizes.at(0); - const int32_t kernel_size = weight_sizes.at(2); - const int32_t stride_size = graph.get_int_list(stride)->at(0); - const int32_t padding_size = graph.get_int_list(padding)->at(0); - const int32_t dilation_size = graph.get_int_list(dilation)->at(0); - const int32_t in_group_size = static_cast(in_channels / groups_val); - const int32_t out_group_size = - static_cast(out_channels / groups_val); - - Kernel1dParams kernel_params = { - kernel_size, - stride_size, - padding_size, - dilation_size, - in_group_size, - out_group_size}; - - const OutputParams out_params = {out_min_val, out_max_val}; - - std::string kernel_name("conv1d"); - if (clamp_out) { - kernel_name += "_clamp"; - } - kernel_name.reserve(kShaderNameReserve); - - add_dtype_suffix(kernel_name, graph.dtype_of(out)); - - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - VK_KERNEL_FROM_STR(kernel_name), - conv1d_gwg, - default_pick_lwg, - // Inputs and Outputs - {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, - // Shader params buffers - { - graph.logical_limits_ubo(out), - graph.sizes_ubo(in), - graph.create_params_buffer(kernel_params), - graph.create_params_buffer(out_params), - }, - // Push Constants - {}, - // Specialization Constants - {graph.hashed_layout_of(out), - graph.hashed_layout_of(in), - graph.hashed_layout_of(arg_weight), - graph.hashed_layout_of(arg_bias)}, - // Resize Args - {weight, stride, padding, dilation}, - // Resizing Logic - resize_conv1d_node)); -} - -void conv(ComputeGraph& graph, const std::vector& args) { - int64_t in_ndim = graph.dim_of(args[0]); - if (in_ndim == 4) { - if (args.size() == 10) { - // ordinary conv2d - return add_conv2d_node( - graph, - args[0], - args[1], - args[2], - args[3], - args[4], - args[5], - args[6], - args[7], - args[8], - /*out_min = */ kDummyValueRef, - /*out_max = */ kDummyValueRef, - args[9], - false); - } else { - // conv2d with clamp - return add_conv2d_node( - graph, - args[0], - args[1], - args[2], - args[3], - args[4], - args[5], - args[6], - args[7], - args[8], - args[9], - args[10], - args[11], - true); - } - } else { - // Conv1d path - if (graph.packed_dim_of(args[0]) == WHCN::kHeightDim) { - // Height-packed: route to optimized conv1d implementations - const auto weight_sizes = graph.sizes_of(args[1]); - const int64_t groups_val = graph.get_int(args[8]); - const bool is_pointwise = weight_sizes.at(2) == 1; - const bool is_depthwise = - groups_val == weight_sizes.at(0) && weight_sizes.at(1) == 1; - - // Build unified 10-arg vector: - // in, weight, bias, stride, padding, dilation, groups, - // output_min, output_max, out - // For non-clamp (args.size() == 10): output_min/max = kDummyValueRef - // For clamp (args.size() == 12): output_min/max from args[9]/args[10] - ValueRef output_min = kDummyValueRef; - ValueRef output_max = kDummyValueRef; - ValueRef out; - if (args.size() == 10) { - out = args[9]; - } else { - output_min = args[9]; - output_max = args[10]; - out = args[11]; - } - - std::vector conv1d_args = { - args[0], - args[1], - args[2], - args[3], - args[4], - args[5], - args[8], - output_min, - output_max, - out}; - - if (is_pointwise) { - VK_GET_OP_FN("et_vk.conv1d_pw.default")(graph, conv1d_args); - } else if (is_depthwise) { - VK_GET_OP_FN("et_vk.conv1d_dw.default")(graph, conv1d_args); - } else { - VK_THROW( - "Height-packed conv1d only supports pointwise (K=1) or " - "depthwise (groups=C)"); - } - return; - } - - // Existing channels-packed fallback - if (args.size() == 10) { - // ordinary conv1d - return add_conv1d_node( - graph, - args[0], - args[1], - args[2], - args[3], - args[4], - args[5], - args[8], - /*out_min = */ kDummyValueRef, - /*out_max = */ kDummyValueRef, - args[9], - false); - } else { - // conv1d with clamp - return add_conv1d_node( - graph, - args[0], - args[1], - args[2], - args[3], - args[4], - args[5], - args[8], - args[9], - args[10], - args[11], - true); - } - } -} - -REGISTER_OPERATORS { - VK_REGISTER_OP(aten.convolution.default, conv); - VK_REGISTER_OP(conv_with_clamp.default, conv); - VK_REGISTER_OP(et_vk.conv_with_clamp.default, conv); -} - -} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h deleted file mode 100644 index e48d1e32a08..00000000000 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h +++ /dev/null @@ -1,216 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once - -#include -#include -#include - -namespace vkcompute { - -inline constexpr int64_t kQ8taConv2dIm2ColScratchBudgetBytes = 16 * 1024 * 1024; -inline constexpr int64_t kQ8taConv2dMaxRowsPerTile = 65535; - -struct Q8taConv2dStreamPlan final { - int64_t aligned_out_width; - int64_t rows_per_tile; - int64_t num_tiles; - int64_t scratch_bytes; - bool feasible; -}; - -Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan( - int64_t batch, - int64_t flattened_kernel_size, - int64_t out_height, - int64_t out_width, - int64_t scratch_budget_bytes); - -// max_buffer_bytes is Adapter::max_buffer_numel(), which returns -// maxStorageBufferRange in bytes (not elements) — directly comparable with -// the byte-denominated scratch budget. -Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan_for_device( - int64_t batch, - int64_t flattened_kernel_size, - int64_t out_height, - int64_t out_width, - uint64_t max_buffer_bytes); - -enum class ActivationType : uint32_t { - kNone = 0, - kRelu = 1, -}; - -ActivationType activation_type_from_string(const std::string& activation); - -bool q8ta_conv2d_check_packed_dim_info(const api::PackedDimInfo& info); - -bool q8ta_conv2d_check_4w4c_packed_dim_info(const api::PackedDimInfo& info); - -ValueRef prepack_quantized_conv2d_weight( - ComputeGraph& graph, - const QuantizationConfig& weight_quant_config, - const ValueRef weight_data, - const ValueRef input, - const ValueRef output, - const ValueRef groups, - const ValueRef kernel_size); - -ValueRef prepack_quantized_conv2d_weight( - ComputeGraph& graph, - const QuantizationConfig& weight_quant_config, - const ValueRef weight_data, - const ValueRef input, - const ValueRef output, - const ValueRef groups, - const ValueRef kernel_size); - -ValueRef prepack_quantized_conv2d_dw_weight( - ComputeGraph& graph, - const QuantizationConfig& weight_quant_config, - const ValueRef weight_data, - const ValueRef kernel_size); - -void add_q8ta_conv2d_dw_node( - ComputeGraph& graph, - const ValueRef packed_int8_input, - const ValueRef input_scale, - const ValueRef input_zp, - const ValueRef packed_weight, - const ValueRef packed_weight_sums, - const ValueRef packed_weight_scales, - const ValueRef output_scale, - const ValueRef output_zp, - const ValueRef bias_data, - const ValueRef packed_bias, - const ValueRef kernel_size, - const ValueRef stride, - const ValueRef padding, - const ValueRef dilation, - const ValueRef groups, - const uint32_t activation_type, - const ValueRef packed_int8_output); - -void add_conv2d_dw_q8ta_q8csw_q8to_4w4c_node( - ComputeGraph& graph, - const ValueRef packed_int8_input, - const ValueRef input_scale, - const ValueRef input_zp, - const ValueRef packed_weight, - const ValueRef packed_weight_sums, - const ValueRef packed_weight_scales, - const ValueRef output_scale, - const ValueRef output_zp, - const ValueRef bias_data, - const ValueRef packed_bias, - const ValueRef kernel_size, - const ValueRef stride, - const ValueRef padding, - const ValueRef dilation, - const ValueRef groups, - const ValueRef packed_int8_output); - -void add_q8ta_conv2d_node( - ComputeGraph& graph, - const ValueRef packed_int8_input, - const ValueRef packed_int8_input_im2col, - const ValueRef input_scale, - const ValueRef input_zp, - const ValueRef packed_weight, - const ValueRef packed_weight_sums, - const ValueRef packed_weight_scales, - const ValueRef output_scale, - const ValueRef output_zp, - const ValueRef bias_data, - const ValueRef packed_bias, - const ValueRef kernel_size, - const ValueRef stride, - const ValueRef padding, - const ValueRef dilation, - const ValueRef groups, - const uint32_t activation_type, - const ValueRef packed_int8_output); - -void add_q8ta_conv2d_pw_node( - ComputeGraph& graph, - const bool use_unsigned_dot, - const ValueRef packed_int8_input, - const ValueRef input_scale, - const ValueRef input_zp, - const ValueRef packed_weight, - const ValueRef packed_weight_sums, - const ValueRef packed_weight_scales, - const ValueRef output_scale, - const ValueRef output_zp, - const ValueRef bias_data, - const ValueRef packed_bias, - const uint32_t activation_type, - const ValueRef packed_int8_output, - const int32_t groups = 1, - const ValueRef conv_input = kDummyValueRef, - const ValueRef kernel_size = kDummyValueRef, - const ValueRef stride = kDummyValueRef, - const ValueRef padding = kDummyValueRef, - const ValueRef dilation = kDummyValueRef, - const bool is_im2col = false, - const ValueRef stream_row_offset_ref = kDummyValueRef, - const ValueRef max_im2col_rows_ref = kDummyValueRef); - -constexpr int64_t kMaxUnsignedDotAccumulatorBytes = 33025; - -bool can_use_unsigned_pw_dot( - const vkapi::Adapter& adapter, - int64_t k_per_group); - -void q8ta_conv2d_pw_impl( - ComputeGraph& graph, - bool use_unsigned_dot, - const std::vector& args); - -void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args); - -std::vector calculate_q8ta_im2col_sizes( - ComputeGraph* graph, - const ValueRef& input, - const ValueRef& output, - const ValueRef& kernel_size, - const ValueRef& groups); - -void add_q8ta_im2col_node( - ComputeGraph& graph, - const ValueRef packed_int8_input, - const ValueRef kernel_size, - const ValueRef stride, - const ValueRef padding, - const ValueRef dilation, - const ValueRef groups, - const ValueRef packed_int8_output, - const ValueRef packed_int8_im2col, - const int32_t zp, - const ValueRef stream_row_offset_ref, - const ValueRef max_im2col_rows_ref = kDummyValueRef); - -void q8ta_conv2d_im2col(ComputeGraph& graph, const std::vector& args); - -void q8ta_conv2d_im2col_impl( - ComputeGraph& graph, - bool use_unsigned_dot, - const std::vector& args); - -void q8ta_conv2d_general( - ComputeGraph& graph, - const std::vector& args); - -// Transposed convolution - -void q8ta_conv2d_transposed( - ComputeGraph& graph, - const std::vector& args); - -} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp index 683b7b5ce07..ea1a87b910d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taLinear.cpp @@ -11,9 +11,9 @@ #include #include -#include #include #include +#include #include namespace vkcompute { diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp index 7a63d33d546..a3944641839 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taLinearGemv.cpp @@ -11,9 +11,9 @@ #include #include -#include #include #include +#include #include namespace vkcompute { diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/Convolution.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/Convolution.cpp new file mode 100644 index 00000000000..4bbb2c56911 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/Convolution.cpp @@ -0,0 +1,84 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include +#include + +#include + +namespace vkcompute { + +ValueRef prepack_biases( + ComputeGraph& graph, + const ValueRef vref, + const ValueRef weight, + const bool transposed, + const utils::StorageType storage_type, + const utils::GPUMemoryLayout memory_layout) { + auto sizes = graph.sizes_of(weight); + const int64_t out_channels = transposed ? sizes.at(1) : sizes.at(0); + + ValueRef v = graph.add_tensor( + {out_channels}, graph.dtype_of(weight), storage_type, memory_layout); + + vkapi::ShaderInfo shader = + get_nchw_to_tensor_shader(graph, v, graph.get_staging_dtype_for(weight)); + + vkapi::ParamsBindList param_buffers = {}; + if (graph.is_buffer_storage(v)) { + param_buffers.append(graph.buffer_meta_ubo(v)); + } else { + param_buffers.append(graph.texture_meta_ubo(v)); + } + + std::vector pcs; + if (graph.is_buffer_storage(v)) { + pcs = {graph.sizes_pc_of(v), graph.strides_pc_of(v), graph.numel_pc_of(v)}; + } + + graph.prepack_nodes().emplace_back(new PrepackNode( + graph, + shader, + graph.create_gwg(v), + graph.create_lwg(v), + vref, + v, + param_buffers, + // Specialization constants + {graph.hashed_layout_of(v)}, + pcs)); + + return v; +} + +void check_conv_args( + ComputeGraph& graph, + const ValueRef in, + const ValueRef out) { + VK_CHECK_COND(graph.packed_dim_of(in) == WHCN::kChannelsDim); + VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kChannelsDim); +} + +void conv(ComputeGraph& graph, const std::vector& args) { + if (graph.dim_of(args[0]) == 4) { + return conv2d_impl(graph, args); + } + return conv1d_impl(graph, args); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(aten.convolution.default, conv); + VK_REGISTER_OP(conv_with_clamp.default, conv); + VK_REGISTER_OP(et_vk.conv_with_clamp.default, conv); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/Convolution.h b/backends/vulkan/runtime/graph/ops/impl/convolution/Convolution.h new file mode 100644 index 00000000000..08bf5bd5c9c --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/Convolution.h @@ -0,0 +1,28 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +ValueRef prepack_biases( + ComputeGraph& graph, + const ValueRef vref, + const ValueRef weight, + const bool transposed, + const utils::StorageType storage_type, + const utils::GPUMemoryLayout memory_layout); + +void check_conv_args( + ComputeGraph& graph, + const ValueRef in, + const ValueRef out); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1d.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1d.cpp new file mode 100644 index 00000000000..1f71e009d32 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1d.cpp @@ -0,0 +1,100 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include + +namespace vkcompute { + +void conv1d_impl(ComputeGraph& graph, const std::vector& args) { + if (graph.packed_dim_of(args[0]) == WHCN::kHeightDim) { + // Height-packed: route to optimized conv1d implementations + const auto weight_sizes = graph.sizes_of(args[1]); + const int64_t groups_val = graph.get_int(args[8]); + const bool is_pointwise = weight_sizes.at(2) == 1; + const bool is_depthwise = + groups_val == weight_sizes.at(0) && weight_sizes.at(1) == 1; + + // Build unified 10-arg vector: + // in, weight, bias, stride, padding, dilation, groups, + // output_min, output_max, out + // For non-clamp (args.size() == 10): output_min/max = kDummyValueRef + // For clamp (args.size() == 12): output_min/max from args[9]/args[10] + ValueRef output_min = kDummyValueRef; + ValueRef output_max = kDummyValueRef; + ValueRef out; + if (args.size() == 10) { + out = args[9]; + } else { + output_min = args[9]; + output_max = args[10]; + out = args[11]; + } + + std::vector conv1d_args = { + args[0], + args[1], + args[2], + args[3], + args[4], + args[5], + args[8], + output_min, + output_max, + out}; + + if (is_pointwise) { + VK_GET_OP_FN("et_vk.conv1d_pw.default")(graph, conv1d_args); + } else if (is_depthwise) { + VK_GET_OP_FN("et_vk.conv1d_dw.default")(graph, conv1d_args); + } else { + VK_THROW( + "Height-packed conv1d only supports pointwise (K=1) or " + "depthwise (groups=C)"); + } + return; + } + + // Existing channels-packed fallback + if (args.size() == 10) { + // ordinary conv1d + return add_conv1d_direct_node( + graph, + args[0], + args[1], + args[2], + args[3], + args[4], + args[5], + args[8], + /*out_min = */ kDummyValueRef, + /*out_max = */ kDummyValueRef, + args[9], + false); + } else { + // conv1d with clamp + return add_conv1d_direct_node( + graph, + args[0], + args[1], + args[2], + args[3], + args[4], + args[5], + args[8], + args[9], + args[10], + args[11], + true); + } +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1d.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1d.h new file mode 100644 index 00000000000..f6030128587 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1d.h @@ -0,0 +1,17 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +void conv1d_impl(ComputeGraph& graph, const std::vector& args); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv1dDW.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDW.cpp similarity index 97% rename from backends/vulkan/runtime/graph/ops/impl/Conv1dDW.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDW.cpp index ba59090c124..58971984821 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv1dDW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDW.cpp @@ -150,7 +150,7 @@ void add_conv1d_dw_node( // Args: in, weight, bias, stride, padding, dilation, groups, // output_min, output_max, out // output_min and output_max may be kDummyValueRef (no clamp). -void conv1d_dw(ComputeGraph& graph, const std::vector& args) { +void conv1d_dw_impl(ComputeGraph& graph, const std::vector& args) { ValueRef in = args[0]; ValueRef weight = args[1]; ValueRef bias = args[2]; @@ -182,7 +182,7 @@ void conv1d_dw(ComputeGraph& graph, const std::vector& args) { } REGISTER_OPERATORS { - VK_REGISTER_OP(et_vk.conv1d_dw.default, conv1d_dw); + VK_REGISTER_OP(et_vk.conv1d_dw.default, conv1d_dw_impl); } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDirect.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDirect.cpp new file mode 100644 index 00000000000..bdd6f8da6e9 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDirect.cpp @@ -0,0 +1,179 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include + +#include + +#include + +namespace vkcompute { + +void resize_conv1d_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& extra_args) { + const ValueRef out = args.at(0).refs.at(0); + const ValueRef self = args.at(1).refs.at(0); + TensorRefPtr weight_ref = graph->get_tref(extra_args.at(0)); + + const int64_t stride_size = graph->get_int_list(extra_args.at(1))->at(0); + const int64_t padding_size = graph->get_int_list(extra_args.at(2))->at(0); + const int64_t dilation_size = graph->get_int_list(extra_args.at(3))->at(0); + + const std::vector& weight_sizes = weight_ref->sizes; + + const std::vector in_sizes = graph->sizes_of(self); + const size_t ndim = in_sizes.size(); + std::vector new_out_sizes(ndim); + + const int64_t kernel_size = weight_sizes.at(2); + const int64_t in_length = in_sizes.at(2); + + new_out_sizes.at(0) = in_sizes.at(0); + new_out_sizes.at(1) = weight_sizes.at(0); + new_out_sizes.at(2) = calc_out_size( + in_length, kernel_size, stride_size, padding_size, dilation_size, false); + + graph->virtual_resize(out, new_out_sizes); +} + +struct OutputParams final { + float out_min; + float out_max; +}; + +// Custom global workgroup size function for conv1d +GlobalWorkGrid conv1d_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + const ValueRef out = args.at(0).refs.at(0); + + return GlobalWorkGrid( + {// out length + graph->size_at(-1, out), + // out channels + static_cast(graph->size_at(-2, out)), + // out batches + utils::div_up_4(graph->size_at(-3, out))}, + kTiledWorkGrid); +} + +void add_conv1d_direct_node( + ComputeGraph& graph, + const ValueRef in, + const ValueRef weight, + const ValueRef bias, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const ValueRef out_min, + const ValueRef out_max, + const ValueRef out, + const bool clamp_out) { + ValueRef arg_weight = prepack_standard( + graph, + weight, + graph.storage_type_of(out), + utils::kChannelsPacked, + /* passthrough = */ false); + ValueRef arg_bias = prepack_biases( + graph, + bias, + weight, + /*transposed = */ false, + /*storage_type = */ utils::kTexture3D, + /*memory_layout = */ utils::kWidthPacked); + + float out_min_val = 0.0f; + float out_max_val = 0.0f; + if (out_min != kDummyValueRef) { + out_min_val = graph.extract_scalar(out_min); + } + if (out_max != kDummyValueRef) { + out_max_val = graph.extract_scalar(out_max); + } + + const int64_t groups_val = graph.get_int(groups); + + const std::vector in_sizes = graph.sizes_of(in); + const std::vector weight_sizes = graph.sizes_of(arg_weight); + const std::vector out_sizes = graph.sizes_of(out); + + check_conv_args(graph, in, out); + + const int32_t in_channels = utils::safe_downcast(in_sizes.at(1)); + const int32_t out_channels = + utils::safe_downcast(weight_sizes.at(0)); + const int32_t kernel_size = utils::safe_downcast(weight_sizes.at(2)); + const int32_t stride_size = + utils::safe_downcast(graph.get_int_list(stride)->at(0)); + const int32_t padding_size = + utils::safe_downcast(graph.get_int_list(padding)->at(0)); + const int32_t dilation_size = + utils::safe_downcast(graph.get_int_list(dilation)->at(0)); + const int32_t in_group_size = + utils::safe_downcast(in_channels / groups_val); + const int32_t out_group_size = + utils::safe_downcast(out_channels / groups_val); + + Kernel1dParams kernel_params = { + kernel_size, + stride_size, + padding_size, + dilation_size, + in_group_size, + out_group_size}; + + const OutputParams out_params = {out_min_val, out_max_val}; + + std::string kernel_name("conv1d"); + if (clamp_out) { + kernel_name += "_clamp"; + } + kernel_name.reserve(kShaderNameReserve); + + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + conv1d_gwg, + default_pick_lwg, + // Inputs and Outputs + {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, + // Shader params buffers + { + graph.logical_limits_ubo(out), + graph.sizes_ubo(in), + graph.create_params_buffer(kernel_params), + graph.create_params_buffer(out_params), + }, + // Push Constants + {}, + // Specialization Constants + {graph.hashed_layout_of(out), + graph.hashed_layout_of(in), + graph.hashed_layout_of(arg_weight), + graph.hashed_layout_of(arg_bias)}, + // Resize Args + {weight, stride, padding, dilation}, + // Resizing Logic + resize_conv1d_node)); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDirect.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDirect.h new file mode 100644 index 00000000000..9800330b9f8 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dDirect.h @@ -0,0 +1,29 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +void add_conv1d_direct_node( + ComputeGraph& graph, + const ValueRef in, + const ValueRef weight, + const ValueRef bias, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const ValueRef out_min, + const ValueRef out_max, + const ValueRef out, + const bool clamp_out); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv1dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dPW.cpp similarity index 98% rename from backends/vulkan/runtime/graph/ops/impl/Conv1dPW.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dPW.cpp index de2fcdbb806..225a5118688 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv1dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv1d/Conv1dPW.cpp @@ -211,7 +211,7 @@ void add_conv1d_pw_node( // Args: in, weight, bias, stride, padding, dilation, groups, // output_min, output_max, out // output_min and output_max may be kDummyValueRef (no clamp). -void conv1d_pw(ComputeGraph& graph, const std::vector& args) { +void conv1d_pw_impl(ComputeGraph& graph, const std::vector& args) { ValueRef in = args[0]; ValueRef weight = args[1]; ValueRef bias = args[2]; @@ -236,7 +236,7 @@ void conv1d_pw(ComputeGraph& graph, const std::vector& args) { } REGISTER_OPERATORS { - VK_REGISTER_OP(et_vk.conv1d_pw.default, conv1d_pw); + VK_REGISTER_OP(et_vk.conv1d_pw.default, conv1d_pw_impl); } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2d.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2d.cpp new file mode 100644 index 00000000000..ff4f8173780 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2d.cpp @@ -0,0 +1,234 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include + +#include + +namespace vkcompute { + +void resize_conv2d_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& extra_args) { + const ValueRef out = args.at(0).refs.at(0); + const ValueRef self = args.at(1).refs.at(0); + + size_t ndim = graph->dim_of(self); + std::vector new_out_sizes(ndim); + const bool transposed = graph->get_bool(extra_args.at(4)); + + std::vector self_sizes = graph->sizes_of(self); + // Batch, Channel + if (ndim == 4) { + new_out_sizes.at(ndim - 4) = self_sizes.at(ndim - 4); + } + + TensorRefPtr weight_ref = graph->get_tref(extra_args.at(0)); + const auto& weight_sizes = weight_ref->sizes; + new_out_sizes.at(ndim - 3) = + transposed ? weight_sizes.at(ndim - 3) : weight_sizes.at(ndim - 4); + + // Height, Width + const auto& new_out_sizes_hw = calc_out_sizes_hw( + *graph, + self_sizes, + extra_args.at(0), + /*kernel_size_only = */ false, + {extra_args.at(1), extra_args.at(2), extra_args.at(3), extra_args.at(5)}, + transposed); + new_out_sizes.at(ndim - 2) = new_out_sizes_hw.at(0); + new_out_sizes.at(ndim - 1) = new_out_sizes_hw.at(1); + + graph->virtual_resize(out, new_out_sizes); +} + +Conv2dMethod get_conv2d_method( + ComputeGraph& graph, + const ValueRef weight, + const int64_t groups, + const bool transposed) { + const auto weight_sizes = graph.sizes_of(weight); + if (!transposed && weight_sizes.at(0) == groups && weight_sizes.at(1) == 1) { + return Conv2dMethod::Depthwise; + } + if (transposed) { + return Conv2dMethod::Transposed; + } + if (weight_sizes.at(2) == 1 && weight_sizes.at(3) == 1) { + return Conv2dMethod::Pointwise; + } + return Conv2dMethod::SlidingWindow; +} + +// Decide whether a SlidingWindow conv2d should be computed via the +// im2col + GEMM path (conv2d_gemm_impl) instead of the direct convolution +// shader. Across 26 configs on Mali-G715 (buffer path) and Adreno SM8650 +// (texture path): FP32 cases were numerically verified against the reference; +// FP16 cases were routing/dispatch-validated only (the reference is float-only +// for the large shapes, so FP16 outputs were not numerically checked). +// +// Only called for SlidingWindow conv2d (1x1 is routed to conv2d_pw and +// Depthwise/Transposed are handled before the call site). +// +// Preconditions (fall back to direct conv if any fail — the im2col path is +// either not applicable or not beneficial): +// - groups == 1 +// - dilation == 1 (all dims) +// +// Selection rule: use im2col once the output channel count is large enough to +// amortize the fixed ~N*K_total im2col gather cost. +constexpr int64_t kIm2colMinCOut = 128; + +// A cheap gather is a second, independent reason to take im2col. The gather +// materializes an N x K_total matrix before the GEMM runs, and a large gather +// is fine when c_out is large because the GEMM reads it back c_out times. When +// c_out is small it is only worth paying if the matrix is small outright, which +// is where the direct shader on Mali loses badly: an 80x80 3x3 conv wants ~4M +// elements, a 640x640 9x9 conv wants ~100M. +constexpr int64_t kIm2colMaxCheapGatherElements = 32 * 1024 * 1024; + +bool should_use_conv2d_im2col( + ComputeGraph& graph, + const ValueRef weight_data, + const int64_t groups_val, + const Kernel2dParams& kernel_params, + const ValueRef out) { + if (groups_val != 1) { + return false; + } + if (kernel_params.dilation[0] != 1 || kernel_params.dilation[1] != 1) { + return false; + } + const auto weight_sizes = graph.sizes_of(weight_data); + const int64_t c_out = weight_sizes.at(0); + const int64_t k_total = + weight_sizes.at(1) * weight_sizes.at(2) * weight_sizes.at(3); + + const auto out_sizes = graph.sizes_of(out); + const size_t ndim = out_sizes.size(); + const int64_t n = out_sizes.at(ndim - 1) * out_sizes.at(ndim - 2); + return c_out >= kIm2colMinCOut || + (graph.device_is_mali() && n * k_total <= kIm2colMaxCheapGatherElements); +} + +void conv2d_impl(ComputeGraph& graph, const std::vector& args) { + const ValueRef in = args[0]; + const ValueRef weight_data = args[1]; + const ValueRef bias = args[2]; + const ValueRef stride = args[3]; + const ValueRef padding = args[4]; + const ValueRef dilation = args[5]; + const ValueRef transposed = args[6]; + const ValueRef output_padding = args[7]; + const ValueRef groups = args[8]; + // conv2d with clamp passes out_min and out_max before the output + const bool clamp_out = args.size() != 10; + const ValueRef out_min = clamp_out ? args[9] : kDummyValueRef; + const ValueRef out_max = clamp_out ? args[10] : kDummyValueRef; + const ValueRef out = clamp_out ? args[11] : args[9]; + + const bool transposed_val = graph.get_bool(transposed); + + float out_min_val = 0.0f; + float out_max_val = 0.0f; + if (out_min != kDummyValueRef) { + out_min_val = graph.extract_scalar(out_min); + } + if (out_max != kDummyValueRef) { + out_max_val = graph.extract_scalar(out_max); + } + + const int64_t groups_val = graph.get_int(groups); + + const Conv2dMethod method = + get_conv2d_method(graph, weight_data, groups_val, transposed_val); + + // Use tiled path for all pointwise conv2d + if (method == Conv2dMethod::Pointwise) { + return conv2d_pw_impl( + graph, + in, + weight_data, + bias, + stride, + padding, + out, + transposed_val, + clamp_out, + out_min_val, + out_max_val); + } + + if (method == Conv2dMethod::Depthwise) { + return conv2d_dw_impl( + graph, + in, + weight_data, + bias, + stride, + padding, + dilation, + out, + clamp_out, + out_min_val, + out_max_val); + } + + const Kernel2dParams kernel_params = create_kernel2d_params( + graph, + weight_data, + /*kernel_size_only = */ false, + stride, + padding, + dilation); + + // SlidingWindow conv2d: route to the im2col + GEMM path when the heuristic + // indicates it is beneficial, falling back to the direct convolution shader + // otherwise. + const bool use_im2col = + method == Conv2dMethod::SlidingWindow && + should_use_conv2d_im2col( + graph, weight_data, groups_val, kernel_params, out); + if (use_im2col) { + return conv2d_gemm_impl( + graph, + in, + weight_data, + bias, + stride, + padding, + dilation, + out, + clamp_out, + out_min_val, + out_max_val); + } + + return conv2d_direct_impl( + graph, + in, + weight_data, + bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + out, + clamp_out, + out_min_val, + out_max_val); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Convolution.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2d.h similarity index 50% rename from backends/vulkan/runtime/graph/ops/impl/Convolution.h rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2d.h index 12e4733e67a..446859c363d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Convolution.h +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2d.h @@ -12,18 +12,12 @@ namespace vkcompute { -ValueRef prepack_biases( - ComputeGraph& graph, - const ValueRef vref, - const ValueRef weight, - const bool transposed, - const utils::StorageType storage_type, - const utils::GPUMemoryLayout memory_layout); - -void check_conv_args( - ComputeGraph& graph, - const ValueRef in, - const ValueRef out); +enum class Conv2dMethod : uint8_t { + Depthwise, + Pointwise, + SlidingWindow, + Transposed, +}; void conv2d_pw_impl( ComputeGraph& graph, @@ -56,26 +50,6 @@ void resize_conv2d_node( const std::vector& args, const std::vector& extra_args); -// `force_direct` overrides the im2col-vs-direct routing heuristic: when true, -// a SlidingWindow conv2d always takes the direct sliding-window path, -// bypassing should_use_conv2d_im2col(). The default (false) preserves the -// production routing exactly. Pointwise / Depthwise / Transposed methods are -// unaffected by this flag. -void add_conv2d_node( - ComputeGraph& graph, - const ValueRef in, - const ValueRef weight_data, - const ValueRef bias, - const ValueRef stride, - const ValueRef padding, - const ValueRef dilation, - const ValueRef transposed, - const ValueRef output_padding, - const ValueRef groups, - const ValueRef out_min, - const ValueRef out_max, - const ValueRef out, - const bool clamp_out, - const bool force_direct = false); +void conv2d_impl(ComputeGraph& graph, const std::vector& args); } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2dDW.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDW.cpp similarity index 98% rename from backends/vulkan/runtime/graph/ops/impl/Conv2dDW.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDW.cpp index a21f6c93e4e..0ac6a56d7ad 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv2dDW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDW.cpp @@ -6,10 +6,11 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include #include +#include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDirect.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDirect.cpp new file mode 100644 index 00000000000..a476d03a900 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDirect.cpp @@ -0,0 +1,402 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include + +#include + +#include + +namespace vkcompute { + +vkapi::ShaderInfo get_conv2d_shader( + ComputeGraph& graph, + const ValueRef out, + const bool prepack_weights, + const Conv2dMethod method, + const ValueRef weight, + const bool clamp_out = false, + const bool stride_equals_dilation = false, + const bool stride_1_padding_0 = false) { + (void)stride_equals_dilation; + std::string kernel_name; + kernel_name.reserve(kShaderNameReserve); + switch (method) { + case Conv2dMethod::Depthwise: + kernel_name = "conv2d_dw"; + break; + case Conv2dMethod::Pointwise: + if (prepack_weights) { + kernel_name = "conv2d"; + } else { + kernel_name = stride_1_padding_0 ? "conv2d_pw_s1p0" : "conv2d_pw"; + } + break; + case Conv2dMethod::SlidingWindow: + kernel_name = "conv2d"; + break; + case Conv2dMethod::Transposed: + kernel_name = "conv_transpose2d"; + break; + } + if (prepack_weights) { + kernel_name += "_prepack_weights"; + } else if (clamp_out) { + kernel_name += "_clamp"; + } + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + + if (prepack_weights) { + add_dtype_suffix(kernel_name, graph.get_staging_dtype_for(weight)); + } + + return VK_KERNEL_FROM_STR(kernel_name); +} + +std::vector get_final_sizes( + const std::vector& original_sizes, + const Conv2dMethod method) { + int64_t batch_padded = utils::align_up_4(utils::val_at(-4, original_sizes)); + int64_t channels_padded = + utils::align_up_4(utils::val_at(-3, original_sizes)); + int64_t height = utils::val_at(-2, original_sizes); + int64_t width = utils::val_at(-1, original_sizes); + + switch (method) { + case Conv2dMethod::Depthwise: + return std::vector{4, batch_padded / 4, height * width}; + case Conv2dMethod::Pointwise: + case Conv2dMethod::SlidingWindow: + return std::vector{ + 4, batch_padded * height / 4, channels_padded * width}; + case Conv2dMethod::Transposed: + return std::vector{ + 4, channels_padded * height / 4, batch_padded * width}; + } +} + +ValueRef prepack_weights( + ComputeGraph& graph, + const ValueRef vref, + const Conv2dMethod method) { + const auto original_sizes = graph.sizes_of(vref); + const auto final_sizes = get_final_sizes(original_sizes, method); + + ValueRef v = graph.add_tensor( + final_sizes, + graph.dtype_of(vref), + utils::kTexture2D, + utils::kChannelsPacked); + + vkapi::ShaderInfo shader = + get_conv2d_shader(graph, v, /*prepack_weights = */ true, method, vref); + + const auto original_sizes_pc = + utils::make_ivec4(original_sizes, /*reverse = */ true); + graph.prepack_nodes().emplace_back(new PrepackNode( + graph, + shader, + graph.create_gwg(v), + graph.create_lwg(v), + vref, + v, + {}, + // Specialization constants + {graph.packed_dim_of(v)}, + {graph.sizes_pc_of(v), + PushConstantDataInfo(&original_sizes_pc, sizeof(original_sizes_pc))})); + + return v; +} + +namespace { + +struct Conv2dParams final { + utils::ivec2 overlay_region; + int in_group_size; +}; + +} // namespace + +struct OutputParams final { + float out_min; + float out_max; +}; + +Conv2dParams create_conv2d_params( + ComputeGraph& graph, + const ValueRef weight, + const Kernel2dParams& p, + const bool transposed) { + const auto& overlay_region = utils::make_ivec2({ + p.kernel_size[0] + (p.kernel_size[0] - 1) * (p.dilation[0] - 1), + p.kernel_size[1] + (p.kernel_size[1] - 1) * (p.dilation[1] - 1), + }); + const auto weight_sizes = graph.sizes_of(weight); + const int32_t in_group_size = utils::safe_downcast( + utils::align_up_4(transposed ? weight_sizes.at(0) : weight_sizes.at(1))); + return {overlay_region, in_group_size}; +} + +void check_conv2d_params(const Kernel2dParams& p, const bool transposed) { + if (transposed) { + if (p.dilation[0] > 1 || p.dilation[1] > 1) { + VK_THROW( + "aten.convolution.default: transposed = true, dilation > 1 is not supported yet!"); + } + } +} + +GlobalWorkGrid create_conv2d_gwg( + ComputeGraph& graph, + const Conv2dMethod method, + const ValueRef out, + const ValueRef weight_data, + const bool stride_equals_dilation) { + (void)weight_data; + (void)stride_equals_dilation; + if (method == Conv2dMethod::Pointwise) { + const utils::uvec3 image_extents = graph.logical_limits_of(out); + return GlobalWorkGrid( + {utils::div_up(image_extents[0u], 1u), + utils::div_up(image_extents[1u], 4u), + image_extents[2u]}, + kTiledWorkGrid); + } else { + return graph.create_gwg(out); + } +} + +// Determines which convolution method a dispatch uses. +// +// Depthwise and transposed convolutions have shader names of their own, but +// the name alone cannot separate pointwise from sliding window: the sliding +// window shader is itself named "conv2d", and a pointwise convolution also +// takes that name when its weights are prepacked. Those two are therefore +// separated by the weight's spatial extent. Shared by the global and local +// workgroup size functions below so that the two cannot disagree about the +// same dispatch. +Conv2dMethod infer_conv2d_method_from_shader( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const ValueRef weight_data) { + const std::string& kernel_name = shader.kernel_name; + // Checked before the plain "conv2d" test below, which "conv2d_dw" and + // "conv2d_pw" would otherwise match too. + if (kernel_name.find("conv2d_dw") != std::string::npos) { + return Conv2dMethod::Depthwise; + } + if (kernel_name.find("conv2d_pw") != std::string::npos) { + return Conv2dMethod::Pointwise; + } + if (kernel_name.find("conv_transpose2d") != std::string::npos) { + return Conv2dMethod::Transposed; + } + const auto& weight_sizes = graph->get_tref(weight_data)->sizes; + if (weight_sizes.at(2) == 1 && weight_sizes.at(3) == 1) { + return Conv2dMethod::Pointwise; + } + return Conv2dMethod::SlidingWindow; +} + +// Custom global workgroup size function for conv2d +GlobalWorkGrid conv2d_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef out = args.at(0).refs.at(0); + const ValueRef weight_data = resize_args.at(0); + + const Conv2dMethod method = + infer_conv2d_method_from_shader(graph, shader, weight_data); + + // Determine stride_equals_dilation from shader name + bool stride_equals_dilation = + shader.kernel_name.find("_sned") == std::string::npos; + + const GlobalWorkGrid wg_size = create_conv2d_gwg( + *graph, method, out, weight_data, stride_equals_dilation); + + if (method == Conv2dMethod::Pointwise) { + utils::uvec3 pointwise_wg_size = {wg_size[0] * wg_size[1], wg_size[2], 1u}; + + if (shader.kernel_name.find("s1p0") != std::string::npos) { + pointwise_wg_size[0] *= 4; + } + return GlobalWorkGrid(pointwise_wg_size, kTiledWorkGrid); + } + + return wg_size; +} + +// Custom local workgroup size function for conv2d +LocalWorkGroup conv2d_lwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const GlobalWorkGrid& gwg, + const std::vector& args, + const std::vector& resize_args) { + (void)args; + + const ValueRef weight_data = resize_args.at(0); + const Conv2dMethod method = + infer_conv2d_method_from_shader(graph, shader, weight_data); + + if (method == Conv2dMethod::Pointwise) { + uint32_t lwg_y = 1; + if (gwg[1] % 8 == 0) { + lwg_y = 8; + } else if (gwg[1] % 4 == 0) { + lwg_y = 4; + } else if (gwg[1] % 2 == 0) { + lwg_y = 2; + } + return LocalWorkGroup(64u / lwg_y, lwg_y, 1u); + } else { + return graph->create_lwg(gwg); + } +} + +void add_conv2d_direct_node( + ComputeGraph& graph, + const ValueRef in, + const ValueRef arg_weight, + const ValueRef arg_bias, + const ValueRef weight_data, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef transposed, + const ValueRef output_padding, + const ValueRef out, + const vkapi::ShaderInfo& shader, + const Kernel2dParams& kernel_params, + const Conv2dParams& extra_params, + const OutputParams& out_params, + const int64_t groups_val) { + vkapi::ParamsBindList param_buffers = { + graph.logical_limits_ubo(out), + graph.sizes_ubo(in), + graph.create_params_buffer(kernel_params), + graph.create_params_buffer(extra_params), + graph.create_params_buffer(out_params), + }; + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + shader, + conv2d_gwg, + conv2d_lwg, + // Inputs and Outputs + {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + {utils::safe_downcast(groups_val)}, + // Resize Args + {weight_data, stride, padding, dilation, transposed, output_padding}, + // Resizing Logic + resize_conv2d_node)); +} + +void conv2d_direct_impl( + ComputeGraph& graph, + const ValueRef in, + const ValueRef weight_data, + const ValueRef bias, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef transposed, + const ValueRef output_padding, + const ValueRef groups, + const ValueRef out, + const bool clamp_out, + const float out_min_val, + const float out_max_val) { + const bool transposed_val = graph.get_bool(transposed); + const int64_t groups_val = graph.get_int(groups); + const Conv2dMethod method = + transposed_val ? Conv2dMethod::Transposed : Conv2dMethod::SlidingWindow; + + const Kernel2dParams kernel_params = create_kernel2d_params( + graph, + weight_data, + /*kernel_size_only = */ false, + stride, + padding, + dilation); + + ValueRef arg_weight = prepack_weights(graph, weight_data, method); + ValueRef arg_bias = prepack_biases( + graph, + bias, + weight_data, + transposed_val, + /* storage_type = */ utils::kTexture2D, + /* memory_layout = */ utils::kWidthPacked); + + const std::vector in_sizes = graph.sizes_of(in); + if (in_sizes.at(0) > 1) { + VK_THROW("conv2d: input batch size > 1 is not supported yet!"); + } + + check_conv_args(graph, in, out); + + Conv2dParams extra_params = + create_conv2d_params(graph, weight_data, kernel_params, transposed_val); + + const bool stride_equals_dilation = + (kernel_params.stride[0] == kernel_params.dilation[0] && + kernel_params.stride[1] == kernel_params.dilation[1]); + + const bool stride_1_padding_0 = + (kernel_params.stride[0] == 1 && kernel_params.stride[1] == 1 && + kernel_params.padding[0] == 0 && kernel_params.padding[1] == 0); + + OutputParams out_params = {out_min_val, out_max_val}; + + check_conv2d_params(kernel_params, transposed_val); + + vkapi::ShaderInfo shader = get_conv2d_shader( + graph, + out, + /*prepack_weights = */ false, + method, + weight_data, + clamp_out, + stride_equals_dilation, + stride_1_padding_0); + + add_conv2d_direct_node( + graph, + in, + arg_weight, + arg_bias, + weight_data, + stride, + padding, + dilation, + transposed, + output_padding, + out, + shader, + kernel_params, + extra_params, + out_params, + groups_val); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDirect.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDirect.h new file mode 100644 index 00000000000..ef3ea388044 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dDirect.h @@ -0,0 +1,33 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +// Dispatches the direct conv2d shader, or the transposed conv2d shader when +// `transposed` is true; the two share weight prepacking and dispatch logic. +void conv2d_direct_impl( + ComputeGraph& graph, + const ValueRef in, + const ValueRef weight_data, + const ValueRef bias, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef transposed, + const ValueRef output_padding, + const ValueRef groups, + const ValueRef out, + const bool clamp_out, + const float out_min_val, + const float out_max_val); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dPW.cpp similarity index 97% rename from backends/vulkan/runtime/graph/ops/impl/Conv2dPW.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dPW.cpp index 3c227448cf2..736bb306f77 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv2dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/Conv2dPW.cpp @@ -6,10 +6,11 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include #include +#include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/ConvolutionUtils.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/ConvolutionUtils.h similarity index 100% rename from backends/vulkan/runtime/graph/ops/impl/ConvolutionUtils.h rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/ConvolutionUtils.h diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2dGemm.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dGemm.cpp similarity index 98% rename from backends/vulkan/runtime/graph/ops/impl/Conv2dGemm.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dGemm.cpp index 6285dc92d4a..55c6582335b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv2dGemm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dGemm.cpp @@ -6,13 +6,13 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include #include -#include -#include #include +#include +#include #include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2dGemm.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dGemm.h similarity index 100% rename from backends/vulkan/runtime/graph/ops/impl/Conv2dGemm.h rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dGemm.h diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2dIm2Col.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dIm2Col.cpp similarity index 98% rename from backends/vulkan/runtime/graph/ops/impl/Conv2dIm2Col.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dIm2Col.cpp index e1bbded799e..e4dcf6592c2 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Conv2dIm2Col.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dIm2Col.cpp @@ -6,7 +6,7 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/Conv2dIm2Col.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dIm2Col.h similarity index 100% rename from backends/vulkan/runtime/graph/ops/impl/Conv2dIm2Col.h rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/im2col/Conv2dIm2Col.h diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2d.cpp similarity index 70% rename from backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2d.cpp index 780dfcbfbf3..290207295ca 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2d.cpp @@ -6,9 +6,13 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include -#include +#include + +#include +#include +#include #include #include @@ -144,4 +148,55 @@ bool should_use_q8ta_conv2d_im2col(const Q8taConv2dRouteParams& params) { (params.in_channels_per_group >= 32 || spatial_out <= 4096); } +void q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { + const ValueRef input = args.at(0); + const ValueRef kernel_size_ref = args.at(9); + const ValueRef groups_ref = args.at(13); + const ValueRef output = args.at(15); + + const int64_t groups = graph.extract_scalar(groups_ref); + // Valid models always carry groups >= 1; fail fast on corrupt input + // instead of dividing channel counts by zero downstream (both this + // dispatcher and q8ta_conv2d_general_impl divide by groups). + VK_CHECK_COND(groups > 0, "q8ta_conv2d requires groups >= 1"); + const int64_t in_channels = graph.size_at(-3, input); + const int64_t in_channels_per_group = in_channels / groups; + const int64_t batch = graph.size_at(-4, input); + + const int64_t H_out = graph.size_at(-2, output); + const int64_t W_out = graph.size_at(-1, output); + const int64_t out_channels = graph.size_at(-3, output); + int64_t kernel_height; + int64_t kernel_width; + { + const auto kernel_size = graph.get_int_list(kernel_size_ref); + kernel_height = kernel_size->at(0); + kernel_width = kernel_size->at(1); + } + + const bool use_im2col = should_use_q8ta_conv2d_im2col({ + graph.device_is_mali(), + graph.can_use_int8_dot_product(), + static_cast(graph.max_buffer_numel()), + batch, + groups, + in_channels_per_group, + out_channels, + kernel_height, + kernel_width, + H_out, + W_out, + }); + + if (use_im2col) { + q8ta_conv2d_im2col(graph, args); + } else { + q8ta_conv2d_general_impl(graph, args); + } +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.q8ta_conv2d.default, q8ta_conv2d); +} + } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2d.h similarity index 84% rename from backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.h rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2d.h index 9f60b6fd372..549ab2b25e8 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.h +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2d.h @@ -8,6 +8,8 @@ #pragma once +#include + #include namespace vkcompute { @@ -30,4 +32,6 @@ struct Q8taConv2dRouteParams final { bool should_use_q8ta_conv2d_im2col(const Q8taConv2dRouteParams& params); +void q8ta_conv2d(ComputeGraph& graph, const std::vector& args); + } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dCommon.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dCommon.cpp new file mode 100644 index 00000000000..bde1492ca18 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dCommon.cpp @@ -0,0 +1,37 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +namespace vkcompute { + +ActivationType activation_type_from_string(const std::string& activation) { + if (activation == "none") { + return ActivationType::kNone; + } else if (activation == "relu") { + return ActivationType::kRelu; + } + VK_THROW("Unknown activation type: ", activation); +} + +bool q8ta_conv2d_check_packed_dim_info(const api::PackedDimInfo& info) { + return info.packed_dim == WHCN::kChannelsDim && + info.packed_dim_block_size == 4 && + info.outer_packed_dim == WHCN::kWidthDim && + (info.outer_packed_dim_block_size == 1 || + info.outer_packed_dim_block_size == 4); +} + +bool q8ta_conv2d_check_4w4c_packed_dim_info(const api::PackedDimInfo& info) { + return info.packed_dim == WHCN::kChannelsDim && + info.packed_dim_block_size == 4 && + info.outer_packed_dim == WHCN::kWidthDim && + info.outer_packed_dim_block_size == 4; +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dCommon.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dCommon.h new file mode 100644 index 00000000000..5c9cb45d55e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dCommon.h @@ -0,0 +1,28 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +#include + +namespace vkcompute { + +enum class ActivationType : uint32_t { + kNone = 0, + kRelu = 1, +}; + +ActivationType activation_type_from_string(const std::string& activation); + +bool q8ta_conv2d_check_packed_dim_info(const api::PackedDimInfo& info); + +bool q8ta_conv2d_check_4w4c_packed_dim_info(const api::PackedDimInfo& info); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDW.cpp similarity index 97% rename from backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDW.cpp index a504d9ab3f6..b0696bb5442 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDW.cpp @@ -6,12 +6,14 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include + +#include #include #include -#include #include +#include #include #include @@ -437,7 +439,9 @@ void add_q8ta_conv2d_dw_node( // High level operator impl // -void q8ta_conv2d_dw(ComputeGraph& graph, const std::vector& args) { +void q8ta_conv2d_dw_impl( + ComputeGraph& graph, + const std::vector& args) { int32_t idx = 0; const ValueRef packed_int8_input = args.at(idx++); const ValueRef input_scale = args.at(idx++); @@ -509,7 +513,7 @@ void q8ta_conv2d_dw(ComputeGraph& graph, const std::vector& args) { } REGISTER_OPERATORS { - VK_REGISTER_OP(et_vk.q8ta_conv2d_dw.default, q8ta_conv2d_dw); + VK_REGISTER_OP(et_vk.q8ta_conv2d_dw.default, q8ta_conv2d_dw_impl); } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDW.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDW.h new file mode 100644 index 00000000000..cbff43acfb9 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDW.h @@ -0,0 +1,61 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +namespace vkcompute { + +ValueRef prepack_quantized_conv2d_dw_weight( + ComputeGraph& graph, + const QuantizationConfig& weight_quant_config, + const ValueRef weight_data, + const ValueRef kernel_size); + +void add_q8ta_conv2d_dw_node( + ComputeGraph& graph, + const ValueRef packed_int8_input, + const ValueRef input_scale, + const ValueRef input_zp, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef output_scale, + const ValueRef output_zp, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef kernel_size, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const uint32_t activation_type, + const ValueRef packed_int8_output); + +void add_conv2d_dw_q8ta_q8csw_q8to_4w4c_node( + ComputeGraph& graph, + const ValueRef packed_int8_input, + const ValueRef input_scale, + const ValueRef input_zp, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef output_scale, + const ValueRef output_zp, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef kernel_size, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const ValueRef packed_int8_output); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDirect.cpp similarity index 84% rename from backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDirect.cpp index 75f65d7a2a0..d1a3edba2de 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDirect.cpp @@ -6,43 +6,19 @@ * LICENSE file in the root directory of this source tree. */ -#include -#include +#include #include #include -#include #include +#include +#include #include #include namespace vkcompute { -ActivationType activation_type_from_string(const std::string& activation) { - if (activation == "none") { - return ActivationType::kNone; - } else if (activation == "relu") { - return ActivationType::kRelu; - } - VK_THROW("Unknown activation type: ", activation); -} - -bool q8ta_conv2d_check_packed_dim_info(const api::PackedDimInfo& info) { - return info.packed_dim == WHCN::kChannelsDim && - info.packed_dim_block_size == 4 && - info.outer_packed_dim == WHCN::kWidthDim && - (info.outer_packed_dim_block_size == 1 || - info.outer_packed_dim_block_size == 4); -} - -bool q8ta_conv2d_check_4w4c_packed_dim_info(const api::PackedDimInfo& info) { - return info.packed_dim == WHCN::kChannelsDim && - info.packed_dim_block_size == 4 && - info.outer_packed_dim == WHCN::kWidthDim && - info.outer_packed_dim_block_size == 4; -} - // // Workgroup size selection functions // @@ -160,7 +136,8 @@ ValueRef prepack_quantized_conv2d_weight( VK_CHECK_COND(weight_quant_config.nbits == 8); VK_CHECK_COND(weight_quant_config.is_symmetric); - const int32_t groups_val = graph.get_int(groups); + const int32_t groups_val = + utils::safe_downcast(graph.get_int(groups)); const int64_t OC = graph.size_at(-3, output); const int64_t IC = graph.size_at(-3, input) / groups_val; @@ -397,7 +374,7 @@ void add_q8ta_conv2d_node( // High level operator impl // -void q8ta_conv2d_general( +void q8ta_conv2d_general_impl( ComputeGraph& graph, const std::vector& args) { int32_t idx = 0; @@ -479,56 +456,8 @@ void q8ta_conv2d_general( packed_int8_output); } -void q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { - const ValueRef input = args.at(0); - const ValueRef kernel_size_ref = args.at(9); - const ValueRef groups_ref = args.at(13); - const ValueRef output = args.at(15); - - const int64_t groups = graph.extract_scalar(groups_ref); - // Valid models always carry groups >= 1; fail fast on corrupt input - // instead of dividing channel counts by zero downstream (both this - // dispatcher and q8ta_conv2d_general divide by groups). - VK_CHECK_COND(groups > 0, "q8ta_conv2d requires groups >= 1"); - const int64_t in_channels = graph.size_at(-3, input); - const int64_t in_channels_per_group = in_channels / groups; - const int64_t batch = graph.size_at(-4, input); - - const int64_t H_out = graph.size_at(-2, output); - const int64_t W_out = graph.size_at(-1, output); - const int64_t out_channels = graph.size_at(-3, output); - int64_t kernel_height; - int64_t kernel_width; - { - const auto kernel_size = graph.get_int_list(kernel_size_ref); - kernel_height = kernel_size->at(0); - kernel_width = kernel_size->at(1); - } - - const bool use_im2col = should_use_q8ta_conv2d_im2col({ - graph.device_is_mali(), - graph.can_use_int8_dot_product(), - static_cast(graph.max_buffer_numel()), - batch, - groups, - in_channels_per_group, - out_channels, - kernel_height, - kernel_width, - H_out, - W_out, - }); - - if (use_im2col) { - q8ta_conv2d_im2col(graph, args); - } else { - q8ta_conv2d_general(graph, args); - } -} - REGISTER_OPERATORS { - VK_REGISTER_OP(et_vk.q8ta_conv2d.default, q8ta_conv2d); - VK_REGISTER_OP(et_vk.q8ta_conv2d_general.default, q8ta_conv2d_general); + VK_REGISTER_OP(et_vk.q8ta_conv2d_general.default, q8ta_conv2d_general_impl); } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDirect.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDirect.h new file mode 100644 index 00000000000..e1960303e1f --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dDirect.h @@ -0,0 +1,50 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +namespace vkcompute { + +ValueRef prepack_quantized_conv2d_weight( + ComputeGraph& graph, + const QuantizationConfig& weight_quant_config, + const ValueRef weight_data, + const ValueRef input, + const ValueRef output, + const ValueRef groups, + const ValueRef kernel_size); + +void add_q8ta_conv2d_node( + ComputeGraph& graph, + const ValueRef packed_int8_input, + const ValueRef packed_int8_input_im2col, + const ValueRef input_scale, + const ValueRef input_zp, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef output_scale, + const ValueRef output_zp, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef kernel_size, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const uint32_t activation_type, + const ValueRef packed_int8_output); + +void q8ta_conv2d_general_impl( + ComputeGraph& graph, + const std::vector& args); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dPW.cpp similarity index 98% rename from backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dPW.cpp index 454a115498b..d4bb4dbab38 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dPW.cpp @@ -6,10 +6,13 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include + +#include #include #include +#include #include #include #include diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dPW.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dPW.h new file mode 100644 index 00000000000..97878ee122c --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/Q8taConv2dPW.h @@ -0,0 +1,53 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +void add_q8ta_conv2d_pw_node( + ComputeGraph& graph, + const bool use_unsigned_dot, + const ValueRef packed_int8_input, + const ValueRef input_scale, + const ValueRef input_zp, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef output_scale, + const ValueRef output_zp, + const ValueRef bias_data, + const ValueRef packed_bias, + const uint32_t activation_type, + const ValueRef packed_int8_output, + const int32_t groups = 1, + const ValueRef conv_input = kDummyValueRef, + const ValueRef kernel_size = kDummyValueRef, + const ValueRef stride = kDummyValueRef, + const ValueRef padding = kDummyValueRef, + const ValueRef dilation = kDummyValueRef, + const bool is_im2col = false, + const ValueRef stream_row_offset_ref = kDummyValueRef, + const ValueRef max_im2col_rows_ref = kDummyValueRef); + +constexpr int64_t kMaxUnsignedDotAccumulatorBytes = 33025; + +bool can_use_unsigned_pw_dot( + const vkapi::Adapter& adapter, + int64_t k_per_group); + +void q8ta_conv2d_pw_impl( + ComputeGraph& graph, + bool use_unsigned_dot, + const std::vector& args); + +void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.cpp similarity index 96% rename from backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.cpp index 062e5ea1d8d..39f9a44b476 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.cpp @@ -6,15 +6,18 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include + +#include +#include +#include #include #include -#include #include #include +#include #include -#include #include #include @@ -372,7 +375,7 @@ void q8ta_conv2d_im2col_impl( full_im2col_sizes.at(3), graph.max_buffer_numel()); if (!stream_plan.feasible) { - q8ta_conv2d_general(graph, args); + q8ta_conv2d_general_impl(graph, args); return; } VK_CHECK_COND( diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.h new file mode 100644 index 00000000000..c1520726111 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/im2col/Q8taConv2dIm2Col.h @@ -0,0 +1,71 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +inline constexpr int64_t kQ8taConv2dIm2ColScratchBudgetBytes = 16 * 1024 * 1024; +inline constexpr int64_t kQ8taConv2dMaxRowsPerTile = 65535; + +struct Q8taConv2dStreamPlan final { + int64_t aligned_out_width; + int64_t rows_per_tile; + int64_t num_tiles; + int64_t scratch_bytes; + bool feasible; +}; + +Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan( + int64_t batch, + int64_t flattened_kernel_size, + int64_t out_height, + int64_t out_width, + int64_t scratch_budget_bytes); + +// max_buffer_bytes is Adapter::max_buffer_numel(), which returns +// maxStorageBufferRange in bytes (not elements) — directly comparable with +// the byte-denominated scratch budget. +Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan_for_device( + int64_t batch, + int64_t flattened_kernel_size, + int64_t out_height, + int64_t out_width, + uint64_t max_buffer_bytes); + +std::vector calculate_q8ta_im2col_sizes( + ComputeGraph* graph, + const ValueRef& input, + const ValueRef& output, + const ValueRef& kernel_size, + const ValueRef& groups); + +void add_q8ta_im2col_node( + ComputeGraph& graph, + const ValueRef packed_int8_input, + const ValueRef kernel_size, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const ValueRef packed_int8_output, + const ValueRef packed_int8_im2col, + const int32_t zp, + const ValueRef stream_row_offset_ref, + const ValueRef max_im2col_rows_ref = kDummyValueRef); + +void q8ta_conv2d_im2col(ComputeGraph& graph, const std::vector& args); + +void q8ta_conv2d_im2col_impl( + ComputeGraph& graph, + bool use_unsigned_dot, + const std::vector& args); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dTransposed.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/transposed/Q8taConv2dTransposed.cpp similarity index 95% rename from backends/vulkan/runtime/graph/ops/impl/Q8taConv2dTransposed.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/transposed/Q8taConv2dTransposed.cpp index 91d745fb1fa..ff5cac10a0b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dTransposed.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/transposed/Q8taConv2dTransposed.cpp @@ -6,13 +6,16 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include + +#include +#include #include #include -#include #include +#include #include #include @@ -228,7 +231,7 @@ void add_q8ta_conv2d_transposed_node( resize_q8ta_conv2d_transposed_node)); } -void q8ta_conv2d_transposed( +void q8ta_conv2d_transposed_impl( ComputeGraph& graph, const std::vector& args) { int32_t idx = 0; @@ -313,7 +316,8 @@ void q8ta_conv2d_transposed( } REGISTER_OPERATORS { - VK_REGISTER_OP(et_vk.q8ta_conv2d_transposed.default, q8ta_conv2d_transposed); + VK_REGISTER_OP( + et_vk.q8ta_conv2d_transposed.default, q8ta_conv2d_transposed_impl); } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/transposed/Q8taConv2dTransposed.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/transposed/Q8taConv2dTransposed.h new file mode 100644 index 00000000000..92607284b00 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/q8ta/transposed/Q8taConv2dTransposed.h @@ -0,0 +1,19 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +void q8ta_conv2d_transposed_impl( + ComputeGraph& graph, + const std::vector& args); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.cpp b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/quantized_legacy/QuantizedConvolution.cpp similarity index 98% rename from backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.cpp rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/quantized_legacy/QuantizedConvolution.cpp index 30415cd4f08..595395b6f82 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/quantized_legacy/QuantizedConvolution.cpp @@ -9,13 +9,15 @@ #include #include -#include -#include #include #include -#include #include #include +#include +#include +#include +#include +#include #include namespace vkcompute { @@ -1341,7 +1343,7 @@ void conv2d_q8ta_q8csw_q8to( // Test operators // -void test_conv2d_q8ta_q8csw_q8to( +void test_conv2d_q8ta_q8csw_q8to_impl( ComputeGraph& graph, const std::vector& args) { int32_t idx = 0; @@ -1407,7 +1409,8 @@ REGISTER_OPERATORS { VK_REGISTER_OP(et_vk.conv2d_q8ta_q8csw.default, conv2d_q8ta_q8csw); VK_REGISTER_OP(et_vk.conv2d_q8csw.default, conv2d_q8csw); VK_REGISTER_OP( - etvk.test_conv2d_q8ta_q8csw_q8to.default, test_conv2d_q8ta_q8csw_q8to); + etvk.test_conv2d_q8ta_q8csw_q8to.default, + test_conv2d_q8ta_q8csw_q8to_impl); VK_REGISTER_OP(et_vk.conv2d_q8ta_q8csw_q8to.default, conv2d_q8ta_q8csw_q8to); VK_REGISTER_OP( et_vk.conv2d_q8ta_q8csw_q8to_dw.default, conv2d_q8ta_q8csw_q8to); diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.h b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/quantized_legacy/QuantizedConvolution.h similarity index 78% rename from backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.h rename to backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/quantized_legacy/QuantizedConvolution.h index 8913e639f0e..6bc453d0d27 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.h +++ b/backends/vulkan/runtime/graph/ops/impl/convolution/conv2d/quantized_legacy/QuantizedConvolution.h @@ -8,7 +8,7 @@ #pragma once -#include +#include namespace vkcompute { diff --git a/backends/vulkan/test/custom_ops/impl/TestConv2d.cpp b/backends/vulkan/test/custom_ops/impl/TestConv2d.cpp index 343f416c6a1..bad30507ac1 100644 --- a/backends/vulkan/test/custom_ops/impl/TestConv2d.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestConv2d.cpp @@ -9,8 +9,8 @@ #include #include -#include -#include +#include +#include #include @@ -32,8 +32,8 @@ void test_conv2d(ComputeGraph& graph, const std::vector& args) { // impl_selector grammar: // "" -> aten.convolution.default (heuristic-routed: // should_use_conv2d_im2col() picks direct vs im2col) - // "direct" -> add_conv2d_node(force_direct=true): forces the direct - // sliding-window path, bypassing the routing heuristic + // "direct" -> conv2d_direct_impl: forces the direct sliding-window + // path, bypassing the routing heuristic // "im2col" -> et_vk.conv2d_gemm.default, auto im2col storage // "im2col_buffer"-> im2col/GEMM, force buffer im2col intermediate // "im2col_tex2d" -> im2col/GEMM, force texture2d im2col intermediate @@ -94,11 +94,10 @@ void test_conv2d(ComputeGraph& graph, const std::vector& args) { // The "direct" selector must reach the exact direct sliding-window dispatch // the heuristic would otherwise pick. The registered op can only route via - // the heuristic, so call add_conv2d_node directly with force_direct=true to - // bypass it (mirroring how the forced-storage variants call - // conv2d_gemm_impl). + // the heuristic, so call conv2d_direct_impl directly to bypass it (mirroring + // how the forced-storage variants call conv2d_gemm_impl). if (impl_selector == "direct") { - add_conv2d_node( + conv2d_direct_impl( graph, input, weight, @@ -109,11 +108,10 @@ void test_conv2d(ComputeGraph& graph, const std::vector& args) { transposed, output_padding, groups, - /*out_min=*/kDummyValueRef, - /*out_max=*/kDummyValueRef, out, /*clamp_out=*/false, - /*force_direct=*/true); + /*out_min_val=*/0.0f, + /*out_max_val=*/0.0f); return; } diff --git a/backends/vulkan/test/custom_ops/impl/TestConv2dDw.cpp b/backends/vulkan/test/custom_ops/impl/TestConv2dDw.cpp index 85c2308b11f..f27a596eca3 100644 --- a/backends/vulkan/test/custom_ops/impl/TestConv2dDw.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestConv2dDw.cpp @@ -9,8 +9,9 @@ #include #include -#include #include +#include +#include #include #include diff --git a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp index e4af2c56b9d..4da454040e7 100644 --- a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp @@ -10,8 +10,9 @@ #include #include -#include #include +#include +#include namespace vkcompute { diff --git a/backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp b/backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp index 63d154b170f..c660430b579 100644 --- a/backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp +++ b/backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp @@ -4,7 +4,7 @@ // This source code is licensed under the BSD-style license found in the // LICENSE file in the root directory of this source tree. -#include +#include #include diff --git a/backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp b/backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp index 6f634a83c27..c1207cb1923 100644 --- a/backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp +++ b/backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp @@ -4,7 +4,7 @@ // This source code is licensed under the BSD-style license found in the // LICENSE file in the root directory of this source tree. -#include +#include #include diff --git a/backends/vulkan/test/custom_ops/test_conv2d.cpp b/backends/vulkan/test/custom_ops/test_conv2d.cpp index 6bec62a4e61..2f21f0532fe 100644 --- a/backends/vulkan/test/custom_ops/test_conv2d.cpp +++ b/backends/vulkan/test/custom_ops/test_conv2d.cpp @@ -526,7 +526,7 @@ static std::vector generate_conv2d_test_cases() { // "" -> heuristic-routed (should_use_conv2d_im2col picks direct on // Adreno for small c_out, im2col on Mali) // "im2col" -> forced im2col/GEMM path - // "direct" -> forced direct sliding-window path (force_direct=true) + // "direct" -> forced direct sliding-window path (conv2d_direct_impl) // Including "direct" guarantees the direct shader gets reference-checked on // BOTH devices — without it, Mali would always route "" to im2col and never // exercise the direct path. diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp index 90fd4e1616b..17045225a2c 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp @@ -12,7 +12,8 @@ #include #include -#include +#include +#include #include #include diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp index 1003c0540ce..5a773f43319 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp @@ -9,7 +9,7 @@ #include #include -#include +#include #include #include diff --git a/backends/vulkan/test/op_tests/conv2d_gemm_dynamic_test.cpp b/backends/vulkan/test/op_tests/conv2d_gemm_dynamic_test.cpp index e5a7216bfe7..f993c453fb0 100644 --- a/backends/vulkan/test/op_tests/conv2d_gemm_dynamic_test.cpp +++ b/backends/vulkan/test/op_tests/conv2d_gemm_dynamic_test.cpp @@ -10,7 +10,7 @@ #include #include -#include +#include #include #include