diff options
Diffstat (limited to 'cutlass-example/cutlass/fragment.h')
| -rw-r--r-- | cutlass-example/cutlass/fragment.h | 278 |
1 files changed, 278 insertions, 0 deletions
diff --git a/cutlass-example/cutlass/fragment.h b/cutlass-example/cutlass/fragment.h new file mode 100644 index 0000000..886b114 --- /dev/null +++ b/cutlass-example/cutlass/fragment.h @@ -0,0 +1,278 @@ +/*************************************************************************************************** + * Copyright (c) 2017-2018, 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 Fragment, a statically-sized array for storing parts of matrices within a + thread's registers. +*/ +#pragma once + +#include <assert.h> +#include <cutlass/shape.h> +#include <cutlass/util/cutlass_math.h> +#include <cutlass/vector.h> + +namespace cutlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/*!@defgroup fragment_concept Fragment Concept +@{ + +\ref fragment_concept is a statically sized array for storing parts of tiles held by individual CUDA +threads. + +@par \ref fragment_concept + Types satisfying \ref fragment_concept define the following members + - <b>Element</b> - type of each access held within the fragment + - <b>kElements</b> - number of elements stored by the fragment + - <b>clear()</b> - overwrites the fragment storage with zeros + - <b>Element & operator[](int i)</b> - by-reference access of the ith element + - <b>Element const & operator[](int i) const</b> - const by-reference access of the ith element +@} +*/ + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/*!@defgroup fragment_iterator_concept Fragment Iterator Concept +@{ + +\ref fragment_iterator_concept provides structured access to the elements within a fragment with an +optional bitcast to the desired access type + +@par \ref fragment_iterator_concept + Types satisfying \ref fragment_iterator_concept define the following members + - <b>AccessType& operator[](int i)</b> - provides access to the ith element of the fragment + - <b>AccessType& at(int d, int h, int w, int c)</b> - applies \ref layout_concept to fragment and +provides access to element at (d, h, w, c) + +@} +*/ + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <int kAlignment_> +struct StorageType { + typedef uint64_t Type; +}; +template <> +struct StorageType<4> { + typedef uint32_t Type; +}; +template <> +struct StorageType<2> { + typedef uint16_t Type; +}; +template <> +struct StorageType<1> { + typedef uint8_t Type; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/** +* @brief A template defining \ref fragment_concept +* @concept{fragment_concept} +*/ +template <typename Element_, int kElements_, size_t kAlignment_ = 16> +struct Fragment : public AlignedStruct<kAlignment_> { + /// Make sure the alignment makes sense wrt the size of elements. + static_assert(kAlignment_ == 16 || kAlignment_ >= sizeof(Element_), "Alignment is too small"); + /// Alignment must be a power of two + static_assert(is_pow2<kAlignment_>::value, "Alignment must be a power of two"); + + /// This class. + typedef Fragment<Element_, kElements_> This_; + /// The element. + typedef Element_ Element; + /// The number of elements. + static int const kElements = kElements_; + + /// Clear a fragment. + CUTLASS_DEVICE void clear() { + // Avoid element-wise access for sub 32b element type + if (kAlignment_ >= 8 && (kElements * sizeof(Element)) % 8 == 0) { + uint64_t* ptr = reinterpret_cast<uint64_t*>(storage); + for (int i = 0; i < (kElements * sizeof(Element)) / 8; ++i) { + ptr[i] = uint64_t(0); + } + } else if (kAlignment_ >= 4 && (kElements * sizeof(Element)) % 4 == 0) { + uint32_t* ptr = reinterpret_cast<uint32_t*>(storage); + for (int i = 0; i < (kElements * sizeof(Element)) / 4; ++i) { + ptr[i] = uint32_t(0); + } + } else if (kAlignment_ >= 2 && (kElements * sizeof(Element)) % 2 == 0) { + uint16_t* ptr = reinterpret_cast<uint16_t*>(storage); + for (int i = 0; i < (kElements * sizeof(Element)) / 2; ++i) { + ptr[i] = uint16_t(0); + } + } else { + for (int i = 0; i < kElements; ++i) { + storage[i] = 0; + } + } + } + + /// The accessor. + CUTLASS_DEVICE Element& operator[](int i) { + assert(i < kElements_); + return reinterpret_cast<Element*>(storage)[i]; + } + + /// The accessor. + CUTLASS_DEVICE Element const& operator[](int i) const { + assert(i < kElements_); + return reinterpret_cast<Element const*>(storage)[i]; + } + + private: + /// Storage type to use for Elements + typedef typename StorageType<kAlignment_>::Type StorageType; + + /// Number of elements in the storage + static int const kStorageCount = + (sizeof(Element_) * kElements_ + sizeof(StorageType) - 1) / sizeof(StorageType); + /// The storage. + StorageType storage[kStorageCount]; + + /// Ensure that there's enough storage for all elements + static_assert(sizeof(StorageType) <= kAlignment_, "StorageType is too big for given alignment"); +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/** +* @brief A template defining \ref fragment_iterator_concept +* @concept{fragment_iterator_concept} +*/ +template <typename Fragment_, typename Iterations_, typename AccessType_> +struct FragmentIterator { + /// This class. + typedef FragmentIterator<Fragment_, Iterations_, AccessType_> This_; + /// The fragment. + typedef Fragment_ Fragment; + /// The number of iterations. + typedef Iterations_ Iterations; + /// The access type. + typedef AccessType_ AccessType; + + /// The element. + typedef typename Fragment::Element Element; + /// The number of elements per access. + static int const kElementsPerAccess = (int)(sizeof(AccessType) / sizeof(Element)); + /// The shape of the the fragment. + typedef typename ShapeMul<Iterations, Shape<1, 1, 1, kElementsPerAccess> >::Shape FragmentShape; + /// The linear strides for iterations. + typedef typename ShapeStrides<FragmentShape, kElementsPerAccess>::Shape Strides; + + /// Ctor. + template <typename OtherFragment_> + CUTLASS_DEVICE FragmentIterator(OtherFragment_& fragment, int offset = 0) + : pointer(reinterpret_cast<Element*>(&fragment[offset])) { + static_assert(OtherFragment_::kElements >= Fragment::kElements, ""); + } + + /// The accessor. + CUTLASS_DEVICE AccessType const& at(int d, int h, int w, int c = 0) const { + int const imm = ComputeOffsetFromStrides<Strides>::get(d, h, w, c); + return reinterpret_cast<AccessType const&>(pointer[imm]); + } + + /// The accessor. + CUTLASS_DEVICE AccessType& at(int d, int h, int w, int c = 0) { + int const imm = ComputeOffsetFromStrides<Strides>::get(d, h, w, c); + return reinterpret_cast<AccessType&>(pointer[imm]); + } + + /// The accessor. + CUTLASS_DEVICE AccessType const& operator[](int i) const { + return reinterpret_cast<AccessType const&>(pointer[i * kElementsPerAccess]); + } + + /// The accessor. + CUTLASS_DEVICE AccessType& operator[](int i) { + return reinterpret_cast<AccessType&>(pointer[i * kElementsPerAccess]); + } + + /// Is the iterator valid? + CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const { return true; } + + /// The pointer. + Element* pointer; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Fragment_, typename Iterations_, typename AccessType_> +struct FragmentConstIterator { + /// This class. + typedef FragmentIterator<Fragment_, Iterations_, AccessType_> This_; + /// The fragment. + typedef Fragment_ Fragment; + /// The number of iterations. + typedef Iterations_ Iterations; + /// The access type. + typedef AccessType_ AccessType; + + /// The element. + typedef typename Fragment::Element Element; + /// The number of elements per access. + static int const kElementsPerAccess = (int)(sizeof(AccessType) / sizeof(Element)); + /// The shape of the the fragment. + typedef typename ShapeMul<Iterations, Shape<1, 1, 1, kElementsPerAccess> >::Shape FragmentShape; + /// The linear strides for iterations. + typedef typename ShapeStrides<FragmentShape, kElementsPerAccess>::Shape IterationsStrides; + + /// Ctor. + template <typename OtherFragment_> + CUTLASS_DEVICE FragmentConstIterator(OtherFragment_& fragment, int offset = 0) + : pointer(reinterpret_cast<Element const*>(&fragment[offset])) { + static_assert(OtherFragment_::kElements >= Fragment::kElements, ""); + } + /// Create from non-constant FragmentIterator + CUTLASS_DEVICE FragmentConstIterator( + FragmentIterator<Fragment_, Iterations_, AccessType_> const& rhs_) + : pointer(reinterpret_cast<Element const*>(rhs_.offset)) {} + + /// The accessor. + CUTLASS_DEVICE AccessType const& at(int d, int h, int w, int c = 0) const { + int const imm = ComputeOffsetFromStrides<IterationsStrides>::get(d, h, w, c); + return reinterpret_cast<AccessType const&>(pointer[imm]); + } + + /// The accessor. + CUTLASS_DEVICE AccessType const& operator[](int i) const { + return reinterpret_cast<AccessType const&>(pointer[i * kElementsPerAccess]); + } + + /// Is the iterator valid? + CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const { return true; } + + /// The pointer. + Element const* pointer; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass |