Split large fully-connected nodes along k.

This is just a `subgraph` rewrite that replaces nodes that look like
```
fully_connected(input_id, weights_id, bias_id)
```
where the data along the `k`th dimension exceeds 16kB, with
```
fully_connected(static_slice(K0, input_id), static_slice(K0, weights_id), bias_id)
  +
fully_connected(static_slice(K1, input_id), static_slice(K1, weights_id), 0)
  +
...
  +
fully_connected(static_slice(Kn, input_id), static_slice(Kn, weights_id), 0)
```
where `static_slice(ival, value_id)` slices `value_id` in the `k` dimension.

PiperOrigin-RevId: 806264392

Author: gonnet

BUILD.bazel

@@ -839,6 +839,7 @@ xnnpack_cc_library(
         ":memory",
         ":microkernel_hdrs",
         ":microkernel_type",
+        ":microkernel_utils",
         ":mutex",
         ":node_type",
         ":operator_type",

src/datatype.c

@@ -1,11 +1,15 @@
-// Copyright 2019 Google LLC
+// Copyright 2019-2025 Google LLC
 //
 // This source code is licensed under the BSD-style license found in the
 // LICENSE file in the root directory of this source tree.
 
 #include "src/xnnpack/datatype.h"
 
+#include <assert.h>
+#include <stddef.h>
+
 #include "include/xnnpack.h"
+#include "src/xnnpack/common.h"
 
 bool xnn_datatype_is_real(enum xnn_datatype t) {
   switch (t) {
@@ -88,6 +92,20 @@ bool xnn_datatype_is_quantized(enum xnn_datatype t) {
   return false;
 }
 
+bool xnn_datatype_is_dynamically_quantized(enum xnn_datatype t) {
+  switch (t) {
+    case xnn_datatype_pqint8:
+    case xnn_datatype_qdint8:
+    case xnn_datatype_qduint8:
+    case xnn_datatype_qpint8:
+      return true;
+    default:
+      return false;
+  }
+  XNN_UNREACHABLE;
+  return false;
+}
+
 bool xnn_datatype_is_channelwise_quantized(enum xnn_datatype t) {
   switch (t) {
     case xnn_datatype_qcint8:

src/operator-run.c

@@ -1494,25 +1494,30 @@ void xnn_compute_pad_5d(struct pad_context* restrict context, size_t i,
   }
 }
 
-void xnn_compute_slice_1d(struct slice_context* restrict context, size_t i) {
-  const void* input =
-      (const void*)((uintptr_t)context->input + i * context->input_stride[0]);
-  void* output =
-      (void*)((uintptr_t)context->output + i * context->output_stride[0]);
+void xnn_compute_slice_1d(struct slice_context* restrict context, size_t offset,
+                          size_t count) {
+  for (size_t i = offset; i < offset + count; i++) {
+    const void* input =
+        (const void*)((uintptr_t)context->input + i * context->input_stride[0]);
+    void* output =
+        (void*)((uintptr_t)context->output + i * context->output_stride[0]);
 
-  context->ukernel(context->contiguous_size, input, output, NULL);
+    context->ukernel(context->contiguous_size, input, output, NULL);
+  }
 }
 
 void xnn_compute_slice_2d(struct slice_context* restrict context, size_t i,
-                          size_t j) {
-  const void* input =
-      (const void*)((uintptr_t)context->input + i * context->input_stride[1] +
-                    j * context->input_stride[0]);
-  void* output =
-      (void*)((uintptr_t)context->output + i * context->output_stride[1] +
-              j * context->output_stride[0]);
+                          size_t offset, size_t count) {
+  for (size_t j = offset; j < offset + count; j++) {
+    const void* input =
+        (const void*)((uintptr_t)context->input + i * context->input_stride[1] +
+                      j * context->input_stride[0]);
+    void* output =
+        (void*)((uintptr_t)context->output + i * context->output_stride[1] +
+                j * context->output_stride[0]);
 
-  context->ukernel(context->contiguous_size, input, output, NULL);
+    context->ukernel(context->contiguous_size, input, output, NULL);
+  }
 }
 
 void xnn_compute_slice_3d(struct slice_context* restrict context, size_t i,

src/operators/fully-connected-nc.c

@@ -2302,9 +2302,9 @@ static enum xnn_status reshape_fully_connected_nc(
   const size_t nc = xnn_gemm_best_tile_size(
       /*num_groups=*/1, /*m=*/batch_size, /*n=*/output_channels,
       /*m_stride=*/
-      fully_connected_op->input_pixel_stride
-          << (packed_lh_config ? packed_lh_config->log2_packed_element_size
-                               : log2_input_element_size),
+      input_channels << (packed_lh_config
+                             ? packed_lh_config->log2_packed_element_size
+                             : log2_input_element_size),
       /*n_stride=*/
       fully_connected_op->weights_stride,
       /*cn_stride=*/1 << log2_output_element_size, mr, nr,
@@ -2362,7 +2362,7 @@ static enum xnn_status reshape_fully_connected_nc(
             .mr = mr_packed,
             .kr = kr,
             .sr = sr,
-            .lhs_stride = input_channels
+            .lhs_stride = fully_connected_op->input_pixel_stride
                           << packed_lh_config->log2_input_element_size,
             .packed_offset_fn = packed_lh_config->offset_fn,
             .pack_lh_ukernel = packed_lh_config->pack_lh_fn,
@@ -2411,6 +2411,7 @@ static enum xnn_status reshape_fully_connected_nc(
                    << log2_output_element_size,
       .cn_stride = nr << log2_output_element_size,
       .log2_csize = log2_output_element_size,
+      .log2_asize = log2_input_element_size,
       .ukernel = gemm_ukernel,
       .mr = mr,
       .nc = output_channels,
@@ -2873,6 +2874,11 @@ static enum xnn_status setup_fully_connected_nc(
   struct gemm_op_context* gemm_context =
       fully_connected_op->dynamic_context.gemm;
 
+  input =
+      (const void*)((uintptr_t)input +
+                    (fully_connected_op->fully_connected.input_element_offset
+                     << gemm_context->gemm.log2_asize));
+
   if (fully_connected_op->num_compute_invocations == 2) {
     gemm_context->pack_lh.lhs = input;
     gemm_context->pack_lh.lhs_packed = workspace;

src/operators/slice-nd.c

@@ -16,6 +16,7 @@
 #include "src/xnnpack/config-types.h"
 #include "src/xnnpack/config.h"
 #include "src/xnnpack/log.h"
+#include "src/xnnpack/math.h"
 #include "src/xnnpack/normalization.h"
 #include "src/xnnpack/operator-type.h"
 #include "src/xnnpack/operator-utils.h"
@@ -221,15 +222,21 @@ static enum xnn_status reshape_slice_nd(
   switch (num_normalized_dims) {
     case 1:
     case 2:
-      slice_op->compute[0].type = xnn_parallelization_type_1d;
-      slice_op->compute[0].task_1d = (pthreadpool_task_1d_t)xnn_compute_slice_1d;
+      slice_op->compute[0].type = xnn_parallelization_type_1d_tile_1d_dynamic;
+      slice_op->compute[0].task_1d_tile_1d_dynamic =
+          (pthreadpool_task_1d_tile_1d_dynamic_t)xnn_compute_slice_1d;
       slice_op->compute[0].range[0] = normalized_output_shape[XNN_MAX_TENSOR_DIMS - 2];
+      slice_op->compute[0].tile[0] =
+          divide_round_up(slice_op->context.slice.contiguous_size, 32 * 1024);
       break;
     case 3:
-      slice_op->compute[0].type = xnn_parallelization_type_2d;
-      slice_op->compute[0].task_2d = (pthreadpool_task_2d_t) xnn_compute_slice_2d;
+      slice_op->compute[0].type = xnn_parallelization_type_2d_tile_1d_dynamic;
+      slice_op->compute[0].task_2d_tile_1d_dynamic =
+          (pthreadpool_task_2d_tile_1d_dynamic_t)xnn_compute_slice_2d;
       slice_op->compute[0].range[0] = normalized_output_shape[XNN_MAX_TENSOR_DIMS - 3];
       slice_op->compute[0].range[1] = normalized_output_shape[XNN_MAX_TENSOR_DIMS - 2];
+      slice_op->compute[0].tile[0] =
+          divide_round_up(slice_op->context.slice.contiguous_size, 32 * 1024);
       break;
     case 4:
       slice_op->compute[0].type = xnn_parallelization_type_3d;

src/runtime.c

@@ -828,6 +828,8 @@ enum xnn_status xnn_create_runtime_with_threadpool(
     threadpool =
         pthreadpool_create_v2(&executor, xnn_threadpool->scheduler_context, 0);
     flags |= XNN_FLAG_RUNTIME_OWNS_THREADPOOL;
+    xnn_log_info("Created pthreadpool from scheduler with %zu threads.",
+                 pthreadpool_get_threads_count(threadpool));
 #endif  // XNN_SLINKY_ENABLED
   }