CUTLASS
CUDA Templates for Linear Algebra Subroutines and Solvers
Public Types | Public Member Functions | Public Attributes | Static Public Attributes | List of all members
cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable > Class Template Reference

Structure to compute the matrix product targeting CUDA cores and SIMT math instructions.

#include <mma_tensor_op.h>

Public Types

using Shape = Shape_
 Shape of warp-level matrix operation (concept: GemmShape) More...
 
using ElementA = ElementA_
 Data type of multiplicand A. More...
 
using LayoutA = LayoutA_
 Layout of multiplicand A. More...
 
using ElementB = ElementB_
 Data type of multiplicand B. More...
 
using LayoutB = LayoutB_
 Layout of multiplicand B. More...
 
using ElementC = ElementC_
 Data type of accumulator matrix C. More...
 
using LayoutC = LayoutC_
 Layout of accumulator matrix C. More...
 
using Policy = Policy_
 Shape of the warp in units of thread (concept: MmaLanePolicySimt) More...
 
using OperatorClass = arch::OpClassTensorOp
 Indicates class of matrix operator. More...
 
using IteratorA = MmaTensorOpMultiplicandTileIterator< MatrixShape< Shape::kM, Shape::kK >, Operand::kA, ElementA, LayoutA, MatrixShape< Policy::Operator::Shape::kM, Policy::Operator::Shape::kK >, Policy::OpDelta::kRow, kThreadCount, kPartitionsK >
 Iterates over the A operand in memory. More...
 
using FragmentA = typename IteratorA::Fragment
 Storage for A tile. More...
 
using IteratorB = MmaTensorOpMultiplicandTileIterator< MatrixShape< Shape::kK, Shape::kN >, Operand::kB, ElementB, LayoutB, MatrixShape< Policy::Operator::Shape::kK, Policy::Operator::Shape::kN >, Policy::OpDelta::kRow, kThreadCount, kPartitionsK >
 Iterates over the B operand in memory. More...
 
using FragmentB = typename IteratorB::Fragment
 Storage for B tile. More...
 
using IteratorC = MmaTensorOpAccumulatorTileIterator< MatrixShape< Shape::kM, Shape::kN >, ElementC, LayoutC, typename Policy::Operator::Shape, typename Policy::OpDelta >
 Iterates over the C operand in memory. More...
 
using FragmentC = typename IteratorC::Fragment
 Storage for C tile. More...
 

Public Member Functions

CUTLASS_DEVICE MmaTensorOp ()
 Ctor. More...
 
CUTLASS_DEVICE void operator() (FragmentC &D, FragmentA const &A, FragmentB const &B, FragmentC const &C, int const &partitionN_idx=0) const
 Performs a warp-level matrix multiply-accumulate operation. More...
 

Public Attributes

Policy::Operator mma
 Underlying matrix multiply operator (concept: arch::Mma) More...
 

Static Public Attributes

static int const kThreadCount = 32
 Number of threads participating in warp-level matrix product. More...
 
static int const kPartitionsK = PartitionsK_
 Number of partitions along K dimension. More...
 
static int const kPartitionsN = PartitionsN_
 PartitionsN indicating how many PartitionsN for multiplicand B. More...
 

Member Typedef Documentation

template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::ElementA = ElementA_
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::ElementB = ElementB_
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::ElementC = ElementC_
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::FragmentA = typename IteratorA::Fragment
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::FragmentB = typename IteratorB::Fragment
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::FragmentC = typename IteratorC::Fragment
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::IteratorA = MmaTensorOpMultiplicandTileIterator< MatrixShape<Shape::kM, Shape::kK>, Operand::kA, ElementA, LayoutA, MatrixShape<Policy::Operator::Shape::kM, Policy::Operator::Shape::kK>, Policy::OpDelta::kRow, kThreadCount, kPartitionsK>
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::IteratorB = MmaTensorOpMultiplicandTileIterator< MatrixShape<Shape::kK, Shape::kN>, Operand::kB, ElementB, LayoutB, MatrixShape<Policy::Operator::Shape::kK, Policy::Operator::Shape::kN>, Policy::OpDelta::kRow, kThreadCount, kPartitionsK>
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::IteratorC = MmaTensorOpAccumulatorTileIterator< MatrixShape<Shape::kM, Shape::kN>, ElementC, LayoutC, typename Policy::Operator::Shape, typename Policy::OpDelta>
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::LayoutA = LayoutA_
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::LayoutB = LayoutB_
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::LayoutC = LayoutC_
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::OperatorClass = arch::OpClassTensorOp
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::Policy = Policy_
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
using cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::Shape = Shape_

Constructor & Destructor Documentation

template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
CUTLASS_DEVICE cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::MmaTensorOp ( )
inline

Member Function Documentation

template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
CUTLASS_DEVICE void cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::operator() ( FragmentC D,
FragmentA const &  A,
FragmentB const &  B,
FragmentC const &  C,
int const &  partitionN_idx = 0 
) const
inline

Member Data Documentation

template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
int const cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::kPartitionsK = PartitionsK_
static
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
int const cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::kPartitionsN = PartitionsN_
static
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
int const cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::kThreadCount = 32
static
template<typename Shape_ , typename ElementA_ , typename LayoutA_ , typename ElementB_ , typename LayoutB_ , typename ElementC_ , typename LayoutC_ , typename Policy_ , int PartitionsK_ = 1, bool AccumulatorsInRowMajor = false, int PartitionsN_ = 1, typename Enable = bool>
Policy::Operator cutlass::gemm::warp::MmaTensorOp< Shape_, ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_, LayoutC_, Policy_, PartitionsK_, AccumulatorsInRowMajor, PartitionsN_, Enable >::mma

The documentation for this class was generated from the following file: