Add Native Extension

Skill: Add a Native C++ Extension & JSI Bindings

Use this guide to add custom, performance-critical native operations in C++ and expose them to TypeScript via React Native JSI.

🚦 Architectural Guidelines

Before writing any C++ code, ensure you adhere to the following principles:

1. Amdahl's Law & Premature Optimization:

   • Evaluate what percentage of total inference/pipeline time the processing step occupies. If the preprocessing/postprocessing step takes < 5% of the total inference budget, write it in pure TypeScript to reduce codebase complexity and maintenance overhead.

2. Destination Tensors & Local Memory:

   • Local Memory is Allowed: You can allocate temporary native C++ memory (such as stack variables, std::vectors, or dynamic memory cleaned up before the function exits) for intermediate calculations.
   • Destination Tensors: If the operation writes dense output, the destination tensor must be pre-allocated by the caller (in TypeScript) and passed as an argument (e.g., sigmoid(src, dst)).
   • Primitive Array Returns: If the operation produces variable-sized non-dense outputs (like bounding box indices in Non-Maximum Suppression (NMS)), return a plain jsi::Array of primitives (like indices or coordinates).
     • Example: nms(boxes, scores, options) returns a jsi::Array of indices (e.g., [0, 4, 12]) rather than a new tensor. This avoids all native memory management overhead for variable-sized outputs.

🚫 Avoid / Anti-Patterns

• Do NOT return implicitly allocated JSI Tensors: Never return newly created TensorHostObject instances from C++. This forces the JavaScript layer to reason about their garbage collection and manual lifetimes, leading to native memory leaks.
• Do NOT define default parameters in C++: Native C++ functions must never define default argument values (e.g. axis = -1). Define all default values explicitly in the TypeScript wrapper layer instead.
• Do NOT perform in-place mutation without safety checks: Never allow inputs and outputs to share the same underlying instance.
• Do NOT throw raw jsi::JSError, std::runtime_error, or std::invalid_argument: They reach JavaScript with no error code. Throw error::X(message) and register through error::guarded(...). See the Error Handling Skill.

🛠️ Step-by-Step Implementation

Step 1: Create the Native Operation Files

Under cpp/extensions/<domain>/, create or modify the header and implementation files for your operations:

1. Header (cpp/extensions/<domain>/operations.h)

Keep the header clean and specify exact JSI install functions:

#pragma once
#include <jsi/jsi.h>

namespace rnexecutorch::extensions::<domain>
{
    void install_customOp(facebook::jsi::Runtime &rt, facebook::jsi::Object &module);
}

2. Source (cpp/extensions/<domain>/operations.cpp)

• Validate, type-check, and extract input/output tensors using tensor::fromJs, specifying expected DType and shape constraints where possible.
• Convert primitive parameters using conversions::asType<T>.
• Prevent in-place mutations (aliasing) using tensor::checkNotSameTensor.
• Lock tensors for thread-safe access using tensor::tryLockShared (for inputs) and tensor::tryLockUnique (for outputs), which also ensure the underlying memory buffer has not been disposed.
• Raise every failure through a code factory, error::X(message), and wrap the registration in error::guarded(...), so the error reaches JavaScript with a code. See the Error Handling Skill.

#include "operations.h"
#include "core/conversions.h"
#include "core/error.h"
#include "core/tensor_helpers.h"
#include <algorithm>

namespace rnexecutorch::extensions::<domain>
{
    namespace jsi = facebook::jsi;
    namespace conversions = rnexecutorch::core::conversions;
    namespace tensor = rnexecutorch::core::tensor;
    // Required under extensions::*; OMIT inside rnexecutorch::core::*, where
    // unqualified `error` already resolves to the sibling namespace.
    namespace error = rnexecutorch::core::error;
    using rnexecutorch::core::types::DType;

    void install_customOp(jsi::Runtime &rt, jsi::Object &module)
    {
        auto name = "customOp";
        auto fnBody = [](jsi::Runtime &rt, const jsi::Value &thisVal, const jsi::Value *args, size_t count) -> jsi::Value
        {
            // 1. Strict argument count validation (No default values here!)
            if (count != 3)
            {
                throw error::InvalidArgument("customOp: Usage: customOp(src, dst, factor)");
            }

            // 2. Validate, extract input/output tensors and check DType/Shape constraints using fromJs
            auto src = tensor::fromJs(rt, "customOp: src", args[0], DType::float32, std::nullopt);
            auto dst = tensor::fromJs(rt, "customOp: dst", args[1], DType::float32, src->shape_);

            // Extract and convert arguments using conversions::asType
            double factor = conversions::asType<double>(rt, "customOp: factor", args[2]);

            // 3. Prevent in-place mutations (aliasing)
            tensor::checkNotSameTensor(rt, "customOp: src", src, "customOp: dst", dst);

            // 4. Lock underlying buffers and verify they aren't disposed
            auto srcLock = tensor::tryLockShared(rt, "customOp: src", src);
            auto dstLock = tensor::tryLockUnique(rt, "customOp: dst", dst);

            // 5. Perform the computation
            const float *srcData = reinterpret_cast<const float *>(src->data_.get());
            float *dstData = reinterpret_cast<float *>(dst->data_.get());
            size_t size = src->numel_;

            for (size_t i = 0; i < size; ++i)
            {
                dstData[i] = srcData[i] * static_cast<float>(factor);
            }

            // Always return the destination tensor (args[1]) as the JSI result
            return jsi::Value(rt, args[1]);
        };

        // error::guarded turns any RnExecuTorchException raised in the native stack into a
        // JS Error carrying `code`. Never register a bare fnBody.
        module.setProperty(rt, name, jsi::Function::createFromHostFunction(rt, jsi::PropNameID::forAscii(rt, name), 3, error::guarded(fnBody)));
    }
}

