Indexer.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2026 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#pragma once
 
#include <sstream>
 
#include "open3d/core/CUDAUtils.h"
#include "open3d/core/Dtype.h"
#include "open3d/core/ShapeUtil.h"
#include "open3d/core/SizeVector.h"
#include "open3d/core/Tensor.h"
#include "open3d/utility/Logging.h"
#include "open3d/utility/MiniVec.h"
 
// The generated "Indexer_ispc.h" header will not be available outside the
// library. Therefore, forward declare all exported ISPC classes.
#ifdef BUILD_ISPC_MODULE
namespace ispc {
struct TensorRef;
struct Indexer;
}  // namespace ispc
#endif
 
namespace open3d {
namespace core {
 
class Indexer;
 
class IndexerIterator;
 
// Maximum number of dimensions of TensorRef.
static constexpr int64_t MAX_DIMS = 5;
 
// Maximum number of inputs of an op.
// MAX_INPUTS shall be >= MAX_DIMS to support advanced indexing.
static constexpr int64_t MAX_INPUTS = 5;
 
// Maximum number of outputs of an op. This number can be increased when
// necessary.
static constexpr int64_t MAX_OUTPUTS = 2;
 
template <int NARGS, typename index_t = uint32_t>
struct OffsetCalculator {
    OffsetCalculator(int dims,
                     const int64_t* sizes,
                     const int64_t* const* strides)
        : dims_(dims) {
        if (dims_ > MAX_DIMS) {
            utility::LogError("tensor has too many (>{}) dims_", MAX_DIMS);
        }
 
        for (int i = 0; i < MAX_DIMS; ++i) {
            if (i < dims_) {
                sizes_[i] = sizes[i];
            } else {
                sizes_[i] = 1;
            }
            for (int arg = 0; arg < NARGS; arg++) {
                strides_[i][arg] = i < dims_ ? strides[arg][i] : 0;
            }
        }
    }
 
    OPEN3D_HOST_DEVICE utility::MiniVec<index_t, NARGS> get(
            index_t linear_idx) const {
        utility::MiniVec<index_t, NARGS> offsets;
#if defined(__CUDA_ARCH__)
#pragma unroll
#endif
        for (int arg = 0; arg < NARGS; arg++) {
            offsets[arg] = 0;
        }
 
#if defined(__CUDA_ARCH__)
#pragma unroll
#endif
        for (int dim = 0; dim < MAX_DIMS; ++dim) {
            if (dim == dims_) {
                break;
            }
            index_t mod = linear_idx % sizes_[dim];
            linear_idx = linear_idx / sizes_[dim];
 
#if defined(__CUDA_ARCH__)
#pragma unroll
#endif
            for (int arg = 0; arg < NARGS; arg++) {
                offsets[arg] += mod * strides_[dim][arg];
            }
        }
        return offsets;
    }
 
    int dims_;
    index_t sizes_[MAX_DIMS];
    index_t strides_[MAX_DIMS][NARGS];
};
 
struct TensorRef {
    // The default copy constructor works on __device__ as well so we don't
    // define it explicitly. shape_[MAX_DIMS] and strides[MAX_DIMS] will be
    // copied fully.
    TensorRef() : data_ptr_(nullptr) {}
 
    TensorRef(const Tensor& t) {
        if (t.NumDims() > MAX_DIMS) {
            utility::LogError("Tensor has too many dimensions {} > {}.",
                              t.NumDims(), MAX_DIMS);
        }
        data_ptr_ = const_cast<void*>(t.GetDataPtr());
        ndims_ = t.NumDims();
        dtype_byte_size_ = t.GetDtype().ByteSize();
        for (int64_t i = 0; i < ndims_; ++i) {
            shape_[i] = t.GetShape(i);
            byte_strides_[i] = t.GetStride(i) * dtype_byte_size_;
        }
        UpdateByteOffsetBounds();
    }
 
