Fix out-of-bounds reads in transpose_a dot kernels

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;
   }

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;

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`.

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];

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) {}

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 {

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<ynn_node::stencil_copy::stencil>{
-                        {/*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) {

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<ynn_node::transpose_a>(&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<ynn_node::transpose_a>(&arg.op);