mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-07 15:47:05 +00:00
added a 1d vector elementwise example
This commit is contained in:
committed by
Philip Maybank
parent
cf34ed5f3d
commit
9edef9e351
22
example/ck_tile/99_toy_example/00_vector_add/CMakeLists.txt
Normal file
22
example/ck_tile/99_toy_example/00_vector_add/CMakeLists.txt
Normal file
@@ -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)
|
||||
@@ -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 <thread>
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
template <typename XDataType, typename YDataType>
|
||||
CK_TILE_HOST void
|
||||
reference_vector_add(const HostTensor<XDataType>& xa_m_n, const HostTensor<XDataType>& xb_m_n, HostTensor<YDataType>& 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<YDataType>(xa_m_n(m, n)) + ck_tile::type_convert<YDataType>(xb_m_n(m, n));
|
||||
}
|
||||
};
|
||||
|
||||
make_ParallelTensorFunctor(f, y_m_n.mDesc.get_lengths()[0])(std::thread::hardware_concurrency());
|
||||
}
|
||||
|
||||
} // namespace ck_tile
|
||||
131
example/ck_tile/99_toy_example/00_vector_add/vector_add.cpp
Normal file
131
example/ck_tile/99_toy_example/00_vector_add/vector_add.cpp
Normal file
@@ -0,0 +1,131 @@
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "reference_vector_add.hpp"
|
||||
#include "vector_add.hpp"
|
||||
#include <cstring>
|
||||
|
||||
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 <typename DataType>
|
||||
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<XDataType> x_host_a({m});
|
||||
ck_tile::HostTensor<XDataType> x_host_b({m});
|
||||
|
||||
ck_tile::HostTensor<YDataType> y_host_ref({m});
|
||||
ck_tile::HostTensor<YDataType> y_host_dev({m});
|
||||
|
||||
ck_tile::FillUniformDistribution<XDataType>{-5.f, 5.f}(x_host_a);
|
||||
ck_tile::FillUniformDistribution<XDataType>{-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<BlockWarps, BlockTile, WarpTile, Vector>;
|
||||
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<XDataType, ComputeDataType, YDataType, Shape>;
|
||||
|
||||
using Kernel = ck_tile::MultiplyVectorKernel<Problem>;
|
||||
|
||||
float ave_time = launch_kernel(ck_tile::stream_config{nullptr, true, 0, warmup, repeat},
|
||||
ck_tile::make_kernel<kBlockSize, kBlockPerCu>(
|
||||
Kernel{},
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0,
|
||||
static_cast<XDataType*>(x_buf_a.GetDeviceBuffer()),
|
||||
static_cast<XDataType*>(x_buf_b.GetDeviceBuffer()),
|
||||
static_cast<YDataType*>(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<XDataType, YDataType>(
|
||||
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<ck_tile::half_t>(arg_parser) ? 0 : -2;
|
||||
}
|
||||
}
|
||||
131
example/ck_tile/99_toy_example/00_vector_add/vector_add.hpp
Normal file
131
example/ck_tile/99_toy_example/00_vector_add/vector_add.hpp
Normal file
@@ -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 BlockWarps, // num warps along seq<M, N>
|
||||
typename BlockTile, // block size, seq<M, N>
|
||||
typename WarpTile, // warp size, seq<M, N>
|
||||
typename Vector> // contiguous pixels(vector size) along seq<M, N>
|
||||
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 <typename XDataType_,
|
||||
typename ComputeDataType_,
|
||||
typename YDataType_,
|
||||
typename BlockShape_>
|
||||
struct MultiplyVectorProblem
|
||||
{
|
||||
using XDataType = remove_cvref_t<XDataType_>;
|
||||
using ComputeDataType = remove_cvref_t<ComputeDataType_>;
|
||||
using YDataType = remove_cvref_t<YDataType_>;
|
||||
using BlockShape = remove_cvref_t<BlockShape_>;
|
||||
};
|
||||
|
||||
struct AddDefaultPolicy
|
||||
{
|
||||
template <typename Problem>
|
||||
CK_TILE_DEVICE static constexpr auto MakeXBlockTileDistribution()
|
||||
{
|
||||
using S = typename Problem::BlockShape;
|
||||
return make_static_tile_distribution(
|
||||
tile_distribution_encoding<
|
||||
sequence<>,
|
||||
tuple<sequence<S::Repeat_M, S::WarpPerBlock_M, S::ThreadPerWarp_M, S::Vector_M>>,
|
||||
tuple<sequence<1>, sequence<1>>,
|
||||
tuple<sequence<1>, sequence<2>>,
|
||||
sequence<1, 1>,
|
||||
sequence<0, 3>>{});
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Problem_, typename Policy_ = AddDefaultPolicy>
|
||||
struct MultiplyVectorKernel
|
||||
{
|
||||
using Problem = ck_tile::remove_cvref_t<Problem_>;
|
||||
using Policy = ck_tile::remove_cvref_t<Policy_>;
|
||||
|
||||
using XDataType = ck_tile::remove_cvref_t<typename Problem::XDataType>;
|
||||
using ComputeDataType = ck_tile::remove_cvref_t<typename Problem::ComputeDataType>;
|
||||
using YDataType = ck_tile::remove_cvref_t<typename Problem::YDataType>;
|
||||
|
||||
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<address_space_enum::global>(
|
||||
p_x_a, make_tuple(M), make_tuple(1), number<S::Vector_M>{});
|
||||
|
||||
const auto x_m_n_b = make_naive_tensor_view<address_space_enum::global>(
|
||||
p_x_b, make_tuple(M), make_tuple(1), number<S::Vector_M>{});
|
||||
|
||||
const auto y_m_n = make_naive_tensor_view<address_space_enum::global>(
|
||||
p_y, make_tuple(M), make_tuple(1), number<S::Vector_M>{});
|
||||
|
||||
const auto iM = get_block_id() * S::Block_M;
|
||||
|
||||
auto x_window_a = make_tile_window(x_m_n_a,
|
||||
make_tuple(number<S::Block_M>{}),
|
||||
{iM},
|
||||
Policy::template MakeXBlockTileDistribution<Problem>());
|
||||
|
||||
auto x_window_b = make_tile_window(x_m_n_b,
|
||||
make_tuple(number<S::Block_M>{}),
|
||||
{iM},
|
||||
Policy::template MakeXBlockTileDistribution<Problem>());
|
||||
|
||||
auto y_window = make_tile_window(y_m_n,
|
||||
make_tuple(number<S::Block_M>{}),
|
||||
{iM},
|
||||
Policy::template MakeXBlockTileDistribution<Problem>());
|
||||
|
||||
// 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<ComputeDataType>(xa[tile_idx]);
|
||||
const auto b_val = type_convert<ComputeDataType>(xb[tile_idx]);
|
||||
y_compute(tile_idx) = a_val + b_val;
|
||||
|
||||
});
|
||||
|
||||
|
||||
// Store results
|
||||
store_tile(y_window, cast_tile<YDataType>(y_compute)); // store the result back to global tensor view
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace ck_tile
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user