Update on "[ET-VK] Introduce BufferMetadata GLSL struct to abstract tensor layout"

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

[ghstack-poisoned]

Author: ssjia

backends/vulkan/runtime/api/containers/Tensor.cpp

@@ -295,16 +295,17 @@ int32_t create_hashed_layout(
 }
 
 size_t calculate_max_ubo_nbytes(
-    const size_t nbytes_per_ubo,
+    const size_t min_nbytes_per_ubo,
     const utils::StorageType storage_type) {
-  // For texture backed tensors, the metadata fields needed are:
-  // sizes, logical limits
-  size_t max_metadata_field_count = 2u;
+  size_t ivec4_ubo_nbytes = utils::align_up(size_t(16), min_nbytes_per_ubo);
+  size_t uvec3_ubo_nbytes = utils::align_up(size_t(12), min_nbytes_per_ubo);
+  size_t int32_ubo_nbytes = utils::align_up(size_t(4), min_nbytes_per_ubo);
   if (storage_type == utils::kBuffer) {
     // sizes, strides, dim order, numel
-    max_metadata_field_count = 4u;
+    return 3 * ivec4_ubo_nbytes + int32_ubo_nbytes;
   }
-  return max_metadata_field_count * nbytes_per_ubo;
+  // sizes, logical limits
+  return ivec4_ubo_nbytes + uvec3_ubo_nbytes;
 }
 
 //
