ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-04-20 06:49:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

[CK Tile] gemm splitk two stage (#2697)

Browse Source 
* Fix a typo

* Use std::variant to call run_gemm_example_with_layouts with the available layout variant combinations

* Use a unified run_gemm_example_prec_type for basic gemm and universal gemm

* Factor out run_gemm_example_prec_type

* Refactor argument parsing in gemm_splitk_two_stage_reduce.cpp

* Parse arguments outside of create_args

* Move the gemm operators to separate structs to facilitate their reuse

* Move the invokers to separate files to facilitate their reuse

* Rename the invoker files for consistency with the examples that use them

* Add fp32 support to the elementwise examples, and produce an error message for unsupported types

* Get rid of four unused variables

* Make two variables const

* Add support for different input-output type combinations in elementwise examples

* Test support for different input and output types in elementwise examples

* Add support for different operations in the elementwise unary tests

* Add support for UnaryConvert in the elementwise unary tests

* Add support for bf16 in elementwise examples, excluding unsupported type combinations

* Make some operator parameters const in ElementWiseKernel

* Remove some unnecessary include statements

* Implement a two-stage GEMM that does a type conversion in the second stage using the elementwise kernel

* Clear workspace instead of output when flushing the cache in SplitKTwoStageInvoker::gemm

* Fix formatting issues reported by clang

* Add back CK_TILE_USE_WMMA related changes

* Use the right prec type for bf16 in the universal GEMM and two stage split K examples

* Add some brackets

* Add some brackets

* Separate the clearing of the GEMM output memory from the cache flushing in the universal GEMM example

* Separate the clearing of the GEMM output memory from the cache flushing in the split K two stage example

* Fix formatting

* No need to call SetZero on ws_m_n_dev_buf here, as clear_gemm_output now does this as part of the kernel preprocessing

* Add fp16 data type to splitk two stage example

* Add preprocessing with optional cache flushing and clearing of output for k_batch > 1 to the basic GEMM example
This commit is contained in:
SamiAario-AMD
2025-09-04 14:33:44 +03:00
committed by GitHub
parent e2d28a92af
commit 1acd8e041c
21 changed files with 1245 additions and 782 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
include/ck_tile/ops/elementwise/kernel/elementwise_kernel.hpp
View File

@@ -23,9 +23,9 @@ struct ElementWiseKernel
static constexpr index_t kBlockSize = Problem::BlockShape::kBlockSize;
template <typename... XDataType, typename Dims>
CK_TILE_DEVICE void operator()(Dims lens,
Dims input_strides,
Dims output_strides,
CK_TILE_DEVICE void operator()(const Dims lens,
const Dims input_strides,
const Dims output_strides,
const tuple<XDataType...>& input_tensors,
YDataType* p_y) const
{