youngkingdom/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

CUTLASS 3.3.0 (#1167)

Browse Source 
* Release 3.3.0

Adds support for mixed precision GEMMs On Hopper and Ampere
Adds support for < 16B aligned GEMMs on Hopper
Enhancements to EVT
Enhancements to Python interface
Enhancements to Sub-byte type handling in CuTe
Several other bug-fixes and performance improvements.

* minor doc update
This commit is contained in:
Pradeep Ramani
2023-11-02 08:09:05 -07:00
committed by GitHub
parent 922fb5108b
commit c008b4aea8
263 changed files with 16214 additions and 5008 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
include/cute/container/array.hpp
View File

@ -41,13 +41,14 @@ namespace cute
template <class T, size_t N>
struct array
{
using value_type = T;
using element_type = T;
using value_type = remove_cv_t<T>;
using size_type = size_t;
using difference_type = ptrdiff_t;
using reference = value_type&;
using const_reference = const value_type&;
using pointer = value_type*;
using const_pointer = const value_type*;
using reference = element_type&;
using const_reference = const element_type&;
using pointer = element_type*;
using const_pointer = const element_type*;
using iterator = pointer;
using const_iterator = const_pointer;
@ -190,20 +191,21 @@ struct array
}
}
value_type __elems_[N > 0 ? N : 1];
element_type __elems_[N];
};
template <class T>
struct array<T, 0>
{
using value_type = T;
using element_type = T;
using value_type = remove_cv_t<T>;
using size_type = size_t;
using difference_type = ptrdiff_t;
using reference = value_type&;
using const_reference = const value_type&;
using pointer = value_type*;
using const_pointer = const value_type*;
using reference = element_type&;
using const_reference = const element_type&;
using pointer = element_type*;
using const_pointer = const element_type*;
using const_iterator = const_pointer;
using iterator = const_iterator;

271
include/cute/container/array_subbyte.hpp
View File

@ -39,11 +39,18 @@
#include <cute/numeric/int.hpp> // sizeof_bits
#include <cute/numeric/integral_constant.hpp>
#include <cute/container/bit_field.hpp> // dummy_type
namespace cute
{
template <class T>
struct is_subbyte {
static constexpr bool value = sizeof_bits_v<T> < 8;
};
template <class T>
constexpr bool is_subbyte_v = is_subbyte<T>::value;
//
// Underlying subbyte storage type
//
@ -53,43 +60,44 @@ using subbyte_storage_type_t = conditional_t<(sizeof_bits_v<T> <= 8), uint8_t,
conditional_t<(sizeof_bits_v<T> <= 32), uint32_t,
conditional_t<(sizeof_bits_v<T> <= 64), uint64_t,
conditional_t<(sizeof_bits_v<T> <= 128), uint128_t,
dummy_type>>>>>;
T>>>>>;
template <class T>
struct subbyte_iterator;
template <class T> struct subbyte_iterator;
template <class, class> struct swizzle_ptr;
//
// subbyte_reference
// Proxy object for sub-byte element references
//
template <class T>
struct subbyte_reference
struct subbyte_reference
{
// Iterator Element type (const or non-const)
using element_type = T;
// Iterator Value type without type qulifier.
// Iterator Value type without type qualifier.
using value_type = remove_cv_t<T>;
// Storage type (const or non-const)
using storage_type = conditional_t<(is_const_v<T>), subbyte_storage_type_t<T> const, subbyte_storage_type_t<T>>;
static_assert(!is_same_v<storage_type, dummy_type>, "Storage type is not supported");
static_assert(sizeof_bits_v<storage_type> % 8 == 0, "Storage type is not supported");
static_assert(sizeof_bits_v<element_type> <= sizeof_bits_v<storage_type>,
"Size of Element must not be greater than Storage.");
// Number of logical elements per stored object
static constexpr uint8_t ElementsPerStoredItem = sizeof_bits_v<storage_type> / sizeof_bits_v<element_type>;
// Bitmask for covering one item
static constexpr storage_type BitMask = storage_type((storage_type(1) << sizeof_bits_v<element_type>) - 1);
private:
// Bitmask for covering one item
static constexpr storage_type BitMask = storage_type(storage_type(-1) >> (sizeof_bits_v<storage_type> - sizeof_bits_v<element_type>));
// Flag for fast branching on straddled elements
static constexpr bool is_storage_unaligned = ((sizeof_bits_v<storage_type> % sizeof_bits_v<element_type>) != 0);
friend class subbyte_iterator<T>;
// Pointer to storage element
storage_type* ptr_ = nullptr;
// Index into elements packed into storage_type element. RI: 0 <= idx_ < ElementsPerStoredItem
// Bit index of value_type starting position within storage_type element.
// RI: 0 <= idx_ < sizeof_bit<storage_type>
uint8_t idx_ = 0;
// Ctor
@ -100,38 +108,73 @@ private:
public:
// Copy Ctor
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
subbyte_reference(subbyte_reference const& other) {
*this = element_type(other);
}
// Copy Assignment
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
subbyte_reference& operator=(subbyte_reference const& other) {
return *this = element_type(other);
}
// Dtor
~subbyte_reference() = default;
// Assignment
template<class T_=element_type>
template <class T_ = element_type>
CUTE_HOST_DEVICE constexpr
enable_if_t<!is_const_v<T_>, subbyte_reference&> operator=(element_type x) {
enable_if_t<!is_const_v<T_>, subbyte_reference&> operator=(element_type x)
{
static_assert(is_same_v<T_, element_type>, "Do not specify template arguments!");
storage_type item = (reinterpret_cast<storage_type const &>(x) & BitMask);
storage_type kUpdateMask = storage_type(~(BitMask << (idx_ * sizeof_bits_v<element_type>)));
*ptr_ = storage_type((*ptr_ & kUpdateMask) | (item << (idx_ * sizeof_bits_v<element_type>)));
storage_type item = (reinterpret_cast<storage_type const&>(x) & BitMask);
// Update the current storage element
storage_type bit_mask_0 = storage_type(BitMask << idx_);
ptr_[0] = storage_type((ptr_[0] & ~bit_mask_0) | (item << idx_));
// If value_type is unaligned with storage_type (static) and this is a straddled value (dynamic)
if (is_storage_unaligned && idx_ + sizeof_bits_v<value_type> > sizeof_bits_v<storage_type>) {
uint8_t straddle_bits = uint8_t(sizeof_bits_v<storage_type> - idx_);
storage_type bit_mask_1 = storage_type(BitMask >> straddle_bits);
// Update the next storage element
ptr_[1] = storage_type((ptr_[1] & ~bit_mask_1) | (item >> straddle_bits));
}
return *this;
}
// Comparison of referenced values
CUTE_HOST_DEVICE constexpr friend
bool operator==(subbyte_reference const& x, subbyte_reference const& y) { return x.get() == y.get(); }
CUTE_HOST_DEVICE constexpr friend
bool operator!=(subbyte_reference const& x, subbyte_reference const& y) { return x.get() != y.get(); }
CUTE_HOST_DEVICE constexpr friend
bool operator< (subbyte_reference const& x, subbyte_reference const& y) { return x.get() < y.get(); }
CUTE_HOST_DEVICE constexpr friend
bool operator> (subbyte_reference const& x, subbyte_reference const& y) { return x.get() > y.get(); }
CUTE_HOST_DEVICE constexpr friend
bool operator<=(subbyte_reference const& x, subbyte_reference const& y) { return x.get() <= y.get(); }
CUTE_HOST_DEVICE constexpr friend
bool operator>=(subbyte_reference const& x, subbyte_reference const& y) { return x.get() >= y.get(); }
// Value
CUTE_HOST_DEVICE
element_type get() const {
element_type get() const
{
if constexpr (is_same_v<bool, value_type>) { // Extract to bool -- potentially faster impl
return bool((*ptr_) & (BitMask << (idx_ * sizeof_bits_v<element_type>)));
return bool((*ptr_) & (BitMask << idx_));
} else { // Extract to element_type
storage_type item = storage_type((*ptr_ >> (idx_ * sizeof_bits_v<element_type>)) & BitMask);
return reinterpret_cast<element_type &>(item);
// Extract from the current storage element
auto item = storage_type((ptr_[0] >> idx_) & BitMask);
// If value_type is unaligned with storage_type (static) and this is a straddled value (dynamic)
if (is_storage_unaligned && idx_ + sizeof_bits_v<value_type> > sizeof_bits_v<storage_type>) {
uint8_t straddle_bits = uint8_t(sizeof_bits_v<storage_type> - idx_);
storage_type bit_mask_1 = storage_type(BitMask >> straddle_bits);
// Extract from the next storage element
item |= storage_type((ptr_[1] & bit_mask_1) << straddle_bits);
}
return reinterpret_cast<element_type&>(item);
}
}
@ -142,63 +185,77 @@ public:
}
};
//
// subbyte_iterator
// Random-access iterator over subbyte references
//
template <class T>
struct subbyte_iterator
struct subbyte_iterator
{
// Iterator Element type (const or non-const)
using element_type = T;
// Iterator Value type without type qulifier.
// Iterator Value type without type qualifier.
using value_type = remove_cv_t<T>;
// Storage type (const or non-const)
using storage_type = conditional_t<(is_const_v<T>), subbyte_storage_type_t<T> const, subbyte_storage_type_t<T>>;
// Reference proxy type
using reference = subbyte_reference<element_type>;
static_assert(!is_same_v<storage_type, dummy_type>, "Storage type is not supported");
static_assert(sizeof_bits_v<storage_type> % 8 == 0, "Storage type is not supported");
static_assert(sizeof_bits_v<element_type> <= sizeof_bits_v<storage_type>,
"Size of Element must not be greater than Storage.");
// Number of logical elements per stored object
static constexpr uint8_t ElementsPerStoredItem = sizeof_bits_v<storage_type> / sizeof_bits_v<element_type>;
private:
template <class, class> friend class swizzle_ptr;
// Pointer to storage element
storage_type* ptr_ = nullptr;
// Index into elements packed into storage_type element. RI: 0 <= idx_ < ElementsPerStoredItem
// Bit index of value_type starting position within storage_type element.
// RI: 0 <= idx_ < sizeof_bit<storage_type>
uint8_t idx_ = 0;
public:
// Ctor
subbyte_iterator() = default;
// Ctor
template <class PointerType>
CUTE_HOST_DEVICE constexpr
subbyte_iterator(PointerType* ptr, uint8_t idx = 0): ptr_(reinterpret_cast<storage_type*>(ptr)), idx_(idx) { }
subbyte_iterator(PointerType* ptr, uint8_t idx = 0) : ptr_(reinterpret_cast<storage_type*>(ptr)), idx_(idx) { }
subbyte_iterator() = default;
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator++() {
++idx_;
if (idx_ == ElementsPerStoredItem) {
++ptr_;
idx_ = 0;
}
reference operator*() const {
return reference(ptr_, idx_);
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator+=(uint64_t k) {
k = sizeof_bits_v<value_type> * k + idx_;
ptr_ += k / sizeof_bits_v<storage_type>;
idx_ = k % sizeof_bits_v<storage_type>;
return *this;
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator--() {
if (idx_) {
--idx_;
} else {
--ptr_;
idx_ = ElementsPerStoredItem - 1;
subbyte_iterator operator+(uint64_t k) const {
return subbyte_iterator(ptr_, idx_) += k;
}
CUTE_HOST_DEVICE constexpr
reference operator[](uint64_t k) const {
return *(*this + k);
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator++() {
idx_ += sizeof_bits_v<value_type>;
if (idx_ >= sizeof_bits_v<storage_type>) {
++ptr_;
idx_ -= sizeof_bits_v<storage_type>;
}
return *this;
}
@ -210,6 +267,17 @@ public:
return ret;
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator--() {
if (idx_ >= sizeof_bits_v<value_type>) {
idx_ -= sizeof_bits_v<value_type>;
} else {
--ptr_;
idx_ += sizeof_bits_v<storage_type> - sizeof_bits_v<value_type>;
}
return *this;
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator operator--(int) {
subbyte_iterator ret(*this);
@ -217,37 +285,45 @@ public:
return ret;
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator& operator+=(uint64_t k) {
k += idx_;
ptr_ += k / ElementsPerStoredItem;
idx_ = k % ElementsPerStoredItem;
return *this;
}
CUTE_HOST_DEVICE constexpr
subbyte_iterator operator+(uint64_t k) const {
return subbyte_iterator(ptr_,idx_) += k;
}
CUTE_HOST_DEVICE constexpr
reference operator*() const {
return reference(ptr_, idx_);
}
CUTE_HOST_DEVICE constexpr
reference operator[](uint64_t k) const {
return *(*this + k);
}
CUTE_HOST_DEVICE constexpr
friend bool operator==(subbyte_iterator const& x, subbyte_iterator const& y) {
CUTE_HOST_DEVICE constexpr friend
bool operator==(subbyte_iterator const& x, subbyte_iterator const& y) {
return x.ptr_ == y.ptr_ && x.idx_ == y.idx_;
}
CUTE_HOST_DEVICE constexpr friend
bool operator< (subbyte_iterator const& x, subbyte_iterator const& y) {
return x.ptr_ < y.ptr_ || (x.ptr_ == y.ptr_ && x.idx_ < y.idx_);
}
CUTE_HOST_DEVICE constexpr friend
bool operator!=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x == y); }
CUTE_HOST_DEVICE constexpr friend
bool operator<=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(y < x); }
CUTE_HOST_DEVICE constexpr friend
bool operator> (subbyte_iterator const& x, subbyte_iterator const& y) { return (y < x); }
CUTE_HOST_DEVICE constexpr friend
bool operator>=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x < y); }
CUTE_HOST_DEVICE constexpr
friend bool operator!=(subbyte_iterator const& x, subbyte_iterator const& y) {
return !(x == y);
// Conversion to raw pointer with loss of subbyte index
CUTE_HOST_DEVICE constexpr friend
T* raw_pointer_cast(subbyte_iterator const& x) {
assert(x.idx_ == 0);
return reinterpret_cast<T*>(x.ptr_);
}
// Conversion to NewT_ with possible loss of subbyte index
template <class NewT_>
CUTE_HOST_DEVICE constexpr friend
auto recast_ptr(subbyte_iterator const& x) {
using NewT = conditional_t<(is_const_v<T>), NewT_ const, NewT_>;
if constexpr (is_subbyte<NewT>::value) { // Making subbyte_iter, preserve the subbyte idx
return subbyte_iterator<NewT>(x.ptr_, x.idx_);
} else { // Not subbyte, assume/assert subbyte idx 0
return reinterpret_cast<NewT*>(raw_pointer_cast(x));
}
CUTE_GCC_UNREACHABLE;
}
CUTE_HOST_DEVICE friend void print(subbyte_iterator x) {
printf("subptr[%db](%p.%u)", int(sizeof_bits<T>::value), x.ptr_, x.idx_);
}
};
@ -281,26 +357,20 @@ struct array_subbyte
// Storage type (const or non-const)
using storage_type = conditional_t<(is_const_v<T>), subbyte_storage_type_t<T> const, subbyte_storage_type_t<T>>;
static_assert(!is_same_v<storage_type, dummy_type>, "Storage type is not supported");
// Number of logical elements per stored object
static constexpr uint8_t ElementsPerStoredItem = sizeof_bits_v<storage_type> / sizeof_bits_v<T>;
// Bitmask for covering one item
static constexpr storage_type BitMask = ((storage_type(1) << sizeof_bits<T>::value) - 1);
// Number of storage elements
static constexpr size_type StorageElements = (N + ElementsPerStoredItem - 1) / ElementsPerStoredItem;
static_assert(sizeof_bits_v<storage_type> % 8 == 0, "Storage type is not supported");
private:
// Number of storage elements, ceil_div
static constexpr size_type StorageElements = (N * sizeof_bits_v<value_type> + sizeof_bits_v<storage_type> - 1) / sizeof_bits_v<storage_type>;
// Internal storage
storage_type storage[StorageElements];
public:
CUTE_HOST_DEVICE constexpr
array_subbyte() { }
array_subbyte() {}
CUTE_HOST_DEVICE constexpr
array_subbyte(array_subbyte const& x) {
@ -334,20 +404,11 @@ public:
}
}
// Efficient fill method
CUTE_HOST_DEVICE constexpr
void fill(T const& value) {
storage_type item = (reinterpret_cast<storage_type const&>(value) & BitMask);
// Reproduce the value over the bits of the storage item
CUTE_UNROLL
for (size_type s = sizeof_bits_v<T>; s < sizeof_bits_v<storage_type>; s *= 2) {
item |= item << s;
}
CUTE_UNROLL
for (size_type i = 0; i < StorageElements; ++i) {
storage[i] = item;
for (size_type i = 0; i < N; ++i) {
at(i) = value;
}
}
@ -428,12 +489,12 @@ public:
CUTE_HOST_DEVICE constexpr
iterator end() {
return iterator(storage + N / ElementsPerStoredItem, N % ElementsPerStoredItem);
return iterator(storage) + N;
}
CUTE_HOST_DEVICE constexpr
const_iterator end() const {
return const_iterator(storage + N / ElementsPerStoredItem, N % ElementsPerStoredItem);
return const_iterator(storage) + N;
}
CUTE_HOST_DEVICE constexpr
@ -509,6 +570,12 @@ T&& get(array_subbyte<T,N>&& a)
namespace CUTE_STL_NAMESPACE
{
template <class T>
struct is_reference<cute::subbyte_reference<T>>
: CUTE_STL_NAMESPACE::true_type
{};
template <class T, size_t N>
struct tuple_size<cute::array_subbyte<T,N>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, N>

12
include/cute/container/bit_field.hpp
View File

@ -72,16 +72,10 @@ struct bit_field
// Number of bits in data_[idx] used for NumBits if straddling, else 0
static constexpr uint32_t bit_hi = (idx + 1 < N) ? (storage_type_bits - bit_lo) : 0;
private:
// MSVC issues warning C4293 ("shift count negative or too big, undefined behavior")
// if we use NumBits directly in the shift expression, even if the shift occurs
// in the branch of a ternary expression where NumBits is known to be less than
// the number of bits of the value being shifted.
static constexpr uint32_t MollifiedNumBits = NumBits > 63u ? 63u : NumBits;
public:
// NumBits mask
static constexpr value_type mask = (NumBits < 64u) ? ((uint64_t(1) << MollifiedNumBits) - 1) : uint64_t(-1);
static constexpr value_type mask = value_type(uint64_t(-1) >> (64u - NumBits));
// NumBits mask for BitStart
static constexpr storage_type mask_lo = storage_type(mask) << bit_lo;
// NumBits mask for leftover bits in data_[idx+1] if straddling, else 0
@ -93,7 +87,7 @@ public:
CUTE_HOST_DEVICE constexpr
value_type get() const {
storage_type result = (data_[idx] & mask_lo) >> bit_lo;
if constexpr (bit_hi) {
if constexpr (bit_hi != 0) {
result |= (data_[idx+1] & mask_hi) << bit_hi;
}
return static_cast<value_type>(result);
@ -104,7 +98,7 @@ public:
void set(value_type x) {
storage_type item = static_cast<storage_type>(x & mask);
data_[idx] = static_cast<storage_type>((data_[idx] & ~mask_lo) | (item << bit_lo));
if constexpr (bit_hi) {
if constexpr (bit_hi != 0) {
data_[idx+1] = static_cast<storage_type>((data_[idx+1] & ~mask_hi) | (item >> bit_hi));
}
}