mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-13 17:55:48 +00:00
This commit is contained in:
Ding, Yi
180
example/ck_tile/21_elementwise/elementwise_example_transpose.cpp
Normal file
180
example/ck_tile/21_elementwise/elementwise_example_transpose.cpp
Normal file
@@ -0,0 +1,180 @@
|
||||
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "ck_tile/ops/elementwise.hpp"
|
||||
#include "ck_tile/host/reference/reference_transpose.hpp"
|
||||
#include "ck_tile/utility/json_dump.hpp"
|
||||
#include "elementwise_common.hpp"
|
||||
|
||||
auto create_args(int argc, char* argv[])
|
||||
{
|
||||
ck_tile::ArgParser arg_parser;
|
||||
arg_parser.insert("m", "1024", "m dimension of input")
|
||||
.insert("n", "1024", "n dimension of input")
|
||||
.insert("stride_in", "-1", "stride for input M dim, if -1 then equal to n")
|
||||
.insert("v", "1", "cpu validation or not")
|
||||
.insert("prec", "fp16", "precision")
|
||||
.insert("warmup", "10", "cold iter")
|
||||
.insert("repeat", "50", "hot iter")
|
||||
.insert("json", "0", "0: No Json, 1: Dump Results in Json format")
|
||||
.insert("jsonfile", "elementwise_transpose.json", "json file name to dump results");
|
||||
|
||||
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)
|
||||
{
|
||||
ck_tile::index_t M = arg_parser.get_int("m");
|
||||
ck_tile::index_t N = arg_parser.get_int("n");
|
||||
ck_tile::index_t stride_in = arg_parser.get_int("stride_in");
|
||||
|
||||
if(stride_in < 0)
|
||||
stride_in = N; // Dense input: stride for M dim is N
|
||||
int do_validation = arg_parser.get_int("v");
|
||||
int warmup = arg_parser.get_int("warmup");
|
||||
int repeat = arg_parser.get_int("repeat");
|
||||
|
||||
if(stride_in < N)
|
||||
{
|
||||
throw std::runtime_error("stride_in must be >= N");
|
||||
}
|
||||
|
||||
using XDataType = DataType;
|
||||
using ComputeDataType = float;
|
||||
using YDataType = DataType;
|
||||
// Use PassThrough operation for transposition (data is moved, not changed)
|
||||
using XElementwiseOperation = ck_tile::element_wise::PassThrough;
|
||||
|
||||
// 1. Initialize the input data on the host (CPU).
|
||||
// Input x_host_a: M x N
|
||||
// Output y_host: N x M (transposed)
|
||||
ck_tile::HostTensor<XDataType> x_host_a({M, N}, {stride_in, 1});
|
||||
// Output tensor y_host will have dimensions N x M.
|
||||
// Assuming dense output, its stride for the N dimension will be M.
|
||||
ck_tile::index_t stride_out_dim0 = M;
|
||||
ck_tile::HostTensor<YDataType> y_host({N, M}, {stride_out_dim0, 1});
|
||||
ck_tile::HostTensor<YDataType> y_validation({N, M}, {stride_out_dim0, 1});
|
||||
|
||||
// The logical shape for the element-wise operation kernel is based on the input tensor's
|
||||
// elements.
|
||||
std::vector<ck_tile::index_t> op_shape_vec = {M, N};
|
||||
auto op_lengths = ck_tile::make_tuple(M, N); // Lens for the kernel
|
||||
|
||||
ck_tile::FillUniformDistribution<XDataType>{0.f, 5.f}(x_host_a);
|
||||
|
||||
// 2. Create device memory buffers
|
||||
ck_tile::DeviceMem x_buf_a(x_host_a.get_element_space_size_in_bytes());
|
||||
ck_tile::DeviceMem y_buf(y_host.get_element_space_size_in_bytes()); // y_host is N x M
|
||||
|
||||
x_buf_a.ToDevice(x_host_a.data());
|
||||
|
||||
// 3. Configure the kernel execution parameters.
|
||||
using BlockTile = ck_tile::sequence<1024>;
|
||||
using BlockWarps = ck_tile::sequence<8>;
|
||||
using WarpTile = ck_tile::sequence<64>;
|
||||
|
||||
using Shape = ck_tile::ElementWiseShape<BlockWarps, BlockTile, WarpTile, XDataType>;
|
||||
|
||||
// Problem definition for a single input tensor
|
||||
using Problem = ck_tile::ElementWisePipelineProblem<XDataType,
|
||||
ComputeDataType,
|
||||
YDataType,
|
||||
Shape,
|
||||
XElementwiseOperation>;
|
||||
|
||||
using Kernel = ck_tile::ElementWiseKernel<Problem, ck_tile::ElementWiseDefaultPolicy>;
|
||||
|
||||
ck_tile::index_t total_elements = M * N;
|
||||
|
||||
const ck_tile::index_t kBlockSize = Kernel::BlockSize();
|
||||
constexpr ck_tile::index_t kBlockPerCu = 1;
|
||||
constexpr ck_tile::index_t elements_per_block = BlockTile::at(ck_tile::number<0>{});
|
||||
ck_tile::index_t kGridSize = (total_elements + elements_per_block - 1) / elements_per_block;
|
||||
|
||||
std::cout << "Input M=" << M << ", N=" << N << ", StrideIn=" << stride_in << std::endl;
|
||||
std::cout << "Output N=" << N << ", M=" << M << ", StrideOut=" << stride_out_dim0 << std::endl;
|
||||
std::cout << "Grid size = " << kGridSize << ", BlockSize = " << kBlockSize << std::endl;
|
||||
std::cout << "Total elements = " << total_elements << std::endl;
|
||||
|
||||
// Input tensors tuple (single input)
|
||||
auto input_tensors = ck_tile::make_tuple(static_cast<XDataType*>(x_buf_a.GetDeviceBuffer()));
|
||||
// Input strides tuple (tuple of tuples, one for each input)
|
||||
auto input_strides = ck_tile::make_tuple(stride_in, 1);
|
||||
// Output strides (for N x M tensor, dense)
|
||||
auto output_strides = ck_tile::make_tuple(1, stride_out_dim0);
|
||||
|
||||
// Check if the kernel configuration is supported
|
||||
if(!Kernel::IsSupportedArgument(op_lengths))
|
||||
{
|
||||
throw std::runtime_error(
|
||||
"The kernel configuration is not supported for the given input size.");
|
||||
}
|
||||
|
||||
// 4. Run the kernel
|
||||
float ave_time = launch_kernel(
|
||||
ck_tile::stream_config{nullptr, true, 0, warmup, repeat},
|
||||
ck_tile::make_kernel<kBlockPerCu>(Kernel{},
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
0, // Shared memory
|
||||
op_lengths, // Logical dimensions for the operation (M, N)
|
||||
input_strides, // Strides for input tensor(s)
|
||||
output_strides, // Strides for output tensor (N, M)
|
||||
input_tensors,
|
||||
static_cast<YDataType*>(y_buf.GetDeviceBuffer())));
|
||||
|
||||
std::cout << "Average time: " << ave_time << " ms" << std::endl;
|
||||
|
||||
// 5. Verify the output
|
||||
bool pass = true;
|
||||
if(do_validation)
|
||||
{
|
||||
y_buf.FromDevice(y_validation.data()); // Copy result from device to y_validation
|
||||
ck_tile::reference_transpose_elementwise<XDataType, YDataType>(
|
||||
x_host_a, y_host); // Compute reference on host
|
||||
pass = ck_tile::check_err(
|
||||
y_validation, y_host, "Transpose Error: Incorrect results!", 0.01, 0.01);
|
||||
}
|
||||
|
||||
if(arg_parser.get_int("json") == 1)
|
||||
{
|
||||
dump_elementwise_json_results(arg_parser.get_str("jsonfile"),
|
||||
arg_parser.get_str("prec"),
|
||||
kGridSize,
|
||||
kBlockSize,
|
||||
ave_time,
|
||||
0,
|
||||
0,
|
||||
"elementwise_transpose");
|
||||
}
|
||||
|
||||
return pass;
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
bool result = true;
|
||||
ck_tile::ArgParser arg_parser;
|
||||
std::tie(result, arg_parser) = create_args(argc, argv);
|
||||
if(!result)
|
||||
return -1;
|
||||
|
||||
try
|
||||
{
|
||||
const auto prec_variant = string_to_datatype(arg_parser.get_str("prec"));
|
||||
return std::visit(
|
||||
[&](auto&& dt) -> int {
|
||||
using DataType = std::decay_t<decltype(dt)>;
|
||||
return run<DataType>(arg_parser);
|
||||
},
|
||||
prec_variant);
|
||||
}
|
||||
catch(const std::exception& e)
|
||||
{
|
||||
std::cerr << "Error: " << e.what() << std::endl;
|
||||
return -3;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user