/***************************************************************************************************
 * 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 Provides a mechanism for packing and unpacking elements smaller than
   one byte
*/
#pragma once

#include "cutlass/numeric_types.h"

namespace cutlass {

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

/// This class provides a mechanism for packing and unpacking elements smaller
/// than one byte. It assumes these sub-byte elements are packed in a
/// traditional C++ numeric type.
///
/// The intended application is to provide a mechanism to indirectly reference
/// elements in memory or Array<> objects whose addresses cannot otherwise be
/// taken since they are smaller than one byte.
///
/// Supports basic pointer arithmetic:
///
/// Example:
///
///   int4b_t *ptr = ...;
///
///   SubbyteReference<int4b_t> ref = ptr;
///   ref += 15;
///
///   int4b_t x = ref;      // load an int4b_t
///   ref = x + 2_s4;      // perform arithmetic on int4b_t and then store
///
template <typename Element_,           /// CUTLASS numeric element type.
          typename Storage_ = uint8_t  /// Underlying storage type. Must be able
                                       /// to hold an integer
                                       ///   number of objects of type Element.
          >
class ConstSubbyteReference {
public:
    using Element = Element_;
    using Storage = Storage_;
    using StoragePointer = Storage const*;

    static_assert(sizeof_bits<Element>::value <= sizeof_bits<Storage>::value,
                  "Size of Element must not be greater than Storage.");

    static_assert(!(sizeof_bits<Storage>::value % sizeof_bits<Element>::value),
                  "Storage must be divisible by Element");

private:
    ///! Number of elements per storage vector
    int const kElementsPerVector =
            sizeof_bits<Storage>::value / sizeof_bits<Element>::value;

    ///! Bit mask
    Storage const kMask =
            ((sizeof_bits<Element>::value < sizeof_bits<Storage>::value)
                     ? (Storage(1) << sizeof_bits<Element>::value) - Storage(1)
                     : ~Storage(0));

private:
    /// Pointer to array containing element
    StoragePointer ptr_;

    /// Offset (in units of elements) from pointer.
    ///
    /// Invariant: must always be in range [0, kElementsPerVector)
    int offset_;

public:
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference() : ptr_(nullptr), offset_(0) {}

    /// Constructor
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference(Element const* ptr,  /// pointer to memory
                          int64_t offset  /// logical offset in units of Element
                          )
            : ptr_(reinterpret_cast<StoragePointer>(ptr)), offset_(0) {
        int64_t offset_in_vectors = offset / kElementsPerVector;
        int64_t offset_in_elements = offset % kElementsPerVector;

        ptr_ += offset_in_vectors;
        offset_ = int(offset_in_elements);
    }

    /// Constructor
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference(Element* ptr = nullptr)
            : ConstSubbyteReference(ptr, 0) {}

    /// Gets storage pointer
    CUTLASS_HOST_DEVICE
    StoragePointer storage_pointer() const { return ptr_; }

    /// Gets element offset within storage vector
    CUTLASS_HOST_DEVICE
    int element_offset() const { return offset_; }

    /// Unpacks an element from memory
    CUTLASS_HOST_DEVICE
    Element get() const {
        Storage item = Storage(
                (*ptr_ >> (offset_ * sizeof_bits<Element>::value)) & kMask);
        return reinterpret_cast<Element const&>(item);
    }

    /// Unpacks an element from memory
    CUTLASS_HOST_DEVICE
    operator Element() const { return get(); }

    /// Adds an offset in units of elements to the reference
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference& operator+=(int offset) {
        offset += offset_;

        int offset_in_vectors = offset / kElementsPerVector;
        int offset_in_elements = offset % kElementsPerVector;

        ptr_ += offset_in_vectors;
        offset_ = offset_in_elements;

        return *this;
    }

    /// Adds an offset in units of elements to the reference
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference& operator+=(long long offset) {
        offset += offset_;

        long long offset_in_vectors = offset / kElementsPerVector;
        int offset_in_elements = int(offset % kElementsPerVector);

        ptr_ += offset_in_vectors;
        offset_ = offset_in_elements;

        return *this;
    }

    /// Adds an offset in units of elements to the reference
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference& operator-=(int offset) {
        int offset_in_vectors = offset / kElementsPerVector;
        int offset_in_elements = offset % kElementsPerVector;

        ptr_ -= offset_in_vectors;
        offset_ -= offset_in_elements;

        if (offset_ < 0) {
            offset_ += kElementsPerVector;
            --ptr_;
        }

        return *this;
    }

    /// Adds an offset in units of elements to the reference
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference& operator-=(long long offset) {
        long long offset_in_vectors = offset / kElementsPerVector;
        int offset_in_elements = int(offset % kElementsPerVector);

        ptr_ -= offset_in_vectors;
        offset_ -= offset_in_elements;

        if (offset_ < 0) {
            offset_ += kElementsPerVector;
            --ptr_;
        }

        return *this;
    }

    /// Returns a reference to an element with a given offset from the current
    /// reference
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference operator+(int offset) const {
        ConstSubbyteReference ref(ptr_, offset_);
        ref += offset;

        return ref;
    }