    void UpdateByteOffsetBounds() {
        int64_t num_elements = 1;
        for (int64_t i = 0; i < ndims_; ++i) {
            num_elements *= shape_[i];
        }
        min_byte_offset_ = 0;
        total_byte_size_ = 0;
        if (num_elements == 0) {
            return;
        }
        int64_t min_o = 0;
        int64_t max_o = 0;
        for (int64_t d = 0; d < ndims_; ++d) {
            const int64_t max_coord = shape_[d] - 1;
            const int64_t stride = byte_strides_[d];
            const int64_t dim_min = stride >= 0 ? 0 : max_coord * stride;
            const int64_t dim_max = stride >= 0 ? max_coord * stride : 0;
            min_o += dim_min;
            max_o += dim_max;
        }
        min_byte_offset_ = min_o;
        total_byte_size_ = max_o - min_o + dtype_byte_size_;
    }
 
    OPEN3D_HOST_DEVICE bool ContainsByteOffset(int64_t offset) const {
        if (total_byte_size_ == 0) {
            return offset == min_byte_offset_;
        }
        return offset >= min_byte_offset_ &&
               offset + dtype_byte_size_ <= min_byte_offset_ + total_byte_size_;
    }
 
    void Permute(const SizeVector& dims) {
        // Check dims are permuntation of [0, 1, 2, ..., n-1]
        if (static_cast<int64_t>(dims.size()) != ndims_) {
            utility::LogError("Number of dimensions mismatch {} != {}.",
                              dims.size(), ndims_);
        }
        std::vector<bool> seen_dims(ndims_, false);
        for (const int64_t& dim : dims) {
            seen_dims[dim] = true;
        }
        if (!std::all_of(seen_dims.begin(), seen_dims.end(),
                         [](bool seen) { return seen; })) {
            utility::LogError(
                    "Permute dims must be a permuntation from 0 to {}.",
                    dims.size() - 1);
        }
 
        // Map to new shape and strides
        SizeVector new_shape(ndims_);
        SizeVector new_byte_strides(ndims_);
        for (int64_t i = 0; i < ndims_; ++i) {
            int64_t old_dim = shape_util::WrapDim(dims[i], ndims_);
            new_shape[i] = shape_[old_dim];
            new_byte_strides[i] = byte_strides_[old_dim];
        }
        for (int64_t i = 0; i < ndims_; ++i) {
            shape_[i] = new_shape[i];
            byte_strides_[i] = new_byte_strides[i];
        }
    }
 
    inline bool IsContiguous() const {
        SizeVector shape(ndims_);
        SizeVector strides(ndims_);
        for (int64_t i = 0; i < ndims_; ++i) {
            shape[i] = shape_[i];
            strides[i] = byte_strides_[i] / dtype_byte_size_;
        }
        return shape_util::DefaultStrides(shape) == strides;
    }
 
    bool operator==(const TensorRef& other) const {
        bool rc = true;
        rc = rc && (data_ptr_ == other.data_ptr_);
        rc = rc && (ndims_ == other.ndims_);
        rc = rc && (dtype_byte_size_ == other.dtype_byte_size_);
        rc = rc && (min_byte_offset_ == other.min_byte_offset_);
        rc = rc && (total_byte_size_ == other.total_byte_size_);
        for (int64_t i = 0; i < ndims_; ++i) {
            rc = rc && (shape_[i] == other.shape_[i]);
            rc = rc && (byte_strides_[i] == other.byte_strides_[i]);
        }
        return rc;
    }
 
    bool operator!=(const TensorRef& other) const { return !(*this == other); }
 
#ifdef BUILD_ISPC_MODULE
    ispc::TensorRef ToISPC() const;
#endif
 
