38 #include "compute_kernel_writer/include/ckw/KernelWriter.h"
43 namespace experimental
45 namespace dynamic_fusion
76 const auto pool_sz_x =
static_cast<int32_t
>(_attributes.
pool_size().
x());
77 const auto pool_sz_y =
static_cast<int32_t
>(_attributes.
pool_size().
y());
78 const auto pad_x =
static_cast<int32_t
>(_attributes.
pad().
left);
79 const auto pad_y =
static_cast<int32_t
>(_attributes.
pad().
top);
80 const auto stride_x =
static_cast<int32_t
>(_attributes.
stride().
x());
81 const auto stride_y =
static_cast<int32_t
>(_attributes.
stride().
y());
87 auto const_pool_sz_x_i32 = writer->declare_constant_tile(ckw::ConstantData({{pool_sz_x}}, ckw::DataType::Int32));
88 auto const_pool_sz_y_i32 = writer->declare_constant_tile(ckw::ConstantData({{pool_sz_y}}, ckw::DataType::Int32));
89 auto const_pad_x_i32 = writer->declare_constant_tile(ckw::ConstantData({{pad_x}}, ckw::DataType::Int32));
90 auto const_pad_y_i32 = writer->declare_constant_tile(ckw::ConstantData({{pad_y}}, ckw::DataType::Int32));
91 auto const_stride_x_i32 = writer->declare_constant_tile(ckw::ConstantData({{stride_x}}, ckw::DataType::Int32));
92 auto const_stride_y_i32 = writer->declare_constant_tile(ckw::ConstantData({{stride_y}}, ckw::DataType::Int32));
93 auto const_src_w_i32 = writer->declare_constant_tile(ckw::ConstantData({{src_w}}, ckw::DataType::Int32));
94 auto const_src_h_i32 = writer->declare_constant_tile(ckw::ConstantData({{src_h}}, ckw::DataType::Int32));
95 auto const_dst_h_i32 = writer->declare_constant_tile(ckw::ConstantData({{dst_h}}, ckw::DataType::Int32));
96 auto const_0_i32 = writer->declare_constant_tile(ckw::ConstantData({{0}}, ckw::DataType::Int32));
97 auto const_pos_1_i32 = writer->declare_constant_tile(ckw::ConstantData({{1}}, ckw::DataType::Int32));
98 auto const_0_fp = writer->declare_constant_tile(ckw::ConstantData({{0.0f}}, dst_dt));
99 auto const_lowest_val_fp =
101 auto const_neg_inf_val_fp = writer->declare_constant_tile(ckw::ConstantData({{-1.0f / 0.0f}}, ckw::DataType::Fp32));
112 const int32_t dst_n0 = root_window.
x().
step();
113 const int32_t dst_m0 = root_window.y().step();
116 const int32_t dst_n0_partial = _dst->
dimension(0) % dst_n0;
117 const int32_t dst_m0_partial = _dst->
dimension(1) % dst_m0;
120 const int32_t dst_shift_back = (dst_n0 - dst_n0_partial) % dst_n0;
122 ckw::TensorSampler sampler_dst;
123 sampler_dst.format(ckw::TensorSamplerFormat::Dim0_Dim1_Dim2);
124 if (dst_n0_partial == 0)
130 sampler_dst.address_mode_x(ckw::TensorSamplerAddressModeX::OverlappingMin);
133 if (dst_m0_partial == 0)
139 sampler_dst.address_mode_y(ckw::TensorSamplerAddressModeY::ClampToBorderMaxOnly);
143 sampler_dst.storage(ckw::TensorStorageType::BufferUint8Ptr);
146 auto tile_dst = writer->declare_tile(
"dst", ckw::TileInfo(dst_dt, dst_m0, dst_n0));
149 writer->op_assign(tile_dst, const_0_fp);
152 dst->init_virtual_tensor(tile_dst, sampler_dst);
158 auto const_dst_n0_i32 = writer->declare_constant_tile(ckw::ConstantData({{dst_n0}}, ckw::DataType::Int32));
159 auto const_dst_m0_i32 = writer->declare_constant_tile(ckw::ConstantData({{dst_m0}}, ckw::DataType::Int32));
160 auto const_shift_back_dst_n0_i32 =
161 writer->declare_constant_tile(ckw::ConstantData({{dst_shift_back}}, ckw::DataType::Int32));
166 ckw::TensorSampler sampler_src;
167 sampler_src.format(ckw::TensorSamplerFormat::Dim0_Dim1_Dim2);
176 const bool is_global_pooling = (pool_sz_x == src_w) && (pool_sz_y == src_h) && (pad_x == 0) && (pad_y == 0);
179 const bool acc_f32 = (dst_dt == ckw::DataType::Fp32) ||
182 const auto acc_dt = acc_f32 ? ckw::DataType::Fp32 : ckw::DataType::Fp16;
184 const bool is_wider_acc = dst_dt != acc_dt;
189 auto tile_gid_0 = writer->declare_tile(
"gid_0", ckw::TileInfo(ckw::DataType::Int32));
190 auto tile_gid_1 = writer->declare_tile(
"gid_1", ckw::TileInfo(ckw::DataType::Int32));
191 auto tile_gid_2 = writer->declare_tile(
"gid_2", ckw::TileInfo(ckw::DataType::Int32));
193 writer->op_get_global_id(tile_gid_0, 0);
194 writer->op_get_global_id(tile_gid_1, 1);
195 writer->op_get_global_id(tile_gid_2, 2);
197 auto tile_cout0 = writer->declare_tile(
"cout0", ckw::TileInfo(ckw::DataType::Int32));
198 auto tile_mout0 = writer->declare_tile(
"mout0", ckw::TileInfo(ckw::DataType::Int32));
199 auto tile_mout1 = writer->declare_tile(
"mout1", ckw::TileInfo(ckw::DataType::Int32));
200 auto tile_bout0 = writer->declare_tile(
"bout0", ckw::TileInfo(ckw::DataType::Int32));
204 const_shift_back_dst_n0_i32, const_0_i32);
206 writer->op_binary(tile_mout1, ckw::BinaryOp::Mod, tile_gid_2, const_dst_h_i32);
207 writer->op_binary(tile_bout0, ckw::BinaryOp::Div, tile_gid_2, const_dst_h_i32);
213 auto tile_res = writer->declare_tile(
"tile_res", ckw::TileInfo(acc_dt, dst_m0, dst_n0));
233 auto tile_src_coord_x_start = writer->declare_tile(
"idx_in_w", ckw::DataType::Int32);
234 writer->op_binary(tile_src_coord_x_start, ckw::BinaryOp::Mul, tile_mout0, const_stride_x_i32);
235 writer->op_binary(tile_src_coord_x_start, ckw::BinaryOp::Sub, tile_src_coord_x_start, const_pad_x_i32);
238 auto tile_src_coord_y_start = writer->declare_tile(
"idx_in_h", ckw::DataType::Int32);
239 writer->op_binary(tile_src_coord_y_start, ckw::BinaryOp::Mul, tile_mout1, const_stride_y_i32);
240 writer->op_binary(tile_src_coord_y_start, ckw::BinaryOp::Sub, tile_src_coord_y_start, const_pad_y_i32);
242 auto tile_neg_src_coord_x_start = writer->declare_tile(
"neg_src_coord_x_start", ckw::DataType::Int32);
243 auto tile_neg_src_coord_y_start = writer->declare_tile(
"neg_src_coord_y_start", ckw::DataType::Int32);
245 writer->op_binary(tile_neg_src_coord_x_start, ckw::BinaryOp::Sub, const_0_i32, tile_src_coord_x_start);
246 writer->op_binary(tile_neg_src_coord_y_start, ckw::BinaryOp::Sub, const_0_i32, tile_src_coord_y_start);
252 auto tile_pool_x_s = writer->declare_tile(
"pool_x_s", ckw::DataType::Int32);
253 auto tile_pool_y_s = writer->declare_tile(
"pool_y_s", ckw::DataType::Int32);
254 auto tile_pool_x_e = writer->declare_tile(
"pool_x_e", ckw::DataType::Int32);
255 auto tile_pool_y_e = writer->declare_tile(
"pool_y_e", ckw::DataType::Int32);
257 writer->op_binary(tile_pool_x_s, ckw::BinaryOp::Max, const_0_i32, tile_neg_src_coord_x_start);
258 writer->op_binary(tile_pool_x_e, ckw::BinaryOp::Add, const_src_w_i32, tile_neg_src_coord_x_start);
259 writer->op_binary(tile_pool_x_e, ckw::BinaryOp::Min, const_pool_sz_x_i32, tile_pool_x_e);
260 writer->op_binary(tile_pool_y_s, ckw::BinaryOp::Max, const_0_i32, tile_neg_src_coord_y_start);
261 writer->op_binary(tile_pool_y_e, ckw::BinaryOp::Add, const_src_h_i32, tile_neg_src_coord_y_start);
262 writer->op_binary(tile_pool_y_e, ckw::BinaryOp::Min, const_pool_sz_y_i32, tile_pool_y_e);
269 auto tile_filter_size = writer->declare_tile(
"filter_size", ckw::DataType::Int32);
272 auto tile_x_diff = writer->declare_tile(
"x_diff", ckw::DataType::Int32);
273 auto tile_y_diff = writer->declare_tile(
"y_diff", ckw::DataType::Int32);
275 writer->op_binary(tile_x_diff, ckw::BinaryOp::Sub, tile_pool_x_e, tile_pool_x_s);
276 writer->op_binary(tile_y_diff, ckw::BinaryOp::Sub, tile_pool_y_e, tile_pool_y_s);
277 writer->op_binary(tile_filter_size, ckw::BinaryOp::Mul, tile_x_diff, tile_y_diff);
281 writer->op_binary(tile_filter_size, ckw::BinaryOp::Mul, const_pool_sz_x_i32, const_pool_sz_y_i32);
284 auto tile_x = writer->declare_tile(
"x", ckw::DataType::Int32);
285 auto tile_y = writer->declare_tile(
"y", ckw::DataType::Int32);
287 if (is_global_pooling)
289 writer->op_assign(tile_y, const_0_i32);
290 writer->op_assign(tile_pool_y_e, const_pool_sz_y_i32);
294 writer->op_assign(tile_y, tile_pool_y_s);
299 tile_y, ckw::BinaryOp::Less, tile_pool_y_e, tile_y, ckw::AssignmentOp::Increment, const_pos_1_i32,
303 if (is_global_pooling)
305 writer->op_assign(tile_x, const_0_i32);
306 writer->op_assign(tile_pool_x_e, const_pool_sz_x_i32);
310 writer->op_assign(tile_x, tile_pool_x_s);
313 auto tile_src_coord_y = writer->declare_tile(
"src_coord_y", ckw::DataType::Int32);
314 writer->op_binary(tile_src_coord_y, ckw::BinaryOp::Add, tile_src_coord_y_start, tile_y);
318 tile_x, ckw::BinaryOp::Less, tile_pool_x_e, tile_x, ckw::AssignmentOp::Increment, const_pos_1_i32,
321 auto tile_src_coord_x = writer->declare_tile(
"src_coord_x", ckw::DataType::Int32);
322 writer->op_binary(tile_src_coord_x, ckw::BinaryOp::Add, tile_src_coord_x_start, tile_x);
325 auto tile_src = writer->declare_tile(
"tile_src", ckw::TileInfo(acc_dt, dst_m0, dst_n0));
330 auto tile_src0 = writer->declare_tile(
"src_tile0", ckw::TileInfo(src_dt, dst_m0, dst_n0));
331 writer->op_load(tile_src0,
src->tensor(), sampler_src, tile_cout0, tile_src_coord_x,
332 tile_src_coord_y, tile_bout0);
337 writer->op_load(tile_src,
src->tensor(), sampler_src, tile_cout0, tile_src_coord_x,
338 tile_src_coord_y, tile_bout0);
344 writer->op_binary(tile_src, ckw::BinaryOp::Mul, tile_src, tile_src);
350 writer->op_binary(tile_res, ckw::BinaryOp::Max, tile_res, tile_src);
354 writer->op_binary(tile_res, ckw::BinaryOp::Add, tile_res, tile_src);
362 auto tile_filter_size_fp = writer->declare_tile(
"filter_size_fp", ckw::TileInfo(acc_dt));
364 writer->op_binary(tile_res, ckw::BinaryOp::Div, tile_res, tile_filter_size_fp);
370 writer->op_unary(tile_res, ckw::UnaryOp::Sqrt, tile_res);
380 writer->op_assign(tile_dst, tile_res);