38 #include "compute_kernel_writer/include/ckw/KernelWriter.h"
44 namespace experimental
46 namespace dynamic_fusion
80 const bool using_bias = _bia !=
nullptr;
96 const auto stride_x =
static_cast<int32_t
>(_attributes.
stride().x());
97 const auto stride_y =
static_cast<int32_t
>(_attributes.
stride().y());
98 const auto pad_x =
static_cast<int32_t
>(_attributes.
pad().left);
99 const auto pad_y =
static_cast<int32_t
>(_attributes.
pad().top);
100 const auto depth_multiplier =
static_cast<int32_t
>(_attributes.
depth_multiplier());
101 const auto dilation_x =
static_cast<int32_t
>(_attributes.
dilation().x());
102 const auto dilation_y =
static_cast<int32_t
>(_attributes.
dilation().y());
103 const auto kernel_size = kernel_width * kernel_height;
106 auto const_kernel_w_i32 = writer->declare_constant_tile(ckw::ConstantData({{kernel_width}}, ckw::DataType::Int32));
107 auto const_kernel_size_i32 =
108 writer->declare_constant_tile(ckw::ConstantData({{kernel_size}}, ckw::DataType::Int32));
109 auto const_dst_h_i32 = writer->declare_constant_tile(ckw::ConstantData({{dst_h}}, ckw::DataType::Int32));
110 auto const_src_w_i32 = writer->declare_constant_tile(ckw::ConstantData({{src_w}}, ckw::DataType::Int32));
111 auto const_src_h_i32 = writer->declare_constant_tile(ckw::ConstantData({{src_h}}, ckw::DataType::Int32));
112 auto const_stride_x_i32 = writer->declare_constant_tile(ckw::ConstantData({{stride_x}}, ckw::DataType::Int32));
113 auto const_stride_y_i32 = writer->declare_constant_tile(ckw::ConstantData({{stride_y}}, ckw::DataType::Int32));
114 auto const_pad_x_i32 = writer->declare_constant_tile(ckw::ConstantData({{pad_x}}, ckw::DataType::Int32));
115 auto const_pad_y_i32 = writer->declare_constant_tile(ckw::ConstantData({{pad_y}}, ckw::DataType::Int32));
116 auto const_0_i32 = writer->declare_constant_tile(ckw::ConstantData({{0}}, ckw::DataType::Int32));
117 auto const_neg_1_i32 = writer->declare_constant_tile(ckw::ConstantData({{-1}}, ckw::DataType::Int32));
118 auto const_depth_multiplier_i32 =
119 writer->declare_constant_tile(ckw::ConstantData({{depth_multiplier}}, ckw::DataType::Int32));
120 auto const_dilation_x_i32 = writer->declare_constant_tile(ckw::ConstantData({{dilation_x}}, ckw::DataType::Int32));
121 auto const_dilation_y_i32 = writer->declare_constant_tile(ckw::ConstantData({{dilation_y}}, ckw::DataType::Int32));
122 auto const_0_fp = writer->declare_constant_tile(ckw::ConstantData({{0.0f}}, dst_dt));
133 const int32_t dst_n0 = root_window.
x().
step();
134 const int32_t dst_m0 = root_window.y().step();
137 const int32_t dst_n0_partial = _dst->
dimension(0) % dst_n0;
138 const int32_t dst_m0_partial = _dst->
dimension(1) % dst_m0;
141 const int32_t dst_shift_back = (dst_n0 - dst_n0_partial) % dst_n0;
143 const int32_t src_m0 = kernel_width + (dst_m0 - 1);
144 const int32_t src_n0 = depth_multiplier > 1 ? 1 : dst_n0;
145 const int32_t wei_m0 = kernel_width;
146 const int32_t wei_n0 = dst_n0;
148 ckw::TensorSampler sampler_dst;
149 sampler_dst.format(ckw::TensorSamplerFormat::Dim0_Dim1_Dim2);
150 if (dst_n0_partial == 0)
156 sampler_dst.address_mode_x(ckw::TensorSamplerAddressModeX::OverlappingMin);
159 if (dst_m0_partial == 0)
165 sampler_dst.address_mode_y(ckw::TensorSamplerAddressModeY::ClampToBorderMaxOnly);
169 sampler_dst.storage(ckw::TensorStorageType::BufferUint8Ptr);
172 auto tile_dst = writer->declare_tile(
"dst", ckw::TileInfo(dst_dt, dst_m0, dst_n0));
175 writer->op_assign(tile_dst, const_0_fp);
178 dst->init_virtual_tensor(tile_dst, sampler_dst);
184 auto const_dst_n0_i32 = writer->declare_constant_tile(ckw::ConstantData({{dst_n0}}, ckw::DataType::Int32));
185 auto const_dst_m0_i32 = writer->declare_constant_tile(ckw::ConstantData({{dst_m0}}, ckw::DataType::Int32));
186 auto const_shift_back_dst_n0_i32 =
187 writer->declare_constant_tile(ckw::ConstantData({{dst_shift_back}}, ckw::DataType::Int32));
193 ckw::TensorSampler sampler_src;
194 sampler_src.format(ckw::TensorSamplerFormat::Dim0_Dim1_Dim2);
196 sampler_src.address_mode_y(ckw::TensorSamplerAddressModeY::SkipLessThanZero);
198 sampler_src.storage(ckw::TensorStorageType::BufferUint8Ptr);
202 ckw::TensorSampler sampler_wei;
203 sampler_wei.format(ckw::TensorSamplerFormat::Dim0_Dim1_Dim2);
209 sampler_wei.storage(ckw::TensorStorageType::Texture2dReadOnly);
213 sampler_wei.storage(ckw::TensorStorageType::BufferUint8Ptr);
217 ckw::TensorSampler sampler_bia;
218 sampler_bia.format(ckw::TensorSamplerFormat::Dim0_Dim1_Dim2);
219 sampler_bia.address_mode_x(sampler_dst.address_mode_x());
222 sampler_bia.storage(ckw::TensorStorageType::BufferUint8Ptr);
232 auto tile_gid_0 = writer->declare_tile(
"gid_0", ckw::TileInfo(ckw::DataType::Int32));
233 auto tile_gid_1 = writer->declare_tile(
"gid_1", ckw::TileInfo(ckw::DataType::Int32));
234 auto tile_gid_2 = writer->declare_tile(
"gid_2", ckw::TileInfo(ckw::DataType::Int32));
236 writer->op_get_global_id(tile_gid_0, 0);
237 writer->op_get_global_id(tile_gid_1, 1);
238 writer->op_get_global_id(tile_gid_2, 2);
240 auto tile_cout0 = writer->declare_tile(
"cout0", ckw::TileInfo(ckw::DataType::Int32));
241 auto tile_mout0 = writer->declare_tile(
"mout0", ckw::TileInfo(ckw::DataType::Int32));
242 auto tile_mout1 = writer->declare_tile(
"mout1", ckw::TileInfo(ckw::DataType::Int32));
243 auto tile_bout0 = writer->declare_tile(
"bout0", ckw::TileInfo(ckw::DataType::Int32));
247 const_shift_back_dst_n0_i32, const_0_i32);
249 writer->op_binary(tile_mout1, ckw::BinaryOp::Mod, tile_gid_2, const_dst_h_i32);
250 writer->op_binary(tile_bout0, ckw::BinaryOp::Div, tile_gid_2, const_dst_h_i32);
252 auto tile_src_ci = writer->declare_tile(
"src_ci", ckw::DataType::Int32);
253 writer->op_binary(tile_src_ci, ckw::BinaryOp::Div, tile_cout0, const_depth_multiplier_i32);
255 auto tile_src_xi = writer->declare_tile(
"src_xi", ckw::DataType::Int32);
256 writer->op_binary(tile_src_xi, ckw::BinaryOp::Mul, tile_mout0, const_stride_x_i32);
257 writer->op_binary(tile_src_xi, ckw::BinaryOp::Sub, tile_src_xi, const_pad_x_i32);
259 auto tile_src_yi = writer->declare_tile(
"src_yi", ckw::DataType::Int32);
260 writer->op_binary(tile_src_yi, ckw::BinaryOp::Mul, tile_mout1, const_stride_y_i32);
261 writer->op_binary(tile_src_yi, ckw::BinaryOp::Sub, tile_src_yi, const_pad_y_i32);
264 auto tile_yk = writer->declare_tile(
"yk", ckw::DataType::Int32);
266 writer->op_assign(tile_yk, const_0_i32);
269 writer->op_for_loop(tile_yk, ckw::BinaryOp::Less, const_kernel_size_i32, tile_yk, ckw::AssignmentOp::Increment, const_kernel_w_i32,
272 auto tile_src = writer->declare_tile(
"a", ckw::TileInfo(
to_ckw(_src->
data_type()), src_m0, src_n0));
273 auto tile_wei = writer->declare_tile(
"b", ckw::TileInfo(
to_ckw(_wei->
data_type()), wei_m0, wei_n0));
275 writer->op_assign(tile_src, const_0_fp);
277 auto tile_x_gte_0 = writer->declare_tile(
"x_gte_0", ckw::TileInfo(ckw::DataType::Int32));
278 auto tile_y_gte_0 = writer->declare_tile(
"y_gte_0", ckw::TileInfo(ckw::DataType::Int32));
279 auto tile_x_lt_w = writer->declare_tile(
"x_lt_w", ckw::TileInfo(ckw::DataType::Int32));
280 auto tile_y_lt_h = writer->declare_tile(
"y_lt_h", ckw::TileInfo(ckw::DataType::Int32));
283 writer->op_binary(tile_y_gte_0, ckw::BinaryOp::GreaterEqual, tile_src_yi, const_0_i32);
284 writer->op_binary(tile_y_lt_h, ckw::BinaryOp::Less, tile_src_yi, const_src_h_i32);
286 auto tile_src_mi = writer->declare_tile(
"src_mi", ckw::TileInfo(ckw::DataType::Int32));
289 for(int32_t xk = 0; xk < src_m0; ++xk)
291 auto const_xk_i32 = writer->declare_constant_tile(ckw::ConstantData({{xk}}, ckw::DataType::Int32));
294 writer->op_binary(tile_src_mi, ckw::BinaryOp::Mul, const_xk_i32, const_dilation_x_i32);
295 writer->op_binary(tile_src_mi, ckw::BinaryOp::Add, tile_src_mi, tile_src_xi);
298 writer->op_binary(tile_x_gte_0, ckw::BinaryOp::GreaterEqual, tile_src_mi, const_0_i32);
299 writer->op_binary(tile_x_lt_w, ckw::BinaryOp::Less, tile_src_mi, const_src_w_i32);
302 writer->op_ternary(tile_src_mi, ckw::TernaryOp::Select, const_neg_1_i32, tile_src_mi, tile_x_gte_0);
303 writer->op_ternary(tile_src_mi, ckw::TernaryOp::Select, const_neg_1_i32, tile_src_mi, tile_x_lt_w);
304 writer->op_ternary(tile_src_mi, ckw::TernaryOp::Select, const_neg_1_i32, tile_src_mi, tile_y_gte_0);
305 writer->op_ternary(tile_src_mi, ckw::TernaryOp::Select, const_neg_1_i32, tile_src_mi, tile_y_lt_h);
307 writer->op_load(tile_src.row(xk),
src->tensor(), sampler_src, tile_src_ci, tile_src_mi, tile_src_yi, tile_bout0);
311 writer->op_load(tile_wei, wei->
tensor(), sampler_wei, tile_cout0, tile_yk, const_0_i32, const_0_i32);
315 auto tile_tmp = writer->declare_tile(
"tmp", ckw::TileInfo(
to_ckw(_src->
data_type()), 1, dst_n0));
316 for(int32_t m0 = 0; m0 < dst_m0; ++m0)
318 for(int32_t xk = 0; xk < kernel_width; ++xk)
320 auto tile_a = tile_src.row(m0 + xk);
321 auto tile_b = tile_wei.row(xk);
322 auto tile_c = tile_dst.row(m0);
324 writer->op_binary(tile_tmp, ckw::BinaryOp::Mul, tile_a, tile_b);
325 writer->op_binary(tile_c, ckw::BinaryOp::Add, tile_c, tile_tmp);
328 writer->op_binary(tile_src_yi, ckw::BinaryOp::Add, tile_src_yi, const_dilation_y_i32);
337 if (!bia->has_tile())
339 auto tile_bia = writer->declare_tile(
"bia", ckw::TileInfo(
to_ckw(_src->data_type()), 1, dst_n0));
340 writer->op_load(tile_bia, bia->tensor(), sampler_bia, tile_cout0, const_0_i32, const_0_i32, const_0_i32);
341 bia->init_virtual_tensor(tile_bia, sampler_bia);
343 auto &tile_bia = bia->tile();
345 writer->op_binary(tile_dst, ckw::BinaryOp::Add, tile_dst, tile_bia);