@@ -517,6 +517,9 @@ struct TransposePitchLinearThreadMap {
|
||||
layout::PitchLinearShape<ThreadMap::Iterations::kStrided,
|
||||
ThreadMap::Iterations::kContiguous>;
|
||||
|
||||
static_assert(Iterations::kContiguous == 1,
|
||||
"Contiguous iteration has to be one to reuse the same shared store function with those that don't need transpose");
|
||||
|
||||
static_assert(Iterations::kCount, "Number of iterations must be non-zero");
|
||||
|
||||
///< Delta betweeen accesses (units of elements, concept: PitchLinearShape)
|
||||
@@ -595,6 +598,9 @@ struct TransposePitchLinearThreadMapSimt {
|
||||
|
||||
static_assert(Iterations::kCount, "Number of iterations must be non-zero");
|
||||
|
||||
static_assert(Iterations::kStrided == 1,
|
||||
"Strided iteration has to be one to reuse the same shared store function with those that don't need transpose");
|
||||
|
||||
/// Shape of access by each thread
|
||||
using ThreadAccessShape = typename ThreadMap::ThreadAccessShape;
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
* Copyright (c) 2017-2021, 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:
|
||||
@@ -26,6 +26,7 @@
|
||||
/*! \file
|
||||
\brief Basic copy routines for tensor views
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
namespace cutlass {
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -46,6 +46,7 @@
|
||||
#include "cutlass/predicate_vector.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "cutlass/tensor_view.h"
|
||||
#include "cutlass/transform/threadblock/predicated_tile_access_iterator_params.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -86,6 +87,7 @@ class PredicatedTileAccessIterator2dThreadTile<Shape_, Element_, layout::PitchLi
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
using StrideIndex = typename Layout::Stride::Index;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using TensorView = TensorView<Element, Layout>;
|
||||
@@ -108,51 +110,31 @@ class PredicatedTileAccessIterator2dThreadTile<Shape_, Element_, layout::PitchLi
|
||||
/// Predicate vector stores mask to guard accesses
|
||||
using Mask = Array<uint32_t, kPredicateWordCount>;
|
||||
|
||||
/// Parameters object is precomputed state and is host-constructible
|
||||
class Params {
|
||||
/// Uses a non-template class
|
||||
struct Params : PredicatedTileAccessIteratorParams {
|
||||
|
||||
public:
|
||||
friend PredicatedTileAccessIterator2dThreadTile;
|
||||
|
||||
private:
|
||||
/// stride of pitch-linear layout (units of Element)
|
||||
int stride_;
|
||||
/// amount (in byte) to increment pointer to move to next access along
|
||||
/// strided dimension
|
||||
int inc_strided_;
|
||||
/// amount (in byte) to increment pointer from last access to first access
|
||||
/// of next tile
|
||||
int inc_next_;
|
||||
/// amount (in byte) to increment pointer from first access of current tile
|
||||
/// to first access of next tile
|
||||
int inc_advance_;
|
||||
|
||||
public:
|
||||
using Base = PredicatedTileAccessIteratorParams;
|
||||
|
||||
// Default ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(): stride_(0), inc_strided_(0), inc_next_(0), inc_advance_(0) { }
|
||||
Params() { }
|
||||
|
||||
/// Construct the Params object given a pitch-linear tensor's layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Layout const &layout) : stride_(layout.stride(0)) {
|
||||
Params(Layout const &layout) :
|
||||
Base(layout.stride(0),
|
||||
MakePredicatedTileAccessIteratorDesc<Shape, Element, Layout, kAdvanceRank, ThreadMap>()()
|
||||
) { }
|
||||
|
||||
inc_strided_ =
|
||||
(stride_ * ThreadMap::Delta::kStrided) * int(sizeof(Element));
|
||||
|
||||
if (kAdvanceRank) {
|
||||
// advance along strided dimension
|
||||
inc_advance_ = Shape::kStrided * stride_ * int(sizeof(Element));
|
||||
} else {
|
||||
// advance along contiguous dimension
|
||||
inc_advance_ = Shape::kContiguous * int(sizeof(Element));
|
||||
}
|
||||
|
||||
inc_next_ = inc_advance_ - (ThreadMap::Iterations::kStrided - 1) *
|
||||
ThreadMap::Delta::kStrided * stride_ *
|
||||
int(sizeof(Element));
|
||||
};
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Base const &base) :
|
||||
Base(base) { }
|
||||
};
|
||||
|
||||
|
||||
private:
|
||||
/// Internal pointer type permits fast address arithmetic
|
||||
using BytePointer = char *;
|
||||
@@ -537,7 +519,12 @@ class PredicatedTileAccessIterator2dThreadTile<Shape_, Element_, layout::ColumnM
|
||||
/// 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::PitchLinear(layout.stride(0))){}
|
||||
|
||||
/// Construct the Params object given a pitch-linear tensor's layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(typename UnderlyingIterator::Params::Base const &base)
|
||||
: params_(base) {}
|
||||
};
|
||||
|
||||
private:
|
||||
@@ -711,7 +698,12 @@ class PredicatedTileAccessIterator2dThreadTile<Shape_, Element_, layout::RowMajo
|
||||
/// 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::PitchLinear(layout.stride(0))){}
|
||||
|
||||
/// Construct the Params object given a pitch-linear tensor's layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(typename UnderlyingIterator::Params::Base const &base)
|
||||
: params_(base) {}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
276
include/cutlass/transform/threadblock/predicated_tile_access_iterator_params.h
Executable file
276
include/cutlass/transform/threadblock/predicated_tile_access_iterator_params.h
Executable file
@@ -0,0 +1,276 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017-2021, 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
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/layout/pitch_linear.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace transform {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Predicated tile access iterator descriptor object containing template dependent state
|
||||
struct PredicatedTileAccessIteratorDesc {
|
||||
|
||||
int element_size_bits;
|
||||
int advance_rank;
|
||||
layout::PitchLinearCoord threadblock_shape;
|
||||
layout::PitchLinearCoord threadmap_iterations;
|
||||
layout::PitchLinearCoord threadmap_delta;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileAccessIteratorDesc() { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileAccessIteratorDesc(
|
||||
int element_size_bits_,
|
||||
int advance_rank_,
|
||||
layout::PitchLinearCoord threadblock_shape_,
|
||||
layout::PitchLinearCoord threadmap_iterations_,
|
||||
layout::PitchLinearCoord threadmap_delta_
|
||||
):
|
||||
element_size_bits(element_size_bits_),
|
||||
advance_rank(advance_rank_),
|
||||
threadblock_shape(threadblock_shape_),
|
||||
threadmap_iterations(threadmap_iterations_),
|
||||
threadmap_delta(threadmap_delta_) { }
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/// Helper template to construct an PredicatedTileAccessIteratorDesc from a template
|
||||
// dependent state
|
||||
template <
|
||||
typename Shape, typename Element, typename Layout,
|
||||
int AdvanceRank, typename ThreadMap>
|
||||
struct MakePredicatedTileAccessIteratorDesc;
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for pitch-linear data.
|
||||
template <
|
||||
typename Shape, typename Element, int AdvanceRank,
|
||||
typename ThreadMap>
|
||||
struct MakePredicatedTileAccessIteratorDesc <
|
||||
Shape, Element, layout::PitchLinear, AdvanceRank, ThreadMap> {
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileAccessIteratorDesc operator()() {
|
||||
|
||||
return PredicatedTileAccessIteratorDesc(
|
||||
sizeof_bits<Element>::value,
|
||||
AdvanceRank,
|
||||
{Shape::kContiguous, Shape::kStrided},
|
||||
{ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided},
|
||||
{ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}
|
||||
);
|
||||
}
|
||||
|
||||
};
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for column-major data.
|
||||
template <
|
||||
typename Shape, typename Element, int AdvanceRank,
|
||||
typename ThreadMap>
|
||||
struct MakePredicatedTileAccessIteratorDesc <
|
||||
Shape, Element, layout::ColumnMajor, AdvanceRank, ThreadMap> {
|
||||
|
||||
static int const kAdvanceRank = AdvanceRank;
|
||||
|
||||
using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc<
|
||||
layout::PitchLinearShape<Shape::kRow, Shape::kColumn>, Element,
|
||||
layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap>;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileAccessIteratorDesc operator()() {
|
||||
|
||||
return UnderlyingMakeOperator()();
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for row-major data.
|
||||
template <
|
||||
typename Shape, typename Element, int AdvanceRank,
|
||||
typename ThreadMap>
|
||||
struct MakePredicatedTileAccessIteratorDesc <
|
||||
Shape, Element, layout::RowMajor, AdvanceRank, ThreadMap> {
|
||||
|
||||
static int const kAdvanceRank = AdvanceRank;
|
||||
|
||||
using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc<
|
||||
layout::PitchLinearShape<Shape::kColumn, Shape::kRow>, Element,
|
||||
layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap>;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileAccessIteratorDesc operator()() {
|
||||
|
||||
return UnderlyingMakeOperator()();
|
||||
}
|
||||
|
||||
};
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for column-major interleaved data.
|
||||
template <
|
||||
typename Shape, typename Element, int AdvanceRank,
|
||||
typename ThreadMap, int InterleavedK>
|
||||
struct MakePredicatedTileAccessIteratorDesc <
|
||||
Shape, Element, layout::ColumnMajorInterleaved<InterleavedK>, AdvanceRank, ThreadMap> {
|
||||
|
||||
static int const kAdvanceRank = AdvanceRank;
|
||||
static int const kInterleavedK = InterleavedK;
|
||||
|
||||
using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc<
|
||||
layout::PitchLinearShape<Shape::kRow * kInterleavedK, Shape::kColumn / kInterleavedK>, Element,
|
||||
layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap>;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileAccessIteratorDesc operator()() {
|
||||
|
||||
return UnderlyingMakeOperator()();
|
||||
}
|
||||
|
||||
};
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for roww-major interleaved data.
|
||||
template <
|
||||
typename Shape, typename Element, int AdvanceRank,
|
||||
typename ThreadMap, int InterleavedK>
|
||||
struct MakePredicatedTileAccessIteratorDesc <
|
||||
Shape, Element, layout::RowMajorInterleaved<InterleavedK>, AdvanceRank, ThreadMap> {
|
||||
|
||||
static int const kAdvanceRank = AdvanceRank;
|
||||
static int const kInterleavedK = InterleavedK;
|
||||
|
||||
using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc<
|
||||
layout::PitchLinearShape<Shape::kColumn * kInterleavedK, Shape::kRow / kInterleavedK>, Element,
|
||||
layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap>;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileAccessIteratorDesc operator()() {
|
||||
|
||||
return UnderlyingMakeOperator()();
|
||||
}
|
||||
|
||||
};
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
//
|
||||
// Parameters struct
|
||||
//
|
||||
|
||||
struct PredicatedTileAccessIteratorParams {
|
||||
|
||||
using Index = int32_t;
|
||||
using LongIndex = int64_t;
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
/// stride of pitch-linear layout (units of Element)
|
||||
LongIndex stride_;
|
||||
/// amount (in byte) to increment pointer to move to next access along
|
||||
/// strided dimension
|
||||
LongIndex inc_strided_;
|
||||
/// amount (in byte) to increment pointer from last access to first access
|
||||
/// of next tile
|
||||
LongIndex inc_next_;
|
||||
/// amount (in byte) to increment pointer from first access of current tile
|
||||
/// to first access of next tile
|
||||
LongIndex inc_advance_;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Status initialize(LongIndex stride, PredicatedTileAccessIteratorDesc desc) {
|
||||
|
||||
stride_ = stride;
|
||||
|
||||
inc_strided_ = (LongIndex(stride_) * desc.threadmap_delta.strided()) *
|
||||
desc.element_size_bits / 8;
|
||||
|
||||
if (desc.advance_rank) {
|
||||
// advance along strided dimension
|
||||
inc_advance_ =
|
||||
desc.threadblock_shape.strided() * LongIndex(stride_) * desc.element_size_bits / 8;
|
||||
} else {
|
||||
// advance along contiguous dimension
|
||||
inc_advance_ = desc.threadblock_shape.contiguous() * desc.element_size_bits / 8;
|
||||
}
|
||||
|
||||
inc_next_ = inc_advance_ - LongIndex(desc.threadmap_iterations.strided() - 1) *
|
||||
desc.threadmap_delta.strided() * LongIndex(stride_) *
|
||||
desc.element_size_bits / 8;
|
||||
|
||||
return Status::kSuccess;
|
||||
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Status initialize(Index stride, PredicatedTileAccessIteratorDesc desc) {
|
||||
return initialize(LongIndex(stride), desc);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileAccessIteratorParams() {
|
||||
initialize(LongIndex(0), PredicatedTileAccessIteratorDesc());
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileAccessIteratorParams(Index stride, PredicatedTileAccessIteratorDesc desc) {
|
||||
initialize(stride, desc);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileAccessIteratorParams(LongIndex stride, PredicatedTileAccessIteratorDesc desc) {
|
||||
initialize(stride, desc);
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace transform
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -189,6 +189,8 @@ class PredicatedTileIterator<Shape_, Element_, layout::PitchLinear, AdvanceRank,
|
||||
/// Parameters object is precomputed state and is host-constructible
|
||||
class Params {
|
||||
public:
|
||||
using Base = typename TileAccessIterator::Params::Base;
|
||||
|
||||
friend PredicatedTileIterator;
|
||||
|
||||
private:
|
||||
@@ -202,6 +204,10 @@ class PredicatedTileIterator<Shape_, Element_, layout::PitchLinear, AdvanceRank,
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params() { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Base const &base)
|
||||
: params_(base) {}
|
||||
};
|
||||
|
||||
private:
|
||||
@@ -449,6 +455,10 @@ public:
|
||||
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {
|
||||
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(typename UnderlyingIterator::Params::Base const &base)
|
||||
: params_(base) {}
|
||||
};
|
||||
|
||||
|
||||
@@ -651,9 +661,12 @@ public:
|
||||
|
||||
/// 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_(base) {}
|
||||
|
||||
};
|
||||
};
|
||||
|
||||
|
||||
@@ -786,6 +799,657 @@ public:
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileIterator for affine rank-2 data.
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
/// WriteableContiguousTileIteratorConcept |
|
||||
/// MaskedTileIteratorConcept
|
||||
///
|
||||
template <typename Shape_, typename Element_, int AdvanceRank,
|
||||
typename ThreadMap_, int AccessSize>
|
||||
class PredicatedTileIterator<Shape_, Element_, layout::AffineRankN<2>, AdvanceRank,
|
||||
ThreadMap_, AccessSize> {
|
||||
public:
|
||||
static_assert(
|
||||
AdvanceRank == 0 || AdvanceRank == 1,
|
||||
"Specialization for pitch-linear iterator may advance along the "
|
||||
"contiguous(rank=0) or strided(rank=1) dimension.");
|
||||
|
||||
using Shape = Shape_;
|
||||
using Element = Element_;
|
||||
using Layout = layout::AffineRankN<2>;
|
||||
static int const kAdvanceRank = AdvanceRank;
|
||||
using ThreadMap = ThreadMap_;
|
||||
|
||||
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 Pointer = Element *;
|
||||
using NonConstPointer = typename platform::remove_const<Element>::type *;
|
||||
|
||||
/// Type used for internal memory accesses
|
||||
using AccessType = AlignedArray<Element, AccessSize, (AccessSize * sizeof_bits<Element>::value / 8)>;
|
||||
|
||||
/// Underlying iterator to compute the addresses
|
||||
using TileAccessIterator =
|
||||
PredicatedTileAccessIterator<Shape, Element, Layout, kAdvanceRank,
|
||||
ThreadMap, AccessType>;
|
||||
|
||||
static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector;
|
||||
|
||||
/// Fragment object to be loaded or stored
|
||||
using Fragment = cutlass::Array<Element, ThreadMap::Iterations::kCount *
|
||||
ThreadMap::kElementsPerAccess>;
|
||||
|
||||
/// Predicate vector stores mask to guard accesses
|
||||
using Mask = typename TileAccessIterator::Mask;
|
||||
|
||||
/// Parameters object is precomputed state and is host-constructible
|
||||
class Params {
|
||||
public:
|
||||
|
||||
friend PredicatedTileIterator;
|
||||
|
||||
private:
|
||||
/// Parameters object
|
||||
typename TileAccessIterator::Params params_;
|
||||
|
||||
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() { }
|
||||
};
|
||||
|
||||
private:
|
||||
/// Internal pointer type permits fast address arithmetic
|
||||
using BytePointer = char *;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Data member to the tile access iterator
|
||||
TileAccessIterator address_iterator_;
|
||||
|
||||
public:
|
||||
/// Constructs a TileIterator from its precomputed state, threadblock offset,
|
||||
/// and thread ID
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileIterator(
|
||||
/// Precomputed parameters object
|
||||
Params const ¶ms,
|
||||
/// Pointer to start of tensor
|
||||
Pointer pointer,
|
||||
/// Extent of tensor
|
||||
TensorCoord extent,
|
||||
/// ID of each participating thread
|
||||
int thread_id,
|
||||
/// Initial offset of threadblock
|
||||
TensorCoord const &threadblock_offset)
|
||||
: address_iterator_(params.params_, pointer, extent, thread_id,
|
||||
threadblock_offset) {}
|
||||
|
||||
/// Construct a PredicatedTileIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileIterator(
|
||||
Params const ¶ms, ///< Precomputed parameters object
|
||||
Pointer pointer, ///< Pointer to start of tensor
|
||||
TensorCoord extent, ///< Extent of tensor
|
||||
int thread_id ///< ID of each participating thread
|
||||
)
|
||||
: PredicatedTileIterator(params, pointer, extent, thread_id,
|
||||
make_Coord(0, 0)) {}
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
address_iterator_.add_pointer_offset(pointer_offset);
|
||||
}
|
||||
|
||||
/// 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
|
||||
PredicatedTileIterator &operator++() {
|
||||
if (kAdvanceRank)
|
||||
address_iterator_.add_tile_offset(make_Coord(0, 1));
|
||||
else
|
||||
address_iterator_.add_tile_offset(make_Coord(1, 0));
|
||||
|
||||
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
|
||||
PredicatedTileIterator operator++(int) {
|
||||
PredicatedTileIterator self(*this);
|
||||
operator++();
|
||||
return self;
|
||||
}
|
||||
|
||||
/// Clears the predicate set efficiently
|
||||
CUTLASS_HOST_DEVICE
|
||||
void clear_mask() { address_iterator_.clear_mask(); }
|
||||
|
||||
/// Clears the predicate set efficiently
|
||||
CUTLASS_HOST_DEVICE
|
||||
void enable_mask() { address_iterator_.enable_mask(); }
|
||||
|
||||
/// Sets the predicate mask, overriding value stored in predicate iterator
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_mask(Mask const &mask) { address_iterator_.set_mask(mask); }
|
||||
|
||||
/// Gets the mask
|
||||
CUTLASS_HOST_DEVICE
|
||||
void get_mask(Mask &mask) { address_iterator_.get_mask(mask); }
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
|
||||
load_with_byte_offset(frag, pointer_offset * sizeof_bits<Element>::value / 8);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) {
|
||||
|
||||
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int v = 0; v < kAccessesPerVector; ++v) {
|
||||
|
||||
int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous);
|
||||
|
||||
address_iterator_.set_iteration_index(idx);
|
||||
char const *byte_ptr = reinterpret_cast<char const *>(address_iterator_.get()) + byte_offset;
|
||||
|
||||
AccessType const *access_ptr = reinterpret_cast<AccessType const *>(byte_ptr);
|
||||
|
||||
cutlass::arch::global_load<AccessType,
|
||||
sizeof(AccessType)
|
||||
>(
|
||||
frag_ptr[idx], access_ptr, address_iterator_.valid());
|
||||
|
||||
++address_iterator_;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) { load_with_byte_offset(frag, 0); }
|
||||
|
||||
/// Store a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
|
||||
store_with_byte_offset(frag, pointer_offset * sizeof_bits<Element>::value / 8);
|
||||
}
|
||||
|
||||
/// Store a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) {
|
||||
address_iterator_.set_iteration_index(0);
|
||||
AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int v = 0; v < kAccessesPerVector; ++v) {
|
||||
|
||||
int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous);
|
||||
|
||||
char *byte_ptr = reinterpret_cast<char *>(address_iterator_.get()) + byte_offset;
|
||||
AccessType *access_ptr = reinterpret_cast<AccessType *>(byte_ptr);
|
||||
|
||||
if (address_iterator_.valid()) {
|
||||
*access_ptr = frag_ptr[idx];
|
||||
}
|
||||
++address_iterator_;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Store a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store(Fragment const &frag) { store_with_byte_offset(frag, 0); }
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileIterator for affine rank 2 column-major data.
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
/// WriteableContiguousTileIteratorConcept |
|
||||
/// MaskedTileIteratorConcept
|
||||
///
|
||||
template <
|
||||
typename Shape_,
|
||||
typename Element_,
|
||||
int AdvanceRank,
|
||||
typename ThreadMap_,
|
||||
int AccessSize
|
||||
>
|
||||
class PredicatedTileIterator<Shape_, Element_, layout::AffineRank2ColumnMajor, AdvanceRank, ThreadMap_, AccessSize> {
|
||||
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::AffineRank2ColumnMajor;
|
||||
static int const kAdvanceRank = AdvanceRank;
|
||||
using ThreadMap = ThreadMap_;
|
||||
|
||||
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 Pointer = Element *;
|
||||
using NonConstPointer = typename platform::remove_const<Element>::type *;
|
||||
|
||||
// Map to the underlying AffineRankN<2> layout
|
||||
using UnderlyingIterator = PredicatedTileIterator<
|
||||
layout::PitchLinearShape<Shape::kRow, Shape::kColumn>,
|
||||
Element,
|
||||
layout::AffineRankN<2>,
|
||||
(kAdvanceRank == 0 ? 0 : 1),
|
||||
ThreadMap,
|
||||
AccessSize
|
||||
>;
|
||||
|
||||
using AccessType = typename UnderlyingIterator::AccessType;
|
||||
|
||||
/// Fragment object to be loaded or stored
|
||||
using Fragment = cutlass::Array<Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>;
|
||||
|
||||
/// Predicate vector stores mask to guard accesses
|
||||
using Mask = typename UnderlyingIterator::Mask;
|
||||
|
||||
/// Parameters object is precomputed state and is host-constructible
|
||||
class Params {
|
||||
private:
|
||||
|
||||
friend PredicatedTileIterator;
|
||||
|
||||
/// Parameters object
|
||||
typename UnderlyingIterator::Params params_;
|
||||
|
||||
public:
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params() { }
|
||||
|
||||
/// 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))) {
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Underlying AffineRankN<2> tile iterator
|
||||
UnderlyingIterator iterator_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileIterator(
|
||||
Params const ¶ms, ///< Precomputed parameters object
|
||||
Pointer pointer, ///< Pointer to start of tensor
|
||||
TensorCoord extent, ///< Extent of tensor
|
||||
int thread_id, ///< ID of each participating thread
|
||||
TensorCoord const &threadblock_offset ///< Initial offset of threadblock
|
||||
):
|
||||
iterator_(
|
||||
params.params_,
|
||||
pointer,
|
||||
layout::PitchLinearCoord(extent.row(), extent.column()),
|
||||
thread_id,
|
||||
layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column())
|
||||
) { }
|
||||
|
||||
/// Construct a PredicatedTileIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileIterator(
|
||||
Params const ¶ms, ///< Precomputed parameters object
|
||||
Pointer pointer, ///< Pointer to start of tensor
|
||||
TensorCoord extent, ///< Extent of tensor
|
||||
int thread_id ///< ID of each participating thread
|
||||
): PredicatedTileIterator(params, pointer, extent, thread_id, make_Coord(0, 0)) { }
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
iterator_.add_pointer_offset(pointer_offset);
|
||||
}
|
||||
|
||||
/// 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
|
||||
PredicatedTileIterator &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
|
||||
PredicatedTileIterator operator++(int) {
|
||||
PredicatedTileIterator self(*this);
|
||||
operator++();
|
||||
return self;
|
||||
}
|
||||
|
||||
/// Clears the predicate set efficiently
|
||||
CUTLASS_HOST_DEVICE
|
||||
void clear_mask() {
|
||||
iterator_.clear_mask();
|
||||
}
|
||||
|
||||
/// Clears the predicate set efficiently
|
||||
CUTLASS_HOST_DEVICE
|
||||
void enable_mask() {
|
||||
iterator_.enable_mask();
|
||||
}
|
||||
|
||||
/// Sets the predicate mask, overriding value stored in predicate iterator
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_mask(Mask const &mask) {
|
||||
iterator_.set_mask(mask);
|
||||
}
|
||||
|
||||
/// Gets the mask
|
||||
CUTLASS_HOST_DEVICE
|
||||
void get_mask(Mask &mask) {
|
||||
iterator_.get_mask(mask);
|
||||
}
|
||||
|
||||
/// 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_with_byte_offset(Fragment &frag, LongIndex byte_offset) {
|
||||
iterator_.load_with_byte_offset(frag, byte_offset);
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
load_with_pointer_offset(frag, 0);
|
||||
}
|
||||
|
||||
/// Store a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
|
||||
iterator_.store_with_pointer_offset(frag, pointer_offset);
|
||||
}
|
||||
|
||||
/// Store a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) {
|
||||
iterator_.store_with_byte_offset(frag, byte_offset);
|
||||
}
|
||||
|
||||
/// Store a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store(Fragment const &frag) {
|
||||
store_with_pointer_offset(frag, 0);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileIterator for affine rank 2 row-major data.
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
/// WriteableContiguousTileIteratorConcept |
|
||||
/// MaskedTileIteratorConcept
|
||||
///
|
||||
template <
|
||||
typename Shape_,
|
||||
typename Element_,
|
||||
int AdvanceRank,
|
||||
typename ThreadMap_,
|
||||
int AccessSize
|
||||
>
|
||||
class PredicatedTileIterator<Shape_, Element_, layout::AffineRank2RowMajor, AdvanceRank, ThreadMap_, AccessSize> {
|
||||
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::AffineRank2RowMajor;
|
||||
static int const kAdvanceRank = AdvanceRank;
|
||||
using ThreadMap = ThreadMap_;
|
||||
|
||||
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 Pointer = Element *;
|
||||
using NonConstPointer = typename platform::remove_const<Element>::type *;
|
||||
|
||||
// Map to the underlying AffineRankN<2> layout
|
||||
using UnderlyingIterator = PredicatedTileIterator<
|
||||
layout::PitchLinearShape<Shape::kColumn, Shape::kRow>,
|
||||
Element,
|
||||
layout::AffineRankN<2>,
|
||||
(kAdvanceRank == 0 ? 1 : 0),
|
||||
ThreadMap,
|
||||
AccessSize
|
||||
>;
|
||||
|
||||
using AccessType = typename UnderlyingIterator::AccessType;
|
||||
|
||||
/// Fragment object to be loaded or stored
|
||||
using Fragment = cutlass::Array<Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>;
|
||||
|
||||
/// Predicate vector stores mask to guard accesses
|
||||
using Mask = typename UnderlyingIterator::Mask;
|
||||
|
||||
/// Parameters object is precomputed state and is host-constructible
|
||||
class Params {
|
||||
private:
|
||||
|
||||
friend PredicatedTileIterator;
|
||||
|
||||
/// Parameters object
|
||||
typename UnderlyingIterator::Params params_;
|
||||
|
||||
public:
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params() { }
|
||||
|
||||
/// Construct the Params object given an AffineRankN<2> tensor's layout
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Layout const &layout): params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))) {}
|
||||
};
|
||||
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Underlying AffineRankN<2> tile iterator
|
||||
UnderlyingIterator iterator_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileIterator(
|
||||
Params const ¶ms, ///< Precomputed parameters object
|
||||
Pointer pointer, ///< Pointer to start of tensor
|
||||
TensorCoord extent, ///< Extent of tensor
|
||||
int thread_id, ///< ID of each participating thread
|
||||
TensorCoord const &threadblock_offset ///< Initial offset of threadblock
|
||||
):
|
||||
iterator_(
|
||||
params.params_,
|
||||
pointer,
|
||||
layout::PitchLinearCoord(extent.column(), extent.row()),
|
||||
thread_id,
|
||||
layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row())
|
||||
) { }
|
||||
|
||||
/// Construct a PredicatedTileIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedTileIterator(
|
||||
Params const ¶ms, ///< Precomputed parameters object
|
||||
Pointer pointer, ///< Pointer to start of tensor
|
||||
TensorCoord extent, ///< Extent of tensor
|
||||
int thread_id ///< ID of each participating thread
|
||||
): PredicatedTileIterator(params, pointer, extent, thread_id, make_Coord(0, 0)) { }
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
iterator_.add_pointer_offset(pointer_offset);
|
||||
}
|
||||
|
||||
/// 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
|
||||
PredicatedTileIterator &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
|
||||
PredicatedTileIterator operator++(int) {
|
||||
PredicatedTileIterator self(*this);
|
||||
operator++();
|
||||
return self;
|
||||
}
|
||||
|
||||
/// Clears the predicate set efficiently
|
||||
CUTLASS_HOST_DEVICE
|
||||
void clear_mask() {
|
||||
iterator_.clear_mask();
|
||||
}
|
||||
|
||||
/// Clears the predicate set efficiently
|
||||
CUTLASS_HOST_DEVICE
|
||||
void enable_mask() {
|
||||
iterator_.enable_mask();
|
||||
}
|
||||
|
||||
/// Sets the predicate mask, overriding value stored in predicate iterator
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_mask(Mask const &mask) {
|
||||
iterator_.set_mask(mask);
|
||||
}
|
||||
|
||||
/// Gets the mask
|
||||
CUTLASS_HOST_DEVICE
|
||||
void get_mask(Mask &mask) {
|
||||
iterator_.get_mask(mask);
|
||||
}
|
||||
|
||||
/// 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_with_byte_offset(Fragment &frag, LongIndex byte_offset) {
|
||||
iterator_.load_with_byte_offset(frag, byte_offset);
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
load_with_pointer_offset(frag, 0);
|
||||
}
|
||||
|
||||
/// Store a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
|
||||
iterator_.store_with_pointer_offset(frag, pointer_offset);
|
||||
}
|
||||
|
||||
/// Store a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) {
|
||||
iterator_.store_with_byte_offset(frag, byte_offset);
|
||||
}
|
||||
|
||||
/// Store a fragment to memory
|
||||
CUTLASS_DEVICE
|
||||
void store(Fragment const &frag) {
|
||||
store_with_pointer_offset(frag, 0);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
/// Specialization of PredicatedTileIterator for interleaved data. It is mapped
|
||||
/// to the congruous layout.
|
||||
///
|
||||
@@ -854,6 +1518,11 @@ class PredicatedTileIterator<Shape_, Element_,
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Layout const &layout)
|
||||
: params_(layout::PitchLinear(layout.stride(0))) {}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(typename UnderlyingIterator::Params::Base const &base)
|
||||
: params_(base) {}
|
||||
|
||||
};
|
||||
|
||||
private:
|
||||
@@ -1035,6 +1704,10 @@ class PredicatedTileIterator<Shape_, Element_,
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Layout const &layout)
|
||||
: params_(layout::PitchLinear(layout.stride(0))) {}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(typename UnderlyingIterator::Params::Base const &base)
|
||||
: params_(base) {}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
@@ -194,6 +194,8 @@ class PredicatedTileIterator2dThreadTile<Shape_, Element_, layout::PitchLinear,
|
||||
/// Parameters object is precomputed state and is host-constructible
|
||||
class Params {
|
||||
public:
|
||||
using Base = typename TileAccessIterator::Params::Base;
|
||||
|
||||
friend PredicatedTileIterator2dThreadTile;
|
||||
|
||||
private:
|
||||
@@ -207,6 +209,10 @@ class PredicatedTileIterator2dThreadTile<Shape_, Element_, layout::PitchLinear,
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params() { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Base const &base)
|
||||
: params_(base) {}
|
||||
};
|
||||
|
||||
private:
|
||||
@@ -443,9 +449,11 @@ public:
|
||||
|
||||
/// 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_(base) {}
|
||||
};
|
||||
|
||||
|
||||
@@ -637,9 +645,11 @@ public:
|
||||
|
||||
/// 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_(base) {}
|
||||
};
|
||||
|
||||
|
||||
|
||||
@@ -74,6 +74,7 @@ class RegularTileAccessIterator<
|
||||
|
||||
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;
|
||||
@@ -89,7 +90,7 @@ class RegularTileAccessIterator<
|
||||
//
|
||||
|
||||
/// Stride value
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
AccessType *pointer_;
|
||||
|
||||
@@ -76,6 +76,7 @@ class RegularTileAccessIterator<
|
||||
|
||||
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;
|
||||
@@ -107,7 +108,7 @@ class RegularTileAccessIterator<
|
||||
//
|
||||
|
||||
/// Stride value
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
AccessType *pointer_[Detail::kPointerCount];
|
||||
@@ -445,6 +446,7 @@ class RegularTileAccessIterator<Shape_, Element_,
|
||||
|
||||
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;
|
||||
@@ -492,7 +494,7 @@ class RegularTileAccessIterator<Shape_, Element_,
|
||||
int sections_per_stage_;
|
||||
|
||||
/// Stride value
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
AccessType *pointer_[Detail::kPointerCount];
|
||||
|
||||
@@ -74,6 +74,7 @@ class RegularTileAccessIterator<
|
||||
|
||||
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;
|
||||
@@ -107,7 +108,7 @@ class RegularTileAccessIterator<
|
||||
//
|
||||
|
||||
/// Stride value
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
AccessType *pointer_;
|
||||
@@ -437,6 +438,7 @@ class RegularTileAccessIterator<
|
||||
|
||||
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;
|
||||
@@ -471,7 +473,7 @@ class RegularTileAccessIterator<
|
||||
//
|
||||
|
||||
/// Stride value
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
AccessType *pointer_;
|
||||
@@ -811,6 +813,7 @@ class RegularTileAccessIterator<
|
||||
|
||||
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;
|
||||
@@ -844,7 +847,7 @@ class RegularTileAccessIterator<
|
||||
//
|
||||
|
||||
/// Stride value
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
AccessType *pointer_;
|
||||
@@ -1175,6 +1178,7 @@ class RegularTileAccessIterator<
|
||||
|
||||
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;
|
||||
@@ -1211,7 +1215,7 @@ class RegularTileAccessIterator<
|
||||
//
|
||||
|
||||
/// Stride value
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
AccessType *pointer_;
|
||||
|
||||
@@ -70,11 +70,14 @@ public:
|
||||
|
||||
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 Fragment = Array<Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>;
|
||||
|
||||
using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess, kAlignment>;
|
||||
|
||||
static_assert(kAdvanceRank == 0 || kAdvanceRank == 1,
|
||||
"Advance rank may only be along the contiguous or strided dimensions.");
|
||||
@@ -84,8 +87,6 @@ private:
|
||||
//
|
||||
// Types
|
||||
//
|
||||
|
||||
using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess, kAlignment>;
|
||||
|
||||
//
|
||||
// Data members
|
||||
@@ -95,7 +96,7 @@ private:
|
||||
uint8_t *pointer_;
|
||||
|
||||
/// Stride quantity
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Amount to increment pointer along strided dimension
|
||||
Index increment_strided_;
|
||||
@@ -242,6 +243,30 @@ public:
|
||||
(coord.contiguous() * Shape::kContiguous + coord.strided() * Shape::kStrided * stride_) / 8;
|
||||
add_pointer_offset(offset);
|
||||
}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(int index) {
|
||||
}
|
||||
|
||||
/// Returns a pointer
|
||||
CUTLASS_HOST_DEVICE
|
||||
AccessType *get() const {
|
||||
#if 0
|
||||
AccessType *access_ptr = pointer_[iteration_strided_ & 1];
|
||||
int stride_idx = (iteration_strided_ & ~1);
|
||||
|
||||
int access_offset = stride_idx * 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_);
|
||||
#endif
|
||||
return reinterpret_cast<AccessType *>(pointer_);
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -281,6 +306,8 @@ public:
|
||||
kAlignment
|
||||
>;
|
||||
|
||||
using AccessType = typename Underlying::AccessType;
|
||||
|
||||
static_assert(kAdvanceRank == 0 || kAdvanceRank == 1,
|
||||
"Advance rank may only be along the row or column dimensions.");
|
||||
|
||||
@@ -364,6 +391,17 @@ public:
|
||||
iterator_.add_tile_offset({coord.column(), coord.row()});
|
||||
}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(int index) {
|
||||
}
|
||||
|
||||
/// Returns a pointer
|
||||
CUTLASS_HOST_DEVICE
|
||||
AccessType *get() const {
|
||||
return iterator_.get();
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -402,6 +440,8 @@ public:
|
||||
ThreadMap
|
||||
>;
|
||||
|
||||
using AccessType = typename Underlying::AccessType;
|
||||
|
||||
static_assert(kAdvanceRank == 0 || kAdvanceRank == 1,
|
||||
"Advance rank may only be along the row or column dimensions.");
|
||||
|
||||
@@ -485,6 +525,17 @@ public:
|
||||
iterator_.add_tile_offset({coord.row(), coord.column()});
|
||||
}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(int index) {
|
||||
}
|
||||
|
||||
/// Returns a pointer
|
||||
CUTLASS_HOST_DEVICE
|
||||
AccessType *get() const {
|
||||
return iterator_.get();
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -79,6 +79,7 @@ public:
|
||||
|
||||
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;
|
||||
@@ -104,13 +105,13 @@ private:
|
||||
uint8_t *pointer_;
|
||||
|
||||
/// Stride quantity
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Amount to increment pointer along strided dimension
|
||||
Index increment_strided_;
|
||||
LongIndex increment_strided_;
|
||||
|
||||
/// Amount to advance pointer between tiles
|
||||
Index increment_advance_;
|
||||
LongIndex increment_advance_;
|
||||
|
||||
public:
|
||||
|
||||
|
||||
@@ -85,6 +85,7 @@ public:
|
||||
|
||||
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;
|
||||
@@ -123,7 +124,7 @@ private:
|
||||
//
|
||||
|
||||
/// Stride value
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
AccessType * pointer_[Detail::kPointerCount];
|
||||
@@ -557,6 +558,7 @@ public:
|
||||
|
||||
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;
|
||||
@@ -595,7 +597,7 @@ private:
|
||||
//
|
||||
|
||||
/// Stride value
|
||||
Index stride_;
|
||||
StrideIndex stride_;
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
AccessType * pointer_[Detail::kPointerCount];
|
||||
|
||||
Reference in New Issue
Block a user