/***************************************************************************************************
 * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 *modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright notice,
 *this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *notice, this list of conditions and the following disclaimer in the
 *documentation and/or other materials provided with the distribution.
 *     * Neither the name of the NVIDIA CORPORATION nor the names of its
 *contributors may be used to endorse or promote products derived from this
 *software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 *AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 *IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY DIRECT,
 *INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 *DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 *OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TOR (INCLUDING
 *NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
 *EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 **************************************************************************************************/
/*! \file
    \brief Defines a structure containing strides and a pointer to tensor data.

    TensorView is derived from TensorRef and contributes bounds to the tensor's
   index space. Thus, it is a complete mathematical object and may be used in
   tensor algorithms. It is decoupled from data storage and is therefore
   lightweight and may be embedded in larger tensor objects or memory
   structures.

    See cutlass/tensor_ref.h for more details about the mapping of the logical
   tensor index space to linear memory.
*/

#pragma once

#if !defined(__CUDACC_RTC__)
#include <cmath>
#endif

#include "cutlass/cutlass.h"
#include "cutlass/tensor_ref.h"

namespace cutlass {

////////////////////////////////////////////////////////////////////////////////////////////////////

template <
        /// Data type of element stored within tensor
        typename Element_,
        /// Maps a Coord<Rank_> in the logical tensor index space to the
        /// internal n-D array
        typename Layout_>
class TensorView : public TensorRef<Element_, Layout_> {
public:
    /// Base tensor reference
    using Base = cutlass::TensorRef<Element_, Layout_>;

    /// Mapping function from logical coordinate to internal n-D array
    using Layout = Layout_;

    /// TensorRef pointing to constant memory
    using ConstTensorRef = typename Base::ConstTensorRef;

    /// Underlying TensorRef type
    using TensorRef = Base;

    /// Data type of individual access
    using Element = Element_;

    /// Reference type to an element
    using Reference = Element&;

    /// Logical rank of tensor index space
    static int const kRank = Layout::kRank;

    /// Index type
    using Index = typename Layout::Index;

    /// Long index used for pointer offsets
    using LongIndex = typename Layout::LongIndex;

    /// Coordinate in logical tensor space
    using TensorCoord = typename Layout::TensorCoord;

    /// Coordinate in storage n-D array
    using Stride = typename Layout::Stride;

    /// TensorView pointing to constant memory
    using ConstTensorView =
            TensorView<typename platform::remove_const<Element>::type const,
                       Layout>;

    /// TensorView pointing to non-constant memory
    using NonConstTensorView =
            TensorView<typename platform::remove_const<Element>::type, Layout>;

    /// Require at least rank=1. Mathematically, a rank=0 tensor would be
    /// considered to be a scalar, but degenerate cases such as these are
    /// difficult to accommodate without extensive C++ metaprogramming or
    /// support for zero-length arrays.
    static_assert(kRank > 0, "Cannot define a zero-rank TensorRef");

private:
    /// View extent
    TensorCoord extent_;

public:
    //
    // Methods
    //

    /// Constructs a TensorView object
    CUTLASS_HOST_DEVICE
    TensorView(TensorCoord const& extent = TensorCoord()) : extent_(extent) {}

    /// Constructs a TensorView object
    CUTLASS_HOST_DEVICE
    TensorView(Element* ptr,          ///< pointer to start of tensor
               Layout const& layout,  ///< layout object containing stride and
                                      ///< mapping function
               TensorCoord const&
                       extent  ///< size of the view in logical coordinates
               )
            : Base(ptr, layout), extent_(extent) {}

    /// Constructs a TensorView object
    CUTLASS_HOST_DEVICE
    TensorView(TensorRef const&
                       ref,  ///< pointer and layout object referencing a tensor
               TensorCoord const& extent  ///< logical size of tensor
               )
            : Base(ref), extent_(extent) {}

    /// Converting constructor from TensorRef to non-constant data.
    CUTLASS_HOST_DEVICE
    TensorView(NonConstTensorView const& view  ///< TensorView to non-const data
               )
            : Base(view), extent_(view.extent_) {}