    /// Returns a reference to an element with a given offset from the current
    /// reference
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference operator+(long long offset) const {
        ConstSubbyteReference ref(ptr_, offset_);
        ref += offset;

        return ref;
    }

    /// Returns a reference to an element with a given offset from the current
    /// reference
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference operator-(int offset) const {
        ConstSubbyteReference ref(ptr_, offset_);
        ref -= offset;

        return ref;
    }

    /// Returns a reference to an element with a given offset from the current
    /// reference
    CUTLASS_HOST_DEVICE
    ConstSubbyteReference operator-=(long long offset) const {
        ConstSubbyteReference ref(ptr_, offset_);
        ref -= offset;

        return ref;
    }

    /// Computes the difference in elements between references
    CUTLASS_HOST_DEVICE
    ptrdiff_t operator-(ConstSubbyteReference ref) const {
        return (ptr_ - ref.ptr_) * kElementsPerVector + (offset_ - ref.offset_);
    }

    /// Explicit cast to int
    CUTLASS_HOST_DEVICE
    explicit operator int() const { return int(get()); }

    /// Explicit cast to signed 64-bit integer
    CUTLASS_HOST_DEVICE
    explicit operator int64_t() const { return int64_t(get()); }

    /// Explicit cast to unsigned 64-bit integer
    CUTLASS_HOST_DEVICE
    explicit operator uint64_t() const { return uint64_t(get()); }

    /// Explicit cast to float
    CUTLASS_HOST_DEVICE
    explicit operator float() const { return float(get()); }

    /// Explicit cast to double
    CUTLASS_HOST_DEVICE
    explicit operator double() const { return double(get()); }
};

template <typename Element_,           /// CUTLASS numeric element type.
          typename Storage_ = uint8_t  /// Underlying storage type. Must be able
                                       /// to hold an integer
                                       ///   number of objects of type Element.
          >
class SubbyteReference {
public:
    using Element = Element_;
    using Storage = Storage_;
    using StoragePointer = Storage*;

    static_assert(sizeof_bits<Element>::value <= sizeof_bits<Storage>::value,
                  "Size of Element must not be greater than Storage.");

    static_assert(!(sizeof_bits<Storage>::value % sizeof_bits<Element>::value),
                  "Storage must be divisible by Element");

private:
    ///! Number of elements per storage vector
    int const kElementsPerVector =
            sizeof_bits<Storage>::value / sizeof_bits<Element>::value;

    ///! Bit mask
    Storage const kMask =
            ((sizeof_bits<Element>::value < sizeof_bits<Storage>::value)
                     ? (Storage(1) << sizeof_bits<Element>::value) - Storage(1)
                     : ~Storage(0));

private:
    /// Pointer to array containing element
    StoragePointer ptr_;

    /// Offset (in units of elements) from pointer.
    ///
    /// Invariant: must always be in range [0, kElementsPerVector)
    int offset_;

public:
    CUTLASS_HOST_DEVICE
    SubbyteReference() : ptr_(nullptr), offset_(0) {}

    /// Constructor
    CUTLASS_HOST_DEVICE
    SubbyteReference(Element* ptr,   /// pointer to memory
                     int64_t offset  /// logical offset in units of Element
                     )
            : ptr_(reinterpret_cast<StoragePointer>(ptr)), offset_(0) {
        int64_t offset_in_vectors = offset / kElementsPerVector;
        int64_t offset_in_elements = offset % kElementsPerVector;

        ptr_ += offset_in_vectors;
        offset_ = int(offset_in_elements);
    }

    /// Constructor
    CUTLASS_HOST_DEVICE
    SubbyteReference(Element* ptr = nullptr) : SubbyteReference(ptr, 0) {}

    /// Gets storage pointer
    CUTLASS_HOST_DEVICE
    StoragePointer storage_pointer() const { return ptr_; }

    /// Gets element offset within storage vector
    CUTLASS_HOST_DEVICE
    int element_offset() const { return offset_; }

    /// Unpacks an element from memory
    CUTLASS_HOST_DEVICE
    Element get() const {
        Storage item = Storage(
                (*ptr_ >> (offset_ * sizeof_bits<Element>::value)) & kMask);
        return reinterpret_cast<Element const&>(item);
    }

    /// Stores an element to memory
    CUTLASS_HOST_DEVICE
    SubbyteReference& set(Element const& x) {
        Storage item = (reinterpret_cast<Storage const&>(x) & kMask);

        Storage kUpdateMask =
                Storage(~(kMask << (offset_ * sizeof_bits<Element>::value)));
        *ptr_ = Storage(
                (*ptr_ & kUpdateMask) |
                Storage(item << (offset_ * sizeof_bits<Element>::value)));

        return *this;
    }

    /// Unpacks an element from memory
    CUTLASS_HOST_DEVICE
    operator Element() const { return get(); }

    /// Stores an element to memory
    CUTLASS_HOST_DEVICE
    SubbyteReference& operator=(Element const& x) { return set(x); }

