41 #ifdef BUILD_ISPC_MODULE
53 class IndexerIterator;
56 static constexpr int64_t MAX_DIMS = 10;
60 static constexpr int64_t MAX_INPUTS = 10;
64 static constexpr int64_t MAX_OUTPUTS = 2;
66 template <
int NARGS,
typename index_t = u
int32_t>
70 const int64_t*
const* strides)
72 if (
dims_ > MAX_DIMS) {
76 for (
int i = 0; i < MAX_DIMS; ++i) {
82 for (
int arg = 0; arg < NARGS; arg++) {
91 #if defined(__CUDA_ARCH__)
94 for (
int arg = 0; arg < NARGS; arg++) {
98 #if defined(__CUDA_ARCH__)
101 for (
int dim = 0; dim < MAX_DIMS; ++dim) {
106 linear_idx = linear_idx /
sizes_[dim];
108 #if defined(__CUDA_ARCH__)
111 for (
int arg = 0; arg < NARGS; arg++) {
112 offsets[arg] += mod *
strides_[dim][arg];
138 for (int64_t i = 0; i <
ndims_; ++i) {
152 if (
static_cast<int64_t
>(dims.size()) !=
ndims_) {
156 std::vector<bool> seen_dims(
ndims_,
false);
157 for (
const int64_t& dim : dims) {
158 seen_dims[dim] =
true;
160 if (!std::all_of(seen_dims.begin(), seen_dims.end(),
161 [](
bool seen) { return seen; })) {
163 "Permute dims must be a permuntation from 0 to {}.",
170 for (int64_t i = 0; i <
ndims_; ++i) {
172 new_shape[i] =
shape_[old_dim];
175 for (int64_t i = 0; i <
ndims_; ++i) {
185 for (int64_t i = 0; i <
ndims_; ++i) {
197 for (int64_t i = 0; i <
ndims_; ++i) {
206 #ifdef BUILD_ISPC_MODULE
208 ispc::TensorRef ToISPC()
const;
246 int64_t num_workloads = 1;
247 for (int64_t i = 0; i <
ndims_; ++i) {
250 return num_workloads;
254 if (workload_idx < 0 || workload_idx >=
NumWorkloads()) {
259 for (int64_t i = 0; i <
ndims_; ++i) {
289 Indexer(
const std::vector<Tensor>& input_tensors,
290 const Tensor& output_tensor,
294 Indexer(
const std::vector<Tensor>& input_tensors,
295 const std::vector<Tensor>& output_tensors,
425 int64_t workload_idx)
const {
441 template <
typename T>
443 int64_t workload_idx)
const {
447 return GetWorkloadDataPtr<T>(
inputs_[input_idx],
468 template <
typename T>
480 int64_t workload_idx)
const {
491 template <
typename T>
493 int64_t workload_idx)
const {
494 return GetWorkloadDataPtr<T>(
outputs_[output_idx],
499 #ifdef BUILD_ISPC_MODULE
501 ispc::Indexer ToISPC()
const;
548 const int64_t* dst_shape);
554 const int64_t* src_shape,
562 int64_t workload_idx)
const {
565 if (workload_idx < 0) {
569 return static_cast<char*
>(tr.
data_ptr_) +
573 for (int64_t i = 0; i <
ndims_; ++i) {
588 template <
typename T>
591 int64_t workload_idx)
const {
594 if (workload_idx < 0) {
598 return static_cast<T*
>(tr.
data_ptr_) + workload_idx;
601 for (int64_t i = 0; i <
ndims_; ++i) {
606 return static_cast<T*
>(
static_cast<void*
>(
669 std::vector<std::unique_ptr<Indexer>>
vec_;
#define OPEN3D_HOST_DEVICE
Definition: CUDAUtils.h:63
#define LogError(...)
Definition: Logging.h:67
int64_t ByteSize() const
Definition: Dtype.h:77
Definition: Indexer.h:280
void UpdateContiguousFlags()
Update input_contiguous_ and output_contiguous_.
Definition: Indexer.cpp:583
bool inputs_contiguous_[MAX_INPUTS]
Array of contiguous flags for all input TensorRefs.
Definition: Indexer.h:622
bool outputs_contiguous_[MAX_OUTPUTS]
Array of contiguous flags for all output TensorRefs.
Definition: Indexer.h:625
OPEN3D_HOST_DEVICE T * GetOutputPtr(int64_t output_idx, int64_t workload_idx) const
Definition: Indexer.h:492
Indexer()
Definition: Indexer.h:282
int64_t num_outputs_
Definition: Indexer.h:613
OPEN3D_HOST_DEVICE char * GetWorkloadDataPtr(const TensorRef &tr, bool tr_contiguous, int64_t workload_idx) const
Definition: Indexer.h:560
static void ReductionRestride(TensorRef &dst, int64_t src_ndims, const int64_t *src_shape, const SizeVector &reduction_dims)
Definition: Indexer.cpp:620
TensorRef outputs_[MAX_OUTPUTS]
Array of output TensorRefs.
Definition: Indexer.h:619
bool IsReductionDim(int64_t dim) const
Returns true if the dim -th dimension is reduced.
Definition: Indexer.h:412
const TensorRef & GetOutput() const
Definition: Indexer.h:403
OPEN3D_HOST_DEVICE T * GetInputPtr(int64_t input_idx, int64_t workload_idx) const
Definition: Indexer.h:442
int64_t NumReductionDims() const
Returns the number of reudction dimensions.
Definition: Indexer.cpp:413
int64_t NumInputs() const
Number of input Tensors.
Definition: Indexer.h:357
OPEN3D_HOST_DEVICE char * GetOutputPtr(int64_t workload_idx) const
Definition: Indexer.h:456
Indexer(const Indexer &)=default
bool IsFinalOutput() const
Definition: Indexer.h:317
const TensorRef & GetOutput(int64_t i) const
Definition: Indexer.h:386
void ReorderDimensions(const SizeVector &reduction_dims)
Definition: Indexer.cpp:509
void CoalesceDimensions()
Definition: Indexer.cpp:443
int64_t master_shape_[MAX_DIMS]
Definition: Indexer.h:638
int64_t NumOutputElements() const
Returns the number of output elements.
Definition: Indexer.cpp:432
int64_t num_inputs_
Number of input and output Tensors.
Definition: Indexer.h:612
bool accumulate_
Definition: Indexer.h:654
bool CanUse32BitIndexing() const
Returns true iff the maximum_offsets in bytes are smaller than 2^31 - 1.
Definition: Indexer.cpp:216
TensorRef inputs_[MAX_INPUTS]
Array of input TensorRefs.
Definition: Indexer.h:616
TensorRef & GetOutput(int64_t i)
Returns output TensorRef.
Definition: Indexer.h:379
bool ShouldAccumulate() const
Definition: Indexer.h:315
const TensorRef & GetInput(int64_t i) const
Definition: Indexer.h:370
int64_t master_strides_[MAX_DIMS]
Definition: Indexer.h:642
Indexer GetPerOutputIndexer(int64_t output_idx) const
Definition: Indexer.cpp:321
int64_t * GetMasterShape()
Definition: Indexer.h:335
Indexer & operator=(const Indexer &)=default
const int64_t * GetMasterShape() const
Definition: Indexer.h:334
std::unique_ptr< Indexer > SplitLargestDim()
Definition: Indexer.cpp:256
int64_t NumWorkloads() const
Definition: Indexer.cpp:424
IndexerIterator SplitTo32BitIndexing() const
Definition: Indexer.cpp:252
OPEN3D_HOST_DEVICE T * GetWorkloadDataPtr(const TensorRef &tr, bool tr_contiguous, int64_t workload_idx) const
Definition: Indexer.h:589
void UpdateMasterStrides()
Update master_strides_ based on master_shape_.
Definition: Indexer.cpp:574
static void BroadcastRestride(TensorRef &src, int64_t dst_ndims, const int64_t *dst_shape)
Definition: Indexer.cpp:593
bool final_output_
Definition: Indexer.h:650
OPEN3D_HOST_DEVICE char * GetOutputPtr(int64_t output_idx, int64_t workload_idx) const
Definition: Indexer.h:479
OPEN3D_HOST_DEVICE char * GetInputPtr(int64_t input_idx, int64_t workload_idx) const
Definition: Indexer.h:424
TensorRef & GetInput(int64_t i)
Returns input TensorRef.
Definition: Indexer.h:363
TensorRef & GetOutput()
Definition: Indexer.h:396
OPEN3D_HOST_DEVICE T * GetOutputPtr(int64_t workload_idx) const
Definition: Indexer.h:469
int64_t ndims_
Indexer's global number of dimensions.
Definition: Indexer.h:645
void ShrinkDim(int64_t dim, int64_t start, int64_t size)
Definition: Indexer.cpp:382
int64_t NumDims() const
Returns number of dimensions of the Indexer.
Definition: Indexer.h:330
int64_t NumOutputs() const
Number of output Tensors.
Definition: Indexer.h:360
const int64_t * GetMasterStrides() const
Definition: Indexer.h:339
Definition: Indexer.h:657
Iterator end() const
Definition: Indexer.cpp:689
Iterator begin() const
Definition: Indexer.cpp:685
IndexerIterator(const Indexer &indexer)
Definition: Indexer.cpp:659
Definition: SizeVector.h:79
T * GetDataPtr()
Definition: Tensor.h:1108
SizeVector GetShape() const
Definition: Tensor.h:1091
int64_t NumDims() const
Definition: Tensor.h:1136
int64_t GetStride(int64_t dim) const
Definition: Tensor.h:1103
Dtype GetDtype() const
Definition: Tensor.h:1128
Definition: Indexer.h:240
OPEN3D_HOST_DEVICE void * GetPtr(int64_t workload_idx) const
Definition: Indexer.h:253
OPEN3D_HOST_DEVICE int64_t NumWorkloads() const
Definition: Indexer.h:245
TensorRef input_
Definition: Indexer.h:269
TensorIterator(const Tensor &tensor)
Definition: Indexer.h:242
int64_t ndims_
Definition: Indexer.h:270
int32_t index_t
Definition: NanoFlannImpl.h:43
int64_t WrapDim(int64_t dim, int64_t max_dim, bool inclusive)
Wrap around negative dim.
Definition: ShapeUtil.cpp:150
SizeVector DefaultStrides(const SizeVector &shape)
Compute default strides for a shape when a tensor is contiguous.
Definition: ShapeUtil.cpp:233
DtypePolicy
Definition: Indexer.h:218
Definition: PinholeCameraIntrinsic.cpp:35
Definition: Indexer.h:659
bool operator!=(const Iterator &other) const
Definition: Indexer.cpp:681
Iterator()
Definition: Indexer.h:660
std::vector< std::unique_ptr< Indexer > > vec_
Definition: Indexer.h:669
Iterator(Iterator &&other)=default
Indexer & operator*() const
Definition: Indexer.cpp:667
bool operator==(const Iterator &other) const
Definition: Indexer.cpp:678
Iterator & operator++()
Definition: Indexer.cpp:669
index_t sizes_[MAX_DIMS]
Definition: Indexer.h:119
int dims_
Definition: Indexer.h:118
OPEN3D_HOST_DEVICE utility::MiniVec< index_t, NARGS > get(index_t linear_idx) const
Definition: Indexer.h:88
OffsetCalculator(int dims, const int64_t *sizes, const int64_t *const *strides)
Definition: Indexer.h:68
index_t strides_[MAX_DIMS][NARGS]
Definition: Indexer.h:120
A minimalistic class that reference a Tensor.
Definition: Indexer.h:124
int64_t dtype_byte_size_
Definition: Indexer.h:213
void Permute(const SizeVector &dims)
Permute (dimension shuffle) the reference to a Tensor.
Definition: Indexer.h:150
int64_t ndims_
Definition: Indexer.h:212
bool operator!=(const TensorRef &other) const
Definition: Indexer.h:204
TensorRef(const Tensor &t)
Definition: Indexer.h:130
TensorRef()
Definition: Indexer.h:128
int64_t shape_[MAX_DIMS]
Definition: Indexer.h:214
bool IsContiguous() const
Returns True if the underlying memory buffer is contiguous.
Definition: Indexer.h:182
void * data_ptr_
Definition: Indexer.h:211
bool operator==(const TensorRef &other) const
Definition: Indexer.h:192
int64_t byte_strides_[MAX_DIMS]
Definition: Indexer.h:215