CUTLASS 2.2 (#96)
Adds support for NVIDIA Ampere Architecture features. CUDA 11 Toolkit recommended.
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
# Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
|
||||
# Copyright (c) 2017-2020, 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:
|
||||
@@ -22,7 +22,7 @@
|
||||
|
||||
include(GNUInstallDirs)
|
||||
|
||||
find_package(Python3 3.6 COMPONENTS Interpreter REQUIRED)
|
||||
find_package(Python3 3.5 COMPONENTS Interpreter REQUIRED)
|
||||
|
||||
add_library(cutlass_library_includes INTERFACE)
|
||||
add_library(nvidia::cutlass::library::includes ALIAS cutlass_library_includes)
|
||||
@@ -59,7 +59,7 @@ cutlass_add_library(
|
||||
src/operation_table.cu
|
||||
src/singleton.cu
|
||||
src/util.cu
|
||||
|
||||
|
||||
)
|
||||
|
||||
file(GLOB_RECURSE GENERATOR_PYTHON_SOURCES CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/scripts/*.py)
|
||||
|
||||
@@ -45,6 +45,9 @@ private:
|
||||
/// Host workspace
|
||||
static int const kHostWorkspaceSize = (4 << 10);
|
||||
|
||||
/// Provider of operations
|
||||
Provider provider_;
|
||||
|
||||
/// CUDA device properties
|
||||
cudaDeviceProp device_;
|
||||
|
||||
@@ -90,6 +93,12 @@ public:
|
||||
/// Gets the current CUDA stream
|
||||
cudaStream_t get_stream() const;
|
||||
|
||||
/// Gets the current provider
|
||||
Provider get_provider() const;
|
||||
|
||||
/// Sets the provider of operations
|
||||
void set_provider(Provider provider);
|
||||
|
||||
/// Gets the device workspace size
|
||||
size_t get_workspace_size() const;
|
||||
|
||||
@@ -149,6 +158,56 @@ public:
|
||||
void * ptr_D, /// Pointer to D matrix
|
||||
int ldd /// Leading dimension of D matrix
|
||||
);
|
||||
|
||||
/// Executes a GEMM computation: D <= alpha * A*B + beta * C.
|
||||
//
|
||||
// Supports batched-strided, batched array or split-K serial or split-K parallel.
|
||||
//
|
||||
Status gemm_universal(
|
||||
|
||||
GemmUniversalMode mode, /// indicates the mode in which the kUniversal GEMM is launched
|
||||
|
||||
int M, /// GEMM M dimension
|
||||
int N, /// GEMM N dimension
|
||||
int K, /// GEMM K dimension
|
||||
|
||||
NumericTypeID element_compute, /// Data type of internal accumulation
|
||||
|
||||
NumericTypeID element_scalar, /// Data type of alpha/beta scalars
|
||||
|
||||
void const *alpha, /// Pointer to alpha scalar
|
||||
|
||||
NumericTypeID element_A, /// Data type of A matrix elements
|
||||
LayoutTypeID layout_A, /// Layout of A matrix
|
||||
ComplexTransform transform_A, /// Complex transformation applied to A matrix - ignored for real-valued matrices
|
||||
|
||||
void const * ptr_A, /// Pointer to A matrix in Global Memory
|
||||
int lda, /// Leading dimension of A matrix
|
||||
|
||||
NumericTypeID element_B, /// Data type of B matrix elements
|
||||
LayoutTypeID layout_B, /// Layout of B matrix
|
||||
ComplexTransform transform_B, /// Complex transformation applied to B matrix - ignored for real-valued matrices
|
||||
|
||||
void const * ptr_B, /// Pointer to B matrix in Global Memory
|
||||
int ldb, /// Leading dimension of B matrix
|
||||
|
||||
void const * beta, /// Pointer to beta scalar
|
||||
|
||||
NumericTypeID element_C, /// Data type of C and D matrices
|
||||
|
||||
void const * ptr_C, /// Pointer to C matrix
|
||||
int ldc, /// Leading dimension of C matrix
|
||||
|
||||
void * ptr_D, /// Pointer to D matrix
|
||||
int ldd, /// Leading dimension of D matrix
|
||||
|
||||
int batch_count = 1, /// Batch count or number of split-K slices
|
||||
|
||||
int64_t batch_stride_A = 0, /// Batch stride of A operand
|
||||
int64_t batch_stride_B = 0, /// Batch stride of B operand
|
||||
int64_t batch_stride_C = 0, /// Batch stride of C operand
|
||||
int64_t batch_stride_D = 0 /// Batch stride of D operand
|
||||
);
|
||||
|
||||
/// Planar complex GEMM
|
||||
///
|
||||
@@ -276,7 +335,6 @@ public:
|
||||
using HandlePtr = std::unique_ptr<Handle>;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace library
|
||||
} // namespace cutlass
|
||||
|
||||
|
||||
@@ -44,6 +44,7 @@
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <cstdint>
|
||||
#include <stdexcept>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
@@ -93,10 +94,14 @@ enum class NumericTypeID {
|
||||
kS32,
|
||||
kS64,
|
||||
kF16,
|
||||
kBF16,
|
||||
kTF32,
|
||||
kF32,
|
||||
kF64,
|
||||
kCF16,
|
||||
kCBF16,
|
||||
kCF32,
|
||||
kCTF32,
|
||||
kCF64,
|
||||
kCS4,
|
||||
kCS8,
|
||||
@@ -120,6 +125,7 @@ enum class ComplexTransform {
|
||||
|
||||
/// Providers
|
||||
enum class Provider {
|
||||
kNone,
|
||||
kCUTLASS,
|
||||
kReferenceHost,
|
||||
kReferenceDevice,
|
||||
@@ -132,6 +138,8 @@ enum class Provider {
|
||||
/// Enumeration indicating the kind of operation
|
||||
enum class OperationKind {
|
||||
kGemm,
|
||||
kEqGemm,
|
||||
kReduction,
|
||||
kInvalid
|
||||
};
|
||||
|
||||
@@ -160,9 +168,11 @@ enum class OpcodeClassID {
|
||||
};
|
||||
|
||||
enum class MathOperationID {
|
||||
kAdd,
|
||||
kMultiplyAdd,
|
||||
kMultiplyAddSaturate,
|
||||
kMultiplyAddComplex,
|
||||
kMultiplyAddGaussianComplex,
|
||||
kXorPopc,
|
||||
kInvalid
|
||||
};
|
||||
@@ -180,12 +190,17 @@ enum class GemmKind {
|
||||
kInvalid
|
||||
};
|
||||
|
||||
/// Mode of GEMM
|
||||
enum class GemmUniversalMode {
|
||||
kGemm,
|
||||
kGemmSplitKParallel,
|
||||
kBatched,
|
||||
kArray,
|
||||
/// Mode of Universal GEMM
|
||||
using GemmUniversalMode = cutlass::gemm::GemmUniversalMode;
|
||||
|
||||
enum class EpilogueKind {
|
||||
kUnknown,
|
||||
kConversion,
|
||||
kLinearCombination,
|
||||
kLinearCombinationClamp,
|
||||
kLinearCombinationPlanarComplex,
|
||||
kLinearCombinationRelu,
|
||||
kLinearCombinationSigmoid,
|
||||
kInvalid
|
||||
};
|
||||
|
||||
@@ -220,6 +235,22 @@ struct MathInstructionDescription {
|
||||
opcode_class(opcode_class),
|
||||
math_operation(math_operation) {}
|
||||
|
||||
// Equality operator
|
||||
inline
|
||||
bool operator==(MathInstructionDescription const& rhs) const{
|
||||
return (
|
||||
(instruction_shape == rhs.instruction_shape) &&
|
||||
(element_accumulator == rhs.element_accumulator) &&
|
||||
(opcode_class == rhs.opcode_class) &&
|
||||
(math_operation == rhs.math_operation));
|
||||
}
|
||||
|
||||
// Inequality operator
|
||||
inline
|
||||
bool operator!=(MathInstructionDescription const& rhs) const {
|
||||
return !(*this == rhs);
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/// Structure describing the tiled structure of a GEMM-like computation
|
||||
@@ -261,6 +292,24 @@ struct TileDescription {
|
||||
math_instruction(math_instruction),
|
||||
minimum_compute_capability(minimum_compute_capability),
|
||||
maximum_compute_capability(maximum_compute_capability) { }
|
||||
|
||||
// Equality operator
|
||||
inline
|
||||
bool operator==(TileDescription const& rhs) const{
|
||||
return (
|
||||
(threadblock_shape == rhs.threadblock_shape) &&
|
||||
(threadblock_stages == rhs.threadblock_stages) &&
|
||||
(warp_count == rhs.warp_count) &&
|
||||
(math_instruction == rhs.math_instruction) &&
|
||||
(minimum_compute_capability == rhs.minimum_compute_capability) &&
|
||||
(maximum_compute_capability == rhs.maximum_compute_capability));
|
||||
}
|
||||
|
||||
// Inequality operator
|
||||
inline
|
||||
bool operator!=(TileDescription const& rhs) const {
|
||||
return !(*this == rhs);
|
||||
}
|
||||
};
|
||||
|
||||
/// High-level description of an operation
|
||||
@@ -379,6 +428,20 @@ struct GemmDescription : public OperationDescription {
|
||||
transform_B(transform_B) {}
|
||||
};
|
||||
|
||||
|
||||
/// Description of all Reduction operations
|
||||
struct ReductionDescription : public OperationDescription {
|
||||
|
||||
/// Describes the data type of workspace
|
||||
NumericTypeID element_workspace;
|
||||
|
||||
/// Describes the data type of final output
|
||||
NumericTypeID element_output;
|
||||
|
||||
/// Describes the data type of the scalars passed to the epilogue
|
||||
NumericTypeID element_epilogue;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -549,6 +612,42 @@ struct GemmArrayArguments {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Universal GEMM supporting multiple split-K modes, multiple batched modes, real and complex
|
||||
//
|
||||
// OperationKind: Gemm
|
||||
// GemmKind: Universal
|
||||
|
||||
struct GemmUniversalConfiguration {
|
||||
|
||||
GemmUniversalMode mode;
|
||||
gemm::GemmCoord problem_size;
|
||||
int batch_count;
|
||||
|
||||
int64_t lda;
|
||||
int64_t ldb;
|
||||
int64_t ldc;
|
||||
int64_t ldd;
|
||||
};
|
||||
|
||||
struct GemmUniversalArguments {
|
||||
|
||||
void const *A;
|
||||
void const *B;
|
||||
void const *C;
|
||||
void *D;
|
||||
|
||||
void const *alpha;
|
||||
void const *beta;
|
||||
ScalarPointerMode pointer_mode;
|
||||
|
||||
int64_t batch_stride_A;
|
||||
int64_t batch_stride_B;
|
||||
int64_t batch_stride_C;
|
||||
int64_t batch_stride_D;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Complex valued GEMM in which real and imaginary parts are separated by a stride
|
||||
//
|
||||
// OperationKind: Gemm
|
||||
@@ -648,7 +747,6 @@ struct GemmPlanarComplexArrayArguments {
|
||||
ScalarPointerMode pointer_mode;
|
||||
};
|
||||
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace library
|
||||
|
||||
@@ -45,6 +45,13 @@ namespace cutlass {
|
||||
namespace library {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// Forward declaration
|
||||
class Manifest;
|
||||
|
||||
// init and insert all cutlass gemm and conv2d op in manifest object (procedurally generated using generator.py)
|
||||
void initialize_all(Manifest &manifest);
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// List of operations
|
||||
using OperationVector = std::vector<std::unique_ptr<Operation>>;
|
||||
|
||||
@@ -29,24 +29,28 @@
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <fstream>
|
||||
#include <iosfwd>
|
||||
#include <unordered_map>
|
||||
#include <algorithm>
|
||||
|
||||
#include "cutlass/library/library.h"
|
||||
#include "cutlass/library/manifest.h"
|
||||
|
||||
#include "cutlass/library/util.h"
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace library {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// Data Structures for Gemm Functional Maps
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tuple uniquely identifying functional behavior
|
||||
/// Tuple uniquely identifying Gemm functional behavior
|
||||
struct GemmFunctionalKey {
|
||||
|
||||
Provider provider;
|
||||
GemmKind gemm_kind;
|
||||
NumericTypeID element_compute;
|
||||
NumericTypeID element_scalar;
|
||||
NumericTypeID element_A;
|
||||
@@ -63,6 +67,8 @@ struct GemmFunctionalKey {
|
||||
|
||||
inline
|
||||
GemmFunctionalKey(
|
||||
Provider provider,
|
||||
GemmKind gemm_kind = GemmKind::kGemm,
|
||||
NumericTypeID element_compute = NumericTypeID::kF32,
|
||||
NumericTypeID element_scalar = NumericTypeID::kF32,
|
||||
NumericTypeID element_A = NumericTypeID::kF16,
|
||||
@@ -73,6 +79,8 @@ struct GemmFunctionalKey {
|
||||
ComplexTransform transform_B = ComplexTransform::kNone,
|
||||
NumericTypeID element_C = NumericTypeID::kF16
|
||||
):
|
||||
provider(provider),
|
||||
gemm_kind(gemm_kind),
|
||||
element_compute(element_compute),
|
||||
element_scalar(element_scalar),
|
||||
element_A(element_A),
|
||||
@@ -87,6 +95,8 @@ struct GemmFunctionalKey {
|
||||
inline
|
||||
bool operator==(GemmFunctionalKey const &rhs) const {
|
||||
return
|
||||
(provider == rhs.provider) &&
|
||||
(gemm_kind == rhs.gemm_kind) &&
|
||||
(element_compute == rhs.element_compute) &&
|
||||
(element_scalar == rhs.element_scalar) &&
|
||||
(element_A == rhs.element_A) &&
|
||||
@@ -104,6 +114,28 @@ struct GemmFunctionalKey {
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
inline
|
||||
std::ostream & operator<<(std::ostream &out, cutlass::library::GemmFunctionalKey const &k) {
|
||||
|
||||
out << "{\n"
|
||||
<< " provider: " << to_string(k.provider) << "\n"
|
||||
<< " gemm_kind: " << to_string(k.gemm_kind) << "\n"
|
||||
<< " element_compute: " << to_string(k.element_compute) << "\n"
|
||||
<< " element_scalar: " << to_string(k.element_scalar) << "\n"
|
||||
<< " element_A: " << to_string(k.element_A) << "\n"
|
||||
<< " layout_A: " << to_string(k.layout_A) << "\n"
|
||||
<< " transform_A: " << to_string(k.transform_A) << "\n"
|
||||
<< " element_B: " << to_string(k.element_B) << "\n"
|
||||
<< " layout_B: " << to_string(k.layout_B) << "\n"
|
||||
<< " transform_B: " << to_string(k.transform_B) << "\n"
|
||||
<< " element_C: " << to_string(k.element_C) << "\n"
|
||||
<< "}";
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Hash function for GemmFunctionalKey
|
||||
@@ -120,15 +152,17 @@ struct GemmFunctionalKeyHasher {
|
||||
IntHash hash;
|
||||
|
||||
return
|
||||
rotl(hash(int(key.element_compute)), 2) ^
|
||||
rotl(hash(int(key.element_scalar)), 3) ^
|
||||
rotl(hash(int(key.element_A)), 4) ^
|
||||
rotl(hash(int(key.layout_A)), 5) ^
|
||||
rotl(hash(int(key.transform_A)), 6) ^
|
||||
rotl(hash(int(key.element_B)), 7) ^
|
||||
rotl(hash(int(key.layout_B)), 8) ^
|
||||
rotl(hash(int(key.transform_B)), 9) ^
|
||||
rotl(hash(int(key.element_C)), 10);
|
||||
rotl(hash(int(key.provider)), 1) ^
|
||||
rotl(hash(int(key.gemm_kind)), 2) ^
|
||||
rotl(hash(int(key.element_compute)), 3) ^
|
||||
rotl(hash(int(key.element_scalar)), 4) ^
|
||||
rotl(hash(int(key.element_A)), 5) ^
|
||||
rotl(hash(int(key.layout_A)), 6) ^
|
||||
rotl(hash(int(key.transform_A)), 7) ^
|
||||
rotl(hash(int(key.element_B)), 8) ^
|
||||
rotl(hash(int(key.layout_B)), 9) ^
|
||||
rotl(hash(int(key.transform_B)), 10) ^
|
||||
rotl(hash(int(key.element_C)), 11);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -172,6 +206,7 @@ using GemmOperationFunctionalMap = std::unordered_map<
|
||||
GemmOperationVectorMap,
|
||||
GemmFunctionalKeyHasher
|
||||
>;
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -179,15 +214,10 @@ using GemmOperationFunctionalMap = std::unordered_map<
|
||||
class OperationTable {
|
||||
public:
|
||||
|
||||
/// Map of all operations of type kGemm and gemm_kind of type kGemm
|
||||
/// Map of all operations of type kGemm
|
||||
// provider (kCUTLASS)
|
||||
GemmOperationFunctionalMap gemm_operations;
|
||||
|
||||
/// Map of all operations of type kGemm and gemm_kind of type kPlanarComplex
|
||||
GemmOperationFunctionalMap gemm_planar_complex_operations;
|
||||
|
||||
/// Map of all operations of type kGemm and gemm_kind of type kPlanarComplexArray
|
||||
GemmOperationFunctionalMap gemm_planar_complex_array_operations;
|
||||
|
||||
public:
|
||||
|
||||
void append(Manifest const &manifest);
|
||||
@@ -202,4 +232,3 @@ public:
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
std::ostream & operator<<(std::ostream &out, cutlass::library::GemmFunctionalKey const &k);
|
||||
|
||||
|
||||
@@ -49,6 +49,9 @@ char const *to_string(Provider provider, bool pretty = false);
|
||||
/// Parses a Provider enumerant from a string
|
||||
template <> Provider from_string<Provider>(std::string const &str);
|
||||
|
||||
/// Converts a GemmKind enumerant to a string
|
||||
char const *to_string(GemmKind type, bool pretty = false);
|
||||
|
||||
/// Converts a NumericType enumerant to a string
|
||||
char const *to_string(OperationKind type, bool pretty = false);
|
||||
|
||||
@@ -111,6 +114,14 @@ char const *to_string(ComplexTransform type, bool pretty = false);
|
||||
template <>
|
||||
ComplexTransform from_string<ComplexTransform>(std::string const &str);
|
||||
|
||||
|
||||
/// Converts a SplitKMode enumerant to a string
|
||||
char const *to_string(SplitKMode split_k_mode, bool pretty = false);
|
||||
|
||||
/// Converts a SplitKMode enumerant from a string
|
||||
template <>
|
||||
SplitKMode from_string<SplitKMode>(std::string const &str);
|
||||
|
||||
/// Lexical cast from int64_t to string
|
||||
std::string lexical_cast(int64_t int_value);
|
||||
|
||||
|
||||
@@ -23,7 +23,7 @@ from library import *
|
||||
class GemmOperation:
|
||||
#
|
||||
def __init__(self, gemm_kind, arch, tile_description, A, B, C, element_epilogue, \
|
||||
epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Cohort):
|
||||
epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity8):
|
||||
|
||||
self.operation_kind = OperationKind.Gemm
|
||||
self.arch = arch
|
||||
@@ -40,6 +40,7 @@ class GemmOperation:
|
||||
def is_complex(self):
|
||||
complex_operators = [
|
||||
MathOperation.multiply_add_complex,
|
||||
MathOperation.multiply_add_complex_gaussian
|
||||
]
|
||||
return self.tile_description.math_instruction.math_operation in complex_operators
|
||||
|
||||
@@ -58,6 +59,8 @@ class GemmOperation:
|
||||
|
||||
#
|
||||
def short_math_name(self):
|
||||
if self.tile_description.math_instruction.math_operation == MathOperation.multiply_add_complex_gaussian:
|
||||
return "g%s" % ShortDataTypeNames[self.accumulator_type()]
|
||||
return ShortDataTypeNames[self.accumulator_type()]
|
||||
|
||||
|
||||
@@ -259,6 +262,135 @@ class EmitGemmInstance:
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class EmitGemmUniversalInstance:
|
||||
''' Responsible for emitting a CUTLASS template definition'''
|
||||
|
||||
def __init__(self):
|
||||
self.gemm_template = """
|
||||
// Gemm operator ${operation_name}
|
||||
using ${operation_name}_base =
|
||||
typename cutlass::gemm::kernel::DefaultGemmUniversal<
|
||||
${element_b}, ${layout_b}, ${transform_b}, ${align_b}, // transposed B operand
|
||||
${element_a}, ${layout_a}, ${transform_a}, ${align_a}, // transposed A operand
|
||||
${element_c}, ${layout_c},
|
||||
${element_accumulator},
|
||||
${opcode_class},
|
||||
${arch},
|
||||
cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
|
||||
cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,
|
||||
cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
|
||||
${epilogue_functor}<
|
||||
${element_c},
|
||||
${epilogue_vector_length},
|
||||
${element_accumulator},
|
||||
${element_epilogue}
|
||||
>,
|
||||
${swizzling_functor},
|
||||
${stages},
|
||||
${math_operation}
|
||||
>::GemmKernel;
|
||||
|
||||
// Define named type
|
||||
struct ${operation_name} :
|
||||
public ${operation_name}_base { };
|
||||
"""
|
||||
self.gemm_template_interleaved = """
|
||||
// Gemm operator ${operation_name}
|
||||
using ${operation_name}_base =
|
||||
typename cutlass::gemm::kernel::DefaultGemmUniversal<
|
||||
${element_a}, ${layout_a}, ${transform_a}, ${align_a},
|
||||
${element_b}, ${layout_b}, ${transform_b}, ${align_b},
|
||||
${element_c}, ${layout_c},
|
||||
${element_accumulator},
|
||||
${opcode_class},
|
||||
${arch},
|
||||
cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
|
||||
cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,
|
||||
cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
|
||||
${epilogue_functor}<
|
||||
${element_c},
|
||||
${epilogue_vector_length},
|
||||
${element_accumulator},
|
||||
${element_epilogue}
|
||||
>,
|
||||
${swizzling_functor},
|
||||
${stages},
|
||||
${math_operation}
|
||||
>::GemmKernel;
|
||||
|
||||
// Define named type
|
||||
struct ${operation_name} :
|
||||
public ${operation_name}_base { };
|
||||
"""
|
||||
|
||||
def emit(self, operation):
|
||||
|
||||
threadblock_shape = operation.tile_description.threadblock_shape
|
||||
warp_count = operation.tile_description.warp_count
|
||||
|
||||
warp_shape = [threadblock_shape[idx] // warp_count[idx] for idx in range(3)]
|
||||
warp_shape[2] = operation.tile_description.threadblock_shape[2]
|
||||
|
||||
epilogue_vector_length = int(min(operation.C.alignment * DataTypeSize[operation.C.element], 128) / DataTypeSize[operation.C.element])
|
||||
|
||||
transpose_layouts = {
|
||||
LayoutType.ColumnMajor: LayoutType.RowMajor,
|
||||
LayoutType.RowMajor: LayoutType.ColumnMajor
|
||||
}
|
||||
|
||||
if operation.A.layout in transpose_layouts.keys() and \
|
||||
operation.B.layout in transpose_layouts.keys() and \
|
||||
operation.C.layout in transpose_layouts.keys():
|
||||
|
||||
instance_layout_A = transpose_layouts[operation.A.layout]
|
||||
instance_layout_B = transpose_layouts[operation.B.layout]
|
||||
instance_layout_C = transpose_layouts[operation.C.layout]
|
||||
|
||||
gemm_template = self.gemm_template
|
||||
else:
|
||||
instance_layout_A, instance_layout_B, instance_layout_C = \
|
||||
(operation.A.layout, operation.B.layout, operation.C.layout)
|
||||
|
||||
gemm_template = self.gemm_template_interleaved
|
||||
#
|
||||
|
||||
values = {
|
||||
'operation_name': operation.procedural_name(),
|
||||
'element_a': DataTypeTag[operation.A.element],
|
||||
'layout_a': LayoutTag[instance_layout_A],
|
||||
'element_b': DataTypeTag[operation.B.element],
|
||||
'layout_b': LayoutTag[instance_layout_B],
|
||||
'element_c': DataTypeTag[operation.C.element],
|
||||
'layout_c': LayoutTag[instance_layout_C],
|
||||
'element_accumulator': DataTypeTag[operation.accumulator_type()],
|
||||
'opcode_class': OpcodeClassTag[operation.tile_description.math_instruction.opcode_class],
|
||||
'arch': "cutlass::arch::Sm%d" % operation.arch,
|
||||
'threadblock_shape_m': str(operation.tile_description.threadblock_shape[0]),
|
||||
'threadblock_shape_n': str(operation.tile_description.threadblock_shape[1]),
|
||||
'threadblock_shape_k': str(operation.tile_description.threadblock_shape[2]),
|
||||
'warp_shape_m': str(warp_shape[0]),
|
||||
'warp_shape_n': str(warp_shape[1]),
|
||||
'warp_shape_k': str(warp_shape[2]),
|
||||
'instruction_shape_m': str(operation.tile_description.math_instruction.instruction_shape[0]),
|
||||
'instruction_shape_n': str(operation.tile_description.math_instruction.instruction_shape[1]),
|
||||
'instruction_shape_k': str(operation.tile_description.math_instruction.instruction_shape[2]),
|
||||
'epilogue_vector_length': str(epilogue_vector_length),
|
||||
'element_epilogue': str(DataTypeTag[operation.element_epilogue]),
|
||||
'epilogue_functor': EpilogueFunctorTag[operation.epilogue_functor],
|
||||
'swizzling_functor': SwizzlingFunctorTag[operation.swizzling_functor],
|
||||
'stages': str(operation.tile_description.stages),
|
||||
'align_a': str(operation.A.alignment),
|
||||
'align_b': str(operation.B.alignment),
|
||||
'transform_a': ComplexTransformTag[operation.A.complex_transform],
|
||||
'transform_b': ComplexTransformTag[operation.B.complex_transform],
|
||||
'math_operation': MathOperationTag[operation.tile_description.math_instruction.math_operation]
|
||||
}
|
||||
|
||||
return SubstituteTemplate(gemm_template, values)
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class EmitGemmPlanarComplexInstance:
|
||||
''' Responsible for emitting a CUTLASS template definition'''
|
||||
@@ -282,12 +414,13 @@ class EmitGemmPlanarComplexInstance:
|
||||
${element_accumulator},
|
||||
${element_epilogue}
|
||||
>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
${stages},
|
||||
${math_operator}
|
||||
>::GemmKernel;
|
||||
|
||||
struct ${operation_name} : public Operation_${operation_name} { };
|
||||
struct ${operation_name} :
|
||||
public Operation_${operation_name} { };
|
||||
"""
|
||||
|
||||
def emit(self, operation):
|
||||
@@ -355,7 +488,7 @@ class EmitGemmPlanarComplexArrayInstance:
|
||||
${element_accumulator},
|
||||
${element_epilogue}
|
||||
>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
${stages},
|
||||
${math_operator}
|
||||
>::GemmArrayKernel;
|
||||
@@ -419,12 +552,14 @@ class EmitGemmConfigurationLibrary:
|
||||
|
||||
self.instance_emitter = {
|
||||
GemmKind.Gemm: EmitGemmInstance,
|
||||
GemmKind.Universal: EmitGemmUniversalInstance,
|
||||
GemmKind.PlanarComplex: EmitGemmPlanarComplexInstance,
|
||||
GemmKind.PlanarComplexArray: EmitGemmPlanarComplexArrayInstance
|
||||
}
|
||||
|
||||
self.gemm_kind_wrappers = {
|
||||
GemmKind.Gemm: 'GemmOperation',
|
||||
GemmKind.Universal: 'GemmUniversalOperation',
|
||||
GemmKind.PlanarComplex: 'GemmPlanarComplexOperation',
|
||||
GemmKind.PlanarComplexArray: 'GemmPlanarComplexArrayOperation'
|
||||
}
|
||||
@@ -436,6 +571,13 @@ class EmitGemmConfigurationLibrary:
|
||||
${compile_guard_start}
|
||||
manifest.append(new ${gemm_kind}<Operation_${operation_name}>("${operation_name}"));
|
||||
${compile_guard_end}
|
||||
""",
|
||||
GemmKind.Universal: """
|
||||
${compile_guard_start}
|
||||
manifest.append(new ${gemm_kind}<
|
||||
cutlass::gemm::device::GemmUniversalAdapter<${operation_name}>
|
||||
>("${operation_name}"));
|
||||
${compile_guard_end}
|
||||
""",
|
||||
GemmKind.PlanarComplex: """
|
||||
${compile_guard_start}
|
||||
@@ -542,3 +684,4 @@ void initialize_${configuration_name}(Manifest &manifest) {
|
||||
|
||||
###################################################################################################
|
||||
###################################################################################################
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ from gemm_operation import *
|
||||
def CudaToolkitVersionSatisfies(semantic_ver_string, major, minor, patch = 0):
|
||||
|
||||
# by default, use the latest CUDA Toolkit version
|
||||
cuda_version = [10, 2, 82]
|
||||
cuda_version = [11, 0, 132]
|
||||
|
||||
# Update cuda_version based on parsed string
|
||||
if semantic_ver_string != '':
|
||||
@@ -36,7 +36,7 @@ def CudaToolkitVersionSatisfies(semantic_ver_string, major, minor, patch = 0):
|
||||
#
|
||||
def CreateGemmOperator(manifest, layouts, tile_descriptions, data_type, \
|
||||
alignment_constraints, complex_transforms = None, epilogue_functor = EpilogueFunctor.LinearCombination, \
|
||||
swizzling_functor = SwizzlingFunctor.Cohort):
|
||||
swizzling_functor = SwizzlingFunctor.Identity8):
|
||||
|
||||
if complex_transforms is None:
|
||||
complex_transforms = [(ComplexTransform.none, ComplexTransform.none),]
|
||||
@@ -61,7 +61,7 @@ def CreateGemmOperator(manifest, layouts, tile_descriptions, data_type, \
|
||||
B = TensorDescription(element_b, layout[1], alignment, complex_transform[1])
|
||||
C = TensorDescription(element_c, layout[2], alignment_c)
|
||||
|
||||
new_operation = GemmOperation(GemmKind.Gemm, tile_description.minimum_compute_capability, \
|
||||
new_operation = GemmOperation(GemmKind.Universal, tile_description.minimum_compute_capability, \
|
||||
tile_description, A, B, C, element_epilogue, epilogue_functor, swizzling_functor)
|
||||
|
||||
manifest.append(new_operation)
|
||||
@@ -466,6 +466,9 @@ def GenerateSM70_WmmaTensorOp_161616(manifest, args):
|
||||
def GenerateSM70(manifest, args):
|
||||
GenerateSM70_TensorOp_884(manifest, args)
|
||||
GenerateSM70_PlanarComplexTensorOp_884(manifest, args)
|
||||
|
||||
# To limit build size, WMMA GEMMs are disabled for now.
|
||||
#
|
||||
#GenerateSM70_WmmaTensorOp_161616(manifest, args)
|
||||
|
||||
###################################################################################################
|
||||
@@ -621,6 +624,11 @@ def GenerateSM75_TensorOp_8816_TN(manifest, args):
|
||||
DataType.s8, DataType.s8, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
MathInstruction( \
|
||||
[8, 8, 16], \
|
||||
DataType.u8, DataType.u8, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
]
|
||||
|
||||
min_cc = 75
|
||||
@@ -654,7 +662,7 @@ def GenerateSM75_TensorOp_8816_TN(manifest, args):
|
||||
data_type_mixed = [
|
||||
math_inst.element_a,
|
||||
math_inst.element_b,
|
||||
math_inst.element_a,
|
||||
DataType.s8,
|
||||
DataType.f32,
|
||||
]
|
||||
|
||||
@@ -687,6 +695,11 @@ def GenerateSM75_TensorOp_8816_Interleaved(manifest, args):
|
||||
DataType.s8, DataType.s8, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
MathInstruction( \
|
||||
[8, 8, 16], \
|
||||
DataType.u8, DataType.u8, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
]
|
||||
|
||||
min_cc = 75
|
||||
@@ -712,8 +725,7 @@ def GenerateSM75_TensorOp_8816_Interleaved(manifest, args):
|
||||
]
|
||||
|
||||
operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp, \
|
||||
SwizzlingFunctor.Identity)
|
||||
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
for op in operations:
|
||||
op.C.alignment = 8
|
||||
@@ -736,6 +748,11 @@ def GenerateSM75_TensorOp_8832_TN(manifest, args):
|
||||
DataType.s4, DataType.s4, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
MathInstruction( \
|
||||
[8, 8, 32], \
|
||||
DataType.u4, DataType.u4, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
]
|
||||
|
||||
min_cc = 75
|
||||
@@ -769,7 +786,7 @@ def GenerateSM75_TensorOp_8832_TN(manifest, args):
|
||||
data_type_mixed = [
|
||||
math_inst.element_a,
|
||||
math_inst.element_b,
|
||||
math_inst.element_a,
|
||||
DataType.s4,
|
||||
DataType.f32,
|
||||
]
|
||||
|
||||
@@ -804,6 +821,11 @@ def GenerateSM75_TensorOp_8832_Interleaved(manifest, args):
|
||||
DataType.s4, DataType.s4, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
MathInstruction( \
|
||||
[8, 8, 32], \
|
||||
DataType.u4, DataType.u4, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
]
|
||||
|
||||
min_cc = 75
|
||||
@@ -832,8 +854,7 @@ def GenerateSM75_TensorOp_8832_Interleaved(manifest, args):
|
||||
]
|
||||
|
||||
operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp, \
|
||||
SwizzlingFunctor.Identity)
|
||||
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
for op in operations:
|
||||
op.C.alignment = 16
|
||||
@@ -911,6 +932,831 @@ def GenerateSM75(manifest, args):
|
||||
###################################################################################################
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_16816(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction( \
|
||||
[16, 8, 16], \
|
||||
DataType.f16, DataType.f16, DataType.f32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add),
|
||||
MathInstruction( \
|
||||
[16, 8, 16], \
|
||||
DataType.f16, DataType.f16, DataType.f16, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add),
|
||||
MathInstruction( \
|
||||
[16, 8, 16], \
|
||||
DataType.bf16, DataType.bf16, DataType.f32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add),
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [8, 4, 2]
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 64], 3, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 64], 3, [2, 1, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 64], 4, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 64], 4, [2, 1, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 32], 10, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 64], 4, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 64], 5, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 256, 64], 3, [1, 4, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [
|
||||
math_inst.element_a,
|
||||
math_inst.element_b,
|
||||
math_inst.element_accumulator,
|
||||
math_inst.element_accumulator,
|
||||
]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints)
|
||||
|
||||
# Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation)
|
||||
if math_inst.element_a != math_inst.element_accumulator:
|
||||
|
||||
data_type_mixed = [
|
||||
math_inst.element_a,
|
||||
math_inst.element_b,
|
||||
math_inst.element_a,
|
||||
math_inst.element_accumulator,
|
||||
]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type_mixed, alignment_constraints)
|
||||
|
||||
#
|
||||
|
||||
#
|
||||
def GenerateSM80_PlanarComplexTensorOp_16816(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
complex_transforms = [
|
||||
(ComplexTransform.none, ComplexTransform.none),
|
||||
(ComplexTransform.conj, ComplexTransform.none),
|
||||
(ComplexTransform.none, ComplexTransform.conj),
|
||||
(ComplexTransform.conj, ComplexTransform.conj)
|
||||
]
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction( \
|
||||
[16, 8, 16], \
|
||||
DataType.f16, DataType.f16, DataType.f32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add),
|
||||
MathInstruction( \
|
||||
[16, 8, 16], \
|
||||
DataType.bf16, DataType.bf16, DataType.f32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add),
|
||||
MathInstruction( \
|
||||
[16, 8, 16], \
|
||||
DataType.f16, DataType.f16, DataType.f16, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add),
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [8, ]
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([ 64, 128, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [
|
||||
math_inst.element_a,
|
||||
math_inst.element_b,
|
||||
math_inst.element_accumulator,
|
||||
math_inst.element_accumulator,
|
||||
]
|
||||
|
||||
CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints, complex_transforms)
|
||||
|
||||
# Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation)
|
||||
if math_inst.element_a != math_inst.element_accumulator:
|
||||
|
||||
data_type_mixed = [
|
||||
math_inst.element_a,
|
||||
math_inst.element_b,
|
||||
math_inst.element_a,
|
||||
math_inst.element_accumulator,
|
||||
]
|
||||
|
||||
CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type_mixed, alignment_constraints, complex_transforms)
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_16832_TN(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction( \
|
||||
[16, 8, 32], \
|
||||
DataType.s8, DataType.s8, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
MathInstruction( \
|
||||
[16, 8, 32], \
|
||||
DataType.u8, DataType.u8, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [16,]
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 64, 128], 3, [4, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 256, 128], 3, [1, 4, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [math_inst.element_a, math_inst.element_b, DataType.s32, DataType.s32]
|
||||
data_type_mixed = [math_inst.element_a, math_inst.element_b, DataType.s8, DataType.f32]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
operations = []
|
||||
|
||||
operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
for op in operations:
|
||||
if op.tile_description.threadblock_shape[1] >= 128:
|
||||
op.C.alignment = 16
|
||||
else:
|
||||
op.C.alignment = 8
|
||||
|
||||
#
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_16832_Interleaved(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajorInterleaved32, LayoutType.RowMajorInterleaved32, LayoutType.ColumnMajorInterleaved32),
|
||||
]
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction( \
|
||||
[16, 8, 32], \
|
||||
DataType.s8, DataType.s8, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
MathInstruction( \
|
||||
[16, 8, 32], \
|
||||
DataType.u8, DataType.u8, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [16,]
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type_mixed = [math_inst.element_a, math_inst.element_b, DataType.s8, DataType.f32]
|
||||
|
||||
operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
for op in operations:
|
||||
op.C.alignment = 8
|
||||
|
||||
#
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_16864_TN(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction( \
|
||||
[16, 8, 64], \
|
||||
DataType.s4, DataType.s4, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
MathInstruction( \
|
||||
[16, 8, 64], \
|
||||
DataType.u4, DataType.u4, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [32,]
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 128, 256], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 256], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 256], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [math_inst.element_a, math_inst.element_b, DataType.s32, DataType.s32]
|
||||
data_type_mixed = [math_inst.element_a, math_inst.element_b, DataType.s4, DataType.f32]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
operations = []
|
||||
|
||||
operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
for op in operations:
|
||||
if op.tile_description.threadblock_shape[1] >= 128:
|
||||
op.C.alignment = 8
|
||||
elif op.tile_description.threadblock_shape[1] == 64:
|
||||
op.C.alignment = 8
|
||||
else:
|
||||
op.C.alignment = 4
|
||||
#
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_16864_Interleaved(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajorInterleaved64, LayoutType.RowMajorInterleaved64, LayoutType.ColumnMajorInterleaved64),
|
||||
]
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction( \
|
||||
[16, 8, 64], \
|
||||
DataType.s4, DataType.s4, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
MathInstruction( \
|
||||
[16, 8, 64], \
|
||||
DataType.u4, DataType.u4, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_saturate),
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [32,]
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type_mixed = [math_inst.element_a, math_inst.element_b, DataType.s4, DataType.f32]
|
||||
|
||||
operations = []
|
||||
|
||||
operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
for op in operations:
|
||||
op.C.alignment = 16
|
||||
#
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_168256(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction( \
|
||||
[16, 8, 256], \
|
||||
DataType.b1, DataType.b1, DataType.s32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.xor_popc),
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [128,]
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([256, 128, 512], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 512], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 512], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 512], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 512], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 512], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 128, 1024], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 1024], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 1024], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 1024], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 1024], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 1024], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [DataType.b1, DataType.b1, DataType.s32, DataType.s32]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp)
|
||||
|
||||
#
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_1688(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction( \
|
||||
[16, 8, 8], \
|
||||
DataType.tf32, DataType.tf32, DataType.f32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add)
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [4, 2, 1]
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 32], 3, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 32], 3, [2, 1, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 32], 4, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 32], 4, [2, 1, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 32], 4, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 32], 5, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 256, 32], 3, [1, 4, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [
|
||||
math_inst.element_a,
|
||||
math_inst.element_b,
|
||||
math_inst.element_accumulator,
|
||||
math_inst.element_accumulator,
|
||||
]
|
||||
|
||||
data_type_mixed = [
|
||||
math_inst.element_a,
|
||||
math_inst.element_b,
|
||||
math_inst.element_a,
|
||||
math_inst.element_accumulator,
|
||||
]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints)
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type_mixed, alignment_constraints)
|
||||
|
||||
#
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_1688_fast_math(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction( \
|
||||
[16, 8, 8], \
|
||||
DataType.tf32, DataType.tf32, DataType.f32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add),
|
||||
MathInstruction( \
|
||||
[16, 8, 8], \
|
||||
DataType.f16, DataType.f16, DataType.f32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_fast_f16),
|
||||
MathInstruction( \
|
||||
[16, 8, 8], \
|
||||
DataType.bf16, DataType.bf16, DataType.f32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_fast_bf16)
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [4, 2, 1]
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 32], 3, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 32], 3, [2, 1, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 32], 4, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 32], 4, [2, 1, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 32], 4, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 32], 5, [1, 2, 2], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 256, 32], 3, [1, 4, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [DataType.f32, DataType.f32, DataType.f32, DataType.f32]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints)
|
||||
|
||||
#
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_1688_complex(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_inst = MathInstruction( \
|
||||
[16, 8, 8], \
|
||||
DataType.f32, DataType.f32, DataType.f32, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_complex)
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
tile_descriptions = [
|
||||
TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, 16], 4, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 32, 16], 4, [2, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [
|
||||
DataType.cf32, DataType.cf32, DataType.cf32, DataType.cf32
|
||||
]
|
||||
|
||||
alignment_constraints = [1,]
|
||||
|
||||
complex_transforms = [
|
||||
(ComplexTransform.none, ComplexTransform.none),
|
||||
(ComplexTransform.conj, ComplexTransform.none),
|
||||
(ComplexTransform.none, ComplexTransform.conj),
|
||||
(ComplexTransform.conj, ComplexTransform.conj)
|
||||
]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints, complex_transforms)
|
||||
#
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_884(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_inst = \
|
||||
MathInstruction( \
|
||||
[8, 8, 4], \
|
||||
DataType.f64, DataType.f64, DataType.f64, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add)
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [1,]
|
||||
|
||||
tile_descriptions = [
|
||||
TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 32, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([16, 32, 16], 5, [1, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 16, 16], 5, [2, 1, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [DataType.f64, DataType.f64, DataType.f64, DataType.f64]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints)
|
||||
#
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_884_complex(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_inst = \
|
||||
MathInstruction( \
|
||||
[8, 8, 4], \
|
||||
DataType.f64, DataType.f64, DataType.f64, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_complex)
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [1,]
|
||||
|
||||
tile_descriptions = [
|
||||
TileDescription([128, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 128, 8], 3, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 64, 8], 3, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64]
|
||||
|
||||
complex_transforms = [
|
||||
(ComplexTransform.none, ComplexTransform.none),
|
||||
(ComplexTransform.conj, ComplexTransform.none),
|
||||
(ComplexTransform.none, ComplexTransform.conj),
|
||||
(ComplexTransform.conj, ComplexTransform.conj)
|
||||
]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints, complex_transforms)
|
||||
|
||||
#
|
||||
def GenerateSM80_TensorOp_884_complex_gaussian(manifest, args):
|
||||
|
||||
if not CudaToolkitVersionSatisfies(args.cuda_version, 11, 0):
|
||||
return
|
||||
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_inst = \
|
||||
MathInstruction( \
|
||||
[8, 8, 4], \
|
||||
DataType.f64, DataType.f64, DataType.f64, \
|
||||
OpcodeClass.TensorOp, \
|
||||
MathOperation.multiply_add_complex_gaussian)
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [1,]
|
||||
|
||||
tile_descriptions = [
|
||||
TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64]
|
||||
|
||||
complex_transforms = [
|
||||
(ComplexTransform.none, ComplexTransform.none),
|
||||
(ComplexTransform.conj, ComplexTransform.none),
|
||||
(ComplexTransform.none, ComplexTransform.conj),
|
||||
(ComplexTransform.conj, ComplexTransform.conj)
|
||||
]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints, complex_transforms)
|
||||
#
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def GenerateSM80_Simt(manifest, args):
|
||||
layouts = [
|
||||
(LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor),
|
||||
(LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor),
|
||||
]
|
||||
|
||||
math_instructions = [
|
||||
MathInstruction( \
|
||||
[1, 1, 1], \
|
||||
DataType.f32, DataType.f32, DataType.f32, \
|
||||
OpcodeClass.Simt, \
|
||||
MathOperation.multiply_add),
|
||||
]
|
||||
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
|
||||
alignment_constraints = [1,]
|
||||
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = [
|
||||
TileDescription([256, 128, 8], 5, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 8], 5, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 8], 5, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([256, 128, 8], 4, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 256, 8], 4, [2, 4, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 128, 8], 4, [4, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 64, 8], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 128, 8], 5, [2, 2, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 64, 64, 8], 5, [2, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([128, 32, 8], 5, [2, 1, 1], math_inst, min_cc, max_cc),
|
||||
TileDescription([ 32, 128, 8], 5, [1, 2, 1], math_inst, min_cc, max_cc),
|
||||
]
|
||||
|
||||
data_type = [
|
||||
math_inst.element_a,
|
||||
math_inst.element_b,
|
||||
math_inst.element_accumulator,
|
||||
math_inst.element_accumulator,
|
||||
]
|
||||
|
||||
CreateGemmOperator(manifest, layouts, tile_descriptions, \
|
||||
data_type, alignment_constraints)
|
||||
#
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
def GenerateSM80(manifest, args):
|
||||
|
||||
GenerateSM80_TensorOp_16816(manifest, args)
|
||||
GenerateSM80_PlanarComplexTensorOp_16816(manifest, args)
|
||||
GenerateSM80_TensorOp_1688(manifest, args)
|
||||
GenerateSM80_TensorOp_1688_fast_math(manifest, args)
|
||||
GenerateSM80_TensorOp_1688_complex(manifest, args)
|
||||
GenerateSM80_TensorOp_884(manifest, args)
|
||||
GenerateSM80_TensorOp_884_complex(manifest, args)
|
||||
GenerateSM80_TensorOp_884_complex_gaussian(manifest, args)
|
||||
GenerateSM80_TensorOp_16832_TN(manifest, args)
|
||||
GenerateSM80_TensorOp_16832_Interleaved(manifest, args)
|
||||
GenerateSM80_TensorOp_16864_TN(manifest, args)
|
||||
GenerateSM80_TensorOp_16864_Interleaved(manifest, args)
|
||||
GenerateSM80_TensorOp_168256(manifest, args)
|
||||
GenerateSM80_Simt(manifest, args)
|
||||
#
|
||||
|
||||
###################################################################################################
|
||||
|
||||
if __name__ == "__main__":
|
||||
@@ -920,7 +1766,7 @@ if __name__ == "__main__":
|
||||
parser.add_argument("--build-dir", default=".", required=False, help="CUTLASS top-level build directory")
|
||||
parser.add_argument("--curr-build-dir", default=".", help="CUTLASS current build directory. cmake files will be emitted in this directory")
|
||||
parser.add_argument("--generator-target", default='library', help="Target of CUTLASS Library Generator.")
|
||||
parser.add_argument("--architectures", default='50;60;61;75', help="Target compute architectures")
|
||||
parser.add_argument("--architectures", default='53;60;61;70;75;80', help="Target compute architectures")
|
||||
parser.add_argument("--kernels", default='', help='Comma delimited list to filter kernels by name.')
|
||||
parser.add_argument("--cuda-version", default="11.0.0", help="Semantic version string of CUDA Toolkit")
|
||||
|
||||
@@ -933,6 +1779,8 @@ if __name__ == "__main__":
|
||||
GenerateSM61(manifest, args)
|
||||
GenerateSM70(manifest, args)
|
||||
GenerateSM75(manifest, args)
|
||||
GenerateSM80(manifest, args)
|
||||
|
||||
if 'library' in args.generator_target.split(','):
|
||||
manifest.emit(GeneratorTarget.Library)
|
||||
|
||||
|
||||
@@ -4,14 +4,32 @@
|
||||
# \brief Generates the CUTLASS Library's instances
|
||||
#
|
||||
|
||||
import enum
|
||||
import re
|
||||
|
||||
###################################################################################################
|
||||
|
||||
import enum
|
||||
|
||||
# The following block implements enum.auto() for Python 3.5 variants that don't include it such
|
||||
# as the default 3.5.2 on Ubuntu 16.04.
|
||||
#
|
||||
# https://codereview.stackexchange.com/questions/177309/reimplementing-pythons-enum-auto-for-compatibility
|
||||
|
||||
try:
|
||||
from enum import auto as enum_auto
|
||||
except ImportError:
|
||||
__cutlass_library_auto_enum = 0
|
||||
def enum_auto() -> int:
|
||||
global __cutlass_library_auto_enum
|
||||
i = __cutlass_library_auto_enum
|
||||
__cutlass_library_auto_enum += 1
|
||||
return i
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class GeneratorTarget(enum.Enum):
|
||||
Library = enum.auto()
|
||||
Library = enum_auto()
|
||||
#
|
||||
GeneratorTargetNames = {
|
||||
GeneratorTarget.Library: 'library'
|
||||
@@ -22,33 +40,37 @@ GeneratorTargetNames = {
|
||||
|
||||
#
|
||||
class DataType(enum.Enum):
|
||||
b1 = enum.auto()
|
||||
u4 = enum.auto()
|
||||
u8 = enum.auto()
|
||||
u16 = enum.auto()
|
||||
u32 = enum.auto()
|
||||
u64 = enum.auto()
|
||||
s4 = enum.auto()
|
||||
s8 = enum.auto()
|
||||
s16 = enum.auto()
|
||||
s32 = enum.auto()
|
||||
s64 = enum.auto()
|
||||
f16 = enum.auto()
|
||||
f32 = enum.auto()
|
||||
f64 = enum.auto()
|
||||
cf16 = enum.auto()
|
||||
cf32 = enum.auto()
|
||||
cf64 = enum.auto()
|
||||
cs4 = enum.auto()
|
||||
cs8 = enum.auto()
|
||||
cs16 = enum.auto()
|
||||
cs32 = enum.auto()
|
||||
cs64 = enum.auto()
|
||||
cu4 = enum.auto()
|
||||
cu8 = enum.auto()
|
||||
cu16 = enum.auto()
|
||||
cu32 = enum.auto()
|
||||
cu64 = enum.auto()
|
||||
b1 = enum_auto()
|
||||
u4 = enum_auto()
|
||||
u8 = enum_auto()
|
||||
u16 = enum_auto()
|
||||
u32 = enum_auto()
|
||||
u64 = enum_auto()
|
||||
s4 = enum_auto()
|
||||
s8 = enum_auto()
|
||||
s16 = enum_auto()
|
||||
s32 = enum_auto()
|
||||
s64 = enum_auto()
|
||||
f16 = enum_auto()
|
||||
bf16 = enum_auto()
|
||||
f32 = enum_auto()
|
||||
tf32 = enum_auto()
|
||||
f64 = enum_auto()
|
||||
cf16 = enum_auto()
|
||||
cbf16 = enum_auto()
|
||||
cf32 = enum_auto()
|
||||
ctf32 = enum_auto()
|
||||
cf64 = enum_auto()
|
||||
cs4 = enum_auto()
|
||||
cs8 = enum_auto()
|
||||
cs16 = enum_auto()
|
||||
cs32 = enum_auto()
|
||||
cs64 = enum_auto()
|
||||
cu4 = enum_auto()
|
||||
cu8 = enum_auto()
|
||||
cu16 = enum_auto()
|
||||
cu32 = enum_auto()
|
||||
cu64 = enum_auto()
|
||||
|
||||
#
|
||||
ShortDataTypeNames = {
|
||||
@@ -74,10 +96,14 @@ DataTypeNames = {
|
||||
DataType.s32: "s32",
|
||||
DataType.s64: "s64",
|
||||
DataType.f16: "f16",
|
||||
DataType.bf16: "bf16",
|
||||
DataType.f32: "f32",
|
||||
DataType.tf32: "tf32",
|
||||
DataType.f64: "f64",
|
||||
DataType.cf16: "cf16",
|
||||
DataType.cbf16: "cbf16",
|
||||
DataType.cf32: "cf32",
|
||||
DataType.ctf32: "ctf32",
|
||||
DataType.cf64: "cf64",
|
||||
DataType.cu4: "cu4",
|
||||
DataType.cu8: "cu8",
|
||||
@@ -104,10 +130,14 @@ DataTypeTag = {
|
||||
DataType.s32: "int32_t",
|
||||
DataType.s64: "int64_t",
|
||||
DataType.f16: "cutlass::half_t",
|
||||
DataType.bf16: "cutlass::bfloat16_t",
|
||||
DataType.f32: "float",
|
||||
DataType.tf32: "cutlass::tfloat32_t",
|
||||
DataType.f64: "double",
|
||||
DataType.cf16: "cutlass::complex<cutlass::half_t>",
|
||||
DataType.cbf16: "cutlass::complex<cutlass::bfloat16_t>",
|
||||
DataType.cf32: "cutlass::complex<float>",
|
||||
DataType.ctf32: "cutlass::complex<cutlass::tfloat32_t>",
|
||||
DataType.cf64: "cutlass::complex<double>",
|
||||
DataType.cu4: "cutlass::complex<cutlass::uint4b_t>",
|
||||
DataType.cu8: "cutlass::complex<cutlass::uint8_t>",
|
||||
@@ -134,10 +164,14 @@ DataTypeSize = {
|
||||
DataType.s32: 32,
|
||||
DataType.s64: 64,
|
||||
DataType.f16: 16,
|
||||
DataType.bf16: 16,
|
||||
DataType.f32: 32,
|
||||
DataType.tf32: 32,
|
||||
DataType.f64: 64,
|
||||
DataType.cf16: 32,
|
||||
DataType.cbf16: 32,
|
||||
DataType.cf32: 64,
|
||||
DataType.ctf32: 32,
|
||||
DataType.cf64: 128,
|
||||
DataType.cu4: 8,
|
||||
DataType.cu8: 16,
|
||||
@@ -155,8 +189,8 @@ DataTypeSize = {
|
||||
|
||||
#
|
||||
class ComplexTransform(enum.Enum):
|
||||
none = enum.auto()
|
||||
conj = enum.auto()
|
||||
none = enum_auto()
|
||||
conj = enum_auto()
|
||||
|
||||
#
|
||||
ComplexTransformTag = {
|
||||
@@ -194,40 +228,47 @@ def get_real_from_complex(complex_type):
|
||||
|
||||
#
|
||||
class ComplexMultiplyOp(enum.Enum):
|
||||
multiply_add = enum.auto()
|
||||
gaussian = enum.auto()
|
||||
multiply_add = enum_auto()
|
||||
gaussian = enum_auto()
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class MathOperation(enum.Enum):
|
||||
multiply_add = enum.auto()
|
||||
multiply_add_saturate = enum.auto()
|
||||
xor_popc = enum.auto()
|
||||
multiply_add_complex = enum.auto()
|
||||
multiply_add = enum_auto()
|
||||
multiply_add_saturate = enum_auto()
|
||||
xor_popc = enum_auto()
|
||||
multiply_add_fast_bf16 = enum_auto()
|
||||
multiply_add_fast_f16 = enum_auto()
|
||||
multiply_add_complex = enum_auto()
|
||||
multiply_add_complex_gaussian = enum_auto()
|
||||
|
||||
#
|
||||
MathOperationTag = {
|
||||
MathOperation.multiply_add: 'cutlass::arch::OpMultiplyAdd',
|
||||
MathOperation.multiply_add_saturate: 'cutlass::arch::OpMultiplyAddSaturate',
|
||||
MathOperation.xor_popc: 'cutlass::arch::OpXorPopc',
|
||||
MathOperation.multiply_add_fast_bf16: 'cutlass::arch::OpMultiplyAddFastBF16',
|
||||
MathOperation.multiply_add_fast_f16: 'cutlass::arch::OpMultiplyAddFastF16',
|
||||
MathOperation.multiply_add_complex: 'cutlass::arch::OpMultiplyAddComplex',
|
||||
MathOperation.multiply_add_complex_gaussian: 'cutlass::arch::OpMultiplyAddGaussianComplex',
|
||||
}
|
||||
|
||||
###################################################################################################
|
||||
|
||||
#
|
||||
class LayoutType(enum.Enum):
|
||||
ColumnMajor = enum.auto()
|
||||
RowMajor = enum.auto()
|
||||
ColumnMajorInterleaved32 = enum.auto()
|
||||
RowMajorInterleaved32 = enum.auto()
|
||||
ColumnMajorInterleaved64 = enum.auto()
|
||||
RowMajorInterleaved64 = enum.auto()
|
||||
TensorNHWC = enum.auto()
|
||||
TensorNCHW = enum.auto()
|
||||
TensorNGHWC = enum.auto()
|
||||
TensorNCxHW32 = enum.auto()
|
||||
TensorNCxHW64 = enum.auto()
|
||||
ColumnMajor = enum_auto()
|
||||
RowMajor = enum_auto()
|
||||
ColumnMajorInterleaved32 = enum_auto()
|
||||
RowMajorInterleaved32 = enum_auto()
|
||||
ColumnMajorInterleaved64 = enum_auto()
|
||||
RowMajorInterleaved64 = enum_auto()
|
||||
TensorNHWC = enum_auto()
|
||||
TensorNCHW = enum_auto()
|
||||
TensorNGHWC = enum_auto()
|
||||
TensorNCxHW32 = enum_auto()
|
||||
TensorNCxHW64 = enum_auto()
|
||||
|
||||
#
|
||||
LayoutTag = {
|
||||
@@ -282,9 +323,9 @@ ShortComplexLayoutNames = {
|
||||
|
||||
#
|
||||
class OpcodeClass(enum.Enum):
|
||||
Simt = enum.auto()
|
||||
TensorOp = enum.auto()
|
||||
WmmaTensorOp = enum.auto()
|
||||
Simt = enum_auto()
|
||||
TensorOp = enum_auto()
|
||||
WmmaTensorOp = enum_auto()
|
||||
|
||||
OpcodeClassNames = {
|
||||
OpcodeClass.Simt: 'simt',
|
||||
@@ -302,7 +343,7 @@ OpcodeClassTag = {
|
||||
|
||||
#
|
||||
class OperationKind(enum.Enum):
|
||||
Gemm = enum.auto()
|
||||
Gemm = enum_auto()
|
||||
#
|
||||
OperationKindNames = {
|
||||
OperationKind.Gemm: 'gemm'
|
||||
@@ -310,7 +351,7 @@ OperationKindNames = {
|
||||
|
||||
#
|
||||
class Target(enum.Enum):
|
||||
library = enum.auto()
|
||||
library = enum_auto()
|
||||
|
||||
ArchitectureNames = {
|
||||
50: 'maxwell',
|
||||
@@ -318,6 +359,7 @@ ArchitectureNames = {
|
||||
61: 'pascal',
|
||||
70: 'volta',
|
||||
75: 'turing',
|
||||
80: 'ampere',
|
||||
}
|
||||
|
||||
###################################################################################################
|
||||
@@ -340,27 +382,27 @@ def SubstituteTemplate(template, values):
|
||||
|
||||
#
|
||||
class GemmKind(enum.Enum):
|
||||
Gemm = enum.auto()
|
||||
Batched = enum.auto()
|
||||
Array = enum.auto()
|
||||
Universal = enum.auto()
|
||||
PlanarComplex = enum.auto()
|
||||
PlanarComplexArray = enum.auto()
|
||||
Gemm = enum_auto()
|
||||
Batched = enum_auto()
|
||||
Array = enum_auto()
|
||||
Universal = enum_auto()
|
||||
PlanarComplex = enum_auto()
|
||||
PlanarComplexArray = enum_auto()
|
||||
|
||||
#
|
||||
GemmKindNames = {
|
||||
GemmKind.Gemm: "gemm",
|
||||
GemmKind.Batched: "gemm_batched",
|
||||
GemmKind.Array: "gemm_array",
|
||||
GemmKind.Universal: "gemm_universal",
|
||||
GemmKind.Universal: "gemm",
|
||||
GemmKind.PlanarComplex: "gemm_planar_complex",
|
||||
GemmKind.PlanarComplexArray: "gemm_planar_complex_array",
|
||||
}
|
||||
|
||||
#
|
||||
class EpilogueFunctor(enum.Enum):
|
||||
LinearCombination = enum.auto()
|
||||
LinearCombinationClamp = enum.auto()
|
||||
LinearCombination = enum_auto()
|
||||
LinearCombinationClamp = enum_auto()
|
||||
|
||||
#
|
||||
EpilogueFunctorTag = {
|
||||
@@ -370,13 +412,17 @@ EpilogueFunctorTag = {
|
||||
|
||||
#
|
||||
class SwizzlingFunctor(enum.Enum):
|
||||
Cohort = enum.auto()
|
||||
Identity = enum.auto()
|
||||
Identity1 = enum_auto()
|
||||
Identity2 = enum_auto()
|
||||
Identity4 = enum_auto()
|
||||
Identity8 = enum_auto()
|
||||
|
||||
#
|
||||
SwizzlingFunctorTag = {
|
||||
SwizzlingFunctor.Cohort: 'cutlass::gemm::threadblock::GemmCohortThreadblockSwizzle<${layout_a}, ${layout_b}>',
|
||||
SwizzlingFunctor.Identity: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle',
|
||||
SwizzlingFunctor.Identity1: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>',
|
||||
SwizzlingFunctor.Identity2: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<2>',
|
||||
SwizzlingFunctor.Identity4: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<4>',
|
||||
SwizzlingFunctor.Identity8: 'cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>',
|
||||
}
|
||||
###################################################################################################
|
||||
|
||||
|
||||
@@ -127,7 +127,7 @@ class Manifest:
|
||||
if args.kernels == 'all':
|
||||
self.kernel_names = []
|
||||
else:
|
||||
self.kernel_names = args.kernels.split(',')
|
||||
self.kernel_names = [x for x in args.kernels.split(',') if x != '']
|
||||
|
||||
self.operation_count = 0
|
||||
self.operations_by_name = {}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
|
||||
* Copyright (c) 2017-2020, 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:
|
||||
@@ -29,13 +29,14 @@
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/gemm/kernel/default_gemm_planar_complex_universal.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm.h"
|
||||
#include "cutlass/gemm/device/gemm_complex.h"
|
||||
#include "cutlass/gemm/device/gemm_batched.h"
|
||||
#include "cutlass/gemm/device/gemm_array.h"
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/gemm/kernel/default_gemm_universal.h"
|
||||
#include "cutlass/gemm/kernel/default_gemm_planar_complex_universal.h"
|
||||
|
||||
#include "cutlass/library/library.h"
|
||||
#include "library_internal.h"
|
||||
@@ -104,10 +105,10 @@ public:
|
||||
MathOperationMap<typename Operator::Operator>::kId;
|
||||
|
||||
description_.tile_description.minimum_compute_capability =
|
||||
ArchMap<typename Operator::ArchTag>::kMin;
|
||||
ArchMap<typename Operator::ArchTag, typename Operator::OperatorClass>::kMin;
|
||||
|
||||
description_.tile_description.maximum_compute_capability =
|
||||
ArchMap<typename Operator::ArchTag>::kMax;
|
||||
ArchMap<typename Operator::ArchTag, typename Operator::OperatorClass>::kMax;
|
||||
|
||||
description_.A = make_TensorDescription<ElementA, LayoutA>(Operator::kAlignmentA);
|
||||
description_.B = make_TensorDescription<ElementB, LayoutB>(Operator::kAlignmentB);
|
||||
@@ -698,6 +699,201 @@ public:
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename Operator_>
|
||||
class GemmUniversalOperation : public GemmOperationBase<Operator_> {
|
||||
public:
|
||||
|
||||
using Operator = Operator_;
|
||||
using ElementA = typename Operator::ElementA;
|
||||
using LayoutA = typename Operator::LayoutA;
|
||||
using ElementB = typename Operator::ElementB;
|
||||
using LayoutB = typename Operator::LayoutB;
|
||||
using ElementC = typename Operator::ElementC;
|
||||
using LayoutC = typename Operator::LayoutC;
|
||||
using ElementAccumulator = typename Operator::ElementAccumulator;
|
||||
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
|
||||
|
||||
using OperatorArguments = typename Operator::Arguments;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructor
|
||||
GemmUniversalOperation(char const *name = "unknown_gemm"):
|
||||
GemmOperationBase<Operator_>(name) {
|
||||
|
||||
this->description_.gemm_kind = GemmKind::kUniversal;
|
||||
}
|
||||
|
||||
protected:
|
||||
|
||||
/// Constructs the arguments structure given the configuration and arguments
|
||||
static Status construct_arguments_(
|
||||
OperatorArguments &operator_args,
|
||||
GemmUniversalConfiguration const *configuration) {
|
||||
|
||||
operator_args.mode = configuration->mode;
|
||||
|
||||
operator_args.problem_size = configuration->problem_size;
|
||||
operator_args.batch_count = configuration->batch_count;
|
||||
|
||||
operator_args.lda = int(configuration->lda);
|
||||
operator_args.ldb = int(configuration->ldb);
|
||||
operator_args.ldc = int(configuration->ldc);
|
||||
operator_args.ldd = int(configuration->ldd);
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
/// Constructs the arguments structure given the configuration and arguments
|
||||
static Status update_arguments_(
|
||||
OperatorArguments &operator_args,
|
||||
GemmUniversalArguments const *arguments) {
|
||||
|
||||
if (arguments->pointer_mode == ScalarPointerMode::kHost) {
|
||||
typename Operator::EpilogueOutputOp::Params params(
|
||||
*static_cast<ElementCompute const *>(arguments->alpha),
|
||||
*static_cast<ElementCompute const *>(arguments->beta)
|
||||
);
|
||||
operator_args.epilogue = params;
|
||||
}
|
||||
else if (arguments->pointer_mode == ScalarPointerMode::kDevice){
|
||||
typename Operator::EpilogueOutputOp::Params params(
|
||||
static_cast<ElementCompute const *>(arguments->alpha),
|
||||
static_cast<ElementCompute const *>(arguments->beta)
|
||||
);
|
||||
operator_args.epilogue = params;
|
||||
}
|
||||
else {
|
||||
return Status::kErrorInvalidProblem;
|
||||
}
|
||||
|
||||
// update arguments
|
||||
operator_args.ptr_A = arguments->A;
|
||||
operator_args.ptr_B = arguments->B;
|
||||
operator_args.ptr_C = arguments->C;
|
||||
operator_args.ptr_D = arguments->D;
|
||||
|
||||
operator_args.batch_stride_A = arguments->batch_stride_A;
|
||||
operator_args.batch_stride_B = arguments->batch_stride_B;
|
||||
operator_args.batch_stride_C = arguments->batch_stride_C;
|
||||
operator_args.batch_stride_D = arguments->batch_stride_D;
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
public:
|
||||
|
||||
/// Returns success if the operation can proceed
|
||||
virtual Status can_implement(
|
||||
void const *configuration_ptr,
|
||||
void const *arguments_ptr) const {
|
||||
|
||||
GemmUniversalConfiguration const *configuration =
|
||||
static_cast<GemmUniversalConfiguration const *>(configuration_ptr);
|
||||
|
||||
GemmUniversalArguments const *arguments =
|
||||
static_cast<GemmUniversalArguments const *>(arguments_ptr);
|
||||
|
||||
OperatorArguments args;
|
||||
|
||||
Status status = construct_arguments_(args, configuration);
|
||||
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
status = update_arguments_(args, arguments);
|
||||
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
return Operator::can_implement(args);
|
||||
}
|
||||
|
||||
/// Gets the host-side workspace
|
||||
virtual uint64_t get_host_workspace_size(
|
||||
void const *configuration) const {
|
||||
|
||||
return sizeof(Operator);
|
||||
}
|
||||
|
||||
/// Gets the device-side workspace
|
||||
virtual uint64_t get_device_workspace_size(
|
||||
void const *configuration_ptr) const {
|
||||
|
||||
OperatorArguments args;
|
||||
|
||||
Status status = construct_arguments_(
|
||||
args,
|
||||
static_cast<GemmUniversalConfiguration const *>(configuration_ptr));
|
||||
|
||||
if (status != Status::kSuccess) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
uint64_t size = Operator::get_workspace_size(args);
|
||||
|
||||
return size;
|
||||
}
|
||||
|
||||
/// Initializes the workspace
|
||||
virtual Status initialize(
|
||||
void const *configuration_ptr,
|
||||
void *host_workspace,
|
||||
void *device_workspace,
|
||||
cudaStream_t stream = nullptr) const {
|
||||
|
||||
OperatorArguments args;
|
||||
|
||||
Status status = construct_arguments_(
|
||||
args,
|
||||
static_cast<GemmUniversalConfiguration const *>(configuration_ptr));
|
||||
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
Operator *op = new (host_workspace) Operator;
|
||||
|
||||
status = op->initialize(args, device_workspace, stream);
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
/// Runs the kernel
|
||||
virtual Status run(
|
||||
void const *arguments_ptr,
|
||||
void *host_workspace,
|
||||
void *device_workspace = nullptr,
|
||||
cudaStream_t stream = nullptr) const {
|
||||
|
||||
OperatorArguments args;
|
||||
|
||||
Status status = update_arguments_(
|
||||
args,
|
||||
static_cast<GemmUniversalArguments const *>(arguments_ptr));
|
||||
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
Operator *op = static_cast<Operator *>(host_workspace);
|
||||
|
||||
status = op->update(args, device_workspace);
|
||||
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
status = op->run(stream);
|
||||
|
||||
return status;
|
||||
}
|
||||
};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename Operator_>
|
||||
|
||||
+206
-7
@@ -26,7 +26,7 @@
|
||||
/*! \file
|
||||
\brief CUTLASS Library handle.
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
#include <cstdint>
|
||||
|
||||
@@ -43,7 +43,8 @@ namespace library {
|
||||
Handle::Handle(
|
||||
cudaStream_t stream,
|
||||
size_t workspace_size
|
||||
):
|
||||
):
|
||||
provider_(Provider::kCUTLASS),
|
||||
stream_(stream),
|
||||
workspace_(nullptr),
|
||||
workspace_size_(0),
|
||||
@@ -95,6 +96,7 @@ Handle::Handle(Handle && handle) {
|
||||
/// Move assignment operator
|
||||
Handle & Handle::operator=(Handle && handle) {
|
||||
|
||||
provider_ = handle.provider_;
|
||||
device_ = handle.device_;
|
||||
workspace_size_ = handle.workspace_size_;
|
||||
workspace_ = handle.workspace_;
|
||||
@@ -121,6 +123,16 @@ cudaStream_t Handle::get_stream() const {
|
||||
return stream_;
|
||||
}
|
||||
|
||||
/// Gets the current provider
|
||||
Provider Handle::get_provider() const {
|
||||
return provider_;
|
||||
}
|
||||
|
||||
/// Sets the provider of operations
|
||||
void Handle::set_provider(Provider provider) {
|
||||
provider_ = provider;
|
||||
}
|
||||
|
||||
/// Gets the device workspace size
|
||||
size_t Handle::get_workspace_size() const {
|
||||
return workspace_size_;
|
||||
@@ -351,6 +363,8 @@ Status Handle::gemm(
|
||||
//
|
||||
|
||||
GemmFunctionalKey key(
|
||||
provider_,
|
||||
GemmKind::kGemm,
|
||||
element_compute,
|
||||
element_scalar,
|
||||
element_A,
|
||||
@@ -457,6 +471,188 @@ Status Handle::gemm(
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Executes a GEMM computation: D <= alpha * A*B + beta * C.
|
||||
//
|
||||
// Supports batched-strided, batched array or split-K serial or split-K parallel.
|
||||
//
|
||||
Status Handle::gemm_universal(
|
||||
|
||||
GemmUniversalMode mode, /// indicates the mode in which the kUniversal GEMM is launched
|
||||
|
||||
int M, /// GEMM M dimension
|
||||
int N, /// GEMM N dimension
|
||||
int K, /// GEMM K dimension
|
||||
|
||||
NumericTypeID element_compute, /// Data type of internal accumulation
|
||||
|
||||
NumericTypeID element_scalar, /// Data type of alpha/beta scalars
|
||||
|
||||
void const *alpha, /// Pointer to alpha scalar
|
||||
|
||||
NumericTypeID element_A, /// Data type of A matrix elements
|
||||
LayoutTypeID layout_A, /// Layout of A matrix
|
||||
ComplexTransform transform_A, /// Complex transformation applied to A matrix - ignored for real-valued matrices
|
||||
|
||||
void const * ptr_A, /// Pointer to A matrix in Global Memory
|
||||
int lda, /// Leading dimension of A matrix
|
||||
|
||||
NumericTypeID element_B, /// Data type of B matrix elements
|
||||
LayoutTypeID layout_B, /// Layout of B matrix
|
||||
ComplexTransform transform_B, /// Complex transformation applied to B matrix - ignored for real-valued matrices
|
||||
|
||||
void const * ptr_B, /// Pointer to B matrix in Global Memory
|
||||
int ldb, /// Leading dimension of B matrix
|
||||
|
||||
void const * beta, /// Pointer to beta scalar
|
||||
|
||||
NumericTypeID element_C, /// Data type of C and D matrices
|
||||
|
||||
void const * ptr_C, /// Pointer to C matrix
|
||||
int ldc, /// Leading dimension of C matrix
|
||||
|
||||
void * ptr_D, /// Pointer to D matrix
|
||||
int ldd, /// Leading dimension of D matrix
|
||||
|
||||
int batch_count, /// Batch count or number of split-K slices
|
||||
|
||||
int64_t batch_stride_A, /// Batch stride of A operand
|
||||
int64_t batch_stride_B, /// Batch stride of B operand
|
||||
int64_t batch_stride_C, /// Batch stride of C operand
|
||||
int64_t batch_stride_D /// Batch stride of D operand
|
||||
) {
|
||||
|
||||
//
|
||||
// Find the operation
|
||||
//
|
||||
|
||||
GemmFunctionalKey key(
|
||||
provider_,
|
||||
GemmKind::kUniversal,
|
||||
element_compute,
|
||||
element_scalar,
|
||||
element_A,
|
||||
layout_A,
|
||||
transform_A,
|
||||
element_B,
|
||||
layout_B,
|
||||
transform_B,
|
||||
element_C
|
||||
);
|
||||
|
||||
auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
|
||||
|
||||
if (operators_it == Singleton::get().operation_table.gemm_operations.end()) {
|
||||
return cutlass::Status::kErrorNotSupported;
|
||||
}
|
||||
|
||||
if (operators_it->second.empty()) {
|
||||
return cutlass::Status::kErrorNotSupported;
|
||||
}
|
||||
|
||||
//
|
||||
// Compute the largest alignment restriction the kernel can satisfy.
|
||||
//
|
||||
|
||||
// Maximum alignment expectation among all kernels (in units of bytes)
|
||||
int const kMaximumAlignmentSize = 16;
|
||||
|
||||
void const *ptr_A_check = ptr_A;
|
||||
void const *ptr_B_check = ptr_B;
|
||||
void const *ptr_C_check = ptr_C;
|
||||
void * ptr_D_check = ptr_D;
|
||||
|
||||
// Ignore alignment of pointers to pointers. We can't check this from the host,
|
||||
// as each batch index has its own pointer in device memory.
|
||||
if (mode == GemmUniversalMode::kArray) {
|
||||
ptr_A_check = nullptr;
|
||||
ptr_B_check = nullptr;
|
||||
ptr_C_check = nullptr;
|
||||
ptr_D_check = nullptr;
|
||||
}
|
||||
|
||||
int alignment = gemm_problem_alignment(
|
||||
M, N, K,
|
||||
element_A, ptr_A_check, lda, 0,
|
||||
element_B, ptr_B_check, ldb, 0,
|
||||
element_C, ptr_C_check, ldc, 0,
|
||||
ptr_D_check, ldd, 0, kMaximumAlignmentSize
|
||||
);
|
||||
|
||||
//
|
||||
// Find the best kernel in descending order of preference.
|
||||
//
|
||||
|
||||
GemmPreferenceKey preference_key(compute_capability(), alignment);
|
||||
|
||||
Operation const *operation = find_gemm_operation(operators_it, preference_key);
|
||||
|
||||
if (!operation) {
|
||||
return cutlass::Status::kErrorNotSupported;
|
||||
}
|
||||
|
||||
last_operation_ = operation;
|
||||
|
||||
//
|
||||
// Configure operation
|
||||
//
|
||||
|
||||
GemmUniversalConfiguration configuration{
|
||||
mode,
|
||||
{M, N, K},
|
||||
batch_count,
|
||||
lda,
|
||||
ldb,
|
||||
ldc,
|
||||
ldd
|
||||
};
|
||||
|
||||
// Query host work space size
|
||||
uint64_t host_workspace_size_needed = operation->get_host_workspace_size(&configuration);
|
||||
|
||||
if (uint64_t(kHostWorkspaceSize) < host_workspace_size_needed) {
|
||||
return cutlass::Status::kErrorNotSupported;
|
||||
}
|
||||
|
||||
char host_workspace[kHostWorkspaceSize];
|
||||
|
||||
// Query device workspace size
|
||||
uint64_t device_workspace_size_needed = operation->get_device_workspace_size(&configuration);
|
||||
|
||||
if (uint64_t(workspace_size_) < device_workspace_size_needed) {
|
||||
return cutlass::Status::kErrorNotSupported;
|
||||
}
|
||||
|
||||
// Initialize host and device workspaces
|
||||
Status status = operation->initialize(
|
||||
&configuration,
|
||||
host_workspace,
|
||||
workspace_,
|
||||
stream_);
|
||||
|
||||
if (status != cutlass::Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
// Run the operator
|
||||
GemmUniversalArguments arguments{
|
||||
ptr_A,
|
||||
ptr_B,
|
||||
ptr_C,
|
||||
ptr_D,
|
||||
alpha,
|
||||
beta,
|
||||
scalar_pointer_mode_,
|
||||
batch_stride_A,
|
||||
batch_stride_B,
|
||||
batch_stride_C,
|
||||
batch_stride_D
|
||||
};
|
||||
|
||||
return operation->run(&arguments, host_workspace, workspace_, stream_);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Planar complex GEMM
|
||||
Status Handle::gemm_planar_complex(
|
||||
|
||||
@@ -522,6 +718,8 @@ Status Handle::gemm_planar_complex(
|
||||
//
|
||||
|
||||
GemmFunctionalKey key(
|
||||
provider_,
|
||||
GemmKind::kPlanarComplex,
|
||||
element_compute,
|
||||
element_scalar,
|
||||
element_A,
|
||||
@@ -533,9 +731,9 @@ Status Handle::gemm_planar_complex(
|
||||
element_C
|
||||
);
|
||||
|
||||
auto operators_it = Singleton::get().operation_table.gemm_planar_complex_operations.find(key);
|
||||
auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
|
||||
|
||||
if (operators_it == Singleton::get().operation_table.gemm_planar_complex_operations.end()) {
|
||||
if (operators_it == Singleton::get().operation_table.gemm_operations.end()) {
|
||||
return cutlass::Status::kErrorNotSupported;
|
||||
}
|
||||
|
||||
@@ -714,6 +912,8 @@ Status Handle::gemm_planar_complex_array(
|
||||
//
|
||||
|
||||
GemmFunctionalKey key(
|
||||
provider_,
|
||||
GemmKind::kPlanarComplexArray,
|
||||
element_compute,
|
||||
element_scalar,
|
||||
element_A,
|
||||
@@ -725,9 +925,9 @@ Status Handle::gemm_planar_complex_array(
|
||||
element_C
|
||||
);
|
||||
|
||||
auto operators_it = Singleton::get().operation_table.gemm_planar_complex_array_operations.find(key);
|
||||
auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
|
||||
|
||||
if (operators_it == Singleton::get().operation_table.gemm_planar_complex_array_operations.end()) {
|
||||
if (operators_it == Singleton::get().operation_table.gemm_operations.end()) {
|
||||
return cutlass::Status::kErrorNotSupported;
|
||||
}
|
||||
|
||||
@@ -837,7 +1037,6 @@ Status Handle::gemm_planar_complex_array(
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace library
|
||||
} // namespace cutlass
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
|
||||
* Copyright (c) 2017-2020, 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:
|
||||
@@ -125,6 +125,14 @@ template <> struct NumericTypeMap<cutlass::complex<double> > {
|
||||
static NumericTypeID const kId = NumericTypeID::kCF64;
|
||||
};
|
||||
|
||||
template <> struct NumericTypeMap<cutlass::bfloat16_t> {
|
||||
static NumericTypeID const kId = NumericTypeID::kBF16;
|
||||
};
|
||||
|
||||
template <> struct NumericTypeMap<cutlass::tfloat32_t> {
|
||||
static NumericTypeID const kId = NumericTypeID::kTF32;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T> struct MathOperationMap {
|
||||
@@ -143,6 +151,10 @@ template <> struct MathOperationMap<cutlass::arch::OpMultiplyAddComplex> {
|
||||
static MathOperationID const kId = MathOperationID::kMultiplyAddComplex;
|
||||
};
|
||||
|
||||
template <> struct MathOperationMap<cutlass::arch::OpMultiplyAddGaussianComplex> {
|
||||
static MathOperationID const kId = MathOperationID::kMultiplyAddGaussianComplex;
|
||||
};
|
||||
|
||||
template <> struct MathOperationMap<cutlass::arch::OpXorPopc> {
|
||||
static MathOperationID const kId = MathOperationID::kXorPopc;
|
||||
};
|
||||
@@ -217,33 +229,43 @@ template <> struct ComplexTransformMap<cutlass::ComplexTransform::kConjugate> {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T> struct ArchMap;
|
||||
template <typename ArchTag, typename OperatorClass> struct ArchMap;
|
||||
|
||||
template <> struct ArchMap<arch::Sm50> {
|
||||
template <> struct ArchMap<arch::Sm50, arch::OpClassSimt> {
|
||||
static int const kMin = 50;
|
||||
static int const kMax = 1024;
|
||||
};
|
||||
|
||||
template <> struct ArchMap<arch::Sm60> {
|
||||
template <> struct ArchMap<arch::Sm60, arch::OpClassSimt> {
|
||||
static int const kMin = 60;
|
||||
static int const kMax = 1024;
|
||||
};
|
||||
|
||||
template <> struct ArchMap<arch::Sm61> {
|
||||
template <> struct ArchMap<arch::Sm61, arch::OpClassSimt> {
|
||||
static int const kMin = 61;
|
||||
static int const kMax = 1024;
|
||||
};
|
||||
|
||||
template <> struct ArchMap<arch::Sm70> {
|
||||
template <> struct ArchMap<arch::Sm70, arch::OpClassWmmaTensorOp> {
|
||||
static int const kMin = 70;
|
||||
static int const kMax = 1024;
|
||||
};
|
||||
|
||||
template <> struct ArchMap<arch::Sm70, arch::OpClassTensorOp> {
|
||||
static int const kMin = 70;
|
||||
static int const kMax = 75;
|
||||
};
|
||||
|
||||
template <> struct ArchMap<arch::Sm75> {
|
||||
template <typename OperatorClass> struct ArchMap<arch::Sm75, OperatorClass> {
|
||||
static int const kMin = 75;
|
||||
static int const kMax = 1024;
|
||||
};
|
||||
|
||||
template <typename OperatorClass> struct ArchMap<arch::Sm80, OperatorClass> {
|
||||
static int const kMin = 80;
|
||||
static int const kMax = 1024;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename Element, typename Layout>
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
|
||||
* Copyright (c) 2017-2020, 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:
|
||||
@@ -37,11 +37,6 @@ namespace library {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// init and insert all cutlass op in manifest object (procedurally generated using generator.py)
|
||||
void initialize_all(Manifest &manifest);
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Top-level initialization
|
||||
Status Manifest::initialize() {
|
||||
|
||||
@@ -49,13 +44,8 @@ Status Manifest::initialize() {
|
||||
operations_.clear();
|
||||
}
|
||||
|
||||
switch(provider_) {
|
||||
case Provider::kCUTLASS:
|
||||
initialize_all(*this); break;
|
||||
|
||||
default:
|
||||
break;
|
||||
}
|
||||
// initialize procedurally generated cutlass op in manifest object
|
||||
initialize_all(*this);
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
@@ -28,30 +28,7 @@
|
||||
instances may be queried.
|
||||
*/
|
||||
|
||||
#include <fstream>
|
||||
|
||||
#include "cutlass/library/library.h"
|
||||
#include "cutlass/library/operation_table.h"
|
||||
#include "cutlass/library/util.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
std::ostream & operator<<(std::ostream &out, cutlass::library::GemmFunctionalKey const &k) {
|
||||
|
||||
out << "{\n"
|
||||
<< " element_compute: " << to_string(k.element_compute) << "\n"
|
||||
<< " element_scalar: " << to_string(k.element_scalar) << "\n"
|
||||
<< " element_A: " << to_string(k.element_A) << "\n"
|
||||
<< " layout_A: " << to_string(k.layout_A) << "\n"
|
||||
<< " transform_A: " << to_string(k.transform_A) << "\n"
|
||||
<< " element_B: " << to_string(k.element_B) << "\n"
|
||||
<< " layout_B: " << to_string(k.layout_B) << "\n"
|
||||
<< " transform_B: " << to_string(k.transform_B) << "\n"
|
||||
<< " element_C: " << to_string(k.element_C) << "\n"
|
||||
<< "}";
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -67,85 +44,38 @@ void OperationTable::append(Manifest const &manifest) {
|
||||
|
||||
OperationDescription const &desc = operation->description();
|
||||
|
||||
// insert all gemm operation into operation table
|
||||
if (desc.kind == OperationKind::kGemm) {
|
||||
GemmDescription const &gemm_desc = static_cast<GemmDescription const &>(desc);
|
||||
|
||||
if (gemm_desc.gemm_kind == GemmKind::kGemm) {
|
||||
|
||||
GemmFunctionalKey functional_key(
|
||||
gemm_desc.tile_description.math_instruction.element_accumulator,
|
||||
gemm_desc.element_epilogue,
|
||||
gemm_desc.A.element,
|
||||
gemm_desc.A.layout,
|
||||
gemm_desc.transform_A,
|
||||
gemm_desc.B.element,
|
||||
gemm_desc.B.layout,
|
||||
gemm_desc.transform_B,
|
||||
gemm_desc.C.element
|
||||
);
|
||||
GemmFunctionalKey functional_key(
|
||||
gemm_desc.provider,
|
||||
gemm_desc.gemm_kind,
|
||||
gemm_desc.tile_description.math_instruction.element_accumulator,
|
||||
gemm_desc.element_epilogue,
|
||||
gemm_desc.A.element,
|
||||
gemm_desc.A.layout,
|
||||
gemm_desc.transform_A,
|
||||
gemm_desc.B.element,
|
||||
gemm_desc.B.layout,
|
||||
gemm_desc.transform_B,
|
||||
gemm_desc.C.element
|
||||
);
|
||||
|
||||
Operation const *op = operation.get();
|
||||
Operation const *op = operation.get();
|
||||
|
||||
int cc = gemm_desc.tile_description.minimum_compute_capability;
|
||||
int cc = gemm_desc.tile_description.minimum_compute_capability;
|
||||
|
||||
int alignment = std::max(std::max(
|
||||
gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
|
||||
int alignment = std::max(std::max(
|
||||
gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
|
||||
|
||||
GemmPreferenceKey preference_key(cc, alignment);
|
||||
GemmPreferenceKey preference_key(cc, alignment);
|
||||
|
||||
gemm_operations[functional_key][preference_key].push_back(op);
|
||||
}
|
||||
else if (gemm_desc.gemm_kind == GemmKind::kPlanarComplex) {
|
||||
|
||||
GemmFunctionalKey functional_key(
|
||||
gemm_desc.tile_description.math_instruction.element_accumulator,
|
||||
gemm_desc.element_epilogue,
|
||||
gemm_desc.A.element,
|
||||
gemm_desc.A.layout,
|
||||
gemm_desc.transform_A,
|
||||
gemm_desc.B.element,
|
||||
gemm_desc.B.layout,
|
||||
gemm_desc.transform_B,
|
||||
gemm_desc.C.element
|
||||
);
|
||||
|
||||
Operation const *op = operation.get();
|
||||
|
||||
int cc = gemm_desc.tile_description.minimum_compute_capability;
|
||||
|
||||
int alignment = std::max(std::max(
|
||||
gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
|
||||
|
||||
GemmPreferenceKey preference_key(cc, alignment);
|
||||
|
||||
gemm_planar_complex_operations[functional_key][preference_key].push_back(op);
|
||||
}
|
||||
else if (gemm_desc.gemm_kind == GemmKind::kPlanarComplexArray) {
|
||||
|
||||
GemmFunctionalKey functional_key(
|
||||
gemm_desc.tile_description.math_instruction.element_accumulator,
|
||||
gemm_desc.element_epilogue,
|
||||
gemm_desc.A.element,
|
||||
gemm_desc.A.layout,
|
||||
gemm_desc.transform_A,
|
||||
gemm_desc.B.element,
|
||||
gemm_desc.B.layout,
|
||||
gemm_desc.transform_B,
|
||||
gemm_desc.C.element
|
||||
);
|
||||
|
||||
Operation const *op = operation.get();
|
||||
|
||||
int cc = gemm_desc.tile_description.minimum_compute_capability;
|
||||
|
||||
int alignment = std::max(std::max(
|
||||
gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
|
||||
|
||||
GemmPreferenceKey preference_key(cc, alignment);
|
||||
|
||||
gemm_planar_complex_array_operations[functional_key][preference_key].push_back(op);
|
||||
}
|
||||
gemm_operations[functional_key][preference_key].push_back(op);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
+203
-1
@@ -45,6 +45,7 @@ static struct {
|
||||
Provider enumerant;
|
||||
}
|
||||
Provider_enumerants[] = {
|
||||
{"none", "None", Provider::kNone},
|
||||
{"cutlass", "CUTLASS", Provider::kCUTLASS},
|
||||
{"host", "reference_host", Provider::kReferenceHost},
|
||||
{"device", "reference_device", Provider::kReferenceDevice},
|
||||
@@ -83,6 +84,38 @@ Provider from_string<Provider>(std::string const &str) {
|
||||
}
|
||||
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
static struct {
|
||||
char const *text;
|
||||
char const *pretty;
|
||||
GemmKind enumerant;
|
||||
}
|
||||
GemmKind_enumerants[] = {
|
||||
{"gemm", "<Gemm>", GemmKind::kGemm},
|
||||
{"batched", "<Batched>", GemmKind::kBatched},
|
||||
{"array", "<Array>", GemmKind::kArray},
|
||||
{"universal", "<Universal>", GemmKind::kUniversal},
|
||||
{"planar_complex", "<PlanarComplex>", GemmKind::kPlanarComplex},
|
||||
{"planar_complex_array", "<PlanarComplexArray>", GemmKind::kPlanarComplexArray},
|
||||
};
|
||||
|
||||
/// Converts a ConvKind enumerant to a string
|
||||
char const *to_string(GemmKind type, bool pretty) {
|
||||
|
||||
for (auto const & possible : GemmKind_enumerants) {
|
||||
if (type == possible.enumerant) {
|
||||
if (pretty) {
|
||||
return possible.pretty;
|
||||
}
|
||||
else {
|
||||
return possible.text;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return pretty ? "Invalid" : "invalid";
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -92,6 +125,7 @@ static struct {
|
||||
OperationKind enumerant;
|
||||
}
|
||||
OperationKind_enumerants[] = {
|
||||
{"eq_gemm", "EqGemm", OperationKind::kEqGemm},
|
||||
{"gemm", "Gemm", OperationKind::kGemm},
|
||||
};
|
||||
|
||||
@@ -194,10 +228,14 @@ NumericTypeID_enumerants[] = {
|
||||
{"s32", "S32", NumericTypeID::kS32},
|
||||
{"s64", "S64", NumericTypeID::kS64},
|
||||
{"f16", "F16", NumericTypeID::kF16},
|
||||
{"bf16", "BF16", NumericTypeID::kBF16},
|
||||
{"f32", "F32", NumericTypeID::kF32},
|
||||
{"tf32", "TF32", NumericTypeID::kTF32},
|
||||
{"f64", "F64", NumericTypeID::kF64},
|
||||
{"cf16", "CF16", NumericTypeID::kCF16},
|
||||
{"cbf16", "CBF16", NumericTypeID::kCBF16},
|
||||
{"cf32", "CF32", NumericTypeID::kCF32},
|
||||
{"ctf32", "CTF32", NumericTypeID::kCTF32},
|
||||
{"cf64", "CF64", NumericTypeID::kCF64},
|
||||
{"cu4", "CU4", NumericTypeID::kCU4},
|
||||
{"cu8", "CU8", NumericTypeID::kCU8},
|
||||
@@ -249,10 +287,14 @@ NumericTypeID from_string<NumericTypeID>(std::string const &str) {
|
||||
int sizeof_bits(NumericTypeID type) {
|
||||
switch (type) {
|
||||
case NumericTypeID::kF16: return 16;
|
||||
case NumericTypeID::kBF16: return 16;
|
||||
case NumericTypeID::kTF32: return 32;
|
||||
case NumericTypeID::kF32: return 32;
|
||||
case NumericTypeID::kF64: return 64;
|
||||
case NumericTypeID::kCF16: return 32;
|
||||
case NumericTypeID::kCBF16: return 32;
|
||||
case NumericTypeID::kCF32: return 64;
|
||||
case NumericTypeID::kCTF32: return 64;
|
||||
case NumericTypeID::kCF64: return 128;
|
||||
case NumericTypeID::kS4: return 4;
|
||||
case NumericTypeID::kS8: return 8;
|
||||
@@ -276,6 +318,8 @@ bool is_complex_type(NumericTypeID type) {
|
||||
case NumericTypeID::kCF16: return true;
|
||||
case NumericTypeID::kCF32: return true;
|
||||
case NumericTypeID::kCF64: return true;
|
||||
case NumericTypeID::kCBF16: return true;
|
||||
case NumericTypeID::kCTF32: return true;
|
||||
default: break;
|
||||
}
|
||||
return false;
|
||||
@@ -287,6 +331,8 @@ NumericTypeID get_real_type(NumericTypeID type) {
|
||||
case NumericTypeID::kCF16: return NumericTypeID::kF16;
|
||||
case NumericTypeID::kCF32: return NumericTypeID::kF32;
|
||||
case NumericTypeID::kCF64: return NumericTypeID::kF64;
|
||||
case NumericTypeID::kCBF16: return NumericTypeID::kBF16;
|
||||
case NumericTypeID::kCTF32: return NumericTypeID::kTF32;
|
||||
default: break;
|
||||
}
|
||||
return type;
|
||||
@@ -314,6 +360,8 @@ bool is_integer_type(NumericTypeID type) {
|
||||
bool is_signed_type(NumericTypeID type) {
|
||||
switch (type) {
|
||||
case NumericTypeID::kF16: return true;
|
||||
case NumericTypeID::kBF16: return true;
|
||||
case NumericTypeID::kTF32: return true;
|
||||
case NumericTypeID::kF32: return true;
|
||||
case NumericTypeID::kF64: return true;
|
||||
case NumericTypeID::kS4: return true;
|
||||
@@ -340,9 +388,13 @@ bool is_unsigned_integer(NumericTypeID type) {
|
||||
bool is_float_type(NumericTypeID type) {
|
||||
switch (type) {
|
||||
case NumericTypeID::kF16: return true;
|
||||
case NumericTypeID::kBF16: return true;
|
||||
case NumericTypeID::kTF32: return true;
|
||||
case NumericTypeID::kF32: return true;
|
||||
case NumericTypeID::kF64: return true;
|
||||
case NumericTypeID::kCF16: return true;
|
||||
case NumericTypeID::kCBF16: return true;
|
||||
case NumericTypeID::kCTF32: return true;
|
||||
case NumericTypeID::kCF32: return true;
|
||||
case NumericTypeID::kCF64: return true;
|
||||
default: break;
|
||||
@@ -431,7 +483,7 @@ OpcodeClassID_enumerants[] = {
|
||||
{"simt", "<simt>", OpcodeClassID::kSimt},
|
||||
{"tensorop", "<tensorop>", OpcodeClassID::kTensorOp},
|
||||
{"wmmatensorop", "<wmmatensorop>", OpcodeClassID::kWmmaTensorOp},
|
||||
{"wmma", "<wmma>", OpcodeClassID::kWmmaTensorOp}
|
||||
{"wmma", "<wmma>", OpcodeClassID::kWmmaTensorOp},
|
||||
};
|
||||
|
||||
/// Converts a OpcodeClassID enumerant to a string
|
||||
@@ -509,6 +561,47 @@ ComplexTransform from_string<ComplexTransform>(std::string const &str) {
|
||||
}
|
||||
|
||||
|
||||
static struct {
|
||||
char const *text;
|
||||
char const *pretty;
|
||||
SplitKMode enumerant;
|
||||
}
|
||||
SplitKMode_enumerants[] = {
|
||||
{"serial", "<serial>", SplitKMode::kSerial},
|
||||
{"parallel", "<parallel>", SplitKMode::kParallel},
|
||||
};
|
||||
|
||||
/// Converts a SplitKMode enumerant to a string
|
||||
char const *to_string(SplitKMode type, bool pretty) {
|
||||
|
||||
for (auto const & possible : SplitKMode_enumerants) {
|
||||
if (type == possible.enumerant) {
|
||||
if (pretty) {
|
||||
return possible.pretty;
|
||||
}
|
||||
else {
|
||||
return possible.text;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return pretty ? "Invalid" : "invalid";
|
||||
}
|
||||
|
||||
/// Converts a SplitKMode enumerant from a string
|
||||
template <>
|
||||
SplitKMode from_string<SplitKMode>(std::string const &str) {
|
||||
|
||||
for (auto const & possible : SplitKMode_enumerants) {
|
||||
if ((str.compare(possible.text) == 0) ||
|
||||
(str.compare(possible.pretty) == 0)) {
|
||||
return possible.enumerant;
|
||||
}
|
||||
}
|
||||
|
||||
return SplitKMode::kInvalid;
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/// Lexical cast a string to a byte array. Returns true if cast is successful or false if invalid.
|
||||
bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string const &str) {
|
||||
@@ -570,6 +663,20 @@ bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string c
|
||||
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(tmp);
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kBF16:
|
||||
{
|
||||
float tmp;
|
||||
ss >> tmp;
|
||||
*reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(tmp);
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kTF32:
|
||||
{
|
||||
float tmp;
|
||||
ss >> tmp;
|
||||
*reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(tmp);
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kF32:
|
||||
{
|
||||
ss >> *reinterpret_cast<float *>(bytes.data());
|
||||
@@ -589,11 +696,29 @@ bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string c
|
||||
x->imag() = static_cast<half_t>(std::imag(tmp));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCBF16:
|
||||
{
|
||||
std::complex<float> tmp;
|
||||
ss >> tmp;
|
||||
cutlass::complex<cutlass::bfloat16_t> *x = reinterpret_cast<cutlass::complex<bfloat16_t> *>(bytes.data());
|
||||
x->real() = static_cast<bfloat16_t>(std::real(tmp));
|
||||
x->imag() = static_cast<bfloat16_t>(std::imag(tmp));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCF32:
|
||||
{
|
||||
ss >> *reinterpret_cast<std::complex<float>*>(bytes.data());
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCTF32:
|
||||
{
|
||||
std::complex<float> tmp;
|
||||
ss >> tmp;
|
||||
cutlass::complex<cutlass::tfloat32_t> *x = reinterpret_cast<cutlass::complex<tfloat32_t> *>(bytes.data());
|
||||
x->real() = static_cast<tfloat32_t>(std::real(tmp));
|
||||
x->imag() = static_cast<tfloat32_t>(std::imag(tmp));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCF64:
|
||||
{
|
||||
ss >> *reinterpret_cast<std::complex<double>*>(bytes.data());
|
||||
@@ -674,6 +799,18 @@ std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type) {
|
||||
ss << tmp;
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kBF16:
|
||||
{
|
||||
float tmp = *reinterpret_cast<bfloat16_t *>(bytes.data());;
|
||||
ss << tmp;
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kTF32:
|
||||
{
|
||||
float tmp = *reinterpret_cast<tfloat32_t *>(bytes.data());;
|
||||
ss << tmp;
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kF32:
|
||||
{
|
||||
ss << *reinterpret_cast<float *>(bytes.data());
|
||||
@@ -696,6 +833,18 @@ std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type) {
|
||||
}
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCBF16:
|
||||
{
|
||||
cutlass::complex<bfloat16_t> const *x =
|
||||
reinterpret_cast<cutlass::complex<bfloat16_t> const *>(bytes.data());
|
||||
|
||||
ss << float(x->real());
|
||||
|
||||
if (x->imag() != cutlass::bfloat16_t()) {
|
||||
ss << "+i" << float(x->imag());
|
||||
}
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCF32:
|
||||
{
|
||||
cutlass::complex<float> const * x = reinterpret_cast<cutlass::complex<float> const *>(bytes.data());
|
||||
@@ -707,6 +856,17 @@ std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type) {
|
||||
}
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCTF32:
|
||||
{
|
||||
cutlass::complex<tfloat32_t> const * x = reinterpret_cast<cutlass::complex<tfloat32_t> const *>(bytes.data());
|
||||
|
||||
ss << float(x->real());
|
||||
|
||||
if (x->imag() != tfloat32_t()) {
|
||||
ss << "+i" << float(x->imag());
|
||||
}
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCF64:
|
||||
{
|
||||
cutlass::complex<double> const * x = reinterpret_cast<cutlass::complex<double> const *>(bytes.data());
|
||||
@@ -780,6 +940,16 @@ bool cast_from_int64(std::vector<uint8_t> &bytes, NumericTypeID type, int64_t sr
|
||||
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kBF16:
|
||||
{
|
||||
*reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(float(src));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kTF32:
|
||||
{
|
||||
*reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(float(src));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kF32:
|
||||
{
|
||||
*reinterpret_cast<float *>(bytes.data()) = static_cast<float>(src);
|
||||
@@ -870,6 +1040,16 @@ bool cast_from_uint64(std::vector<uint8_t> &bytes, NumericTypeID type, uint64_t
|
||||
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kBF16:
|
||||
{
|
||||
*reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(float(src));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kTF32:
|
||||
{
|
||||
*reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(float(src));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kF32:
|
||||
{
|
||||
*reinterpret_cast<float *>(bytes.data()) = static_cast<float>(src);
|
||||
@@ -961,6 +1141,16 @@ bool cast_from_double(std::vector<uint8_t> &bytes, NumericTypeID type, double sr
|
||||
*reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kBF16:
|
||||
{
|
||||
*reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(float(src));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kTF32:
|
||||
{
|
||||
*reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(float(src));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kF32:
|
||||
{
|
||||
*reinterpret_cast<float *>(bytes.data()) = static_cast<float>(src);
|
||||
@@ -978,11 +1168,23 @@ bool cast_from_double(std::vector<uint8_t> &bytes, NumericTypeID type, double sr
|
||||
x->imag() = static_cast<half_t>(float(0));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCBF16:
|
||||
{
|
||||
cutlass::complex<cutlass::bfloat16_t> *x = reinterpret_cast<cutlass::complex<bfloat16_t> *>(bytes.data());
|
||||
x->real() = static_cast<bfloat16_t>(bfloat16_t(src));
|
||||
x->imag() = static_cast<bfloat16_t>(bfloat16_t(0));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCF32:
|
||||
{
|
||||
*reinterpret_cast<std::complex<float>*>(bytes.data()) = std::complex<float>(float(src), float(0));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCTF32:
|
||||
{
|
||||
*reinterpret_cast<std::complex<tfloat32_t>*>(bytes.data()) = std::complex<tfloat32_t>(tfloat32_t(src), tfloat32_t(0));
|
||||
}
|
||||
break;
|
||||
case NumericTypeID::kCF64:
|
||||
{
|
||||
*reinterpret_cast<std::complex<double>*>(bytes.data()) = std::complex<double>(src, double(0));
|
||||
|
||||
Reference in New Issue
Block a user