library.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 <vector>
#include <string>
#include <cstdint>
#include <cuda_runtime.h>

#include "cutlass/cutlass.h"
#include "cutlass/matrix_coord.h"
#include "cutlass/tensor_coord.h"
#include "cutlass/layout/tensor.h"

#include "cutlass/gemm/gemm.h"

namespace cutlass {
namespace library {


enum class LayoutTypeID {
  kUnknown,
  kColumnMajor,
  kRowMajor,
  kColumnMajorInterleavedK4,
  kRowMajorInterleavedK4,
  kColumnMajorInterleavedK16,
  kRowMajorInterleavedK16,
  kTensorNCHW,
  kTensorNHWC,
  kInvalid
};
  
enum class NumericTypeID {
  kUnknown,
  kVoid,
  kB1,
  kU4,
  kU8,
  kU16,
  kU32,
  kU64,
  kS4,
  kS8,
  kS16,
  kS32,
  kS64,
  kF16,
  kF32,
  kF64,
  kCF16,
  kCF32,
  kCF64,
  kCS4,
  kCS8,
  kCS16,
  kCS32,
  kCS64,
  kCU4,
  kCU8,
  kCU16,
  kCU32,
  kCU64,
  kInvalid
};

enum class ComplexTransform {
  kNone,
  kConjugate
};

enum class OperationKind {
  kGemm,
  kInvalid
};

enum class ScalarPointerMode {
  kHost,
  kDevice,
  kInvalid
};

enum class SplitKMode {
  kNone,
  kSerial,
  kParallel,
  kParallelSerial,
  kInvalid
};

enum class OpcodeClassID {
  kSimt,
  kTensorOp,
  kWmmaTensorOp,
  kInvalid
};


enum class GemmKind {
  kGemm,
  kBatched,
  kArray,
  kPlanarComplex,
  kPlanarComplexBatched,
  kInvalid
};


template <typename T> T from_string(std::string const &);

char const *to_string(OperationKind type, bool pretty = false);

template <> OperationKind from_string<OperationKind>(std::string const &str);

char const *to_string(NumericTypeID type, bool pretty = false);

template <> NumericTypeID from_string<NumericTypeID>(std::string const &str);

int sizeof_bits(NumericTypeID type);

bool is_complex_type(NumericTypeID type);

NumericTypeID get_real_type(NumericTypeID type);

bool is_integer_type(NumericTypeID type);

bool is_signed_type(NumericTypeID type);

bool is_signed_integer(NumericTypeID type);

bool is_unsigned_integer(NumericTypeID type);

bool is_float_type(NumericTypeID type);

char const *to_string(Status status, bool pretty = false);

char const *to_string(LayoutTypeID layout, bool pretty = false);

template <> LayoutTypeID from_string<LayoutTypeID>(std::string const &str);

int get_layout_stride_rank(LayoutTypeID layout_id);

char const *to_string(OpcodeClassID type, bool pretty = false);

template <>
OpcodeClassID from_string<OpcodeClassID>(std::string const &str);

std::string lexical_cast(int64_t int_value);

bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string const &str);

std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type);

bool cast_from_int64(std::vector<uint8_t> &bytes, NumericTypeID type, int64_t src);

bool cast_from_uint64(std::vector<uint8_t> &bytes, NumericTypeID type, uint64_t src);

bool cast_from_double(std::vector<uint8_t> &bytes, NumericTypeID type, double src);


struct MathInstructionDescription {

  cutlass::gemm::GemmCoord instruction_shape;

  NumericTypeID element_accumulator;

  OpcodeClassID opcode_class;

  //
  // Methods
  //

  MathInstructionDescription(
    cutlass::gemm::GemmCoord instruction_shape = cutlass::gemm::GemmCoord(),
    NumericTypeID element_accumulator = NumericTypeID::kInvalid,
    OpcodeClassID opcode_class = OpcodeClassID::kInvalid
  ):
    instruction_shape(instruction_shape), element_accumulator(element_accumulator), opcode_class(opcode_class) {}

};

struct TileDescription {

  cutlass::gemm::GemmCoord threadblock_shape;

  int threadblock_stages;

  cutlass::gemm::GemmCoord warp_count;

  MathInstructionDescription math_instruction;

  int minimum_compute_capability;

  int maximum_compute_capability;

  //
  // Methods
  //

  TileDescription(
    cutlass::gemm::GemmCoord threadblock_shape = cutlass::gemm::GemmCoord(),
    int threadblock_stages = 0,
    cutlass::gemm::GemmCoord warp_count = cutlass::gemm::GemmCoord(),
    MathInstructionDescription math_instruction = MathInstructionDescription(),
    int minimum_compute_capability = 0,
    int maximum_compute_capability = 0
  ):
    threadblock_shape(threadblock_shape), 
    threadblock_stages(threadblock_stages), 
    warp_count(warp_count),
    math_instruction(math_instruction),
    minimum_compute_capability(minimum_compute_capability),
    maximum_compute_capability(maximum_compute_capability) { }
};

struct OperationDescription {

  char const * name;

  OperationKind kind;

  TileDescription tile_description;