    void* data_ptr_;
    int64_t ndims_ = 0;
    int64_t dtype_byte_size_ = 0;
    int64_t min_byte_offset_ = 0;
    int64_t total_byte_size_ = 0;
    int64_t shape_[MAX_DIMS];
    int64_t byte_strides_[MAX_DIMS];
};
 
enum class DtypePolicy {
    NONE,        // Do not check. Expects the kernel to handle the conversion.
                 // E.g. in Copy kernel with type casting.
    ALL_SAME,    // All inputs and outputs to to have the same dtype.
    INPUT_SAME,  // All inputs have the same dtype.
    INPUT_SAME_OUTPUT_BOOL  // All inputs have the same dtype. Outputs
                            // have bool dtype.
};
 
class TensorIterator {
public:
    TensorIterator(const Tensor& tensor)
        : input_(TensorRef(tensor)), ndims_(tensor.NumDims()) {
        is_contiguous_ = true;
        int64_t expected_byte_stride = input_.dtype_byte_size_;
        for (int64_t d = ndims_ - 1; d >= 0; --d) {
            if (input_.byte_strides_[d] != expected_byte_stride) {
                is_contiguous_ = false;
                break;
            }
            expected_byte_stride *= input_.shape_[d];
        }
    }
 
    OPEN3D_HOST_DEVICE int64_t NumWorkloads() const {
        int64_t num_workloads = 1;
        for (int64_t i = 0; i < ndims_; ++i) {
            num_workloads *= input_.shape_[i];
        }
        return num_workloads;
    }
 
    OPEN3D_HOST_DEVICE void* GetPtr(int64_t workload_idx) const {
        if (workload_idx < 0 || workload_idx >= NumWorkloads()) {
            return nullptr;
        }
        int64_t offset;
        if (is_contiguous_) {
            offset = workload_idx * input_.dtype_byte_size_;
        } else {
            offset = 0;
            int64_t remaining = workload_idx;
            for (int64_t d = ndims_ - 1; d >= 0; --d) {
                const int64_t coord = remaining % input_.shape_[d];
                remaining /= input_.shape_[d];
                offset += coord * input_.byte_strides_[d];
            }
        }
        OPEN3D_ASSERT(input_.ContainsByteOffset(offset),
                      "Index operation data pointer is out of range.");
        return static_cast<void*>(static_cast<char*>(input_.data_ptr_) +
                                  offset);
    }
 
protected:
    TensorRef input_;
    int64_t ndims_;
    bool is_contiguous_;
};
 
class Indexer {
public:
    Indexer() {}
    Indexer(const Indexer&) = default;
    Indexer& operator=(const Indexer&) = default;
 
    Indexer(const std::vector<Tensor>& input_tensors,
            const Tensor& output_tensor,
            DtypePolicy dtype_policy = DtypePolicy::ALL_SAME,
            const SizeVector& reduction_dims = {});
 
    Indexer(const std::vector<Tensor>& input_tensors,
            const std::vector<Tensor>& output_tensors,
            DtypePolicy dtype_policy = DtypePolicy::ALL_SAME,
            const SizeVector& reduction_dims = {});
 
    bool CanUse32BitIndexing() const;
 
    IndexerIterator SplitTo32BitIndexing() const;
 
    std::unique_ptr<Indexer> SplitLargestDim();
 
    Indexer GetPerOutputIndexer(int64_t output_idx) const;
 
    bool ShouldAccumulate() const { return accumulate_; }
 
    bool IsFinalOutput() const { return final_output_; }
 
    void ShrinkDim(int64_t dim, int64_t start, int64_t size);
 
    int64_t NumReductionDims() const;
 
    int64_t NumDims() const { return ndims_; }
 
    const int64_t* GetPrimaryShape() const { return primary_shape_; }
    int64_t* GetPrimaryShape() { return primary_shape_; }
 
    const int64_t* GetPrimaryStrides() const { return primary_strides_; }
 
    int64_t NumWorkloads() const;
 
    int64_t NumOutputElements() const;
 
    int64_t NumInputs() const { return num_inputs_; }
 
