From e61490877ad9bf609f9ab207c04f17cc4b98d8c3 Mon Sep 17 00:00:00 2001 From: Aleksei Ivakhnenko Date: Mon, 29 Jun 2026 21:00:03 +0200 Subject: [PATCH 1/2] feat(miopen): fused single-layer LSTM forward-inference kernel (flagged) Kokoro-82M TTS inference on gfx1201 (navi48) is dominated by the bidirectional nn.LSTM in the decoder: at batch=1 the per-timestep recurrence issues ~45k tiny serial GEMM + elementwise dispatches, leaving the GPU idle ~94% of the LSTM and making the workload launch-latency-bound rather than compute-bound. Add a cooperative-grid fused inference kernel (one launch per direction) that walks the whole sequence on-device with a grid-wide barrier per timestep, reusing MIOpen's existing batched input-projection workspace and adding only the recurrent W_hh*h_{t-1} term and the LSTM cell update. This collapses the per-timestep host dispatch loop in RNNForwardInferencePacked. Strictly opt-in behind MIOPEN_DEBUG_RNN_FUSED_INFERENCE and gated to the validated case (single layer, LSTM, miopenRNNdefault, fp32, no dropout, packed, linear input, uniform batch); every other configuration falls back to the existing path unchanged. Default behavior is identical to before -- the flag defaults off because enabling such fusion unconditionally could change results for other models. The kernel honors the caller's initial hidden/cell state (hx/cx) and writes the final hy/cy, matching the stock path; uniform batch>1 is supported. Test: test/gtest/lstm_fused_inference_test.cpp drives miopenRNNForwardInference twice in-process (flag off = stock reference, flag on = fused) and requires the GPU output and hy/cy to agree to 1e-5, giving the comparison real teeth. Covers uni+bi, N=1 and uniform batch>1, fp32, including the single-timestep edge case. The debug flag is registered in the MIOPEN_BUILD_TESTING env registry so ScopedEnvironment can toggle it. Real Kokoro LSTM outputs match the stock path to ~1e-6. JIRA ID : AIRADSW-142 --- projects/miopen/src/CMakeLists.txt | 1 + projects/miopen/src/env_debug.cpp | 2 + .../miopen/src/include/miopen/rnn_util.hpp | 24 ++ .../src/kernels/MIOpenRNNFusedInference.cpp | 175 +++++++++++ projects/miopen/src/ocl/rnnocl.cpp | 71 +++++ projects/miopen/src/rnn/rnn_util.cpp | 82 ++++++ .../test/gtest/lstm_fused_inference_test.cpp | 275 ++++++++++++++++++ 7 files changed, 630 insertions(+) create mode 100644 projects/miopen/src/kernels/MIOpenRNNFusedInference.cpp create mode 100644 projects/miopen/test/gtest/lstm_fused_inference_test.cpp diff --git a/projects/miopen/src/CMakeLists.txt b/projects/miopen/src/CMakeLists.txt index d442837b8637..cff3a1c7fcb4 100644 --- a/projects/miopen/src/CMakeLists.txt +++ b/projects/miopen/src/CMakeLists.txt @@ -585,6 +585,7 @@ if( MIOPEN_BACKEND MATCHES "OpenCL" OR MIOPEN_BACKEND STREQUAL "HIPOC" OR MIOPEN kernels/MIOpenConvFFT.cpp kernels/MIOpenRNNHiddenStateUpdate.cpp + kernels/MIOpenRNNFusedInference.cpp kernels/bugzilla_34765_detect.s kernels/dummy_kernel.s kernels/conv3x3.s diff --git a/projects/miopen/src/env_debug.cpp b/projects/miopen/src/env_debug.cpp index b6827cb1e088..a5a6f6b2fd27 100644 --- a/projects/miopen/src/env_debug.cpp +++ b/projects/miopen/src/env_debug.cpp @@ -61,6 +61,7 @@ MIOPEN_DECLARE_ENV_VAR_BOOL(MIOPEN_DEBUG_CONV_WINOGRAD) MIOPEN_DECLARE_ENV_VAR_STR(MIOPEN_DEBUG_FIND_ONLY_SOLVER) MIOPEN_DECLARE_ENV_VAR_BOOL(MIOPEN_DEBUG_HIP_KERNELS) MIOPEN_DECLARE_ENV_VAR_BOOL(MIOPEN_DEBUG_OPENCL_CONVOLUTIONS) +MIOPEN_DECLARE_ENV_VAR_BOOL(MIOPEN_DEBUG_RNN_FUSED_INFERENCE) MIOPEN_DECLARE_ENV_VAR_BOOL(MIOPEN_DEBUG_WORKAROUND_ISSUE_2492) MIOPEN_DECLARE_ENV_VAR_BOOL(MIOPEN_ENABLE_LOGGING) MIOPEN_DECLARE_ENV_VAR_BOOL(MIOPEN_ENABLE_LOGGING_CMD) @@ -204,6 +205,7 @@ const LibEnvVar& FindEnvVariable(std::string_view name) {MIOPEN_DEBUG_FIND_ONLY_SOLVER.GetName(), MIOPEN_DEBUG_FIND_ONLY_SOLVER}, {MIOPEN_DEBUG_HIP_KERNELS.GetName(), MIOPEN_DEBUG_HIP_KERNELS}, {MIOPEN_DEBUG_OPENCL_CONVOLUTIONS.GetName(), MIOPEN_DEBUG_OPENCL_CONVOLUTIONS}, + {MIOPEN_DEBUG_RNN_FUSED_INFERENCE.GetName(), MIOPEN_DEBUG_RNN_FUSED_INFERENCE}, {MIOPEN_DEBUG_TUNING_ITERATIONS_MAX.GetName(), MIOPEN_DEBUG_TUNING_ITERATIONS_MAX}, {MIOPEN_DEBUG_WORKAROUND_ISSUE_2492.GetName(), MIOPEN_DEBUG_WORKAROUND_ISSUE_2492}, {MIOPEN_ENABLE_LOGGING.GetName(), MIOPEN_ENABLE_LOGGING}, diff --git a/projects/miopen/src/include/miopen/rnn_util.hpp b/projects/miopen/src/include/miopen/rnn_util.hpp index fadc1c837924..1a19daa0220a 100644 --- a/projects/miopen/src/include/miopen/rnn_util.hpp +++ b/projects/miopen/src/include/miopen/rnn_util.hpp @@ -131,6 +131,30 @@ void LSTMForwardHiddenStateUpdate(const Handle& handle, std::size_t activ_cell_offset, std::size_t hidden_offset); +void RNNFusedLSTMInferenceLoop(const Handle& handle, + Data_t workSpace, + ConstData_t w, + ConstData_t cx, + ConstData_t hx, + Data_t hy, + Data_t cy, + int seqLen, + int max_batch, + int hy_h, + int hy_stride, + int wei_len, + int hid_off, + int uni_stride, + int bi, + int ri, + int reverse, + std::size_t wei_shift_dir, + std::size_t hcx_offset, + bool use_cx, + bool use_hx, + bool use_hy, + bool use_cy); + void LSTMBackwardHiddenStateUpdate(const Handle& handle, miopenDataType_t rnn_data_type, bool is_seq_begin, diff --git a/projects/miopen/src/kernels/MIOpenRNNFusedInference.cpp b/projects/miopen/src/kernels/MIOpenRNNFusedInference.cpp new file mode 100644 index 000000000000..d23fb773c176 --- /dev/null +++ b/projects/miopen/src/kernels/MIOpenRNNFusedInference.cpp @@ -0,0 +1,175 @@ +// Copyright © Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +// +// Fused single-layer LSTM forward-INFERENCE recurrent loop (one cooperative +// launch per direction). Consumes MIOpen's existing workSpace layout: the +// batched input-projection GEMM + bias have already written the i,f,o,c gate +// pre-activations (from x) into workSpace. This kernel adds the recurrent +// contribution W_hh * h_{t-1} and applies the LSTM cell update, walking the +// whole sequence on-device with a grid-wide barrier per timestep — replacing +// the host-side per-timestep loop (~2 launches/step) that otherwise dominates +// at small batch. +// +// Supports UNIFORM batch (all sequences same length, max_batch rows per timestep). +// The recurrence is independent across batch, so a single grid barrier per timestep +// still suffices; work is mapped over (gate_row, batch) and (hidden, batch). +// +// GATED OFF BY DEFAULT in host code (MIOPEN_DEBUG_RNN_FUSED_INFERENCE). Only +// engaged for: single layer, LSTM, miopenRNNdefault, fp32, no dropout, packed, +// uniform batch across time. Falls back otherwise. +// +// workSpace layout. Within timestep t the max_batch rows are contiguous at +// row index (t*max_batch + b); each row is hy_stride floats. For batch b, +// direction ri (gate order matches LSTMForwardHiddenStateUpdate: i,f,o,c): +// rb = (t*max_batch + b)*hy_stride +// base = rb + ri*wei_len +// i/f/o gate = base + {0,1,2}*hy_h (sigmoid) ; c gate = base + 3*hy_h (tanh) +// hidden = rb + hid_off + ri*hy_h (hid_off = bi*hy_h*5) +// cell = rb + bi*wei_len + ri*hy_h +// hx/cx (initial hidden/cell) and hy/cy (final) = hcx_offset + b*hy_h +// hy_stride = hy_h*bi*6 ; wei_len = 4*hy_h ; W_hh packed in w at wei_shift_dir. + +#ifndef MIOPEN_HIP_RUNTIME_COMPILE +#include +#endif +#include + +namespace cg = cooperative_groups; + +#ifndef MIO_RNN_FINF_WAVE +#define MIO_RNN_FINF_WAVE 32 +#endif + +__device__ __forceinline__ float finf_sigmoid(float x) { return 1.0f / (1.0f + __expf(-x)); } + +// One cooperative grid; handles uniform batch (max_batch rows/timestep). +// workSpace : MIOpen RNN workspace (gates pre-filled with input projection+bias) +// w : packed weights; W_hh for this direction begins at wei_shift_dir +// hx/cx : initial hidden/cell (h_{-1}/c_{-1}); skipped when use_hx/use_cx==0 +// hy/cy : final hidden/cell out (h_{T-1}/c_{T-1}); written when use_hy/use_cy +// max_batch : rows per timestep (B); recurrence is independent across batch +// hy_h : hidden size H +// hy_stride : workspace row stride (= H*bi*6) +// wei_len : 4*H +// hid_off : bi*H*5 +// uni_stride: H (W_hh row stride) +// bi : 1 or 2 +// ri : direction index (0 fwd, 1 rev) for column placement +extern "C" __global__ void RNNFusedLSTMInfer(float* __restrict__ workSpace, + const float* __restrict__ w, + const float* __restrict__ cx, + const float* __restrict__ hx, + float* __restrict__ hy, + float* __restrict__ cy, + int seqLen, + int max_batch, + int hy_h, + int hy_stride, + int wei_len, + int hid_off, + int uni_stride, + int bi, + int ri, + int reverse, + long long wei_shift_dir, + long long hcx_offset, + int use_cx, + int use_hx, + int use_hy, + int use_cy) +{ + cg::grid_group grid = cg::this_grid(); + const int lane = threadIdx.x % MIO_RNN_FINF_WAVE; + const int wave_in_block = threadIdx.x / MIO_RNN_FINF_WAVE; + const int waves_per_block = blockDim.x / MIO_RNN_FINF_WAVE; + const int global_wave = blockIdx.x * waves_per_block + wave_in_block; + const int total_waves = gridDim.x * waves_per_block; + const int gtid = blockIdx.x * blockDim.x + threadIdx.x; + const int gnth = gridDim.x * blockDim.x; + const long long wh_dir = wei_shift_dir + (long long)ri * wei_len * uni_stride; + + // h/c carried in the workspace hidden/cell slots. Rows for timestep t are + // (t*max_batch + b) for b in [0,max_batch). At s==0 the previous hidden + // state is the caller-provided hx (h_{-1}); for s>0 it is the workspace + // hidden slot of the previously processed timestep. + for(int s = 0; s < seqLen; ++s) + { + const int t = reverse ? (seqLen - 1 - s) : s; + const int t_prev = reverse ? (t + 1) : (t - 1); + + // recurrent: for each (gate row g, batch b): gates[t,b,g] += W_hh[g,:].h_prev[b,:] + // One wave per (g,b) pair, coalesced lane reduction over k, shuffle-combine. + // h_prev is hx at s==0 (skipped when hx is absent, i.e. h_{-1}=0) and the + // prior timestep's workspace hidden otherwise. The guard is uniform across + // the grid (s and use_hx are the same for every thread) so the grid.sync() + // is reached collectively. + const bool has_prev = (s > 0) || use_hx; + if(has_prev) + { + const int ng = 4 * hy_h; + for(int gb = global_wave; gb < ng * max_batch; gb += total_waves) + { + const int g = gb % ng; + const int b = gb / ng; + const long long base = (long long)(t * max_batch + b) * hy_stride + ri * wei_len; + const float* hprev = + (s == 0) ? (hx + hcx_offset + (long long)b * hy_h) + : (workSpace + (long long)(t_prev * max_batch + b) * hy_stride + + hid_off + ri * hy_h); + const float* wh = w + wh_dir + (long long)g * uni_stride; + float acc = 0.f; + for(int k = lane; k < hy_h; k += MIO_RNN_FINF_WAVE) + acc += wh[k] * hprev[k]; + for(int o = MIO_RNN_FINF_WAVE / 2; o > 0; o >>= 1) + acc += __shfl_down(acc, o); + if(lane == 0) + workSpace[base + g] += acc; + } + grid.sync(); + } + + // cell update + activations; one thread per (hidden unit j, batch b). + // MIOpen gate order is i,f,o,c: +0=i +1=f +2=o (sigmoid), +3=c (tanh). + for(int jb = gtid; jb < hy_h * max_batch; jb += gnth) + { + const int j = jb % hy_h; + const int b = jb / hy_h; + const long long rb = (long long)(t * max_batch + b) * hy_stride; + const long long base = rb + ri * wei_len; + const long long hid = rb + hid_off + ri * hy_h; + const long long cell = rb + (long long)bi * wei_len + ri * hy_h; + + float ig = finf_sigmoid(workSpace[base + 0 * hy_h + j]); + float fg = finf_sigmoid(workSpace[base + 1 * hy_h + j]); + float og = finf_sigmoid(workSpace[base + 2 * hy_h + j]); + float cg_ = tanhf(workSpace[base + 3 * hy_h + j]); + float c_prev = 0.f; + if(s == 0) + { + if(use_cx) + c_prev = cx[hcx_offset + (long long)b * hy_h + j]; + } + else + { + c_prev = workSpace[(long long)(t_prev * max_batch + b) * hy_stride + + (long long)bi * wei_len + ri * hy_h + j]; + } + float c_new = fg * c_prev + ig * cg_; + float h_new = og * tanhf(c_new); + workSpace[cell + j] = c_new; + workSpace[hid + j] = h_new; + + // Final processed timestep for this direction carries the recurrent + // state out to hy/cy (offset matches the stock copy-out for uniform + // batch: hcx_offset + b*hy_h). + if(s == seqLen - 1) + { + if(use_hy) + hy[hcx_offset + (long long)b * hy_h + j] = h_new; + if(use_cy) + cy[hcx_offset + (long long)b * hy_h + j] = c_new; + } + } + grid.sync(); + } +} diff --git a/projects/miopen/src/ocl/rnnocl.cpp b/projects/miopen/src/ocl/rnnocl.cpp index 14ead06556ba..56cd4afd7c1d 100644 --- a/projects/miopen/src/ocl/rnnocl.cpp +++ b/projects/miopen/src/ocl/rnnocl.cpp @@ -40,6 +40,13 @@ MIOPEN_DECLARE_ENV_VAR_BOOL(MIOPEN_RNNFWD_EXP) MIOPEN_DECLARE_ENV_VAR_UINT64(MIOPEN_RNNFWD_MS_DISPATCH) MIOPEN_DECLARE_ENV_VAR_UINT64(MIOPEN_RNN_MS_STREAM_CNT) +// EXPERIMENTAL, OFF BY DEFAULT. Fused single-layer LSTM forward-inference path: +// runs the recurrent timestep loop in one cooperative kernel launch per +// direction instead of ~2 host-launched kernels per timestep. Only valid for +// single-layer LSTM, fp32, packed, batch=1, no dropout (Kokoro-class workloads). +// Enabling it globally would break other RNN shapes — keep it gated. +MIOPEN_DECLARE_ENV_VAR_BOOL(MIOPEN_DEBUG_RNN_FUSED_INFERENCE) + namespace miopen { namespace { @@ -70,6 +77,27 @@ bool RNNForwardMSIsFast(const int seqLen) return false; } +// EXPERIMENTAL fused single-layer LSTM inference gate. OFF unless the env var is +// set. Restricts to the case the fused cooperative kernel supports & was +// validated for: single layer, LSTM/default algo, no dropout, packed inputs, +// linear input mode, fp32. The kernel handles uniform batch (constant batch +// across timesteps); per-timestep uniformity of in_n is checked at dispatch. +bool RNNFusedInferenceIsSupported(const RNNDescriptor& desc, + miopenDataType_t dtype, + bool use_dropout) +{ +#if MIOPEN_BACKEND_HIP + return env::enabled(MIOPEN_DEBUG_RNN_FUSED_INFERENCE) && desc.rnnMode == miopenLSTM && + desc.algoMode == miopenRNNdefault && desc.nLayers == 1 && !use_dropout && + desc.inputMode != miopenRNNskip && dtype == miopenFloat; +#else + std::ignore = desc; + std::ignore = dtype; + std::ignore = use_dropout; + return false; +#endif +} + void checkGemmStatusAndLog(miopenStatus_t gemm_status) { if(gemm_status != miopenStatusSuccess) @@ -1815,6 +1843,49 @@ void RNNDescriptor::RNNForwardInferencePacked(const Handle& handle, } } + // EXPERIMENTAL fused single-layer LSTM inference (gated, batch=1). Runs the + // recurrent timestep loop in one cooperative launch per direction, then + // skips the host-side per-timestep loop below. The input-projection GEMM + + // bias above already populated the workSpace gates; the fused kernel only + // adds W_hh*h_{t-1} and the LSTM cell update. cx carry handled in-kernel. + // The fused kernel assumes UNIFORM batch (constant in_n across timesteps); + // packed variable-length sequences (decreasing in_n) fall back to the loop. + const bool uniform_batch = + std::all_of(in_n.begin(), in_n.end(), [&](int v) { return v == in_n.at(0); }); + if(uniform_batch && RNNFusedInferenceIsSupported(*this, wDesc.GetType(), false)) + { + const int max_batch = in_n.at(0); + const size_t wei_shift_dir = static_cast(in_h) * wei_stride + + static_cast(li) * (bi * hy_h + hy_h) * wei_stride; + for(int ri = 0; ri < bi; ri++) + { + const size_t hcx_off = hx_shift + static_cast(ri) * hy_n * hy_h; + RNNFusedLSTMInferenceLoop(handle, + workSpace, + w, + cx, + hx, + hy, + cy, + seqLen, + max_batch, + hy_h, + hy_stride, + wei_len, + static_cast(hid_off), + uni_stride, + bi, + ri, + /*reverse=*/ri, + wei_shift_dir, + hcx_off, + cx != nullptr, + hx != nullptr, + hy != nullptr, + cy != nullptr); + } + continue; + } // from hidden state int bacc = 0; int baccbi = batch_n; diff --git a/projects/miopen/src/rnn/rnn_util.cpp b/projects/miopen/src/rnn/rnn_util.cpp index ab4717770615..70c960d55ae9 100644 --- a/projects/miopen/src/rnn/rnn_util.cpp +++ b/projects/miopen/src/rnn/rnn_util.cpp @@ -790,4 +790,86 @@ void FillSeqTensorByPaddingMarker(const Handle& handle, } } +// Fused single-layer LSTM forward-inference recurrent loop (one cooperative +// launch per direction). Replaces the host-side per-timestep +// CallGemm + LSTMForwardHiddenStateUpdate loop. fp32, uniform batch, packed. +// Engaged only behind MIOPEN_DEBUG_RNN_FUSED_INFERENCE (see rnnocl.cpp gate). +void RNNFusedLSTMInferenceLoop(const Handle& handle, + Data_t workSpace, + ConstData_t w, + ConstData_t cx, + ConstData_t hx, + Data_t hy, + Data_t cy, + int seqLen, + int max_batch, + int hy_h, + int hy_stride, + int wei_len, + int hid_off, + int uni_stride, + int bi, + int ri, + int reverse, + std::size_t wei_shift_dir, + std::size_t hcx_offset, + bool use_cx, + bool use_hx, + bool use_hy, + bool use_cy) +{ + std::string program_name = "MIOpenRNNFusedInference.cpp"; + std::string kernel_name = "RNNFusedLSTMInfer"; + + // Cooperative launch: blocks must all be co-resident. Use one block per CU. + const std::size_t threads = 256; + const std::size_t blocks = handle.GetMaxComputeUnits(); + + std::string network_config = "rnnfusedinf-fp32-h" + std::to_string(hy_h) + "-b" + + std::to_string(blocks) + "-bi" + std::to_string(bi); + + const std::vector vld{threads, 1, 1}; + const std::vector vgd{threads * blocks, 1, 1}; + + auto existing = handle.GetKernelsImpl(kernel_name, network_config); + Kernel k; + if(!existing.empty()) + { + k = existing.front(); + } + else + { + const auto build_params = KernelBuildParameters{ + {"MIO_RNN_FINF_WAVE", static_cast(handle.GetWavefrontWidth())}}; + const std::string params = build_params.GenerateFor(kbp::HIP{}); + // AddKernel compiles & caches; fetch the Kernel handle back for a + // cooperative Run (AddKernel's returned invoke is non-cooperative). + handle.AddKernel(kernel_name, network_config, program_name, kernel_name, vld, vgd, params); + k = handle.GetKernelsImpl(kernel_name, network_config).front(); + } + + handle.Run(k, /*coop_launch=*/true)(workSpace, + w, + cx, + hx, + hy, + cy, + seqLen, + max_batch, + hy_h, + hy_stride, + wei_len, + hid_off, + uni_stride, + bi, + ri, + reverse, + static_cast(wei_shift_dir), + static_cast(hcx_offset), + static_cast(use_cx), + static_cast(use_hx), + static_cast(use_hy), + static_cast(use_cy)); +} + } // namespace miopen diff --git a/projects/miopen/test/gtest/lstm_fused_inference_test.cpp b/projects/miopen/test/gtest/lstm_fused_inference_test.cpp new file mode 100644 index 000000000000..ea173088acc5 --- /dev/null +++ b/projects/miopen/test/gtest/lstm_fused_inference_test.cpp @@ -0,0 +1,275 @@ +/******************************************************************************* + * + * MIT License + * + * Copyright (c) 2025 Advanced Micro Devices, Inc. + * + * 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. + * + *******************************************************************************/ + +// Correctness coverage for the experimental fused single-layer LSTM inference +// path, gated by MIOPEN_DEBUG_RNN_FUSED_INFERENCE (see src/ocl/rnnocl.cpp). +// +// The fused path only replaces the per-timestep hidden-state recurrence inside +// RNNForwardInferencePacked; the batched input projection (the GEMM that +// dominates throughput) is identical to the stock path. The natural oracle is +// therefore the stock path itself: run miopenRNNForwardInference twice on the +// same inputs in the same process -- once with the flag off (reference) and once +// with it on (fused) -- and require the two GPU outputs to agree to a tight +// tolerance. A bug in the fused recurrence (gate order, cell update, direction +// offset, batch indexing) shifts the fused output and trips the assertion, so +// the comparison has teeth that the float-epsilon CPU reference in lstm.hpp does +// not. The shared LSTM_test base is reused only for its tensor-setup fields. + +#include "lstm.hpp" +#include "gtest_common.hpp" +#include "rnn_util.hpp" +#include "workspace.hpp" +#include "verify.hpp" + +MIOPEN_LIB_ENV_VAR(MIOPEN_DEBUG_RNN_FUSED_INFERENCE) + +namespace { + +struct FusedInferCase +{ + int dirMode; // 0 = unidirectional, 1 = bidirectional + int batchSize; // uniform batch across the sequence + int seqLength; + int inVecLen; + int hiddenSize; + + friend std::ostream& operator<<(std::ostream& os, const FusedInferCase& tc) + { + return os << "dir-mode:" << tc.dirMode << " batch:" << tc.batchSize + << " seq:" << tc.seqLength << " in-vec:" << tc.inVecLen + << " hidden:" << tc.hiddenSize; + } +}; + +std::vector GetTestCases() +{ + return std::vector{ + // clang-format off + // dir batch seq invec hidden + // N=1 launch-latency path that the fusion targets (Kokoro decoder shape). + FusedInferCase{0, 1, 16, 128, 128}, + FusedInferCase{1, 1, 16, 128, 128}, + FusedInferCase{0, 1, 25, 256, 256}, + FusedInferCase{1, 1, 25, 256, 256}, + // Non-square in/hidden, still N=1. + FusedInferCase{0, 1, 8, 64, 128}, + FusedInferCase{1, 1, 8, 64, 128}, + // batch > 1 (the fused path also handles a uniform batch). + FusedInferCase{0, 4, 10, 64, 64}, + FusedInferCase{1, 4, 10, 64, 64}, + FusedInferCase{0, 16, 5, 32, 48}, + FusedInferCase{1, 16, 5, 32, 48}, + // single timestep edge case. + FusedInferCase{0, 2, 1, 32, 32}, + FusedInferCase{1, 2, 1, 32, 32} + // clang-format on + }; +} + +template +struct GPU_LSTM_fused_infer_test : LSTM_test, testing::TestWithParam +{ +protected: + void SetUp() override + { + const auto tc = GetParam(); + + this->dataType = miopen_type{}; + this->batchSize = tc.batchSize; + this->seqLength = tc.seqLength; + this->batchSeq = std::vector(tc.seqLength, tc.batchSize); + this->inVecLen = tc.inVecLen; + this->hiddenSize = tc.hiddenSize; + this->numLayers = 1; + this->inputMode = 0; // linear input (skip mode is unsupported by the fused path) + this->biasMode = 1; // with bias + this->dirMode = tc.dirMode; + this->algoMode = 0; // miopenRNNdefault + this->nohx = false; + this->nocx = false; + this->nohy = false; + this->nocy = false; + } + + // One forward-inference pass. Returns {output, hy, cy} read back to host. + std::tuple, std::vector, std::vector> + RunInference(miopenRNNDescriptor_t rnnDesc, + const std::vector& input, + const std::vector& hx, + const std::vector& cx, + const std::vector& weights) + { + auto&& handle = get_handle(); + + const auto inVecReal = (this->inputMode != 0) ? this->hiddenSize : this->inVecLen; + + std::vector inputCPPDescs; + std::vector inputDescs; + createTensorDescArray(inputCPPDescs, inputDescs, this->batchSeq, inVecReal, this->dataType); + + std::vector outputCPPDescs; + std::vector outputDescs; + createTensorDescArray(outputCPPDescs, + outputDescs, + this->batchSeq, + this->hiddenSize * ((this->dirMode != 0) ? 2 : 1), + this->dataType); + + size_t workspace_size = 0; + miopenGetRNNWorkspaceSize( + &handle, rnnDesc, this->seqLength, inputDescs.data(), &workspace_size); + Workspace wspace{workspace_size}; + + size_t out_sz = 0; + miopenGetRNNInputTensorSize(&handle, rnnDesc, this->seqLength, outputDescs.data(), &out_sz); + std::vector output(out_sz / sizeof(T)); + + std::vector hlens{ + this->numLayers * ((this->dirMode != 0) ? 2 : 1), this->batchSeq[0], this->hiddenSize}; + miopen::TensorDescriptor hiddenDesc(this->dataType, hlens); + + std::vector wlen{static_cast(weights.size())}; + miopen::TensorDescriptor weightDesc(this->dataType, wlen); + + std::vector hy(hx.size(), T(0)); + std::vector cy(cx.size(), T(0)); + + auto input_dev = handle.Write(input); + auto output_dev = handle.Write(output); + auto weights_dev = handle.Write(weights); + auto hx_dev = handle.Write(hx); + auto cx_dev = handle.Write(cx); + auto hy_dev = handle.Write(hy); + auto cy_dev = handle.Write(cy); + + miopenRNNForwardInference(&handle, + rnnDesc, + this->seqLength, + inputDescs.data(), + input_dev.get(), + &hiddenDesc, + hx_dev.get(), + &hiddenDesc, + cx_dev.get(), + &weightDesc, + weights_dev.get(), + outputDescs.data(), + output_dev.get(), + &hiddenDesc, + hy_dev.get(), + &hiddenDesc, + cy_dev.get(), + wspace.ptr(), + wspace.size()); + + return {handle.Read(output_dev, output.size()), + handle.Read(hy_dev, hy.size()), + handle.Read(cy_dev, cy.size())}; + } + + void RunFusedInference() + { + auto&& handle = get_handle(); + + RNNDescGuard rnnDesc; + DestroyInternalRnnDropoutDesc(rnnDesc); + miopenSetRNNDescriptor(rnnDesc, + this->hiddenSize, + this->numLayers, + miopenRNNInputMode_t(this->inputMode), + miopenRNNDirectionMode_t(this->dirMode), + miopenLSTM, + miopenRNNBiasMode_t(this->biasMode), + miopenRNNAlgo_t(this->algoMode), + this->dataType); + + const auto inVecReal = (this->inputMode != 0) ? this->hiddenSize : this->inVecLen; + + const std::size_t in_sz = getSuperTensorSize(this->batchSeq, + this->seqLength, + inVecReal, + this->hiddenSize, + 0, + this->dirMode != 0, + true, + 0); + std::vector input(in_sz); + for(auto& v : input) + v = prng::gen_descreet_unsigned(this->dataScale, 100); + + const std::size_t hx_sz = ((this->dirMode != 0) ? 2ULL : 1ULL) * this->hiddenSize * + this->batchSize * this->numLayers; + std::vector hx(hx_sz); + std::vector cx(hx_sz); + for(auto& v : hx) + v = prng::gen_descreet_unsigned(this->dataScale, 100); + for(auto& v : cx) + v = prng::gen_descreet_unsigned(this->dataScale, 100); + + std::vector inlens{this->batchSeq.at(0), inVecReal}; + auto firstInputDesc = miopen::TensorDescriptor(this->dataType, inlens); + size_t wei_bytes = 0; + miopenGetRNNParamsSize(&handle, rnnDesc, &firstInputDesc, &wei_bytes, this->dataType); + std::vector weights(wei_bytes / sizeof(T)); + for(auto& v : weights) + v = prng::gen_descreet_uniform_sign(this->dataScale, 100); + + // Reference: stock per-timestep path (flag explicitly off). + std::tuple, std::vector, std::vector> ref; + { + ScopedEnvironment fused_env(MIOPEN_DEBUG_RNN_FUSED_INFERENCE, false); + ref = RunInference(rnnDesc, input, hx, cx, weights); + } + + // Under test: fused cooperative-grid path (flag on). + std::tuple, std::vector, std::vector> fused; + { + ScopedEnvironment fused_env(MIOPEN_DEBUG_RNN_FUSED_INFERENCE, true); + fused = RunInference(rnnDesc, input, hx, cx, weights); + } + + DestroyInternalRnnDropoutDesc(rnnDesc); + + // Both paths share the input-projection GEMM and differ only in the + // recurrence arithmetic, so they should agree to a few ULPs of fp32 + // accumulation noise. rms_range normalizes by magnitude; 1e-5 is far + // below the ~1e-2 shift any gate-order/offset bug produces. + constexpr double kTol = 1e-5; + + ASSERT_EQ(std::get<0>(ref).size(), std::get<0>(fused).size()); + EXPECT_LT(miopen::rms_range(std::get<0>(ref), std::get<0>(fused)), kTol) << "output"; + EXPECT_LT(miopen::rms_range(std::get<1>(ref), std::get<1>(fused)), kTol) << "hy"; + EXPECT_LT(miopen::rms_range(std::get<2>(ref), std::get<2>(fused)), kTol) << "cy"; + } +}; + +} // namespace + +using GPU_LSTM_fused_infer_FP32 = GPU_LSTM_fused_infer_test; + +TEST_P(GPU_LSTM_fused_infer_FP32, FloatTest) { RunFusedInference(); } + +INSTANTIATE_TEST_SUITE_P(Full, GPU_LSTM_fused_infer_FP32, testing::ValuesIn(GetTestCases())); From f199760ce495d83caaa7d257d68e478974dfba66 Mon Sep 17 00:00:00 2001 From: ialex2004 Date: Thu, 30 Jul 2026 13:49:38 +0200 Subject: [PATCH 2/2] fix(miopen): address fused LSTM inference review feedback Keep the flagged path side-effect-free and cooperative-launch-safe, account for profiling, and strengthen the GPU test to prove fused dispatch while checking every API result. Co-Authored-By: Claude Opus 4.7 --- projects/miopen/src/hip/rnnocl.cpp | 81 +++++++------ .../test/gtest/lstm_fused_inference_test.cpp | 112 ++++++++++-------- 2 files changed, 108 insertions(+), 85 deletions(-) diff --git a/projects/miopen/src/hip/rnnocl.cpp b/projects/miopen/src/hip/rnnocl.cpp index 46161fef0722..a61e514ca971 100644 --- a/projects/miopen/src/hip/rnnocl.cpp +++ b/projects/miopen/src/hip/rnnocl.cpp @@ -82,18 +82,18 @@ bool RNNForwardMSIsFast(const int seqLen) // validated for: single layer, LSTM/default algo, no dropout, packed inputs, // linear input mode, fp32. The kernel handles uniform batch (constant batch // across timesteps); per-timestep uniformity of in_n is checked at dispatch. -bool RNNFusedInferenceIsSupported(const RNNDescriptor& desc, - miopenDataType_t dtype, - bool use_dropout) +bool RNNFusedInferenceIsSupported(const Handle& handle, + const RNNDescriptor& desc, + miopenDataType_t dtype) { #if MIOPEN_BACKEND_HIP - return env::enabled(MIOPEN_DEBUG_RNN_FUSED_INFERENCE) && desc.rnnMode == miopenLSTM && - desc.algoMode == miopenRNNdefault && desc.nLayers == 1 && !use_dropout && + return env::enabled(MIOPEN_DEBUG_RNN_FUSED_INFERENCE) && handle.CooperativeLaunchSupported() && + desc.rnnMode == miopenLSTM && desc.algoMode == miopenRNNdefault && desc.nLayers == 1 && desc.inputMode != miopenRNNskip && dtype == miopenFloat; #else + std::ignore = handle; std::ignore = desc; std::ignore = dtype; - std::ignore = use_dropout; return false; #endif } @@ -1851,41 +1851,46 @@ void RNNDescriptor::RNNForwardInferencePacked(const Handle& handle, // adds W_hh*h_{t-1} and the LSTM cell update. cx carry handled in-kernel. // The fused kernel assumes UNIFORM batch (constant in_n across timesteps); // packed variable-length sequences (decreasing in_n) fall back to the loop. - const bool uniform_batch = - std::all_of(in_n.begin(), in_n.end(), [&](int v) { return v == in_n.at(0); }); - if(uniform_batch && RNNFusedInferenceIsSupported(*this, wDesc.GetType(), false)) + if(RNNFusedInferenceIsSupported(handle, *this, wDesc.GetType())) { - const int max_batch = in_n.at(0); - const size_t wei_shift_dir = static_cast(in_h) * wei_stride + - static_cast(li) * (bi * hy_h + hy_h) * wei_stride; - for(int ri = 0; ri < bi; ri++) + const bool uniform_batch = + std::all_of(in_n.begin(), in_n.end(), [&](int v) { return v == in_n.at(0); }); + if(uniform_batch) { - const size_t hcx_off = hx_shift + static_cast(ri) * hy_n * hy_h; - RNNFusedLSTMInferenceLoop(handle, - workSpace, - w, - cx, - hx, - hy, - cy, - seqLen, - max_batch, - hy_h, - hy_stride, - wei_len, - static_cast(hid_off), - uni_stride, - bi, - ri, - /*reverse=*/ri, - wei_shift_dir, - hcx_off, - cx != nullptr, - hx != nullptr, - hy != nullptr, - cy != nullptr); + const int max_batch = in_n.at(0); + const size_t wei_shift_dir = + static_cast(in_h) * wei_stride + + static_cast(li) * (bi * hy_h + hy_h) * wei_stride; + for(int ri = 0; ri < bi; ri++) + { + const size_t hcx_off = hx_shift + static_cast(ri) * hy_n * hy_h; + RNNFusedLSTMInferenceLoop(handle, + workSpace, + w, + cx, + hx, + hy, + cy, + seqLen, + max_batch, + hy_h, + hy_stride, + wei_len, + static_cast(hid_off), + uni_stride, + bi, + ri, + /*reverse=*/ri, + wei_shift_dir, + hcx_off, + cx != nullptr, + hx != nullptr, + hy != nullptr, + cy != nullptr); + profileRNNkernels(handle, 1, ctime); + } + continue; } - continue; } // from hidden state int bacc = 0; diff --git a/projects/miopen/test/gtest/lstm_fused_inference_test.cpp b/projects/miopen/test/gtest/lstm_fused_inference_test.cpp index ea173088acc5..63c14bdcf8db 100644 --- a/projects/miopen/test/gtest/lstm_fused_inference_test.cpp +++ b/projects/miopen/test/gtest/lstm_fused_inference_test.cpp @@ -114,13 +114,14 @@ struct GPU_LSTM_fused_infer_test : LSTM_test, testing::TestWithParamnocy = false; } - // One forward-inference pass. Returns {output, hy, cy} read back to host. - std::tuple, std::vector, std::vector> - RunInference(miopenRNNDescriptor_t rnnDesc, - const std::vector& input, - const std::vector& hx, - const std::vector& cx, - const std::vector& weights) + using InferenceResult = std::tuple, std::vector, std::vector>; + + void RunInference(miopenRNNDescriptor_t rnnDesc, + const std::vector& input, + const std::vector& hx, + const std::vector& cx, + const std::vector& weights, + InferenceResult& result) { auto&& handle = get_handle(); @@ -139,12 +140,15 @@ struct GPU_LSTM_fused_infer_test : LSTM_test, testing::TestWithParamdataType); size_t workspace_size = 0; - miopenGetRNNWorkspaceSize( - &handle, rnnDesc, this->seqLength, inputDescs.data(), &workspace_size); + ASSERT_EQ(miopenStatusSuccess, + miopenGetRNNWorkspaceSize( + &handle, rnnDesc, this->seqLength, inputDescs.data(), &workspace_size)); Workspace wspace{workspace_size}; size_t out_sz = 0; - miopenGetRNNInputTensorSize(&handle, rnnDesc, this->seqLength, outputDescs.data(), &out_sz); + ASSERT_EQ(miopenStatusSuccess, + miopenGetRNNInputTensorSize( + &handle, rnnDesc, this->seqLength, outputDescs.data(), &out_sz)); std::vector output(out_sz / sizeof(T)); std::vector hlens{ @@ -165,46 +169,50 @@ struct GPU_LSTM_fused_infer_test : LSTM_test, testing::TestWithParamseqLength, - inputDescs.data(), - input_dev.get(), - &hiddenDesc, - hx_dev.get(), - &hiddenDesc, - cx_dev.get(), - &weightDesc, - weights_dev.get(), - outputDescs.data(), - output_dev.get(), - &hiddenDesc, - hy_dev.get(), - &hiddenDesc, - cy_dev.get(), - wspace.ptr(), - wspace.size()); - - return {handle.Read(output_dev, output.size()), - handle.Read(hy_dev, hy.size()), - handle.Read(cy_dev, cy.size())}; + ASSERT_EQ(miopenStatusSuccess, + miopenRNNForwardInference(&handle, + rnnDesc, + this->seqLength, + inputDescs.data(), + input_dev.get(), + &hiddenDesc, + hx_dev.get(), + &hiddenDesc, + cx_dev.get(), + &weightDesc, + weights_dev.get(), + outputDescs.data(), + output_dev.get(), + &hiddenDesc, + hy_dev.get(), + &hiddenDesc, + cy_dev.get(), + wspace.ptr(), + wspace.size())); + + result = {handle.Read(output_dev, output.size()), + handle.Read(hy_dev, hy.size()), + handle.Read(cy_dev, cy.size())}; } void RunFusedInference() { auto&& handle = get_handle(); + if(!handle.CooperativeLaunchSupported()) + GTEST_SKIP() << "Cooperative launch is not supported"; RNNDescGuard rnnDesc; DestroyInternalRnnDropoutDesc(rnnDesc); - miopenSetRNNDescriptor(rnnDesc, - this->hiddenSize, - this->numLayers, - miopenRNNInputMode_t(this->inputMode), - miopenRNNDirectionMode_t(this->dirMode), - miopenLSTM, - miopenRNNBiasMode_t(this->biasMode), - miopenRNNAlgo_t(this->algoMode), - this->dataType); + ASSERT_EQ(miopenStatusSuccess, + miopenSetRNNDescriptor(rnnDesc, + this->hiddenSize, + this->numLayers, + miopenRNNInputMode_t(this->inputMode), + miopenRNNDirectionMode_t(this->dirMode), + miopenLSTM, + miopenRNNBiasMode_t(this->biasMode), + miopenRNNAlgo_t(this->algoMode), + this->dataType)); const auto inVecReal = (this->inputMode != 0) ? this->hiddenSize : this->inVecLen; @@ -232,24 +240,34 @@ struct GPU_LSTM_fused_infer_test : LSTM_test, testing::TestWithParam inlens{this->batchSeq.at(0), inVecReal}; auto firstInputDesc = miopen::TensorDescriptor(this->dataType, inlens); size_t wei_bytes = 0; - miopenGetRNNParamsSize(&handle, rnnDesc, &firstInputDesc, &wei_bytes, this->dataType); + ASSERT_EQ( + miopenStatusSuccess, + miopenGetRNNParamsSize(&handle, rnnDesc, &firstInputDesc, &wei_bytes, this->dataType)); std::vector weights(wei_bytes / sizeof(T)); for(auto& v : weights) v = prng::gen_descreet_uniform_sign(this->dataScale, 100); + const auto fused_kernel_name = "RNNFusedLSTMInfer"; + const auto network_config = "rnnfusedinf-fp32-h" + std::to_string(this->hiddenSize) + "-b" + + std::to_string(handle.GetMaxComputeUnits()) + "-bi" + + std::to_string(this->dirMode != 0 ? 2 : 1); + handle.ClearKernels(fused_kernel_name, network_config); + // Reference: stock per-timestep path (flag explicitly off). - std::tuple, std::vector, std::vector> ref; + InferenceResult ref; { ScopedEnvironment fused_env(MIOPEN_DEBUG_RNN_FUSED_INFERENCE, false); - ref = RunInference(rnnDesc, input, hx, cx, weights); + RunInference(rnnDesc, input, hx, cx, weights, ref); } + ASSERT_TRUE(handle.GetKernelsImpl(fused_kernel_name, network_config).empty()); // Under test: fused cooperative-grid path (flag on). - std::tuple, std::vector, std::vector> fused; + InferenceResult fused; { ScopedEnvironment fused_env(MIOPEN_DEBUG_RNN_FUSED_INFERENCE, true); - fused = RunInference(rnnDesc, input, hx, cx, weights); + RunInference(rnnDesc, input, hx, cx, weights, fused); } + ASSERT_FALSE(handle.GetKernelsImpl(fused_kernel_name, network_config).empty()); DestroyInternalRnnDropoutDesc(rnnDesc);