AdvancedIndexing.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018 www.open3d.org
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.
// ----------------------------------------------------------------------------

#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 "
                    "indexd_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)));
            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