pytorch/test/cpp/api/tensor_options.cpp
Edward Yang 6bdbad93b9 Refactor Device to not depend on Backend. (#10478)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/10478

- Removed Backend constructor from Device, and fixed all
  use-sites to use DeviceType::CPU instead of kCPU, or
  use a new function backendToDeviceType to perform
  the conversion.
- New method device_type() on Type; it gives you the
  underlying device type, e.g., CPU for SparseCPU.
- We add backward compatibility for kCPU/kCUDA uses,
  by introducing a new special type which is implicitly
  convertible to both DeviceType and Backend.  As long as
  you don't define a function that's overloaded on both
  DeviceType and Backend (but not on BackendOrDeviceType),
  the implicit conversions will ensure that uses
  of at::Device(at::kCPU) keep working. We fixed use-sites in
  the library, but did NOT fix sites in the test code, so that
  we can exercise this BC code.

Reviewed By: Yangqing

Differential Revision: D9301861

fbshipit-source-id: 9a9d88620500715c7b37e655b4fd761f6dd72716
2018-08-18 17:39:14 -07:00

137 lines
4.0 KiB
C++

#include "catch.hpp"
#include <torch/tensor.h>
#include <ATen/Context.h>
#include <ATen/Functions.h>
#include <ATen/OptionsGuard.h>
#include <ATen/TensorOptions.h>
#include <vector>
#include <string>
using namespace at;
// A macro so we don't lose location information when an assertion fails.
#define REQUIRE_OPTIONS(device_, index_, type_, layout_) \
REQUIRE(options.device().type() == Device((device_), (index_)).type()); \
REQUIRE(options.device().index() == Device((device_), (index_)).index()); \
REQUIRE(options.dtype() == (type_)); \
REQUIRE(options.layout() == (layout_))
#define REQUIRE_TENSOR_OPTIONS(device_, index_, type_, layout_) \
REQUIRE(tensor.device().type() == Device((device_), (index_)).type()); \
REQUIRE(tensor.device().index() == Device((device_), (index_)).index()); \
REQUIRE(tensor.type().scalarType() == (type_)); \
REQUIRE(tensor.type().layout() == (layout_))
TEST_CASE("TensorOptions/DefaultsToTheRightValues") {
TensorOptions options;
REQUIRE_OPTIONS(kCPU, -1, kFloat, kStrided);
}
TEST_CASE("TensorOptions/ReturnsTheCorrectType") {
auto options = TensorOptions().device(kCPU).dtype(kInt).layout(kSparse);
REQUIRE(options.type() == getType(Backend::SparseCPU, kInt));
}
TEST_CASE("TensorOptions/UtilityFunctionsReturnTheRightTensorOptions") {
auto options = dtype(kInt);
REQUIRE_OPTIONS(kCPU, -1, kInt, kStrided);
options = layout(kSparse);
REQUIRE_OPTIONS(kCPU, -1, kFloat, kSparse);
options = device({kCUDA, 1});
REQUIRE_OPTIONS(kCUDA, 1, kFloat, kStrided);
options = device_index(1);
REQUIRE_OPTIONS(kCUDA, 1, kFloat, kStrided);
options = dtype(kByte).layout(kSparse).device({kCUDA, 2}).device_index(3);
REQUIRE_OPTIONS(kCUDA, 3, kByte, kSparse);
}
TEST_CASE("TensorOptions/ConstructsWellFromCPUTypes") {
TensorOptions options;
REQUIRE_OPTIONS(kCPU, -1, kFloat, kStrided);
options = TensorOptions({kCPU, 0});
REQUIRE_OPTIONS(kCPU, 0, kFloat, kStrided);
options = TensorOptions(kInt);
REQUIRE_OPTIONS(kCPU, -1, kInt, kStrided);
options = TensorOptions(getType(Backend::SparseCPU, kFloat));
REQUIRE_OPTIONS(kCPU, -1, kFloat, kSparse);
options = TensorOptions(getType(Backend::SparseCPU, kByte));
REQUIRE_OPTIONS(kCPU, -1, kByte, kSparse);
}
TEST_CASE("TensorOptions/ConstructsWellFromCPUTensors") {
auto options = TensorOptions(empty(5, kDouble));
REQUIRE_OPTIONS(kCPU, -1, kDouble, kStrided);
options = TensorOptions(empty(5, getType(Backend::SparseCPU, kByte)));
REQUIRE_OPTIONS(kCPU, -1, kByte, kSparse);
}
TEST_CASE("TensorOptions/ConstructsWellFromVariables") {
auto options = TensorOptions(torch::empty(5));
REQUIRE_OPTIONS(kCPU, -1, kFloat, kStrided);
REQUIRE(!options.requires_grad());
options = TensorOptions(torch::empty(5, at::requires_grad()));
REQUIRE_OPTIONS(kCPU, -1, kFloat, kStrided);
REQUIRE(!options.requires_grad());
}
TEST_CASE("Device/ParsesCorrectlyFromString") {
Device device("cpu:0");
REQUIRE(device == Device(kCPU, 0));
device = Device("cpu");
REQUIRE(device == Device(kCPU));
device = Device("cuda:123");
REQUIRE(device == Device(kCUDA, 123));
device = Device("cuda");
REQUIRE(device == Device(kCUDA));
std::vector<std::string> badnesses = {
"", "cud:1", "cuda:", "cpu::1", ":1", "3", "tpu:4", "??"};
for (const auto& badness : badnesses) {
REQUIRE_THROWS(Device(badness));
}
}
TEST_CASE("OptionsGuard") {
Tensor tensor;
{
OptionsGuard guard(TensorOptions{});
tensor = at::empty({10});
}
REQUIRE_TENSOR_OPTIONS(kCPU, -1, kFloat, kStrided);
{
OptionsGuard guard(TensorOptions().dtype(kInt));
tensor = at::empty({10});
}
REQUIRE_TENSOR_OPTIONS(kCPU, -1, kInt, kStrided);
{
OptionsGuard guard(TensorOptions().dtype(kInt).layout(kSparse));
tensor = at::empty({10});
}
REQUIRE_TENSOR_OPTIONS(kCPU, -1, kInt, kSparse);
{
OptionsGuard guard(requires_grad(true));
tensor = torch::empty({10});
}
REQUIRE_TENSOR_OPTIONS(kCPU, -1, kFloat, kStrided);
REQUIRE(tensor.requires_grad());
}