Open3D (C++ API)  0.18.0
AdvancedIndexing.h
Go to the documentation of this file.
1 // ----------------------------------------------------------------------------
2 // - Open3D: www.open3d.org -
3 // ----------------------------------------------------------------------------
4 // Copyright (c) 2018-2023 www.open3d.org
5 // SPDX-License-Identifier: MIT
6 // ----------------------------------------------------------------------------
7 
8 #pragma once
9 
10 #include <vector>
11 
12 #include "open3d/core/Indexer.h"
13 #include "open3d/core/SizeVector.h"
14 #include "open3d/core/Tensor.h"
15 
16 namespace open3d {
17 namespace core {
18 
21 public:
23  const std::vector<Tensor>& index_tensors)
24  : tensor_(tensor), index_tensors_(ExpandBoolTensors(index_tensors)) {
25  RunPreprocess();
26  }
27 
28  inline Tensor GetTensor() const { return tensor_; }
29 
30  inline std::vector<Tensor> GetIndexTensors() const {
31  return index_tensors_;
32  }
33 
34  inline SizeVector GetOutputShape() const { return output_shape_; }
35 
36  inline SizeVector GetIndexedShape() const { return indexed_shape_; }
37 
38  inline SizeVector GetIndexedStrides() const { return indexed_strides_; }
39 
43  static bool IsIndexSplittedBySlice(
44  const std::vector<Tensor>& index_tensors);
45 
48  static std::pair<Tensor, std::vector<Tensor>> ShuffleIndexedDimsToFront(
49  const Tensor& tensor, const std::vector<Tensor>& index_tensors);
50 
53  static std::pair<std::vector<Tensor>, SizeVector>
54  ExpandToCommonShapeExceptZeroDim(const std::vector<Tensor>& index_tensors);
55 
56  // Replace indexed dimensions with stride 0 and the size of the result
57  // tensor.
58  //
59  // The offset in these dimensions is computed by the kernel using
60  // the index tensor's values and the stride of the tensor. The new shape is
61  // not meaningful. It's used to make the shape compatible with the result
62  // tensor.
63  //
64  // Effectively, we throw away the tensor's shape and strides for the sole
65  // purpose of element-wise iteration for the Indexer. The tensor's original
66  // strides are stored in indexed_shape_ and indexed_strides_,
67  // which are passed to fancy indexing kernels.
68  static Tensor RestrideTensor(const Tensor& tensor,
69  int64_t dims_before,
70  int64_t dims_indexed,
71  SizeVector replacement_shape);
72 
73  // Add dimensions of size 1 to an index tensor so that it can be broadcast
74  // to the result shape and iterated over element-wise like the result tensor
75  // and the restrided src.
76  static Tensor RestrideIndexTensor(const Tensor& index_tensor,
77  int64_t dims_before,
78  int64_t dims_after);
79 
80 protected:
82  void RunPreprocess();
83 
85  static std::vector<Tensor> ExpandBoolTensors(
86  const std::vector<Tensor>& index_tensors);
87 
91 
93  std::vector<Tensor> index_tensors_;
94 
97 
101 
105 };
106 
117 public:
118  enum class AdvancedIndexerMode { SET, GET };
119 
121  const Tensor& dst,
122  const std::vector<Tensor>& index_tensors,
123  const SizeVector& indexed_shape,
124  const SizeVector& indexed_strides,
126  : mode_(mode) {
127  if (indexed_shape.size() != indexed_strides.size()) {
129  "Internal error: indexed_shape's ndim {} does not equal to "
130  "indexed_strides' ndim {}",
131  indexed_shape.size(), indexed_strides.size());
132  }
133  num_indices_ = indexed_shape.size();
134 
135  // Initialize Indexer
136  std::vector<Tensor> inputs;
137  inputs.push_back(src);
138  for (const Tensor& index_tensor : index_tensors) {
139  if (index_tensor.NumDims() != 0) {
140  inputs.push_back(index_tensor);
141  }
142  }
143  indexer_ = Indexer({inputs}, dst, DtypePolicy::NONE);
144 
145  // Fill shape and strides
146  if (num_indices_ != static_cast<int64_t>(indexed_strides.size())) {
148  "Internal error: indexed_shape's ndim {} does not equal to "
149  "indexd_strides' ndim {}",
150  num_indices_, indexed_strides.size());
151  }
152  for (int64_t i = 0; i < num_indices_; ++i) {
153  indexed_shape_[i] = indexed_shape[i];
154  indexed_strides_[i] = indexed_strides[i];
155  }
156 
157  // Check dtypes
158  if (src.GetDtype() != dst.GetDtype()) {
160  "src's dtype {} is not the same as dst's dtype {}.",
161  src.GetDtype().ToString(), dst.GetDtype().ToString());
162  }
164  }
165 
166  inline OPEN3D_HOST_DEVICE char* GetInputPtr(int64_t workload_idx) const {
167  char* ptr = indexer_.GetInputPtr(0, workload_idx);
168  ptr += GetIndexedOffset(workload_idx) * element_byte_size_ *
170  return ptr;
171  }
172 
173  inline OPEN3D_HOST_DEVICE char* GetOutputPtr(int64_t workload_idx) const {
174  char* ptr = indexer_.GetOutputPtr(workload_idx);
175  ptr += GetIndexedOffset(workload_idx) * element_byte_size_ *
177  return ptr;
178  }
179 
180  inline OPEN3D_HOST_DEVICE int64_t
181  GetIndexedOffset(int64_t workload_idx) const {
182  int64_t offset = 0;
183  for (int64_t i = 0; i < num_indices_; ++i) {
184  int64_t index = *(reinterpret_cast<int64_t*>(
185  indexer_.GetInputPtr(i + 1, workload_idx)));
186  OPEN3D_ASSERT(index >= -indexed_shape_[i] &&
187  index < indexed_shape_[i] && "Index out of bounds.");
188  index += indexed_shape_[i] * (index < 0);
189  offset += index * indexed_strides_[i];
190  }
191  return offset;
192  }
193 
194  int64_t NumWorkloads() const { return indexer_.NumWorkloads(); }
195 
196 protected:
199  int64_t num_indices_;
201  int64_t indexed_shape_[MAX_DIMS];
202  int64_t indexed_strides_[MAX_DIMS];
203 };
204 
205 } // namespace core
206 } // namespace open3d
#define OPEN3D_HOST_DEVICE
Definition: CUDAUtils.h:44
#define LogError(...)
Definition: Logging.h:48
#define OPEN3D_ASSERT(...)
Definition: Macro.h:48
This class is based on PyTorch's aten/src/ATen/native/Indexing.cpp.
Definition: AdvancedIndexing.h:20
void RunPreprocess()
Preprocess tensor and index tensors.
Definition: AdvancedIndexing.cpp:110
static Tensor RestrideTensor(const Tensor &tensor, int64_t dims_before, int64_t dims_indexed, SizeVector replacement_shape)
Definition: AdvancedIndexing.cpp:85
static bool IsIndexSplittedBySlice(const std::vector< Tensor > &index_tensors)
Definition: AdvancedIndexing.cpp:17
std::vector< Tensor > GetIndexTensors() const
Definition: AdvancedIndexing.h:30
static std::pair< Tensor, std::vector< Tensor > > ShuffleIndexedDimsToFront(const Tensor &tensor, const std::vector< Tensor > &index_tensors)
Definition: AdvancedIndexing.cpp:41
SizeVector output_shape_
Output shape.
Definition: AdvancedIndexing.h:96
SizeVector GetIndexedStrides() const
Definition: AdvancedIndexing.h:38
std::vector< Tensor > index_tensors_
The processed index tensors.
Definition: AdvancedIndexing.h:93
Tensor tensor_
Definition: AdvancedIndexing.h:90
static std::vector< Tensor > ExpandBoolTensors(const std::vector< Tensor > &index_tensors)
Expand boolean tensor to integer index.
Definition: AdvancedIndexing.cpp:230
static std::pair< std::vector< Tensor >, SizeVector > ExpandToCommonShapeExceptZeroDim(const std::vector< Tensor > &index_tensors)
Definition: AdvancedIndexing.cpp:63
AdvancedIndexPreprocessor(const Tensor &tensor, const std::vector< Tensor > &index_tensors)
Definition: AdvancedIndexing.h:22
static Tensor RestrideIndexTensor(const Tensor &index_tensor, int64_t dims_before, int64_t dims_after)
Definition: AdvancedIndexing.cpp:100
SizeVector indexed_shape_
Definition: AdvancedIndexing.h:100
SizeVector GetIndexedShape() const
Definition: AdvancedIndexing.h:36
SizeVector GetOutputShape() const
Definition: AdvancedIndexing.h:34
Tensor GetTensor() const
Definition: AdvancedIndexing.h:28
SizeVector indexed_strides_
Definition: AdvancedIndexing.h:104
Definition: AdvancedIndexing.h:116
Indexer indexer_
Definition: AdvancedIndexing.h:197
AdvancedIndexer(const Tensor &src, const Tensor &dst, const std::vector< Tensor > &index_tensors, const SizeVector &indexed_shape, const SizeVector &indexed_strides, AdvancedIndexerMode mode)
Definition: AdvancedIndexing.h:120
int64_t element_byte_size_
Definition: AdvancedIndexing.h:200
int64_t NumWorkloads() const
Definition: AdvancedIndexing.h:194
int64_t num_indices_
Definition: AdvancedIndexing.h:199
OPEN3D_HOST_DEVICE char * GetOutputPtr(int64_t workload_idx) const
Definition: AdvancedIndexing.h:173
int64_t indexed_strides_[MAX_DIMS]
Definition: AdvancedIndexing.h:202
AdvancedIndexerMode mode_
Definition: AdvancedIndexing.h:198
OPEN3D_HOST_DEVICE char * GetInputPtr(int64_t workload_idx) const
Definition: AdvancedIndexing.h:166
int64_t indexed_shape_[MAX_DIMS]
Definition: AdvancedIndexing.h:201
OPEN3D_HOST_DEVICE int64_t GetIndexedOffset(int64_t workload_idx) const
Definition: AdvancedIndexing.h:181
AdvancedIndexerMode
Definition: AdvancedIndexing.h:118
std::string ToString() const
Definition: Dtype.h:64
int64_t ByteSize() const
Definition: Dtype.h:58
Definition: Indexer.h:261
OPEN3D_HOST_DEVICE char * GetOutputPtr(int64_t workload_idx) const
Definition: Indexer.h:437
int64_t NumWorkloads() const
Definition: Indexer.cpp:406
OPEN3D_HOST_DEVICE char * GetInputPtr(int64_t input_idx, int64_t workload_idx) const
Definition: Indexer.h:405
Definition: SizeVector.h:69
size_t size() const
Definition: SmallVector.h:119
Definition: Tensor.h:32
Dtype GetDtype() const
Definition: Tensor.h:1163
int offset
Definition: FilePCD.cpp:45
const char const char value recording_handle imu_sample recording_handle uint8_t size_t data_size k4a_record_configuration_t config target_format k4a_capture_t capture_handle k4a_imu_sample_t imu_sample playback_handle k4a_logging_message_cb_t void min_level device_handle k4a_imu_sample_t timeout_in_ms capture_handle capture_handle capture_handle image_handle temperature_c k4a_image_t image_handle uint8_t image_handle image_handle image_handle image_handle image_handle timestamp_usec white_balance image_handle k4a_device_configuration_t config device_handle char size_t serial_number_size bool int32_t int32_t int32_t int32_t k4a_color_control_mode_t default_mode mode
Definition: K4aPlugin.cpp:678
Definition: PinholeCameraIntrinsic.cpp:16