ANIKET SHIVAM
@@ -0,0 +1,375 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 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 Templates calculating the address and predicates to the load of scale and bias vectors.
|
||||
|
||||
This iterator uses masks to guard out-of-bounds accesses.
|
||||
|
||||
It can be used to load the gamma and beta vectors of layernorm which is loop variant.
|
||||
|
||||
A precomputed "Params" object minimizes the amount of state that must be
|
||||
stored in registers, and integer addition is used to advance the pointer
|
||||
through memory.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/coord.h"
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/layout/pitch_linear.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/predicate_vector.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "cutlass/tensor_view.h"
|
||||
#include "cutlass/conv/threadblock/conv2d_params.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace transform {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// PredicatedScaleBiasVectorAccessIterator
|
||||
///
|
||||
template <typename ThreadblockShape,
|
||||
typename Element,
|
||||
typename Layout>
|
||||
class PredicatedScaleBiasVectorAccessIterator;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for fprop pitch-linear data.
|
||||
///
|
||||
template <typename ThreadblockShape_, typename Element_>
|
||||
class PredicatedScaleBiasVectorAccessIterator<ThreadblockShape_,
|
||||
Element_,
|
||||
layout::PitchLinear> {
|
||||
public:
|
||||
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using Element = Element_;
|
||||
using Layout = layout::PitchLinear;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using TensorView = TensorView<Element, Layout>;
|
||||
using TensorCoord = typename Layout::TensorCoord;
|
||||
|
||||
using ConstPointer = const Element *;
|
||||
using NonConstPointer = typename platform::remove_const<Element>::type *;
|
||||
|
||||
static int const kElementsPerAccess = 128 / sizeof_bits<Element>::value;
|
||||
static int const kThreads = ThreadblockShape::kContiguous / kElementsPerAccess;
|
||||
|
||||
using AccessType = AlignedArray<Element, kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
/// Internal pointer type permits fast address arithmetic
|
||||
using BytePointer = char *;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
BytePointer pointer_;
|
||||
|
||||
TensorCoord thread_offset_;
|
||||
|
||||
int problem_size_k_;
|
||||
|
||||
/// Used for out-of-order visitation
|
||||
bool is_residue_tile_;
|
||||
|
||||
bool guard_;
|
||||
|
||||
TensorCoord::Index residue_size_;
|
||||
|
||||
public:
|
||||
/// Constructs a TileIterator from its precomputed state, threadblock offset,
|
||||
/// and thread ID
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIterator(
|
||||
/// Extent of tensor
|
||||
int problem_size_k,
|
||||
/// Pointer to the start of the scale vector
|
||||
ConstPointer scale_pointer,
|
||||
/// Pointer to the start of the bias vector
|
||||
ConstPointer bias_pointer,
|
||||
/// ID of each participating thread
|
||||
int thread_id,
|
||||
/// Initial offset of threadblock
|
||||
TensorCoord const &threadblock_offset) {
|
||||
pointer_ = (thread_id < kThreads)
|
||||
? reinterpret_cast<BytePointer>(
|
||||
const_cast<NonConstPointer>(scale_pointer))
|
||||
: reinterpret_cast<BytePointer>(
|
||||
const_cast<NonConstPointer>(bias_pointer));
|
||||
|
||||
// Per-thread offset in logical coordinates of tensor
|
||||
int thread_base = (thread_id < kThreads) ? 0 : kThreads;
|
||||
|
||||
problem_size_k_ = problem_size_k;
|
||||
|
||||
is_residue_tile_ = true;
|
||||
|
||||
residue_size_ = (problem_size_k_ - threadblock_offset.contiguous()) % ThreadblockShape::kContiguous;
|
||||
|
||||
if (residue_size_ == 0) {
|
||||
residue_size_ = ThreadblockShape::kContiguous;
|
||||
}
|
||||
|
||||
guard_ = ((thread_id - thread_base) * kElementsPerAccess) < residue_size_;
|
||||
|
||||
thread_offset_ =
|
||||
threadblock_offset +
|
||||
TensorCoord((thread_id - thread_base) * kElementsPerAccess, 0);
|
||||
|
||||
set_iteration_index(0);
|
||||
}
|
||||
|
||||
/// Construct a PredicatedTileAccessIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIterator(
|
||||
/// Extent of tensor
|
||||
int problem_size_k,
|
||||
/// Pointer to start of scale vector
|
||||
ConstPointer scale_pointer,
|
||||
/// Pointer to start of scale vector
|
||||
ConstPointer bias_pointer,
|
||||
///< ID of each participating thread
|
||||
int thread_id)
|
||||
: PredicatedScaleBiasVectorAccessIterator(problem_size_k,
|
||||
scale_pointer, bias_pointer,
|
||||
thread_id, make_Coord(0, 0)) {}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(int index) {}
|
||||
|
||||
/// Advances an iterator along logical dimensions of matrix in units of whole threadblock tiles
|
||||
CUTLASS_DEVICE
|
||||
void add_tile_offset(
|
||||
TensorCoord const &tile_offset) {
|
||||
|
||||
guard_ = threadIdx.x < kThreads * 2;
|
||||
|
||||
TensorCoord offset = is_residue_tile_ ?
|
||||
TensorCoord(residue_size_ + ThreadblockShape::kContiguous * (tile_offset.contiguous() - 1), 0)
|
||||
: TensorCoord(ThreadblockShape::kContiguous * tile_offset.contiguous(), 0);
|
||||
|
||||
thread_offset_ =
|
||||
thread_offset_ +
|
||||
offset;
|
||||
|
||||
is_residue_tile_ = false;
|
||||
}
|
||||
|
||||
/// Returns a pointer
|
||||
CUTLASS_HOST_DEVICE
|
||||
AccessType *get() const {
|
||||
|
||||
return reinterpret_cast<AccessType *>(
|
||||
pointer_ +
|
||||
(thread_offset_.contiguous() * sizeof_bits<Element>::value / 8));
|
||||
}
|
||||
|
||||
/// Increment and return an instance to self.
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIterator &operator++() {
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Increment and return an instance to self.
|
||||
CUTLASS_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIterator operator++(int) {
|
||||
PredicatedScaleBiasVectorAccessIterator self(*this);
|
||||
operator++();
|
||||
return self;
|
||||
}
|
||||
|
||||
/// Clears the predicate set efficiently
|
||||
CUTLASS_HOST_DEVICE
|
||||
void clear_mask(bool enable = true) {
|
||||
guard_ &= (!enable);
|
||||
}
|
||||
|
||||
/// Returns whether access is valid or not
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool valid() {
|
||||
return guard_;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for row-major data.
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
/// WriteableContiguousTileIteratorConcept |
|
||||
/// MaskedTileIteratorConcept
|
||||
///
|
||||
template <typename ThreadblockShape_,
|
||||
typename Element_>
|
||||
class PredicatedScaleBiasVectorAccessIterator<ThreadblockShape_,
|
||||
Element_,
|
||||
layout::RowMajor> {
|
||||
public:
|
||||
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using Element = Element_;
|
||||
using Layout = layout::RowMajor;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using TensorView = TensorView<Element, Layout>;
|
||||
using TensorCoord = typename Layout::TensorCoord;
|
||||
|
||||
using ConstPointer = const Element *;
|
||||
using NonConstPointer = typename platform::remove_const<Element>::type *;
|
||||
|
||||
using UnderlyingIterator = PredicatedScaleBiasVectorAccessIterator<
|
||||
layout::PitchLinearShape<ThreadblockShape::kColumn, ThreadblockShape::kRow>,
|
||||
Element,
|
||||
layout::PitchLinear>;
|
||||
|
||||
using AccessType = typename UnderlyingIterator::AccessType;
|
||||
static int const kElementsPerAccess = UnderlyingIterator::kElementsPerAccess;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Underlying pitch-linear tile iterator
|
||||
UnderlyingIterator iterator_;
|
||||
|
||||
public:
|
||||
/// Constructs a TileIterator from its precomputed state, threadblock offset,
|
||||
/// and thread ID
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIterator(
|
||||
///< Extent of tensor
|
||||
int problem_size_k,
|
||||
///< Pointer to the start of the scale vector
|
||||
ConstPointer scale_pointer,
|
||||
///< Pointer to the start of the bias vector
|
||||
ConstPointer bias_pointer,
|
||||
///< ID of each participating thread
|
||||
int thread_id,
|
||||
///< Initial offset of threadblock
|
||||
TensorCoord const &threadblock_offset)
|
||||
: iterator_(problem_size_k, scale_pointer, bias_pointer,
|
||||
thread_id,
|
||||
layout::PitchLinearCoord(threadblock_offset.column(),
|
||||
threadblock_offset.row())) {}
|
||||
|
||||
/// Construct a PredicatedTileAccessIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIterator(
|
||||
int problem_size_k, ///< Extent of tensor
|
||||
ConstPointer scale_pointer, ///< Pointer to the start of the scale vector
|
||||
ConstPointer bias_pointer, ///< Pointer to the start of the bias vector
|
||||
int thread_id ///< ID of each participating thread
|
||||
)
|
||||
: PredicatedScaleBiasVectorAccessIterator(problem_size_k,
|
||||
scale_pointer, bias_pointer,
|
||||
thread_id, make_Coord(0, 0)) {}
|
||||
|
||||
/// Advances an iterator along logical dimensions of matrix in units of whole
|
||||
/// threadblock tiles
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_tile_offset(TensorCoord const &tile_offset) {
|
||||
iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()});
|
||||
}
|
||||
|
||||
/// Returns a pointer
|
||||
CUTLASS_HOST_DEVICE
|
||||
AccessType *get() const {
|
||||
return reinterpret_cast<AccessType *>(iterator_.get());
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
///
|
||||
/// The first time this method is called, predicates are updated, and the
|
||||
/// iterator's internal pointer is reverted to the first "steady state" tile.
|
||||
/// Subsequent calls are lightweight and must only update the internal
|
||||
/// pointer.
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIterator &operator++() {
|
||||
++iterator_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
///
|
||||
/// The first time this method is called, predicates are updated, and the
|
||||
/// iterator's internal pointer is reverted to the first "steady state" tile.
|
||||
/// Subsequent calls are lightweight and must only update the internal
|
||||
/// pointer.
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIterator operator++(int) {
|
||||
PredicatedScaleBiasVectorAccessIterator self(*this);
|
||||
operator++();
|
||||
return self;
|
||||
}
|
||||
|
||||
/// Clears the predicate set efficiently
|
||||
CUTLASS_HOST_DEVICE
|
||||
void clear_mask(bool enable = true) {
|
||||
iterator_.clear_mask(enable);
|
||||
}
|
||||
|
||||
/// Returns whether access is valid or not
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool valid() {
|
||||
return iterator_.valid();
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace transform
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,328 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 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 Templates calculating the address and predicates to the load of scale and bias vectors.
|
||||
|
||||
This iterator uses masks to guard out-of-bounds accesses.
|
||||
|
||||
This can be used to load var and mean vectors in layernorm which is loop invariant.
|
||||
|
||||
A precomputed "Params" object minimizes the amount of state that must be
|
||||
stored in registers, and integer addition is used to advance the pointer
|
||||
through memory.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/coord.h"
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/layout/pitch_linear.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/predicate_vector.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "cutlass/tensor_view.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace transform {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// PredicatedScaleBiasVectorIterator
|
||||
///
|
||||
template <typename WarpShape,
|
||||
typename Element,
|
||||
typename Layout>
|
||||
class PredicatedScaleBiasVectorIterator;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileIterator for wgrad pitch-linear data.
|
||||
///
|
||||
template <typename WarpShape_, typename Element_>
|
||||
class PredicatedScaleBiasVectorIterator<WarpShape_,
|
||||
Element_,
|
||||
layout::PitchLinear> {
|
||||
public:
|
||||
|
||||
using WarpShape = WarpShape_;
|
||||
using Element = Element_;
|
||||
using Layout = layout::PitchLinear;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using TensorView = TensorView<Element, Layout>;
|
||||
using TensorCoord = typename Layout::TensorCoord;
|
||||
|
||||
using ConstPointer = const Element *;
|
||||
using NonConstPointer = typename platform::remove_const<Element>::type *;
|
||||
|
||||
static int const kElementsPerAccess = 1;
|
||||
|
||||
using AccessType = AlignedArray<Element, kElementsPerAccess>;
|
||||
|
||||
static int const kIterations = WarpShape::kContiguous / 8;
|
||||
|
||||
/// Fragment object to be loaded or stored
|
||||
using Fragment = cutlass::Array<__half2, 2 * kIterations * kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
ConstPointer scale_pointer_;
|
||||
ConstPointer bias_pointer_;
|
||||
|
||||
/// Size of tensor
|
||||
int problem_size_;
|
||||
|
||||
int32_t thread_offset_;
|
||||
|
||||
public:
|
||||
/// Constructs a TileIterator from its precomputed state, threadblock offset,
|
||||
/// and thread ID
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorIterator(
|
||||
/// Extent of tensor
|
||||
int problem_size,
|
||||
/// Pointer to the start of the scale vector
|
||||
ConstPointer scale_pointer,
|
||||
/// Pointer to the start of the bias vector
|
||||
ConstPointer bias_pointer,
|
||||
/// ID of each participating thread
|
||||
int thread_id,
|
||||
/// Initial offset of threadblock
|
||||
TensorCoord const &threadblock_offset)
|
||||
: problem_size_(problem_size),
|
||||
scale_pointer_(scale_pointer),
|
||||
bias_pointer_(bias_pointer) {
|
||||
|
||||
thread_offset_ = threadblock_offset.contiguous() + (thread_id % 32) / 4;
|
||||
}
|
||||
|
||||
/// Construct a PredicatedTileIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorIterator(
|
||||
/// Extent of tensor
|
||||
int problem_size,
|
||||
/// Pointer to start of scale vector
|
||||
ConstPointer scale_pointer,
|
||||
/// Pointer to start of scale vector
|
||||
ConstPointer bias_pointer,
|
||||
///< ID of each participating thread
|
||||
int thread_id)
|
||||
: PredicatedScaleBiasVectorIterator(problem_size,
|
||||
scale_pointer, bias_pointer,
|
||||
thread_id, make_Coord(0, 0)) {}
|
||||
|
||||
/// Advances an iterator along logical dimensions of matrix in units of whole warp tiles
|
||||
CUTLASS_DEVICE
|
||||
void add_tile_offset(
|
||||
TensorCoord const &tile_offset) {
|
||||
|
||||
thread_offset_ += (WarpShape::kContiguous * tile_offset.contiguous());
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
|
||||
|
||||
frag.fill(__float2half2_rn(0.0f));
|
||||
__half2 *frag_ptr = reinterpret_cast<__half2 *>(&frag);
|
||||
|
||||
// load scale
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int c = 0; c < kIterations; ++c) {
|
||||
|
||||
cutlass::arch::global_load<
|
||||
__half,
|
||||
sizeof(AccessType)
|
||||
>(
|
||||
frag_ptr[c * 2].x,
|
||||
scale_pointer_ + thread_offset_ + c * 8,
|
||||
(thread_offset_ + c * 8) < problem_size_
|
||||
);
|
||||
}
|
||||
|
||||
// load bias
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int c = 0; c < kIterations; ++c) {
|
||||
|
||||
cutlass::arch::global_load<
|
||||
__half,
|
||||
sizeof(AccessType)
|
||||
>(
|
||||
frag_ptr[c * 2 + 1].x,
|
||||
bias_pointer_ + thread_offset_ + c * 8,
|
||||
(thread_offset_ + c * 8) < problem_size_
|
||||
);
|
||||
}
|
||||
|
||||
// duplicate scale
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int c = 0; c < kIterations; ++c) {
|
||||
frag_ptr[c * 2].y = frag_ptr[c * 2].x;
|
||||
}
|
||||
|
||||
// duplicate bias
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int c = 0; c < kIterations; ++c) {
|
||||
frag_ptr[c * 2 + 1].y = frag_ptr[c * 2 + 1].x;
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
load_with_pointer_offset(frag, 0);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileIterator for row-major data.
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
/// WriteableContiguousTileIteratorConcept |
|
||||
/// MaskedTileIteratorConcept
|
||||
///
|
||||
template <typename WarpShape_,
|
||||
typename Element_>
|
||||
class PredicatedScaleBiasVectorIterator<WarpShape_,
|
||||
Element_,
|
||||
layout::RowMajor> {
|
||||
public:
|
||||
|
||||
using WarpShape = WarpShape_;
|
||||
using Element = Element_;
|
||||
using Layout = layout::RowMajor;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using TensorView = TensorView<Element, Layout>;
|
||||
using TensorCoord = typename Layout::TensorCoord;
|
||||
|
||||
using ConstPointer = const Element *;
|
||||
using NonConstPointer = typename platform::remove_const<Element>::type *;
|
||||
|
||||
using UnderlyingIterator = PredicatedScaleBiasVectorIterator<
|
||||
layout::PitchLinearShape<WarpShape::kColumn, WarpShape::kRow>,
|
||||
Element,
|
||||
layout::PitchLinear>;
|
||||
|
||||
using AccessType = typename UnderlyingIterator::AccessType;
|
||||
static int const kElementsPerAccess = UnderlyingIterator::kElementsPerAccess;
|
||||
using Fragment = typename UnderlyingIterator::Fragment;
|
||||
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Underlying pitch-linear tile iterator
|
||||
UnderlyingIterator iterator_;
|
||||
|
||||
public:
|
||||
/// Constructs a TileIterator from its precomputed state, threadblock offset,
|
||||
/// and thread ID
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorIterator(
|
||||
///< Extent of tensor
|
||||
int problem_size,
|
||||
///< Pointer to the start of the scale vector
|
||||
ConstPointer scale_pointer,
|
||||
///< Pointer to the start of the bias vector
|
||||
ConstPointer bias_pointer,
|
||||
///< ID of each participating thread
|
||||
int thread_id,
|
||||
///< Initial offset of threadblock
|
||||
TensorCoord const &threadblock_offset)
|
||||
: iterator_(problem_size, scale_pointer, bias_pointer,
|
||||
thread_id,
|
||||
layout::PitchLinearCoord(threadblock_offset.column(),
|
||||
threadblock_offset.row())) {}
|
||||
|
||||
/// Construct a PredicatedTileIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorIterator(
|
||||
int problem_size, ///< Extent of tensor
|
||||
ConstPointer scale_pointer, ///< Pointer to the start of the scale vector
|
||||
ConstPointer bias_pointer, ///< Pointer to the start of the bias vector
|
||||
int thread_id ///< ID of each participating thread
|
||||
)
|
||||
: PredicatedScaleBiasVectorIterator(problem_size,
|
||||
scale_pointer, bias_pointer,
|
||||
thread_id, make_Coord(0, 0)) {}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(int index) { iterator_.set_iteration_index(index); }
|
||||
|
||||
/// Advances an iterator along logical dimensions of matrix in units of whole
|
||||
/// threadblock tiles
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_tile_offset(TensorCoord const &tile_offset) {
|
||||
iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()});
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
|
||||
iterator_.load_with_pointer_offset(frag, pointer_offset);
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
iterator_.load(frag);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace transform
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,253 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 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 Templates implementing computing the addresses of storing of small
|
||||
scale and bias vectors in the shared memory.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/pitch_linear.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/matrix_coord.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace transform {
|
||||
namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// RegularScaleBiasVectorAccessIterator
|
||||
///
|
||||
template <typename Shape, typename Element, typename Layout>
|
||||
class RegularScaleBiasVectorAccessIterator;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tile iterator specialized for congruous arrangements for TensorOps
|
||||
///
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
/// WriteableContiguousTileIteratorConcept
|
||||
///
|
||||
template <typename Shape_, typename Element_>
|
||||
class RegularScaleBiasVectorAccessIterator<Shape_, Element_, layout::PitchLinear> {
|
||||
public:
|
||||
|
||||
using Shape = Shape_;
|
||||
using Element = Element_;
|
||||
using Layout = layout::PitchLinear;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using TensorCoord = typename Layout::TensorCoord;
|
||||
|
||||
/// Element type per access
|
||||
static int const kElementsPerAccess = 128 / sizeof_bits<Element>::value;
|
||||
static int const kThreads = Shape::kContiguous / kElementsPerAccess;
|
||||
using AccessType = Array<Element, kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Internal pointer
|
||||
AccessType *pointer_;
|
||||
|
||||
/// Internal byte offset
|
||||
Index byte_offset_;
|
||||
|
||||
public:
|
||||
/// Construct a TileIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularScaleBiasVectorAccessIterator(
|
||||
TensorRef scale_bias_ref, ///< Pointer to the start of the scale and bias
|
||||
///< vector
|
||||
int thread_id ///< ID of each participating thread
|
||||
)
|
||||
: byte_offset_(0) {
|
||||
// Per-thread offset in logical coordinates of tensor
|
||||
int thread_offset = thread_id * kElementsPerAccess;
|
||||
|
||||
// initialize pointer
|
||||
pointer_ =
|
||||
reinterpret_cast<AccessType *>(scale_bias_ref.data() + thread_offset);
|
||||
|
||||
set_iteration_index(0);
|
||||
}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(int index) {}
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
byte_offset_ += pointer_offset * sizeof(Element);
|
||||
}
|
||||
|
||||
/// Returns a pointer
|
||||
CUTLASS_DEVICE
|
||||
AccessType *get() const {
|
||||
|
||||
char *access_byte_ptr =
|
||||
reinterpret_cast<char *>(pointer_);
|
||||
|
||||
return reinterpret_cast<AccessType *>(access_byte_ptr + byte_offset_);
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularScaleBiasVectorAccessIterator &operator++() { return *this; }
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularScaleBiasVectorAccessIterator operator++(int) {
|
||||
RegularScaleBiasVectorAccessIterator prev(*this);
|
||||
this->operator++();
|
||||
|
||||
return prev;
|
||||
}
|
||||
|
||||
/// Adds a tile offset in the unit of tile.
|
||||
CUTLASS_DEVICE
|
||||
void add_tile_offset(TensorCoord const &coord) {
|
||||
// Multiply by 2 because we store scale and bias belong to the same stage
|
||||
// next to each other.
|
||||
add_pointer_offset(coord.contiguous() * Shape::kContiguous * 2);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tile iterator specialized for row major layouts
|
||||
///
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
/// WriteableContiguousTileIteratorConcept
|
||||
///
|
||||
template <typename Shape_, typename Element_>
|
||||
class RegularScaleBiasVectorAccessIterator<
|
||||
Shape_, Element_,
|
||||
layout::RowMajor> {
|
||||
public:
|
||||
|
||||
using Shape = Shape_;
|
||||
using Element = Element_;
|
||||
using Layout = layout::RowMajor;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using TensorCoord = typename Layout::TensorCoord;
|
||||
|
||||
/// Underlying iterator type
|
||||
using UnderlyingIterator = RegularScaleBiasVectorAccessIterator<
|
||||
layout::PitchLinearShape<Shape::kColumn, Shape::kRow>, Element,
|
||||
layout::PitchLinear>;
|
||||
|
||||
using AccessType = typename UnderlyingIterator::AccessType;
|
||||
|
||||
private:
|
||||
|
||||
/// Underlying iterator
|
||||
UnderlyingIterator iterator_;
|
||||
|
||||
public:
|
||||
/// Construct a TileIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularScaleBiasVectorAccessIterator(
|
||||
TensorRef scale_bias_ref, ///< Pointer to the start of the scale and bias
|
||||
///< vector
|
||||
int thread_id ///< ID of each participating thread
|
||||
)
|
||||
: iterator_({scale_bias_ref.data(), scale_bias_ref.stride()}, thread_id) {
|
||||
}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(int index) { iterator_.set_iteration_index(index); }
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
iterator_.add_pointer_offset(pointer_offset);
|
||||
}
|
||||
|
||||
/// Returns a pointer
|
||||
CUTLASS_HOST_DEVICE
|
||||
AccessType *get() const {
|
||||
return reinterpret_cast<AccessType *>(iterator_.get());
|
||||
}
|
||||
|
||||
/// Adds a tile offset
|
||||
CUTLASS_DEVICE
|
||||
void add_tile_offset(TensorCoord const &coord) {
|
||||
iterator_.add_tile_offset({coord.column(), coord.row()});
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularScaleBiasVectorAccessIterator &operator++() {
|
||||
++iterator_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularScaleBiasVectorAccessIterator operator++(int) {
|
||||
RegularScaleBiasVectorAccessIterator prev(*this);
|
||||
++iterator_;
|
||||
|
||||
return prev;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace transform
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
Reference in New Issue
Block a user