@@ -835,7 +836,7 @@ const vkapi::BufferBindInfo vTensor::numel_ubo() {
 }
 
 const vkapi::BufferBindInfo vTensor::buffer_meta_ubo() {
-  size_t ubo_nbytes = std::max(sizeof(BufferMetadata), min_nbytes_per_ubo_);
+  size_t ubo_nbytes = sizeof(BufferMetadata);
   if (!buffer_meta_.buffer()) {
     BufferMetadata data(sizes_, dim_order_, strides_, numel_);
     buffer_meta_ = ParamsBuffer(storage_->context_, data);

backends/vulkan/runtime/api/containers/Tensor.h

@@ -548,7 +548,7 @@ class vTensor final {
     if (!uniforms_.buffer()) {
       uniforms_ = ParamsBuffer(storage_->context_, max_ubo_nbytes_, true);
     }
-    size_t ubo_nbytes = std::max(sizeof(T), min_nbytes_per_ubo_);
+    size_t ubo_nbytes = utils::align_up(sizeof(data), min_nbytes_per_ubo_);
     if (*param_buffer_offset == kUniformOffsetUnset) {
       VK_CHECK_COND(
           (uniforms_size_ + ubo_nbytes) <= max_ubo_nbytes_,

examples/models/llava/runner/llava_runner.h

@@ -42,7 +42,7 @@ class ET_EXPERIMENTAL LlavaRunner {
       const float temperature = 0.8f)
       : temperature_(temperature),
         module_(std::make_unique<Module>(model_path, Module::LoadMode::File)),
-        io_manager_(std::make_unique<IOManager>()),
+        io_manager_(std::make_unique<IOManager>(*module_)),
         tokenizer_path_(tokenizer_path) {
     ET_LOG(
         Info,

examples/models/qwen3/README.md

@@ -45,7 +45,7 @@ python -m extension.llm.export.export_llm \
 ### Example run
 With ExecuTorch pybindings:
 ```
-python -m examples.models.llama.runner.native
+python -m examples.models.llama.runner.native \
   --model qwen3_0_6b \
   --pte qwen3_0_6b.pte \
   --tokenizer ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/a9c98e602b9d36d2a2f7ba1eb0f5f31e4e8e5143/tokenizer.json \
@@ -59,9 +59,9 @@ python -m examples.models.llama.runner.native
 
 With ExecuTorch's sample c++ runner (see the Llama README's [Step 3: Run on your computer to validate](../llama/README.md#step-3-run-on-your-computer-to-validate) to build the runner):
 ```
-cmake-out/examples/models/llama/llama_main
-  --model_path qwen3_0_6b.pte
-  --tokenizer_path ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/a9c98e602b9d36d2a2f7ba1eb0f5f31e4e8e5143/tokenizer.json
+cmake-out/examples/models/llama/llama_main \
+  --model_path qwen3_0_6b.pte \
+  --tokenizer_path ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/a9c98e602b9d36d2a2f7ba1eb0f5f31e4e8e5143/tokenizer.json \
   --prompt="Who is the president of the US?"
 ```
 

extension/android/BUCK

@@ -10,6 +10,7 @@ non_fbcode_target(_kind = fb_android_library,
         "executorch_android/src/main/java/org/pytorch/executorch/DType.java",
         "executorch_android/src/main/java/org/pytorch/executorch/EValue.java",
         "executorch_android/src/main/java/org/pytorch/executorch/ExecuTorchRuntime.java",
+        "executorch_android/src/main/java/org/pytorch/executorch/ExecutorchRuntimeException.java",
         "executorch_android/src/main/java/org/pytorch/executorch/MethodMetadata.java",
         "executorch_android/src/main/java/org/pytorch/executorch/Module.java",
         "executorch_android/src/main/java/org/pytorch/executorch/Tensor.java",

extension/android/CMakeLists.txt

@@ -71,6 +71,7 @@ executorch_target_link_options_shared_lib(executorch)
 
 add_library(
   executorch_jni SHARED jni/jni_layer.cpp jni/log.cpp jni/jni_layer_runtime.cpp
+                        jni/jni_helper.cpp
 )
 
 set(link_libraries)

extension/android/executorch_android/src/main/java/org/pytorch/executorch/ExecutorchRuntimeException.java

@@ -0,0 +1,125 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+package org.pytorch.executorch;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+public class ExecutorchRuntimeException extends RuntimeException {
+  // Error code constants - keep in sync with runtime/core/error.h
+  // System errors
+  public static final int OK = 0x00;
+  public static final int INTERNAL = 0x01;
+  public static final int INVALID_STATE = 0x02;
+  public static final int END_OF_METHOD = 0x03;
+
+  // Logical errors
+  public static final int NOT_SUPPORTED = 0x10;
+  public static final int NOT_IMPLEMENTED = 0x11;
+  public static final int INVALID_ARGUMENT = 0x12;
+  public static final int INVALID_TYPE = 0x13;
+  public static final int OPERATOR_MISSING = 0x14;
+  public static final int REGISTRATION_EXCEEDING_MAX_KERNELS = 0x15;
+  public static final int REGISTRATION_ALREADY_REGISTERED = 0x16;
+
+  // Resource errors
+  public static final int NOT_FOUND = 0x20;
+  public static final int MEMORY_ALLOCATION_FAILED = 0x21;
+  public static final int ACCESS_FAILED = 0x22;
+  public static final int INVALID_PROGRAM = 0x23;
+  public static final int INVALID_EXTERNAL_DATA = 0x24;
+  public static final int OUT_OF_RESOURCES = 0x25;
+
+  // Delegate errors
+  public static final int DELEGATE_INVALID_COMPATIBILITY = 0x30;
+  public static final int DELEGATE_MEMORY_ALLOCATION_FAILED = 0x31;
+  public static final int DELEGATE_INVALID_HANDLE = 0x32;
+
+  private static final Map<Integer, String> ERROR_CODE_MESSAGES;
+
+  static {
+    Map<Integer, String> map = new HashMap<>();
+
+    // System errors
+    map.put(OK, "Operation successful");
+    map.put(INTERNAL, "Internal error");
+    map.put(INVALID_STATE, "Invalid state");
+    map.put(END_OF_METHOD, "End of method reached");
+    // Logical errors
+    map.put(NOT_SUPPORTED, "Operation not supported");
+    map.put(NOT_IMPLEMENTED, "Operation not implemented");
+    map.put(INVALID_ARGUMENT, "Invalid argument");
+    map.put(INVALID_TYPE, "Invalid type");
+    map.put(OPERATOR_MISSING, "Operator missing");
+    map.put(REGISTRATION_EXCEEDING_MAX_KERNELS, "Exceeded max kernels");
+    map.put(REGISTRATION_ALREADY_REGISTERED, "Kernel already registered");
+    // Resource errors
+    map.put(NOT_FOUND, "Resource not found");
+    map.put(MEMORY_ALLOCATION_FAILED, "Memory allocation failed");
+    map.put(ACCESS_FAILED, "Access failed");
+    map.put(INVALID_PROGRAM, "Invalid program");
+    map.put(INVALID_EXTERNAL_DATA, "Invalid external data");
+    map.put(OUT_OF_RESOURCES, "Out of resources");
+    // Delegate errors
+    map.put(DELEGATE_INVALID_COMPATIBILITY, "Delegate invalid compatibility");
+    map.put(DELEGATE_MEMORY_ALLOCATION_FAILED, "Delegate memory allocation failed");
+    map.put(DELEGATE_INVALID_HANDLE, "Delegate invalid handle");
+    ERROR_CODE_MESSAGES = Collections.unmodifiableMap(map);
+  }
+
+  static class ErrorHelper {
+    static String formatMessage(int errorCode, String details) {
+      String baseMessage = ERROR_CODE_MESSAGES.get(errorCode);
+      if (baseMessage == null) {
+        baseMessage = "Unknown error code 0x" + Integer.toHexString(errorCode);
+      }
+      return "[Executorch Error 0x"
+          + Integer.toHexString(errorCode)
+          + "] "
+          + baseMessage
+          + ": "
+          + details;
+    }
+  }
+
+  private final int errorCode;
+
+  public ExecutorchRuntimeException(int errorCode, String details) {
+    super(ErrorHelper.formatMessage(errorCode, details));
+    this.errorCode = errorCode;
+  }
+
+  public int getErrorCode() {
+    return errorCode;
+  }
+
+  // Idiomatic Java exception for invalid arguments.
+  public static class ExecutorchInvalidArgumentException extends IllegalArgumentException {
+    private final int errorCode = INVALID_ARGUMENT;
+
+    public ExecutorchInvalidArgumentException(String details) {
+      super(ErrorHelper.formatMessage(INVALID_ARGUMENT, details));
+    }
+
+    public int getErrorCode() {
+      return errorCode;
+    }
+  }
+
+  // Factory method to create an exception of the appropriate subclass.
+  public static RuntimeException makeExecutorchException(int errorCode, String details) {
+    switch (errorCode) {
+      case INVALID_ARGUMENT:
+        return new ExecutorchInvalidArgumentException(details);
+      default:
+        return new ExecutorchRuntimeException(errorCode, details);
+    }
+  }
+}

extension/android/jni/BUCK

@@ -7,6 +7,14 @@ load(":build_defs.bzl", "ET_JNI_COMPILER_FLAGS")
 
 oncall("executorch")
 
+# Define the common JNI source files
+shared_srcs = [
+    "jni_layer.cpp",
+    "jni_layer_runtime.cpp",
+    "jni_helper.cpp",
+    "log.cpp",
+]
+
 non_fbcode_target(_kind = executorch_generated_lib,
     name = "generated_op_lib_optimized",
     custom_ops_aten_kernel_deps = [
@@ -28,7 +36,7 @@ non_fbcode_target(_kind = executorch_generated_lib,
 
 non_fbcode_target(_kind = fb_android_cxx_library,
     name = "executorch_jni",
-    srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp"],
+    srcs = shared_srcs,
     allow_jni_merging = False,
     compiler_flags = ET_JNI_COMPILER_FLAGS,
     soname = "libexecutorch.$(ext)",
@@ -49,7 +57,7 @@ non_fbcode_target(_kind = fb_android_cxx_library,
 
 non_fbcode_target(_kind = fb_android_cxx_library,
     name = "executorch_jni_full",
-    srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp"],
+    srcs = shared_srcs,
     allow_jni_merging = False,
     compiler_flags = ET_JNI_COMPILER_FLAGS,
     soname = "libexecutorch.$(ext)",
@@ -71,7 +79,7 @@ non_fbcode_target(_kind = fb_android_cxx_library,
 
 non_fbcode_target(_kind = fb_android_cxx_library,
     name = "executorch_training_jni",
-    srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp", "jni_layer_training.cpp"],
+    srcs = shared_srcs + ["jni_layer_training.cpp"],
     allow_jni_merging = False,
     compiler_flags = ET_JNI_COMPILER_FLAGS + [
         "-DEXECUTORCH_BUILD_EXTENSION_TRAINING",
@@ -98,11 +106,9 @@ non_fbcode_target(_kind = fb_android_cxx_library,
 
 non_fbcode_target(_kind = fb_android_cxx_library,
     name = "executorch_llama_jni",
-    srcs = [
-        "jni_layer.cpp",
-        "jni_layer_llama.cpp",
-        "jni_layer_runtime.cpp",
-    ],
+    exclude_files = ["log.cpp"]
+    shared_srcs_filtered = [f for f in shared_srcs if f not in exclude_files]
+    srcs = shared_srcs_filtered + ["jni_layer_llama.cpp"]
     allow_jni_merging = False,
     compiler_flags = ET_JNI_COMPILER_FLAGS + [
         "-DEXECUTORCH_BUILD_LLAMA_JNI",
@@ -145,6 +151,10 @@ runtime.export_file(
     name = "jni_layer_runtime.cpp",
 )
 
+runtime.export_file(
+    name = "jni_helper.cpp",
+)
+
 runtime.cxx_library(
     name = "jni_headers",
     exported_headers = [

extension/android/jni/jni_helper.cpp

@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include "jni_helper.h"
+
+namespace executorch::jni_helper {
+
+void throwExecutorchException(uint32_t errorCode, const std::string& details) {
+  // Get the current JNI environment
+  auto env = facebook::jni::Environment::current();
+
+  // Find the Java ExecutorchRuntimeException class
+  static auto exceptionClass = facebook::jni::findClassLocal(
+      "org/pytorch/executorch/ExecutorchRuntimeException");
+
+  // Find the static factory method: makeExecutorchException(int, String)
+  static auto makeExceptionMethod = exceptionClass->getStaticMethod<
+      facebook::jni::local_ref<facebook::jni::JThrowable>(
+          int, facebook::jni::alias_ref<facebook::jni::JString>)>(
+      "makeExecutorchException",
+      "(ILjava/lang/String;)Lorg/pytorch/executorch/ExecutorchRuntimeException;");
+
+  auto jDetails = facebook::jni::make_jstring(details);
+  // Call the factory method to create the exception object
+  auto exception = makeExceptionMethod(exceptionClass, errorCode, jDetails);
+  facebook::jni::throwNewJavaException(exception.get());
+}
+
+} // namespace executorch::jni_helper

extension/android/jni/jni_helper.h

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <fbjni/fbjni.h>
+#include <string>
+
+namespace executorch::jni_helper {
+
+/**
+ * Throws a Java ExecutorchRuntimeException corresponding to the given error
+ * code and details. Uses the Java factory method
+ * ExecutorchRuntimeException.makeExecutorchException(int, String).
+ *
+ * @param errorCode The error code from the C++ Executorch runtime.
+ * @param details Additional details to include in the exception message.
+ */
+void throwExecutorchException(uint32_t errorCode, const std::string& details);
+
+} // namespace executorch::jni_helper