    /// Stores an element to memory
    CUTLASS_HOST_DEVICE
    SubbyteReference& operator=(SubbyteReference const& x) {
        return set(x.get());
    }

    /// Stores an element to memory
    CUTLASS_HOST_DEVICE
    SubbyteReference& operator=(
            ConstSubbyteReference<Element, Storage> const& x) {
        return set(x.get());
    }

    /// Adds an offset in units of elements to the reference
    CUTLASS_HOST_DEVICE
    SubbyteReference& operator+=(int offset) {
        offset += offset_;

        int offset_in_vectors = offset / kElementsPerVector;
        int offset_in_elements = offset % kElementsPerVector;

        ptr_ += offset_in_vectors;
        offset_ = offset_in_elements;

        return *this;
    }

    /// Adds an offset in units of elements to the reference
    CUTLASS_HOST_DEVICE
    SubbyteReference& operator+=(long long offset) {
        offset += offset_;

        long long offset_in_vectors = offset / kElementsPerVector;
        int offset_in_elements = int(offset % kElementsPerVector);

        ptr_ += offset_in_vectors;
        offset_ = offset_in_elements;

        return *this;
    }

    /// Adds an offset in units of elements to the reference
    CUTLASS_HOST_DEVICE
    SubbyteReference& operator-=(int offset) {
        int offset_in_vectors = offset / kElementsPerVector;
        int offset_in_elements = offset % kElementsPerVector;

        ptr_ -= offset_in_vectors;
        offset_ -= offset_in_elements;

        if (offset_ < 0) {
            offset_ += kElementsPerVector;
            --ptr_;
        }

        return *this;
    }

    /// Adds an offset in units of elements to the reference
    CUTLASS_HOST_DEVICE
    SubbyteReference& operator-=(long long offset) {
        long long offset_in_vectors = offset / kElementsPerVector;
        int offset_in_elements = int(offset % kElementsPerVector);

        ptr_ -= offset_in_vectors;
        offset_ -= offset_in_elements;

        if (offset_ < 0) {
            offset_ += kElementsPerVector;
            --ptr_;
        }

        return *this;
    }

    /// Returns a reference to an element with a given offset from the current
    /// reference
    CUTLASS_HOST_DEVICE
    SubbyteReference operator+(int offset) const {
        SubbyteReference ref(ptr_, offset_);
        ref += offset;

        return ref;
    }

    /// Returns a reference to an element with a given offset from the current
    /// reference
    CUTLASS_HOST_DEVICE
    SubbyteReference operator+(long long offset) const {
        SubbyteReference ref(ptr_, offset_);
        ref += offset;

        return ref;
    }

    /// Returns a reference to an element with a given offset from the current
    /// reference
    CUTLASS_HOST_DEVICE
    SubbyteReference operator-(int offset) const {
        SubbyteReference ref(ptr_, offset_);
        ref -= offset;

        return ref;
    }

    /// Returns a reference to an element with a given offset from the current
    /// reference
    CUTLASS_HOST_DEVICE
    SubbyteReference operator-=(long long offset) const {
        SubbyteReference ref(ptr_, offset_);
        ref -= offset;

        return ref;
    }

    /// Computes the difference in elements between references
    CUTLASS_HOST_DEVICE
    ptrdiff_t operator-(SubbyteReference ref) const {
        return (ptr_ - ref.ptr_) * kElementsPerVector + (offset_ - ref.offset_);
    }

    /// Explicit cast to int
    CUTLASS_HOST_DEVICE
    explicit operator int() const { return int(get()); }

    /// Explicit cast to signed 64-bit integer
    CUTLASS_HOST_DEVICE
    explicit operator int64_t() const { return int64_t(get()); }

    /// Explicit cast to unsigned 64-bit integer
    CUTLASS_HOST_DEVICE
    explicit operator uint64_t() const { return uint64_t(get()); }

    /// Explicit cast to float
    CUTLASS_HOST_DEVICE
    explicit operator float() const { return float(get()); }

    /// Explicit cast to double
    CUTLASS_HOST_DEVICE
    explicit operator double() const { return double(get()); }
};

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

template <typename Element, bool subbyte = (sizeof_bits<Element>::value < 8)>
struct ReferenceFactory;

template <typename Element>
struct ReferenceFactory<Element, false> {
    CUTLASS_HOST_DEVICE
    static Element& get(Element* ptr, int64_t offset) { return ptr[offset]; }

    CUTLASS_HOST_DEVICE
    static Element const& get(Element const* ptr, int64_t offset) {
        return ptr[offset];
    }
};

template <typename Element>
struct ReferenceFactory<Element, true> {
    CUTLASS_HOST_DEVICE
    static SubbyteReference<Element> get(Element* ptr, int64_t offset) {
        return SubbyteReference<Element>(ptr, offset);
    }

    CUTLASS_HOST_DEVICE
    static ConstSubbyteReference<Element> get(Element const* ptr,
                                              int64_t offset) {
        return ConstSubbyteReference<Element>(ptr, offset);
    }
};

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

}  // namespace cutlass
