|
| 1 | +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. |
| 2 | +// SPDX-License-Identifier: MIT |
| 3 | + |
| 4 | +#include "ck_tile/host.hpp" |
| 5 | +#include "ck_tile/ops/reduce.hpp" |
| 6 | +#include "ck_tile/utility/json_dump.hpp" |
| 7 | +#include <cstring> |
| 8 | + |
| 9 | +template <typename T> |
| 10 | +struct DataTypeTraits; |
| 11 | + |
| 12 | +template <> |
| 13 | +struct DataTypeTraits<ck_tile::half_t> |
| 14 | +{ |
| 15 | + static constexpr const char* name = "fp16"; |
| 16 | +}; |
| 17 | + |
| 18 | +template <> |
| 19 | +struct DataTypeTraits<ck_tile::bf16_t> |
| 20 | +{ |
| 21 | + static constexpr const char* name = "bf16"; |
| 22 | +}; |
| 23 | + |
| 24 | +auto create_args(int argc, char* argv[]) |
| 25 | +{ |
| 26 | + ck_tile::ArgParser arg_parser; |
| 27 | + arg_parser.insert("n", "32", "n dimension") |
| 28 | + .insert("h", "19", "h dimension") |
| 29 | + .insert("w", "7", "w dimension") |
| 30 | + .insert("c", "512", "c dimension") |
| 31 | + .insert("v", "1", "cpu validation or not") |
| 32 | + .insert("prec", "fp16", "precision") |
| 33 | + .insert("warmup", "5", "cold iter") |
| 34 | + .insert("repeat", "20", "hot iter") |
| 35 | + .insert("json", "0", "0: No Json, 1: Dump Results in Json format") |
| 36 | + .insert("jsonfile", "multi_reduce_multiblock.json", "json file name to dump results"); |
| 37 | + |
| 38 | + bool result = arg_parser.parse(argc, argv); |
| 39 | + return std::make_tuple(result, arg_parser); |
| 40 | +} |
| 41 | + |
| 42 | +template <typename DataType> |
| 43 | +bool run(const ck_tile::ArgParser& arg_parser) |
| 44 | +{ |
| 45 | + using XDataType = DataType; |
| 46 | + using ComputeDataType = float; |
| 47 | + using YDataType = float; |
| 48 | + |
| 49 | + ck_tile::index_t N = arg_parser.get_int("n"); |
| 50 | + ck_tile::index_t H = arg_parser.get_int("h"); |
| 51 | + ck_tile::index_t W = arg_parser.get_int("w"); |
| 52 | + ck_tile::index_t C = arg_parser.get_int("c"); |
| 53 | + int do_validation = arg_parser.get_int("v"); |
| 54 | + int warmup = arg_parser.get_int("warmup"); |
| 55 | + int repeat = arg_parser.get_int("repeat"); |
| 56 | + |
| 57 | + // Validate input dimensions |
| 58 | + const ck_tile::index_t kept_dim_len_prod = N * C; |
| 59 | + const ck_tile::index_t reduce_total_length = H * W; |
| 60 | + |
| 61 | + if(kept_dim_len_prod == 0) |
| 62 | + { |
| 63 | + std::cerr << "Warning: Product of kept dimensions is zero (N=" << N << ", C=" << C |
| 64 | + << ", product=" << kept_dim_len_prod << ")." << std::endl; |
| 65 | + std::cerr << "This will result in an empty output tensor." << std::endl; |
| 66 | + return false; |
| 67 | + } |
| 68 | + |
| 69 | + if(reduce_total_length == 0) |
| 70 | + { |
| 71 | + std::cerr << "Warning: Product of reduce dimensions is zero (H=" << H << ", W=" << W |
| 72 | + << ", product=" << reduce_total_length << ")." << std::endl; |
| 73 | + std::cerr << "This will result in an empty reduction with no data to process." << std::endl; |
| 74 | + std::cerr << "The kernel will exit early without performing any computation." << std::endl; |
| 75 | + return false; |
| 76 | + } |
| 77 | + |
| 78 | + std::vector<ck_tile::index_t> problem_shape = {N, H, W, C}; |
| 79 | + std::vector<ck_tile::index_t> strides(4); |
| 80 | + strides[0] = H * W * C; |
| 81 | + strides[1] = W * C; |
| 82 | + strides[2] = C; |
| 83 | + strides[3] = 1; |
| 84 | + |
| 85 | + // Define reduction specification: |
| 86 | + constexpr auto kept_dim = ck_tile::sequence<0, 3>{}; // Which dimension to keep |
| 87 | + constexpr auto reduce_dims = ck_tile::sequence<1, 2>{}; // Which dimensions to reduce |
| 88 | + |
| 89 | + ck_tile::HostTensor<XDataType> x_host(problem_shape, strides); |
| 90 | + ck_tile::HostTensor<YDataType> y_host_add_ref({N, C}, {C, 1}); |
| 91 | + ck_tile::HostTensor<YDataType> y_host_max_ref({N, C}, {C, 1}); |
| 92 | + auto y_host_ref_tuple = ck_tile::make_tuple(y_host_add_ref, y_host_max_ref); |
| 93 | + |
| 94 | + ck_tile::HostTensor<YDataType> y_host_add_dev({N, C}, {C, 1}); |
| 95 | + ck_tile::HostTensor<YDataType> y_host_max_dev({N, C}, {C, 1}); |
| 96 | + auto y_host_dev_tuple = ck_tile::make_tuple(y_host_add_dev, y_host_max_dev); |
| 97 | + |
| 98 | + const auto number_operations = y_host_dev_tuple.size(); |
| 99 | + |
| 100 | + std::vector<YDataType> h(number_operations * N * C); |
| 101 | + |
| 102 | + auto y_buf_size = number_operations * |
| 103 | + y_host_dev_tuple.at(ck_tile::number<0>{}).get_element_space_size_in_bytes(); |
| 104 | + ck_tile::DeviceMem y_buf(y_buf_size); |
| 105 | + |
| 106 | + const auto output_tensor_offset = N * C; |
| 107 | + |
| 108 | + // Operations: one doing a sum reduction, the other computing the mean square |
| 109 | + // In the case of mean square: |
| 110 | + // 1. The element wise operation squares each element before reduction |
| 111 | + // 2. The reduction operation sum the squared element |
| 112 | + // 3. The accumulator element wise operation divides the result by the total number of reduced |
| 113 | + // elements (intra block operation) |
| 114 | + // 4. The partial result is updated across blocks using inter block reduction, a sum. |
| 115 | + auto reduce_ops = |
| 116 | + ck_tile::make_tuple(ck_tile::ReduceOp::Add{}, ck_tile::ReduceOp::Add{}); // reductions |
| 117 | + auto elementwise_ops = ck_tile::make_tuple(ck_tile::element_wise::PassThrough{}, |
| 118 | + ck_tile::element_wise::UnarySquare{}); // Elementwise |
| 119 | + // ops |
| 120 | + auto accumulator_elementwise_ops = ck_tile::make_tuple( |
| 121 | + ck_tile::element_wise::PassThrough{}, |
| 122 | + ck_tile::element_wise::UnaryDivide{ |
| 123 | + reduce_total_length}); // Accumulator Elementwise ops on reduction, intra block |
| 124 | + auto inter_block_reduce_ops = ck_tile::make_tuple( |
| 125 | + ck_tile::ReduceOp::Add{}, ck_tile::ReduceOp::Add{}); // Inter block reduction |
| 126 | + |
| 127 | + ck_tile::FillUniformDistribution<XDataType>{-5.f, 5.f}(x_host); |
| 128 | + |
| 129 | + ck_tile::DeviceMem x_buf(x_host.get_element_space_size_in_bytes()); |
| 130 | + |
| 131 | + x_buf.ToDevice(x_host.data()); |
| 132 | + |
| 133 | + using BlockWarps = ck_tile::sequence<4, 1>; |
| 134 | + using BlockTile = ck_tile::sequence<128, 128>; |
| 135 | + using WarpTile = ck_tile::sequence<32, 128>; |
| 136 | + using ThreadTile = ck_tile::sequence<8, 8>; |
| 137 | + |
| 138 | + constexpr ck_tile::index_t kBlockPerCu = 1; |
| 139 | + |
| 140 | + using Shape = ck_tile::Reduce2dShape<BlockWarps, BlockTile, WarpTile, ThreadTile>; |
| 141 | + using Problem = ck_tile::Reduce2dProblem<XDataType, |
| 142 | + ComputeDataType, |
| 143 | + YDataType, |
| 144 | + Shape, |
| 145 | + decltype(reduce_ops), |
| 146 | + decltype(kept_dim), |
| 147 | + decltype(reduce_dims), |
| 148 | + 4>; |
| 149 | + |
| 150 | + using Kernel = ck_tile::MultiReduceMultiblock<Problem>; |
| 151 | + |
| 152 | + // Determine block group size for multi-block reduction |
| 153 | + // block_group_size records how many blocks participate to a reduction (input data dependent) |
| 154 | + // , for efficiency reasons this size if limited to a maximum of 128. If this is not sufficient |
| 155 | + // to process the whole reduction, each thread will to process multiple thread tile |
| 156 | + // a num_block_tile_iterations times |
| 157 | + auto [num_block_tile_iterations, block_group_size] = |
| 158 | + typename Kernel::TilePartitioner{reduce_total_length}.GetBlockGroupParams(); |
| 159 | + |
| 160 | + const ck_tile::index_t kBlockSize = Kernel::BlockSize(); |
| 161 | + ck_tile::index_t kGridSize = |
| 162 | + ((kept_dim_len_prod + Shape::Block_M - 1) / Shape::Block_M) * block_group_size; |
| 163 | + |
| 164 | + std::cout << "Block group size: " << block_group_size |
| 165 | + << ", Num block tile iterations: " << num_block_tile_iterations |
| 166 | + << ", Reduce total length: " << reduce_total_length << std::endl; |
| 167 | + std::cout << "grid size " << kGridSize << ", block size " << kBlockSize << std::endl; |
| 168 | + |
| 169 | + // Create input tensor shape and strides |
| 170 | + auto input_shape = |
| 171 | + ck_tile::make_tuple(problem_shape[0], problem_shape[1], problem_shape[2], problem_shape[3]); |
| 172 | + auto input_strides = ck_tile::make_tuple(strides[0], strides[1], strides[2], strides[3]); |
| 173 | + |
| 174 | + if(!Kernel::IsSupportedArgument( |
| 175 | + C, input_strides)) // output tensor's continuous dimension and input strides |
| 176 | + { |
| 177 | + throw std::runtime_error("Wrong! Arguments not supported!\n"); |
| 178 | + } |
| 179 | + |
| 180 | + // Init the output data with identity values respective to each reduce op |
| 181 | + ck_tile::static_for<0, number_operations, 1>{}([&](auto i) { |
| 182 | + constexpr auto op = reduce_ops.at(i); |
| 183 | + const auto identity_val = op.template GetIdentityValue<YDataType>(); |
| 184 | + const auto output_number_elements = N * C; |
| 185 | + std::fill(h.begin() + i * output_number_elements, |
| 186 | + h.begin() + (i + 1) * output_number_elements, |
| 187 | + identity_val); |
| 188 | + }); |
| 189 | + |
| 190 | + auto clear_output_buffer = [&]() { y_buf.ToDevice(h.data()); }; |
| 191 | + |
| 192 | + float ave_time = launch_kernel_time_mask( |
| 193 | + ck_tile::stream_config{nullptr, true, 0, warmup, repeat}, |
| 194 | + clear_output_buffer, |
| 195 | + ck_tile::make_kernel<kBlockPerCu>(Kernel{}, |
| 196 | + kGridSize, |
| 197 | + kBlockSize, |
| 198 | + 0, |
| 199 | + static_cast<XDataType*>(x_buf.GetDeviceBuffer()), |
| 200 | + static_cast<YDataType*>(y_buf.GetDeviceBuffer()), |
| 201 | + input_shape, |
| 202 | + input_strides, |
| 203 | + kept_dim, |
| 204 | + reduce_dims, |
| 205 | + output_tensor_offset, |
| 206 | + elementwise_ops, |
| 207 | + accumulator_elementwise_ops, |
| 208 | + inter_block_reduce_ops) |
| 209 | + |
| 210 | + ); |
| 211 | + |
| 212 | + std::size_t num_btype = sizeof(XDataType) * N * C * H * W + sizeof(YDataType) * N * C; |
| 213 | + |
| 214 | + float gb_per_sec = num_btype / 1.E6 / ave_time; |
| 215 | + |
| 216 | + std::cout << "Perf: " << ave_time << " ms, " << gb_per_sec << " GB/s" << std::endl; |
| 217 | + |
| 218 | + bool pass = true; |
| 219 | + |
| 220 | + if(do_validation) |
| 221 | + { |
| 222 | + // reference |
| 223 | + ck_tile::reference_multiple_reduce_multiblock<XDataType, ComputeDataType, YDataType>( |
| 224 | + x_host, |
| 225 | + y_host_ref_tuple, |
| 226 | + reduce_ops, |
| 227 | + kept_dim, |
| 228 | + reduce_dims, |
| 229 | + elementwise_ops, |
| 230 | + accumulator_elementwise_ops, |
| 231 | + inter_block_reduce_ops, |
| 232 | + block_group_size); |
| 233 | + std::cout << "Read " << y_buf_size / 10 << " Bytes from the device" << std::endl; |
| 234 | + |
| 235 | + // Transfer data from device and check error for each operation |
| 236 | + y_buf.FromDevice(h.data()); |
| 237 | + ck_tile::static_for<0, number_operations, 1>{}([&](auto i) { |
| 238 | + std::memcpy(y_host_dev_tuple.get(ck_tile::number<i>{}).data(), |
| 239 | + h.data() + i * output_tensor_offset, |
| 240 | + output_tensor_offset * sizeof(YDataType)); |
| 241 | + std::cout << "Checking operation " << i << ": " << std::endl; |
| 242 | + |
| 243 | + bool pass_op = ck_tile::check_err(y_host_dev_tuple.get(ck_tile::number<i>{}), |
| 244 | + y_host_ref_tuple.get(ck_tile::number<i>{})); |
| 245 | + |
| 246 | + if(pass_op) |
| 247 | + { |
| 248 | + std::cout << "✅ valid results for this operation" << std::endl; |
| 249 | + } |
| 250 | + pass &= pass_op; |
| 251 | + }); |
| 252 | + |
| 253 | + std::cout << "valid:" << (pass ? "y" : "n") << std::flush << std::endl; |
| 254 | + } |
| 255 | + |
| 256 | + return pass; |
| 257 | +} |
| 258 | + |
| 259 | +int main(int argc, char* argv[]) |
| 260 | +{ |
| 261 | + auto [result, arg_parser] = create_args(argc, argv); |
| 262 | + if(!result) |
| 263 | + return -1; |
| 264 | + |
| 265 | + const std::string data_type = arg_parser.get_str("prec"); |
| 266 | + |
| 267 | + if(data_type == "fp16") |
| 268 | + { |
| 269 | + return run<ck_tile::half_t>(arg_parser) ? 0 : -2; |
| 270 | + } |
| 271 | +} |
0 commit comments