Step 2: Register in Extension and Core JSI Installs

1. Extension Register (cpp/extensions/<domain>/install.cpp):

   #include "install.h"
   #include "operations.h"
   
   namespace rnexecutorch::extensions::<domain>
   {
       void install(facebook::jsi::Runtime &rt, facebook::jsi::Object &module)
       {
           facebook::jsi::Object subModule(rt);
           install_customOp(rt, subModule);
           module.setProperty(rt, "<domain>", subModule);
       }
   }
   

2. Core Register (cpp/RnExecutorch.cpp):

   #include "extensions/<domain>/install.h"
   // ... inside rnexecutorch::install ...
   rnexecutorch::extensions::<domain>::install(jsiRuntime, module);
   

Step 3: TypeScript Bridge & Wrappers

Where the wrapper file goes — general vs. model-specific ops:

• Domain-general ops (reusable across every task in the domain, e.g. image_ops, box_ops) go in the domain's shared op files (src/extensions/<domain>/ops.ts), re-exported from src/extensions/<domain>/index.ts.
• Model- or task-specific ops (only meaningful to one model/pipeline, e.g. FSMN-VAD framing) must go under src/extensions/<domain>/utils/<name>.ts — one file per model/task, named for it (e.g. vadUtils.ts, supertonicUtils.ts). Keep them out of the shared ops.ts so it stays a home for genuinely reusable ops, and group each model's helpers together under utils/. Re-export them from the domain index.ts too so power users can reach them.

Then, in that wrapper file:

• Use the rnexecutorchJsi Symbol: You must import and interact with native bindings using the rnexecutorchJsi symbol exported from src/native/bridge.ts. Do not reference the global __rnexecutorch_jsi__ directly throughout your wrapper files.
• Expose the TypeScript wrapper.
• Handle default values here instead of the C++ layer.
• Mark wrapper functions with the "worklet"; directive.

import { rnexecutorchJsi } from '../native/bridge';
import { type Tensor } from '../core/tensor';

/**
 * Applies a custom operation scaling the src tensor by factor.
 * @param src Input Tensor.
 * @param dst Pre-allocated Destination Tensor.
 * @param factor Scale factor. Defaults to 1.0.
 */
export function customOp(src: Tensor, dst: Tensor, factor: number = 1.0): Tensor {
  'worklet';
  return rnexecutorchJsi.<domain>.customOp(src, dst, factor);
}

📋 Verification Checklist

When adding a native extension, verify that:

• You only implemented in C++ if the operation takes > 5% of the total inference budget.
• No JSI Tensors are implicitly allocated and returned in the C++ code.
• Input and output tensors are extracted and validated using tensor::fromJs with strict expected DType and shape constraints where possible.
• Primitive arguments are extracted and converted using conversions::asType<T>.
• In-place mutation is explicitly prevented using tensor::checkNotSameTensor.
• Input and output tensors are locked using tensor::tryLockShared and tensor::tryLockUnique respectively.
• No default parameter values are defined in the C++ header/source files.
• Every failure is raised through a factory, error::X(message); no raw jsi::JSError / std::runtime_error / std::invalid_argument was introduced.
• The host function is registered through error::guarded(fnBody), not a bare fnBody.
• #include "core/error.h" and the namespace error = ... alias sit at file scope, outside any #if defined(__ANDROID__) / #elif defined(__APPLE__) branch.
• The custom operation install function is registered in both the domain install function and core cpp/RnExecutorch.cpp.
• The TypeScript wrapper lives in the right place: domain-general ops in the shared ops.ts, model-/task-specific ops under src/extensions/<domain>/utils/<name>.ts.
• The TypeScript wrapper imports and uses rnexecutorchJsi instead of the global __rnexecutorch_jsi__.
• The TypeScript wrapper is marked with the "worklet"; directive and defines all default parameter values.