Skip to content

Commit c6d8d52

Browse files
cyyevercyyever
committed
Replace JNI dtype maps with constexpr functions, make Stats scaling factor static
1 parent 5c9cd91 commit c6d8d52

4 files changed

Lines changed: 107 additions & 58 deletions

File tree

‎extension/android/jni/jni_layer.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -61,15 +61,15 @@ class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
6161
// Java wrapper currently only supports contiguous tensors.
6262

6363
const auto scalarType = tensor.scalar_type();
64-
if (scalar_type_to_java_dtype.count(scalarType) == 0) {
64+
const int jdtype = scalar_type_to_java_dtype(scalarType);
65+
if (jdtype < 0) {
6566
std::stringstream ss;
6667
ss << "executorch::aten::Tensor scalar type "
6768
<< static_cast<int>(scalarType) << " is not supported on java side";
6869
jni_helper::throwExecutorchException(
6970
static_cast<uint32_t>(Error::InvalidArgument), ss.str().c_str());
7071
return nullptr;
7172
}
72-
int jdtype = scalar_type_to_java_dtype.at(scalarType);
7373

7474
const auto& tensor_shape = tensor.sizes();
7575
std::vector<jlong> tensor_shape_vec;
@@ -130,14 +130,14 @@ class TensorHybrid : public facebook::jni::HybridClass<TensorHybrid> {
130130
numel *= shapeArr[i];
131131
}
132132
JNIEnv* jni = facebook::jni::Environment::current();
133-
if (java_dtype_to_scalar_type.count(jdtype) == 0) {
133+
const ScalarType scalar_type = java_dtype_to_scalar_type(jdtype);
134+
if (scalar_type == ScalarType::Undefined) {
134135
std::stringstream ss;
135136
ss << "Unknown Tensor jdtype: [" << jdtype << "]";
136137
jni_helper::throwExecutorchException(
137138
static_cast<uint32_t>(Error::InvalidArgument), ss.str().c_str());
138139
return nullptr;
139140
}
140-
ScalarType scalar_type = java_dtype_to_scalar_type.at(jdtype);
141141
const jlong dataCapacity = jni->GetDirectBufferCapacity(jbuffer.get());
142142
if (dataCapacity < 0) {
143143
std::stringstream ss;

‎extension/android/jni/jni_layer_constants.h‎

Lines changed: 101 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,6 @@
66
* LICENSE file in the root directory of this source tree.
77
*/
88

9-
#include <unordered_map>
10-
119
#include <executorch/runtime/core/exec_aten/exec_aten.h>
1210

1311
namespace executorch::extension {
@@ -39,58 +37,109 @@ constexpr static int kTensorDTypeBits16 = 22;
3937

4038
using executorch::aten::ScalarType;
4139

42-
const std::unordered_map<ScalarType, int> scalar_type_to_java_dtype = {
43-
{ScalarType::Byte, kTensorDTypeUInt8},
44-
{ScalarType::Char, kTensorDTypeInt8},
45-
{ScalarType::Short, kTensorDTypeInt16},
46-
{ScalarType::Int, kTensorDTypeInt32},
47-
{ScalarType::Long, kTensorDTypeInt64},
48-
{ScalarType::Half, kTensorDTypeHalf},
49-
{ScalarType::Float, kTensorDTypeFloat},
50-
{ScalarType::Double, kTensorDTypeDouble},
40+
// Returns the Java dtype code for a ScalarType, or -1 if unsupported.
41+
constexpr int scalar_type_to_java_dtype(ScalarType scalar_type) {
42+
switch (scalar_type) {
43+
case ScalarType::Byte:
44+
return kTensorDTypeUInt8;
45+
case ScalarType::Char:
46+
return kTensorDTypeInt8;
47+
case ScalarType::Short:
48+
return kTensorDTypeInt16;
49+
case ScalarType::Int:
50+
return kTensorDTypeInt32;
51+
case ScalarType::Long:
52+
return kTensorDTypeInt64;
53+
case ScalarType::Half:
54+
return kTensorDTypeHalf;
55+
case ScalarType::Float:
56+
return kTensorDTypeFloat;
57+
case ScalarType::Double:
58+
return kTensorDTypeDouble;
5159
// These types are not supported yet
52-
// {ScalarType::ComplexHalf, kTensorDTypeComplexHalf},
53-
// {ScalarType::ComplexFloat, kTensorDTypeComplexFloat},
54-
// {ScalarType::ComplexDouble, kTensorDTypeComplexDouble},
55-
{ScalarType::Bool, kTensorDTypeBool},
56-
{ScalarType::QInt8, kTensorDTypeQint8},
57-
{ScalarType::QUInt8, kTensorDTypeQuint8},
58-
{ScalarType::QInt32, kTensorDTypeQint32},
59-
{ScalarType::BFloat16, kTensorDTypeBFloat16},
60-
{ScalarType::QUInt4x2, kTensorDTypeQuint4x2},
61-
{ScalarType::QUInt2x4, kTensorDTypeQuint2x4},
62-
{ScalarType::Bits1x8, kTensorDTypeBits1x8},
63-
{ScalarType::Bits2x4, kTensorDTypeBits2x4},
64-
{ScalarType::Bits4x2, kTensorDTypeBits4x2},
65-
{ScalarType::Bits8, kTensorDTypeBits8},
66-
{ScalarType::Bits16, kTensorDTypeBits16},
67-
};
60+
// case ScalarType::ComplexHalf:
61+
// case ScalarType::ComplexFloat:
62+
// case ScalarType::ComplexDouble:
63+
case ScalarType::Bool:
64+
return kTensorDTypeBool;
65+
case ScalarType::QInt8:
66+
return kTensorDTypeQint8;
67+
case ScalarType::QUInt8:
68+
return kTensorDTypeQuint8;
69+
case ScalarType::QInt32:
70+
return kTensorDTypeQint32;
71+
case ScalarType::BFloat16:
72+
return kTensorDTypeBFloat16;
73+
case ScalarType::QUInt4x2:
74+
return kTensorDTypeQuint4x2;
75+
case ScalarType::QUInt2x4:
76+
return kTensorDTypeQuint2x4;
77+
case ScalarType::Bits1x8:
78+
return kTensorDTypeBits1x8;
79+
case ScalarType::Bits2x4:
80+
return kTensorDTypeBits2x4;
81+
case ScalarType::Bits4x2:
82+
return kTensorDTypeBits4x2;
83+
case ScalarType::Bits8:
84+
return kTensorDTypeBits8;
85+
case ScalarType::Bits16:
86+
return kTensorDTypeBits16;
87+
default:
88+
return -1;
89+
}
90+
}
6891

69-
const std::unordered_map<int, ScalarType> java_dtype_to_scalar_type = {
70-
{kTensorDTypeUInt8, ScalarType::Byte},
71-
{kTensorDTypeInt8, ScalarType::Char},
72-
{kTensorDTypeInt16, ScalarType::Short},
73-
{kTensorDTypeInt32, ScalarType::Int},
74-
{kTensorDTypeInt64, ScalarType::Long},
75-
{kTensorDTypeHalf, ScalarType::Half},
76-
{kTensorDTypeFloat, ScalarType::Float},
77-
{kTensorDTypeDouble, ScalarType::Double},
92+
// Returns the ScalarType for a Java dtype code, or ScalarType::Undefined if
93+
// unsupported.
94+
constexpr ScalarType java_dtype_to_scalar_type(int java_dtype) {
95+
switch (java_dtype) {
96+
case kTensorDTypeUInt8:
97+
return ScalarType::Byte;
98+
case kTensorDTypeInt8:
99+
return ScalarType::Char;
100+
case kTensorDTypeInt16:
101+
return ScalarType::Short;
102+
case kTensorDTypeInt32:
103+
return ScalarType::Int;
104+
case kTensorDTypeInt64:
105+
return ScalarType::Long;
106+
case kTensorDTypeHalf:
107+
return ScalarType::Half;
108+
case kTensorDTypeFloat:
109+
return ScalarType::Float;
110+
case kTensorDTypeDouble:
111+
return ScalarType::Double;
78112
// These types are not supported yet
79-
// {kTensorDTypeComplexHalf, ScalarType::ComplexHalf},
80-
// {kTensorDTypeComplexFloat, ScalarType::ComplexFloat},
81-
// {kTensorDTypeComplexDouble, ScalarType::ComplexDouble},
82-
{kTensorDTypeBool, ScalarType::Bool},
83-
{kTensorDTypeQint8, ScalarType::QInt8},
84-
{kTensorDTypeQuint8, ScalarType::QUInt8},
85-
{kTensorDTypeQint32, ScalarType::QInt32},
86-
{kTensorDTypeBFloat16, ScalarType::BFloat16},
87-
{kTensorDTypeQuint4x2, ScalarType::QUInt4x2},
88-
{kTensorDTypeQuint2x4, ScalarType::QUInt2x4},
89-
{kTensorDTypeBits1x8, ScalarType::Bits1x8},
90-
{kTensorDTypeBits2x4, ScalarType::Bits2x4},
91-
{kTensorDTypeBits4x2, ScalarType::Bits4x2},
92-
{kTensorDTypeBits8, ScalarType::Bits8},
93-
{kTensorDTypeBits16, ScalarType::Bits16},
94-
};
113+
// case kTensorDTypeComplexHalf:
114+
// case kTensorDTypeComplexFloat:
115+
// case kTensorDTypeComplexDouble:
116+
case kTensorDTypeBool:
117+
return ScalarType::Bool;
118+
case kTensorDTypeQint8:
119+
return ScalarType::QInt8;
120+
case kTensorDTypeQuint8:
121+
return ScalarType::QUInt8;
122+
case kTensorDTypeQint32:
123+
return ScalarType::QInt32;
124+
case kTensorDTypeBFloat16:
125+
return ScalarType::BFloat16;
126+
case kTensorDTypeQuint4x2:
127+
return ScalarType::QUInt4x2;
128+
case kTensorDTypeQuint2x4:
129+
return ScalarType::QUInt2x4;
130+
case kTensorDTypeBits1x8:
131+
return ScalarType::Bits1x8;
132+
case kTensorDTypeBits2x4:
133+
return ScalarType::Bits2x4;
134+
case kTensorDTypeBits4x2:
135+
return ScalarType::Bits4x2;
136+
case kTensorDTypeBits8:
137+
return ScalarType::Bits8;
138+
case kTensorDTypeBits16:
139+
return ScalarType::Bits16;
140+
default:
141+
return ScalarType::Undefined;
142+
}
143+
}
95144

96145
} // namespace executorch::extension

‎extension/llm/runner/pybindings.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -310,7 +310,7 @@ PYBIND11_MODULE(_llm_runner, m) {
310310

311311
// Bind Stats
312312
py::class_<Stats>(m, "Stats")
313-
.def_readonly(
313+
.def_readonly_static(
314314
"SCALING_FACTOR_UNITS_PER_SECOND",
315315
&Stats::SCALING_FACTOR_UNITS_PER_SECOND)
316316
.def_readonly("model_load_start_ms", &Stats::model_load_start_ms)

‎extension/llm/runner/stats.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ namespace llm {
2020

2121
struct ET_EXPERIMENTAL Stats {
2222
// Scaling factor for timestamps - in this case, we use ms.
23-
const long SCALING_FACTOR_UNITS_PER_SECOND = 1000;
23+
static constexpr long SCALING_FACTOR_UNITS_PER_SECOND = 1000;
2424
// Time stamps for the different stages of the execution
2525
// model_load_start_ms: Start of model loading.
2626
long model_load_start_ms = 0;

0 commit comments

Comments
 (0)