mirror of
https://github.com/pytorch/pytorch.git
synced 2025-11-18 09:34:57 +08:00
0ec53beaeb
This PR moves the implementations of Tensor accessor classes to headeronly with the following modifications:
- Add ArrayRef and IndexBoundsCheck template parameters to refactor out the usages of `IntArrayRef` and `TORCH_CHECK_INDEX` from Tensor accessor implementations.
- Eliminate usage of `c10::irange` as it is not headeronly-compatible.
- Introduce `torch::headeronly::{TensorAccessorBase,TensorAccessor, GenericPackedTensorAccessorBase, GenericPackedTensorAccessor}` that are headeronly-equivalent to `at::{TensorAccessorBase,TensorAccessor, GenericPackedTensorAccessorBase, GenericPackedTensorAccessor}`. Both these sets of template classes use original implementations from `torch::headeronly::detail` that have new template parameters `ArrayRefCls` and `IndexBoundsCheck` to facilitate `at` and `torch::headeronly` implementations of ArrayRef and checking indices.
TODO:
- ~when https://github.com/pytorch/pytorch/pull/164991 lands, eliminate the placeholder class HeaderOnlyArrayRef~ UPDATE: done.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/166855
Approved by: https://github.com/janeyx99
51 lines
1.4 KiB
C++
51 lines
1.4 KiB
C++
#include <gtest/gtest.h>
|
|
#include <torch/headeronly/core/TensorAccessor.h>
|
|
#include <string>
|
|
|
|
TEST(TestAccessor, HeaderOnlyTensorAccessor) {
|
|
std::vector<int32_t> v = {11, 12, 13, 21, 22, 23};
|
|
std::vector<int64_t> sizes = {2, 3};
|
|
std::vector<int64_t> strides = {3, 1};
|
|
|
|
auto acc = torch::headeronly::HeaderOnlyTensorAccessor<int32_t, 2>(
|
|
v.data(), sizes.data(), strides.data());
|
|
EXPECT_EQ(acc[0][0], 11);
|
|
EXPECT_EQ(acc[0][1], 12);
|
|
EXPECT_EQ(acc[0][2], 13);
|
|
EXPECT_EQ(acc[1][0], 21);
|
|
EXPECT_EQ(acc[1][1], 22);
|
|
EXPECT_EQ(acc[1][2], 23);
|
|
}
|
|
|
|
TEST(TestAccessor, HeaderOnlyGenericPackedTensorAccessor) {
|
|
std::vector<int32_t> v = {11, 12, 13, 21, 22, 23};
|
|
std::vector<int64_t> sizes = {2, 3};
|
|
std::vector<int64_t> strides = {3, 1};
|
|
|
|
auto acc =
|
|
torch::headeronly::HeaderOnlyGenericPackedTensorAccessor<int32_t, 2>(
|
|
v.data(), sizes.data(), strides.data());
|
|
EXPECT_EQ(acc[0][0], 11);
|
|
EXPECT_EQ(acc[0][1], 12);
|
|
EXPECT_EQ(acc[0][2], 13);
|
|
EXPECT_EQ(acc[1][0], 21);
|
|
EXPECT_EQ(acc[1][1], 22);
|
|
EXPECT_EQ(acc[1][2], 23);
|
|
|
|
auto tacc = acc.transpose(0, 1);
|
|
EXPECT_EQ(tacc[0][0], 11);
|
|
EXPECT_EQ(tacc[0][1], 21);
|
|
EXPECT_EQ(tacc[1][0], 12);
|
|
EXPECT_EQ(tacc[1][1], 22);
|
|
EXPECT_EQ(tacc[2][0], 13);
|
|
EXPECT_EQ(tacc[2][1], 23);
|
|
|
|
try {
|
|
acc.transpose(0, 2);
|
|
} catch (const std::exception& e) {
|
|
EXPECT_TRUE(
|
|
std::string(e.what()).find("HeaderOnlyIndexBoundsCheck") !=
|
|
std::string::npos);
|
|
}
|
|
}
|