Update on "Use c10 version of half/bfloat16 in executorch"

Accomplished by importing relevant files from c10 into
executorch/runtime/core/portable_type/c10, and then using `using` in
the top-level ExecuTorch headers. This approach should keep the
ExecuTorch build hermetic for embedded use cases. In the future, we
should add a CI job to ensure the c10 files stay identical to the
PyTorch ones.

Differential Revision: [D66106969](https://our.internmc.facebook.com/intern/diff/D66106969/)

[ghstack-poisoned]

Author: swolchok

.lintrunner.toml

@@ -303,7 +303,7 @@ include_patterns = [
     # 'exir/**/*.py',
     # 'extension/**/*.py',
     'kernels/**/*.py',
-    # 'profiler/**/*.py',
+    'profiler/**/*.py',
     'runtime/**/*.py',
     'scripts/**/*.py',
     # 'test/**/*.py',
@@ -314,6 +314,7 @@ exclude_patterns = [
     'third-party/**',
     '**/third-party/**',
     'scripts/check_binary_dependencies.py',
+    'profiler/test/test_profiler_e2e.py',
 ]
 command = [
     'python',

.mypy.ini

@@ -40,6 +40,9 @@ follow_untyped_imports = True
 [mypy-executorch.kernels.*]
 follow_untyped_imports = True
 
+[mypy-executorch.profiler.*]
+follow_untyped_imports = True
+
 [mypy-executorch.runtime.*]
 follow_untyped_imports = True
 

backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl

@@ -35,7 +35,10 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
  * output at a single output location.
  */
 void main() {
-  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+  const ivec3 pos = ivec3(
+    gl_GlobalInvocationID.x % out_limits.x,
+    (gl_GlobalInvocationID.x / out_limits.x) % out_limits.y,
+    gl_GlobalInvocationID.x / (out_limits.x * out_limits.y));
 
   if (any(greaterThanEqual(pos, out_limits))) {
     return;

backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.glsl

@@ -14,6 +14,8 @@
 
 #define TILE_SIZE ${TILE_SIZE}
 
+#define BATCH_SIZE_Y ${BATCH_SIZE_Y}
+
 #define op(X, A, B) ${OPERATOR}
 
 #include "indexing_utils.h"
@@ -39,9 +41,20 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
  * output at a single output location.
  */
 void main() {
-  const u16vec3 pos = u16vec3(gl_GlobalInvocationID);
+  // y divided up by batch size is used to determine 3d position
+  // since work size is calculated by x * ((y + B_Y - 1) / B_Y) * z
+  const uint out_limits_y_scaled = (out_limits.y + BATCH_SIZE_Y - 1) / BATCH_SIZE_Y;
+
+  u16vec3 pos = u16vec3(
+    gl_GlobalInvocationID.x % out_limits.x,
+    ((gl_GlobalInvocationID.x / out_limits.x) % out_limits_y_scaled),
+    gl_GlobalInvocationID.x / (out_limits.x * out_limits_y_scaled));
 
-  if (any(greaterThanEqual(pos, out_limits))) {
+  // scale pos.y by batch size, because that's the top pixel to be processed
+  pos.y *= uint16_t(BATCH_SIZE_Y);
+
+  // do not process if top pixel does not fit within the output range
+  if (any(greaterThanEqual(u16vec3(pos.x, pos.y, pos.z), out_limits))) {
     return;
   }
 
@@ -54,18 +67,47 @@ void main() {
   const u16vec2 start = ipos;
   const u16vec2 end = ipos + u16vec2(overlay_region.xy);
 
-  VEC4_T sum = texelFetch(t_bias, u16vec2(pos.z, 0), 0);
+  // sum outputs
+  VEC4_T sum[BATCH_SIZE_Y];
+
+  sum[0] = texelFetch(t_bias, u16vec2(pos.z, 0), 0);
+  for (int i = 1; i < BATCH_SIZE_Y; i++) {
+    sum[i] = sum[0];
+  }
+
+  // array to store input texels
+  VEC4_T in_texels[TILE_SIZE];
+
+  // array to store kernel data of previous y
+  VEC4_T prev_kernel_line[TILE_SIZE];
+
   uint16_t kx = uint16_t(0);
-  for (uint16_t y = start.y, i = uint16_t(0); i < uint16_t(TILE_SIZE); y += uint16_t(dilation.y), i++) {
+  for (uint16_t y = start.y, i = uint16_t(0); i < uint16_t(TILE_SIZE + BATCH_SIZE_Y - 1); y += uint16_t(dilation.y), i++) {
     for (uint16_t x = start.x, j = uint16_t(0); j < uint16_t(TILE_SIZE); x += uint16_t(dilation.x), j++) {
-      // The weight kernel was rearranged such that every NxN filter is
-      // flattened to fit in one row. Each filter was then stacked on top of
-      // each other vertically.
-      const vec4 in_texel = texelFetch(t_in, u16vec3(x, y, pos.z), 0);
-      sum = fma(in_texel, texelFetch(t_kernel, u16vec2(kx, pos.z), 0), sum);
-      kx++;
+      in_texels[int(j)] = texelFetch(t_in, u16vec3(x, y, pos.z), 0);
+    }
+
+    // from 2nd iteration onwards accumulate dot product in 2nd sum
+    // based on kernel line data fetched in previous iteration and input texel from this iteration
+    if (i > uint16_t(0)) {
+      for (uint16_t j = uint16_t(0); j < uint16_t(TILE_SIZE); j++) {
+        sum[1] = fma(in_texels[int(j)], prev_kernel_line[int(j)], sum[1]);
+      }
+    }
+
+    // accumulate dot product in 1st sum only until tile size
+    if (i < uint16_t(TILE_SIZE)) {
+      for (uint16_t j = uint16_t(0); j < uint16_t(TILE_SIZE); j++, kx++) {
+        prev_kernel_line[int(j)] = texelFetch(t_kernel, u16vec2(kx, pos.z), 0);
+        sum[0] = fma(in_texels[int(j)], prev_kernel_line[int(j)], sum[0]);
+      }
     }
   }
 
-  imageStore(t_out, pos, op(sum, out_min, out_max));
+  for (int i = 0; i < BATCH_SIZE_Y; i++) {
+    if (any(greaterThanEqual(u16vec3(pos.x, pos.y + i, pos.z), out_limits))) {
+      continue;
+    }
+    imageStore(t_out, u16vec3(pos.x, pos.y + i, pos.z), op(sum[i], out_min, out_max));
+  }
 }

backends/vulkan/runtime/graph/ops/glsl/conv2d_dw_output_tile.yaml

@@ -10,6 +10,7 @@ conv2d_dw_output_tile:
     NDIM: 3
     DTYPE: float
     TILE_SIZE: 3
+    BATCH_SIZE_Y: 2
   generate_variant_forall:
     DTYPE:
       - VALUE: half

backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl

@@ -34,32 +34,42 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
 #extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
 
+// shared memory to hold calculated positions, this would reduce register usage thus improving performance.
+shared u16vec2 pos_shared[gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroupSize.z * TILE_SIZE * TILE_SIZE];
+
 /*
  * Computes a 2D pointwise convolution of an NxN output tile. Calculating an
  * output tile for pointwise convolution is more efficient because the kernel
  * size is only 1x1, making it easier to re-use loaded texels from t_kernel.
  */
 void main() {
-  const u16vec3 gpos = u16vec3(gl_GlobalInvocationID);
+  const uvec2 out_limits_scaled = (out_limits.xy + TILE_SIZE - 1) / TILE_SIZE;
+  const uint shared_mem_stride = gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroupSize.z;
+
+  const u16vec3 gpos = u16vec3(
+    gl_GlobalInvocationID.x % out_limits_scaled.x,
+    (gl_GlobalInvocationID.x / out_limits_scaled.x) % out_limits_scaled.y,
+    gl_GlobalInvocationID.x / (out_limits_scaled.x * out_limits_scaled.y));
 
   // Output position for TILE_SIZE = 2
   // +--------+--------+
   // | pos[0] | pos[1] |
   // +--------+--------+
   // | pos[2] | pos[3] |
   // +--------+--------+
-  u16vec3 pos[TILE_SIZE * TILE_SIZE];
+  u16vec2 pos[TILE_SIZE * TILE_SIZE];
   for (int y = 0, i = 0; y < TILE_SIZE; ++y) {
     for (int x = 0; x < TILE_SIZE; ++x) {
-      pos[i] = u16vec3(
-          gpos.x * TILE_SIZE + x, gpos.y * TILE_SIZE + y, gpos.z);
+      pos[i] = u16vec2(
+          gpos.x * TILE_SIZE + x, gpos.y * TILE_SIZE + y);
+      pos_shared[(shared_mem_stride * i) + gl_LocalInvocationIndex] = pos[i];
       i++;
     }
   }
 
   // If the top left position is out of bounds, then this invocation will have
   // no work to do.
-  if (any(greaterThanEqual(pos[0], out_limits))) {
+  if (any(greaterThanEqual(u16vec3(pos[0], gpos.z), out_limits))) {
     return;
   }
 
@@ -68,7 +78,7 @@ void main() {
   // the top-left element is in a region added by padding.
   u16vec2 ipos[TILE_SIZE * TILE_SIZE];
   for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) {
-    ipos[i] = pos[i].xy * u16vec2(stride) - u16vec2(padding);
+    ipos[i] = pos[i] * u16vec2(stride) - u16vec2(padding);
   }
 
   vec4 sum[TILE_SIZE * TILE_SIZE];
@@ -133,8 +143,9 @@ void main() {
   }
 
   for (int i = 0; i < TILE_SIZE * TILE_SIZE; ++i) {
-    if (all(lessThan(pos[i], out_limits))) {
-      imageStore(t_out, pos[i], op(sum[i], out_min, out_max));
+    const u16vec2 pos = pos_shared[(shared_mem_stride * i) + gl_LocalInvocationIndex];
+    if (all(lessThan(u16vec3(pos, gpos.z), out_limits))) {
+      imageStore(t_out, u16vec3(pos, gpos.z), op(sum[i], out_min, out_max));
     }
   }
 }

backends/vulkan/runtime/graph/ops/impl/Convolution.cpp

@@ -296,6 +296,12 @@ utils::uvec3 create_conv2d_global_wg_size(
         utils::div_up(image_extents[0u], 2u),
         utils::div_up(image_extents[1u], 2u),
         image_extents[2u]};
+  } else if (method == Conv2dMethod::Depthwise) {
+    const utils::uvec3 image_extents = graph.logical_limits_of(out);
+    return {
+        utils::div_up(image_extents[0u], 1u),
+        utils::div_up(image_extents[1u], 2u),
+        image_extents[2u]};
   } else {
     return graph.create_global_wg_size(out);
   }
@@ -370,11 +376,17 @@ void add_conv2d_node(
       weight_data,
       clamp_out);
 
+  utils::uvec3 wg_size = create_conv2d_global_wg_size(graph, method, out);
+
+  if (method == Conv2dMethod::Pointwise || method == Conv2dMethod::Depthwise) {
+    wg_size = {wg_size[0] * wg_size[1] * wg_size[2], 1, 1};
+  }
+
   graph.execute_nodes().emplace_back(new DispatchNode(
       graph,
       shader,
-      create_conv2d_global_wg_size(graph, method, out),
-      graph.create_local_wg_size(out),
+      wg_size,
+      graph.create_local_wg_size(wg_size),
       // Inputs and Outputs
       {{out, vkapi::MemoryAccessType::WRITE},
        {{in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},

docs/TARGETS

@@ -0,0 +1,46 @@
+load("@fbcode_macros//build_defs:native_rules.bzl", "buck_filegroup", "buck_sh_test")
+load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary")
+
+oncall("pytorch_r2p")
+
+python_binary(
+    name = "sphinx",
+    main_module = "sphinx.cmd.build",
+    par_style = "xar",
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir:lib",
+        "//executorch/devtools:lib",
+        "//executorch/exir/backend/test:backend_with_compiler_demo",
+        "//executorch/exir/backend/test:op_partitioner_demo",
+        "//executorch/devtools/bundled_program/serialize:lib",
+        "fbsource//third-party/pypi/ipykernel:ipykernel",
+        "fbsource//third-party/pypi/jupyter-client:jupyter-client",
+        "fbsource//third-party/pypi/jupytext:jupytext",
+        "fbsource//third-party/pypi/nbsphinx:nbsphinx",
+        "fbsource//third-party/pypi/pytorch-sphinx-theme:pytorch-sphinx-theme",
+        "fbsource//third-party/pypi/sphinx:sphinx",
+        "fbsource//third-party/pypi/breathe:breathe",
+        "fbsource//third-party/pypi/sphinx-copybutton:sphinx-copybutton",
+        "fbsource//third-party/pypi/sphinx-design:sphinx-design",
+        "fbsource//third-party/pypi/sphinx-gallery:sphinx-gallery",
+        "fbsource//third-party/pypi/matplotlib:matplotlib",
+        "fbsource//third-party/pypi/myst-parser:myst-parser",
+    ],
+)
+
+buck_filegroup(
+    name = "source",
+    srcs = glob(["source/**/*"]),
+)
+
+buck_sh_test(
+    name = "doctest",
+    args = [
+        "-M",
+        "doctest",
+        "$(location :source)/source/",
+        "/tmp/sphinxbuild",
+    ],
+    test = ":sphinx",
+)

docs/source/conf.py

@@ -44,6 +44,8 @@
 import os
 import sys
 
+FBCODE = "fbcode" in os.getcwd()
+
 extensions = [
     "breathe",
     "sphinx.ext.autodoc",
@@ -60,9 +62,13 @@
     "myst_parser",
     "sphinx_design",
     "sphinx_gallery.gen_gallery",
-    "executorch_custom_versions",
 ]
 
+if not FBCODE:
+    extensions += [
+        "executorch_custom_versions",
+    ]
+
 this_file_dir = os.path.abspath(os.path.dirname(__file__))
 doxygen_xml_dir = os.path.join(
     os.path.dirname(this_file_dir),  # {repo_root}/docs/

examples/qualcomm/oss_scripts/llama3_2/TARGETS

@@ -9,6 +9,9 @@ python_binary(
     name = "llama",
     srcs = ["llama.py"],
     main_function = "executorch.examples.qualcomm.oss_scripts.llama3_2.llama.main",
+    preload_deps = [
+        "//executorch/extension/llm/custom_ops:model_sharding_py",
+    ],
     deps = [
         "//executorch/examples/qualcomm/oss_scripts/llama2:static_llama",
         "//caffe2:torch",