    /// Updates the pointer and layout object
    CUTLASS_HOST_DEVICE
    void reset(Element* ptr, Layout const& layout, TensorCoord const& extent) {
        Base::reset(ptr, layout);
        this->resize(extent);
    }

    /// Updates the pointer
    CUTLASS_HOST_DEVICE
    void reset(Element* ptr) { Base::reset(ptr); }

    /// Changes the size of the view without affecting pointer or layout
    CUTLASS_HOST_DEVICE
    void resize(TensorCoord const& extent) { this->extent_ = extent; }

    /// Returns the extent of the view (the size along each logical dimension).
    CUTLASS_HOST_DEVICE
    TensorCoord const& extent() const { return extent_; }

    /// Returns the extent along a particular logical dimension.
    CUTLASS_HOST_DEVICE
    Index extent(int dim) const { return extent_.at(dim); }

    /// Returns the number of logical elements
    CUTLASS_HOST_DEVICE
    LongIndex size() const { return extent_.product(); }

    /// Determines whether a location is within a tensor
    CUTLASS_HOST_DEVICE
    bool contains(TensorCoord const& coord) const {
        CUTLASS_PRAGMA_UNROLL
        for (int dim = 0; dim < kRank; ++dim) {
            if (!(coord[dim] >= 0 && coord[dim] < extent(dim))) {
                return false;
            }
        }
        return true;
    }

    /// Returns a TensorRef pointing to the first element of the tensor.
    CUTLASS_HOST_DEVICE
    TensorRef ref() const { return TensorRef(this->data(), this->layout()); }

    /// Returns a TensorRef pointing to the first element of the tensor.
    CUTLASS_HOST_DEVICE
    ConstTensorRef const_ref() const {
        return ConstTensorRef(this->data(), this->layout());
    }

    /// Returns a TensorView to const data
    CUTLASS_HOST_DEVICE
    ConstTensorView const_view() const {
        return ConstTensorView(const_ref(), extent_);
    }

    /// Returns a Tensor_view given location and size quantities
    CUTLASS_HOST_DEVICE
    TensorView subview(TensorCoord extent,  ///< extent of the resulting view
                       TensorCoord const& location =
                               TensorCoord()  ///< resulting view's origin
                                              ///< within the old view
                       ) const {
        return TensorView(ref(), extent.clamp(extent_ - location))
                .add_coord_offset(location);
    }

    /// Returns the number of scalar elements needed to store tensor.
    CUTLASS_HOST_DEVICE
    size_t capacity() const { return Base::layout().capacity(extent_); }

    /// Returns a TensorView offset by a given amount
    CUTLASS_HOST_DEVICE
    TensorView operator+(
            TensorCoord const&
                    b  ///< offset in the logical coordinate space of the tensor
            ) const {
        TensorView result(*this);
        result.add_pointer_offset(this->offset(b));
        return result;
    }

    /// Returns a TensorRef offset by a given amount
    CUTLASS_HOST_DEVICE
    TensorView& operator+=(
            TensorCoord const&
                    b  ///< offset in the logical coordinate space of the tensor
    ) {
        this->add_pointer_offset(this->offset(b));
        return *this;
    }

    /// Returns a TensorRef offset by a given amount
    CUTLASS_HOST_DEVICE
    TensorView operator-(
            TensorCoord const&
                    b  ///< offset in the logical coordinate space of the tensor
            ) const {
        TensorRef result(*this);
        result.add_pointer_offset(-this->offset(b));
        return result;
    }

    /// Returns a TensorRef offset by a given amount
    CUTLASS_HOST_DEVICE
    TensorView& operator-=(
            TensorCoord const&
                    b  ///< offset in the logical coordinate space of the tensor
    ) {
        this->add_pointer_offset(-this->offset(b));
        return *this;
    }
};

////////////////////////////////////////////////////////////////////////////////////////////////////

/// Constructs a TensorRef, deducing types from arguments.
template <typename Element, typename Layout>
CUTLASS_HOST_DEVICE TensorView<Element, Layout> make_TensorView(
        Element* ptr, Layout const& layout,
        typename Layout::TensorCoord const& extent) {
    return TensorView<Element, Layout>(ptr, layout, extent);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

}  // namespace cutlass