    int64_t NumOutputs() const { return num_outputs_; }
 
    TensorRef& GetInput(int64_t i) {
        if (i >= num_inputs_ || i < 0) {
            utility::LogError("0 <= i < {} required, however, i = {}.",
                              num_inputs_, i);
        }
        return inputs_[i];
    }
    const TensorRef& GetInput(int64_t i) const {
        if (i >= num_inputs_ || i < 0) {
            utility::LogError("0 <= i < {} required, however, i = {}.",
                              num_inputs_, i);
        }
        return inputs_[i];
    }
 
    TensorRef& GetOutput(int64_t i) {
        if (i >= num_outputs_ || i < 0) {
            utility::LogError("0 <= i < {} required, however, i = {}.",
                              num_outputs_, i);
        }
        return outputs_[i];
    }
    const TensorRef& GetOutput(int64_t i) const {
        if (i >= num_outputs_ || i < 0) {
            utility::LogError("0 <= i < {} required, however, i = {}.",
                              num_outputs_, i);
        }
        return outputs_[i];
    }
 
    TensorRef& GetOutput() {
        if (num_outputs_ > 1) {
            utility::LogError("num_outputs_ == {} > 0, use GetOutput(i)",
                              num_outputs_);
        }
        return GetOutput(0);
    }
    const TensorRef& GetOutput() const {
        if (num_outputs_ > 1) {
            utility::LogError("num_outputs_ == {} > 0, use GetOutput(i)",
                              num_outputs_);
        }
        return GetOutput(0);
    }
 
    bool IsReductionDim(int64_t dim) const {
        // All outputs have the same shape and reduction dims. Even if they
        // don't have the same initial strides, the reduced strides are always
        // set to 0. Thus it is okay to use outputs_[0].
        return outputs_[0].byte_strides_[dim] == 0 && primary_shape_[dim] > 1;
    }
 
    OPEN3D_HOST_DEVICE char* GetInputPtr(int64_t input_idx,
                                         int64_t workload_idx) const {
        if (input_idx < 0 || input_idx >= num_inputs_) {
            return nullptr;
        }
        return GetWorkloadDataPtr(inputs_[input_idx],
                                  inputs_contiguous_[input_idx], workload_idx);
    }
 
    template <typename T>
    OPEN3D_HOST_DEVICE T* GetInputPtr(int64_t input_idx,
                                      int64_t workload_idx) const {
        if (input_idx < 0 || input_idx >= num_inputs_) {
            return nullptr;
        }
        return GetWorkloadDataPtr<T>(inputs_[input_idx],
                                     inputs_contiguous_[input_idx],
                                     workload_idx);
    }
 
    OPEN3D_HOST_DEVICE char* GetOutputPtr(int64_t workload_idx) const {
        return GetWorkloadDataPtr(outputs_[0], outputs_contiguous_[0],
                                  workload_idx);
    }
 
    template <typename T>
    OPEN3D_HOST_DEVICE T* GetOutputPtr(int64_t workload_idx) const {
        return GetWorkloadDataPtr<T>(outputs_[0], outputs_contiguous_[0],
                                     workload_idx);
    }
 
    OPEN3D_HOST_DEVICE char* GetOutputPtr(int64_t output_idx,
                                          int64_t workload_idx) const {
        return GetWorkloadDataPtr(outputs_[output_idx],
                                  outputs_contiguous_[output_idx],
                                  workload_idx);
    }
 
    template <typename T>
    OPEN3D_HOST_DEVICE T* GetOutputPtr(int64_t output_idx,
                                       int64_t workload_idx) const {
        return GetWorkloadDataPtr<T>(outputs_[output_idx],
                                     outputs_contiguous_[output_idx],
                                     workload_idx);
    }
 
#ifdef BUILD_ISPC_MODULE
    ispc::Indexer ToISPC() const;
#endif
 
protected:
    void CoalesceDimensions();
 
