diff --git a/src/targets/gpu/jit/winograd_conv.cpp b/src/targets/gpu/jit/winograd_conv.cpp index d8aca2b23df..dd3c70c2b9d 100644 --- a/src/targets/gpu/jit/winograd_conv.cpp +++ b/src/targets/gpu/jit/winograd_conv.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -47,26 +48,49 @@ using namespace migraphx::gpu::gen; // NOLINT static std::string post_input_cast(const module& pm) { // Pointwise submodule params are named x0, x1, ...; x0 is arg 0, which the - // fusion wires to the winograd conv output. + // fusion wires to the winograd conv output. Its type is the conv's natural + // output precision (half for the fp16 kernel, float for the fp32 kernel), + // which is the base type the post-op computes at. auto x0 = pm.get_parameter("x0"); if(x0 == pm.end()) return "half"; + std::string base = shape::cpp_type(x0->get_shape().type()); // Only treat a *leading* convert as the post-op's compute type, i.e. when // the conv result feeds exactly one op and that op is a convert to a type - // wider than the conv's half output. A convert that appears later (after - // an add/activation/etc.) must still run at half precision first. + // wider than the conv output. A convert that appears later (after an + // add/activation/etc.) must still run at conv precision first. const auto& users = x0->outputs(); if(users.size() != 1) - return "half"; + return base; auto user = users.front(); if(user->name() != "convert") - return "half"; + return base; auto t = user->get_shape().type(); - if(shape{t}.type_size() <= shape{shape::half_type}.type_size()) - return "half"; + if(shape{t}.type_size() <= x0->get_shape().type_size()) + return base; return shape::cpp_type(t); } +// Which host transform prefuse baked into the fp32 winograd weight literal. The +// choice is encoded in the literal's shape; classify it so the kernel template +// args are named instead of re-derived from magic dimension sizes. +enum class winograd_fp32_weight_layout +{ + full_u, // [4, 4, K, C] (u, v, k, c; C innermost) -- NCHW + full_u_vinner, // [4, K, C, 4] (u, k, c, v; v innermost) -- NHWC, coalesced load + sstore, // [3, 4, K, C] (i, v, k, c) -- v-half g*G^T, 25% smaller +}; + +static winograd_fp32_weight_layout winograd_fp32_weight_layout_of(const shape& w) +{ + const auto& l = w.lens(); + assert(l.size() == 4); + if(l[0] == 3) + return winograd_fp32_weight_layout::sstore; + return l[1] == 4 ? winograd_fp32_weight_layout::full_u + : winograd_fp32_weight_layout::full_u_vinner; +} + // NOLINTNEXTLINE static const char* const winograd_conv_kernel = R"__migraphx__( #include @@ -96,12 +120,135 @@ MIGRAPHX_GLOBAL void ${kernel}(${params}) )__migraphx__"; +// fp32 FMA/DPP kernel (winograd_conv_f23_fp32). Configured by nw (waves), +// ko (output channels per lane) and tiles (winograd tiles per quad). +// NOLINTNEXTLINE +static const char* const winograd_conv_fp32_kernel = R"__migraphx__( +#include +#include +#include +#include +#include + +namespace migraphx { + +${preamble} + +extern "C" { + +MIGRAPHX_GLOBAL void ${kernel}(${params}) +{ + transform_args(make_tensors(), rotate_last())(${args})( + [](auto output, auto x, auto u, auto... inputs) { + winograd_conv_f23_fp32<${nw}, ${ko}, ${tiles}, ${sk}, ${pipe}, ${cu}, ${sstore}, ${nhwc}, ${vinner}, ${conv_cast}>( + ${post}, output, x, u, inputs...); + }); +} + +} + +} // namespace migraphx + +)__migraphx__"; + struct winograd_conv_compiler : compiler { std::vector names() const { return {"gpu::winograd_conv", "winograd_conv"}; } + // fp32 FMA/DPP kernel: lane%4 = winograd v-column, quad = TILES tiles, + // each wave = 8 quads. Launch covers (tile groups) x (output-channel blocks), + // where a tile group is one workgroup's worth of quads. + operation compile_op_fp32(context& ctx, const std::vector& inputs, const value& v) const + { + hip_compile_options options; + const auto& out_s = inputs.back(); + options.inputs = inputs; + options.output = out_s; + options.virtual_inputs = inputs; + options.kernel_name = v.get("kernel", std::string{"winograd_conv_fp32_kernel"}); + + const auto nw = v.get("nw", std::size_t{4}); + const auto ko = v.get("ko", std::size_t{8}); + const auto tiles = v.get("tiles", std::size_t{2}); + // sk = within-WG channel-split factor (must divide nw). sk>1 has nw/sk + // NT-groups whose sk waves split the channel contraction. + const auto sk = v.get("sk", std::size_t{1}); + // pipe = software-pipeline the input transform into the FMA loop (true) vs + // the simple transform-then-FMA path (false). pipe costs a second live + // v_reg, so the tuner only offers it on small (non-spilling) solutions. + const bool pipe = v.get("pipe", false); + // cu = channel-unroll (weight load width): 4=b128, 2=b64, 1=b32. Smaller + // cu shrinks the pipelined double-buffer (finer-grained pipeline) at the + // cost of more, narrower weight loads. + const auto cu = v.get("cu", std::size_t{4}); + + if(nw == 0 or ko == 0 or tiles == 0 or sk == 0) + MIGRAPHX_THROW("winograd_conv_fp32: nw/ko/tiles/sk must be non-zero"); + if(sk > nw or (nw % sk) != 0) + MIGRAPHX_THROW("winograd_conv_fp32: sk must be a non-zero divisor of nw"); + if(cu != 1 and cu != 2 and cu != 4) + MIGRAPHX_THROW("winograd_conv_fp32: cu must be 1, 2, or 4"); + if(ko * tiles > 32) + MIGRAPHX_THROW("winograd_conv_fp32: ko*tiles must be <= 32"); + // - sstore: v-half-transformed S=[3,4,K,C]; the kernel finishes the + // register-only G u-transform, cutting weight loads+bytes to 3/4 (for + // weight-bandwidth-bound shapes). + // - vinner: full U laid out v-innermost so the 4 v_col lanes coalesce (the + // gated NHWC configs); else the plain C-innermost full U. + const auto wlayout = winograd_fp32_weight_layout_of(inputs.at(1)); + const bool sstore = wlayout == winograd_fp32_weight_layout::sstore; + const bool vinner = wlayout == winograd_fp32_weight_layout::full_u_vinner; + // NHWC: the conv input has its channel axis innermost (stride 1), so the + // kernel loads CU contiguous channels with one b128 instead of CU b32. + const bool nhwc = inputs.front().strides().at(1) == 1; + + // Only nw/sk NT-groups' worth of distinct tiles are covered per WG. + const std::size_t quads_per_wg = 8 * (nw / sk); + const std::size_t block_size = nw * 32; + + const auto& out_lens = out_s.lens(); + assert(out_lens.size() == 4); + const auto n = out_lens[0]; + const auto out_c = out_lens[1]; + const auto out_h = out_lens[2]; + const auto out_w = out_lens[3]; + const auto tiles_h = (out_h + 1) / 2; + const auto tiles_w = (out_w + 1) / 2; + const auto nt_total = n * tiles_h * tiles_w; + + // One workgroup per (tile_group, k_block); its nw/sk NT-groups cover a + // contiguous run of tiles for that k_block. + const auto k_blocks = integer_divide_ceil(out_c, ko); + const auto quad_groups = integer_divide_ceil(nt_total, tiles); + const auto tile_blocks = integer_divide_ceil(quad_groups, quads_per_wg); + const auto num_blocks = k_blocks * tile_blocks; + + options.set_launch_params(v, num_blocks * block_size, block_size); + + auto src = interpolate_string(winograd_conv_fp32_kernel, + {{"kernel", options.kernel_name}, + {"params", enum_params(inputs.size(), "void * private_p")}, + {"args", enum_params(inputs.size(), "private_p")}, + {"nw", std::to_string(nw)}, + {"ko", std::to_string(ko)}, + {"tiles", std::to_string(tiles)}, + {"sk", std::to_string(sk)}, + {"pipe", pipe ? "true" : "false"}, + {"cu", std::to_string(cu)}, + {"sstore", sstore ? "true" : "false"}, + {"nhwc", nhwc ? "true" : "false"}, + {"vinner", vinner ? "true" : "false"}, + {"post", v.get("post", std::string{"op::id{}"})}, + {"conv_cast", v.get("conv_cast", std::string{"float"})}, + {"preamble", v.get("preamble", std::string{})}}); + + return compile_hip_code_object(ctx, src, options); + } + operation compile_op(context& ctx, const std::vector& inputs, const value& v) const { + if(inputs.front().type() == shape::float_type) + return compile_op_fp32(ctx, inputs, v); hip_compile_options options; const auto& out_s = inputs.back(); options.inputs = inputs; @@ -142,8 +289,8 @@ struct winograd_conv_compiler : compiler const auto tiles_w = (out_w + 1) / 2; const auto nt_total = n * tiles_h * tiles_w; - const auto k_wg_blocks = (out_c + bk_wg - 1) / bk_wg; - const auto t_blocks = (nt_total + bt - 1) / bt; + const auto k_wg_blocks = integer_divide_ceil(out_c, bk_wg); + const auto t_blocks = integer_divide_ceil(nt_total, bt); const auto num_blocks = k_wg_blocks * t_blocks; options.set_launch_params(v, num_blocks * block_size, block_size); @@ -189,6 +336,69 @@ struct winograd_conv_compiler : compiler auto shapes = to_shapes(ins->inputs()); tc.problem = to_value(shapes); + // fp32 FMA/DPP configs: nw (waves), ko (out-channels/lane), tiles + // (winograd tiles/quad). ko*tiles is kept <= 32 (accumulators/lane = + // 4*ko*tiles <= 128) to bound register spilling. + if(shapes.front().type() == shape::float_type) + { + // Larger ko amortizes the (out-channel-independent) input transform + // over more output channels -- important when out_c is large; larger + // tiles amortizes the shared weight load -- good for large in_c and + // small out_c. + tc.solutions.push_back({{"nw", 2}, {"ko", 8}, {"tiles", 2}}); + tc.solutions.push_back({{"nw", 4}, {"ko", 8}, {"tiles", 2}}); + tc.solutions.push_back({{"nw", 8}, {"ko", 8}, {"tiles", 2}}); + tc.solutions.push_back({{"nw", 4}, {"ko", 8}, {"tiles", 1}}); + tc.solutions.push_back({{"nw", 8}, {"ko", 8}, {"tiles", 1}}); + tc.solutions.push_back({{"nw", 4}, {"ko", 8}, {"tiles", 4}}); + tc.solutions.push_back({{"nw", 2}, {"ko", 16}, {"tiles", 1}}); + tc.solutions.push_back({{"nw", 4}, {"ko", 16}, {"tiles", 1}}); + tc.solutions.push_back({{"nw", 4}, {"ko", 16}, {"tiles", 2}}); + tc.solutions.push_back({{"nw", 2}, {"ko", 16}, {"tiles", 2}}); + tc.solutions.push_back({{"nw", 4}, {"ko", 32}, {"tiles", 1}}); + // pipe=1: software-pipeline the input transform into the FMA loop so + // the DPP transform ops interleave with the FMAs instead of clustering + // ahead (helps FMA-throughput-bound shapes -- large out_c x spatial). + // The pipeline needs a second live v_reg, which spills for tiles>=4 or + // ko>=16, so only the small ko=8/tiles<=2 solutions are offered; the + // tuner keeps pipe=1 only where it beats the simple path. + tc.solutions.push_back({{"nw", 2}, {"ko", 8}, {"tiles", 2}, {"pipe", true}}); + tc.solutions.push_back({{"nw", 4}, {"ko", 8}, {"tiles", 2}, {"pipe", true}}); + tc.solutions.push_back({{"nw", 8}, {"ko", 8}, {"tiles", 2}, {"pipe", true}}); + tc.solutions.push_back({{"nw", 4}, {"ko", 8}, {"tiles", 1}, {"pipe", true}}); + tc.solutions.push_back({{"nw", 8}, {"ko", 8}, {"tiles", 1}, {"pipe", true}}); + // Finer-grained pipeline: cu=2 (b64 weight load) halves the pipelined + // double-buffer, so ko can go to 16 without spilling -- amortizing the + // (non-dual-issue) DPP transforms over more FMAs. Costs more, narrower + // weight loads; the tuner keeps it only where the trade pays off. + tc.solutions.push_back({{"nw", 4}, {"ko", 8}, {"tiles", 2}, {"pipe", true}, {"cu", 2}}); + tc.solutions.push_back( + {{"nw", 4}, {"ko", 16}, {"tiles", 1}, {"pipe", true}, {"cu", 2}}); + tc.solutions.push_back( + {{"nw", 2}, {"ko", 16}, {"tiles", 2}, {"pipe", true}, {"cu", 2}}); + tc.solutions.push_back( + {{"nw", 4}, {"ko", 16}, {"tiles", 2}, {"pipe", true}, {"cu", 2}}); + // Channel-split (sk>1): nw/sk NT-groups whose sk waves split the + // channel contraction and reduce partial M through LDS. Helps shapes + // with few tiles + many channels (small spatial, large in_c/out_c), + // where the plain path leaves waves idle. LDS = nw*32*4*tiles*ko + // floats, so keep tiles=1, ko<=8. Only OFFERED when tiles are scarce + // (small nt_total): on tile-rich shapes the plain path already fills + // the machine, and offering sk configs there only adds tuner noise. + const auto& out_lens = shapes.back().lens(); + assert(out_lens.size() == 4); + const auto nt_total = out_lens[0] * ((out_lens[2] + 1) / 2) * ((out_lens[3] + 1) / 2); + if(nt_total < 256) + { + tc.solutions.push_back({{"nw", 4}, {"ko", 8}, {"tiles", 1}, {"sk", 2}}); + tc.solutions.push_back({{"nw", 4}, {"ko", 8}, {"tiles", 1}, {"sk", 4}}); + tc.solutions.push_back({{"nw", 8}, {"ko", 8}, {"tiles", 1}, {"sk", 2}}); + tc.solutions.push_back({{"nw", 8}, {"ko", 8}, {"tiles", 1}, {"sk", 4}}); + tc.solutions.push_back({{"nw", 8}, {"ko", 8}, {"tiles", 1}, {"sk", 8}}); + } + return tc; + } + // Wave32 WMMA configs. CB must be a multiple of WMMA K (16). KW is // the number of K_blocks (BK=16 each) processed per workgroup. // V values live in per-lane registers, so LDS budget diff --git a/src/targets/gpu/kernels/include/migraphx/kernels/buffer_load.hpp b/src/targets/gpu/kernels/include/migraphx/kernels/buffer_load.hpp new file mode 100644 index 00000000000..49fd009ebea --- /dev/null +++ b/src/targets/gpu/kernels/include/migraphx/kernels/buffer_load.hpp @@ -0,0 +1,78 @@ +/* + * The MIT License (MIT) + * + * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ +#ifndef MIGRAPHX_GUARD_KERNELS_BUFFER_LOAD_HPP +#define MIGRAPHX_GUARD_KERNELS_BUFFER_LOAD_HPP + +#include +#include +#include + +namespace migraphx { + +// gfx12 buffer-resource word 3 (from composable_kernel): makes raw buffer loads +// return 0 for out-of-range byte offsets, so bounds/halo checks collapse to an +// offset select against a sentinel instead of a per-load branch. +constexpr uint32_t oob_buffer_rsrc_word3 = 0x31004000; + +// Build an out-of-bounds-tolerant buffer descriptor for a read-only pointer (the +// resource is only ever loaded from; the const_cast is required by the builtin's +// non-const pointer parameter). +template +__device__ inline __amdgpu_buffer_rsrc_t make_oob_buffer_rsrc(const T* p, uint32_t byte_count) +{ + auto* base = const_cast(p); // NOLINT(cppcoreguidelines-pro-type-const-cast) + return __builtin_amdgcn_make_buffer_rsrc(base, 0, byte_count, oob_buffer_rsrc_word3); +} + +// Raw buffer load of N contiguous T (N*sizeof(T) must be 2/4/8/16 bytes -> +// b16/b32/b64/b128; gfx12 tolerates 4-byte alignment). OOB bytes read as 0. +template +__device__ inline vec buffer_load_vec(__amdgpu_buffer_rsrc_t rsrc, int byte_offset) +{ + constexpr index_int bytes = N * sizeof(T); + static_assert(bytes == 2 or bytes == 4 or bytes == 8 or bytes == 16, + "buffer_load_vec width must be 2, 4, 8, or 16 bytes"); + if constexpr(bytes == 16) + return bit_cast>(__builtin_amdgcn_raw_buffer_load_b128(rsrc, byte_offset, 0, 0)); + else if constexpr(bytes == 8) + return bit_cast>(__builtin_amdgcn_raw_buffer_load_b64(rsrc, byte_offset, 0, 0)); + else if constexpr(bytes == 4) + return bit_cast>(__builtin_amdgcn_raw_buffer_load_b32(rsrc, byte_offset, 0, 0)); + else + return bit_cast>(__builtin_amdgcn_raw_buffer_load_b16(rsrc, byte_offset, 0, 0)); +} + +// Raw buffer load of a single T (2- or 4-byte element). OOB reads as 0. +template +__device__ inline T buffer_load(__amdgpu_buffer_rsrc_t rsrc, int byte_offset) +{ + static_assert(sizeof(T) == 2 or sizeof(T) == 4, "buffer_load element must be 2 or 4 bytes"); + if constexpr(sizeof(T) == 2) + return bit_cast(__builtin_amdgcn_raw_buffer_load_b16(rsrc, byte_offset, 0, 0)); + else + return bit_cast(__builtin_amdgcn_raw_buffer_load_b32(rsrc, byte_offset, 0, 0)); +} + +} // namespace migraphx +#endif // MIGRAPHX_GUARD_KERNELS_BUFFER_LOAD_HPP diff --git a/src/targets/gpu/kernels/include/migraphx/kernels/winograd_conv.hpp b/src/targets/gpu/kernels/include/migraphx/kernels/winograd_conv.hpp index 3485a798d6c..201039450c3 100644 --- a/src/targets/gpu/kernels/include/migraphx/kernels/winograd_conv.hpp +++ b/src/targets/gpu/kernels/include/migraphx/kernels/winograd_conv.hpp @@ -26,6 +26,7 @@ #include #include +#include #include #include #include @@ -36,11 +37,6 @@ namespace migraphx { -// gfx12 buffer-resource word 3 constant (from composable_kernel). -// Setting this in the SRD makes raw_buffer_load_* return 0 for OOB accesses, -// which lets us drop the per-element bounds checks in the input transform. -constexpr uint32_t buffer_rsrc_3rd_dword_gfx12 = 0x31004000; - // Quad of WMMAs in a single inline-asm block. Forces the compiler to issue // them back-to-back (each is 8-cycle wait state but to a DIFFERENT // accumulator, so the next can issue ~1 cycle later). The compiler is then @@ -124,37 +120,6 @@ __device__ inline void wmma_octet_asm(vec a0, "v"(b7)); } -__device__ inline auto make_input_buffer_rsrc(const half* p, uint32_t byte_count) -{ - // The builtin takes a non-const base pointer, but the input tensor is const. - auto* base = const_cast(p); // NOLINT(cppcoreguidelines-pro-type-const-cast) - return __builtin_amdgcn_make_buffer_rsrc(base, 0, byte_count, buffer_rsrc_3rd_dword_gfx12); -} - -// Lane-indexed raw buffer load of a single fp16. OOB returns 0. -__device__ inline half buffer_load_half(__amdgpu_buffer_rsrc_t rsrc, int byte_offset) -{ - uint16_t v = __builtin_amdgcn_raw_buffer_load_b16(rsrc, byte_offset, 0, 0); - return bit_cast(v); -} - -// Lane-indexed raw buffer load of 4 fp16 = 8 bytes. OOB bytes return 0. -// Caller is responsible for alignment (byte_offset divisible by 4 to avoid -// faulting; gfx12 buffer loads tolerate 4-byte alignment for b64). -__device__ inline vec buffer_load_half4(__amdgpu_buffer_rsrc_t rsrc, int byte_offset) -{ - auto v = __builtin_amdgcn_raw_buffer_load_b64(rsrc, byte_offset, 0, 0); - return bit_cast>(v); -} - -// Lane-indexed raw buffer load of 8 fp16 = 16 bytes (b128). OOB bytes return 0. -// Used by the NHWC input load, where 8 contiguous channels are read at once. -__device__ inline vec buffer_load_half8(__amdgpu_buffer_rsrc_t rsrc, int byte_offset) -{ - auto v = __builtin_amdgcn_raw_buffer_load_b128(rsrc, byte_offset, 0, 0); - return bit_cast>(v); -} - // F(2x2, 3x3) Winograd transforms used inline by the WMMA path. // B^T (input): | 1 0 -1 0 | A^T (output): | 1 1 1 0 | // | 0 1 1 0 | | 0 1 -1 -1 | @@ -363,12 +328,12 @@ __device__ void winograd_conv_f23_wmma(F f, Output output, Input x, Weights u, I const auto x_sh = x_shape.strides; const auto* x_data = x.data(); const uint32_t x_byte_count = static_cast(x_shape.element_space()) * sizeof(half); - auto x_rsrc = make_input_buffer_rsrc(x_data, x_byte_count); + auto x_rsrc = make_oob_buffer_rsrc(x_data, x_byte_count); const auto* u_data = u.data(); const uint32_t u_byte_count = static_cast(u.get_shape().element_space()) * sizeof(half); - auto u_rsrc = make_input_buffer_rsrc(u_data, u_byte_count); + auto u_rsrc = make_oob_buffer_rsrc(u_data, u_byte_count); // U layout: [4 or 3, 3, K, C] -- strides for byte offset computation. const auto u_sh = u.get_shape().strides; @@ -473,7 +438,7 @@ __device__ void winograd_conv_f23_wmma(F f, Output output, Input x, Weights u, I constexpr int j = j_val; const bool ok = active and hi[i] and wj[j]; const int32_t off = ok ? base_off + i * sh_b + j * sw_b : oob_byte; - auto v8 = buffer_load_half8(x_rsrc, off); + auto v8 = buffer_load_vec(x_rsrc, off); // Last channel block may be partial: zero c >= C, which // would otherwise read the next pixel's channels. if(c_partial) @@ -510,7 +475,7 @@ __device__ void winograd_conv_f23_wmma(F f, Output output, Input x, Weights u, I { repeat_c<4>([&](auto i) { const int32_t row_off = hi[i] ? off + static_cast(i) * sh_b : oob_byte; - auto row = buffer_load_half4(x_rsrc, row_off); + auto row = buffer_load_vec(x_rsrc, row_off); d[i * 4 + 0] = wj[0] ? row.x : hzero; d[i * 4 + 1] = wj[1] ? row.y : hzero; d[i * 4 + 2] = wj[2] ? row.z : hzero; @@ -525,7 +490,7 @@ __device__ void winograd_conv_f23_wmma(F f, Output output, Input x, Weights u, I (hi[i] and wj[j]) ? off + static_cast(i) * sh_b + static_cast(j) * sw_b : oob_byte; - d[i * 4 + j] = buffer_load_half(x_rsrc, e_off); + d[i * 4 + j] = buffer_load(x_rsrc, e_off); }); }); } diff --git a/src/targets/gpu/kernels/include/migraphx/kernels/winograd_conv_fp32.hpp b/src/targets/gpu/kernels/include/migraphx/kernels/winograd_conv_fp32.hpp new file mode 100644 index 00000000000..4ff3a79f058 --- /dev/null +++ b/src/targets/gpu/kernels/include/migraphx/kernels/winograd_conv_fp32.hpp @@ -0,0 +1,640 @@ +/* + * The MIT License (MIT) + * + * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ +#ifndef MIGRAPHX_GUARD_KERNELS_WINOGRAD_CONV_FP32_HPP +#define MIGRAPHX_GUARD_KERNELS_WINOGRAD_CONV_FP32_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace migraphx { + +// FMA + DPP Winograd F(2x2, 3x3) for fp32 on gfx12 (RDNA4, wave32), modeled on +// MIOpen's Conv_Winograd_v40_6_0_gfx12_fp32_f2x3_stride1 (an FMA/DPP kernel, +// NOT a matrix-core/WMMA kernel). +// +// The winograd 4x4 input tile is laid out so that lane%4 = the winograd +// "v" column (the tile's W axis, positions 0..3) and the 4 winograd "u" rows +// (the tile's H axis) live in registers. With this mapping the winograd +// elementwise product M[u,v] = sum_c U_c[u,v] * V_c[u,v] is entirely +// lane-local (each lane owns one v column), so the channel contraction is a +// plain per-lane FMA accumulation with no cross-lane traffic. Only the input +// and output transforms need to cross the v axis, and those use intra-quad +// cross-lane shuffles (DPP for the input transform, ds_swizzle/bpermute for the +// output). The base contraction uses no shared memory (LDS); only the optional +// SK>1 channel-split reduce allocates any. +// +// Transforms (canonical Lavin-Gray F(2,3)): +// B^T = | 1 0 -1 0 | A^T = | 1 1 1 0 | G = | 1 0 0 | +// | 0 1 1 0 | | 0 1 -1 -1 | | .5 .5 .5 | +// | 0 -1 1 0 | | .5 -.5 .5 | +// | 0 1 0 -1 | | 0 0 1 | +// U = G g G^T is precomputed on the host as a [4,4,K,C] literal (u,v,k,c). +// The v-axis DPP butterfly (quad_perm:[2,2,1,1]) can only realize the input +// transform's v=3 column with a sign flip (d3-d1 instead of d1-d3); the host +// weight bakes a matching negation into U[:,3,:,:] so the product is exact. +// +// Input/weight reads use the shared gfx12 OOB buffer loads (buffer_load.hpp): +// make_oob_buffer_rsrc, buffer_load (b32), and buffer_load_vec +// (b128/b64/b32) for a CU-wide channel-unrolled load. + +// Input transform, v axis (across the 4 lanes of a quad). Given this lane's raw +// datum d for one tile row, returns P = B^T applied along the v (W) axis. +// lane0: d0 - d2 lane1: d1 + d2 lane2: d2 - d1 lane3: d3 - d1 +// The quad_perm:[2,2,1,1] shuffle broadcasts lane2 to {0,1} and lane1 to {2,3} -- +// the pair of neighbours the 4-point B^T needs -- and shuf_sign encodes the +/- +// on that neighbour (self coefficient is always +1); lane3 uses -1, which yields +// the sign-variant d3-d1 that the host weight compensates for. +__device__ inline float wino_f23_bt_v(float d, float shuf_sign) +{ + // Fused butterfly: acc = shuf_sign*dpp(d) + d in ONE v_fmac_f32_dpp. The + // compiler cannot emit this: GCNDPPCombine has an explicit TODO that discards + // MAC/FMA (the fmac DPP form has no "old" operand slot), so the intrinsic + // dpp_mov lowers to mov_dpp + cndmask + add (3 VALU) instead. The hand-written + // fused op is 1 VALU. The asm is deliberately non-volatile so it stays + // schedulable -- the surrounding input loads still software-pipeline, which a + // volatile block would serialize. quad_perm only sources in-quad lanes, so + // bound_ctrl:1 (required for the fused encoding) changes no result. + float acc = d; + asm("v_fmac_f32_dpp %[acc], %[d], %[sign] quad_perm:[2,2,1,1] row_mask:0xf " + "bank_mask:0xf bound_ctrl:1" + : [acc] "+v"(acc) + : [d] "v"(d), [sign] "v"(shuf_sign)); + return acc; +} + +// Input transform, u axis (across the 4 registers P[0..3]). B^T along H. +__device__ inline array wino_f23_bt_u(const array& p) +{ + return {p[0] - p[2], p[1] + p[2], p[2] - p[1], p[1] - p[3]}; +} + +// FMA + DPP Winograd F(2x2, 3x3) kernel. +// NW : waves per workgroup (each wave = 32 lanes = 8 quads = 8*TILES tiles). +// KO : output channels held per lane (register K tile). +// TILES : winograd tiles processed per quad (amortizes the weight load). +// SK : within-workgroup channel-split factor. The NW waves form NW/SK +// NT-groups; the SK waves of a group cover the SAME tiles but split +// the channel contraction (each does 1/SK of the channels), then +// reduce their partial M accumulators through LDS. SK=1 is the plain +// no-split path (no LDS). SK>1 fills otherwise-idle waves and cuts +// per-wave input traffic on shapes with few tiles + many channels. +// PIPE : false = simple transform-then-FMA per channel block; true = software- +// pipeline the next block's input transform into the current block's loop so +// the (non-dual-issue) DPP transform ops interleave with the dual-issue +// FMAs instead of clustering ahead of them. PIPE=1 costs a second live +// v_reg -- a tuner-selected option that wins on FMA-throughput-bound +// shapes; gated to small solutions since it spills for large TILES/KO. +// CU : channel-unroll (1/2/4) = weight load width (b32/b64/b128). CU=4 reads +// 4 channels per (u,k) with one b128. A smaller CU halves/quarters the +// pipelined double-buffer (v_cur+v_next) and the weight vector, trading +// narrower weight loads for higher occupancy -- so the pipeline can run +// at a higher KO without spilling. Tuner-selected with PIPE. +// +// The channel contraction is unrolled by CU so the shared weight can be read with +// one vector load per (u,k) covering CU channels, and so the load latency of the +// per-channel input transform is pipelined across CU channels. +// +// PostInput / F / Inputs...: fused pointwise post-op, same contract as the +// fp16 kernel -- f(cast(y), inputs[idx]...) is applied at each output position, +// collapsing to a plain cast when F = op::id{} and Inputs... is empty. +// NOLINTBEGIN(readability-function-size): one fused winograd transform+FMA+writeback kernel +template +__device__ void +winograd_conv_f23_fp32(F f, Output output, Input x, Weights weights, Inputs... inputs) +{ + static_assert(KO >= 1, "KO must be >= 1"); + static_assert(TILES >= 1, "TILES must be >= 1"); + static_assert(SK >= 1 and SK <= NW and (NW % SK) == 0, "SK must divide NW"); + static_assert(CU == 1 or CU == 2 or CU == 4, "CU must be 1, 2, or 4"); + static_assert(MIGRAPHX_WAVEFRONTSIZE == 32, "winograd_conv_f23_fp32 requires wave32"); + + auto idx = make_index(); + auto out_shape = output.get_shape(); + auto x_shape = x.get_shape(); + auto w_shape = weights.get_shape(); + + const auto out_c = out_shape.lens[1]; + const auto out_h = out_shape.lens[2]; + const auto out_w = out_shape.lens[3]; + const auto n = out_shape.lens[0]; + const auto in_c = x_shape.lens[1]; + const auto in_h = x_shape.lens[2]; + const auto in_w = x_shape.lens[3]; + + const auto tiles_w = (out_w + 1) / 2; + const auto tiles_h = (out_h + 1) / 2; + const auto nt_total = n * tiles_h * tiles_w; + + constexpr index_int quads_per_wave = 32 / 4; // 8 quads per wave + constexpr index_int nt_groups = NW / SK; // independent tile groups per WG + constexpr index_int quads_per_wg = quads_per_wave * nt_groups; // tiles-worth of quads + const auto k_blocks = (out_c + KO - 1) / KO; + + // One workgroup per k_block; its nt_groups NT-groups cover a contiguous run + // of tiles (the SK waves of a group split the channel contraction, not the + // tiles). Consecutive workgroups cycle the k_block (idx.group % k_blocks) so + // concurrently-scheduled workgroups tend to share input tiles (same + // tile_group, different output channels) -> input-cache reuse. + const index_int k_block = idx.group % k_blocks; + const index_int tile_group = idx.group / k_blocks; + const index_int k_base = k_block * KO; + + const index_int lane = idx.local % 32; + const index_int wave_id = idx.local / 32; + const index_int wave_nt_idx = wave_id / SK; // which NT-group (tile range) + const index_int wave_sk_part = wave_id % SK; // which channel subset + const index_int v_col = lane % 4; + const index_int quad_in_wave = lane / 4; + const index_int quad_id = + tile_group * quads_per_wg + wave_nt_idx * quads_per_wave + quad_in_wave; + + // Shuffle-sign for the input v-axis butterfly (see wino_f23_bt_v). + const float in_shuf_sign = (v_col == 1) ? 1.0f : -1.0f; + + // Per-tile geometry for the TILES tiles this quad owns. + array n_arr{}; + array th_arr{}; + array tw_arr{}; + array active_arr{}; + repeat_c([&](auto tt) { + constexpr index_int t = tt; + const index_int nt = quad_id * TILES + t; + const bool active = nt < nt_total; + const auto tile = + active ? array{n, tiles_h, tiles_w}.multi(nt) : array{}; + n_arr[t] = tile[0]; + th_arr[t] = tile[1]; + tw_arr[t] = tile[2]; + active_arr[t] = active; + }); + + // ---- Buffer resources + precomputed base offsets ---- + const auto x_str = x_shape.strides; // {sn, sc, sh, sw} + const uint32_t x_byte_count = static_cast(x_shape.element_space()) * sizeof(float); + auto x_rsrc = make_oob_buffer_rsrc(x.data(), x_byte_count); + const int32_t x_oob = static_cast(x_byte_count); + const index_int c_stride_x = x_str[1]; + // The NHWC path assumes the channel axis is innermost (packed, per-channel + // step 1 float); the JIT sets NHWC from strides()[1] == 1. + MIGRAPHX_ASSERT(not NHWC or c_stride_x == 1); + + const auto w_str = w_shape.strides; // {su, sv, sk, sc} + const uint32_t w_byte_count = static_cast(w_shape.element_space()) * sizeof(float); + auto w_rsrc = make_oob_buffer_rsrc(weights.data(), w_byte_count); + const int32_t w_oob = static_cast(w_byte_count); + // Innermost weight dim (stride 1) is C for NCHW [u,v,k,c] and v for NHWC + // [u,k,c,v]; either way dim 3 is packed. NHWC's v-innermost layout makes the 4 + // v_col lanes read consecutive floats (coalesced); its channel load is then + // w_str[w_c_dim]-strided (see fma_block). + MIGRAPHX_ASSERT(w_str[3] == 1); + + // Per (tile,row) input byte offset for channel 0 of this lane's column; the + // channel stride is added in the loop. OOB rows/cols use the sentinel so the + // buffer load returns 0 (the winograd zero-padding halo). + array, TILES> x_off{}; + repeat_c([&](auto tt) { + constexpr index_int t = tt; + const int h0 = static_cast(2 * th_arr[t]) - 1; + const int w0 = static_cast(2 * tw_arr[t]) - 1; + const int ww = w0 + static_cast(v_col); + const bool w_in = active_arr[t] and ww >= 0 and ww < static_cast(in_w); + repeat_c<4>([&](auto aa) { + constexpr int a = aa; + const int hh = h0 + a; + const bool ok = w_in and hh >= 0 and hh < static_cast(in_h); + x_off[t][a] = ok ? static_cast((n_arr[t] * x_str[0] + + static_cast(hh) * x_str[2] + + static_cast(ww) * x_str[3]) * + sizeof(float)) + : x_oob; + }); + }); + + // Weight byte-offset bases for this lane's v column (channel 0). Computed + // inline in the loop rather than precomputed per (u,k) so we don't pin + // 4*KO offset registers live across the channel loop. The host stores U either + // packed [u,v,k,c] (c innermost) or -- for the gated NHWC configs where it + // helps -- v-innermost [u,k,c,v] (so the 4 v_col lanes coalesce). VINNER (set + // by the JIT from the weight layout) picks the v/k/c stride-dim indices + // accordingly (u is always dim 0). + constexpr bool w_vinner = VINNER; + const index_int w_v_dim = w_vinner ? 3 : 1; // v-dim stride index + const index_int w_k_dim = w_vinner ? 1 : 2; // k-dim stride index + const index_int w_c_dim = w_vinner ? 2 : 3; // c-dim stride index + const int32_t w_lane_base = + static_cast((v_col * w_str[w_v_dim] + k_base * w_str[w_k_dim]) * sizeof(float)); + const int32_t w_u_stride = w_str[0] * sizeof(float); + const int32_t w_k_stride = w_str[w_k_dim] * sizeof(float); + const int32_t w_c_stride = w_str[w_c_dim] * sizeof(float); + auto w_byte_off = [&](index_int u, index_int k) { + return (k_base + k < out_c) ? (w_lane_base + static_cast(u) * w_u_stride + + static_cast(k) * w_k_stride) + : w_oob; + }; + + // Accumulators M[u][t][k]. + array, TILES>, 4> m{}; + + using v_reg_t = array, 4>, TILES>; + + // Apply the winograd input transform (B^T along the v axis then the u axis) to + // one tile's 4 raw H-column data and store the result into vr[t][*][cu]. Shared + // by the NCHW per-channel path and the NHWC block path, which differ only in + // where the 4 data come from (a scattered b32 load vs a preloaded b128 vector). + auto store_transformed = + [&](v_reg_t& vr, index_int t, index_int cu, const array& draw) { + array p{}; + repeat_c<4>([&](auto aa) { + constexpr int a = aa; + p[a] = wino_f23_bt_v(draw[a], in_shuf_sign); + }); + const auto vu = wino_f23_bt_u(p); + repeat_c<4>([&](auto uu) { vr[t][uu][cu] = vu[uu]; }); + }; + + // Input transform of ONE channel cu of the block at c0 (all TILES tiles) into + // vr. This is the DPP-heavy part (4 fmac_dpp per tile). cu is a compile-time + // integral_constant so it selects the fixed v_reg slot. + auto transform_chan = [&](v_reg_t& vr, index_int c0, index_int nchan, auto cc) { + constexpr index_int cu = cc; + if(cu >= nchan) + return; + const int32_t coff = static_cast((c0 + cu) * c_stride_x * sizeof(float)); + repeat_c([&](auto tt) { + constexpr index_int t = tt; + array draw{}; + repeat_c<4>([&](auto aa) { + constexpr int a = aa; + draw[a] = buffer_load(x_rsrc, x_off[t][a] + coff); + }); + store_transformed(vr, t, cu, draw); + }); + }; + + // Full block transform (all CU channels). For NHWC (channels innermost, + // stride 1) the CU channels are contiguous, so load them with one b128 per + // (tile, a-row) -- 4x fewer input loads than the per-channel b32 path -- then + // apply the same per-channel v/u transform. (NCHW channels are H*W apart, so + // it stays per-channel.) + // + // The NHWC weight load is scattered (lane == v_col, so the 4 lanes read + // v-slices K*C apart). The host mitigates that by storing U as v-innermost + // [u,k,c,v] so lanes v_col=0..3 read consecutive floats, gated to out_c>=128 + // and in_c<=out_c since v-innermost makes the channel load strided. The + // input-side scatter is inherent to the layout and is left as-is. + auto transform_block = [&](index_int c0, index_int nchan) { + v_reg_t vr{}; + if constexpr(NHWC) + { + const int32_t coff = static_cast(c0 * sizeof(float)); // c_stride_x == 1 + repeat_c([&](auto tt) { + constexpr index_int t = tt; + array, 4> d4{}; // [a][cu], CU contiguous channels via b128 + repeat_c<4>([&](auto aa) { + constexpr int a = aa; + d4[a] = buffer_load_vec(x_rsrc, x_off[t][a] + coff); + }); + repeat_c([&](auto cc) { + constexpr index_int cu = cc; + if(cu >= nchan) + return; + array draw{}; + repeat_c<4>([&](auto aa) { + constexpr int a = aa; + draw[a] = d4[a][cu]; + }); + store_transformed(vr, t, cu, draw); + }); + }); + } + else + { + repeat_c([&](auto cc) { transform_chan(vr, c0, nchan, cc); }); + } + return vr; + }; + + // Weight load (b128 over CU channels) + FMA accumulate for the block whose + // transform is in v_cur. The FMA nest is channel-outer (cu), then (k, t): each + // cu-iteration updates KO*TILES *independent* accumulators, so consecutive + // FMAs are dependency-free and the matrix pipe stays busy. This u's KO weight + // vectors are loaded up front so they can feed the cu loop. + // + // For PIPE!=0 and has_next, the NEXT block's transform is software-pipelined + // in: u-iteration u transforms next-block channel cu=u into v_next, then a + // sched_barrier pins it so the compiler can neither hoist/cluster the DPP ahead + // of the FMA stream nor over-pipeline into a spill. This interleaves the + // non-dual-issue DPP ops with the dual-issuable FMAs (MIOpen-style). The whole + // weave is compiled out for PIPE==0 (v_next is then an unused dummy). + auto fma_block = [&](const v_reg_t& v_cur, + index_int c0, + index_int nchan, + bool has_next, + v_reg_t& v_next, + index_int c_next, + index_int nchan_next) { + // Weight channel stride in bytes: 1 float (NCHW, C innermost) or w_str[w_c_dim] + // (NHWC, v-innermost -> the channel dim is strided by the 4 v's). + const int32_t coff_w = static_cast(c0) * w_c_stride; + repeat_c<4>([&](auto uu) { + constexpr index_int u = uu; + array, KO> wv{}; + repeat_c([&](auto kk) { + constexpr index_int k = kk; + const int32_t w_off_base = w_byte_off(u, k) + coff_w; + if(not w_vinner and nchan == CU) + { + // C-innermost weight, full block: 4 contiguous channels in b128. + wv[k] = buffer_load_vec(w_rsrc, w_off_base); + } + else + { + // v-innermost weight (w_c_stride-strided channels, coalesced + // across the 4 v_col lanes) or a partial block: per-channel loads. + repeat_c([&](auto cc) { + constexpr index_int cu = cc; + wv[k][cu] = + (cu < nchan) + ? buffer_load( + w_rsrc, w_off_base + static_cast(cu) * w_c_stride) + : 0.0f; + }); + } + }); + repeat_c([&](auto cc) { + constexpr index_int cu = cc; + if(cu >= nchan) + return; + repeat_c([&](auto kk) { + constexpr index_int k = kk; + repeat_c([&](auto tt) { + constexpr index_int t = tt; + m[u][t][k] += v_cur[t][u][cu] * wv[k][cu]; + }); + }); + }); + if constexpr(PIPE) + { + if(has_next) + transform_chan(v_next, c_next, nchan_next, uu); + __builtin_amdgcn_sched_barrier(0); + } + }); + }; + + // S-store contraction (weight is the v-half-transformed S=[3,4,K,C]). k-outer: + // per output channel k, load this lane's 3 S[i][v_col] values (vs 4 U[u][v_col] + // for full U) and finish U = G S with a register-only u-transform (U0=S0, + // U1=.5(S0+S1+S2), U2=.5(S0-S1+S2), U3=S2; v=3 already negated in S). 25% fewer + // weight loads/bytes -- for weight-bandwidth-bound shapes. w_byte_off(i,k) + // reuses the full-U offset formula with i in 0..2 (S's dim0 stride == U's). + auto fma_block_sstore = [&](const v_reg_t& v_cur, index_int c0, index_int nchan) { + const int32_t coff_w = static_cast(c0 * sizeof(float)); + repeat_c([&](auto kk) { + constexpr index_int k = kk; + array, 3> s{}; + repeat_c<3>([&](auto ii) { + constexpr index_int i = ii; + const int32_t w_off_base = w_byte_off(i, k) + coff_w; + if(nchan == CU) + { + s[i] = buffer_load_vec(w_rsrc, w_off_base); + } + else + { + repeat_c([&](auto cc) { + constexpr index_int cu = cc; + s[i][cu] = + (cu < nchan) + ? buffer_load( + w_rsrc, w_off_base + static_cast(cu * sizeof(float))) + : 0.0f; + }); + } + }); + array, 4> uwv{}; + repeat_c([&](auto cc) { + constexpr index_int cu = cc; + const float s0 = s[0][cu]; + const float s1 = s[1][cu]; + const float s2 = s[2][cu]; + const float a = s0 + s2; + uwv[0][cu] = s0; + uwv[1][cu] = 0.5f * (a + s1); + uwv[2][cu] = 0.5f * (a - s1); + uwv[3][cu] = s2; + }); + repeat_c([&](auto cc) { + constexpr index_int cu = cc; + if(cu >= nchan) + return; + repeat_c<4>([&](auto uz) { + constexpr index_int u = uz; + repeat_c([&](auto tt) { + constexpr index_int t = tt; + m[u][t][k] += v_cur[t][u][cu] * uwv[u][cu]; + }); + }); + }); + }); + }; + + // Channel loop, CU channels per step. With SK>1 the CU-blocks are split + // round-robin across the SK waves of this NT-group (each wave sums 1/SK of the + // channels into its own partial M). + if constexpr(SSTORE) + { + // Simple per-block loop with the S-store (v-half) weight contraction. + for(index_int cb = wave_sk_part; cb * CU < in_c; cb += SK) + { + const index_int c = cb * CU; + const index_int avail = in_c - c; + const index_int nchan = avail < CU ? avail : CU; + v_reg_t v_cur = transform_block(c, nchan); + fma_block_sstore(v_cur, c, nchan); + } + } + else if constexpr(PIPE) + { + // Software-pipelined: transform block N+1 while FMA-accumulating block N. + // Costs a second live v_reg (higher VGPR) -- a tuner-selected option gated + // to small (non-spilling) solutions; wins on FMA-throughput-bound shapes. + index_int cb = wave_sk_part; + if(cb * CU < in_c) + { + const index_int avail0 = in_c - cb * CU; + v_reg_t v_cur = transform_block(cb * CU, avail0 < CU ? avail0 : CU); + while(cb * CU < in_c) + { + const index_int c_cur = cb * CU; + const index_int avail = in_c - c_cur; + const index_int nchan = avail < CU ? avail : CU; + const index_int cb_next = cb + SK; + const bool has_next = cb_next * CU < in_c; + const index_int c_next = cb_next * CU; + const index_int avail_nx = has_next ? (in_c - c_next) : 0; + const index_int nchan_nx = avail_nx < CU ? avail_nx : CU; + v_reg_t v_next{}; + fma_block(v_cur, c_cur, nchan, has_next, v_next, c_next, nchan_nx); + v_cur = v_next; + cb = cb_next; + } + } + } + else + { + // Simple: transform each block fully, then FMA it (no cross-block overlap). + // v_scratch is never touched (the weave is compiled out) -- DCE removes it. + v_reg_t v_scratch{}; + for(index_int cb = wave_sk_part; cb * CU < in_c; cb += SK) + { + const index_int c = cb * CU; + const index_int avail = in_c - c; + const index_int nchan = avail < CU ? avail : CU; + v_reg_t v_cur = transform_block(c, nchan); + fma_block(v_cur, c, nchan, false, v_scratch, 0, 0); + } + } + + // ---- Split-c cross-wave reduce (SK>1): sum the SK partial M accumulators + // of this NT-group through LDS; the wave_sk_part==0 wave ends up with the + // full M and does the output transform + writeback. ---- + constexpr index_int m_per_lane = 4 * TILES * KO; + constexpr index_int red_len = (SK > 1) ? (NW * 32 * m_per_lane) : 1; + __shared__ uninitialized_buffer m_reduce; + if constexpr(SK > 1) + { + const index_int lane_base = (wave_id * 32 + lane) * m_per_lane; + repeat_c<4>([&](auto uu) { + constexpr index_int u = uu; + repeat_c([&](auto tt) { + constexpr index_int t = tt; + repeat_c([&](auto kk) { + constexpr index_int k = kk; + constexpr index_int off = u * (TILES * KO) + t * KO + k; + m_reduce[lane_base + off] = m[u][t][k]; + }); + }); + }); + __syncthreads(); + // Only wave_sk_part==0 sums the SK partials and writes back; the others are + // done once they have synced. + if(wave_sk_part != 0) + return; + for(index_int s = 1; s < SK; ++s) + { + const index_int s_base = ((wave_nt_idx * SK + s) * 32 + lane) * m_per_lane; + repeat_c<4>([&](auto uu) { + constexpr index_int u = uu; + repeat_c([&](auto tt) { + constexpr index_int t = tt; + repeat_c([&](auto kk) { + constexpr index_int k = kk; + constexpr index_int off = u * (TILES * KO) + t * KO + k; + m[u][t][k] += m_reduce[s_base + off]; + }); + }); + }); + } + } + + // ---- Output transform A^T M A + writeback ---- + using out_type = typename Output::type; + repeat_c([&](auto tt) { + constexpr index_int t = tt; + if(not active_arr[t]) + return; + const index_int oh0 = 2 * th_arr[t]; + const index_int ow0 = 2 * tw_arr[t]; + repeat_c([&](auto kk) { + constexpr index_int k2 = kk; + const index_int k = k_base + k2; + if(k >= out_c) + return; + + // u-axis reduce (registers): N[i] for i = 0,1. + const float m0 = m[0][t][k2]; + const float m1 = m[1][t][k2]; + const float m2 = m[2][t][k2]; + const float m3 = m[3][t][k2]; + float n_row0 = m0 + m1 + m2; // i = 0 + float n_row1 = m1 - m2 - m3; // i = 1 + + // v-axis reduce (cross-lane within the quad). Every lane must join + // the gathers, but only lanes 0 and 1 (output columns 0 and 1) store, + // so bail the other two before forming the results. + const float g0_1 = readlane_xor<1>(n_row0); + const float g0_2 = readlane_xor<2>(n_row0); + const float g0_3 = readlane_xor<3>(n_row0); + const float g1_1 = readlane_xor<1>(n_row1); + const float g1_2 = readlane_xor<2>(n_row1); + const float g1_3 = readlane_xor<3>(n_row1); + if(v_col > 1) + return; + + const index_int ow = ow0 + v_col; + if(ow >= out_w) + return; + // On lane0 the "+n1+n2" form equals Y[i][col0]; on lane1 the + // "-n3-n2" form equals Y[i][col1]. + const float y_i0 = (v_col == 0) ? (n_row0 + g0_1 + g0_2) : (n_row0 - g0_3 - g0_2); + const float y_i1 = (v_col == 0) ? (n_row1 + g1_1 + g1_2) : (n_row1 - g1_3 - g1_2); + auto store = [&](index_int oh, float y) { + const array oid{n_arr[t], k, oh, ow}; + output[oid] = static_cast(f(static_cast(y), inputs[oid]...)); + }; + store(oh0, y_i0); + if(oh0 + 1 < out_h) + store(oh0 + 1, y_i1); + }); + }); +} +// NOLINTEND(readability-function-size) + +} // namespace migraphx + +#endif // MIGRAPHX_GUARD_KERNELS_WINOGRAD_CONV_FP32_HPP diff --git a/src/targets/gpu/prefuse_ops.cpp b/src/targets/gpu/prefuse_ops.cpp index 18c44db7f17..80ca18b171f 100644 --- a/src/targets/gpu/prefuse_ops.cpp +++ b/src/targets/gpu/prefuse_ops.cpp @@ -46,6 +46,7 @@ namespace gpu { MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_ENABLE_LAYERNORM_FUSION); MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_DISABLE_MLIR); MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_ENABLE_WINOGRAD); +MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_WINOGRAD_FP32_SSTORE); MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_DISABLE_WINOGRAD); MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_WINOGRAD_FULL_TRANSFORM); @@ -365,8 +366,14 @@ struct winograd_conv const auto& x_shape = inputs[0]; const auto& u_shape = inputs[1]; auto x_lens = x_shape.lens(); - // u_shape is [4 or 3, 3, K, C]; lens()[2] is the output channel count. - auto out_c = u_shape.lens()[2]; + const auto& u_lens = u_shape.lens(); + // The weight encodes the output channel count K; its axis depends on the + // store layout: fp16 [4|3, 3, K, C] and fp32 full-U / S-store [4|3, 4, K, C] + // put K at dim 2, while the fp32 v-innermost store [4, K, C, 4] (the NHWC + // coalesced weight load) puts K at dim 1. v-inner is the fp32 layout whose + // dim 1 is not the 4-wide v axis (matches the kernel's lens[1] != 4 test). + const bool vinner = x_shape.type() == shape::float_type and u_lens[1] != 4; + const auto out_c = vinner ? u_lens[1] : u_lens[2]; std::vector out_lens = {x_lens[0], out_c, x_lens[2], x_lens[3]}; return shape::from_permutation(x_shape.type(), out_lens, output_layout); } @@ -427,6 +434,108 @@ literal compute_winograd_weights_f23(const argument& w_arg, bool full_transform) return literal{w_shape, data}; } +// Winograd F(2,3) filter matrix G (4x3): rows [1,0,0], [.5,.5,.5], [.5,-.5,.5], +// [0,0,1]. Shared by the fp32 full-U and S-store weight transforms below. +constexpr std::array, 4> winograd_f23_gmat{ + {{1.0f, 0.0f, 0.0f}, {0.5f, 0.5f, 0.5f}, {0.5f, -0.5f, 0.5f}, {0.0f, 0.0f, 1.0f}}}; + +// Precompute the FULL Winograd filter transform U = G g G^T for the fp32 FMA/DPP +// kernel, stored as an [4, 4, K, C] literal (indices u, v, k, c; C innermost). +// The fp32 kernel does not transform the weight in-kernel (unlike the fp16 +// path), so the whole 4x4 winograd weight is materialized here. +// +// The kernel realizes the input transform's v=3 column with a sign flip +// (d3-d1 instead of d1-d3) because that is the form a single quad DPP butterfly +// can produce; a matching negation of U[:,3,:,:] here makes the elementwise +// product U*V exact. +// vinner=true stores U packed as [u, k, c, v] (v physically INNERMOST) instead of +// the NCHW [u, v, k, c] (c innermost). In NHWC the fp32 kernel's 4 v_col lanes then +// read 4 *consecutive* floats -> the weight load coalesces (lane==v_col otherwise +// scatters the weight across v-slices K*C apart, thrashing the cache with the +// input). The kernel selects the v/k/c stride indices by layout (NHWC template). +literal compute_winograd_weights_f23_fp32(const argument& w_arg, bool vinner = false) +{ + const auto& sh = w_arg.get_shape(); + auto out_c = sh.lens()[0]; + auto in_c = sh.lens()[1]; + shape u_shape = vinner ? shape{shape::float_type, {4, out_c, in_c, 4}} // [u,k,c,v] + : shape{shape::float_type, {4, 4, out_c, in_c}}; // [u,v,k,c] + auto widx = [&](std::size_t u, std::size_t v, std::size_t k, std::size_t c) { + return vinner ? u_shape.index({u, k, c, v}) : u_shape.index({u, v, k, c}); + }; + + std::vector data(u_shape.elements(), 0.0f); + const auto& gmat = winograd_f23_gmat; + w_arg.visit([&](auto w_view) { + dfor(out_c, in_c)([&](auto k, auto c) { + float g[3][3]; + dfor(std::size_t{3}, + std::size_t{3})([&](auto i, auto j) { g[i][j] = w_view(k, c, i, j); }); + + // Gg (4x3): (Gg)[u][j] = sum_i G[u][i] g[i][j]. + float gg[4][3]; + dfor(std::size_t{4}, std::size_t{3})([&](auto u, auto j) { + gg[u][j] = gmat[u][0] * g[0][j] + gmat[u][1] * g[1][j] + gmat[u][2] * g[2][j]; + }); + // U[u][v] = sum_j (Gg)[u][j] G[v][j], with the v=3 column negated. + dfor(std::size_t{4}, std::size_t{4})([&](auto u, auto v) { + float uv = gg[u][0] * gmat[v][0] + gg[u][1] * gmat[v][1] + gg[u][2] * gmat[v][2]; + data[widx(u, v, k, c)] = (v == 3) ? -uv : uv; + }); + }); + }); + return literal{u_shape, data}; +} + +// S-store: precompute the v-half-transformed winograd weight S = g G^T, stored as +// [3, 4, K, C] (i, v, k, c; C innermost). This is 12 values/(k,c) vs full U's 16 +// (25% less weight DRAM), and each lane loads only its v_col's 3 values (vs 4 for +// U). The kernel finishes U = G S with a register-only u-transform (u lives in +// registers, so no cross-lane traffic -- unlike a u-half g-store, which would +// leave the kernel a cross-lane v-transform). The +// v=3 column is negated here to match the input butterfly's d3-d1 sign, exactly +// as the full-U store does. +literal compute_winograd_weights_f23_fp32_sstore(const argument& w_arg) +{ + const auto& sh = w_arg.get_shape(); + auto out_c = sh.lens()[0]; + auto in_c = sh.lens()[1]; + shape s_shape{shape::float_type, {3, 4, out_c, in_c}}; + + std::vector data(s_shape.elements(), 0.0f); + const auto& gmat = winograd_f23_gmat; + w_arg.visit([&](auto w_view) { + dfor(out_c, in_c)([&](auto k, auto c) { + float g[3][3]; + dfor(std::size_t{3}, + std::size_t{3})([&](auto i, auto j) { g[i][j] = w_view(k, c, i, j); }); + // S[i][v] = sum_j g[i][j] G[v][j], with the v=3 column negated. + dfor(std::size_t{3}, std::size_t{4})([&](auto i, auto v) { + float sv = g[i][0] * gmat[v][0] + g[i][1] * gmat[v][1] + g[i][2] * gmat[v][2]; + data[s_shape.index({i, v, k, c})] = (v == 3) ? -sv : sv; + }); + }); + }); + return literal{s_shape, data}; +} + +// Look up an exact (C, K, H, W) entry in a per-shape override table (each entry +// must expose in_ch/out_ch/height/width fields); returns the entry or nullptr. +// Shared by the fp16/fp32 profitability heuristics and the S-store table. +template +const typename Table::value_type* find_shape_override(const Table& table, + std::size_t in_ch, + std::size_t out_ch, + std::size_t height, + std::size_t width) +{ + auto it = std::find_if(table.begin(), table.end(), [&](const auto& o) { + return std::tie(o.in_ch, o.out_ch, o.height, o.width) == + std::tie(in_ch, out_ch, height, width); + }); + return it == table.end() ? nullptr : &*it; +} + // Measured per-shape overrides: exact (C, K, H, W) convolutions where the // analytic heuristic below mispredicts the winograd-vs-default winner by more // than 10% (using the better of the two weight stores). These are @@ -501,13 +610,7 @@ bool winograd_f23_profitable( if(nhwc and min_ch >= 224) return false; - // NOLINTNEXTLINE(readability-qualified-auto) - const auto ovr = std::find_if( - winograd_f23_overrides.begin(), winograd_f23_overrides.end(), [&](const auto& o) { - return std::tie(o.in_ch, o.out_ch, o.height, o.width) == - std::tie(in_ch, out_ch, height, width); - }); - if(ovr != winograd_f23_overrides.end()) + if(const auto* ovr = find_shape_override(winograd_f23_overrides, in_ch, out_ch, height, width)) return ovr->use_winograd; if(in_ch < 16 and (max_ch > 16 or in_ch < 8)) @@ -559,6 +662,141 @@ bool winograd_f23_full_transform(std::size_t in_ch, return true; } +struct winograd_f23_sstore_shape +{ + std::size_t in_ch; + std::size_t out_ch; + std::size_t height; + std::size_t width; +}; + +// Shapes in the spatial-16..64 high-channel band that prefer the S-store. In this +// band the S-vs-full-U preference is micro-architecturally non-monotonic -- +// neighbouring channel counts flip it (512->512 vs 515->512 at 16x16, 192->191 vs +// 192->192 at 64x64, 768->383 vs 384->384 at 32x32) -- so a smooth rule cannot +// separate them without regressing the shapes that need the full U. Hence an +// explicit table, like the fp16 path. +constexpr std::array winograd_f23_sstore_overrides{{ + {512, 512, 16, 16}, + {512, 512, 24, 24}, + {195, 192, 64, 64}, + {768, 383, 32, 32}, + {384, 383, 32, 32}, + {384, 191, 64, 64}, + {192, 191, 64, 64}, +}}; + +// Choose the fp32 winograd weight encoding: S-store (v-half-transformed g*G^T, +// [3,4,K,C]) vs the full U ([4,4,K,C]). S-store cuts weight loads AND bytes 25% +// and finishes U with a cheap register-only u-transform, so it suits the +// weight-load-dominated shapes: high channels with small spatial. Elsewhere its +// extra register FMA and k-outer's lower ILP cost more than the bandwidth saved. +// +// Two selection paths: (1) a smooth zone (very small spatial, high channels) +// where S-store is uniformly preferable; (2) the override table above for the +// non-monotonic spatial-16..64 band. +bool winograd_f23_use_sstore(std::size_t in_ch, + std::size_t out_ch, + std::size_t height, + std::size_t width) +{ + if(std::min(in_ch, out_ch) >= 256 and std::min(height, width) <= 12) + return true; + return find_shape_override(winograd_f23_sstore_overrides, in_ch, out_ch, height, width) != + nullptr; +} + +// Per-shape overrides for the fp32 F(2,3) heuristic below: exact (C, K, H, W) +// convs where the smooth rule mispredicts winograd-vs-MLIR. Two groups need a +// table (like the fp16 path): high-channel square convs in the spatial-16..32 +// band where the S-store flips some but not their neighbours, and awkward-out_ch +// convs at large spatial (see the per-entry notes below). Reuses +// winograd_f23_shape (same C/K/H/W + use_winograd fields as the fp16 table). +constexpr std::array winograd_f23_fp32_overrides{{ + // Exact-square high-channel convs at their native spatial: rocBLAS/MLIR is + // tuned best exactly here, and the smooth min_ch>=224 rule (which holds for + // the off-square and larger-channel neighbours) mispredicts these. + {512, 512, 16, 16, false}, // square, the S-store cannot close the gap + {256, 256, 32, 32, false}, // square + // Awkward output-channel count (95, not a multiple of the KO tile) at large + // spatial: the winograd kernel wastes a partial output-channel block that the + // many tiles then pay for repeatedly, while MLIR does not. min(C,K)=95 < 128, + // so the smooth rule would otherwise keep them (96->96 at the same spatial is + // fine). + {96, 95, 128, 128, false}, + {192, 95, 128, 128, false}, +}}; + +// Heuristic for when the fp32 FMA/DPP F(2,3) winograd kernel beats the default +// (rocMLIR implicit-GEMM) lowering on gfx12, for 3x3/pad-1/stride-1 convs with +// the kernel's own weight-store selection (S-store / v-inner) active. Structure +// mirrors the fp16 winograd_f23_profitable, but the thresholds differ: the fp32 +// kernel has 2.25x fewer MACs than MLIR yet a heavier input/weight scatter, so it +// wins the compute-bound low/mid-channel shapes and loses the memory-bandwidth- +// bound high-channel large-spatial ones. +// - NHWC: rocMLIR's channels-last GEMM reads the input fully coalesced and wins +// almost everywhere; the winograd kernel's C-strided input scatter only pays +// off at tiny spatial + high channels (spatial<=16, min_ch>=256), so NHWC is +// gated to that narrow zone. +// - NCHW: winograd wins broadly. Excluded regions: +// * C*K >= 700k: bandwidth-bound big GEMMs MLIR owns (1280-channel convs). +// * min(C,K) >= 224: only small spatial (<=32) wins (2.25x fewer MACs); +// mid/large spatial is input/output-transform + weight-expansion bound. +// * min(C,K) >= 128 at spatial >= 128: transform-bound, loses. +// * spatial >= 512 with out_ch > 32: the 4x input-tile re-read dominates +// unless the output fits a single KO block. +// Output-collapse (out_ch <= 3) and tiny-input (in_ch < 16) shapes are handled +// layout-independently up front. MIGRAPHX_ENABLE/DISABLE_WINOGRAD override it. +bool winograd_f23_fp32_profitable( + std::size_t in_ch, std::size_t out_ch, std::size_t height, std::size_t width, bool nhwc) +{ + const auto spatial = std::min(height, width); + const auto min_ch = std::min(in_ch, out_ch); + const auto max_ch = std::max(in_ch, out_ch); + + // The next three checks are layout-independent (they hold for both NCHW and + // NHWC), so they run before the layout split. + + // Tiny input channels (RGB-style stems) have too little contraction to + // amortize the winograd transforms, so they lose to a plain GEMM. (Mirrors + // the fp16 rule.) + if(in_ch < 16 and (max_ch > 16 or in_ch < 8)) + return false; + + // Output-collapse convs (out_ch <= 3, e.g. a segmentation/prediction head): + // the winograd output transform is nearly free on so few output channels while + // MLIR still runs a full small GEMM, so winograd wins in both layouts. + if(out_ch <= 3) + return true; + + // Per-shape overrides (each listed shape loses in both layouts). + if(const auto* ovr = + find_shape_override(winograd_f23_fp32_overrides, in_ch, out_ch, height, width)) + return ovr->use_winograd; + + // NHWC: rocMLIR's coalesced channels-last GEMM wins almost everywhere; the + // winograd kernel's C-strided NHWC input scatter (no host layout freedom, so + // it can't coalesce) only pays off at high channels with tiny spatial, where + // the re-read footprint is small and cached. + if(nhwc) + return min_ch >= 256 and spatial <= 16; + + // NCHW: winograd wins broadly. + if(in_ch * out_ch >= 700000) + return false; + if(min_ch >= 224) + return spatial <= 32; + if(min_ch >= 128 and spatial >= 128) + return false; + // Very large spatial (the 4x winograd input-tile re-read dominates): only a + // single-KO-block output (out_ch <= 32) survives it. (Low/mid channel counts + // still win up to 256x256; min(C,K)>=128 at large spatial is already excluded + // above.) + if(spatial >= 512 and out_ch > 32) + return false; + return true; +} + MIGRAPHX_PRED_MATCHER(conv_winograd_f23, instruction_ref ins) { if(ins->name() != "convolution") @@ -581,9 +819,9 @@ MIGRAPHX_PRED_MATCHER(conv_winograd_f23, instruction_ref ins) if(x_lens.size() != 4) return false; auto x_type = ins->inputs().front()->get_shape().type(); - // Kernel currently only supports half_type (fp16). The fp32 path was - // never wired through the buffer-resource-based loads. - if(x_type != shape::half_type) + // fp16 uses the WMMA kernel; fp32 uses the FMA/DPP kernel. Other types are + // unsupported. + if(x_type != shape::half_type and x_type != shape::float_type) return false; if(ins->inputs().front()->get_shape().dynamic() or ins->inputs().back()->get_shape().dynamic()) return false; @@ -601,8 +839,14 @@ MIGRAPHX_PRED_MATCHER(conv_winograd_f23, instruction_ref ins) // the same test the kernel uses to pick its NHWC path. layout_convolution // runs before this pass, so the strides already reflect the chosen layout. const bool nhwc = ins->inputs().front()->get_shape().strides()[1] == 1; - return enabled(MIGRAPHX_ENABLE_WINOGRAD{}) or - winograd_f23_profitable(w_lens[1], w_lens[0], x_lens[2], x_lens[3], nhwc); + // MIGRAPHX_ENABLE_WINOGRAD forces winograd on every eligible shape (bypassing + // the heuristic); otherwise the per-shape, per-dtype perf heuristic decides. + // fp16 uses the WMMA kernel's heuristic, fp32 the FMA/DPP kernel's. + if(enabled(MIGRAPHX_ENABLE_WINOGRAD{})) + return true; + if(x_type == shape::float_type) + return winograd_f23_fp32_profitable(w_lens[1], w_lens[0], x_lens[2], x_lens[3], nhwc); + return winograd_f23_profitable(w_lens[1], w_lens[0], x_lens[2], x_lens[3], nhwc); } struct find_winograd_f23 @@ -617,15 +861,48 @@ struct find_winograd_f23 auto w_arg = weights->eval(); - auto w_lens = weights->get_shape().lens(); // [K, C, 3, 3] - auto x_lens = input->get_shape().lens(); // [N, C, H, W] - const bool ft = winograd_f23_full_transform(w_lens[1], w_lens[0], x_lens[2], x_lens[3]); // Match the output layout layout_convolution chose for this conv, so // the op is a drop-in replacement (no surrounding transpose changes). auto out_layout = find_permutation(ins->get_shape()); - auto u_lit = compute_winograd_weights_f23(w_arg, ft); - auto u_ins = m.add_literal(u_lit); + // fp32 uses the FMA/DPP kernel with the host-transformed weight (full U, or + // the S-store half selected below); fp16 uses the WMMA kernel with the T or + // g weight store. full_transform is unused for fp32 (precision is derived + // from the input type); false is just the required ctor argument. + if(input->get_shape().type() == shape::float_type) + { + // Pick the weight encoding: S-store (v-half g*G^T [3,4,K,C], 25% less + // weight DRAM+loads) on the weight-load-significant shapes, else the + // full U [4,4,K,C]. The JIT routes to the S path by the weight's first + // dim (3 vs 4). MIGRAPHX_WINOGRAD_FP32_SSTORE forces S on. + // NHWC uses the full U laid out v-innermost so the weight load coalesces + // (S-store stays NCHW-only). + const auto& x_lens = input->get_shape().lens(); // [N, C, H, W] + const auto& w_lens = weights->get_shape().lens(); // [K, C, 3, 3] + const bool nhwc = input->get_shape().strides()[1] == 1; + const bool use_s = + not nhwc and (enabled(MIGRAPHX_WINOGRAD_FP32_SSTORE{}) or + winograd_f23_use_sstore(w_lens[1], w_lens[0], x_lens[2], x_lens[3])); + // v-innermost weight coalesces the NHWC weight load, relieving the + // input/weight cache thrash -- but its strided (b32) channel load adds + // issue overhead that regresses shapes where the weight is small and + // cached (low out_c) or the input dominates (channel-reducing). Gate to + // where the weight is substantial and not channel-reducing. + const auto out_c = w_lens[0]; + const auto in_c = w_lens[1]; + const bool vinner = nhwc and out_c >= 128 and in_c <= out_c; + auto u_lit = use_s ? compute_winograd_weights_f23_fp32_sstore(w_arg) + : compute_winograd_weights_f23_fp32(w_arg, vinner); + m.replace_instruction( + ins, winograd_conv{false, out_layout}, input, m.add_literal(u_lit)); + return; + } + + auto w_lens = weights->get_shape().lens(); // [K, C, 3, 3] + auto x_lens = input->get_shape().lens(); // [N, C, H, W] + const bool ft = winograd_f23_full_transform(w_lens[1], w_lens[0], x_lens[2], x_lens[3]); + auto u_lit = compute_winograd_weights_f23(w_arg, ft); + auto u_ins = m.add_literal(u_lit); m.replace_instruction(ins, winograd_conv{ft, out_layout}, input, u_ins); } diff --git a/test/gpu/winograd_conv_shape.cpp b/test/gpu/winograd_conv_shape.cpp new file mode 100644 index 00000000000..d7618ae52bb --- /dev/null +++ b/test/gpu/winograd_conv_shape.cpp @@ -0,0 +1,188 @@ +/* + * The MIT License (MIT) + * + * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include +#include +#include +#include +#include +#include // links migraphx_gpu so gpu::winograd_conv is registered + +// gpu::winograd_conv::compute_shape derives the output channel count K from the +// weight literal's shape, whose axis order depends on the store layout the host +// picked: +// fp16 T-store [4, 3, K, C] and g-store [3, 3, K, C] -> K at dim 2 +// fp32 full-U [4, 4, K, C] and S-store [3, 4, K, C] -> K at dim 2 +// fp32 v-innermost [4, K, C, 4] (NHWC coalesced weight) -> K at dim 1 +// The output takes the batch/spatial dims of the input and the winograd op's +// output_layout permutation. These tests exercise every (dtype, layout) pair with +// channel-expanding, channel-reducing, and square convs -- the v-inner regression +// (in_c != out_c returning in_c) was masked because every prior test was square. + +using lens_t = std::vector; + +// Run compute_shape by placing the op in a module; conv->get_shape() is its result. +static migraphx::shape wino_shape(const migraphx::shape& x, + const migraphx::shape& u, + const std::vector& layout = {0, 1, 2, 3}) +{ + migraphx::module m; + auto xp = m.add_parameter("x", x); + auto up = m.add_literal(migraphx::generate_literal(u)); + auto conv = m.add_instruction( + migraphx::make_op("gpu::winograd_conv", {{"output_layout", layout}}), xp, up); + return conv->get_shape(); +} + +// ---- fp32 full-U weight [4, 4, K, C] (K at dim 2) ---- +TEST_CASE(fp32_full_u_expand) +{ + auto s = wino_shape({migraphx::shape::float_type, {1, 128, 16, 16}}, + {migraphx::shape::float_type, {4, 4, 256, 128}}); + EXPECT(s.lens() == lens_t{1, 256, 16, 16}); + EXPECT(s.type() == migraphx::shape::float_type); +} + +TEST_CASE(fp32_full_u_reduce) +{ + auto s = wino_shape({migraphx::shape::float_type, {1, 256, 16, 16}}, + {migraphx::shape::float_type, {4, 4, 64, 256}}); + EXPECT(s.lens() == lens_t{1, 64, 16, 16}); +} + +TEST_CASE(fp32_full_u_square) +{ + auto s = wino_shape({migraphx::shape::float_type, {1, 192, 32, 32}}, + {migraphx::shape::float_type, {4, 4, 192, 192}}); + EXPECT(s.lens() == lens_t{1, 192, 32, 32}); +} + +// ---- fp32 S-store weight [3, 4, K, C] (K at dim 2) ---- +TEST_CASE(fp32_sstore_expand) +{ + auto s = wino_shape({migraphx::shape::float_type, {1, 128, 16, 16}}, + {migraphx::shape::float_type, {3, 4, 256, 128}}); + EXPECT(s.lens() == lens_t{1, 256, 16, 16}); +} + +TEST_CASE(fp32_sstore_square) +{ + auto s = wino_shape({migraphx::shape::float_type, {1, 512, 8, 8}}, + {migraphx::shape::float_type, {3, 4, 512, 512}}); + EXPECT(s.lens() == lens_t{1, 512, 8, 8}); +} + +// ---- fp32 v-innermost weight [4, K, C, 4] (K at dim 1) -- the regression ---- +TEST_CASE(fp32_vinner_expand) +{ + auto s = wino_shape({migraphx::shape::float_type, {1, 128, 16, 16}}, + {migraphx::shape::float_type, {4, 256, 128, 4}}); + EXPECT(s.lens() == lens_t{1, 256, 16, 16}); // K = 256, not C = 128 + EXPECT(s.type() == migraphx::shape::float_type); +} + +TEST_CASE(fp32_vinner_reduce) +{ + auto s = wino_shape({migraphx::shape::float_type, {1, 256, 16, 16}}, + {migraphx::shape::float_type, {4, 64, 256, 4}}); + EXPECT(s.lens() == lens_t{1, 64, 16, 16}); +} + +TEST_CASE(fp32_vinner_square) +{ + auto s = wino_shape({migraphx::shape::float_type, {1, 256, 64, 64}}, + {migraphx::shape::float_type, {4, 256, 256, 4}}); + EXPECT(s.lens() == lens_t{1, 256, 64, 64}); +} + +// ---- fp16 T-store weight [4, 3, K, C] (K at dim 2; dim1 == 3 must not trip the +// fp32-only v-inner test, which keys on dim1 != 4) ---- +TEST_CASE(fp16_tstore_expand) +{ + auto s = wino_shape({migraphx::shape::half_type, {1, 128, 16, 16}}, + {migraphx::shape::half_type, {4, 3, 256, 128}}); + EXPECT(s.lens() == lens_t{1, 256, 16, 16}); + EXPECT(s.type() == migraphx::shape::half_type); +} + +TEST_CASE(fp16_tstore_reduce) +{ + auto s = wino_shape({migraphx::shape::half_type, {1, 512, 16, 16}}, + {migraphx::shape::half_type, {4, 3, 128, 512}}); + EXPECT(s.lens() == lens_t{1, 128, 16, 16}); +} + +TEST_CASE(fp16_tstore_square) +{ + auto s = wino_shape({migraphx::shape::half_type, {1, 64, 128, 128}}, + {migraphx::shape::half_type, {4, 3, 64, 64}}); + EXPECT(s.lens() == lens_t{1, 64, 128, 128}); +} + +// ---- fp16 g-store weight [3, 3, K, C] (K at dim 2) ---- +TEST_CASE(fp16_gstore_expand) +{ + auto s = wino_shape({migraphx::shape::half_type, {1, 128, 16, 16}}, + {migraphx::shape::half_type, {3, 3, 256, 128}}); + EXPECT(s.lens() == lens_t{1, 256, 16, 16}); + EXPECT(s.type() == migraphx::shape::half_type); +} + +// ---- output_layout: NCHW is standard/packed; NHWC puts the channel axis innermost ---- +TEST_CASE(nchw_output_layout) +{ + auto s = wino_shape({migraphx::shape::float_type, {1, 128, 16, 16}}, + {migraphx::shape::float_type, {4, 4, 256, 128}}, + {0, 1, 2, 3}); + EXPECT(s.lens() == lens_t{1, 256, 16, 16}); + EXPECT(s.standard()); +} + +TEST_CASE(nhwc_output_layout_fp32_vinner) +{ + auto s = wino_shape({migraphx::shape::float_type, {1, 128, 16, 16}}, + {migraphx::shape::float_type, {4, 256, 128, 4}}, + {0, 2, 3, 1}); + EXPECT(s.lens() == lens_t{1, 256, 16, 16}); + EXPECT(s.strides()[1] == 1); // channels-last +} + +TEST_CASE(nhwc_output_layout_fp16) +{ + auto s = wino_shape({migraphx::shape::half_type, {1, 128, 16, 16}}, + {migraphx::shape::half_type, {4, 3, 256, 128}}, + {0, 2, 3, 1}); + EXPECT(s.lens() == lens_t{1, 256, 16, 16}); + EXPECT(s.strides()[1] == 1); +} + +// ---- batch and spatial dims are carried from the input ---- +TEST_CASE(batch_and_spatial_preserved) +{ + auto s = wino_shape({migraphx::shape::float_type, {4, 128, 30, 40}}, + {migraphx::shape::float_type, {4, 256, 128, 4}}); + EXPECT(s.lens() == lens_t{4, 256, 30, 40}); +} + +int main(int argc, const char* argv[]) { test::run(argc, argv); } diff --git a/test/verify/test_conv_3x3_winograd_fp32.cpp b/test/verify/test_conv_3x3_winograd_fp32.cpp new file mode 100644 index 00000000000..4384db2a3bc --- /dev/null +++ b/test/verify/test_conv_3x3_winograd_fp32.cpp @@ -0,0 +1,55 @@ +/* + * The MIT License (MIT) + * + * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include "verify_program.hpp" +#include +#include +#include +#include + +// fp32 F(2,3) winograd (the FMA/DPP kernel). This 20->36 @15x15 conv is selected +// by the fp32 winograd heuristic, so it exercises the kernel by default. Odd +// spatial size exercises the boundary tiles (halo padding); the channel count is +// not a multiple of the per-lane output block so the partial-KO store path is +// covered. MIGRAPHX_DISABLE_WINOGRAD forces the default lowering for comparison. +struct test_conv_3x3_winograd_fp32 : verify_program +{ + migraphx::program create_program() const + { + migraphx::program p; + auto* mm = p.get_main_module(); + auto x = mm->add_parameter("x", {migraphx::shape::float_type, {1, 20, 15, 15}}); + // Winograd matcher requires can_eval() on weights -> add as a literal. + auto w = mm->add_literal( + migraphx::generate_literal({migraphx::shape::float_type, {36, 20, 3, 3}}, 1)); + auto y = mm->add_instruction( + migraphx::make_op("convolution", + {{"padding", {1, 1}}, {"stride", {1, 1}}, {"dilation", {1, 1}}}), + x, + w); + mm->add_return({y}); + return p; + } + std::string section() const { return "conv"; } +};