Caffe2 - C++ API
A deep learning, cross platform ML framework
xla_tensor_test.cpp
1 #include <gtest/gtest.h>
2 
3 #include <ATen/ATen.h>
4 
5 using namespace at;
6 
7 void XLAFree(void *ptr) {
8  free(ptr);
9 }
10 
11 void* XLAMalloc(ptrdiff_t size) {
12  return malloc(size);
13 }
14 
15 struct XLAAllocator final : public at::Allocator {
16  at::DataPtr allocate(size_t size) const override {
17  auto* ptr = XLAMalloc(size);
18  return {ptr, ptr, &XLAFree, at::DeviceType::XLA};
19  }
20  at::DeleterFnPtr raw_deleter() const override {
21  return &XLAFree;
22  }
23 };
24 
25 TEST(XlaTensorTest, TestNoStorage) {
26  XLAAllocator allocator;
27  auto storage = Storage(caffe2::TypeMeta::Make<float>(), 0, &allocator, true);
28  auto tensor_impl = c10::make_intrusive<TensorImpl, UndefinedTensorImpl>(
29  std::move(storage),
30  XLATensorId(),
31  /*is_variable=*/false);
32  at::Tensor t(std::move(tensor_impl));
33  ASSERT_TRUE(t.device() == DeviceType::XLA);
34 }
Flush-To-Zero and Denormals-Are-Zero mode.