mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-04-20 06:49:15 +00:00
Add transpose scale amax example (#1547)
* Add transpose scale amax example * fixes * Tune reduce instance
This commit is contained in:
Bartłomiej Kocot
@@ -5,3 +5,4 @@ add_example_executable(example_elementwise_permute_4D_fp32_col elementwise_permu
|
||||
add_example_executable(example_elementwise_permute_4D_fp16_col elementwise_permute_4D_fp16_col.cpp)
|
||||
add_example_executable(example_elementwise_binary_4D_fp16 elementwise_binary_4D_fp16.cpp)
|
||||
add_example_executable(example_elementwise_trinary_4D_fp16 elementwise_trinary_4D_fp16.cpp)
|
||||
add_example_executable(elementwise_scale_permute_amax_2D_fp16_fp8 elementwise_scale_permute_amax_2D_fp16_fp8.cpp)
|
||||
|
||||
@@ -0,0 +1,247 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#include <iostream>
|
||||
#include <cstdlib>
|
||||
|
||||
#include "ck/ck.hpp"
|
||||
#include "ck/tensor_operation/gpu/element/binary_element_wise_operation.hpp"
|
||||
#include "ck/tensor_operation/gpu/device/impl/device_elementwise_dynamic_vector_dims_impl.hpp"
|
||||
#include "ck/tensor_operation/gpu/device/impl/device_reduce_multiblock.hpp"
|
||||
#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp"
|
||||
|
||||
#include "ck/library/reference_tensor_operation/cpu/reference_elementwise.hpp"
|
||||
|
||||
#include "ck/library/utility/algorithm.hpp"
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
#include "ck/utility/reduction_enums.hpp"
|
||||
|
||||
using F16 = ck::half_t;
|
||||
using F32 = float;
|
||||
using F8 = ck::f8_t;
|
||||
|
||||
using InputDataType = F16;
|
||||
using ScaleDataType = F32;
|
||||
using OutputDataType = F8;
|
||||
|
||||
static constexpr ck::index_t NumDim = 2;
|
||||
|
||||
constexpr ck::ReduceTensorOp ReduceOpId = ck::ReduceTensorOp::MAX;
|
||||
constexpr bool PropagateNan = true;
|
||||
constexpr bool OutputIndex = false;
|
||||
|
||||
using ReduceOperation = typename ck::reduce_binary_operator<ReduceOpId>::opType;
|
||||
|
||||
struct ScalePassThrough
|
||||
{
|
||||
ScalePassThrough(const float alpha = 1.f) : alpha_(alpha) {}
|
||||
|
||||
__host__ __device__ constexpr void
|
||||
operator()(OutputDataType& y0, OutputDataType& y1, const InputDataType& x0) const
|
||||
{
|
||||
y0 = ck::type_convert<OutputDataType>(ck::type_convert<ScaleDataType>(x0) * alpha_);
|
||||
y1 = y0;
|
||||
}
|
||||
|
||||
const ScaleDataType alpha_;
|
||||
};
|
||||
|
||||
using PassThrough = ck::tensor_operation::element_wise::PassThrough;
|
||||
using UnaryAbs = ck::tensor_operation::element_wise::UnaryAbs;
|
||||
|
||||
using DeviceElementwisePermuteInstance = ck::tensor_operation::device::DeviceElementwiseImpl<
|
||||
ck::Tuple<InputDataType>, // InDataTypeTuple
|
||||
ck::Tuple<OutputDataType, OutputDataType>, // OutDataTypeTuple
|
||||
ScalePassThrough, // Elementwise
|
||||
NumDim, // NumDim
|
||||
256, // BlockSize
|
||||
128, // M0PerBlock
|
||||
128, // M1PerBlock
|
||||
8, // M0PerThread
|
||||
8, // M1PerThread
|
||||
ck::Sequence<1, 0>, // ThreadClusterArrangeOrder
|
||||
ck::Sequence<8>, // InScalarPerVectorSeq
|
||||
ck::Sequence<8, 1>>; // OutScalarPerVectorSeq
|
||||
|
||||
using DeviceReduceInstance =
|
||||
ck::tensor_operation::device::DeviceReduceMultiBlock<OutputDataType,
|
||||
OutputDataType,
|
||||
OutputDataType,
|
||||
NumDim,
|
||||
NumDim,
|
||||
ReduceOperation,
|
||||
UnaryAbs,
|
||||
PassThrough,
|
||||
ck::InMemoryDataOperationEnum::Set,
|
||||
PropagateNan,
|
||||
OutputIndex,
|
||||
false, // HaveIndexInputIfOutputIndex
|
||||
1024, // BlockSize
|
||||
1, // MThreadClusterSize
|
||||
1024, // KThreadClusterSize
|
||||
1, // MThreadSliceSize
|
||||
16, // KThreadSliceSize
|
||||
1, // InSrcVectorDim
|
||||
16, // InSrceVectorSize
|
||||
1>; // OutDstVectorSize
|
||||
|
||||
void reference_scale_permute_amax(Tensor<InputDataType>& input,
|
||||
Tensor<OutputDataType>& host_output_scaled_casted_transposed,
|
||||
Tensor<OutputDataType>& host_output_scaled_casted,
|
||||
Tensor<OutputDataType>& host_output_amax,
|
||||
const float scale)
|
||||
{
|
||||
ScalePassThrough out_element_op(scale);
|
||||
const ck::index_t M = input.GetLengths()[0];
|
||||
const ck::index_t K = input.GetLengths()[1];
|
||||
for(ck::index_t m = 0; m < M; m++)
|
||||
{
|
||||
for(ck::index_t k = 0; k < K; k++)
|
||||
{
|
||||
OutputDataType y0, y1;
|
||||
out_element_op(y0, y1, input(m, k));
|
||||
|
||||
host_output_scaled_casted(m, k) = y0;
|
||||
host_output_scaled_casted_transposed(m, k) = y1;
|
||||
const OutputDataType y_fabs =
|
||||
ck::type_convert<OutputDataType>(ck::math::abs(ck::type_convert<float>(y0)));
|
||||
host_output_amax(0) = ck::math::max(y_fabs, host_output_amax(0));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
bool do_verification = true;
|
||||
bool time_kernel = true;
|
||||
|
||||
const float scale = 2.f;
|
||||
|
||||
ck::index_t M = 1024;
|
||||
ck::index_t K = 1024;
|
||||
|
||||
if(argc == 3)
|
||||
{
|
||||
M = std::stoi(argv[1]);
|
||||
K = std::stoi(argv[2]);
|
||||
}
|
||||
|
||||
std::array<ck::index_t, 2> dims = {M, K};
|
||||
std::array<ck::index_t, 2> in_strides = {K, 1};
|
||||
std::array<ck::index_t, 2> out_strides = {1, M};
|
||||
|
||||
Tensor<InputDataType> input(dims, in_strides);
|
||||
Tensor<OutputDataType> output_scaled_casted_transposed(dims, out_strides);
|
||||
Tensor<OutputDataType> output_scaled_casted(dims, in_strides);
|
||||
Tensor<OutputDataType> output_amax({1});
|
||||
|
||||
input.GenerateTensorValue(GeneratorTensor_3<InputDataType>{0.0, 1.0});
|
||||
|
||||
DeviceMem input_dev_buf(sizeof(InputDataType) * input.mDesc.GetElementSpaceSize());
|
||||
DeviceMem output_scaled_casted_transposed_dev_buf(
|
||||
sizeof(OutputDataType) * output_scaled_casted_transposed.mDesc.GetElementSpaceSize());
|
||||
DeviceMem output_scaled_casted_dev_buf(sizeof(OutputDataType) *
|
||||
output_scaled_casted.mDesc.GetElementSpaceSize());
|
||||
DeviceMem output_amax_dev_buf(sizeof(OutputDataType) * output_amax.mDesc.GetElementSpaceSize());
|
||||
|
||||
input_dev_buf.ToDevice(input.mData.data());
|
||||
|
||||
std::array<const void*, 1> inputs = {input_dev_buf.GetDeviceBuffer()};
|
||||
std::array<void*, 2> outputs = {output_scaled_casted_transposed_dev_buf.GetDeviceBuffer(),
|
||||
output_scaled_casted_dev_buf.GetDeviceBuffer()};
|
||||
|
||||
std::cout << "Input: " << input.mDesc << std::endl;
|
||||
std::cout << "Scale: " << scale << std::endl;
|
||||
std::cout << "Output scaled casted transposed: " << output_scaled_casted_transposed.mDesc
|
||||
<< std::endl;
|
||||
std::cout << "Output scaled casted: " << output_scaled_casted.mDesc << std::endl;
|
||||
std::cout << "Output amax: " << output_amax.mDesc << std::endl;
|
||||
|
||||
auto launch_transpose_scale = [&]() {
|
||||
auto transposeScale = DeviceElementwisePermuteInstance{};
|
||||
auto argument = transposeScale.MakeArgumentPointer(dims,
|
||||
{in_strides},
|
||||
{out_strides, in_strides},
|
||||
inputs,
|
||||
outputs,
|
||||
ScalePassThrough{scale});
|
||||
|
||||
if(!transposeScale.IsSupportedArgument(argument.get()))
|
||||
{
|
||||
throw std::runtime_error(
|
||||
"The runtime parameters seems not supported by the device instance, exiting!");
|
||||
};
|
||||
|
||||
auto transposeScale_invoker_ptr = transposeScale.MakeInvokerPointer();
|
||||
return transposeScale_invoker_ptr->Run(argument.get(), StreamConfig{nullptr, time_kernel});
|
||||
};
|
||||
|
||||
auto launch_reduce = [&]() {
|
||||
auto reduce = DeviceReduceInstance{};
|
||||
auto reduce_argument_ptr =
|
||||
reduce.MakeArgumentPointer(dims,
|
||||
in_strides,
|
||||
{1}, // Output Lengths
|
||||
{1}, // Output Strides
|
||||
{0, 1}, // Reduce Dims
|
||||
static_cast<double>(1.f),
|
||||
static_cast<double>(0.f),
|
||||
output_scaled_casted_dev_buf.GetDeviceBuffer(),
|
||||
nullptr,
|
||||
output_amax_dev_buf.GetDeviceBuffer(),
|
||||
nullptr,
|
||||
UnaryAbs{},
|
||||
PassThrough{});
|
||||
|
||||
if(!reduce.IsSupportedArgument(reduce_argument_ptr.get()))
|
||||
{
|
||||
throw std::runtime_error(
|
||||
"The runtime parameters seems not supported by the device instance, exiting!");
|
||||
};
|
||||
|
||||
auto invoker_ptr = reduce.MakeInvokerPointer();
|
||||
|
||||
return invoker_ptr->Run(reduce_argument_ptr.get(), StreamConfig{nullptr, time_kernel});
|
||||
};
|
||||
|
||||
float ave_time = launch_transpose_scale();
|
||||
ave_time += launch_reduce();
|
||||
std::cout << "Perf: " << ave_time << " ms" << std::endl;
|
||||
bool pass = true;
|
||||
|
||||
if(do_verification)
|
||||
{
|
||||
Tensor<OutputDataType> host_output_scaled_casted_transposed(dims, out_strides);
|
||||
Tensor<OutputDataType> host_output_scaled_casted(dims, in_strides);
|
||||
Tensor<OutputDataType> host_output_amax({1});
|
||||
|
||||
reference_scale_permute_amax(input,
|
||||
host_output_scaled_casted_transposed,
|
||||
host_output_scaled_casted,
|
||||
host_output_amax,
|
||||
scale);
|
||||
|
||||
output_scaled_casted_transposed_dev_buf.FromDevice(
|
||||
output_scaled_casted_transposed.mData.data());
|
||||
output_scaled_casted_dev_buf.FromDevice(output_scaled_casted.mData.data());
|
||||
output_amax_dev_buf.FromDevice(output_amax.mData.data());
|
||||
|
||||
pass &= ck::utils::check_err(output_scaled_casted_transposed.mData,
|
||||
host_output_scaled_casted_transposed.mData,
|
||||
"Error: Incorrect results scaled transposed",
|
||||
1e-3,
|
||||
1e-3);
|
||||
pass &= ck::utils::check_err(output_scaled_casted.mData,
|
||||
host_output_scaled_casted.mData,
|
||||
"Error: Incorrect results scaled",
|
||||
1e-3,
|
||||
1e-3);
|
||||
pass &= ck::utils::check_err(
|
||||
output_amax.mData, host_output_amax.mData, "Error: Incorrect results amax", 1e-3, 1e-3);
|
||||
}
|
||||
|
||||
return pass ? 0 : 1;
|
||||
}
|
||||
Reference in New Issue
Block a user