From be7ae5ce2042bb70b85cae6919de6ac1741fba85 Mon Sep 17 00:00:00 2001 From: Frank Barchard Date: Thu, 3 Sep 2026 20:55:40 -0700 Subject: [PATCH] Fix out-of-bounds reads in transpose_a dot kernels Dot kernels with transposed A read rows in chunks of `tile_m`. Previously, the output buffer of `transpose_a` did not declare this row padding, relying on incidental allocation slack. With Slinky memory pooling enabled, buffers are tightly packed, causing kernels to read past the end of the buffer (b/556792224). This change: 1. Adds `tile_m` dimension alignment to `transpose_a` outputs so Slinky allocates the required row padding. 2. Prevents rewriting `transpose_a(stencil_copy)` into `stencil_copy(transpose_a)` when `tile_m > 1`, avoiding broken convolution sliding on ARM. 3. Fixes stride and schedule calculations in `define_transpose_a` for broadcast and multi-reduction dimensions. PiperOrigin-RevId: 976076613 --- ynnpack/kernels/dot/dot.cc | 5 +- ynnpack/kernels/dot/dot.h | 1 + ynnpack/subgraph/dot.cc | 221 ++++++++++++++++++--------- ynnpack/subgraph/dot.h | 11 +- ynnpack/subgraph/fusion.cc | 14 +- ynnpack/subgraph/subgraph.cc | 3 +- ynnpack/subgraph/subgraph.h | 6 +- ynnpack/subgraph/test/fusion_copy.cc | 5 +- ynnpack/subgraph/test/matchers.h | 7 + 9 files changed, 188 insertions(+), 85 deletions(-) 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);