    // Permute reduction dimensions to front.
    // TODO: Sort the dimensions based on strides in ascending orderto improve
    // thread coalescing.
    void ReorderDimensions(const SizeVector& reduction_dims);
 
    void UpdatePrimaryStrides();
 
    void UpdateContiguousFlags();
 
    static void BroadcastRestride(TensorRef& src,
                                  int64_t dst_ndims,
                                  const int64_t* dst_shape);
 
    static void ReductionRestride(TensorRef& dst,
                                  int64_t src_ndims,
                                  const int64_t* src_shape,
                                  const SizeVector& reduction_dims);
 
    OPEN3D_HOST_DEVICE char* GetWorkloadDataPtr(const TensorRef& tr,
                                                bool tr_contiguous,
                                                int64_t workload_idx) const {
        // For 0-sized input reduction op, the output Tensor
        // workload_idx == 1 > NumWorkloads() == 0.
        if (workload_idx < 0) {
            return nullptr;
        }
 
        int64_t offset = 0;
        bool use_linear = tr_contiguous;
        for (int64_t i = 0; i < ndims_; ++i) {
            if (tr.byte_strides_[i] == 0) {
                use_linear = false;
                break;
            }
        }
        if (use_linear) {
            int64_t tr_elements = 1;
            for (int64_t i = 0; i < tr.ndims_; ++i) {
                tr_elements *= tr.shape_[i];
            }
            int64_t primary_elements = 1;
            for (int64_t i = 0; i < ndims_; ++i) {
                primary_elements *= primary_shape_[i];
            }
            if (tr_elements != primary_elements) {
                use_linear = false;
            }
        }
        if (use_linear) {
            offset = workload_idx * tr.dtype_byte_size_;
        } else {
            for (int64_t i = 0; i < ndims_; ++i) {
                offset += workload_idx / primary_strides_[i] *
                          tr.byte_strides_[i];
                workload_idx = workload_idx % primary_strides_[i];
            }
        }
 
        OPEN3D_ASSERT(tr.ContainsByteOffset(offset),
                      "Index operation data pointer is out of range.");
        return static_cast<char*>(tr.data_ptr_) + offset;
    }
 
    template <typename T>
    OPEN3D_HOST_DEVICE T* GetWorkloadDataPtr(const TensorRef& tr,
                                             bool tr_contiguous,
                                             int64_t workload_idx) const {
        // See note of this function.
        // If sizeof(T) == tr.dtype_byte_size_, then we can just static cast the
        // byte pointer.
        return static_cast<T*>(static_cast<void*>(
                GetWorkloadDataPtr(tr, tr_contiguous, workload_idx)));
    }
 
    int64_t num_inputs_ = 0;
    int64_t num_outputs_ = 0;
 
    TensorRef inputs_[MAX_INPUTS];
 
    TensorRef outputs_[MAX_OUTPUTS];
 
    bool inputs_contiguous_[MAX_INPUTS];
 
    bool outputs_contiguous_[MAX_OUTPUTS];
 
    int64_t primary_shape_[MAX_DIMS];
 
    int64_t primary_strides_[MAX_DIMS];
 
    int64_t ndims_ = 0;
 
    bool final_output_ = true;
 
    bool accumulate_ = false;
};
 
class IndexerIterator {
public:
    struct Iterator {
        Iterator() {};
        Iterator(const Indexer& indexer);
        Iterator(Iterator&& other) = default;
 
        Indexer& operator*() const;
        Iterator& operator++();
        bool operator==(const Iterator& other) const;
        bool operator!=(const Iterator& other) const;
 
        std::vector<std::unique_ptr<Indexer>> vec_;
    };
 
    IndexerIterator(const Indexer& indexer);
 
    Iterator begin() const;
    Iterator end() const;
 
private:
    const Indexer& indexer_;
};
 
}  // namespace core
}  // namespace open3d