From 9edef9e351a414b7b4bcd2c8e3ccb71882061850 Mon Sep 17 00:00:00 2001 From: AviralGoelAMD Date: Tue, 15 Apr 2025 14:12:04 +0000 Subject: [PATCH] added a 1d vector elementwise example --- .../00_vector_add/CMakeLists.txt | 22 +++ .../00_vector_add/reference_vector_add.hpp | 28 ++++ .../00_vector_add/vector_add.cpp | 131 ++++++++++++++++++ .../00_vector_add/vector_add.hpp | 131 ++++++++++++++++++ example/ck_tile/99_toy_example/CMakeLists.txt | 1 + 5 files changed, 313 insertions(+) create mode 100644 example/ck_tile/99_toy_example/00_vector_add/CMakeLists.txt create mode 100644 example/ck_tile/99_toy_example/00_vector_add/reference_vector_add.hpp create mode 100644 example/ck_tile/99_toy_example/00_vector_add/vector_add.cpp create mode 100644 example/ck_tile/99_toy_example/00_vector_add/vector_add.hpp diff --git a/example/ck_tile/99_toy_example/00_vector_add/CMakeLists.txt b/example/ck_tile/99_toy_example/00_vector_add/CMakeLists.txt new file mode 100644 index 0000000000..b6a7886d12 --- /dev/null +++ b/example/ck_tile/99_toy_example/00_vector_add/CMakeLists.txt @@ -0,0 +1,22 @@ +set(EXAMPLE_REDUCE "vector_add") +# not using add_example_executable() to add this target, since we don't want this to have +# to be included in "make all/install/check" +message("adding example ${EXAMPLE_REDUCE}") + +add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL vector_add.cpp) +target_include_directories(${EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR}) +set(EXAMPLE_REDUCE_COMPILE_OPTIONS) + +# generate assembly +# list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS -v --save-temps -Wno-gnu-line-marker) + +# NOTE: we turn off undefined-func-template to let source compile without explicit declare function specializations +list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal) + +target_compile_options(${EXAMPLE_REDUCE} PRIVATE ${EXAMPLE_REDUCE_COMPILE_OPTIONS}) + +# TODO: we have to turn off this global prop, otherwise the progress bar generated +# by cmake will print too many files, execvp: /bin/sh: Argument list too long +# however, this property may affect global +# TODO: consider codegen a makefile by us +set_property(GLOBAL PROPERTY RULE_MESSAGES OFF) diff --git a/example/ck_tile/99_toy_example/00_vector_add/reference_vector_add.hpp b/example/ck_tile/99_toy_example/00_vector_add/reference_vector_add.hpp new file mode 100644 index 0000000000..72b822c8ed --- /dev/null +++ b/example/ck_tile/99_toy_example/00_vector_add/reference_vector_add.hpp @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/host/host_tensor.hpp" +#include + +namespace ck_tile { + +template +CK_TILE_HOST void +reference_vector_add(const HostTensor& xa_m_n, const HostTensor& xb_m_n, HostTensor& y_m_n) +{ + auto f = [&](auto m) { + const int N = 1; + + for(int n = 0; n < N; ++n) + { + y_m_n(m, n) = ck_tile::type_convert(xa_m_n(m, n)) + ck_tile::type_convert(xb_m_n(m, n)); + } + }; + + make_ParallelTensorFunctor(f, y_m_n.mDesc.get_lengths()[0])(std::thread::hardware_concurrency()); +} + +} // namespace ck_tile diff --git a/example/ck_tile/99_toy_example/00_vector_add/vector_add.cpp b/example/ck_tile/99_toy_example/00_vector_add/vector_add.cpp new file mode 100644 index 0000000000..ab51a5393b --- /dev/null +++ b/example/ck_tile/99_toy_example/00_vector_add/vector_add.cpp @@ -0,0 +1,131 @@ +#include "ck_tile/host.hpp" +#include "reference_vector_add.hpp" +#include "vector_add.hpp" +#include + +auto create_args(int argc, char* argv[]) +{ + ck_tile::ArgParser arg_parser; + arg_parser.insert("m", "256000000", "m dimension") + .insert("v", "1", "cpu validation or not") + .insert("prec", "fp16", "precision") + .insert("warmup", "5", "cold iter") + .insert("repeat", "20", "hot iter"); + + bool result = arg_parser.parse(argc, argv); + return std::make_tuple(result, arg_parser); +} + +template +bool run(const ck_tile::ArgParser& arg_parser) +{ + using XDataType = DataType; + using ComputeDataType = float; + using YDataType = DataType; + + ck_tile::index_t m = arg_parser.get_int("m"); + int do_validation = arg_parser.get_int("v"); + int warmup = arg_parser.get_int("warmup"); + int repeat = arg_parser.get_int("repeat"); + + ck_tile::HostTensor x_host_a({m}); + ck_tile::HostTensor x_host_b({m}); + + ck_tile::HostTensor y_host_ref({m}); + ck_tile::HostTensor y_host_dev({m}); + + ck_tile::FillUniformDistribution{-5.f, 5.f}(x_host_a); + ck_tile::FillUniformDistribution{-5.f, 5.f}(x_host_b); + + ck_tile::DeviceMem x_buf_a(x_host_a.get_element_space_size_in_bytes()); + ck_tile::DeviceMem x_buf_b(x_host_b.get_element_space_size_in_bytes()); + ck_tile::DeviceMem y_buf(y_host_dev.get_element_space_size_in_bytes()); + + x_buf_a.ToDevice(x_host_a.data()); + x_buf_b.ToDevice(x_host_b.data()); + + // using BlockTile = ck_tile::sequence<8192>; + // using BlockWarps = ck_tile::sequence<4>; + // using WarpTile = ck_tile::sequence<512>; // 8 * 64 = 512 + // using Vector = ck_tile::sequence<8>; // 8 * 16 = 128 bytes + + + // constexpr ck_tile::index_t kBlockSize = 256; + // constexpr ck_tile::index_t kBlockPerCu = 1; + + + using BlockTile = ck_tile::sequence<8192>; + using BlockWarps = ck_tile::sequence<8>; + using WarpTile = ck_tile::sequence<64>; + using Vector = ck_tile::sequence<1>; + + constexpr ck_tile::index_t kBlockSize = 512; + constexpr ck_tile::index_t kBlockPerCu = 1; + + + ck_tile::index_t kGridSize = (m / BlockTile::at(ck_tile::number<0>{})); + std::cout << "block x-size = " << BlockTile::at(ck_tile::number<0>{}) << std::endl; + std::cout << "grid size " << kGridSize << std::endl; + + + + using Shape = ck_tile::MultiplyVector; + std::cout << "Problem Shape:: M = " << m << std::endl; + std::cout << "BlockTile: " << BlockTile::at(ck_tile::number<0>{}) << std::endl; + std::cout << "Number of Blocks in Grid: " << m / BlockTile::at(ck_tile::number<0>{}) << std::endl; + std::cout << "BlockWarps: " << BlockWarps::at(ck_tile::number<0>{}) << std::endl; + std::cout << "WarpTile: " << WarpTile::at(ck_tile::number<0>{}) << std::endl; + std::cout << "Vector: " << Vector::at(ck_tile::number<0>{}) << std::endl; + std::cout << "Repeat: " << Shape::Repeat_M << std::endl; + std::cout << "Threads per Block: " << kBlockSize << std::endl; + std::cout << "ThreadBlocks per CU: " << kBlockPerCu << std::endl; + using Problem = + ck_tile::MultiplyVectorProblem; + + using Kernel = ck_tile::MultiplyVectorKernel; + + float ave_time = launch_kernel(ck_tile::stream_config{nullptr, true, 0, warmup, repeat}, + ck_tile::make_kernel( + Kernel{}, + kGridSize, + kBlockSize, + 0, + static_cast(x_buf_a.GetDeviceBuffer()), + static_cast(x_buf_b.GetDeviceBuffer()), + static_cast(y_buf.GetDeviceBuffer()), + m)); + + std::size_t num_btype = sizeof(XDataType) * m + sizeof(YDataType) * m; + + float gb_per_sec = num_btype / 1.E6 / ave_time; + + std::cout << "Perf: " << ave_time << " ms, " << gb_per_sec << " GB/s" << std::endl; + + bool pass = true; + + if(do_validation) + { + ck_tile::reference_vector_add( + x_host_a, x_host_b, y_host_ref); + y_buf.FromDevice(y_host_dev.mData.data()); + pass = ck_tile::check_err(y_host_dev, y_host_ref); + + std::cout << "valid:" << (pass ? "y" : "n") << std::flush << std::endl; + } + + return pass; +} + +int main(int argc, char* argv[]) +{ + auto [result, arg_parser] = create_args(argc, argv); + if(!result) + return -1; + + const std::string data_type = arg_parser.get_str("prec"); + + if(data_type == "fp16") + { + return run(arg_parser) ? 0 : -2; + } +} diff --git a/example/ck_tile/99_toy_example/00_vector_add/vector_add.hpp b/example/ck_tile/99_toy_example/00_vector_add/vector_add.hpp new file mode 100644 index 0000000000..12c97b41ae --- /dev/null +++ b/example/ck_tile/99_toy_example/00_vector_add/vector_add.hpp @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/common.hpp" + +namespace ck_tile { + +template + typename BlockTile, // block size, seq + typename WarpTile, // warp size, seq + typename Vector> // contiguous pixels(vector size) along seq +struct MultiplyVector +{ + static constexpr index_t Block_M = BlockTile::at(number<0>{}); + //static constexpr index_t Block_N = BlockTile::at(number<1>{}); + + static constexpr index_t Warp_M = WarpTile::at(number<0>{}); + //static constexpr index_t Warp_N = WarpTile::at(number<1>{}); + + static constexpr index_t Vector_M = Vector::at(number<0>{}); + //static constexpr index_t Vector_N = Vector::at(number<1>{}); + + static constexpr index_t WarpPerBlock_M = BlockWarps::at(number<0>{}); + //static constexpr index_t WarpPerBlock_N = BlockWarps::at(number<1>{}); + + static constexpr index_t ThreadPerWarp_M = Warp_M / Vector_M; + //static constexpr index_t ThreadPerWarp_N = Warp_N / Vector_N; + + static constexpr index_t Repeat_M = Block_M / (WarpPerBlock_M * Warp_M); + //static constexpr index_t Repeat_N = Block_N / (WarpPerBlock_N * Warp_N); + + static constexpr index_t BlockSize = + warpSize * reduce_on_sequence(BlockWarps{}, multiplies{}, number<1>{}); +}; + +template +struct MultiplyVectorProblem +{ + using XDataType = remove_cvref_t; + using ComputeDataType = remove_cvref_t; + using YDataType = remove_cvref_t; + using BlockShape = remove_cvref_t; +}; + +struct AddDefaultPolicy +{ + template + CK_TILE_DEVICE static constexpr auto MakeXBlockTileDistribution() + { + using S = typename Problem::BlockShape; + return make_static_tile_distribution( + tile_distribution_encoding< + sequence<>, + tuple>, + tuple, sequence<1>>, + tuple, sequence<2>>, + sequence<1, 1>, + sequence<0, 3>>{}); + } +}; + +template +struct MultiplyVectorKernel +{ + using Problem = ck_tile::remove_cvref_t; + using Policy = ck_tile::remove_cvref_t; + + using XDataType = ck_tile::remove_cvref_t; + using ComputeDataType = ck_tile::remove_cvref_t; + using YDataType = ck_tile::remove_cvref_t; + + CK_TILE_DEVICE void operator()(const XDataType* p_x_a, const XDataType* p_x_b, YDataType* p_y, index_t M) const + { + using S = typename Problem::BlockShape; + + const auto x_m_n_a = make_naive_tensor_view( + p_x_a, make_tuple(M), make_tuple(1), number{}); + + const auto x_m_n_b = make_naive_tensor_view( + p_x_b, make_tuple(M), make_tuple(1), number{}); + + const auto y_m_n = make_naive_tensor_view( + p_y, make_tuple(M), make_tuple(1), number{}); + + const auto iM = get_block_id() * S::Block_M; + + auto x_window_a = make_tile_window(x_m_n_a, + make_tuple(number{}), + {iM}, + Policy::template MakeXBlockTileDistribution()); + + auto x_window_b = make_tile_window(x_m_n_b, + make_tuple(number{}), + {iM}, + Policy::template MakeXBlockTileDistribution()); + + auto y_window = make_tile_window(y_m_n, + make_tuple(number{}), + {iM}, + Policy::template MakeXBlockTileDistribution()); + + // Load tile data + const auto xa = load_tile(x_window_a); // load tile data from global tensor view, load from where? what? how many? logical memory layout? all are defined in x_window_a + const auto xb = load_tile(x_window_b); + auto y_compute = load_tile(y_window); + + + + // Process the vector multiplication + constexpr auto spans = decltype(xa)::get_distributed_spans(); // shape of the tile + sweep_tile_span(spans[number<0>{}], [&](auto idx) { // iterate over the tile // idx+=4 + const auto tile_idx = make_tuple(idx); + const auto a_val = type_convert(xa[tile_idx]); + const auto b_val = type_convert(xb[tile_idx]); + y_compute(tile_idx) = a_val + b_val; + + }); + + + // Store results + store_tile(y_window, cast_tile(y_compute)); // store the result back to global tensor view + } +}; + +} // namespace ck_tile diff --git a/example/ck_tile/99_toy_example/CMakeLists.txt b/example/ck_tile/99_toy_example/CMakeLists.txt index cadbb1c06f..49c3677fea 100644 --- a/example/ck_tile/99_toy_example/CMakeLists.txt +++ b/example/ck_tile/99_toy_example/CMakeLists.txt @@ -2,6 +2,7 @@ include_directories(AFTER ${CMAKE_CURRENT_LIST_DIR} ) +add_subdirectory(00_vector_add) add_subdirectory(01_add) add_subdirectory(02_gemm) add_subdirectory(03_flash_attention_fwd)