1 #include <gtest/gtest.h> 4 #include <ATen/NativeFunctions.h> 5 #include <ATen/ExtensionBackendRegistration.h> 13 auto tensor_impl = c10::make_intrusive<TensorImpl, UndefinedTensorImpl>(
15 caffe2::TypeMeta::Make<float>(), 0,
at::DataPtr(
nullptr,
Device(DeviceType::MSNPU, 1)),
nullptr,
false),
18 return Tensor(std::move(tensor_impl));
31 TEST(BackendExtensionTest, TestRegisterOp) {
32 EXPECT_ANY_THROW(empty({5, 5}, at::kMSNPU));
33 register_extension_backend_op(
35 "empty(IntArrayRef size, TensorOptions options) -> Tensor", &empty_override);
36 Tensor a = empty({5, 5}, at::kMSNPU);
39 ASSERT_EQ(a.
dtype(), caffe2::TypeMeta::Make<float>());
40 ASSERT_EQ(test_int, 1);
42 EXPECT_ANY_THROW(empty_like(a, at::kMSNPU));
43 register_extension_backend_op(
45 "empty_like(Tensor self, TensorOptions options) -> Tensor", &empty_like_override);
46 Tensor b = empty_like(a, at::kMSNPU);
47 ASSERT_EQ(test_int, 2);
49 EXPECT_ANY_THROW(add(a, b));
50 register_extension_backend_op(
52 "add(Tensor self, Tensor other, Scalar alpha) -> Tensor", &add_override);
54 ASSERT_EQ(test_int, 3);
57 Tensor d = empty({5, 5}, at::kCPU);
62 register_extension_backend_op(
64 "empty(IntArrayRef size, TensorOptions options) -> Tensor", &empty_override)
caffe2::TypeMeta dtype() const noexcept
Returns a Tensor's dtype (TypeMeta). Defined in TensorMethods.h.
Represents a a compute device on which a tensor is located.
Device device() const
Returns a Tensor's device.
Flush-To-Zero and Denormals-Are-Zero mode.
DeviceIndex index() const noexcept
Returns the optional index.
DeviceType type() const noexcept
Returns the type of device this is.