diff --git a/ynnpack/kernels/dot/dot.cc b/ynnpack/kernels/dot/dot.cc index 62ca931a567..44d0dce2008 100644 --- a/ynnpack/kernels/dot/dot.cc +++ b/ynnpack/kernels/dot/dot.cc @@ -379,9 +379,8 @@ struct optimizer { return; } result = { - kernel, block_m, block_n, - block_k, tile_n, tile_k, - flags, dot_cost_k, result.max_block_n, + kernel, block_m, block_n, block_k, tile_m, + tile_n, tile_k, flags, dot_cost_k, result.max_block_n, }; kernel_used = name; } diff --git a/ynnpack/kernels/dot/dot.h b/ynnpack/kernels/dot/dot.h index f78923d1a84..fcf0a7600dc 100644 --- a/ynnpack/kernels/dot/dot.h +++ b/ynnpack/kernels/dot/dot.h @@ -99,6 +99,7 @@ struct dot_kernel { int block_m = 0; int block_n = 0; int block_k = 0; + int tile_m = 0; int tile_n = 0; int tile_k = 0; uint32_t flags = 0; diff --git a/ynnpack/subgraph/dot.cc b/ynnpack/subgraph/dot.cc index 9b3c978aecf..3625c14c914 100644 --- a/ynnpack/subgraph/dot.cc +++ b/ynnpack/subgraph/dot.cc @@ -41,6 +41,7 @@ #include "slinky/runtime/buffer.h" #include "slinky/runtime/evaluate.h" #include "slinky/runtime/expr.h" +#include "slinky/runtime/print.h" #include "slinky/runtime/stmt.h" using slinky::index_t; @@ -286,12 +287,13 @@ auto make_dot_impl(dot_type type, bool consistent_arithmetic, bool symmetric_b, index_t k3_extent = num_k_dims >= 3 ? r_k3.extent() : 1; // Learn what we need to know about m, n, k1, k2, k3 before slicing them. - const int a_k1_dim = transposed_a ? 1 : 0; + const int a_k1_dim = transposed_a ? 2 : 0; const slinky::dim& init_c_m = init_c.dim(1); const slinky::dim& init_c_n = init_c.dim(0); const slinky::dim& c_m = c.dim(1); const slinky::dim& c_n = c.dim(0); const slinky::dim& a_k1i = transposed_a ? a.dim(0) : dummy_dim; + const slinky::dim& a_mi = transposed_a ? a.dim(1) : dummy_dim; const slinky::dim& a_k1o = a.dim(a_k1_dim); const slinky::dim& a_k2 = num_k_dims >= 2 ? a.dim(a_k1_dim + 1) : dummy_dim; const slinky::dim& a_k3 = num_k_dims >= 3 ? a.dim(a_k1_dim + 2) : dummy_dim; @@ -314,7 +316,7 @@ auto make_dot_impl(dot_type type, bool consistent_arithmetic, bool symmetric_b, // If a is transposed, then the k dimension has been reshaped to have // `tile_k` values in each element. const index_t a_tile_k = a_k1i.extent(); - const index_t a_stride_m = a_m.stride(); + const index_t a_stride_m = transposed_a ? a_mi.stride() : a_m.stride(); const index_t a_stride_k3 = a_k3.stride(); const index_t a_stride_k2 = a_k2.stride(); const index_t a_stride_k1 = a_k1o.stride() / a_tile_k; @@ -388,6 +390,7 @@ auto make_dot_impl(dot_type type, bool consistent_arithmetic, bool symmetric_b, assert(!c_n.is_folded()); assert(!a_m.is_folded(c_m.min(), c_m.max())); assert(!a_k1i.is_folded()); + assert(!a_mi.is_folded()); assert(!a_k1o.is_folded()); assert(!a_k2.is_folded()); assert(!a_k3.is_folded()); @@ -427,7 +430,13 @@ auto make_dot_impl(dot_type type, bool consistent_arithmetic, bool symmetric_b, for (size_t i = 0; i < a_k1_dim + num_k_dims; ++i) { a.slice(0); } - a.slice(0, slinky::in_bounds{c_m.min()}); + if (transposed_a) { + const index_t tile_m = a_mi.extent(); + assert(c_m.min() % tile_m == 0); + a.slice(0, slinky::in_bounds{c_m.min() / tile_m}); + } else { + a.slice(0, slinky::in_bounds{c_m.min()}); + } if (pack_b) { // If b is packed, we must slice b at blocks of n. assert(c_n.min() % block_n == 0); @@ -785,21 +794,30 @@ auto make_transpose_a_impl(int m_dim) { const slinky::dim& input_k = input.dim(0); const slinky::dim& input_m = input.dim(m_dim); const slinky::dim& output_ki = output.dim(0); - const slinky::dim& output_ko = output.dim(1); - const slinky::dim& output_m = output.dim(m_dim + 1); + const slinky::dim& output_mi = output.dim(1); + const slinky::dim& output_ko = output.dim(2); + const slinky::dim& output_mo = output.dim(m_dim + 2); const index_t tile_k = output_ki.extent(); + const index_t tile_m = output_mi.extent(); const index_t elem_size = input.elem_size; - assert(output_m.extent() == 1 || output_m.stride() == elem_size * tile_k); - (void)output_m; + assert(output_ki.min() == 0); + assert(output_mi.min() == 0); + assert(output_ki.extent() == 1 || output_ki.stride() == elem_size); + assert(output_mi.extent() == 1 || output_mi.stride() == elem_size * tile_k); + assert(output_mo.extent() == 1 || + output_mo.stride() == elem_size * tile_k * tile_m); // We need the intersection of the input and output bounds. - const index_t m = std::max( - 0, std::min(output_m.end(), input_m.end()) - output_m.min()); - assert(input_k.min() <= output_ko.min() * tile_k); + const index_t m_begin = output_mo.begin() * tile_m; + const index_t m_end = output_mo.end() * tile_m; + const index_t m = + std::max(0, std::min(m_end, input_m.end()) - m_begin); + const index_t k_begin = output_ko.begin() * tile_k; + const index_t k_end = output_ko.end() * tile_k; + assert(input_k.min() <= k_begin); const index_t k = - std::max(0, std::min(output_ko.end() * tile_k, input_k.end()) - - output_ko.min() * tile_k); + std::max(0, std::min(k_end, input_k.end()) - k_begin); // We're transposing columns of the input to rows of the output, but // doing tile_k of them at a time. @@ -809,9 +827,9 @@ auto make_transpose_a_impl(int m_dim) { const index_t input_m_stride = input_m.stride(); const index_t output_ko_stride = output_ko.stride(); - input.slice(0, slinky::in_bounds{output_ko.min() * tile_k}); - input.slice(m_dim - 1, output_m.min()); - output.slice({0, 1, static_cast(m_dim + 1)}); + input.slice(0, slinky::in_bounds{k_begin}); + input.slice(m_dim - 1, slinky::in_bounds{m_begin}); + output.slice({0, 1, 2, static_cast(m_dim + 2)}); slinky::for_each_element( [=, &p](void* output, const void* input) { @@ -825,28 +843,33 @@ auto make_transpose_a_impl(int m_dim) { } // namespace -// Packing means transposing -// a(k, m, ...) => a([0, tile_k), m, k/tile_k, ...) -void define_transpose_a(ynn_subgraph& subgraph, ynn_node& node, index_t tile_k, - int m_dim, uint32_t input_a_id, uint32_t output_id) { +// a(k, m, ...) => a([0, tile_k), [0, tile_m), k/tile_k, ..., m/tile_m, ...) +void define_transpose_a(ynn_subgraph& subgraph, ynn_node& node, index_t tile_m, + index_t tile_k, int m_dim, uint32_t input_a_id, + uint32_t output_id) { const ynn_value& a = subgraph.value(input_a_id); ynn_value& output = subgraph.get_output_value(&output_id, a.type); output.type = a.type; slinky::expr k = a.extent(0); + slinky::expr m = a.extent(m_dim); output.extents = a.extents; - while (output.extents.size() < 2) { + while (output.extents.size() <= static_cast(m_dim)) { output.extents.push_back(slinky::expr{}); } output.extents[0] = slinky::simplify(slinky::ceil_div(k, tile_k)); - output.extents.insert(output.extents.begin(), tile_k); + output.extents[m_dim] = + slinky::simplify(slinky::ceil_div(m, tile_m)); + output.extents.insert(output.extents.begin(), {tile_k, tile_m}); node.inputs = {input_a_id}; node.outputs = {output.id}; - node.op = ynn_node::transpose_a{static_cast(tile_k), m_dim}; + node.op = ynn_node::transpose_a{static_cast(tile_m), + static_cast(tile_k), m_dim}; node.create = [](const ynn_node& node, ynn_runtime& runtime) { const ynn_node::transpose_a& op = std::get(node.op); + const index_t tile_m = op.tile_m; const index_t tile_k = op.tile_k; const int m_dim = op.m_dim; const ynn_runtime_value& input = runtime.value(node.inputs[0]); @@ -854,33 +877,53 @@ void define_transpose_a(ynn_subgraph& subgraph, ynn_node& node, index_t tile_k, slinky::expr elem_size = input.buffer->elem_size(); output.make_buffer(runtime, elem_size); - output.buffer->dim(0).stride = elem_size; - output.buffer->dim(m_dim + 1).stride = elem_size * tile_k; - output.buffer->dim(1).stride = - elem_size * tile_k * output.buffer->dim(m_dim + 1).extent(); + slinky::expr stride = elem_size; + output.buffer->dim(0).stride = stride; + stride *= tile_k; + output.buffer->dim(1).stride = stride; + stride *= tile_m; + output.buffer->dim(m_dim + 2).stride = stride; + stride *= output.buffer->dim(m_dim + 2).extent(); + for (int d = 3; d <= m_dim + 1; ++d) { + output.buffer->dim(d).stride = stride; + stride *= output.buffer->dim(d).extent(); + } + output.buffer->dim(2).stride = stride; + stride *= output.buffer->dim(2).extent(); + for (size_t d = m_dim + 3; d < output.buffer->rank(); ++d) { + if (slinky::prove_true(output.buffer->dim(d).is_broadcast())) continue; + output.buffer->dim(d).stride = stride; + stride *= output.buffer->dim(d).extent(); + } // Don't allow folding of dimensions we transpose. - output.buffer->dim(0).fold_factor = slinky::dim::unfolded; - output.buffer->dim(m_dim + 1).fold_factor = slinky::dim::unfolded; - output.buffer->dim(1).fold_factor = slinky::dim::unfolded; + for (size_t d = 0; d < output.buffer->rank(); ++d) { + if (slinky::prove_true(output.buffer->dim(d).is_broadcast())) continue; + output.buffer->dim(d).fold_factor = slinky::dim::unfolded; + } // Split + Transpose std::vector dims = runtime.globals.make_dims(output.buffer->rank()); - slinky::expr ko = dims[1]; + slinky::expr ko = dims[2]; + slinky::expr mo = dims[m_dim + 2]; slinky::func::input func_input = {input.buffer}; - func_input.bounds = { - slinky::min_extent(ko * tile_k, tile_k), - }; - for (size_t i = 2; i < dims.size(); ++i) { - func_input.bounds.push_back(slinky::point(dims[i])); + func_input.bounds.resize(input.buffer->rank()); + func_input.bounds[0] = slinky::min_extent(ko * tile_k, tile_k); + for (int i = 1; i < m_dim; ++i) { + func_input.bounds[i] = slinky::point(dims[i + 2]); + } + func_input.bounds[m_dim] = slinky::min_extent(mo * tile_m, tile_m); + for (size_t i = m_dim + 1; i < input.buffer->rank(); ++i) { + func_input.bounds[i] = slinky::point(dims[i + 2]); } - // This transpose handles padding the input up to tile_k. - func_input.input_crop = { - all_bounds(input.extent(0)), - }; + + // This transpose handles padding the input up to tile_k and tile_m. + func_input.input_crop.resize(input.buffer->rank()); + func_input.input_crop[0] = all_bounds(input.extent(0)); + func_input.input_crop[m_dim] = all_bounds(input.extent(m_dim)); slinky::call_stmt::attributes attrs; attrs.name = "transpose_a"; @@ -888,6 +931,25 @@ void define_transpose_a(ynn_subgraph& subgraph, ynn_node& node, index_t tile_k, {std::move(func_input)}, {{output.buffer, dims}}, std::move(attrs)); + std::vector given_splits = {output.physical_extent(0), + output.physical_extent(1)}; + auto sched = + runtime.make_schedule(dims, output.physical_extents(), + output.buffer->elem_size(), given_splits); + sched->loop_splits[0].step_is_required = true; + sched->loop_splits[1].step_is_required = true; + + if (m_dim > 1) { + sched->force_root = true; + } + + sched->input_scheduler_bounds.resize(1); + sched->input_scheduler_bounds[0].resize(m_dim + 1); + sched->input_scheduler_bounds[0][m_dim] = slinky::point(mo); + + func.user_data() = sched.get(); + runtime.scheduling_info_storage.push_back(std::move(sched)); + runtime.funcs.push_back(std::move(func)); return ynn_status_success; }; @@ -895,11 +957,13 @@ void define_transpose_a(ynn_subgraph& subgraph, ynn_node& node, index_t tile_k, namespace { -uint32_t define_transpose_a(ynn_subgraph& subgraph, index_t tile_k, - int32_t m_dim, uint32_t input_a_id) { +uint32_t define_transpose_a(ynn_subgraph& subgraph, index_t tile_m, + index_t tile_k, int32_t m_dim, + uint32_t input_a_id) { ynn_node node; ynn_value& output = subgraph.new_internal_value(); - ynn::define_transpose_a(subgraph, node, tile_k, m_dim, input_a_id, output.id); + ynn::define_transpose_a(subgraph, node, tile_m, tile_k, m_dim, input_a_id, + output.id); subgraph.add_node(std::move(node)); return output.id; } @@ -1360,8 +1424,8 @@ ynn_status define_dot(ynn_subgraph& subgraph, size_t num_k_dims, // The kernel we want to use has a transposed a. // By definition, `m_dim` is the first dimension after the k dims. const int m_dim = num_k_dims; - node.inputs[0] = - define_transpose_a(subgraph, kernel.tile_k, m_dim, input_a_id); + node.inputs[0] = define_transpose_a(subgraph, kernel.tile_m, kernel.tile_k, + m_dim, input_a_id); } // If we're using an unpacked kernel, we'll be reading columns of B, make sure @@ -1372,8 +1436,9 @@ ynn_status define_dot(ynn_subgraph& subgraph, size_t num_k_dims, : std::max(YNN_CACHE_LINE_SIZE / b_elem_size, unpacked_kernel.block_n); node.create = [consistent_arithmetic, symmetric_b, pack_b, transpose_a, - block_n_unpacked, tile_k = kernel.tile_k]( - const ynn_node& node, ynn_runtime& runtime) { + block_n_unpacked, tile_k = kernel.tile_k, + tile_m = kernel.tile_m](const ynn_node& node, + ynn_runtime& runtime) { const ynn_node::dot& op = std::get(node.op); const size_t num_k_dims = op.num_k_dims; ynn_runtime_value& input_a = runtime.value(node.inputs[0]); @@ -1413,7 +1478,7 @@ ynn_status define_dot(ynn_subgraph& subgraph, size_t num_k_dims, all_dims.push_back(r_dim); reduction_dims.push_back(r_dim); - const int a_k_dim = transpose_a ? 1 : 0; + const int a_k_dim = transpose_a ? 2 : 0; slinky::expr k_extent = input_a.extent(a_k_dim + d); if (transpose_a && d == 0) { // When A is transposed, its K1 dimension is split into blocks of size @@ -1437,21 +1502,20 @@ ynn_status define_dot(ynn_subgraph& subgraph, size_t num_k_dims, } // A: We need all of the k dims, i is elementwise. - const int num_a_k_dims = num_k_dims + (transpose_a ? 1 : 0); + const int num_a_k_dims = num_k_dims + (transpose_a ? 2 : 0); slinky::box_expr a_bounds(std::min(input_a.rank(), num_a_k_dims)); if (transpose_a) { - a_bounds[0] = all_bounds(tile_k); - for (size_t d = 1; d < a_bounds.size(); ++d) { - int k_idx = d - 1; - if (k_idx == 0) { - // Since the reduction dimension represents the total number of - // elements, we need to divide by the block size (tile_k) to get the - // corresponding block index for the transposed A buffer. - a_bounds[d] = - slinky::point(slinky::simplify(reduction_dims[k_idx] / tile_k)); - } else { - a_bounds[d] = slinky::point(reduction_dims[k_idx]); - } + a_bounds[0] = all_bounds(input_a.physical_extent(0)); + a_bounds[1] = all_bounds(input_a.physical_extent(1)); + a_bounds[2] = slinky::point(slinky::simplify(reduction_dims[0] / tile_k)); + for (size_t d = 1; d < num_k_dims; ++d) { + a_bounds[2 + d] = slinky::point(reduction_dims[d]); + } + if (output_dims.size() >= 2) { + slinky::var i = output_dims[1]; + a_bounds.push_back(slinky::point(i) / tile_m); + } else { + a_bounds.push_back(slinky::point(0)); } } else { for (size_t d = 0; d < a_bounds.size(); ++d) { @@ -1480,23 +1544,33 @@ ynn_status define_dot(ynn_subgraph& subgraph, size_t num_k_dims, // C: Elementwise slinky::box_expr c_bounds; if (input_c.rank() >= 1) { - c_bounds.push_back( - elementwise_bounds(output_dims[0], input_c.physical_extent(0))); + c_bounds.push_back(make_broadcast_bounds(j, input_c.physical_extent(0), + output.extent(0))); } // Batch dims are elementwise too. for (size_t i = 1; i < output_dims.size(); ++i) { - if (i + num_a_k_dims - 1 < input_a.rank()) { - a_bounds.push_back(elementwise_bounds( - output_dims[i], input_a.physical_extent(i + num_a_k_dims - 1))); + if (transpose_a) { + if (i >= 2 && i + num_a_k_dims - 1 < input_a.rank()) { + a_bounds.push_back(make_broadcast_bounds( + output_dims[i], input_a.physical_extent(i + num_a_k_dims - 1), + output.extent(i))); + } + } else { + if (i + num_a_k_dims - 1 < input_a.rank()) { + a_bounds.push_back(make_broadcast_bounds( + output_dims[i], input_a.physical_extent(i + num_a_k_dims - 1), + output.extent(i))); + } } if (i >= 2 && i + 2 + num_k_dims - 1 < packed_b.rank()) { - b_bounds.push_back(elementwise_bounds( - output_dims[i], packed_b.physical_extent(i + 2 + num_k_dims - 1))); + b_bounds.push_back(make_broadcast_bounds( + output_dims[i], packed_b.physical_extent(i + 2 + num_k_dims - 1), + output.extent(i))); } if (i < input_c.rank()) { - c_bounds.push_back( - elementwise_bounds(output_dims[i], input_c.physical_extent(i))); + c_bounds.push_back(make_broadcast_bounds( + output_dims[i], input_c.physical_extent(i), output.extent(i))); } } @@ -1609,11 +1683,18 @@ ynn_status define_dot(ynn_subgraph& subgraph, size_t num_k_dims, // indices (j / block_n), which breaks the scheduler's source region // inference. Declare a virtual 1-to-1 mapping with `j` instead, so the // pack (and anything feeding it) can be fused with loops derived from j. + sched->input_scheduler_bounds.resize(2); if (pack_b) { - sched->input_scheduler_bounds.resize(2); sched->input_scheduler_bounds[1].resize(4); sched->input_scheduler_bounds[1][3] = slinky::point(j); } + if (transpose_a) { + sched->input_scheduler_bounds[0].resize(num_a_k_dims + 1); + if (output_dims.size() >= 2) { + slinky::var i = output_dims[1]; + sched->input_scheduler_bounds[0][num_a_k_dims] = slinky::point(i); + } + } func.user_data() = sched.get(); runtime.scheduling_info_storage.push_back(std::move(sched)); diff --git a/ynnpack/subgraph/dot.h b/ynnpack/subgraph/dot.h index 108e7148ad7..8c05a48a090 100644 --- a/ynnpack/subgraph/dot.h +++ b/ynnpack/subgraph/dot.h @@ -17,8 +17,15 @@ namespace ynn { void define_transpose_a(ynn_subgraph& subgraph, ynn_node& node, - slinky::index_t tile_k, int m_dim, uint32_t input_a_id, - uint32_t output_id); + slinky::index_t tile_m, slinky::index_t tile_k, + int m_dim, uint32_t input_a_id, uint32_t output_id); + +inline void define_transpose_a(ynn_subgraph& subgraph, ynn_node& node, + slinky::index_t tile_k, int m_dim, + uint32_t input_a_id, uint32_t output_id) { + define_transpose_a(subgraph, node, /*tile_m=*/1, tile_k, m_dim, input_a_id, + output_id); +} // Returns true if dots of type uint8 x `b_type` are faster than dots of type // int8 x `b_type`. diff --git a/ynnpack/subgraph/fusion.cc b/ynnpack/subgraph/fusion.cc index faeb0793383..cfb3fffed0d 100644 --- a/ynnpack/subgraph/fusion.cc +++ b/ynnpack/subgraph/fusion.cc @@ -952,8 +952,12 @@ bool rewrite_transpose_stencil_copy(ynn_subgraph& subgraph, ynn_node& node, if (transpose_a == nullptr) { return false; } + const int tile_m = transpose_a->tile_m; const int tile_k = transpose_a->tile_k; const int m_dim = transpose_a->m_dim; + if (tile_m > 1) { + return false; + } auto producer_it = analysis.producers.find(node.inputs[0]); if (producer_it == analysis.producers.end()) { @@ -1008,11 +1012,11 @@ bool rewrite_transpose_stencil_copy(ynn_subgraph& subgraph, ynn_node& node, YNN_LOG_DEBUG() << "Rewriting transpose_a(stencil_copy(x)) to " "stencil_copy(transpose_a(x))"; - // transpose_a inserts a new dimension 0, update our stencil dimensions to - // account for this. + // transpose_a inserts new dimensions at 0 and 1, update our stencil + // dimensions to account for this. for (ynn_node::stencil_copy::stencil& stencil : stencil_copy.stencils) { - stencil.axis++; - stencil.new_axis++; + stencil.axis += 2; + stencil.new_axis += 2; } uint32_t stencil_input_id = stencil_node->inputs[0]; @@ -1023,7 +1027,7 @@ bool rewrite_transpose_stencil_copy(ynn_subgraph& subgraph, ynn_node& node, // Replace stencil_copy(x) with transpose_a'(x), reusing the stencil_node's x // input and y output. - ynn::define_transpose_a(subgraph, *stencil_node, tile_k, new_m_dim, + ynn::define_transpose_a(subgraph, *stencil_node, tile_m, tile_k, new_m_dim, stencil_input_id, stencil_output_id); uint32_t output_id = node.outputs[0]; diff --git a/ynnpack/subgraph/subgraph.cc b/ynnpack/subgraph/subgraph.cc index 1ebcc674dbf..393ff048a5e 100644 --- a/ynnpack/subgraph/subgraph.cc +++ b/ynnpack/subgraph/subgraph.cc @@ -1349,7 +1349,8 @@ void print(std::ostream& os, const ynn_node::iota& op) {} void print(std::ostream& os, const ynn_node::pack_b& op) {} void print(std::ostream& os, const ynn_node::transpose_a& op) { - os << "tile_k=" << op.tile_k << " m_dim=" << op.m_dim; + os << "tile_m=" << op.tile_m << " tile_k=" << op.tile_k + << " m_dim=" << op.m_dim; } void print(std::ostream& os, const ynn_node::dequantize_dot& op) {} diff --git a/ynnpack/subgraph/subgraph.h b/ynnpack/subgraph/subgraph.h index 31d72fb690c..374962ef37c 100644 --- a/ynnpack/subgraph/subgraph.h +++ b/ynnpack/subgraph/subgraph.h @@ -522,13 +522,15 @@ struct ynn_node { friend bool operator<(const pack_b&, const pack_b&) { return false; } }; struct transpose_a { + size_t tile_m = 1; size_t tile_k; int32_t m_dim; friend bool operator==(const transpose_a& a, const transpose_a& b) { - return a.tile_k == b.tile_k && a.m_dim == b.m_dim; + return a.tile_m == b.tile_m && a.tile_k == b.tile_k && a.m_dim == b.m_dim; } friend bool operator<(const transpose_a& a, const transpose_a& b) { - return std::tie(a.tile_k, a.m_dim) < std::tie(b.tile_k, b.m_dim); + return std::tie(a.tile_m, a.tile_k, a.m_dim) < + std::tie(b.tile_m, b.tile_k, b.m_dim); } }; struct get_tensor_shape { diff --git a/ynnpack/subgraph/test/fusion_copy.cc b/ynnpack/subgraph/test/fusion_copy.cc index 6b40bc2c94a..e4e59b06d3b 100644 --- a/ynnpack/subgraph/test/fusion_copy.cc +++ b/ynnpack/subgraph/test/fusion_copy.cc @@ -143,7 +143,7 @@ TEST(fusion, transpose_stencil_copy) { EXPECT_THAT(ProducerOf(z_id, subgraph), AllOf(IsStencilCopy(std::vector{ - {/*axis=*/2, /*new_axis=*/3, /*extent=*/3, + {/*axis=*/3, /*new_axis=*/4, /*extent=*/3, /*stride=*/1, /*dilation=*/1}}), InputsAre(y_id, YNN_INVALID_VALUE_ID))); @@ -155,7 +155,8 @@ TEST(fusion, transpose_stencil_copy) { // m_dim was 1 (size 10). // new_axis (2) > m_dim (1), so m_dim is not decremented. EXPECT_THAT(ProducerOf(y_id, subgraph), - AllOf(IsTransposeA(/*tile_k=*/4, /*m_dim=*/1), InputsAre(x_id))); + AllOf(IsTransposeA(/*tile_m=*/1, /*tile_k=*/4, /*m_dim=*/1), + InputsAre(x_id))); } TEST(fusion, transpose_stencil_copy_grouped) { diff --git a/ynnpack/subgraph/test/matchers.h b/ynnpack/subgraph/test/matchers.h index e9170e4f83e..7c6992822ab 100644 --- a/ynnpack/subgraph/test/matchers.h +++ b/ynnpack/subgraph/test/matchers.h @@ -165,6 +165,13 @@ MATCHER_P(IsUnary, op_type, "") { // // Example: // EXPECT_THAT(ProducerOf(y_id, subgraph), IsTransposeA(16, 2)); +MATCHER_P3(IsTransposeA, tile_m, tile_k, m_dim, "") { + const ynn_node::transpose_a* transpose = + std::get_if(&arg.op); + return transpose && transpose->tile_m == tile_m && + transpose->tile_k == tile_k && transpose->m_dim == m_dim; +} + MATCHER_P2(IsTransposeA, tile_k, m_dim, "") { const ynn_node::transpose_a* transpose = std::get_if(&arg.op);