Update on "[ET-VK][ez] Fix partitioner logic of finding keepdim arg of reduce ops"

Title says it all. For reduce ops, their signature are not all alike so some extra legwork needs to be done to identify specific arguments that need to be checked.

Also included a small update to partitioner logging to improve debuggability.

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

[ghstack-poisoned]

Author: ssjia

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