1 #include <algorithm>
2
3 #include <torch/csrc/lazy/core/ops/utils.h>
4 #include <torch/csrc/lazy/core/permutation_util.h>
5 #include <torch/csrc/lazy/core/tensor_util.h>
6 #include <torch/csrc/lazy/core/util.h>
7
8 namespace torch {
9 namespace lazy {
10
StrideIsSupported(c10::ArrayRef<int64_t> stride)11 bool StrideIsSupported(c10::ArrayRef<int64_t> stride) {
12 std::vector<int64_t> sorted_stride(stride.begin(), stride.end());
13 std::sort(sorted_stride.begin(), sorted_stride.end());
14 return stride.empty() || sorted_stride.front() == 1;
15 }
16
GetArrayStridePermutation(c10::ArrayRef<int64_t> stride)17 std::vector<int64_t> GetArrayStridePermutation(c10::ArrayRef<int64_t> stride) {
18 std::vector<int64_t> permutation = Iota<int64_t>(stride.size());
19 std::sort(permutation.begin(), permutation.end(), [&](int64_t a, int64_t b) {
20 return stride[a] > stride[b];
21 });
22 return permutation;
23 }
24
MakeDiagonalShape(const Shape & shape,int64_t offset,int64_t dim1,int64_t dim2)25 Shape MakeDiagonalShape(
26 const Shape& shape,
27 int64_t offset,
28 int64_t dim1,
29 int64_t dim2) {
30 std::vector<int64_t> dimensions;
31 for (const auto dim : c10::irange(shape.dim())) {
32 if (dim != dim1 && dim != dim2) {
33 dimensions.push_back(shape.size(dim));
34 }
35 }
36 int64_t dsize = 0;
37 if (offset >= 0) {
38 dsize = std::max<int64_t>(
39 std::min(shape.size(dim1), shape.size(dim2) - offset), 0);
40 } else {
41 dsize = std::max<int64_t>(
42 std::min(shape.size(dim1) + offset, shape.size(dim2)), 0);
43 }
44 dimensions.push_back(dsize);
45 return Shape(shape.scalar_type(), dimensions);
46 }
47
MakePermuteShape(const Shape & source_shape,c10::ArrayRef<int64_t> permutation)48 Shape MakePermuteShape(
49 const Shape& source_shape,
50 c10::ArrayRef<int64_t> permutation) {
51 return Shape(
52 source_shape.scalar_type(),
53 PermuteDimensions(permutation, source_shape.sizes()));
54 }
55
MakeSelectShape(const Shape & shape,int64_t dim,int64_t start,int64_t end,int64_t stride)56 Shape MakeSelectShape(
57 const Shape& shape,
58 int64_t dim,
59 int64_t start,
60 int64_t end,
61 int64_t stride) {
62 int64_t effective_stride = GetStride(start, end, stride);
63 Shape select_shape(shape);
64 select_shape.set_size(
65 dim, (end - start + effective_stride - 1) / effective_stride);
66 return select_shape;
67 }
68
GetStride(int64_t start,int64_t end,int64_t stride)69 int64_t GetStride(int64_t start, int64_t end, int64_t stride) {
70 if (stride == 0) {
71 TORCH_CHECK_EQ(start, end);
72 stride = 1;
73 }
74 return stride;
75 }
76
77 // This is almost like at::inferSqueezeGeometry, but that requires a Tensor
78 // input and also computes new strides. This logic seems correct.
BuildSqueezedDimensions(c10::ArrayRef<int64_t> dimensions,int64_t squeeze_dim)79 std::vector<int64_t> BuildSqueezedDimensions(
80 c10::ArrayRef<int64_t> dimensions,
81 int64_t squeeze_dim) {
82 std::vector<int64_t> output_dimensions;
83 for (const auto i : c10::irange(dimensions.size())) {
84 int64_t dim = dimensions[i];
85 if (dim != 1 ||
86 (static_cast<int64_t>(i) != squeeze_dim && squeeze_dim >= 0)) {
87 output_dimensions.push_back(dim);
88 }
89 }
90 return output_dimensions;
91 }
92
BuildUnsqueezedDimensions(c10::ArrayRef<int64_t> dimensions,int64_t squeeze_dim)93 std::vector<int64_t> BuildUnsqueezedDimensions(
94 c10::ArrayRef<int64_t> dimensions,
95 int64_t squeeze_dim) {
96 std::vector<int64_t> output_dimensions(
97 dimensions.cbegin(), dimensions.cend());
98 output_dimensions.insert(output_dimensions.begin() + squeeze_dim, 1);
99 return output_dimensions;
100 }
101
102 } // namespace lazy
103 } // namespace torch
104