Skip to content

Commit 809a1fd

Browse files
authored
Add portable rand kernel implementation (#11127)
### Summary Add a portable operator implementation for ATen rand. This is a core op that we previously have not had a kernel for. ### Test plan I've added operator-level tests for rand. It's relatively sparse, but validates that the mean and stdev are reasonably sane, as well as that the operator functions correctly for different dtypes, ranks, and sizes.
1 parent 9f74777 commit 809a1fd

File tree

7 files changed

+163
-0
lines changed

7 files changed

+163
-0
lines changed

kernels/aten/functions.yaml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,8 @@
315315

316316
- op: prod.out
317317

318+
- op: rand.out
319+
318320
- op: reciprocal.out
319321

320322
- op: relu.out

kernels/portable/cpu/op_rand.cpp

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
#include <c10/util/irange.h>
9+
10+
#include <executorch/kernels/portable/cpu/scalar_utils.h>
11+
#include <executorch/runtime/kernel/kernel_includes.h>
12+
13+
#include <random>
14+
15+
namespace torch {
16+
namespace executor {
17+
namespace native {
18+
19+
using executorch::aten::IntArrayRef;
20+
using Tensor = executorch::aten::Tensor;
21+
using ScalarType = executorch::aten::ScalarType;
22+
23+
Tensor&
24+
rand_out(KernelRuntimeContext& ctx, const IntArrayRef sizes, Tensor& out) {
25+
(void)ctx;
26+
27+
std::mt19937 gen((std::random_device())());
28+
std::uniform_real_distribution<double> dist(0.0, 1.0);
29+
30+
// Resize for dynamic shape
31+
ET_KERNEL_CHECK_MSG(
32+
ctx,
33+
resize_tensor(out, sizes) == Error::Ok,
34+
InvalidArgument,
35+
out,
36+
"Failed to resize output tensor.");
37+
38+
ET_SWITCH_FLOATHBF16_TYPES(out.scalar_type(), ctx, "randn.out", CTYPE, [&] {
39+
auto data_out = out.mutable_data_ptr<CTYPE>();
40+
for (const auto i : c10::irange(out.numel())) {
41+
data_out[i] = static_cast<CTYPE>(dist(gen));
42+
}
43+
});
44+
45+
return out;
46+
}
47+
48+
} // namespace native
49+
} // namespace executor
50+
} // namespace torch

kernels/portable/functions.yaml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -713,6 +713,12 @@
713713
- arg_meta: null
714714
kernel_name: torch::executor::prod_out
715715

716+
- op: rand.out
717+
kernels:
718+
- arg_meta: null
719+
kernel_name: torch::executor::rand_out
720+
tags: nondeterministic_seeded
721+
716722
- op: reciprocal.out
717723
kernels:
718724
- arg_meta: null

kernels/test/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,7 @@ set(all_test_sources
197197
"op_permute_copy_test.cpp"
198198
"op_pixel_shuffle_test.cpp"
199199
"op_prod_test.cpp"
200+
"op_rand_test.cpp"
200201
"op_reciprocal_test.cpp"
201202
"op_relu_test.cpp"
202203
"op_remainder_test.cpp"

kernels/test/op_rand_test.cpp

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <c10/util/irange.h>
10+
#include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
11+
#include <executorch/kernels/test/TestUtil.h>
12+
#include <executorch/runtime/core/exec_aten/exec_aten.h>
13+
#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
14+
#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
15+
16+
#include <gtest/gtest.h>
17+
18+
#include <cmath>
19+
#include <numeric>
20+
21+
using executorch::aten::IntArrayRef;
22+
using executorch::aten::ScalarType;
23+
using executorch::aten::Tensor;
24+
using torch::executor::testing::TensorFactory;
25+
26+
class OpRandTest : public OperatorTest {
27+
protected:
28+
void op_rand_out(const IntArrayRef sizes, Tensor& out) {
29+
torch::executor::aten::rand_outf(context_, sizes, out);
30+
}
31+
32+
template <typename CTYPE, ScalarType DTYPE>
33+
void test_rand(std::vector<int64_t>& sizes) {
34+
TensorFactory<DTYPE> tf;
35+
36+
// Tensor factory wants int32 scales, op kernel wants int64.
37+
std::vector<int32_t> sizes_i32;
38+
std::transform(
39+
sizes.begin(),
40+
sizes.end(),
41+
std::back_inserter(sizes_i32),
42+
[](int64_t s) { return static_cast<int32_t>(s); });
43+
Tensor out = tf.zeros(sizes_i32);
44+
45+
IntArrayRef sizes_ref(sizes.data(), sizes.size());
46+
op_rand_out(sizes_ref, out);
47+
48+
// Check mean and standard deviation. To avoid flaky CI, test pretty
49+
// loosely.
50+
auto out_data = out.const_data_ptr<CTYPE>();
51+
double mean =
52+
std::accumulate(
53+
out_data,
54+
out_data + out.numel(),
55+
0.0,
56+
[](double acc, CTYPE n) { return acc + static_cast<double>(n); }) /
57+
out.numel();
58+
double var = std::accumulate(
59+
out_data,
60+
out_data + out.numel(),
61+
0.0,
62+
[=](double acc, CTYPE n) {
63+
return acc + std::pow(static_cast<double>(n) - mean, 2);
64+
}) /
65+
out.numel();
66+
auto stdev = std::sqrt(var);
67+
68+
// These are very rough thresholds. A better test implementation would
69+
// probably do a proper statistical test to compare the generated empirical
70+
// data to the reference distribution, but this should do.
71+
72+
// Expected mean is 0.5
73+
EXPECT_NEAR(mean, 0.5, 5.0 / std::sqrt(out.numel()));
74+
// Expected stdev is 1/sqrt(12) ~= 0.289
75+
EXPECT_NEAR(stdev, 1.0 / std::sqrt(12), 0.1);
76+
EXPECT_GT(stdev, 0);
77+
}
78+
};
79+
80+
TEST_F(OpRandTest, SmokeTest) {
81+
std::vector<int64_t> sizes = {2, 3, 4, 128};
82+
83+
#define TEST_ENTRY(ctype, dtype) test_rand<ctype, ScalarType::dtype>(sizes);
84+
ET_FORALL_FLOATHBF16_TYPES(TEST_ENTRY);
85+
#undef TEST_ENTRY
86+
}
87+
88+
TEST_F(OpRandTest, Rank) {
89+
std::vector<int64_t> sizes = {1024};
90+
91+
for (int64_t i = 0; i < 4; i++) {
92+
sizes.push_back(i + 1);
93+
test_rand<float, executorch::aten::ScalarType::Float>(sizes);
94+
}
95+
}

kernels/test/targets.bzl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,6 +285,7 @@ def define_common_targets():
285285
_common_op_test("op_pixel_unshuffle_test", ["aten", "portable"])
286286
_common_op_test("op_pow_test", ["aten", "portable"])
287287
_common_op_test("op_prod_test", ["aten", "portable"])
288+
_common_op_test("op_rand_test", ["aten", "portable"])
288289
_common_op_test("op_reciprocal_test", ["aten", "portable"])
289290
_common_op_test("op_relu_test", ["aten", "portable"])
290291
_common_op_test("op_remainder_test", ["aten", "portable"])

shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -973,6 +973,14 @@ ATEN_OPS = (
973973
"//executorch/kernels/portable/cpu/util:reduce_util",
974974
],
975975
),
976+
op_target(
977+
name = "op_rand",
978+
deps = [
979+
":scalar_utils",
980+
"//executorch/runtime/core/exec_aten/util:scalar_type_util",
981+
"//executorch/runtime/core/exec_aten/util:tensor_util",
982+
]
983+
),
976984
op_target(
977985
name = "op_reciprocal",
978986
deps = [

0 commit comments

Comments
 (0)