namespace c10 {
namespace {
DeviceType parse_type(const std::string& device_string) {
- static const std::array<std::pair<std::string, DeviceType>, 7> types = {{
+ static const std::array<std::pair<std::string, DeviceType>, 8> types = {{
{"cpu", DeviceType::CPU},
{"cuda", DeviceType::CUDA},
{"mkldnn", DeviceType::MKLDNN},
{"opencl", DeviceType::OPENCL},
{"ideep", DeviceType::IDEEP},
{"hip", DeviceType::HIP},
+ {"msnpu", DeviceType::MSNPU},
}};
auto device = std::find_if(
types.begin(),
return device->second;
}
AT_ERROR(
- "Expected one of cpu, cuda, mkldnn, opengl, opencl, ideep, or hip device type at start of device string: ", device_string);
+ "Expected one of cpu, cuda, mkldnn, opengl, opencl, ideep, hip, msnpu device type at start of device string: ", device_string);
}
} // namespace
--- /dev/null
+#include <torch/extension.h>
+
+#include <ATen/ExtensionBackendRegistration.h>
+
+using namespace at;
+
+static int test_int;
+
+Tensor get_dummy_tensor() {
+ auto tensor_impl = c10::make_intrusive<TensorImpl, UndefinedTensorImpl>(
+ Storage(
+ caffe2::TypeMeta::Make<float>(), 0, at::DataPtr(nullptr, Device(DeviceType::MSNPU, 1)), nullptr, false),
+ MSNPUTensorId(),
+ false);
+ return Tensor(std::move(tensor_impl));
+}
+
+Tensor zeros_override(IntList size, const TensorOptions & options) {
+ test_int = 0;
+ return get_dummy_tensor();
+}
+
+Tensor add_override(const Tensor & a, const Tensor & b , Scalar c) {
+ test_int = 1;
+ return get_dummy_tensor();
+}
+
+Tensor sum_override(const Tensor & self) {
+ test_int = 2;
+ return get_dummy_tensor();
+}
+
+// needed for sum backwards
+Tensor expand_override(const Tensor & self, IntList size, bool implicit) {
+ return get_dummy_tensor();
+}
+
+
+Tensor kl_div_override(
+ const Tensor & self, const Tensor & target, int64_t reduction) {
+ test_int = 3;
+ return get_dummy_tensor();
+}
+
+Tensor kl_div_backward_override(
+ const Tensor & grad_output,
+ const Tensor & self,
+ const Tensor & target,
+ int64_t reduction) {
+ test_int = 4;
+ return get_dummy_tensor();
+}
+
+// numel and ones_like are needed for autograd backwards
+int64_t numel_override(const Tensor & self) {
+ return 1;
+}
+
+Tensor ones_like_override(const Tensor & self, const TensorOptions & options) {
+ return get_dummy_tensor();
+}
+
+void init_msnpu_extension() {
+ register_extension_backend_op(
+ Backend::MSNPU,
+ "zeros(IntList size, TensorOptions options) -> Tensor", &zeros_override);
+ register_extension_backend_op(
+ Backend::MSNPU,
+ "add(Tensor self, Tensor other, Scalar alpha) -> Tensor", &add_override);
+ register_extension_backend_op(
+ Backend::MSNPU,
+ "sum(Tensor self) -> Tensor", &sum_override);
+ register_extension_backend_op(
+ Backend::MSNPU,
+ "expand(Tensor self, IntList size, bool implicit) -> Tensor",
+ &expand_override);
+ register_extension_backend_op(
+ Backend::MSNPU,
+ "kl_div(Tensor self, Tensor target, int64_t reduction) -> Tensor",
+ &kl_div_override);
+ register_extension_backend_op(
+ Backend::MSNPU,
+ "kl_div_backward(Tensor grad_output, Tensor self, Tensor target, int64_t reduction) -> Tensor",
+ &kl_div_backward_override);
+ register_extension_backend_op(
+ Backend::MSNPU,
+ "numel(Tensor self) -> int64_t", &numel_override);
+ register_extension_backend_op(
+ Backend::MSNPU,
+ "ones_like(Tensor self, TensorOptions options) -> Tensor",
+ &ones_like_override);
+}
+
+int get_test_int() {
+ return test_int;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("init_msnpu_extension", &init_msnpu_extension);
+ m.def("get_test_int", &get_test_int);
+}
CppExtension(
'torch_test_cpp_extension.cpp', ['extension.cpp'],
extra_compile_args=CXX_FLAGS),
+ CppExtension(
+ 'torch_test_cpp_extension.msnpu', ['msnpu_extension.cpp'],
+ extra_compile_args=CXX_FLAGS),
]
if torch.cuda.is_available() and CUDA_HOME is not None:
try:
import torch_test_cpp_extension.cpp as cpp_extension
+ import torch_test_cpp_extension.msnpu as msnpu_extension
except ImportError:
warnings.warn(
"test_cpp_extensions.py cannot be invoked directly. Run "
torch.set_default_dtype(initial_default)
+class TestMSNPUTensor(common.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ msnpu_extension.init_msnpu_extension()
+
+ def test_unregistered(self):
+ a = torch.empty(5, 5, device='cpu')
+ with self.assertRaisesRegex(RuntimeError, "No function registered"):
+ b = torch.empty(5, 5, device='msnpu')
+
+ def test_zeros(self):
+ a = torch.zeros(5, 5, device='cpu')
+ self.assertEqual(a.device, torch.device('cpu'))
+ self.assertEqual(a.sum(), 0)
+
+ b = torch.zeros(5, 5, device='msnpu')
+ self.assertEqual(msnpu_extension.get_test_int(), 0)
+
+ def test_add(self):
+ a = torch.zeros(5, 5, device='msnpu')
+ self.assertEqual(msnpu_extension.get_test_int(), 0)
+
+ b = torch.zeros(5, 5, device='msnpu')
+ self.assertEqual(msnpu_extension.get_test_int(), 0)
+
+ c = torch.add(a, b)
+ self.assertEqual(msnpu_extension.get_test_int(), 1)
+
+ def test_backwards(self):
+ a = torch.zeros(5, 5, device='msnpu', requires_grad=True)
+ self.assertEqual(msnpu_extension.get_test_int(), 0)
+
+ b = torch.zeros(5, 5, device='msnpu')
+ self.assertEqual(msnpu_extension.get_test_int(), 0)
+
+ c = torch.kl_div(a, b)
+ self.assertEqual(msnpu_extension.get_test_int(), 3)
+
+ d = c.sum()
+ self.assertEqual(msnpu_extension.get_test_int(), 2)
+
+ d.backward()
+ self.assertEqual(msnpu_extension.get_test_int(), 4)
+
+
if __name__ == "__main__":
common.run_tests()
'`{}` type is not supported in python_default_init'.format(typename)
unpack_with_default = unpack_with_default_methods.get(typename)
default_expr = arg.get('python_default_init')
- # TODO: Type currently maps to ScalarType, figure out a cleaner solution
- if typename == 'const Type &':
- default_expr += '.scalarType()'
expr = 'r.{}({}, {})'.format(unpack_with_default, arg_index, default_expr)
else:
unpack = unpack_methods.get(typename, typename.lower())
if (is_factory_function and not has_type_input_arg) or has_options_arg:
default_type = get_type_default(declaration)
- py_default_dtype = 'self.type()' if is_like_function_with_options else None
+ py_default_dtype = 'self.scalar_type()' if is_like_function_with_options else None
dtype_arg = {
'default': default_type,
'dynamic_type': 'Type',
// for now, let's look these up by Backend; we could create our own enum in the future.
registerLayoutObject((THPLayout*)strided_layout, at::Backend::CPU);
registerLayoutObject((THPLayout*)strided_layout, at::Backend::CUDA);
+ registerLayoutObject((THPLayout*)strided_layout, at::Backend::MSNPU);
PyObject *sparse_coo_layout = THPLayout_New(at::Layout::Sparse, "torch.sparse_coo");
Py_INCREF(sparse_coo_layout);