1 #include "kernel_runtime_context.h"
2 #include "operator_registry.h"
3
4 #include <gtest/gtest.h>
5
6 namespace torch {
7 namespace executor {
8
9 // add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)
TEST(OperatorRegistrationTest,Add)10 TEST(OperatorRegistrationTest, Add) {
11 EValue values[4];
12 values[0] = EValue(at::ones({2, 3}));
13 values[1] = EValue(at::ones({2, 3}));
14 values[2] = EValue(int64_t(1));
15 values[3] = EValue(at::zeros({2, 3}));
16 ASSERT_TRUE(hasKernelFn("aten::add.out"));
17 auto op = getKernelFn("aten::add.out");
18
19 EValue* kernel_values[4];
20 for (size_t i = 0; i < 4; i++) {
21 kernel_values[i] = &values[i];
22 }
23 KernelRuntimeContext context{};
24 op(context, kernel_values);
25 at::Tensor expected = at::ones({2, 3});
26 expected = at::fill(expected, 2);
27 ASSERT_TRUE(expected.equal(kernel_values[3]->toTensor()));
28
29 }
30
31 // custom::add_3.out(Tensor a, Tensor b, Tensor c, *, Tensor(a!) out) -> Tensor(a!)
TEST(OperatorRegistrationTest,CustomAdd3)32 TEST(OperatorRegistrationTest, CustomAdd3) {
33 EValue values[4];
34 values[0] = EValue(at::ones({2, 3}));
35 values[1] = EValue(at::ones({2, 3}));
36 values[2] = EValue(at::ones({2, 3}));
37 values[3] = EValue(at::zeros({2, 3}));
38 ASSERT_TRUE(hasKernelFn("custom::add_3.out"));
39 auto op = getKernelFn("custom::add_3.out");
40
41 EValue* kernel_values[4];
42 for (size_t i = 0; i < 4; i++) {
43 kernel_values[i] = &values[i];
44 }
45 KernelRuntimeContext context{};
46 op(context, kernel_values);
47 at::Tensor expected = at::ones({2, 3});
48 expected = at::fill(expected, 3);
49 ASSERT_TRUE(expected.equal(kernel_values[3]->toTensor()));
50
51 }
52 } // namespace executor
53 } // namespace torch
54