1391 lines
40 KiB
C++
1391 lines
40 KiB
C++
/***************************************************************************************************
|
|
* Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: BSD-3-Clause
|
|
*
|
|
* Redistribution and use in source and binary forms, with or without
|
|
* modification, are permitted provided that the following conditions are met:
|
|
*
|
|
* 1. Redistributions of source code must retain the above copyright notice, this
|
|
* list of conditions and the following disclaimer.
|
|
*
|
|
* 2. 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.
|
|
*
|
|
* 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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/cutlass.h"
|
|
#include "cutlass/integer_subbyte.h"
|
|
#include "cutlass/fast_math.h"
|
|
|
|
namespace cutlass {
|
|
|
|
namespace detail {
|
|
// This is an implementation detail of cutlass::SubbyteReference and.
|
|
// cutlass::HostTensor. For a given logical element type Element,
|
|
// and its corresponding storage (physical) element type StorageUnit,
|
|
// it computes quantities that help with managing allocations.
|
|
//
|
|
// CUTLASS uses a hidden "ContainerUnitType" or StorageUnit type to support
|
|
// packed arrays of subbyte types such as int4. Element is the "logical" type
|
|
// for computations, while CUTLASS uses StorageUnit as the element type
|
|
// of a packed array of Element. If Element is not a subbyte type,
|
|
// then the corresponding StorageUnit type is just Element itself.
|
|
//
|
|
// The ContainerType is always calculated as an array StorageUnit type (the StorageUnit
|
|
// is always a byte for subbyte types),
|
|
// and its number of bits is the lcm of the subbyte type's number of bits and 8.
|
|
// Below are some examples for different subbyte types.
|
|
//
|
|
// * Subbyte Type=int2, ContainerType=StorageUnit[1] (StorageUnit=uint8_t)
|
|
// * Subbyte Type=int4, ContainerType=StorageUnit[1] (StorageUnit=uint8_t)
|
|
template<class Element, class StorageUnit>
|
|
struct StorageContainerCalculator {
|
|
// kContainerTypeNumBits: The number of bits needed for ContainerType
|
|
static constexpr int kContainerTypeNumBits = (sizeof_bits<Element>::value < 8) ? cutlass::lcm_cxx11(sizeof_bits<Element>::value, sizeof_bits<StorageUnit>::value) : sizeof_bits<Element>::value;
|
|
static_assert(kContainerTypeNumBits % sizeof_bits<Element>::value == 0, "The bits of ContainerType should be divisible by the element's number of bits");
|
|
// kContainerTypeNumLogicalElements: The number of logical Element instance(s) that can be stored per ContainerType instance
|
|
static constexpr int kContainerTypeNumLogicalElements = kContainerTypeNumBits / sizeof_bits<Element>::value;
|
|
/// 3. kContainerTypeNumBytes: The number of bytes per ContainerType instance
|
|
static constexpr int kContainerTypeNumBytes = kContainerTypeNumBits / 8;
|
|
/// 4. kContainerTypeNumBytes: The number of base StorageUnit in the ContainerType
|
|
static constexpr int kContainerTypeNumStorageUnit = kContainerTypeNumBits / sizeof_bits<StorageUnit>::value;
|
|
|
|
static_assert(kContainerTypeNumBits != 0, "kContainerTypeNumBits can not be zero");
|
|
static_assert(kContainerTypeNumLogicalElements != 0, "kContainerTypeNumLogicalElements can not be zero");
|
|
static_assert(kContainerTypeNumBytes != 0, "kContainerTypeNumBytes can not be zero");
|
|
};
|
|
}
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// 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 = void
|
|
>
|
|
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_ = /// Underlying storage type. Must be able to hold an integer
|
|
/// number of objects of type Element.
|
|
|
|
#if defined(__CUDA_ARCH__) /// Default size depends on width of atomicCas() overloads.
|
|
#if (__CUDA_ARCH__ >= 700) ///
|
|
uint16_t
|
|
#else
|
|
uint32_t
|
|
#endif
|
|
#else
|
|
uint8_t
|
|
#endif
|
|
,
|
|
class = void
|
|
>
|
|
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 storage pointer
|
|
CUTLASS_HOST_DEVICE
|
|
Element * operator&() const {
|
|
return reinterpret_cast<Element *>(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 {
|
|
uint8_t const* byte_ptr = reinterpret_cast<uint8_t const*>(ptr_);
|
|
// Convert offset in elements to offset in bytes
|
|
constexpr int elements_per_byte = cutlass::sizeof_bits<uint8_t>::value / cutlass::sizeof_bits<Element>::value;
|
|
byte_ptr += offset_ / elements_per_byte;
|
|
// Offset of element within a byte
|
|
int byte_offset = offset_ % elements_per_byte;
|
|
uint8_t item = uint8_t((*byte_ptr >> (byte_offset * cutlass::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_ * cutlass::sizeof_bits<Element>::value)));
|
|
Storage new_bits = Storage(item << (offset_ * cutlass::sizeof_bits<Element>::value));
|
|
|
|
#if defined(__CUDA_ARCH__)
|
|
|
|
//
|
|
// Homebrew read-modify-write
|
|
//
|
|
Storage original;
|
|
Storage updated;
|
|
|
|
do {
|
|
|
|
original = (*ptr_);
|
|
|
|
updated = Storage((original & kUpdateMask) | new_bits);
|
|
|
|
original = atomicCAS(ptr_, original, updated);
|
|
|
|
} while (updated != original);
|
|
|
|
#else
|
|
|
|
Storage original = (*ptr_);
|
|
Storage updated = Storage((original & kUpdateMask) | new_bits);
|
|
*ptr_ = updated;
|
|
|
|
#endif
|
|
|
|
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 T> using _war = T;
|
|
template <
|
|
typename Element_, /// CUTLASS numeric element type.
|
|
typename Storage_ /// Underlying basic storage type.
|
|
>
|
|
class SubbyteReference<Element_, Storage_,
|
|
typename platform::enable_if<sizeof_bits<Storage_>::value % sizeof_bits<Element_>::value != 0>::type> {
|
|
public:
|
|
|
|
using Element = Element_;
|
|
/// Note: It's possible that StorageUnit is not divisible by Element.
|
|
/// For example, an Element instance might be stored across 2 StorageUnit instances.
|
|
/// Thus, CUTLASS needs a storage vector to hold an integer number of Element instances.
|
|
|
|
using StorageUnit = Storage_;
|
|
private:
|
|
using StorageContainerCalculator = cutlass::detail::StorageContainerCalculator<Element, StorageUnit>;
|
|
public:
|
|
static int const kBitsStoredVec = StorageContainerCalculator::kContainerTypeNumBits;
|
|
static int const kNumStorageUnitPerStoredVec = StorageContainerCalculator::kContainerTypeNumStorageUnit;
|
|
|
|
using StorageVec = StorageUnit[kNumStorageUnitPerStoredVec];
|
|
using StorageVecPointer = StorageVec *;
|
|
|
|
using CudaAtomicType = typename platform::conditional<
|
|
sizeof_bits<StorageUnit>::value == 16,
|
|
uint32_t,
|
|
uint64_t
|
|
>::type;
|
|
|
|
static_assert(sizeof_bits<Element>::value <= sizeof_bits<StorageVec>::value,
|
|
"Size of Element must not be greater than StorageVec.");
|
|
|
|
static_assert(!(sizeof_bits<StorageVec>::value % sizeof_bits<Element>::value),
|
|
"StorageVec must be divisible by Element");
|
|
|
|
private:
|
|
|
|
///! Number of elements per storage vector
|
|
int const kElementsPerVector = sizeof_bits<StorageVec>::value / sizeof_bits<Element>::value;
|
|
|
|
///! Bit mask for storage unit.
|
|
StorageUnit const kMask = (StorageUnit(1) << sizeof_bits<Element>::value) - StorageUnit(1);
|
|
|
|
/// Pointer to array containing element
|
|
_war<StorageVecPointer> ptr_;
|
|
|
|
/// Offset (in units of elements) from pointer.
|
|
///
|
|
/// Invariant: must always be in range [0, kElementsPerVector)
|
|
int offset_;
|
|
|
|
/// Element may be stored across 2 storage unit.
|
|
/// Low storage unit index in StorageVec
|
|
/// High storage unit index in StorageVec
|
|
int low_storage_unit_idx_;
|
|
int high_storage_unit_idx_;
|
|
|
|
/// Full Mask to extract the entire element
|
|
uint64_t full_element_mask_;
|
|
|
|
/// Mask to extract the Element from Low storage unit and High storage unit.
|
|
StorageUnit low_storage_mask_;
|
|
StorageUnit high_storage_mask_;
|
|
|
|
/// Start bit index inside the storage unit.
|
|
int start_bit_idx_;
|
|
|
|
private:
|
|
|
|
CUTLASS_HOST_DEVICE
|
|
void update_element_status() {
|
|
int num_bits = offset_ * sizeof_bits<Element>::value;
|
|
|
|
start_bit_idx_ = num_bits % sizeof_bits<StorageUnit>::value;
|
|
|
|
low_storage_unit_idx_ = num_bits / sizeof_bits<StorageUnit>::value;
|
|
high_storage_unit_idx_ = sizeof_bits<StorageUnit>::value - (start_bit_idx_) < sizeof_bits<Element>::value
|
|
? low_storage_unit_idx_ + 1 : low_storage_unit_idx_;
|
|
|
|
full_element_mask_ = uint64_t(kMask) << start_bit_idx_;
|
|
low_storage_mask_ = StorageUnit(full_element_mask_ & ~StorageUnit(0));
|
|
high_storage_mask_ = StorageUnit((full_element_mask_ >> sizeof_bits<StorageUnit>::value) & ~StorageUnit(0));
|
|
}
|
|
|
|
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<StorageVecPointer>(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);
|
|
|
|
update_element_status();
|
|
}
|
|
|
|
/// Constructor
|
|
CUTLASS_HOST_DEVICE
|
|
SubbyteReference(
|
|
Element *ptr = nullptr
|
|
): SubbyteReference(ptr, 0) { }
|
|
|
|
/// Gets StorageVec pointer
|
|
CUTLASS_HOST_DEVICE
|
|
StorageVecPointer storage_pointer() const {
|
|
return ptr_;
|
|
}
|
|
|
|
/// Gets StorageVec pointer
|
|
CUTLASS_HOST_DEVICE
|
|
Element * operator&() const {
|
|
return reinterpret_cast<Element *>(ptr_);
|
|
}
|
|
|
|
/// Gets element offset within StorageVec vector
|
|
CUTLASS_HOST_DEVICE
|
|
int element_offset() const {
|
|
return offset_;
|
|
}
|
|
|
|
/// Unpacks an element from memory
|
|
CUTLASS_HOST_DEVICE
|
|
Element get() const {
|
|
StorageUnit low_bits = (*ptr_)[low_storage_unit_idx_] & low_storage_mask_;
|
|
StorageUnit high_bits = low_storage_unit_idx_ != high_storage_unit_idx_ ? (*ptr_)[high_storage_unit_idx_] & high_storage_mask_ : 0;
|
|
|
|
uint64_t full_item = ((uint64_t)high_bits << sizeof_bits<StorageUnit>::value) | low_bits;
|
|
uint8_t result = uint8_t(full_item >> start_bit_idx_);
|
|
|
|
return reinterpret_cast<Element const &>(result);
|
|
}
|
|
|
|
/// Stores an element to memory
|
|
CUTLASS_HOST_DEVICE
|
|
SubbyteReference & set(Element const &x) {
|
|
|
|
uint64_t item = static_cast<uint64_t>((reinterpret_cast<uint8_t const &>(x) & kMask)) << start_bit_idx_;
|
|
|
|
StorageUnit low_new_bits = StorageUnit(item & ~StorageUnit(0));
|
|
StorageUnit high_new_bits = StorageUnit(item >> sizeof_bits<StorageUnit>::value);
|
|
|
|
StorageUnit const kLowUpdateMask = StorageUnit((~full_element_mask_) & (~StorageUnit(0)));
|
|
StorageUnit const kHighUpdateMask = StorageUnit(((~full_element_mask_) >> sizeof_bits<StorageUnit>::value) & (~StorageUnit(0)));
|
|
|
|
#if defined(__CUDA_ARCH__)
|
|
//
|
|
// Homebrew read-modify-write
|
|
//
|
|
if(high_storage_unit_idx_ != low_storage_unit_idx_){
|
|
/// Only need update 2 storage unit at once.
|
|
/// consider misaligned address issue, we need to do atomicCAS twice
|
|
StorageUnit original_low_bits, original_high_bits, update_low_bits, update_high_bits;
|
|
do {
|
|
original_low_bits = ((*ptr_)[low_storage_unit_idx_]);
|
|
update_low_bits = (original_low_bits & kLowUpdateMask) | low_new_bits;
|
|
original_low_bits = atomicCAS(&((*ptr_)[low_storage_unit_idx_]), original_low_bits, update_low_bits);
|
|
} while (update_low_bits != original_low_bits);
|
|
do {
|
|
original_high_bits = ((*ptr_)[high_storage_unit_idx_]);
|
|
update_high_bits = (original_high_bits & kHighUpdateMask) | high_new_bits;
|
|
original_high_bits = atomicCAS(&((*ptr_)[high_storage_unit_idx_]), original_high_bits, update_high_bits);
|
|
} while (update_high_bits != original_high_bits);
|
|
}
|
|
else {
|
|
/// Only need update 1 storage unit.
|
|
StorageUnit original, updated;
|
|
do {
|
|
original = ((*ptr_)[low_storage_unit_idx_]);
|
|
|
|
updated = (original & kLowUpdateMask) | low_new_bits;
|
|
|
|
original = atomicCAS(&((*ptr_)[low_storage_unit_idx_]), original, updated);
|
|
|
|
} while (updated != original);
|
|
}
|
|
#else
|
|
|
|
|
|
StorageUnit update_low_bits = ((*ptr_)[low_storage_unit_idx_] & kLowUpdateMask) | low_new_bits;
|
|
StorageUnit update_high_bits = ((*ptr_)[high_storage_unit_idx_] & kHighUpdateMask) | high_new_bits;
|
|
|
|
(*ptr_)[low_storage_unit_idx_] = update_low_bits;
|
|
|
|
if(low_storage_unit_idx_ != high_storage_unit_idx_)
|
|
(*ptr_)[high_storage_unit_idx_] = update_high_bits;
|
|
#endif
|
|
|
|
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, StorageVec> 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;
|
|
|
|
update_element_status();
|
|
|
|
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;
|
|
|
|
update_element_status();
|
|
|
|
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_;
|
|
}
|
|
|
|
update_element_status();
|
|
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_;
|
|
}
|
|
|
|
update_element_status();
|
|
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 T> using _war = T;
|
|
template <
|
|
typename Element_, /// CUTLASS numeric element type.
|
|
typename Storage_ /// Underlying storage type. Must be able to hold an integer
|
|
>
|
|
class ConstSubbyteReference<Element_, Storage_,
|
|
typename platform::enable_if<sizeof_bits<Storage_>::value % sizeof_bits<Element_>::value != 0>::type> {
|
|
public:
|
|
|
|
using Element = Element_;
|
|
///! Note: Storage unit could not be divisibale by Element,
|
|
/// Type element may be stored across 2 storage units, so need a storage vector to hold integer
|
|
/// number of objects of type Element.
|
|
using StorageUnit = Storage_;
|
|
static int const kBitsStoredVec = cutlass::lcm_cxx11(sizeof_bits<Element>::value, sizeof_bits<StorageUnit>::value);
|
|
static int const kNumStorageUnitPerStoredVec = kBitsStoredVec / sizeof_bits<StorageUnit>::value;
|
|
|
|
using StorageVec = StorageUnit[kNumStorageUnitPerStoredVec];
|
|
using StorageVecPointer = StorageVec const *;
|
|
|
|
using CudaAtomicType = typename platform::conditional<
|
|
sizeof_bits<StorageUnit>::value == 16,
|
|
uint32_t,
|
|
uint64_t
|
|
>::type;
|
|
|
|
static_assert(sizeof_bits<Element>::value <= sizeof_bits<StorageVec>::value,
|
|
"Size of Element must not be greater than StorageVec.");
|
|
|
|
static_assert(!(sizeof_bits<StorageVec>::value % sizeof_bits<Element>::value),
|
|
"StorageVec must be divisible by Element");
|
|
|
|
private:
|
|
|
|
///! Number of elements per storage vector
|
|
int const kElementsPerVector = sizeof_bits<StorageVec>::value / sizeof_bits<Element>::value;
|
|
|
|
///! Bit mask for storage unit.
|
|
StorageUnit const kMask = (StorageUnit(1) << sizeof_bits<Element>::value) - StorageUnit(1);
|
|
|
|
/// Pointer to array containing element
|
|
_war<StorageVecPointer> ptr_;
|
|
|
|
/// Offset (in units of elements) from pointer.
|
|
///
|
|
/// Invariant: must always be in range [0, kElementsPerVector)
|
|
int offset_;
|
|
|
|
/// Element may be stored across 2 storage unit.
|
|
/// Low storage unit index in StorageVec
|
|
/// High storage unit index in StorageVec
|
|
int low_storage_unit_idx_;
|
|
int high_storage_unit_idx_;
|
|
|
|
/// Full Mask to extract the entire element
|
|
uint64_t full_element_mask_;
|
|
|
|
/// Mask to extract the Element from Low storage unit and High storage unit.
|
|
StorageUnit low_storage_mask_;
|
|
StorageUnit high_storage_mask_;
|
|
|
|
/// Start bit index inside the storage unit.
|
|
int start_bit_idx_;
|
|
|
|
private:
|
|
|
|
CUTLASS_HOST_DEVICE
|
|
void update_element_status() {
|
|
int num_bits = offset_ * sizeof_bits<Element>::value;
|
|
|
|
start_bit_idx_ = num_bits % sizeof_bits<StorageUnit>::value;
|
|
|
|
low_storage_unit_idx_ = num_bits / sizeof_bits<StorageUnit>::value;
|
|
high_storage_unit_idx_ = sizeof_bits<StorageUnit>::value - (start_bit_idx_) < sizeof_bits<Element>::value
|
|
? low_storage_unit_idx_ + 1 : low_storage_unit_idx_;
|
|
|
|
full_element_mask_ = uint64_t(kMask) << start_bit_idx_;
|
|
low_storage_mask_ = StorageUnit(full_element_mask_ & ~StorageUnit(0));
|
|
high_storage_mask_ = StorageUnit((full_element_mask_ >> sizeof_bits<StorageUnit>::value) & ~StorageUnit(0));
|
|
}
|
|
|
|
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<StorageVecPointer>(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);
|
|
|
|
update_element_status();
|
|
}
|
|
|
|
/// Constructor
|
|
CUTLASS_HOST_DEVICE
|
|
ConstSubbyteReference(
|
|
Element *ptr = nullptr
|
|
): ConstSubbyteReference(ptr, 0) { }
|
|
|
|
/// Gets storage pointer
|
|
CUTLASS_HOST_DEVICE
|
|
StorageVecPointer 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 {
|
|
StorageUnit low_bits = (*ptr_)[low_storage_unit_idx_] & low_storage_mask_;
|
|
StorageUnit high_bits = low_storage_unit_idx_ != high_storage_unit_idx_ ? (*ptr_)[high_storage_unit_idx_] & high_storage_mask_ : 0;
|
|
|
|
uint64_t full_item = ((uint64_t)high_bits << sizeof_bits<StorageUnit>::value) | low_bits;
|
|
uint8_t result = uint8_t(full_item >> start_bit_idx_);
|
|
|
|
return reinterpret_cast<Element const &>(result);
|
|
}
|
|
|
|
/// 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;
|
|
|
|
update_element_status();
|
|
|
|
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;
|
|
|
|
update_element_status();
|
|
|
|
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_;
|
|
}
|
|
|
|
update_element_status();
|
|
|
|
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_;
|
|
}
|
|
|
|
update_element_status();
|
|
|
|
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, bool subbyte = (sizeof_bits<Element>::value < 8)>
|
|
struct ReferenceFactory;
|
|
|
|
template <typename Element>
|
|
struct ReferenceFactory<Element, false> {
|
|
|
|
///! Number of elements per storage vector
|
|
static int const kElementsPerVector = 1;
|
|
|
|
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];
|
|
}
|
|
|
|
CUTLASS_HOST_DEVICE
|
|
static Element *add_pointer_offset(Element *ptr, int64_t offset) {
|
|
return ptr + offset;
|
|
}
|
|
|
|
CUTLASS_HOST_DEVICE
|
|
static Element const *add_pointer_offset(Element const *ptr, int64_t offset) {
|
|
return ptr + offset;
|
|
}
|
|
};
|
|
|
|
template <typename Element>
|
|
struct ReferenceFactory<Element, true> {
|
|
|
|
//
|
|
// Static methods
|
|
//
|
|
|
|
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);
|
|
}
|
|
|
|
/// Helper to add an offset in number of elements, assuming this offset is divisible
|
|
/// by the vector size.
|
|
CUTLASS_HOST_DEVICE
|
|
static Element *add_pointer_offset(Element *ptr, int64_t offset_in_elements) {
|
|
|
|
return ptr + offset_in_elements * sizeof_bits<Element>::value / sizeof(Element) / 8;
|
|
}
|
|
|
|
/// Helper to add an offset in number of elements, assuming this offset is divisible
|
|
/// by the vector size.
|
|
CUTLASS_HOST_DEVICE
|
|
static Element const *add_pointer_offset(Element const *ptr, int64_t offset_in_elements) {
|
|
|
|
return ptr + offset_in_elements * sizeof_bits<Element>::value / sizeof(Element) / 8;
|
|
}
|
|
};
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
} // namespace cutlass
|