wmma_gemm_global_tile.h
Go to the documentation of this file.
/***************************************************************************************************
* Copyright (c) 2017-2018, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
* * Redistributions of source code must retain the above copyright notice, this list of
* conditions and the following disclaimer.
* * Redistributions in binary form must reproduce the above copyright notice, this list of
* conditions and the following disclaimer in the documentation and/or other materials
* provided with the distribution.
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
* to endorse or promote products derived from this software without specific prior written
* permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include <cutlass/gemm/gemm_global_tile.h>
namespace cutlass {
namespace gemm {
template <typename Scalar_, typename Tile_, typename Threads_, int kAccessSize_>
struct WmmaGemmGlobalIteratorCdTraits : public GemmGlobalTileTraits<GemmOperand::kC,
MatrixLayout::kColumnMajor,
Scalar_,
Tile_,
Threads_,
kAccessSize_> {
typedef GemmGlobalTileTraits<GemmOperand::kC,
MatrixLayout::kColumnMajor,
Scalar_,
Tile_,
Threads_,
kAccessSize_>
Base;
typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> Delta;
struct ThreadOffset {
CUTLASS_HOST_DEVICE
Coord<4> operator()() const {
int thread_offset_h = threadIdx.x / Base::Threads::kW;
int thread_offset_w = threadIdx.x % Base::Threads::kW * Base::ThreadsDelta::kW;
return make_Coord(0, thread_offset_h, thread_offset_w, 0);
}
};
};
template <typename TileTraits_, typename Index_ = int>
struct WmmaGemmGlobalIteratorCd : public TileIteratorBase<TileTraits_,
typename TileTraits_::Scalar,
IteratorAdvance::kH,
MemorySpace::kGlobal,
Index_> {
typedef WmmaGemmGlobalIteratorCd<TileTraits_, Index_> This_;
typedef TileTraits_ Traits;
typedef TileIteratorBase<Traits,
typename TileTraits_::Scalar,
IteratorAdvance::kH,
MemorySpace::kGlobal,
Index_>
Base;
typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> ImmediateOffsetStrides;
static MatrixLayout::Kind const kLayout = TileTraits_::kLayout;
typedef typename TileTraits_::Scalar Scalar;
typedef typename TileTraits_::Pointer Pointer;
typedef typename TileTraits_::Threads Threads;
typedef Index_ Index;
typedef typename TileTraits_::ThreadOffset ThreadOffset;
struct Params {
Pointer pointer;
Index stride_h;
Index inc_h, inc_advance;
Index predicate_offset;
Index predicate_inc_h, predicate_inc_advance;
CUTLASS_HOST_DEVICE int initialize(
Pointer pointer, Index ld, Index n, Index epilogue_stride_w, Index epilogue_delta_w) {
// The pointer.
this->pointer = pointer;
// Setup the base stride. One "group of threads" per column.
stride_h = ld;
// Each thread output 1 column per iteration. .
inc_h = ld * TileTraits_::Threads::kH;
inc_advance = inc_h + epilogue_stride_w;
predicate_offset = n;
predicate_inc_h = TileTraits_::Threads::kH;
predicate_inc_advance = predicate_inc_h + epilogue_delta_w;
// It worked.
return 0;
}
};
Params params;
Coord<4> thread_offset;
CUTLASS_DEVICE WmmaGemmGlobalIteratorCd() {}
CUTLASS_DEVICE WmmaGemmGlobalIteratorCd(Params const& params,
const Coord<3>& bounds,
const Coord<3>& block,
int const pointer_offset = 0,
int const pred_offset = 0,
ThreadOffset thread_offset_func = ThreadOffset())
: params(params) {
thread_offset = thread_offset_func();
// Each warp works on a different column of the tile.
int const h = thread_offset[1] + block[1];
// Each lane writes a different element.
int const w = thread_offset[2] + block[2];
// Setup the pointer.
this->params.pointer += ((h * params.stride_h + w) + pointer_offset);
// Prepare the vector of predicates.
for (int i = 0; i < Base::Iterations::kW; ++i) {
predicates.set(i, w + i * Base::Delta::kW < bounds[2]);
}
this->params.predicate_offset -= (h + pred_offset);
}
CUTLASS_DEVICE void inc_c() {}
CUTLASS_DEVICE void inc_w() {}
CUTLASS_DEVICE void inc_h() {
params.pointer += params.inc_h;
params.predicate_offset -= params.predicate_inc_h;
}
CUTLASS_DEVICE void inc_d() {}
CUTLASS_DEVICE void inc_advance() {
params.pointer += params.inc_advance;
params.predicate_offset -= params.predicate_inc_advance;
}
CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const {
return predicates.at(w) && params.predicate_offset > 0;
}
CUTLASS_HOST_DEVICE
Pointer data() { return params.pointer; }
CUTLASS_HOST_DEVICE
Pointer const data() const { return params.pointer; }
cutlass::PredicateVector<Base::Iterations::kW> predicates;
};
} // namespace gemm
} // namespace cutlass
cutlass::gemm::WmmaGemmGlobalIteratorCd::Threads
TileTraits_::Threads Threads
The threads.
Definition: wmma_gemm_global_tile.h:97
cutlass
Definition: convert.h:33
gemm_global_tile.h
Defines iterators for efficiently loading and storing to global memory.
cutlass::gemm::GemmGlobalTileTraits
Definition: gemm_global_tile.h:70
cutlass::PredicateVector::at
CUTLASS_HOST_DEVICE bool at(int idx) const
Accesses a bit within the predicate vector.
Definition: predicate_vector.h:356
cutlass::gemm::WmmaGemmGlobalIteratorCd::inc_d
CUTLASS_DEVICE void inc_d()
Increment the pointer in the D dimension.
Definition: wmma_gemm_global_tile.h:177
cutlass::MemorySpace::kGlobal
Definition: load_store.h:43
cutlass::make_Coord
CUTLASS_HOST_DEVICE Coord< 1 > make_Coord(int _0)
Helper to make a 2-element coordinate.
Definition: coord.h:241
cutlass::gemm::WmmaGemmGlobalIteratorCd::Params::stride_h
Index stride_h
The stride in the H dimension to setup the thread in the block.
Definition: wmma_gemm_global_tile.h:108
cutlass::gemm::WmmaGemmGlobalIteratorCd::inc_w
CUTLASS_DEVICE void inc_w()
Increment the pointer in the W dimension.
Definition: wmma_gemm_global_tile.h:170
cutlass::gemm::WmmaGemmGlobalIteratorCd::Index
Index_ Index
The index.
Definition: wmma_gemm_global_tile.h:99
cutlass::gemm::WmmaGemmGlobalIteratorCd::Scalar
TileTraits_::Scalar Scalar
The scalar.
Definition: wmma_gemm_global_tile.h:93
cutlass::IteratorAdvance::kH
Definition: tile_iterator.h:62
cutlass::GemmOperand::kC
Definition: matrix_traits.h:43
cutlass::gemm::WmmaGemmGlobalIteratorCd::params
Params params
Definition: wmma_gemm_global_tile.h:136
cutlass::gemm::WmmaGemmGlobalIteratorCd::Params::predicate_inc_h
Index predicate_inc_h
The strides to increment the predicate offset.
Definition: wmma_gemm_global_tile.h:114
cutlass::gemm::WmmaGemmGlobalIteratorCd::Params::pointer
Pointer pointer
The pointer.
Definition: wmma_gemm_global_tile.h:106
cutlass::gemm::WmmaGemmGlobalIteratorCd::data
CUTLASS_HOST_DEVICE Pointer const data() const
Definition: wmma_gemm_global_tile.h:194
cutlass::gemm::WmmaGemmGlobalIteratorCd::inc_advance
CUTLASS_DEVICE void inc_advance()
Increment the pointer to move to the next iteration.
Definition: wmma_gemm_global_tile.h:179
cutlass::gemm::WmmaGemmGlobalIteratorCd::Params
The params.
Definition: wmma_gemm_global_tile.h:104
cutlass::gemm::WmmaGemmGlobalIteratorCd::Params::inc_h
Index inc_h
The strides to increment the pointer.
Definition: wmma_gemm_global_tile.h:110
cutlass::gemm::WmmaGemmGlobalIteratorCd::Base
TileIteratorBase< Traits, typename TileTraits_::Scalar, IteratorAdvance::kH, MemorySpace::kGlobal, Index_ > Base
The base class.
Definition: wmma_gemm_global_tile.h:86
cutlass::gemm::WmmaGemmGlobalIteratorCd::WmmaGemmGlobalIteratorCd
CUTLASS_DEVICE WmmaGemmGlobalIteratorCd()
Ctor.
Definition: wmma_gemm_global_tile.h:141
cutlass::gemm::WmmaGemmGlobalIteratorCd::Params::predicate_offset
Index predicate_offset
The column offset to compute the predicate for the columns.
Definition: wmma_gemm_global_tile.h:112
cutlass::gemm::WmmaGemmGlobalIteratorCdTraits::ThreadOffset::operator()
CUTLASS_HOST_DEVICE Coord< 4 > operator()() const
Definition: wmma_gemm_global_tile.h:59
cutlass::gemm::WmmaGemmGlobalIteratorCd
Definition: wmma_gemm_global_tile.h:71
cutlass::gemm::WmmaGemmGlobalIteratorCd::Params::predicate_inc_advance
Index predicate_inc_advance
Definition: wmma_gemm_global_tile.h:114
cutlass::gemm::WmmaGemmGlobalIteratorCd::Pointer
TileTraits_::Pointer Pointer
The pointer.
Definition: wmma_gemm_global_tile.h:95
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:46
cutlass::gemm::WmmaGemmGlobalIteratorCdTraits::Delta
Shape< 0, 0, Base::Delta::kW, Base::Delta::kC > Delta
Override the strides in each dimension between different loads/stores.
Definition: wmma_gemm_global_tile.h:54
cutlass::PredicateVector< Base::Iterations::kW >
cutlass::Shape
A Shape implementing Layout Concept describing the dimensions of a cube.
Definition: shape.h:64
cutlass::gemm::WmmaGemmGlobalIteratorCd::thread_offset
Coord< 4 > thread_offset
Definition: wmma_gemm_global_tile.h:138
cutlass::gemm::WmmaGemmGlobalIteratorCd::Params::inc_advance
Index inc_advance
Definition: wmma_gemm_global_tile.h:110
cutlass::gemm::WmmaGemmGlobalIteratorCd::kLayout
static MatrixLayout::Kind const kLayout
The layout.
Definition: wmma_gemm_global_tile.h:90
cutlass::Coord< 4 >
cutlass::gemm::WmmaGemmGlobalIteratorCdTraits
Definition: wmma_gemm_global_tile.h:38
cutlass::TileIteratorBase
Iterator for accessing a stripmined tile in memory.
Definition: tile_iterator.h:102
cutlass::gemm::WmmaGemmGlobalIteratorCd::ThreadOffset
TileTraits_::ThreadOffset ThreadOffset
The thread offset functor.
Definition: wmma_gemm_global_tile.h:101
cutlass::MatrixLayout::kColumnMajor
Definition: matrix_traits.h:36
cutlass::gemm::WmmaGemmGlobalIteratorCd::data
CUTLASS_HOST_DEVICE Pointer data()
Returns the raw pointer.
Definition: wmma_gemm_global_tile.h:191
cutlass::Shape::kW
static int const kW
The width of the cube.
Definition: shape.h:70
cutlass::PredicateVector::set
CUTLASS_HOST_DEVICE void set(int idx, bool value=true)
Set a bit within the predicate vector.
Definition: predicate_vector.h:364
cutlass::MatrixLayout::Kind
Kind
Definition: matrix_traits.h:36
cutlass::gemm::WmmaGemmGlobalIteratorCdTraits::Base
GemmGlobalTileTraits< GemmOperand::kC, MatrixLayout::kColumnMajor, Scalar_, Tile_, Threads_, kAccessSize_ > Base
The base class.
Definition: wmma_gemm_global_tile.h:51
cutlass::gemm::WmmaGemmGlobalIteratorCd::WmmaGemmGlobalIteratorCd
CUTLASS_DEVICE WmmaGemmGlobalIteratorCd(Params const &params, const Coord< 3 > &bounds, const Coord< 3 > &block, int const pointer_offset=0, int const pred_offset=0, ThreadOffset thread_offset_func=ThreadOffset())
Ctor.
Definition: wmma_gemm_global_tile.h:144
cutlass::gemm::WmmaGemmGlobalIteratorCd::This_
WmmaGemmGlobalIteratorCd< TileTraits_, Index_ > This_
This class.
Definition: wmma_gemm_global_tile.h:77
cutlass::gemm::WmmaGemmGlobalIteratorCd::predicates
cutlass::PredicateVector< Base::Iterations::kW > predicates
The predicates for the row.
Definition: wmma_gemm_global_tile.h:197
cutlass::gemm::WmmaGemmGlobalIteratorCd::ImmediateOffsetStrides
Shape< 0, 0, Base::Delta::kW, Base::Delta::kC > ImmediateOffsetStrides
Override the strides in each dimension between different loads/stores.
Definition: wmma_gemm_global_tile.h:88
cutlass::gemm::WmmaGemmGlobalIteratorCdTraits::ThreadOffset
Computes the thread offset in (H, W) based on thread ID.
Definition: wmma_gemm_global_tile.h:57
cutlass::gemm::WmmaGemmGlobalIteratorCd::inc_c
CUTLASS_DEVICE void inc_c()
Increment the pointer in the C dimension.
Definition: wmma_gemm_global_tile.h:168
cutlass::gemm::WmmaGemmGlobalIteratorCd::inc_h
CUTLASS_DEVICE void inc_h()
Increment the pointer in the H dimension.
Definition: wmma_gemm_global_tile.h:172
cutlass::gemm::WmmaGemmGlobalIteratorCd::Traits
TileTraits_ Traits
The traits.
Definition: wmma_gemm_global_tile.h:79
cutlass::gemm::WmmaGemmGlobalIteratorCd::valid
CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const
Test the predicate.
Definition: wmma_gemm_global_tile.h:185
cutlass::gemm::WmmaGemmGlobalIteratorCd::Params::initialize
CUTLASS_HOST_DEVICE int initialize(Pointer pointer, Index ld, Index n, Index epilogue_stride_w, Index epilogue_delta_w)
Setup the params.
Definition: wmma_gemm_global_tile.h:117