AdvancedIndexing.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2023 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#pragma once
 
#include <vector>
 
#include "open3d/core/Indexer.h"
#include "open3d/core/SizeVector.h"
#include "open3d/core/Tensor.h"
 
namespace open3d {
namespace core {
 
class AdvancedIndexPreprocessor {
public:
    AdvancedIndexPreprocessor(const Tensor& tensor,
                              const std::vector<Tensor>& index_tensors)
        : tensor_(tensor), index_tensors_(ExpandBoolTensors(index_tensors)) {
        RunPreprocess();
    }
 
    inline Tensor GetTensor() const { return tensor_; }
 
    inline std::vector<Tensor> GetIndexTensors() const {
        return index_tensors_;
    }
 
    inline SizeVector GetOutputShape() const { return output_shape_; }
 
    inline SizeVector GetIndexedShape() const { return indexed_shape_; }
 
    inline SizeVector GetIndexedStrides() const { return indexed_strides_; }
 
    static bool IsIndexSplittedBySlice(
            const std::vector<Tensor>& index_tensors);
 
    static std::pair<Tensor, std::vector<Tensor>> ShuffleIndexedDimsToFront(
            const Tensor& tensor, const std::vector<Tensor>& index_tensors);
 
    static std::pair<std::vector<Tensor>, SizeVector>
    ExpandToCommonShapeExceptZeroDim(const std::vector<Tensor>& index_tensors);
 
    // Replace indexed dimensions with stride 0 and the size of the result
    // tensor.
    //
    // The offset in these dimensions is computed by the kernel using
    // the index tensor's values and the stride of the tensor. The new shape is
    // not meaningful. It's used to make the shape compatible with the result
    // tensor.
    //
    // Effectively, we throw away the tensor's shape and strides for the sole
    // purpose of element-wise iteration for the Indexer. The tensor's original
    // strides are stored in indexed_shape_ and indexed_strides_,
    // which are passed to fancy indexing kernels.
    static Tensor RestrideTensor(const Tensor& tensor,
                                 int64_t dims_before,
                                 int64_t dims_indexed,
                                 SizeVector replacement_shape);
 
    // Add dimensions of size 1 to an index tensor so that it can be broadcast
    // to the result shape and iterated over element-wise like the result tensor
    // and the restrided src.
    static Tensor RestrideIndexTensor(const Tensor& index_tensor,
                                      int64_t dims_before,
                                      int64_t dims_after);
 
protected:
    void RunPreprocess();
 
    static std::vector<Tensor> ExpandBoolTensors(
            const std::vector<Tensor>& index_tensors);
 
    Tensor tensor_;
 
    std::vector<Tensor> index_tensors_;
 
    SizeVector output_shape_;
 
    SizeVector indexed_shape_;
 
    SizeVector indexed_strides_;
};
 
class AdvancedIndexer {
public:
    enum class AdvancedIndexerMode { SET, GET };
 
    AdvancedIndexer(const Tensor& src,
                    const Tensor& dst,
                    const std::vector<Tensor>& index_tensors,
                    const SizeVector& indexed_shape,
                    const SizeVector& indexed_strides,
                    AdvancedIndexerMode mode)
        : mode_(mode) {
        if (indexed_shape.size() != indexed_strides.size()) {
            utility::LogError(
                    "Internal error: indexed_shape's ndim {} does not equal to "
                    "indexed_strides' ndim {}",
                    indexed_shape.size(), indexed_strides.size());
        }
        num_indices_ = indexed_shape.size();
 
        // Initialize Indexer
        std::vector<Tensor> inputs;
        inputs.push_back(src);
        for (const Tensor& index_tensor : index_tensors) {
            if (index_tensor.NumDims() != 0) {
                inputs.push_back(index_tensor);
            }
        }
        indexer_ = Indexer({inputs}, dst, DtypePolicy::NONE);
 
        // Fill shape and strides
        if (num_indices_ != static_cast<int64_t>(indexed_strides.size())) {
            utility::LogError(
                    "Internal error: indexed_shape's ndim {} does not equal to "
                    "indexd_strides' ndim {}",
                    num_indices_, indexed_strides.size());
        }
        for (int64_t i = 0; i < num_indices_; ++i) {
            indexed_shape_[i] = indexed_shape[i];
            indexed_strides_[i] = indexed_strides[i];
        }
 
        // Check dtypes
        if (src.GetDtype() != dst.GetDtype()) {
            utility::LogError(
                    "src's dtype {} is not the same as dst's dtype {}.",
                    src.GetDtype().ToString(), dst.GetDtype().ToString());
        }
        element_byte_size_ = src.GetDtype().ByteSize();
    }
 
    inline OPEN3D_HOST_DEVICE char* GetInputPtr(int64_t workload_idx) const {
        char* ptr = indexer_.GetInputPtr(0, workload_idx);
        ptr += GetIndexedOffset(workload_idx) * element_byte_size_ *
               (mode_ == AdvancedIndexerMode::GET);
        return ptr;
    }
 
    inline OPEN3D_HOST_DEVICE char* GetOutputPtr(int64_t workload_idx) const {
        char* ptr = indexer_.GetOutputPtr(workload_idx);
        ptr += GetIndexedOffset(workload_idx) * element_byte_size_ *
               (mode_ == AdvancedIndexerMode::SET);
        return ptr;
    }
 
    inline OPEN3D_HOST_DEVICE int64_t
    GetIndexedOffset(int64_t workload_idx) const {
        int64_t offset = 0;
        for (int64_t i = 0; i < num_indices_; ++i) {
            int64_t index = *(reinterpret_cast<int64_t*>(
                    indexer_.GetInputPtr(i + 1, workload_idx)));
            OPEN3D_ASSERT(index >= -indexed_shape_[i] &&
                          index < indexed_shape_[i] && "Index out of bounds.");
            index += indexed_shape_[i] * (index < 0);
            offset += index * indexed_strides_[i];
        }
        return offset;
    }
 
    int64_t NumWorkloads() const { return indexer_.NumWorkloads(); }
 
protected:
    Indexer indexer_;
    AdvancedIndexerMode mode_;
    int64_t num_indices_;
    int64_t element_byte_size_;
    int64_t indexed_shape_[MAX_DIMS];
    int64_t indexed_strides_[MAX_DIMS];
};
 
}  // namespace core
}  // namespace open3d