Skip to content

[RNE Rewrite] Use std::span instead of raw pointers in native C++ ops + tokenizer cosmetic refactor #1315

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
msluszniak wants to merge 1 commit into
base: rne-rewrite
Choose a base branch
from
Open

[RNE Rewrite] Use std::span instead of raw pointers in native C++ ops + tokenizer cosmetic refactor #1315

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
49 changes: 23 additions & 26 deletions packages/react-native-executorch/cpp/extensions/cv/box_ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,12 @@
#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <format>
#include <numeric>
#include <optional>
#include <shared_mutex>
#include <span>
#include <stdexcept>
#include <utility>
#include <vector>
Expand Down Expand Up @@ -60,9 +62,13 @@ NmsType parseNmsType(const std::string &s) {
throw std::invalid_argument("unsupported nmsType '" + s + "'");
}

std::array<float, 4> decodeToXyxy(
float a, float b, float c, float d,
BoxFormat format) {
constexpr size_t kBoxCoords = 4;

std::array<float, kBoxCoords> decodeToXyxy(std::span<const float, kBoxCoords> box, BoxFormat format) {
const float a = box[0];
const float b = box[1];
const float c = box[2];
const float d = box[3];
switch (format) {
case BoxFormat::XYXY:
return {a, b, c, d};
Expand Down Expand Up @@ -104,14 +110,20 @@ void install_nms(jsi::Runtime &rt, jsi::Object &module) {
}

std::int32_t numAnchors = scores->shape_[0];
const auto *boxesPtr = reinterpret_cast<const float *>(boxes->data_.get());
const auto *scoresPtr = reinterpret_cast<const float *>(scores->data_.get());
const std::span<const float> boxesData(reinterpret_cast<const float *>(boxes->data_.get()), boxes->numel_);
const std::span<const float> scoresData(reinterpret_cast<const float *>(scores->data_.get()), scores->numel_);

// Coordinates of anchor `i` occupy the i-th 4-element row of `boxesData`.
auto boxAt = [boxesData](std::int32_t index) -> std::span<const float, kBoxCoords> {
return std::span<const float, kBoxCoords>(
boxesData.subspan(static_cast<size_t>(index) * kBoxCoords, kBoxCoords));
};

std::vector<std::pair<std::int32_t, float>> candidates;
candidates.reserve(static_cast<size_t>(numAnchors));

for (size_t idx = 0; std::cmp_less(idx, numAnchors); ++idx) {
float score = scoresPtr[idx];
for (std::int32_t idx = 0; idx < numAnchors; ++idx) {
const float score = scoresData[static_cast<size_t>(idx)];

if (score >= confidenceThreshold) {
candidates.emplace_back(idx, score);
Expand All @@ -134,12 +146,7 @@ void install_nms(jsi::Runtime &rt, jsi::Object &module) {

std::int32_t idxI = candidates[i].first;

auto [xminA, yminA, xmaxA, ymaxA] = decodeToXyxy(
boxesPtr[idxI * 4 + 0],
boxesPtr[idxI * 4 + 1],
boxesPtr[idxI * 4 + 2],
boxesPtr[idxI * 4 + 3],
boxFormat);
auto [xminA, yminA, xmaxA, ymaxA] = decodeToXyxy(boxAt(idxI), boxFormat);

const float areaA = (xmaxA - xminA) * (ymaxA - yminA);

Expand All @@ -152,12 +159,7 @@ void install_nms(jsi::Runtime &rt, jsi::Object &module) {

std::int32_t idxJ = candidates[j].first;

auto [xminB, yminB, xmaxB, ymaxB] = decodeToXyxy(
boxesPtr[idxJ * 4 + 0],
boxesPtr[idxJ * 4 + 1],
boxesPtr[idxJ * 4 + 2],
boxesPtr[idxJ * 4 + 3],
boxFormat);
auto [xminB, yminB, xmaxB, ymaxB] = decodeToXyxy(boxAt(idxJ), boxFormat);

const float areaB = (xmaxB - xminB) * (ymaxB - yminB);

Expand Down Expand Up @@ -220,7 +222,7 @@ void install_restrictToBox(jsi::Runtime &rt, jsi::Object &module) {
auto dstLock = tensor::tryLockUnique(rt, "restrictToBox: dst", dst);

auto boxVec = conversions::asVector<float>(rt, "restrictToBox: boxTuple", args[2]);
if (boxVec.size() != 4) {
if (boxVec.size() != kBoxCoords) {
throw jsi::JSError(rt, "restrictToBox: boxTuple must contain exactly 4 coordinates");
}

Expand All @@ -232,12 +234,7 @@ void install_restrictToBox(jsi::Runtime &rt, jsi::Object &module) {
throw jsi::JSError(rt, std::format("restrictToBox: {}", e.what()));
}

float a = boxVec[0];
float b = boxVec[1];
float c = boxVec[2];
float d = boxVec[3];

auto [xmin, ymin, xmax, ymax] = decodeToXyxy(a, b, c, d, boxFormat);
auto [xmin, ymin, xmax, ymax] = decodeToXyxy(std::span<const float, kBoxCoords>(boxVec), boxFormat);

int32_t H = src->shape_[0];
int32_t W = src->shape_[1];
Expand Down
41 changes: 24 additions & 17 deletions packages/react-native-executorch/cpp/extensions/cv/image_ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@

#include <algorithm>
#include <cmath>
#include <cstring>
#include <cstddef>
#include <cstdint>
#include <format>
#include <optional>
#include <span>
#include <stdexcept>
#include <utility>

Expand Down Expand Up @@ -261,10 +263,12 @@ void install_toChannelsFirst(jsi::Runtime &rt, jsi::Object &module) {

const size_t hw = static_cast<size_t>(srcH) * static_cast<size_t>(srcW);
const size_t elemSize = rnexecutorch::core::types::elementSize(src->dtype_);
uint8_t *dstPtr = dst->data_.get();
const size_t planeBytes = hw * elemSize;
const std::span<uint8_t> dstBytes(dst->data_.get(), dst->size_);

for (size_t i = 0; std::cmp_less(i, srcC); ++i) {
std::memcpy(dstPtr + i * hw * elemSize, channels[i].data, hw * elemSize);
const std::span<const uint8_t> plane(channels[i].data, planeBytes);
std::ranges::copy(plane, dstBytes.subspan(i * planeBytes, planeBytes).begin());
}
} catch (const std::exception &e) {
throw jsi::JSError(rt, "toChannelsFirst: " + std::string(e.what()));
Expand Down Expand Up @@ -299,11 +303,12 @@ void install_toChannelsLast(jsi::Runtime &rt, jsi::Object &module) {

const size_t hw = static_cast<size_t>(srcH) * static_cast<size_t>(srcW);
const size_t elemSize = rnexecutorch::core::types::elementSize(src->dtype_);
uint8_t *srcPtr = src->data_.get();
const size_t planeBytes = hw * elemSize;
const std::span<uint8_t> srcBytes(src->data_.get(), src->size_);

std::vector<::cv::Mat> channels;
for (size_t i = 0; std::cmp_less(i, srcC); ++i) {
channels.emplace_back(srcH, srcW, cvDepth, srcPtr + i * hw * elemSize);
channels.emplace_back(srcH, srcW, cvDepth, srcBytes.subspan(i * planeBytes, planeBytes).data());
}

::cv::Mat dstMat(srcH, srcW, CV_MAKETYPE(cvDepth, srcC), dst->data_.get());
Expand Down Expand Up @@ -362,13 +367,18 @@ void install_normalize(jsi::Runtime &rt, jsi::Object &module) {

const size_t srcElemSize = rnexecutorch::core::types::elementSize(src->dtype_);
const size_t dstElemSize = rnexecutorch::core::types::elementSize(dst->dtype_);
uint8_t *srcPtr = src->data_.get();
uint8_t *dstPtr = dst->data_.get();
const std::span<uint8_t> srcBytes(src->data_.get(), src->size_);
const std::span<uint8_t> dstBytes(dst->data_.get(), dst->size_);

const size_t plane = static_cast<size_t>(h) * static_cast<size_t>(w);
const size_t srcPlaneBytes = plane * srcElemSize;
const size_t dstPlaneBytes = plane * dstElemSize;

for (size_t ch = 0; std::cmp_less(ch, c); ++ch) {
const ::cv::Mat srcChannel(h, w, srcDepthType, srcPtr + ch * plane * srcElemSize);
::cv::Mat dstChannel(h, w, dstDepthType, dstPtr + ch * plane * dstElemSize);
const ::cv::Mat srcChannel(h, w, srcDepthType,
srcBytes.subspan(ch * srcPlaneBytes, srcPlaneBytes).data());
::cv::Mat dstChannel(h, w, dstDepthType,
dstBytes.subspan(ch * dstPlaneBytes, dstPlaneBytes).data());

srcChannel.convertTo(dstChannel, dstDepthType, alpha[ch], beta[ch]);
}
Expand Down Expand Up @@ -411,21 +421,18 @@ void install_applyColormap(jsi::Runtime &rt, jsi::Object &module) {
}
}

const size_t pixels = src->numel_;

const auto *srcData = reinterpret_cast<const int32_t *>(src->data_.get());
uint8_t *dstData = dst->data_.get();
const std::span<const int32_t> srcData(reinterpret_cast<const int32_t *>(src->data_.get()), src->numel_);
const std::span<uint8_t> dstData(dst->data_.get(), dst->numel_);

for (size_t i = 0; i < pixels; ++i) {
for (size_t i = 0; i < srcData.size(); ++i) {
const int32_t idx = srcData[i];
if (idx < 0 || std::cmp_greater_equal(idx, numColors)) {
throw jsi::JSError(rt, "applyColormap: tensor contains class index (" +
std::to_string(idx) + ") that exceeds provided colormap size (" +
std::to_string(numColors) + ")");
}
for (size_t c = 0; c < numRgbaChannels; ++c) {
dstData[i * numRgbaChannels + c] = lut[static_cast<size_t>(idx)][c];
}
std::ranges::copy(lut[static_cast<size_t>(idx)],
dstData.subspan(i * numRgbaChannels, numRgbaChannels).begin());
}

return jsi::Value(rt, args[1]);
Expand Down
54 changes: 34 additions & 20 deletions packages/react-native-executorch/cpp/extensions/math/operations.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,11 @@

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <optional>
#include <span>
#include <utility>

#include "core/tensor.h"
Expand All @@ -30,13 +33,12 @@ void install_sigmoid(jsi::Runtime &rt, jsi::Object &module) {
auto srcLock = tensor::tryLockShared(rt, "sigmoid: src", src);
auto dstLock = tensor::tryLockUnique(rt, "sigmoid: dst", dst);

const auto countElements = src->numel_;
const auto *srcData = reinterpret_cast<const float *>(src->data_.get());
auto *dstData = reinterpret_cast<float *>(dst->data_.get());
const std::span<const float> srcData(reinterpret_cast<const float *>(src->data_.get()), src->numel_);
const std::span<float> dstData(reinterpret_cast<float *>(dst->data_.get()), dst->numel_);

for (size_t i = 0; i < countElements; ++i) {
dstData[i] = 1.0f / (1.0f + std::exp(-srcData[i]));
}
std::ranges::transform(srcData, dstData.begin(), [](const float value) {
return 1.0f / (1.0f + std::exp(-value));
});

return jsi::Value(rt, args[1]);
};
Expand Down Expand Up @@ -75,8 +77,8 @@ void install_softmax(jsi::Runtime &rt, jsi::Object &module) {
}
const auto axisIdx = static_cast<size_t>(axis);

const auto *srcData = reinterpret_cast<const float *>(src->data_.get());
auto *dstData = reinterpret_cast<float *>(dst->data_.get());
const std::span<const float> srcData(reinterpret_cast<const float *>(src->data_.get()), src->numel_);
const std::span<float> dstData(reinterpret_cast<float *>(dst->data_.get()), dst->numel_);

const auto axisDim = static_cast<size_t>(src->shape_[axisIdx]);
if (axisDim == 0) {
Expand All @@ -93,24 +95,30 @@ void install_softmax(jsi::Runtime &rt, jsi::Object &module) {
inner *= static_cast<size_t>(src->shape_[i]);
}

// Elements along `axis` are strided by `inner`, so a lane spans from its
// first element to its last: (axisDim - 1) * inner + 1 elements.
const size_t laneSize = (axisDim - 1) * inner + 1;

for (size_t outerIndex = 0; outerIndex < outer; ++outerIndex) {
for (size_t innerIndex = 0; innerIndex < inner; ++innerIndex) {
const size_t base = outerIndex * axisDim * inner + innerIndex;
const auto srcLane = srcData.subspan(base, laneSize);
const auto dstLane = dstData.subspan(base, laneSize);

float maxValue = -std::numeric_limits<float>::infinity();
for (size_t axisIndex = 0; axisIndex < axisDim; ++axisIndex) {
maxValue = std::max(maxValue, srcData[base + axisIndex * inner]);
maxValue = std::max(maxValue, srcLane[axisIndex * inner]);
}

float sum = 0.0f;
for (size_t axisIndex = 0; axisIndex < axisDim; ++axisIndex) {
const float value = std::exp(srcData[base + axisIndex * inner] - maxValue);
dstData[base + axisIndex * inner] = value;
const float value = std::exp(srcLane[axisIndex * inner] - maxValue);
dstLane[axisIndex * inner] = value;
sum += value;
}

for (size_t axisIndex = 0; axisIndex < axisDim; ++axisIndex) {
dstData[base + axisIndex * inner] /= sum;
dstLane[axisIndex * inner] /= sum;
}
}
}
Expand Down Expand Up @@ -154,7 +162,7 @@ void install_argmax(jsi::Runtime &rt, jsi::Object &module) {
throw jsi::JSError(rt, "argmax: dst shape must match src shape but with axis dimension 1");
}

const auto *srcData = reinterpret_cast<const float *>(src->data_.get());
const std::span<const float> srcData(reinterpret_cast<const float *>(src->data_.get()), src->numel_);

const auto axisDim = static_cast<size_t>(src->shape_[axisIdx]);
if (axisDim == 0) {
Expand All @@ -170,16 +178,22 @@ void install_argmax(jsi::Runtime &rt, jsi::Object &module) {
inner *= static_cast<size_t>(src->shape_[i]);
}

auto *dstData = reinterpret_cast<int32_t *>(dst->data_.get());
const std::span<int32_t> dstData(reinterpret_cast<int32_t *>(dst->data_.get()), dst->numel_);

// Elements along `axis` are strided by `inner`, so a lane spans from its
// first element to its last: (axisDim - 1) * inner + 1 elements.
const size_t laneSize = (axisDim - 1) * inner + 1;

// DO NOT swap loop order. This structure intentionally prioritizes the
// most common case (axis = -1, inner = 1) for sequential access.
for (size_t o = 0; o < outer; ++o) {
for (size_t i = 0; i < inner; ++i) {
const auto srcLane = srcData.subspan(o * axisDim * inner + i, laneSize);

float maxVal = -std::numeric_limits<float>::infinity();
int32_t maxIdx = 0;
for (size_t d = 0; d < axisDim; ++d) {
const float val = srcData[o * axisDim * inner + d * inner + i];
const float val = srcLane[d * inner];
if (val > maxVal) {
maxVal = val;
maxIdx = static_cast<int32_t>(d);
Expand Down Expand Up @@ -210,12 +224,12 @@ void install_threshold(jsi::Runtime &rt, jsi::Object &module) {

auto thresholdVal = conversions::asType<float>(rt, "threshold: threshold", args[2]);

const auto *srcData = reinterpret_cast<const float *>(src->data_.get());
auto *dstData = reinterpret_cast<float *>(dst->data_.get());
const std::span<const float> srcData(reinterpret_cast<const float *>(src->data_.get()), src->numel_);
const std::span<float> dstData(reinterpret_cast<float *>(dst->data_.get()), dst->numel_);

for (size_t i = 0; i < src->numel_; ++i) {
dstData[i] = (srcData[i] >= thresholdVal) ? 1.0f : 0.0f;
}
std::ranges::transform(srcData, dstData.begin(), [thresholdVal](const float value) {
return (value >= thresholdVal) ? 1.0f : 0.0f;
});

return jsi::Value(rt, args[1]);
};
Expand Down
Loading