  //
  // Methods
  //
  OperationDescription(
    char const * name = "unknown",
    OperationKind kind = OperationKind::kInvalid, 
    TileDescription const & tile_description = TileDescription()
  ):
    name(name), kind(kind), tile_description(tile_description) { }
};

struct TensorDescription {

  NumericTypeID element;

  LayoutTypeID layout;

  int alignment;

  int log_extent_range;

  int log_stride_range;

  //
  // Methods
  //
  TensorDescription(
    NumericTypeID element = NumericTypeID::kInvalid,
    LayoutTypeID layout = LayoutTypeID::kInvalid,
    int alignment = 1,
    int log_extent_range = 24,
    int log_stride_range = 24
  ):
    element(element), 
    layout(layout), 
    alignment(alignment), 
    log_extent_range(log_extent_range), 
    log_stride_range(log_stride_range) { }
};


struct GemmDescription : public OperationDescription {

  GemmKind gemm_kind;
  
  TensorDescription A;

  TensorDescription B;

  TensorDescription C;

  NumericTypeID element_epilogue;

  SplitKMode split_k_mode;

  ComplexTransform transform_A;

  ComplexTransform transform_B;

  //
  // Methods
  //

  GemmDescription(
    GemmKind gemm_kind = GemmKind::kGemm,
    TensorDescription const &A = TensorDescription(),
    TensorDescription const &B = TensorDescription(),
    TensorDescription const &C = TensorDescription(),
    NumericTypeID element_epilogue = NumericTypeID::kInvalid,
    SplitKMode split_k_mode = SplitKMode::kNone,
    ComplexTransform transform_A = ComplexTransform::kNone,
    ComplexTransform transform_B = ComplexTransform::kNone
  ):
    gemm_kind(gemm_kind),
    A(A),
    B(B),
    C(C),
    element_epilogue(element_epilogue),
    split_k_mode(split_k_mode),
    transform_A(transform_A),
    transform_B(transform_B) {} 
};


class Operation {
public:

  virtual ~Operation() { }

  virtual OperationDescription const & description() const = 0;

  virtual Status can_implement(
    void const *configuration, 
    void const *arguments) const = 0;
  
  virtual uint64_t get_host_workspace_size(
    void const *configuration) const = 0;
  
  virtual uint64_t get_device_workspace_size(
    void const *configuration) const = 0;
  
  virtual Status initialize(
    void const *configuration, 
    void *host_workspace, 
    void *device_workspace, 
    cudaStream_t stream = nullptr) const = 0;

  virtual Status run(
    void const *arguments,
    void *host_workspace, 
    void *device_workspace = nullptr, 
    cudaStream_t stream = nullptr) const = 0;
};


//
// OperationKind: Gemm
// GemmKind:      Gemm
//
struct GemmConfiguration {

  gemm::GemmCoord problem_size;

  int64_t lda;

  int64_t ldb;

  int64_t ldc;

  int64_t ldd;

  int split_k_slices;
};

struct GemmArguments {

  void const *A;

  void const *B;

  void const *C;

  void *D;

  void const *alpha;

  void const *beta;

  ScalarPointerMode pointer_mode;
};


//
// OperationKind: Gemm
// GemmKind:      Batched

struct GemmBatchedConfiguration {

  gemm::GemmCoord problem_size;

  int64_t lda;

  int64_t ldb;

  int64_t ldc;

  int64_t ldd;

  int64_t batch_stride_A;

  int64_t batch_stride_B;

  int64_t batch_stride_C;

  int64_t batch_stride_D;

  int batch_count;
};

using GemmBatchedArguments = GemmArguments;


//
// OperationKind: Gemm
// GemmKind:      Array

struct GemmArrayConfiguration {

  gemm::GemmCoord problem_size;

  int64_t const *lda;
  int64_t const *ldb;
  int64_t const *ldc;
  int64_t const *ldd;

  int batch_count;
};

struct GemmArrayArguments {
  void const * const *A;
  void const * const *B;
  void const * const *C;
  void * const *D;
  void const *alpha;
  void const *beta;
  ScalarPointerMode pointer_mode;  
};


//
// OperationKind: Gemm
// GemmKind:      Planar complex

struct GemmPlanarComplexConfiguration {

  gemm::GemmCoord problem_size;

  int64_t lda;
  int64_t ldb;
  int64_t ldc;
  int64_t ldd;

  int64_t imag_stride_A;
  int64_t imag_stride_B;
  int64_t imag_stride_C;
  int64_t imag_stride_D;
};

using GemmPlanarComplexArgments = GemmArguments;


//
// OperationKind: Gemm
// GemmKind:      Planar complex batched
//
struct GemmPlanarComplexBatchedConfiguration {

  gemm::GemmCoord problem_size;

  int64_t lda;
  int64_t ldb;
  int64_t ldc;
  int64_t ldd;

  int64_t imag_stride_A;
  int64_t imag_stride_B;
  int64_t imag_stride_C;
  int64_t imag_stride_D;

  int64_t batched_stride_A;
  int64_t batched_stride_B;
  int64_t batched_stride_C;
  int64_t batched_stride_D;
};

using GemmPlanarComplexBatchedArguments = GemmArguments;


} // namespace library
} // namespace cutlass