include/cutlass/gemm/gemm.h
Go to the documentation of this file.
/***************************************************************************************************
* Copyright (c) 2017-2019, 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/cutlass.h"
#include "cutlass/coord.h"
namespace cutlass {
namespace gemm {
enum class Operand {
kA,
kB,
kC,
kD
};
template <
int M = 1,
int N = 1,
int K = 1
>
struct GemmShape {
static int const kM = M;
static int const kN = N;
static int const kK = K;
static int const kMN = M * N;
static int const kMK = M * K;
static int const kKN = N * K;
static int const kMNK = M * N * K;
static int const kCount = kMNK;
//
// Static member functions
//
CUTLASS_HOST_DEVICE
static Coord<3> toCoord() {
return make_Coord(kM, kN, kK);
}
};
template <
typename Shape
>
using GemmShapeTranspose = GemmShape<Shape::kN, Shape::kM, Shape::kK>;
struct GemmCoord : public Coord<3, int> {
typedef int Index;
typedef Coord<3, Index> Base;
static int const kM = 0;
static int const kN = 1;
static int const kK = 2;
//
// Methods
//
CUTLASS_HOST_DEVICE
GemmCoord() { }
CUTLASS_HOST_DEVICE
GemmCoord(Coord<3, Index> const &coord): Base(make_Coord(coord[0], coord[1], coord[2])) { }
CUTLASS_HOST_DEVICE
GemmCoord(Index m, Index n, Index k): Base(make_Coord(m, n, k)) { }
CUTLASS_HOST_DEVICE
Index const & m() const { return this->at(kM); }
CUTLASS_HOST_DEVICE
Index & m() { return this->at(kM); }
CUTLASS_HOST_DEVICE
Index const & n() const { return this->at(kN); }
CUTLASS_HOST_DEVICE
Index & n() { return this->at(kN); }
CUTLASS_HOST_DEVICE
Index const & k() const { return this->at(kK); }
CUTLASS_HOST_DEVICE
Index & k() { return this->at(kK); }
CUTLASS_HOST_DEVICE
Coord<3> mnk() const {
return make_Coord(m(), n(), k());
}
CUTLASS_HOST_DEVICE
Coord<3> knm() const {
return make_Coord(k(), n(), m());
}
CUTLASS_HOST_DEVICE
Coord<2> nm() const {
return make_Coord(n(), m());
}
CUTLASS_HOST_DEVICE
Coord<2> mn() const {
return make_Coord(m(), n());
}
CUTLASS_HOST_DEVICE
Coord<2> mk() const {
return make_Coord(m(), k());
}
CUTLASS_HOST_DEVICE
Coord<2> km() const {
return make_Coord(k(), m());
}
CUTLASS_HOST_DEVICE
Coord<2> nk() const {
return make_Coord(n(), k());
}
CUTLASS_HOST_DEVICE
Coord<2> kn() const {
return make_Coord(k(), n());
}
//
// Coord operators
//
CUTLASS_HOST_DEVICE
GemmCoord operator+(Base const& b) const {
return GemmCoord(Base::operator+(b));
}
CUTLASS_HOST_DEVICE
GemmCoord operator-(Base const& b) const {
return GemmCoord(Base::operator-(b));
}
CUTLASS_HOST_DEVICE
GemmCoord operator*(Base const& b) const {
return GemmCoord(Base::operator*(b));
}
CUTLASS_HOST_DEVICE
GemmCoord operator/(Base const& b) const {
return GemmCoord(Base::operator/(b));
}
CUTLASS_HOST_DEVICE
GemmCoord& operator+=(Base const& b) {
Base::operator+=(b);
return *this;
}
CUTLASS_HOST_DEVICE
GemmCoord& operator-=(Base const& b) {
Base::operator-=(b);
return *this;
}
CUTLASS_HOST_DEVICE
GemmCoord& operator*=(Base const& b) {
Base::operator*=(b);
return *this;
}
CUTLASS_HOST_DEVICE
GemmCoord& operator/=(Base const& b) {
Base::operator/=(b);
return *this;
}
};
struct BatchedGemmCoord : public Coord<4, int> {
typedef int Index;
typedef Coord<4, Index> Base;
static int const kM = 0;
static int const kN = 1;
static int const kK = 2;
static int const kBatch = 3;
//
// Methods
//
CUTLASS_HOST_DEVICE
BatchedGemmCoord() { }
CUTLASS_HOST_DEVICE
BatchedGemmCoord(Base const &coord): Base(coord) { }
CUTLASS_HOST_DEVICE
BatchedGemmCoord(Index m, Index n, Index k, Index b): Base(make_Coord(m, n, k, b)) { }
CUTLASS_HOST_DEVICE
Index const & m() const { return this->at(kM); }
CUTLASS_HOST_DEVICE
Index & m() { return this->at(kM); }
CUTLASS_HOST_DEVICE
Index const & n() const { return this->at(kN); }
CUTLASS_HOST_DEVICE
Index & n() { return this->at(kN); }
CUTLASS_HOST_DEVICE
Index const & k() const { return this->at(kK); }
CUTLASS_HOST_DEVICE
Index & k() { return this->at(kK); }
CUTLASS_HOST_DEVICE
Index const & batch() const { return this->at(kBatch); }
CUTLASS_HOST_DEVICE
Index & batch() { return this->at(kBatch); }
CUTLASS_HOST_DEVICE
GemmCoord mnk() const {
return GemmCoord(m(), n(), k());
}
CUTLASS_HOST_DEVICE
Coord<4> mnkb() const {
return make_Coord(m(), n(), k(), batch());
}
//
// Coord operators
//
CUTLASS_HOST_DEVICE
BatchedGemmCoord operator+(Base const& b) const {
return BatchedGemmCoord(Base::operator+(b));
}
CUTLASS_HOST_DEVICE
BatchedGemmCoord operator-(Base const& b) const {
return BatchedGemmCoord(Base::operator-(b));
}
CUTLASS_HOST_DEVICE
BatchedGemmCoord operator*(Base const& b) const {
return BatchedGemmCoord(Base::operator*(b));
}
CUTLASS_HOST_DEVICE
BatchedGemmCoord operator/(Base const& b) const {
return BatchedGemmCoord(Base::operator/(b));
}
CUTLASS_HOST_DEVICE
BatchedGemmCoord& operator+=(Base const& b) {
Base::operator+=(b);
return *this;
}
CUTLASS_HOST_DEVICE
BatchedGemmCoord& operator-=(Base const& b) {
Base::operator-=(b);
return *this;
}
CUTLASS_HOST_DEVICE
BatchedGemmCoord& operator*=(Base const& b) {
Base::operator*=(b);
return *this;
}
CUTLASS_HOST_DEVICE
BatchedGemmCoord& operator/=(Base const& b) {
Base::operator/=(b);
return *this;
}
};
} // namespace gemm
} // namespace cutlass
cutlass::gemm::BatchedGemmCoord::mnkb
CUTLASS_HOST_DEVICE Coord< 4 > mnkb() const
Obtains a Coord<4> from BatchedGemmCoord.
Definition: include/cutlass/gemm/gemm.h:336
cutlass::gemm::BatchedGemmCoord::Base
Coord< 4, Index > Base
Base type is a Coord of rank=4.
Definition: include/cutlass/gemm/gemm.h:266
cutlass::gemm::BatchedGemmCoord::m
CUTLASS_HOST_DEVICE Index & m()
Returns reference to the GEMM M coordinate.
Definition: include/cutlass/gemm/gemm.h:302
cutlass
Definition: aligned_buffer.h:35
cutlass::gemm::BatchedGemmCoord::operator/
CUTLASS_HOST_DEVICE BatchedGemmCoord operator/(Base const &b) const
Element-wise division.
Definition: include/cutlass/gemm/gemm.h:364
cutlass::gemm::GemmCoord::operator/=
CUTLASS_HOST_DEVICE GemmCoord & operator/=(Base const &b)
In-place division.
Definition: include/cutlass/gemm/gemm.h:250
cutlass::gemm::GemmCoord::Index
int Index
Integer-valued index.
Definition: include/cutlass/gemm/gemm.h:97
cutlass::gemm::GemmCoord::GemmCoord
CUTLASS_HOST_DEVICE GemmCoord(Coord< 3, Index > const &coord)
Constructs from Coord<3> and a batch.
Definition: include/cutlass/gemm/gemm.h:121
cutlass::gemm::BatchedGemmCoord::mnk
CUTLASS_HOST_DEVICE GemmCoord mnk() const
Obtains a GemmCoord from BatchedGemmCoord.
Definition: include/cutlass/gemm/gemm.h:330
cutlass::gemm::GemmCoord::operator+
CUTLASS_HOST_DEVICE GemmCoord operator+(Base const &b) const
Element-wise addition.
Definition: include/cutlass/gemm/gemm.h:205
coord.h
A Coord is a coordinate of arbitrary rank into a tensor or matrix.
cutlass::make_Coord
CUTLASS_HOST_DEVICE Coord< 1 > make_Coord(int _0)
Helper to make a 2-element coordinate.
Definition: coord.h:387
cutlass::gemm::GemmCoord::m
CUTLASS_HOST_DEVICE Index & m()
Returns reference to the GEMM M coordinate.
Definition: include/cutlass/gemm/gemm.h:133
cutlass::gemm::Operand
Operand
GEMM operand enumeration: D = A * B + C.
Definition: include/cutlass/gemm/gemm.h:39
cutlass::gemm::GemmCoord
Definition: include/cutlass/gemm/gemm.h:94
cutlass::gemm::GemmCoord::mn
CUTLASS_HOST_DEVICE Coord< 2 > mn() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:171
cutlass::operator/=
CUTLASS_HOST_DEVICE half_t & operator/=(half_t &lhs, half_t const &rhs)
Definition: half.h:684
cutlass::gemm::BatchedGemmCoord::operator-
CUTLASS_HOST_DEVICE BatchedGemmCoord operator-(Base const &b) const
Element-wise subtraction.
Definition: include/cutlass/gemm/gemm.h:352
cutlass::gemm::BatchedGemmCoord::batch
CUTLASS_HOST_DEVICE Index & batch()
Returns reference to the GEMM batch coordinate.
Definition: include/cutlass/gemm/gemm.h:326
cutlass::gemm::GemmCoord::n
CUTLASS_HOST_DEVICE Index const & n() const
Returns the GEMM N coordinate.
Definition: include/cutlass/gemm/gemm.h:137
cutlass::gemm::GemmCoord::nm
CUTLASS_HOST_DEVICE Coord< 2 > nm() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:165
cutlass::gemm::BatchedGemmCoord::operator+
CUTLASS_HOST_DEVICE BatchedGemmCoord operator+(Base const &b) const
Element-wise addition.
Definition: include/cutlass/gemm/gemm.h:346
cutlass::gemm::GemmCoord::k
CUTLASS_HOST_DEVICE Index & k()
Returns reference to the GEMM K coordinate.
Definition: include/cutlass/gemm/gemm.h:149
cutlass::gemm::BatchedGemmCoord::n
CUTLASS_HOST_DEVICE Index & n()
Returns reference to the GEMM N coordinate.
Definition: include/cutlass/gemm/gemm.h:310
cutlass::gemm::GemmCoord::operator/
CUTLASS_HOST_DEVICE GemmCoord operator/(Base const &b) const
Element-wise division.
Definition: include/cutlass/gemm/gemm.h:223
cutlass::gemm::GemmCoord::GemmCoord
CUTLASS_HOST_DEVICE GemmCoord(Index m, Index n, Index k)
Helper to construct from a K, N, M, batch variables.
Definition: include/cutlass/gemm/gemm.h:125
cutlass::gemm::BatchedGemmCoord::Index
int Index
Integer-valued index.
Definition: include/cutlass/gemm/gemm.h:263
cutlass::operator+=
CUTLASS_HOST_DEVICE half_t & operator+=(half_t &lhs, half_t const &rhs)
Definition: half.h:654
cutlass::gemm::GemmCoord::nk
CUTLASS_HOST_DEVICE Coord< 2 > nk() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:189
cutlass::gemm::BatchedGemmCoord::BatchedGemmCoord
CUTLASS_HOST_DEVICE BatchedGemmCoord(Base const &coord)
Constructs from Coord<4>
Definition: include/cutlass/gemm/gemm.h:290
cutlass::gemm::GemmCoord::k
CUTLASS_HOST_DEVICE Index const & k() const
Returns the GEMM K coordinate.
Definition: include/cutlass/gemm/gemm.h:145
cutlass::gemm::Operand::kC
B multiplicand.
cutlass::operator-=
CUTLASS_HOST_DEVICE half_t & operator-=(half_t &lhs, half_t const &rhs)
Definition: half.h:664
cutlass::gemm::GemmCoord::Base
Coord< 3, Index > Base
Base type is a Coord of rank=4.
Definition: include/cutlass/gemm/gemm.h:100
cutlass::gemm::Operand::kA
cutlass::gemm::GemmShape::toCoord
static CUTLASS_HOST_DEVICE Coord< 3 > toCoord()
Returns a Coord object.
Definition: include/cutlass/gemm/gemm.h:76
cutlass::gemm::GemmCoord::km
CUTLASS_HOST_DEVICE Coord< 2 > km() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:183
cutlass::gemm::BatchedGemmCoord
Definition: include/cutlass/gemm/gemm.h:260
cutlass::gemm::GemmCoord::mnk
CUTLASS_HOST_DEVICE Coord< 3 > mnk() const
Obtains a Coord<3> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:153
cutlass::gemm::GemmCoord::operator-
CUTLASS_HOST_DEVICE GemmCoord operator-(Base const &b) const
Element-wise subtraction.
Definition: include/cutlass/gemm/gemm.h:211
cutlass::gemm::BatchedGemmCoord::batch
CUTLASS_HOST_DEVICE Index const & batch() const
Returns the GEMM batch coordinate.
Definition: include/cutlass/gemm/gemm.h:322
cutlass::gemm::BatchedGemmCoord::operator*=
CUTLASS_HOST_DEVICE BatchedGemmCoord & operator*=(Base const &b)
In-place multiplication.
Definition: include/cutlass/gemm/gemm.h:384
cutlass::gemm::BatchedGemmCoord::operator*
CUTLASS_HOST_DEVICE BatchedGemmCoord operator*(Base const &b) const
Element-wise multiplication.
Definition: include/cutlass/gemm/gemm.h:358
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:89
cutlass::gemm::BatchedGemmCoord::k
CUTLASS_HOST_DEVICE Index const & k() const
Returns the GEMM K coordinate.
Definition: include/cutlass/gemm/gemm.h:314
cutlass::gemm::GemmShape
Shape of a matrix multiply-add operation.
Definition: include/cutlass/gemm/gemm.h:57
cutlass::gemm::GemmCoord::mk
CUTLASS_HOST_DEVICE Coord< 2 > mk() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:177
cutlass::gemm::BatchedGemmCoord::BatchedGemmCoord
CUTLASS_HOST_DEVICE BatchedGemmCoord()
Default ctor.
Definition: include/cutlass/gemm/gemm.h:286
cutlass::gemm::BatchedGemmCoord::k
CUTLASS_HOST_DEVICE Index & k()
Returns reference to the GEMM K coordinate.
Definition: include/cutlass/gemm/gemm.h:318
cutlass::gemm::Operand::kD
Source accumulator.
cutlass::gemm::GemmCoord::knm
CUTLASS_HOST_DEVICE Coord< 3 > knm() const
Obtains a Coord<3> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:159
cutlass::gemm::GemmCoord::operator+=
CUTLASS_HOST_DEVICE GemmCoord & operator+=(Base const &b)
In-place addition.
Definition: include/cutlass/gemm/gemm.h:229
cutlass::operator*=
CUTLASS_HOST_DEVICE half_t & operator*=(half_t &lhs, half_t const &rhs)
Definition: half.h:674
cutlass::Coord< 3 >
cutlass::gemm::GemmCoord::GemmCoord
CUTLASS_HOST_DEVICE GemmCoord()
Default ctor.
Definition: include/cutlass/gemm/gemm.h:117
cutlass::gemm::BatchedGemmCoord::operator/=
CUTLASS_HOST_DEVICE BatchedGemmCoord & operator/=(Base const &b)
In-place division.
Definition: include/cutlass/gemm/gemm.h:391
cutlass::gemm::GemmCoord::operator-=
CUTLASS_HOST_DEVICE GemmCoord & operator-=(Base const &b)
In-place subtraction.
Definition: include/cutlass/gemm/gemm.h:236
cutlass::gemm::BatchedGemmCoord::operator+=
CUTLASS_HOST_DEVICE BatchedGemmCoord & operator+=(Base const &b)
In-place addition.
Definition: include/cutlass/gemm/gemm.h:370
cutlass::gemm::GemmCoord::kn
CUTLASS_HOST_DEVICE Coord< 2 > kn() const
Obtains a Coord<2> from GemmCoord.
Definition: include/cutlass/gemm/gemm.h:195
cutlass::gemm::BatchedGemmCoord::n
CUTLASS_HOST_DEVICE Index const & n() const
Returns the GEMM N coordinate.
Definition: include/cutlass/gemm/gemm.h:306
cutlass::gemm::BatchedGemmCoord::BatchedGemmCoord
CUTLASS_HOST_DEVICE BatchedGemmCoord(Index m, Index n, Index k, Index b)
Helper to construct from a K, N, M, and batch variables.
Definition: include/cutlass/gemm/gemm.h:294
cutlass::gemm::GemmCoord::m
CUTLASS_HOST_DEVICE Index const & m() const
Returns the GEMM M coordinate.
Definition: include/cutlass/gemm/gemm.h:129
cutlass::gemm::GemmCoord::operator*=
CUTLASS_HOST_DEVICE GemmCoord & operator*=(Base const &b)
In-place multiplication.
Definition: include/cutlass/gemm/gemm.h:243
cutlass::gemm::Operand::kB
A multiplicand.
cutlass::gemm::BatchedGemmCoord::operator-=
CUTLASS_HOST_DEVICE BatchedGemmCoord & operator-=(Base const &b)
In-place subtraction.
Definition: include/cutlass/gemm/gemm.h:377
cutlass.h
Basic include for CUTLASS.
cutlass::gemm::GemmCoord::n
CUTLASS_HOST_DEVICE Index & n()
Returns reference to the GEMM N coordinate.
Definition: include/cutlass/gemm/gemm.h:141
cutlass::gemm::GemmCoord::operator*
CUTLASS_HOST_DEVICE GemmCoord operator*(Base const &b) const
Element-wise multiplication.
Definition: include/cutlass/gemm/gemm.h:217
cutlass::gemm::BatchedGemmCoord::m
CUTLASS_HOST_DEVICE Index const & m() const
Returns the GEMM M coordinate.
Definition: include/cutlass/gemm/gemm.h:298