gemv.h
Go to the documentation of this file.
/***************************************************************************************************
* Copyright (c) 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/array.h"
#include "cutlass/numeric_types.h"
#include "cutlass/matrix_shape.h"
#include "cutlass/gemm/gemm.h"
namespace cutlass {
namespace gemm {
namespace threadblock {
template <
class Core_ //< GemvCore
>
class Gemv {
public:
using Shape = typename Core_::Shape;
using Operator = typename Core_::Operator;
using IteratorA = typename Core_::IteratorA;
using IteratorB = typename Core_::IteratorB;
using IteratorC = typename Core_::IteratorC;
using FragmentA = typename IteratorA::Fragment;
using FragmentB = typename IteratorB::Fragment;
using FragmentC = typename Operator::FragmentC;
using ThreadShape = typename Core_::ThreadShape;
public:
CUTLASS_DEVICE
Gemv() { }
CUTLASS_DEVICE
void operator()(
GemmCoord const &problem_size,
FragmentC &accum,
IteratorA iterator_A,
IteratorB iterator_B,
FragmentC const &src_accum) {
//
// Prologue
//
FragmentA frag_A;
FragmentB frag_B;
frag_A.clear();
frag_B.clear();
iterator_A.load(frag_A);
iterator_B.load(frag_B);
++iterator_A;
++iterator_B;
//
// Mainloop
//
Operator thread_mma;
int gemm_k = problem_size.k();
if (gemm_k < Shape::kK)
{
iterator_A.clear_mask();
iterator_B.clear_mask();
}
// iterate over K to accumulate result
CUTLASS_GEMM_LOOP
for (; gemm_k > 0; gemm_k -= Shape::kK) {
thread_mma(accum, frag_A, frag_B, accum);
iterator_A.load(frag_A);
iterator_B.load(frag_B);
++iterator_A;
++iterator_B;
if (gemm_k < Shape::kK)
{
iterator_A.clear_mask();
iterator_B.clear_mask();
}
}
}
};
} // namespace threadblock
} // namespace gemm
} // namespace cutlass
cutlass
Definition: aligned_buffer.h:35
cutlass::gemm::threadblock::Gemv::Gemv
CUTLASS_DEVICE Gemv()
Definition: gemv.h:80
cutlass::gemm::threadblock::Gemv::FragmentB
typename IteratorB::Fragment FragmentB
Fragment of operand B loaded from global memory.
Definition: gemv.h:70
cutlass::gemm::GemmCoord
Definition: include/cutlass/gemm/gemm.h:94
gemm.h
Defines common types used for all GEMM-like operators.
cutlass::gemm::threadblock::Gemv::Operator
typename Core_::Operator Operator
The MMA operator that computes GEMV.
Definition: gemv.h:55
cutlass::gemm::GemmCoord::k
CUTLASS_HOST_DEVICE Index const & k() const
Returns the GEMM K coordinate.
Definition: include/cutlass/gemm/gemm.h:145
array.h
Statically sized array of elements that accommodates all CUTLASS-supported numeric types and is safe ...
cutlass::gemm::threadblock::Gemv
Structure to compute the matrix-vector product using SIMT math instructions.
Definition: gemv.h:50
matrix_shape.h
Defines a Shape template for matrix tiles.
numeric_types.h
Top-level include for all CUTLASS numeric types.
cutlass::gemm::threadblock::Gemv::ThreadShape
typename Core_::ThreadShape ThreadShape
Shape of the per-thread GEMV operation.
Definition: gemv.h:76
cutlass::gemm::threadblock::Gemv::IteratorA
typename Core_::IteratorA IteratorA
Iterates over A in global memory.
Definition: gemv.h:58
cutlass::gemm::threadblock::Gemv::FragmentC
typename Operator::FragmentC FragmentC
Fragment of operand accumulator loaded/stored to global memory.
Definition: gemv.h:73
cutlass::gemm::threadblock::Gemv::IteratorB
typename Core_::IteratorB IteratorB
Iterates over B in global memory.
Definition: gemv.h:61
CUTLASS_GEMM_LOOP
#define CUTLASS_GEMM_LOOP
Definition: cutlass.h:112
cutlass::gemm::threadblock::Gemv::FragmentA
typename IteratorA::Fragment FragmentA
Fragment of operand A loaded from global memory.
Definition: gemv.h:67
cutlass::gemm::threadblock::Gemv::operator()
CUTLASS_DEVICE void operator()(GemmCoord const &problem_size, FragmentC &accum, IteratorA iterator_A, IteratorB iterator_B, FragmentC const &src_accum)
Definition: gemv.h:83
cutlass.h
Basic include for CUTLASS.
cutlass::gemm::threadblock::Gemv::Shape
typename Core_::Shape Shape
Definition: gemv.h:52
cutlass::gemm::threadblock::Gemv::IteratorC
typename Core_::IteratorC IteratorC
Fragment of operand C loaded from global memory.
Definition: gemv.h:64