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

releaase 2.11 (#703)

Browse Source
This commit is contained in:
Aditya Atluri
2022-11-19 06:02:15 -08:00
committed by GitHub
parent 3c90f6aea6
commit c975e2ccbb
329 changed files with 47332 additions and 10607 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
include/cutlass/transform/pitch_linear_thread_map.h
View File

@@ -88,20 +88,15 @@ struct PitchLinearStripminedThreadMap {
static_assert(!(Shape::kContiguous % kElementsPerAccess), "");
static_assert(!((Shape::kContiguous * Shape::kStrided) % (kThreads * kElementsPerAccess)),
"Shape must be divisible thread count.");
/// Shape of the tile in units of vectors
using ShapeVec = layout::PitchLinearShape<
Shape::kContiguous / kElementsPerAccess,
Shape::kStrided
>;
static_assert(
(Threads < ShapeVec::kContiguous && !(ShapeVec::kContiguous % kThreads)) ||
(!(kThreads % ShapeVec::kContiguous) && !(ShapeVec::kStrided % (kThreads / ShapeVec::kContiguous))),
"Shape must be divisible by number of iterations of each thread."
);
static_assert((Threads < ShapeVec::kContiguous && !(ShapeVec::kContiguous % kThreads)) ||
(!(kThreads % ShapeVec::kContiguous)),
"Shape must be divisible by number of iterations of each thread.");
};
/// Number of iterations by each thread
@@ -112,11 +107,12 @@ struct PitchLinearStripminedThreadMap {
// Redo the comparison here to work around divide by zero compiler
// error. The compiler evaluates both path of platform::conditional.
(Threads >= Detail::ShapeVec::kContiguous
? Detail::ShapeVec::kStrided /
? (Detail::ShapeVec::kStrided + (kThreads / Detail::ShapeVec::kContiguous - 1)) /
(kThreads / Detail::ShapeVec::kContiguous)
: 0)>,
layout::PitchLinearShape<Detail::ShapeVec::kContiguous / kThreads,
Detail::ShapeVec::kStrided>>::type;
/// Interval between accesses along each dimension of the tensor's logical coordinate space
/// (in units of Elements)
@@ -132,6 +128,13 @@ struct PitchLinearStripminedThreadMap {
>
>::type;
/// Shape of the tile in units of vectors
using StorageShape = typename platform::conditional<
Threads >= Detail::ShapeVec::kContiguous,
layout::PitchLinearShape<Shape::kContiguous,
Iterations::kStrided*(kThreads / Detail::ShapeVec::kContiguous)>,
layout::PitchLinearShape<Shape::kContiguous, Shape::kStrided>>::type;
/// Maps thread ID to a coordinate offset within the tensor's logical coordinate space
/// (in units of Elements)
CUTLASS_HOST_DEVICE

199
include/cutlass/transform/threadblock/ell_iterator.h Normal file
View File

@@ -0,0 +1,199 @@
/***************************************************************************************************
* 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 Ell iterator for matrix of indices (ellColInd matrix)
*/
#pragma once
namespace cutlass {
namespace transform {
namespace threadblock {
namespace ell{
constexpr unsigned int SmemPow = 8;
constexpr unsigned int SmemStages = 2;
constexpr unsigned int SmemSize = 1 << SmemPow;
constexpr unsigned int SmemMask = (SmemSize*SmemStages-1);
class SharedStorage{
public:
Array<int, SmemSize*SmemStages> array;
};
class Iterator{
public:
using Layout = layout::PitchLinear;
using LongIndex = typename Layout::LongIndex;
private:
const int *gmem_col_idx_;
int *smem_col_idx_;
const int block_size_;
const int base_idx_;
const int k_shape_;
const int ell_increment_;
const int array_length_;
int col_idx_base_;
int residue_;
int counter_;
int pow2_;
int residue_shape_;
int smem_offset_;
int smem_stage_;
int gmem_offset_;
int lane_;
bool is_pow2_;
bool is_residue_tile_;
public:
CUTLASS_DEVICE
void load_ell_indices(){
for(int i=threadIdx.x; i<SmemSize; i+=blockDim.x){
int idx = (gmem_offset_+i < array_length_) ? gmem_offset_+i : array_length_-1;
int gmem_col_idx = gmem_col_idx_[idx] - base_idx_;
smem_col_idx_[i + smem_stage_ * SmemSize] =
(gmem_col_idx >= 0) ? gmem_col_idx : -1;
}
gmem_offset_ += SmemSize;
smem_stage_ ^= 1;
}
CUTLASS_DEVICE
Iterator(
SharedStorage& shared_storage_base,
const int* col_idx,
const int& block_size,
const int& base_idx,
const int k_shape,
const int& problem_size_k,
const int& ell_stride,
const int& thread_idx)
: residue_(0),
counter_(0),
smem_offset_(0),
smem_stage_(0),
gmem_offset_(0),
block_size_(block_size),
base_idx_(base_idx),
k_shape_(k_shape),
ell_increment_(ell_stride * block_size),
array_length_((problem_size_k + block_size_ - 1) / block_size_),
residue_shape_(problem_size_k % k_shape_),
is_residue_tile_(residue_shape_ != 0),
smem_col_idx_(reinterpret_cast<int*>(&shared_storage_base.array)),
gmem_col_idx_(const_cast<int*>(col_idx)),
lane_(thread_idx % 32) {
load_ell_indices();
__syncthreads();
is_pow2_ = ((block_size_ & (block_size_ - 1)) == 0);
if( is_pow2_ && k_shape <= block_size_ ) lane_ = 0;
col_idx_base_ = smem_col_idx_[(smem_offset_ + lane_) & SmemMask] * ell_increment_;
pow2_ = 0;
while(block_size_ >> (pow2_ + 1)) ++pow2_;
}
CUTLASS_DEVICE
int get_blocksize(){
return block_size_;
}
CUTLASS_DEVICE
Iterator &operator++(){
if(is_residue_tile_){
residue_ += residue_shape_;
is_residue_tile_ = false;
} else {
residue_ += k_shape_;
}
if(residue_ < block_size_){
return *this;
}
if((array_length_ > SmemSize) && (((smem_offset_ >> SmemPow) & 1) != smem_stage_))
load_ell_indices();
if(residue_ == block_size_){
++smem_offset_;
counter_ += ell_increment_;
residue_ = 0;
col_idx_base_ = smem_col_idx_[(smem_offset_ + lane_) & SmemMask] * ell_increment_ - counter_;
return *this;
}
if(is_pow2_){
smem_offset_ += residue_ >> pow2_;
counter_ += (residue_ >> pow2_) * ell_increment_;
residue_ = residue_ & ((1 << pow2_) - 1);
}
else {
smem_offset_ += residue_ / block_size_;
counter_ += (residue_ / block_size_) * ell_increment_;
residue_ %= block_size_;
}
col_idx_base_ = smem_col_idx_[(smem_offset_ + lane_) & SmemMask] * ell_increment_ - counter_;
return *this;
}
CUTLASS_DEVICE
LongIndex get_offset(const int& idx) {
int num_jump_tiles;
if(is_pow2_)
num_jump_tiles = (idx + residue_) >> pow2_;
else
num_jump_tiles = (idx + residue_) / block_size_;
int tmp = __shfl_sync(0xffffffff, col_idx_base_, num_jump_tiles);
return tmp - num_jump_tiles * ell_increment_;
}
CUTLASS_DEVICE
LongIndex get_offset_fast() {
return col_idx_base_;
}
};
}
}
}
}

1350
include/cutlass/transform/threadblock/ell_predicated_tile_access_iterator.h Normal file
View File

File diff suppressed because it is too large Load Diff

1315
include/cutlass/transform/threadblock/ell_predicated_tile_iterator.h Normal file
View File

File diff suppressed because it is too large Load Diff

77
include/cutlass/transform/threadblock/predicated_tile_access_iterator.h
View File

@@ -221,7 +221,10 @@ class PredicatedTileAccessIteratorPredicates {
set_iteration_index(0);
}
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// Default constructor
PredicatedTileAccessIteratorPredicates() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorPredicates(
@@ -360,9 +363,8 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::PitchLinear,
using Base = PredicatedTileAccessIteratorParams;
// Default ctor
CUTLASS_HOST_DEVICE
Params() { }
/// Default constructor
Params() = default;
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
@@ -388,7 +390,7 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::PitchLinear,
UnderlyingPredicates the_predicates;
/// Parameters object with precomputed internal state
Params const &params_;
Params params_;
/// Internal pointer to first access of tile
BytePointer pointer_;
@@ -419,7 +421,10 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::PitchLinear,
}
public:
/// Default constructor
PredicatedTileAccessIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
@@ -699,9 +704,8 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::ColumnMajor,
public:
/// Default ctor
CUTLASS_HOST_DEVICE
Params() { }
/// Default constructor
Params() = default;
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
@@ -723,6 +727,10 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::ColumnMajor,
UnderlyingIterator iterator_;
public:
/// Default constructor
PredicatedTileAccessIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
@@ -883,9 +891,8 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::RowMajor,
public:
/// Default ctor
CUTLASS_HOST_DEVICE
Params() { }
/// Default constructor
Params() = default;
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
@@ -907,6 +914,10 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::RowMajor,
UnderlyingIterator iterator_;
public:
/// Default constructor
PredicatedTileAccessIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
@@ -1122,7 +1133,7 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::AffineRankN<2>,
//
/// Parameters object with precomputed internal state
Params const &params_;
Params params_;
/// Internal pointer to first access of tile
BytePointer pointer_;
@@ -1144,6 +1155,10 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::AffineRankN<2>,
}
public:
/// Default constructor
PredicatedTileAccessIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
@@ -1371,9 +1386,8 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::AffineRank2ColumnMa
public:
/// Default ctor
CUTLASS_HOST_DEVICE
Params() { }
/// Default constructor
Params() = default;
/// Construct the Params object given an AffineRankN<2> tensor's layout
CUTLASS_HOST_DEVICE
@@ -1390,6 +1404,10 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::AffineRank2ColumnMa
UnderlyingIterator iterator_;
public:
/// Default constructor
PredicatedTileAccessIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
@@ -1550,9 +1568,8 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::AffineRank2RowMajor
public:
/// Default ctor
CUTLASS_HOST_DEVICE
Params() { }
/// Default constructor
Params() = default;
/// Construct the Params object given an AffineRankN<2> tensor's layout
CUTLASS_HOST_DEVICE
@@ -1569,6 +1586,10 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::AffineRank2RowMajor
UnderlyingIterator iterator_;
public:
/// Default constructor
PredicatedTileAccessIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
@@ -1733,8 +1754,9 @@ class PredicatedTileAccessIterator<Shape_, Element_,
typename UnderlyingIterator::Params params_;
public:
CUTLASS_HOST_DEVICE
Params() {}
/// Default constructor
Params() = default;
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
@@ -1755,6 +1777,10 @@ class PredicatedTileAccessIterator<Shape_, Element_,
UnderlyingIterator iterator_;
public:
/// Default constructor
PredicatedTileAccessIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
@@ -1919,8 +1945,9 @@ class PredicatedTileAccessIterator<Shape_, Element_,
typename UnderlyingIterator::Params params_;
public:
CUTLASS_HOST_DEVICE
Params() {}
/// Default constructor
Params() = default;
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
@@ -1941,6 +1968,10 @@ class PredicatedTileAccessIterator<Shape_, Element_,
UnderlyingIterator iterator_;
public:
/// Default constructor
PredicatedTileAccessIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE

100
include/cutlass/transform/threadblock/predicated_tile_iterator.h
View File

@@ -207,13 +207,13 @@ class PredicatedTileIterator<Shape_, Element_, layout::PitchLinear, AdvanceRank,
public:
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout) : params_(layout) { }
CUTLASS_HOST_DEVICE
Params() { }
Params(Layout const &layout) : params_(layout) {}
/// Default constructor
Params() = default;
CUTLASS_HOST_DEVICE
Params(Base const &base)
Params(Base const &base)
: params_(base) {}
};
@@ -230,6 +230,10 @@ class PredicatedTileIterator<Shape_, Element_, layout::PitchLinear, AdvanceRank,
TileAccessIterator address_iterator_;
public:
/// Default constructor
PredicatedTileIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
@@ -457,18 +461,17 @@ public:
typename UnderlyingIterator::Params params_;
public:
CUTLASS_HOST_DEVICE
Params() { }
/// Default constructor
Params() = default;
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {
}
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0)))
{}
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
@@ -484,6 +487,9 @@ private:
public:
/// Default constructor
PredicatedTileIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID
CUTLASS_HOST_DEVICE
PredicatedTileIterator(
@@ -670,16 +676,16 @@ public:
typename UnderlyingIterator::Params params_;
public:
CUTLASS_HOST_DEVICE
Params() { }
/// Default constructor
Params() = default;
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {}
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
@@ -695,6 +701,9 @@ private:
public:
/// Default constructor
PredicatedTileIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID
CUTLASS_HOST_DEVICE
PredicatedTileIterator(
@@ -878,10 +887,10 @@ class PredicatedTileIterator<Shape_, Element_, layout::AffineRankN<2>, AdvanceRa
public:
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout) : params_(layout) { }
CUTLASS_HOST_DEVICE
Params() { }
Params(Layout const &layout) : params_(layout) {}
/// Default constructor
Params() = default;
};
private:
@@ -897,6 +906,10 @@ class PredicatedTileIterator<Shape_, Element_, layout::AffineRankN<2>, AdvanceRa
TileAccessIterator address_iterator_;
public:
/// Default constructor
PredicatedTileIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
@@ -1123,15 +1136,14 @@ public:
typename UnderlyingIterator::Params params_;
public:
CUTLASS_HOST_DEVICE
Params() { }
/// Default constructor
Params() = default;
/// Construct the Params object given an AffineRankN<2> tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout): params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))) {
}
Params(Layout const &layout): params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1)))
{}
};
private:
@@ -1145,6 +1157,9 @@ private:
public:
/// Default constructor
PredicatedTileIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID
CUTLASS_HOST_DEVICE
PredicatedTileIterator(
@@ -1329,9 +1344,9 @@ public:
typename UnderlyingIterator::Params params_;
public:
CUTLASS_HOST_DEVICE
Params() { }
/// Default constructor
Params() = default;
/// Construct the Params object given an AffineRankN<2> tensor's layout
CUTLASS_HOST_DEVICE
@@ -1350,6 +1365,9 @@ private:
public:
/// Default constructor
PredicatedTileIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID
CUTLASS_HOST_DEVICE
PredicatedTileIterator(
@@ -1530,8 +1548,9 @@ class PredicatedTileIterator<Shape_, Element_,
typename UnderlyingIterator::Params params_;
public:
CUTLASS_HOST_DEVICE
Params() {}
/// Default constructor
Params() = default;
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
@@ -1539,7 +1558,7 @@ class PredicatedTileIterator<Shape_, Element_,
: params_(layout::PitchLinear(layout.stride(0))) {}
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
@@ -1553,6 +1572,10 @@ class PredicatedTileIterator<Shape_, Element_,
UnderlyingIterator iterator_;
public:
/// Default constructor
PredicatedTileIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
@@ -1701,7 +1724,7 @@ class PredicatedTileIterator<Shape_, Element_,
using AccessType = typename UnderlyingIterator::AccessType;
/// Fragment object to be loaded or stored
using Fragment = cutlass::Array<Element, ThreadMap::Iterations::kCount *
ThreadMap::kElementsPerAccess>;
@@ -1718,8 +1741,9 @@ class PredicatedTileIterator<Shape_, Element_,
typename UnderlyingIterator::Params params_;
public:
CUTLASS_HOST_DEVICE
Params() {}
/// Default constructor
Params() = default;
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
@@ -1727,7 +1751,7 @@ class PredicatedTileIterator<Shape_, Element_,
: params_(layout::PitchLinear(layout.stride(0))) {}
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
@@ -1740,6 +1764,10 @@ class PredicatedTileIterator<Shape_, Element_,
UnderlyingIterator iterator_;
public:
/// Default constructor
PredicatedTileIterator() = default;
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE

587
include/cutlass/transform/threadblock/regular_tile_access_iterator_pitch_linear_direct_conv.h Normal file
View File

@@ -0,0 +1,587 @@
/***************************************************************************************************
* 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 tiles
from pitch-linear rank=2 tensors.
*/
#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"
#include "cutlass/transform/threadblock/regular_tile_access_iterator.h"
////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace transform {
namespace threadblock {
////////////////////////////////////////////////////////////////////////////////
template <typename Shape, typename Element, typename Layout, int AdvanceRank,
typename ThreadMap,
bool Dynamic_iterations = false,
int Alignment =
sizeof_bits<Element>::value* ThreadMap::kElementsPerAccess / 8
>
class RegularTileAccessIteratorDirectConv;
////////////////////////////////////////////////////////////////////////////////
/// Tile iterator specialized for congruous arrangements for TensorOps with dynamic_iterations OFF
///
///
/// Satisfies: ForwardTileIteratorConcept |
/// ReadableContiguousTileIteratorConcept |
/// WriteableContiguousTileIteratorConcept
///
template <typename Shape_, typename Element_, int AdvanceRank,
typename ThreadMap_, int Alignment>
class RegularTileAccessIteratorDirectConv<
Shape_, Element_,
layout::PitchLinear,
AdvanceRank, ThreadMap_, false, Alignment> {
public:
static_assert(
AdvanceRank == 0 || AdvanceRank == 1,
"Specialization for pitch-linear iterator may along advance along the "
"contiguous(rank=0) or strided(rank=1) dimension.");
using Shape = Shape_;
using Element = Element_;
using Layout = layout::PitchLinear;
static int const kAdvanceRank = AdvanceRank;
static int const kAlignment = Alignment;
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
using ThreadMap = ThreadMap_;
/// Element type per access
using AccessType = Array<Element, ThreadMap::kElementsPerAccess>;
private:
//
// Data members
//
/// Stride value
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType *pointer_;
/// Internal byte offset
Index byte_offset_;
/// Iteration in the contiguous dimension
int iteration_contiguous_;
/// Iteration in the strided dimension
int iteration_strided_;
public:
/// Construct a TileIterator with zero threadblock offset
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv(TensorRef ref, ///< Pointer to start of tensor
int thread_id ///< ID of each participating thread
)
: stride_(ref.stride(0) / ThreadMap::kElementsPerAccess),
byte_offset_(0) {
layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id);
// initialize pointer
pointer_ = reinterpret_cast<AccessType *>(ref.data() + ref.offset(thread_offset_base));
set_iteration_index(0);
}
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_index(int index) {
iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous;
iteration_strided_ = index / ThreadMap::Iterations::kContiguous;
}
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_num(int num) {
//Do nothing
}
/// 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 {
AccessType *access_ptr = pointer_;
int access_offset = iteration_strided_ * ThreadMap::Delta::kStrided * stride_ +
iteration_contiguous_ * ThreadMap::Delta::kContiguous /
ThreadMap::kElementsPerAccess;
char *access_byte_ptr =
reinterpret_cast<char *>(access_ptr + access_offset);
return reinterpret_cast<AccessType *>(access_byte_ptr + byte_offset_);
}
/// Advances to the next tile in memory.
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv &operator++() {
++iteration_contiguous_;
if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous)
return *this;
// Enter here only if (iteration_contiguous_ ==
// ThreadMap::Iteration::kContiguous)
iteration_contiguous_ = 0;
++iteration_strided_;
if (iteration_strided_ < ThreadMap::Iterations::kStrided) {
return *this;
}
// Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided)
// which means we enter the next tile.
iteration_strided_ = 0;
return *this;
}
/// Advances to the next tile in memory.
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv operator++(int) {
RegularTileAccessIteratorDirectConv prev(*this);
this->operator++();
return prev;
}
/// Adds a tile offset in the unit of tile.
CUTLASS_DEVICE
void add_tile_offset(TensorCoord const &coord) {
add_pointer_offset(coord.contiguous() * Shape::kContiguous +
coord.strided() * ThreadMap::Iterations::kStrided *
ThreadMap::Delta::kStrided * stride_ * ThreadMap::kElementsPerAccess);
}
};
////////////////////////////////////////////////////////////////////////////////
/// Tile iterator specialized for congruous arrangements for TensorOps with dynamic_iterations ON
///
///
/// Satisfies: ForwardTileIteratorConcept |
/// ReadableContiguousTileIteratorConcept |
/// WriteableContiguousTileIteratorConcept
///
template <typename Shape_, typename Element_, int AdvanceRank,
typename ThreadMap_, int Alignment>
class RegularTileAccessIteratorDirectConv<
Shape_, Element_,
layout::PitchLinear,
AdvanceRank, ThreadMap_,true, Alignment> {
public:
static_assert(
AdvanceRank == 0 || AdvanceRank == 1,
"Specialization for pitch-linear iterator may along advance along the "
"contiguous(rank=0) or strided(rank=1) dimension.");
using Shape = Shape_;
using Element = Element_;
using Layout = layout::PitchLinear;
static int const kAdvanceRank = AdvanceRank;
static int const kAlignment = Alignment;
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
using ThreadMap = ThreadMap_;
/// Element type per access
using AccessType = Array<Element, ThreadMap::kElementsPerAccess>;
private:
//
// Data members
//
/// Stride value
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType *pointer_;
/// Internal byte offset
Index byte_offset_;
/// Iteration in the contiguous dimension
int iteration_contiguous_;
/// Iteration in the strided dimension
int iteration_strided_;
/// Total iterattions in the strided dimension: Dynamic value
int total_iteration_strided_;
public:
/// Construct a TileIterator with zero threadblock offset
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv(TensorRef ref, ///< Pointer to start of tensor
int thread_id ///< ID of each participating thread
)
: stride_(ref.stride(0) / ThreadMap::kElementsPerAccess),
byte_offset_(0) {
layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id);
// initialize pointer
pointer_ = reinterpret_cast<AccessType *>(ref.data() + ref.offset(thread_offset_base));
set_iteration_index(0);
}
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_index(int index) {
iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous;
iteration_strided_ = index / ThreadMap::Iterations::kContiguous;
}
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_num(int num) {
total_iteration_strided_ = num;
}
/// 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 {
AccessType *access_ptr = pointer_;
int access_offset = iteration_strided_ * ThreadMap::Delta::kStrided * stride_ +
iteration_contiguous_ * ThreadMap::Delta::kContiguous /
ThreadMap::kElementsPerAccess;
char *access_byte_ptr =
reinterpret_cast<char *>(access_ptr + access_offset);
return reinterpret_cast<AccessType *>(access_byte_ptr + byte_offset_);
}
/// Advances to the next tile in memory.
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv &operator++() {
++iteration_contiguous_;
if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous)
return *this;
// Enter here only if (iteration_contiguous_ ==
// ThreadMap::Iteration::kContiguous)
iteration_contiguous_ = 0;
++iteration_strided_;
if (iteration_strided_ < total_iteration_strided_) {
return *this;
}
// Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided)
// which means we enter the next tile.
iteration_strided_ = 0;
return *this;
}
/// Advances to the next tile in memory.
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv operator++(int) {
RegularTileAccessIteratorDirectConv prev(*this);
this->operator++();
return prev;
}
/// Adds a tile offset in the unit of tile.
CUTLASS_DEVICE
void add_tile_offset(TensorCoord const &coord) {
add_pointer_offset(coord.contiguous() * Shape::kContiguous +
coord.strided() * total_iteration_strided_ * ThreadMap::Delta::kStrided * stride_ *
ThreadMap::kElementsPerAccess);
}
};
////////////////////////////////////////////////////////////////////////////////
/// Tile iterator specialized for column major layouts
///
///
/// Satisfies: ForwardTileIteratorConcept |
/// ReadableContiguousTileIteratorConcept |
/// WriteableContiguousTileIteratorConcept
///
template <typename Shape_, typename Element_, int AdvanceRank,
typename ThreadMap_,bool Dynamic_iterations, int Alignment >
class RegularTileAccessIteratorDirectConv<
Shape_, Element_,
layout::ColumnMajor,
AdvanceRank, ThreadMap_, Dynamic_iterations , Alignment> {
public:
static_assert(
AdvanceRank == 0 || AdvanceRank == 1,
"Specialization for pitch-linear iterator may along advance along the "
"contiguous(rank=0) or strided(rank=1) dimension.");
using Shape = Shape_;
using Element = Element_;
using Layout = layout::ColumnMajor;
static int const kAdvanceRank = AdvanceRank;
static int const kAlignment = Alignment;
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
using ThreadMap = ThreadMap_;
/// Underlying iterator type
using UnderlyingIterator = RegularTileAccessIteratorDirectConv<
layout::PitchLinearShape<Shape::kRow, Shape::kColumn>, Element,
layout::PitchLinear,
(kAdvanceRank == 0 ? 0 : 1),
ThreadMap_,
Dynamic_iterations>;
using AccessType = typename UnderlyingIterator::AccessType;
private:
/// Underlying iterator
UnderlyingIterator iterator_;
public:
/// Construct a TileIterator with zero threadblock offset
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv(TensorRef ref, ///< Pointer to start of tensor
int thread_id ///< ID of each participating thread
)
: iterator_({ref.data(), ref.stride()}, thread_id) {}
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_index(int index) { iterator_.set_iteration_index(index); }
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_num(int num) {
iterator_.set_iteration_num(num);
}
/// 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.row(), coord.column()});
}
/// Advances to the next tile in memory.
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv &operator++() {
++iterator_;
return *this;
}
/// Advances to the next tile in memory.
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv operator++(int) {
RegularTileAccessIteratorDirectConv prev(*this);
++iterator_;
return prev;
}
};
////////////////////////////////////////////////////////////////////////////////
/// Tile iterator specialized for row major layouts
///
///
/// Satisfies: ForwardTileIteratorConcept |
/// ReadableContiguousTileIteratorConcept |
/// WriteableContiguousTileIteratorConcept
///
template <typename Shape_, typename Element_, int AdvanceRank,
typename ThreadMap_,bool Dynamic_iterations, int Alignment>
class RegularTileAccessIteratorDirectConv<
Shape_, Element_,
layout::RowMajor,
AdvanceRank, ThreadMap_, Dynamic_iterations, Alignment> {
public:
static_assert(
AdvanceRank == 0 || AdvanceRank == 1,
"Specialization for pitch-linear iterator may along advance along the "
"contiguous(rank=0) or strided(rank=1) dimension.");
using Shape = Shape_;
using Element = Element_;
using Layout = layout::RowMajor;
static int const kAdvanceRank = AdvanceRank;
static int const kAlignment = Alignment;
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
using ThreadMap = ThreadMap_;
/// Underlying iterator type
using UnderlyingIterator = RegularTileAccessIteratorDirectConv<
layout::PitchLinearShape<Shape::kColumn, Shape::kRow>, Element,
layout::PitchLinear,
(kAdvanceRank == 0 ? 1 : 0),
ThreadMap_,
Dynamic_iterations>;
using AccessType = typename UnderlyingIterator::AccessType;
private:
/// Underlying iterator
UnderlyingIterator iterator_;
public:
/// Construct a TileIterator with zero threadblock offset
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv(TensorRef ref, ///< Pointer to start of tensor
int thread_id ///< ID of each participating thread
)
: iterator_({ref.data(), ref.stride()}, thread_id) {}
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_index(int index) { iterator_.set_iteration_index(index); }
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_num(int num) {
iterator_.set_iteration_num(num);
}
/// 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
RegularTileAccessIteratorDirectConv &operator++() {
++iterator_;
return *this;
}
/// Advances to the next tile in memory.
CUTLASS_HOST_DEVICE
RegularTileAccessIteratorDirectConv operator++(int) {
RegularTileAccessIteratorDirectConv prev(*this);
++iterator_;
return prev;
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace threadblock
} // namespace transform
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////