Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 16 additions & 123 deletions packages/react-native-executorch/cpp/core/model.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include <vector>

#include "dtype.h"
#include "model_companion_parser.h"
#include "tensor_helpers.h"

#include <jsi/jsi.h>
Expand All @@ -21,7 +22,6 @@
namespace {
namespace jsi = facebook::jsi;
namespace types = rnexecutorch::core::types;
namespace tensor = rnexecutorch::core::tensor;

template <typename T>
T unwrap(const std::string &ctx, executorch::runtime::Result<T> result) {
Expand Down Expand Up @@ -70,124 +70,6 @@ jsi::Object tensorMetaToJs(jsi::Runtime &rt, const executorch::runtime::TensorIn
return jsTensorMeta;
}

/**
* Parses compile-time dynamic dimension constraints for the inputs of a module
* method.
*
* @note This is a temporary workaround. ExecuTorch (.pte) model metadata
* natively serializes only the static upper-bound limits of dynamic/symbolic
* dimensions, and does not expose the active dynamic range (min, max, step) at
* runtime.
*
* To use this feature, the .pte model must be exported with a companion method
* named "get_dynamic_dims_<methodName>" (e.g., "get_dynamic_dims_forward").
*
* Python export requirements:
* 1. The companion method must take no arguments.
* 2. It must return a list of outputs matching the number of Tag::Tensor inputs
* of the target method (scalar inputs are excluded).
* 3. Each output must be a 2D int32 tensor of shape [rank, 3], where each row
* represents [min, max, step] constraints for the corresponding dimension of
* that input tensor.
*
* @param module The ExecuTorch extension module to query.
* @param methodName The name of the target module method (e.g. "forward").
* @return A vector of SymbolicShape objects, or std::nullopt if the companion
* method is not defined. If returned, the vector's size is guaranteed
* to equal methodMeta.num_inputs(), containing parsed SymbolicShapes
* for tensor inputs and empty SymbolicShapes for non-tensor inputs.
*/
std::optional<std::vector<tensor::SymbolicShape>>
parseDynamicInputShapes(executorch::extension::Module &module, const std::string &methodName) {
using executorch::aten::ScalarType;

const auto getDynamicShapesMethodName = std::format("get_dynamic_dims_{}", methodName);
const auto ctx = getDynamicShapesMethodName + ": ";

auto methodNames = unwrap(ctx + "failed to get method names", module.method_names());
if (methodName == getDynamicShapesMethodName ||
!methodNames.contains(getDynamicShapesMethodName)) {
return std::nullopt;
}

auto methodMeta = unwrap(std::format("{}failed to get meta for method '{}'", ctx, methodName),
module.method_meta(methodName));

size_t expectedTensorInputs = 0;
for (size_t i = 0; i < methodMeta.num_inputs(); ++i) {
auto tag = unwrap(std::format("{}failed to get tag for input [{}]", ctx, i), methodMeta.input_tag(i));
if (tag == executorch::runtime::Tag::Tensor) {
expectedTensorInputs++;
}
}

auto result = unwrap(ctx + "failed to execute", module.execute(getDynamicShapesMethodName));
if (result.size() != expectedTensorInputs) {
throw std::runtime_error(std::format("{}number of outputs returned ({}) does not match the number of "
"tensor inputs declared by method '{}' ({})",
ctx, result.size(), methodName, expectedTensorInputs));
}

std::vector<tensor::SymbolicShape> dynamicShapes;
dynamicShapes.reserve(methodMeta.num_inputs());
size_t tensorIndex = 0;

for (size_t i = 0; i < methodMeta.num_inputs(); ++i) {
auto tag = unwrap(std::format("{}failed to get tag for input [{}]", ctx, i),
methodMeta.input_tag(i));

if (tag != executorch::runtime::Tag::Tensor) {
// Emplace an empty shape for non-tensor inputs to maintain a 1-to-1
// index alignment between dynamicShapes and the model's inputs vector.
dynamicShapes.emplace_back();
continue;
}

const auto &out = result.at(tensorIndex);
if (!out.isTensor()) {
throw std::runtime_error(std::format("{}output[{}] is not a tensor", ctx, tensorIndex));
}

auto inputMeta = unwrap(std::format("{}failed to get tensor meta for input [{}]", ctx, i),
methodMeta.input_tensor_meta(i));
const auto rank = inputMeta.sizes().size();
const auto shapeTensor = out.toTensor();

if (shapeTensor.dim() != 2 || shapeTensor.size(1) != 3 ||
shapeTensor.size(0) != static_cast<ssize_t>(rank) ||
shapeTensor.scalar_type() != ScalarType::Int) {
throw std::runtime_error(std::format("{}output[{}] expected to be a 2D int32_t tensor of shape [{}, 3]",
ctx, tensorIndex, rank));
}

const auto *shape = shapeTensor.const_data_ptr<int32_t>();
tensor::SymbolicShape symbolicShape;
symbolicShape.reserve(rank);

for (size_t axis = 0; axis < rank; ++axis) {
const auto minDim = shape[axis * 3 + 0];
const auto maxDim = shape[axis * 3 + 1];
const auto step = shape[axis * 3 + 2];
if (minDim < 0 || maxDim < minDim || step < 1) {
throw std::runtime_error(std::format("{}output[{}], axis {} is invalid: "
"expected 0 <= min <= max and step >= 1 but got [{}, {}, {}]",
ctx, tensorIndex, axis, minDim, maxDim, step));
}
if (maxDim > inputMeta.sizes()[axis]) {
throw std::runtime_error(std::format("{}output[{}], axis {} max dimension ({}) "
"exceeds model metadata upper limit ({})",
ctx, tensorIndex, axis, maxDim, inputMeta.sizes()[axis]));
}

symbolicShape.emplace_back(tensor::RangeDim{.min = minDim, .max = maxDim, .step = step});
}

dynamicShapes.push_back(std::move(symbolicShape));
++tensorIndex;
}

return dynamicShapes;
}
} // namespace

namespace rnexecutorch::core::model {
Expand All @@ -199,6 +81,7 @@ using rnexecutorch::core::tensor::TensorHostObject;
ModelHostObject::ModelHostObject(const std::string &modelPath)
: modelPath_(modelPath),
etModule_(std::make_unique<executorch::extension::Module>(modelPath)) {

auto error = etModule_->load();
if (!etModule_->is_loaded()) {
const std::string errorMsg = executorch::runtime::to_string(error);
Expand All @@ -208,8 +91,18 @@ ModelHostObject::ModelHostObject(const std::string &modelPath)
auto methodNames = unwrap("Failed to get method names", etModule_->method_names());
for (const auto &methodName : methodNames) {
auto dynamicShapes = parseDynamicInputShapes(*etModule_, methodName);
auto enumeratedShapes = parseEnumeratedInputShapes(*etModule_, methodName);

if (dynamicShapes && enumeratedShapes) {
throw std::runtime_error(std::format(
"Method '{}' cannot define both get_dynamic_dims_{} and get_enumerated_dims_{}",
methodName, methodName, methodName));
}
if (dynamicShapes) {
dynamicInputShapes_.emplace(methodName, std::move(*dynamicShapes));
inputShapeConstraints_.emplace(methodName, std::move(*dynamicShapes));
}
if (enumeratedShapes) {
inputShapeConstraints_.emplace(methodName, std::move(*enumeratedShapes));
}
}
}
Expand Down Expand Up @@ -364,9 +257,9 @@ jsi::Value ModelHostObject::get(jsi::Runtime &rt, const jsi::PropNameID &name) {
auto expectedDtype = fromScalarType(rt, ctx, tensorMeta.scalar_type());

std::shared_ptr<TensorHostObject> tensorHostObject;
if (self->dynamicInputShapes_.contains(methodName)) {
auto expectedShape = self->dynamicInputShapes_[methodName][i];
tensorHostObject = tensor::fromJs(rt, ctx, val, expectedDtype, expectedShape);
if (self->inputShapeConstraints_.contains(methodName)) {
auto shapeConstraint = self->inputShapeConstraints_[methodName][i];
tensorHostObject = tensor::fromJs(rt, ctx, val, expectedDtype, shapeConstraint);
} else {
tensorHostObject = tensor::fromJs(rt, ctx, val, expectedDtype, tensorMeta.sizes());
}
Expand Down
2 changes: 1 addition & 1 deletion packages/react-native-executorch/cpp/core/model.h
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ class ModelHostObject : public jsi::HostObject,
std::unique_ptr<executorch::extension::Module> etModule_;
std::mutex mutex_;

std::unordered_map<std::string, std::vector<core::tensor::SymbolicShape>> dynamicInputShapes_;
std::unordered_map<std::string, std::vector<core::tensor::ShapeConstraint>> inputShapeConstraints_;
};

void install_loadModel(jsi::Runtime &rt, jsi::Object &module);
Expand Down
207 changes: 207 additions & 0 deletions packages/react-native-executorch/cpp/core/model_companion_parser.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,207 @@
#include "model_companion_parser.h"

#include <cstddef>
#include <cstdint>
#include <format>
#include <optional>
#include <stdexcept>
#include <string>
#include <vector>

#include <executorch/runtime/core/tag.h>

namespace {

template <typename T>
T unwrap(const std::string &ctx, executorch::runtime::Result<T> result) {
if (!result.ok()) {
throw std::runtime_error(
std::format("{}: {}", ctx, executorch::runtime::to_string(result.error())));
}
return std::move(result.get());
}

} // namespace

namespace rnexecutorch::core::model {

std::optional<std::vector<tensor::ShapeConstraint>>
parseDynamicInputShapes(executorch::extension::Module &module, const std::string &methodName) {
using executorch::aten::ScalarType;

const auto getDynamicShapesMethodName = std::format("get_dynamic_dims_{}", methodName);
const auto ctx = getDynamicShapesMethodName + ": ";

auto methodNames = unwrap(ctx + "failed to get method names", module.method_names());
if (methodName == getDynamicShapesMethodName ||
!methodNames.contains(getDynamicShapesMethodName)) {
return std::nullopt;
}

auto methodMeta = unwrap(std::format("{}failed to get meta for method '{}'", ctx, methodName),
module.method_meta(methodName));

size_t expectedTensorInputs = 0;
for (size_t i = 0; i < methodMeta.num_inputs(); ++i) {
auto tag = unwrap(std::format("{}failed to get tag for input [{}]", ctx, i), methodMeta.input_tag(i));
if (tag == executorch::runtime::Tag::Tensor) {
expectedTensorInputs++;
}
}

auto result = unwrap(ctx + "failed to execute", module.execute(getDynamicShapesMethodName));
if (result.size() != expectedTensorInputs) {
throw std::runtime_error(std::format("{}number of outputs returned ({}) does not match the number of "
"tensor inputs declared by method '{}' ({})",
ctx, result.size(), methodName, expectedTensorInputs));
}

std::vector<tensor::ShapeConstraint> dynamicShapes;
dynamicShapes.reserve(methodMeta.num_inputs());
size_t tensorIndex = 0;

for (size_t i = 0; i < methodMeta.num_inputs(); ++i) {
auto tag = unwrap(std::format("{}failed to get tag for input [{}]", ctx, i),
methodMeta.input_tag(i));

if (tag != executorch::runtime::Tag::Tensor) {
dynamicShapes.emplace_back();
continue;
}

const auto &out = result.at(tensorIndex);
if (!out.isTensor()) {
throw std::runtime_error(std::format("{}output[{}] is not a tensor", ctx, tensorIndex));
}

auto inputMeta = unwrap(std::format("{}failed to get tensor meta for input [{}]", ctx, i),
methodMeta.input_tensor_meta(i));
const auto rank = inputMeta.sizes().size();
const auto shapeTensor = out.toTensor();

if (shapeTensor.dim() != 2 || shapeTensor.size(1) != 3 ||
shapeTensor.size(0) != static_cast<ssize_t>(rank) ||
shapeTensor.scalar_type() != ScalarType::Int) {
throw std::runtime_error(std::format("{}output[{}] expected to be a 2D int32_t tensor of shape [{}, 3]",
ctx, tensorIndex, rank));
}

const auto *shape = shapeTensor.const_data_ptr<int32_t>();
tensor::SymbolicShape symbolicShape;
symbolicShape.reserve(rank);

for (size_t axis = 0; axis < rank; ++axis) {
const auto minDim = shape[axis * 3 + 0];
const auto maxDim = shape[axis * 3 + 1];
const auto step = shape[axis * 3 + 2];
if (minDim < 0 || maxDim < minDim || step < 1) {
throw std::runtime_error(std::format("{}output[{}], axis {} is invalid: "
"expected 0 <= min <= max and step >= 1 but got [{}, {}, {}]",
ctx, tensorIndex, axis, minDim, maxDim, step));
}
if (maxDim > inputMeta.sizes()[axis]) {
throw std::runtime_error(std::format("{}output[{}], axis {} max dimension ({}) "
"exceeds model metadata upper limit ({})",
ctx, tensorIndex, axis, maxDim, inputMeta.sizes()[axis]));
}
Comment on lines +102 to +106

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about the analogical check for minDim?


symbolicShape.emplace_back(tensor::RangeDim{.min = minDim, .max = maxDim, .step = step});
}

dynamicShapes.emplace_back(std::move(symbolicShape));
++tensorIndex;
}

return dynamicShapes;
}

std::optional<std::vector<tensor::ShapeConstraint>>
parseEnumeratedInputShapes(executorch::extension::Module &module, const std::string &methodName) {
using executorch::aten::ScalarType;

const auto getEnumeratedShapesMethodName = std::format("get_enumerated_dims_{}", methodName);
const auto ctx = getEnumeratedShapesMethodName + ": ";

auto methodNames = unwrap(ctx + "failed to get method names", module.method_names());
if (methodName == getEnumeratedShapesMethodName ||
!methodNames.contains(getEnumeratedShapesMethodName)) {
return std::nullopt;
}

auto methodMeta = unwrap(std::format("{}failed to get meta for method '{}'", ctx, methodName),
module.method_meta(methodName));

size_t expectedTensorInputs = 0;
for (size_t i = 0; i < methodMeta.num_inputs(); ++i) {
auto tag = unwrap(std::format("{}failed to get tag for input [{}]", ctx, i),
methodMeta.input_tag(i));
if (tag == executorch::runtime::Tag::Tensor) {
expectedTensorInputs++;
}
}

auto result = unwrap(ctx + "failed to execute", module.execute(getEnumeratedShapesMethodName));
if (result.size() != expectedTensorInputs) {
throw std::runtime_error(std::format("{}number of outputs returned ({}) does not match the number of "
"tensor inputs declared by method '{}' ({})",
ctx, result.size(), methodName, expectedTensorInputs));
}
Comment on lines +120 to +148

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we should abstract common part of these two into function in anonymous helper?


std::vector<tensor::ShapeConstraint> enumeratedShapes;
enumeratedShapes.reserve(methodMeta.num_inputs());
size_t tensorIndex = 0;

for (size_t i = 0; i < methodMeta.num_inputs(); ++i) {
auto tag = unwrap(std::format("{}failed to get tag for input [{}]", ctx, i),
methodMeta.input_tag(i));

if (tag != executorch::runtime::Tag::Tensor) {
enumeratedShapes.emplace_back();
continue;
}

const auto &out = result.at(tensorIndex);
if (!out.isTensor()) {
throw std::runtime_error(std::format("{}output[{}] is not a tensor", ctx, tensorIndex));
}

auto inputMeta = unwrap(std::format("{}failed to get tensor meta for input [{}]", ctx, i),
methodMeta.input_tensor_meta(i));
const auto rank = inputMeta.sizes().size();
const auto shapeTensor = out.toTensor();

if (shapeTensor.dim() != 2 ||
shapeTensor.size(1) != static_cast<ssize_t>(rank) ||
shapeTensor.scalar_type() != ScalarType::Int) {
throw std::runtime_error(std::format("{}output[{}] expected to be a 2D int32_t tensor of shape [num_shapes, {}]",
ctx, tensorIndex, rank));
}

const auto numShapes = static_cast<size_t>(shapeTensor.size(0));
const auto *shapeData = shapeTensor.const_data_ptr<int32_t>();
std::vector<tensor::SymbolicShape> alternatives;
alternatives.reserve(numShapes);

for (size_t s = 0; s < numShapes; ++s) {
tensor::SymbolicShape singleShape;
singleShape.reserve(rank);
for (size_t axis = 0; axis < rank; ++axis) {
const auto dimVal = shapeData[s * rank + axis];
if (dimVal < 0) {
throw std::runtime_error(std::format("{}output[{}], shape {}, axis {} is invalid: "
"dimension cannot be negative ({})",
ctx, tensorIndex, s, axis, dimVal));
}
singleShape.emplace_back(dimVal);
}
alternatives.push_back(std::move(singleShape));
}

enumeratedShapes.emplace_back(std::move(alternatives));
++tensorIndex;
}

return enumeratedShapes;
}

} // namespace rnexecutorch::core::model
Loading