mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-04-19 22:39:03 +00:00
Add optimized copy to ck wrapper (#1126)
* Add optimized copy to ck wrapper * Example optimizations * Fixes * Move img2col test to client example * Refactor example * Fix docs * Fixes * Fix * Fixes * Fixes * Fixes * Fixes * Fixes --------- Co-authored-by: zjing14 <zhangjing14@gmail.com>
This commit is contained in:
Bartłomiej Kocot
180
client_example/25_wrapper/wrapper_img2col.cpp
Normal file
180
client_example/25_wrapper/wrapper_img2col.cpp
Normal file
@@ -0,0 +1,180 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
#include <numeric>
|
||||
#include <cstdlib>
|
||||
#include <iomanip>
|
||||
#include <iostream>
|
||||
#include <initializer_list>
|
||||
#include <vector>
|
||||
|
||||
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
|
||||
|
||||
#include "ck/host_utility/kernel_launch.hpp"
|
||||
#include "ck/utility/common_header.hpp"
|
||||
#include "ck/wrapper/layout.hpp"
|
||||
#include "ck/wrapper/tensor.hpp"
|
||||
#include "ck/wrapper/operations/copy.hpp"
|
||||
|
||||
static constexpr ck::index_t NumDimSpatial = 3;
|
||||
using DataType = float;
|
||||
using InputLayout = ck::tensor_layout::convolution::NDHWGC;
|
||||
|
||||
struct SimpleDeviceMem
|
||||
{
|
||||
SimpleDeviceMem() = delete;
|
||||
|
||||
SimpleDeviceMem(std::size_t mem_size) : p_mem_{}
|
||||
{
|
||||
(void)hipMalloc(static_cast<void**>(&p_mem_), mem_size);
|
||||
}
|
||||
|
||||
void* GetDeviceBuffer() { return p_mem_; }
|
||||
|
||||
~SimpleDeviceMem() { (void)hipFree(p_mem_); }
|
||||
|
||||
void* p_mem_;
|
||||
};
|
||||
|
||||
// Test copy from Global to Global through LDS and VGPR
|
||||
template <typename InputTensor,
|
||||
typename OutputTensor,
|
||||
typename BlockShape,
|
||||
typename ThreadLayoutShape>
|
||||
__global__ void DeviceImageToColumnPad0(InputTensor input_tensor,
|
||||
OutputTensor output_tensor,
|
||||
const BlockShape tile_shape,
|
||||
const ThreadLayoutShape thread_layout)
|
||||
{
|
||||
const ck::index_t block_idx = static_cast<ck::index_t>(blockIdx.x);
|
||||
|
||||
// Get local tiles for global memory
|
||||
auto input_local_tile = ck::wrapper::make_local_tile(input_tensor, tile_shape, block_idx);
|
||||
auto output_local_tile = ck::wrapper::make_local_tile(output_tensor, tile_shape, block_idx);
|
||||
|
||||
// Get partition per thread
|
||||
const auto input_local_partition =
|
||||
ck::wrapper::make_local_partition(input_local_tile, thread_layout, threadIdx.x);
|
||||
auto output_local_partition =
|
||||
ck::wrapper::make_local_partition(output_local_tile, thread_layout, threadIdx.x);
|
||||
|
||||
// Perform copy
|
||||
using DimAccessOrder = ck::Tuple<ck::Number<0>, ck::Number<1>>;
|
||||
constexpr ck::index_t vector_dim = 1;
|
||||
constexpr ck::index_t scalar_per_vector = 4;
|
||||
ck::wrapper::copy<DimAccessOrder, vector_dim, scalar_per_vector>(input_local_partition,
|
||||
output_local_partition);
|
||||
}
|
||||
|
||||
void PerformImageToColumnPad0(const ck::index_t G,
|
||||
const ck::index_t N,
|
||||
const ck::index_t Di,
|
||||
const ck::index_t Hi,
|
||||
const ck::index_t Wi,
|
||||
const ck::index_t Do,
|
||||
const ck::index_t Ho,
|
||||
const ck::index_t Wo,
|
||||
const ck::index_t C,
|
||||
const ck::index_t Z,
|
||||
const ck::index_t Y,
|
||||
const ck::index_t X,
|
||||
std::array<ck::index_t, NumDimSpatial> filter_strides,
|
||||
std::array<ck::index_t, NumDimSpatial> filter_dilations)
|
||||
{
|
||||
const ck::index_t ZYXC = Z * Y * X * C;
|
||||
const ck::index_t GC = G * C;
|
||||
|
||||
// shape: (G, (Wo, Ho, Do, N)), (C, X, Y, Z))
|
||||
const auto shape = ck::make_tuple(ck::make_tuple(G, ck::make_tuple(Wo, Ho, Do, N)),
|
||||
ck::make_tuple(C, X, Y, Z));
|
||||
const auto in_strides =
|
||||
ck::make_tuple(ck::make_tuple(C,
|
||||
ck::make_tuple(filter_strides[2] * GC,
|
||||
filter_strides[1] * Wi * GC,
|
||||
filter_strides[0] * Hi * Wi * GC,
|
||||
Di * Hi * Wi * GC)),
|
||||
ck::make_tuple(1,
|
||||
filter_dilations[2] * GC,
|
||||
filter_dilations[1] * Wi * GC,
|
||||
filter_dilations[0] * Hi * Wi * GC));
|
||||
const auto in_layout = ck::wrapper::make_layout(shape, in_strides);
|
||||
|
||||
const auto out_strides = ck::make_tuple(
|
||||
ck::make_tuple(
|
||||
ZYXC,
|
||||
ck::make_tuple(ZYXC * G, Wo * ZYXC * G, Ho * Wo * ZYXC * G, Do * Ho * Wo * ZYXC * G)),
|
||||
ck::make_tuple(1, C, X * C, Y * X * C));
|
||||
const auto out_layout = ck::wrapper::make_layout(shape, out_strides);
|
||||
|
||||
const ck::index_t input_size = N * Di * Hi * Wi * GC;
|
||||
// Global memory buffers
|
||||
SimpleDeviceMem in_buf(input_size * sizeof(DataType));
|
||||
SimpleDeviceMem out_buf(ck::wrapper::size(out_layout) * sizeof(DataType));
|
||||
|
||||
// User can choose appropriate number of threads and sizes per block
|
||||
const auto thread_layout = ck::make_tuple(ck::Number<8>{}, ck::Number<16>{});
|
||||
// This example doesn't support padding, user should select tile sizes
|
||||
// which divides the shape completely
|
||||
const auto tile_shape = ck::make_tuple(ck::Number<32>{}, ck::Number<64>{});
|
||||
|
||||
// Create buffers for global memory
|
||||
auto input_tensor_global = ck::wrapper::make_tensor<ck::wrapper::MemoryTypeEnum::Global>(
|
||||
static_cast<const DataType*>(in_buf.GetDeviceBuffer()), in_layout);
|
||||
auto output_tensor_global = ck::wrapper::make_tensor<ck::wrapper::MemoryTypeEnum::Global>(
|
||||
static_cast<DataType*>(out_buf.GetDeviceBuffer()), out_layout);
|
||||
|
||||
const ck::index_t grid_size = ck::math::integer_divide_ceil(ck::wrapper::size<0>(in_layout),
|
||||
ck::wrapper::size<0>(tile_shape)) *
|
||||
ck::math::integer_divide_ceil(ck::wrapper::size<1>(in_layout),
|
||||
ck::wrapper::size<1>(tile_shape));
|
||||
|
||||
const auto kernel = DeviceImageToColumnPad0<decltype(input_tensor_global),
|
||||
decltype(output_tensor_global),
|
||||
decltype(tile_shape),
|
||||
decltype(thread_layout)>;
|
||||
const float avg_time = launch_and_time_kernel(StreamConfig{nullptr, true},
|
||||
kernel,
|
||||
dim3(grid_size),
|
||||
dim3(ck::wrapper::size(thread_layout)),
|
||||
0,
|
||||
input_tensor_global,
|
||||
output_tensor_global,
|
||||
tile_shape,
|
||||
thread_layout);
|
||||
|
||||
std::size_t num_btype = G * N * Do * Ho * Wo * ZYXC * 2 * sizeof(DataType);
|
||||
float gb_per_sec = num_btype / 1.E6 / avg_time;
|
||||
std::cout << "Perf: " << std::setw(10) << avg_time << " ms, " << gb_per_sec << " GB/s, "
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
constexpr ck::index_t G = 4; // number of groups
|
||||
constexpr ck::index_t N = 32; // batch
|
||||
constexpr ck::index_t C = 64; // input channel (per group)
|
||||
constexpr ck::index_t Z = 3; // filter D
|
||||
constexpr ck::index_t Y = 3; // filter H
|
||||
constexpr ck::index_t X = 3; // filter W
|
||||
constexpr ck::index_t Di = 9; // input D
|
||||
constexpr ck::index_t Hi = 9; // input H
|
||||
constexpr ck::index_t Wi = 7; // input W
|
||||
constexpr ck::index_t Do = 7; // output D
|
||||
constexpr ck::index_t Ho = 7; // output H
|
||||
constexpr ck::index_t Wo = 5; // output W
|
||||
PerformImageToColumnPad0(G,
|
||||
N,
|
||||
Di,
|
||||
Hi,
|
||||
Wi,
|
||||
Do,
|
||||
Ho,
|
||||
Wo,
|
||||
C,
|
||||
Z,
|
||||
Y,
|
||||
X,
|
||||
{1, 1, 1} /*filter_strides*/,
|
||||
{1, 1, 1} /*filter_dilations*/);
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user