gemm_operand.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/matrix_traits.h>
#include <cutlass/reshape_tile.h>
#include <cutlass/util/platform.h>
namespace cutlass {
namespace gemm {
template <GemmOperand::Kind kOperand_, MatrixLayout::Kind kLayout_>
struct GemmOperandTraitsAb {
static const bool Congruous =
(kOperand_ == GemmOperand::kA ^ kLayout_ == MatrixLayout::kRowMajor);
};
template <typename GemmOperand::Kind kOperand_, typename Tile_>
struct GetExtent;
template <typename Tile_>
struct GetExtent<GemmOperand::kA, Tile_> {
static const int kExtent = Tile_::kW;
};
template <typename Tile_>
struct GetExtent<GemmOperand::kB, Tile_> {
static const int kExtent = Tile_::kH;
};
template <typename ThreadBlockTile_, GemmOperand::Kind Usage, MatrixLayout::Kind Layout>
struct GemmMultiplicandTraits {
// Only defined for A or B
static_assert(Usage == GemmOperand::kA || Usage == GemmOperand::kB,
"MultiplicandTileShape defined only for A or B operands.");
typedef ThreadBlockTile_ ThreadBlockTile;
static GemmOperand::Kind const kUsage = Usage;
static MatrixLayout::Kind const kLayout = Layout;
// True if K is the strided dimension
static bool const kKstrided = (kUsage == GemmOperand::kA ^ kLayout == MatrixLayout::kRowMajor);
typedef typename platform::conditional<
kKstrided,
Shape<1, ThreadBlockTile::kD, GetExtent<Usage, ThreadBlockTile>::kExtent>,
Shape<1, GetExtent<Usage, ThreadBlockTile>::kExtent, ThreadBlockTile::kD> >::type Shape;
};
template <GemmOperand::Kind operand, bool Kstrided = true>
struct ProjectOperand;
template <bool Kstrided>
struct ProjectOperand<GemmOperand::kA, Kstrided> {
CUTLASS_HOST_DEVICE
static Coord<3> project(Coord<3> const &coord) {
if (Kstrided) {
return make_Coord(0, coord[0], coord[2]);
} else {
return make_Coord(0, coord[2], coord[0]);
}
}
};
template <bool Kstrided>
struct ProjectOperand<GemmOperand::kB, Kstrided> {
CUTLASS_HOST_DEVICE
static Coord<3> project(Coord<3> const &coord) {
if (Kstrided) {
return make_Coord(0, coord[0], coord[1]);
} else {
return make_Coord(0, coord[1], coord[0]);
}
}
};
template <>
struct ProjectOperand<GemmOperand::kC, true> {
CUTLASS_HOST_DEVICE
static Coord<3> project(Coord<3> const &coord) { return make_Coord(0, coord[1], coord[2]); }
};
template <>
struct ProjectOperand<GemmOperand::kD, true> {
CUTLASS_HOST_DEVICE
static Coord<3> project(Coord<3> const &coord) { return make_Coord(0, coord[1], coord[2]); }
};
} // namespace gemm
} // namespace cutlass
cutlass::gemm::GemmMultiplicandTraits::kKstrided
static bool const kKstrided
Definition: gemm_operand.h:82
cutlass::gemm::ProjectOperand< GemmOperand::kB, Kstrided >::project
static CUTLASS_HOST_DEVICE Coord< 3 > project(Coord< 3 > const &coord)
Definition: gemm_operand.h:115
cutlass
Definition: convert.h:33
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::ProjectOperand< GemmOperand::kC, true >::project
static CUTLASS_HOST_DEVICE Coord< 3 > project(Coord< 3 > const &coord)
Definition: gemm_operand.h:128
platform.h
C++ features that may be otherwise unimplemented for CUDA device functions.
cutlass::gemm::GemmMultiplicandTraits::ThreadBlockTile
ThreadBlockTile_ ThreadBlockTile
Shape of GEMM thread block tile (K, N, M)
Definition: gemm_operand.h:70
cutlass::gemm::GemmMultiplicandTraits::Shape
platform::conditional< kKstrided, Shape< 1, ThreadBlockTile::kD, GetExtent< Usage, ThreadBlockTile >::kExtent >, Shape< 1, GetExtent< Usage, ThreadBlockTile >::kExtent, ThreadBlockTile::kD > >::type Shape
Map the ThreadBlockShape onto (kH, kW) dimensions for A and B operand.
Definition: gemm_operand.h:88
cutlass::MatrixLayout::kRowMajor
Definition: matrix_traits.h:36
reshape_tile.h
Defines a type for restructuring a tile.
cutlass::gemm::GemmMultiplicandTraits
Definition: gemm_operand.h:67
cutlass::gemm::GemmOperandTraitsAb::Congruous
static const bool Congruous
Definition: gemm_operand.h:43
cutlass::GemmOperand::kB
Definition: matrix_traits.h:43
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:46
cutlass::platform::conditional
std::conditional (true specialization)
Definition: platform.h:343
static_assert
#define static_assert(__e, __m)
Definition: platform.h:145
cutlass::gemm::GemmMultiplicandTraits::kLayout
static MatrixLayout::Kind const kLayout
Layout of tile.
Definition: gemm_operand.h:79
cutlass::gemm::ProjectOperand< GemmOperand::kA, Kstrided >::project
static CUTLASS_HOST_DEVICE Coord< 3 > project(Coord< 3 > const &coord)
Definition: gemm_operand.h:102
cutlass::Shape
A Shape implementing Layout Concept describing the dimensions of a cube.
Definition: shape.h:64
cutlass::gemm::ProjectOperand
Definition: gemm_operand.h:96
cutlass::gemm::GetExtent
Definition: gemm_operand.h:50
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:48
cutlass::GemmOperand
Gemm operand - D = A * B + C.
Definition: matrix_traits.h:42
cutlass::MatrixLayout::Kind
Kind
Definition: matrix_traits.h:36
cutlass::gemm::ProjectOperand< GemmOperand::kD, true >::project
static CUTLASS_HOST_DEVICE Coord< 3 > project(Coord< 3 > const &coord)
Definition: gemm_operand.h:135
cutlass::GemmOperand::Kind
Kind
Definition: matrix_traits.h:43
cutlass::GemmOperand::kA
Definition: matrix_traits.h:43
cutlass::gemm::GemmMultiplicandTraits::kUsage
static GemmOperand::Kind const kUsage
Identifies multiplicand.
Definition: gemm_operand.h:76
matrix_traits.h
Defines properties of matrices used to denote layout and operands to GEMM kernels.
cutlass::gemm::GemmOperandTraitsAb
Helper to describe attributes of GEMM matrix operands.
Definition: gemm_operand.h:42