|
1 | 1 | #include "torch_xla/csrc/ops/as_strided.h" |
2 | 2 |
|
3 | 3 | #include <algorithm> |
4 | | - |
5 | | -#include <torch/csrc/lazy/core/util.h> |
6 | | - |
| 4 | +#include <cstdint> |
| 5 | +#include <numeric> |
| 6 | +#include <string> |
| 7 | +#include <vector> |
| 8 | + |
| 9 | +#include <ATen/core/aten_interned_strings.h> |
| 10 | +#include <torch/csrc/lazy/core/hash.h> |
| 11 | +#include <torch/csrc/lazy/core/ir.h> |
| 12 | + |
| 13 | +#include "absl/status/status.h" |
| 14 | +#include "absl/status/statusor.h" |
| 15 | +#include "absl/strings/str_cat.h" |
| 16 | +#include "absl/strings/str_join.h" |
| 17 | +#include "absl/types/span.h" |
| 18 | +#include "xla/hlo/builder/xla_builder.h" |
| 19 | +#include "xla/permutation_util.h" |
| 20 | +#include "xla/shape.h" |
7 | 21 | #include "xla/shape_util.h" |
8 | | -#include "xla/util.h" |
9 | 22 |
|
10 | | -#include "torch_xla/csrc/data_ops.h" |
11 | 23 | #include "torch_xla/csrc/helpers.h" |
| 24 | +#include "torch_xla/csrc/ir.h" |
12 | 25 | #include "torch_xla/csrc/lowering_context.h" |
13 | | -#include "torch_xla/csrc/runtime/util.h" |
14 | 26 | #include "torch_xla/csrc/shape_helper.h" |
15 | | -#include "torch_xla/csrc/tensor_util.h" |
16 | | -#include "torch_xla/csrc/torch_util.h" |
| 27 | +#include "torch_xla/csrc/status.h" |
17 | 28 |
|
18 | 29 | namespace torch_xla { |
19 | 30 | namespace { |
20 | 31 |
|
21 | | -xla::XlaOp LowerAsStrided(xla::XlaOp input, absl::Span<const int64_t> size, |
22 | | - absl::Span<const int64_t> stride, |
23 | | - int64_t storage_offset) { |
24 | | - const xla::Shape& input_shape = ShapeHelper::ShapeOfXlaOp(input); |
25 | | - int64_t input_element_count = xla::ShapeUtil::ElementsIn(input_shape); |
26 | | - int64_t slice_size = torch_xla::runtime::util::Multiply<int64_t>(size); |
27 | | - XLA_CHECK_LE(storage_offset + slice_size, input_element_count); |
28 | | - |
29 | | - xla::XlaOp off_input = input; |
30 | | - if (storage_offset > 0 || slice_size < input_element_count) { |
31 | | - xla::XlaOp r1_input = XlaHelpers::Flatten(input); |
32 | | - off_input = xla::SliceInDim(r1_input, storage_offset, |
33 | | - storage_offset + slice_size, 1, 0); |
34 | | - } |
35 | | - |
36 | | - std::vector<int64_t> permutation = xla::InversePermutation( |
37 | | - AsStrided::GetArrayStridePermutation(stride, size)); |
38 | | - std::vector<int64_t> new_sizes = xla::PermuteInverse(size, permutation); |
39 | | - xla::XlaOp reshaped_input = XlaHelpers::DynamicReshape(off_input, new_sizes); |
40 | | - return xla::IsIdentityPermutation(permutation) |
41 | | - ? reshaped_input |
42 | | - : xla::Transpose(reshaped_input, permutation); |
| 32 | +xla::Shape AsStridedOutputShape(const torch::lazy::Value& input, |
| 33 | + absl::Span<const int64_t> size) { |
| 34 | + return xla::ShapeUtil::MakeShape(GetXlaShape(input).element_type(), size); |
43 | 35 | } |
44 | 36 |
|
45 | 37 | } // namespace |
46 | 38 |
|
47 | | -AsStrided::AsStrided(const torch::lazy::Value& input, std::vector<int64_t> size, |
48 | | - std::vector<int64_t> stride, int64_t storage_offset) |
| 39 | +AsStrided::AsStrided(const torch::lazy::Value& input, |
| 40 | + const std::vector<int64_t>& size, |
| 41 | + const std::vector<int64_t>& stride, int64_t storage_offset) |
49 | 42 | : XlaNode( |
50 | 43 | torch::lazy::OpKind(at::aten::as_strided), {input}, |
51 | | - [&]() { |
52 | | - return xla::ShapeUtil::MakeShape(GetXlaShape(input).element_type(), |
53 | | - size); |
54 | | - }, |
55 | | - /*num_outputs=*/1, torch::lazy::MHash(size, stride, storage_offset)), |
56 | | - size_(std::move(size)), |
57 | | - stride_(std::move(stride)), |
58 | | - storage_offset_(storage_offset) {} |
| 44 | + AsStridedOutputShape(input, size), |
| 45 | + /* num_outputs= */ 1, |
| 46 | + /* hash_seed= */ torch::lazy::MHash(size, stride, storage_offset)), |
| 47 | + size_(size), |
| 48 | + stride_(stride), |
| 49 | + storage_offset_(storage_offset) { |
| 50 | + // Make sure `input` has enough elements to fit the given spec. |
| 51 | + XLA_CHECK_OK(CheckSpecFitsInput(input)); |
| 52 | +} |
59 | 53 |
|
60 | 54 | std::string AsStrided::ToString() const { |
61 | | - std::stringstream ss; |
62 | | - ss << XlaNode::ToString() << ", size=(" << absl::StrJoin(size_, ", ") |
63 | | - << "), stride=(" << absl::StrJoin(stride_, ", ") |
64 | | - << "), storage_offset=" << storage_offset_; |
65 | | - return ss.str(); |
| 55 | + return absl::StrCat(XlaNode::ToString(), ", size=(", |
| 56 | + absl::StrJoin(size_, ", "), "), stride=(", |
| 57 | + absl::StrJoin(stride_, ", "), |
| 58 | + "), storage_offset=", storage_offset_); |
66 | 59 | } |
67 | 60 |
|
68 | 61 | torch::lazy::NodePtr AsStrided::Clone(torch::lazy::OpList operands) const { |
69 | 62 | return torch_xla::MakeNode<AsStrided>(operands. at( 0), size_, stride_, |
70 | 63 | storage_offset_); |
71 | 64 | } |
72 | 65 |
|
73 | | -XlaOpVector AsStrided::Lower(LoweringContext* loctx) const { |
74 | | - xla::XlaOp input = loctx->GetOutputOp(operand(0)); |
75 | | - return ReturnOp(LowerAsStrided(input, size_, stride_, storage_offset_), |
76 | | - loctx); |
| 66 | +absl::StatusOr<XlaOpVector> AsStrided::SafeLower(LoweringContext* loctx) const { |
| 67 | + XLA_ASSIGN_OR_RETURN(xla::XlaOp input, loctx->SafeGetOutputOp(operand(0))); |
| 68 | + |
| 69 | + XLA_ASSIGN_OR_RETURN(const xla::Shape* absl_nonnull input_shape_ptr, |
| 70 | + GetShape(input)); |
| 71 | + |
| 72 | + int64_t input_element_count = xla::ShapeUtil::ElementsIn(*input_shape_ptr); |
| 73 | + int64_t spec_element_count = GetSpecElementCount(); |
| 74 | + |
| 75 | + // Preprocess `input` so that it: |
| 76 | + // 1. Starts from `storage_offset_` |
| 77 | + // 2. Has the same element count as the spec |
| 78 | + // |
| 79 | + // This preprocessing should only be done if: |
| 80 | + // 1. There's actually a `storage_offset_` to start from; or |
| 81 | + // 2. The element count of the input is different from the spec |
| 82 | + if (storage_offset_ > 0 || input_element_count != spec_element_count) { |
| 83 | + XLA_ASSIGN_OR_RETURN(xla::XlaOp flattened, XlaHelpers::SafeFlatten(input)); |
| 84 | + input = xla::SliceInDim(flattened, storage_offset_, |
| 85 | + storage_offset_ + spec_element_count, 1, 0); |
| 86 | + } |
| 87 | + |
| 88 | + // Since PyTorch/XLA has no concept of strides in a tensor (i.e. all tensors |
| 89 | + // are contiguous), we need a way to compute a contiguous tensor that accesses |
| 90 | + // the same elements that a similarly spec'd strided tensor would. In order to |
| 91 | + // do that, we need to: |
| 92 | + // |
| 93 | + // 1. Reshape the `input`, so that the dimensions with larger strides come |
| 94 | + // first. This should yield the correct contiguous tensor, but with |
| 95 | + // permuted dimensions. |
| 96 | + std::vector<int64_t> permutation = |
| 97 | + xla::InversePermutation(GetDescendingOrderPermutation(stride_)); |
| 98 | + std::vector<int64_t> permuted_sizes = xla::PermuteInverse(size_, permutation); |
| 99 | + XLA_ASSIGN_OR_RETURN(xla::XlaOp input_reshaped_with_permuted_sizes, |
| 100 | + XlaHelpers::SafeDynamicReshape(input, permuted_sizes)); |
| 101 | + |
| 102 | + // 2. Reverse the dimension permutation we did on `input` in the previous |
| 103 | + // step. |
| 104 | + xla::XlaOp output = |
| 105 | + xla::IsIdentityPermutation(permutation) |
| 106 | + ? input_reshaped_with_permuted_sizes |
| 107 | + : xla::Transpose(input_reshaped_with_permuted_sizes, permutation); |
| 108 | + |
| 109 | + return ReturnOp(output, loctx); |
| 110 | +} |
| 111 | + |
| 112 | +int64_t AsStrided::GetSpecElementCount() const { |
| 113 | + return runtime::util::Multiply<int64_t>(size_); |
| 114 | +} |
| 115 | + |
| 116 | +absl::Status AsStrided::CheckSpecFitsInput(xla::XlaOp input) const { |
| 117 | + XLA_ASSIGN_OR_RETURN(const xla::Shape* absl_nonnull shape_ptr, |
| 118 | + GetShape(input)); |
| 119 | + XLA_RETURN_IF_ERROR( |
| 120 | + CheckSpecFitsInputImpl(*shape_ptr, xla::ShapeUtil::ElementsIn(*shape_ptr), |
| 121 | + GetSpecElementCount())); |
| 122 | + return absl::OkStatus(); |
| 123 | +} |
| 124 | + |
| 125 | +absl::Status AsStrided::CheckSpecFitsInput( |
| 126 | + const torch::lazy::Value& input) const { |
| 127 | + const xla::Shape& shape = GetXlaShape(input); |
| 128 | + XLA_RETURN_IF_ERROR(CheckSpecFitsInputImpl( |
| 129 | + shape, xla::ShapeUtil::ElementsIn(shape), GetSpecElementCount())); |
| 130 | + return absl::OkStatus(); |
| 131 | +} |
| 132 | + |
| 133 | +absl::Status AsStrided::CheckSpecFitsInputImpl( |
| 134 | + const xla::Shape& input_shape, int64_t input_element_count, |
| 135 | + int64_t spec_element_count) const { |
| 136 | + if (input_element_count < storage_offset_ + spec_element_count) { |
| 137 | + return XLA_ERROR_WITH_LOCATION(absl::InternalError(absl::StrCat( |
| 138 | + "as_strided(): expected input ", input_shape.ToString(), |
| 139 | + " (elements=", input_element_count, |
| 140 | + ") to have enough elements to fit the given spec of size=[", |
| 141 | + absl::StrJoin(size_, /* separator= */ ", "), "], stride=[", |
| 142 | + absl::StrJoin(stride_, /* separator= */ ", "), "], and storage_offset=", |
| 143 | + storage_offset_, " (elements=", spec_element_count, ")"))); |
| 144 | + } |
| 145 | + return absl::OkStatus(); |
77 | 146 | } |
78 | 147 |
|
79 | | -bool AsStrided::StrideIsSupported(const xla::Shape& input_shape, |
80 | | - absl::Span<const int64_t> size, |
81 | | - absl::Span<const int64_t> stride, |
82 | | - int64_t storage_offset) { |
| 148 | +bool IsAsStridedWithStrideSupported(absl::Span<const int64_t> stride) { |
83 | 149 | std::vector<int64_t> sorted_stride(stride.begin(), stride.end()); |
84 | 150 | std::sort(sorted_stride.begin(), sorted_stride.end()); |
85 | 151 | return stride.empty() || sorted_stride.front() == 1; |
86 | 152 | } |
87 | 153 |
|
88 | | -std::vector<int64_t> AsStrided::GetArrayStridePermutation( |
89 | | - absl::Span<const int64_t> stride, absl::Span<const int64_t> size) { |
90 | | - std::vector<int64_t> permutation = torch::lazy::Iota<int64_t>(stride.size()); |
| 154 | +std::vector<int64_t> GetDescendingOrderPermutation( |
| 155 | + absl::Span<const int64_t> v) { |
| 156 | + std::vector<int64_t> permutation(v.size()); |
| 157 | + std::iota(permutation.begin(), permutation.end(), 0); |
91 | 158 | std::sort(permutation.begin(), permutation.end(), |
92 | | - [&](int64_t a, int64_t b) { return stride[a] > stride[b]; }); |
| 159 | + [&](int64_t a, int64_t b) { return v[a] > v[b]; }); |
93 | 160 | return permutation; |
94 | 161 | } |
95 | 162 |
|
|
0 commit comments