mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-14 11:07:44 +00:00
359f664b25
[CK_TILE] Implement RTC API for a subset of FMHA functionality for MGX (#6086) ## Motivation Introduce a wrapper for the FmhaFwdKernel, for use in real time compilation in MIGraphX. ## Technical Details The intent of the API is to provide multiple instances of the FmhaFwdKernelWrapper, suitable for a particular problem definition. At the moment the wrapper only supports bias and causal masking, feature expansion will come in a future pr. The usage pattern is, in short: 1. Define fmha_fwd::Problem (input dimensions, data type, etc) 2. Fetch Solutions for target architecture (currently only gfx942) based on Problem. The solutions contain a map of template -> template parameter and can be converted to a string representing the full instantiation of FmhFwdKernelWrapper e.g. `ck_tile::FmhaFwdWrapper<ck_tile::fp16_t, 128, 64, 16, 32, 32, 32, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, false, true, false, true, true, true, true, ck_tile::FmhaPipelineTag::QR>` 3. The instance can then be used in an RTC kernel. The kernel needs to: * Construct a Descriptor (containing descriptions of all input tensors) * Call IsValid() on the descriptor to check if the instance is applicable. Note that this is constexpr by design so that it can fail the kernel compilation as a signal that the kernel is not applicable. * Pass the descriptor and input pointers to the wrapper Run method. A more detailed example of usage can be found in codegen/test/fmh_fwd.cpp Beside work on creating the wrapper and the supporting API, the PR also contains some changes necessary to enable compilation with HIPRTC. The contents of the CK tile headers are embedded in a binary file which is used to pass the header files as strings to HIPRTC. Many of the ck tile headers contain host only code which leads to compilation failures. ck_tile_headers_preprocessor goes through the embedded headers and removes the bodies of host only functions, thereby eliminating the compilation failures. ## Test Plan <!-- Explain any relevant testing done to verify this PR. --> ## Test Result <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
93 lines
1.8 KiB
C++
93 lines
1.8 KiB
C++
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#pragma once
|
|
|
|
#include <cstdlib>
|
|
#include <vector>
|
|
#include <string>
|
|
#include "ck/host/types.hpp"
|
|
#include "ck/host/device_fmha_fwd/problem.hpp"
|
|
|
|
namespace ck {
|
|
namespace host {
|
|
namespace device_fmha_fwd {
|
|
|
|
// Derived from fmha_fwd.py FmhaFwdTileSize.
|
|
struct TileConfig
|
|
{
|
|
// Block tile
|
|
std::size_t bm0;
|
|
std::size_t bn0;
|
|
std::size_t bk0;
|
|
std::size_t bn1;
|
|
std::size_t bk1;
|
|
std::size_t bk0max;
|
|
|
|
// Gemm0 block warps
|
|
std::size_t rm0;
|
|
std::size_t rn0;
|
|
std::size_t rk0;
|
|
|
|
// Gemm1 block warps
|
|
std::size_t rm1;
|
|
std::size_t rn1;
|
|
std::size_t rk1;
|
|
|
|
// Gemm0 warp tile
|
|
std::size_t wm0;
|
|
std::size_t wn0;
|
|
std::size_t wk0;
|
|
|
|
// Gemm1 warp tile
|
|
std::size_t wm1;
|
|
std::size_t wn1;
|
|
std::size_t wk1;
|
|
};
|
|
|
|
struct PipelineConfig
|
|
{
|
|
std::string name;
|
|
bool pad_m;
|
|
bool pad_n;
|
|
bool pad_k;
|
|
bool pad_o;
|
|
};
|
|
|
|
struct Operation
|
|
{
|
|
TileConfig tile = {};
|
|
|
|
std::string pipeline = "qr_async";
|
|
|
|
bool is_causal = false;
|
|
bool is_v_rowmajor = true;
|
|
bool has_bias = false;
|
|
DataType dtype = DataType::Half;
|
|
|
|
bool pad_m = true; // pad seqlen_q
|
|
bool pad_n = true; // pad seqlen_k
|
|
bool pad_k = true; // pad hdim_q
|
|
bool pad_o = true; // pad hdim_v
|
|
|
|
static std::vector<Operation> CreateOperations(const Problem& prob, const std::string& arch);
|
|
|
|
Solution ToSolution() const;
|
|
};
|
|
|
|
struct HdimBucketResult
|
|
{
|
|
std::size_t bucket_hdim = 0;
|
|
std::size_t bucket_hdim_v = 0;
|
|
std::vector<TileConfig> tiles;
|
|
};
|
|
|
|
HdimBucketResult
|
|
GetTileConfigsForHdim(const std::string& arch, DataType dtype, std::size_t K, std::size_t O);
|
|
|
|
bool IsSupportedArch(const std::string& arch);
|
|
|
|
} // namespace device_fmha_fwd
|
|
} // namespace host
|
|
} // namespace ck
|