CUTLASS 2.2 (#96)
Adds support for NVIDIA Ampere Architecture features. CUDA 11 Toolkit recommended.
This commit is contained in:
Andrew Kerr
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -52,6 +52,10 @@ struct Sm72 {
|
||||
struct Sm75 {
|
||||
static int const kMinComputeCapability = 75;
|
||||
};
|
||||
struct Sm80 {
|
||||
static int const kMinComputeCapability = 80;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace arch
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief Directives related to cache operations
|
||||
*/
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
namespace cutlass {
|
||||
namespace arch {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Controls PTX cache operations
|
||||
struct CacheOperation {
|
||||
enum Kind {
|
||||
/// Cache at all levels - accessed again
|
||||
Always,
|
||||
/// Cache at global level
|
||||
Global,
|
||||
/// Streaming - likely to be accessed once
|
||||
Streaming,
|
||||
/// Indicates the line will not be used again
|
||||
LastUse,
|
||||
/// Don't cache, and fetch again
|
||||
Volatile,
|
||||
/// Write back at all coherent levels
|
||||
WriteBack,
|
||||
/// Write through to system memory
|
||||
WriteThrough
|
||||
};
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace arch
|
||||
} // namespace cutlass
|
||||
@@ -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:
|
||||
@@ -28,13 +28,271 @@
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
namespace cutlass {
|
||||
namespace arch {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
/// Fragment type to store loaded data
|
||||
typename AccessType,
|
||||
/// The bytes of loading
|
||||
int LoadBytes
|
||||
>
|
||||
struct global_load;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Specializations
|
||||
//
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename AccessType
|
||||
>
|
||||
struct global_load<AccessType,
|
||||
32
|
||||
> {
|
||||
CUTLASS_DEVICE
|
||||
global_load(AccessType &D, void const *ptr, bool pred_guard) {
|
||||
uint4 *data = reinterpret_cast<uint4 *>(&D);
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %9, 0;\n"
|
||||
" mov.b32 %0, %10;\n"
|
||||
" mov.b32 %1, %11;\n"
|
||||
" mov.b32 %2, %12;\n"
|
||||
" mov.b32 %3, %13;\n"
|
||||
" mov.b32 %4, %14;\n"
|
||||
" mov.b32 %5, %15;\n"
|
||||
" mov.b32 %6, %16;\n"
|
||||
" mov.b32 %7, %17;\n"
|
||||
" @p ld.global.v4.u32 {%0, %1, %2, %3}, [%8];\n"
|
||||
" @p ld.global.v4.u32 {%4, %5, %6, %7}, [%18];\n"
|
||||
"}\n"
|
||||
: "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w),
|
||||
"=r"(data[1].x), "=r"(data[1].y), "=r"(data[1].z), "=r"(data[1].w)
|
||||
: "l"(ptr), "r"((int)pred_guard), "r"(data[0].x), "r"(data[0].y),
|
||||
"r"(data[0].z), "r"(data[0].w), "r"(data[1].x), "r"(data[1].y),
|
||||
"r"(data[1].z), "r"(data[1].w), "l"(((uint8_t *)ptr) + 16));
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
template <typename AccessType
|
||||
>
|
||||
struct global_load<AccessType,
|
||||
16
|
||||
> {
|
||||
CUTLASS_DEVICE
|
||||
global_load(AccessType &D, void const *ptr, bool pred_guard) {
|
||||
uint4 &data = reinterpret_cast<uint4 &>(D);
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %5, 0;\n"
|
||||
" mov.b32 %0, %6;\n"
|
||||
" mov.b32 %1, %7;\n"
|
||||
" mov.b32 %2, %8;\n"
|
||||
" mov.b32 %3, %9;\n"
|
||||
" @p ld.global.v4.u32 {%0, %1, %2, %3}, [%4];\n"
|
||||
"}\n"
|
||||
: "=r"(data.x), "=r"(data.y), "=r"(data.z), "=r"(data.w)
|
||||
: "l"(ptr), "r"((int)pred_guard), "r"(data.x), "r"(data.y), "r"(data.z), "r"(data.w));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename AccessType
|
||||
>
|
||||
struct global_load<AccessType,
|
||||
8
|
||||
> {
|
||||
CUTLASS_DEVICE
|
||||
global_load(AccessType &D, void const *ptr, bool pred_guard) {
|
||||
uint2 &data = reinterpret_cast<uint2 &>(D);
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %3, 0;\n"
|
||||
" mov.b32 %0, %4;\n"
|
||||
" mov.b32 %1, %5;\n"
|
||||
" @p ld.global.v2.u32 {%0, %1}, [%2];\n"
|
||||
"}\n"
|
||||
: "=r"(data.x), "=r"(data.y)
|
||||
: "l"(ptr), "r"((int)pred_guard), "r"(data.x), "r"(data.y));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename AccessType
|
||||
>
|
||||
struct global_load<AccessType,
|
||||
4
|
||||
> {
|
||||
CUTLASS_DEVICE
|
||||
global_load(AccessType &D, void const *ptr, bool pred_guard) {
|
||||
unsigned &data = reinterpret_cast<unsigned &>(D);
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %2, 0;\n"
|
||||
" mov.b32 %0, %3;\n"
|
||||
" @p ld.global.u32 %0, [%1];\n"
|
||||
"}\n"
|
||||
: "=r"(data)
|
||||
: "l"(ptr), "r"((int)pred_guard), "r"(data));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename AccessType
|
||||
>
|
||||
struct global_load<AccessType,
|
||||
2
|
||||
> {
|
||||
CUTLASS_DEVICE
|
||||
global_load(AccessType &D, void const *ptr, bool pred_guard) {
|
||||
uint16_t &data = reinterpret_cast<uint16_t &>(D);
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %2, 0;\n"
|
||||
" mov.b16 %0, %3;\n"
|
||||
" @p ld.global.u16 %0, [%1];\n"
|
||||
"}\n"
|
||||
: "=h"(data)
|
||||
: "l"(ptr), "r"((int)pred_guard), "h"(data));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename AccessType
|
||||
>
|
||||
struct global_load<AccessType,
|
||||
1
|
||||
> {
|
||||
CUTLASS_DEVICE
|
||||
global_load(AccessType &D, void const *ptr, bool pred_guard) {
|
||||
if (pred_guard) D = *(reinterpret_cast<AccessType const *>(ptr));
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
/// Fragment type to store loaded data
|
||||
typename AccessType,
|
||||
/// The bytes of loading
|
||||
int LoadBytes
|
||||
>
|
||||
struct global_store;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Specializations
|
||||
//
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename AccessType>
|
||||
struct global_store<AccessType, 32> {
|
||||
CUTLASS_DEVICE
|
||||
global_store(AccessType const &D, void *ptr, bool pred_guard) {
|
||||
uint4 const *data = reinterpret_cast<uint4 const *>(&D);
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %5, 0;\n"
|
||||
" @p st.global.v4.u32 [%0], {%1, %2, %3, %4};\n"
|
||||
" @p st.global.v4.u32 [%6], {%7, %8, %9, %10};\n"
|
||||
"}\n"
|
||||
:
|
||||
: "l"(ptr), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z),
|
||||
"r"(data[0].w), "r"((int)pred_guard), "l"(((uint8_t *)ptr) + 16),
|
||||
"r"(data[1].x), "r"(data[1].y), "r"(data[1].z), "r"(data[1].w));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename AccessType>
|
||||
struct global_store<AccessType, 16> {
|
||||
CUTLASS_DEVICE
|
||||
global_store(AccessType const &D, void *ptr, bool pred_guard) {
|
||||
uint4 const &data = reinterpret_cast<uint4 const &>(D);
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %5, 0;\n"
|
||||
" @p st.global.v4.u32 [%0], {%1, %2, %3, %4};\n"
|
||||
"}\n"
|
||||
:
|
||||
: "l"(ptr), "r"(data.x), "r"(data.y), "r"(data.z), "r"(data.w), "r"((int)pred_guard));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename AccessType>
|
||||
struct global_store<AccessType, 8> {
|
||||
CUTLASS_DEVICE
|
||||
global_store(AccessType const &D, void *ptr, bool pred_guard) {
|
||||
uint2 const &data = reinterpret_cast<uint2 const &>(D);
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %3, 0;\n"
|
||||
" @p st.global.v2.u32 [%0], {%1, %2};\n"
|
||||
"}\n"
|
||||
:
|
||||
: "l"(ptr), "r"(data.x), "r"(data.y), "r"((int)pred_guard));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename AccessType>
|
||||
struct global_store<AccessType, 4> {
|
||||
CUTLASS_DEVICE
|
||||
global_store(AccessType const &D, void *ptr, bool pred_guard) {
|
||||
uint32_t const &data = reinterpret_cast<uint32_t const &>(D);
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %2, 0;\n"
|
||||
" @p st.global.u32 [%0], %1;\n"
|
||||
"}\n"
|
||||
:
|
||||
: "l"(ptr), "r"(data), "r"((int)pred_guard));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename AccessType>
|
||||
struct global_store<AccessType, 2> {
|
||||
CUTLASS_DEVICE
|
||||
global_store(AccessType const &D, void *ptr, bool pred_guard) {
|
||||
uint16_t const &data = reinterpret_cast<uint16_t const &>(D);
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %2, 0;\n"
|
||||
" @p st.global.u16 [%0], %1;\n"
|
||||
"}\n"
|
||||
:
|
||||
: "l"(ptr), "h"(data), "r"((int)pred_guard));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename AccessType>
|
||||
struct global_store<AccessType, 1> {
|
||||
CUTLASS_DEVICE
|
||||
global_store(AccessType const &D, void *ptr, bool pred_guard) {
|
||||
if (pred_guard) *(reinterpret_cast<AccessType *>(ptr)) = D;
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
} // namespace arch
|
||||
} // namespace cutlass
|
||||
@@ -42,4 +300,6 @@ namespace arch {
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#include "memory_sm75.h"
|
||||
#include "memory_sm80.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -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:
|
||||
@@ -50,20 +50,20 @@ inline __device__ void ldsm(Array<unsigned, MatrixCount> & D, void const* ptr);
|
||||
//
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#if ! defined(CUDA_LDMATRIX_SUPPORTED)
|
||||
#define CUDA_LDMATRIX_SUPPORTED ((__CUDACC_VER_MAJOR__ == 10) && (__CUDACC_VER_MINOR__ >= 2))
|
||||
#if (__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2) || (__CUDACC_VER_MAJOR__ >= 11)
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750)
|
||||
#define CUDA_LDMATRIX_ACTIVATED 1
|
||||
#endif
|
||||
|
||||
#if ! defined(CUDA_LDMATRIX_ENABLED)
|
||||
#define CUDA_LDMATRIX_ENABLED CUDA_LDMATRIX_SUPPORTED
|
||||
#endif
|
||||
|
||||
#if CUDA_LDMATRIX_ENABLED && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750)
|
||||
#define CUDA_LDMATRIX_ACTIVATED 1
|
||||
#define CUDA_LDMATRIX_SUPPORTED 1
|
||||
#endif
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/*
|
||||
#if ! defined(CUDA_NVVM_GET_SMEM_POINTER_SUPPORTED) && (__CUDACC_VER_MAJOR__ > 10)
|
||||
#define CUDA_NVVM_GET_SMEM_POINTER_SUPPORTED 1
|
||||
#endif
|
||||
#if ! defined(CUDA_NVVM_GET_SMEM_POINTER_SUPPORTED)
|
||||
#define CUDA_NVVM_GET_SMEM_POINTER_SUPPORTED ((__CUDACC_VER_MAJOR__ == 10) && (__CUDACC_VER_MINOR__ >= 1))
|
||||
#endif
|
||||
@@ -71,8 +71,9 @@ inline __device__ void ldsm(Array<unsigned, MatrixCount> & D, void const* ptr);
|
||||
#if ! defined(CUDA_NVVM_GET_SMEM_POINTER_ENABLED)
|
||||
#define CUDA_NVVM_GET_SMEM_POINTER_ENABLED CUDA_NVVM_GET_SMEM_POINTER_SUPPORTED
|
||||
#endif
|
||||
*/
|
||||
|
||||
#if CUDA_NVVM_GET_SMEM_POINTER_ENABLED
|
||||
#if (__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2)
|
||||
extern "C" {
|
||||
//
|
||||
// This NVVM intrinsic is subject to change in future versions of CUDA.
|
||||
@@ -85,19 +86,49 @@ inline __device__ void ldsm(Array<unsigned, MatrixCount> & D, void const* ptr);
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#if CUDA_NVVM_GET_SMEM_POINTER_ENABLED
|
||||
/// CUTLASS helper to get SMEM pointer
|
||||
inline __device__ unsigned cutlass_get_smem_pointer(void *ptr) {
|
||||
|
||||
// We prefer to use the new CVTA intrinsics if they are available, otherwise we will fall back to
|
||||
// the previous internal intrinsics if they are available.
|
||||
#if (defined(__CUDA_ARCH__) && __CUDACC_VER_MAJOR__ >= 11)
|
||||
//
|
||||
// This NVVM intrinsic converts an address in shared memory to a plain
|
||||
// unsigned integer. This is necessary to pass to shared memory instructions
|
||||
// in inline PTX.
|
||||
//
|
||||
// In CUDA 11 and beyond, this replaces __nvvm_get_smem_pointer() [only available in 10.2].
|
||||
//
|
||||
//__device__ size_t __cvta_generic_to_shared(void* ptr);
|
||||
|
||||
/// CUTLASS helper to get SMEM pointer
|
||||
inline __device__ unsigned cutlass_get_smem_pointer(void const *ptr) {
|
||||
return __nvvm_get_smem_pointer(const_cast<void *>(ptr));
|
||||
}
|
||||
return static_cast<unsigned>(__cvta_generic_to_shared(ptr));
|
||||
|
||||
/// CUTLASS helper to get SMEM pointer
|
||||
inline __device__ unsigned cutlass_get_smem_pointer(void *ptr) {
|
||||
return __nvvm_get_smem_pointer(ptr);
|
||||
}
|
||||
#elif (defined(__CUDA_ARCH__) && __CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2)
|
||||
|
||||
return __nvvm_get_smem_pointer(ptr);
|
||||
|
||||
#elif defined(__CUDA_ARCH__)
|
||||
|
||||
uint32_t smem_ptr;
|
||||
|
||||
asm(
|
||||
"{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }\n"
|
||||
: "=r"(smem_ptr) : "l"(ptr));
|
||||
|
||||
return smem_ptr;
|
||||
|
||||
#else
|
||||
|
||||
return 0;
|
||||
#endif
|
||||
}
|
||||
|
||||
/// CUTLASS helper to get SMEM pointer
|
||||
inline __device__ unsigned cutlass_get_smem_pointer(void const *ptr) {
|
||||
return cutlass_get_smem_pointer(const_cast<void *>(ptr));
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <>
|
||||
@@ -235,5 +266,6 @@ inline __device__ void ldsm<layout::ColumnMajor, 4>(
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace arch
|
||||
} // namespace cutlass
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Architecture-specific operators on memory added for SM80
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/arch/memory_sm75.h"
|
||||
#include "cutlass/arch/cache_operation.h"
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
#define CUDA_CP_ASYNC_ACTIVATED 1
|
||||
#else
|
||||
#define CUDA_CP_ASYNC_ACTIVATED 0
|
||||
#endif
|
||||
|
||||
namespace cutlass {
|
||||
namespace arch {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Initiates an asynchronous copy from global memory to shared memory.
|
||||
///
|
||||
/// LDGSTS
|
||||
///
|
||||
template <
|
||||
/// Size of the access in bytes
|
||||
int SizeInBytes,
|
||||
/// Cache operation
|
||||
CacheOperation::Kind cache_op = CacheOperation::Always>
|
||||
struct cp_async;
|
||||
|
||||
/// Initiates an asynchronous copy from global memory to shared memory. Rather than predicate
|
||||
/// the entire transfer, zeros are written to SMEM if the guard predicate is false.
|
||||
///
|
||||
/// LDGSTS
|
||||
///
|
||||
template <
|
||||
/// Size of the access in bytes
|
||||
int SizeInBytes,
|
||||
/// Cache operation
|
||||
CacheOperation::Kind cache_op = CacheOperation::Always>
|
||||
struct cp_async_zfill;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization
|
||||
template <
|
||||
/// Size of the access in bytes
|
||||
int SizeInBytes>
|
||||
struct cp_async<SizeInBytes, CacheOperation::Always> {
|
||||
/// Copy
|
||||
CUTLASS_DEVICE
|
||||
cp_async(void *smem_ptr, void const *global_ptr, bool pred_guard = true) {
|
||||
#if CUDA_CP_ASYNC_ACTIVATED
|
||||
|
||||
unsigned smem_int_ptr = cutlass_get_smem_pointer(smem_ptr);
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %0, 0;\n"
|
||||
" @p cp.async.ca.shared.global [%1], [%2], %3;\n"
|
||||
"}\n" ::"r"((int)pred_guard),
|
||||
"r"(smem_int_ptr), "l"(global_ptr), "n"(SizeInBytes));
|
||||
|
||||
#else
|
||||
using AccessType = Array<uint8_t, SizeInBytes>;
|
||||
|
||||
if (pred_guard) {
|
||||
*static_cast<AccessType *>(smem_ptr) = *static_cast<AccessType const *>(global_ptr);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
/// Partial specialization
|
||||
template <
|
||||
/// Size of the access in bytes
|
||||
int SizeInBytes>
|
||||
struct cp_async_zfill<SizeInBytes, CacheOperation::Always> {
|
||||
/// Copy with zero fill
|
||||
CUTLASS_DEVICE
|
||||
cp_async_zfill(void *smem_ptr, void const *global_ptr, bool pred_guard) {
|
||||
#if CUDA_CP_ASYNC_ACTIVATED
|
||||
|
||||
unsigned smem_int_ptr = cutlass_get_smem_pointer(smem_ptr);
|
||||
int src_in_bytes = (pred_guard ? SizeInBytes : 0);
|
||||
|
||||
asm volatile(
|
||||
"cp.async.ca.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(global_ptr), "n"(SizeInBytes), "r"(src_in_bytes));
|
||||
|
||||
#else
|
||||
using AccessType = Array<uint8_t, SizeInBytes>;
|
||||
|
||||
if (pred_guard) {
|
||||
*static_cast<AccessType *>(smem_ptr) = *static_cast<AccessType const *>(global_ptr);
|
||||
}
|
||||
else {
|
||||
AccessType zeros;
|
||||
zeros.clear();
|
||||
*static_cast<AccessType *>(smem_ptr) = zeros;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization
|
||||
template <
|
||||
/// Size of the access in bytes
|
||||
int SizeInBytes>
|
||||
struct cp_async<SizeInBytes, CacheOperation::Global> {
|
||||
/// Copy
|
||||
CUTLASS_DEVICE
|
||||
cp_async(void *smem_ptr, void const *global_ptr, bool pred_guard = true) {
|
||||
#if CUDA_CP_ASYNC_ACTIVATED
|
||||
|
||||
static_assert(SizeInBytes == 16,
|
||||
"cp.async only supports CacheOperation::Global when access size is 16B.");
|
||||
|
||||
unsigned smem_int_ptr = cutlass_get_smem_pointer(smem_ptr);
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred p;\n"
|
||||
" setp.ne.b32 p, %0, 0;\n"
|
||||
" @p cp.async.cg.shared.global [%1], [%2], %3;\n"
|
||||
"}\n" ::"r"((int)pred_guard),
|
||||
"r"(smem_int_ptr), "l"(global_ptr), "n"(SizeInBytes));
|
||||
|
||||
#else
|
||||
using AccessType = Array<uint8_t, SizeInBytes>;
|
||||
|
||||
if (pred_guard) {
|
||||
*static_cast<AccessType *>(smem_ptr) = *static_cast<AccessType const *>(global_ptr);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
/// Partial specialization
|
||||
template <
|
||||
/// Size of the access in bytes
|
||||
int SizeInBytes>
|
||||
struct cp_async_zfill<SizeInBytes, CacheOperation::Global> {
|
||||
/// Copy with zero fill
|
||||
CUTLASS_DEVICE
|
||||
cp_async_zfill(void *smem_ptr, void const *global_ptr, bool pred_guard = true) {
|
||||
#if CUDA_CP_ASYNC_ACTIVATED
|
||||
|
||||
static_assert(SizeInBytes == 16,
|
||||
"cp.async only supports CacheOperation::Global when access size is 16B.");
|
||||
|
||||
unsigned smem_int_ptr = cutlass_get_smem_pointer(smem_ptr);
|
||||
int src_in_bytes = (pred_guard ? SizeInBytes : 0);
|
||||
|
||||
asm volatile(
|
||||
"cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
|
||||
"l"(global_ptr), "n"(SizeInBytes), "r"(src_in_bytes));
|
||||
|
||||
#else
|
||||
using AccessType = Array<uint8_t, SizeInBytes>;
|
||||
|
||||
if (pred_guard) {
|
||||
*static_cast<AccessType *>(smem_ptr) = *static_cast<AccessType const *>(global_ptr);
|
||||
}
|
||||
else {
|
||||
AccessType zeros;
|
||||
zeros.clear();
|
||||
*static_cast<AccessType *>(smem_ptr) = zeros;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Establishes an ordering w.r.t previously issued cp.async instructions. Does not block.
|
||||
CUTLASS_DEVICE
|
||||
void cp_async_fence() {
|
||||
#if CUDA_CP_ASYNC_ACTIVATED
|
||||
asm volatile("cp.async.commit_group;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Blocks until all but <N> previous cp.async.commit_group operations have committed.
|
||||
template <int N>
|
||||
CUTLASS_DEVICE void cp_async_wait() {
|
||||
#if CUDA_CP_ASYNC_ACTIVATED
|
||||
asm volatile("cp.async.wait_group %0;\n" ::"n"(N));
|
||||
#endif
|
||||
}
|
||||
|
||||
/// Blocks until all previous cp.async.commit_group operations have committed.
|
||||
template <>
|
||||
CUTLASS_DEVICE void cp_async_wait<0>() {
|
||||
#if CUDA_CP_ASYNC_ACTIVATED
|
||||
asm volatile("cp.async.wait_all;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace arch
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -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:
|
||||
@@ -51,11 +51,26 @@ struct OpMultiplyAddSaturate;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tag indicating the input is converted to a narrower type (BF16)
|
||||
struct OpMultiplyAddFastBF16;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tag indicating the input is converted to a narrower type (F16)
|
||||
struct OpMultiplyAddFastF16;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tag indicating the complex multiply-add operation
|
||||
struct OpMultiplyAddComplex;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tag indicating the gaussian complex multiply-add operation
|
||||
struct OpMultiplyAddGaussianComplex;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tag indicating the inner product is defined by (XOR, POPC)
|
||||
struct OpXorPopc;
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -120,8 +120,7 @@ struct Wmma<
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// WMMA template structure defines nvcuda::wmma::fragments and static assert for
|
||||
// wmma native instruction sizes supported for cutlass::uint1b_t (experimental::b1)
|
||||
// (nvcuda::wmma targetting SASS instruction BMMA)
|
||||
// wmma native instruction sizes supported for cutlass::uint1b_t (experimental::b1).
|
||||
//
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <
|
||||
|
||||
@@ -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:
|
||||
@@ -167,7 +167,7 @@ public:
|
||||
class const_iterator {
|
||||
|
||||
/// Pointer to object
|
||||
T *ptr_;
|
||||
const T *ptr_;
|
||||
|
||||
public:
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -0,0 +1,461 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*!
|
||||
\file
|
||||
\brief Defines a proxy class for storing non-standard 16-bit floating point values with
|
||||
8 bits of exponent and 7 bit of mantissa.
|
||||
*/
|
||||
#pragma once
|
||||
|
||||
#if !defined(__CUDACC_RTC__)
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
#include <cstdint>
|
||||
#endif
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
namespace cutlass {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Floating-point type with 8 bits of exponent and 7 bits of mantissa.
|
||||
struct alignas(2) bfloat16_t {
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Storage type
|
||||
uint16_t storage;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructs from an unsigned short
|
||||
CUTLASS_HOST_DEVICE
|
||||
static bfloat16_t bitcast(uint16_t x) {
|
||||
bfloat16_t h;
|
||||
h.storage = x;
|
||||
return h;
|
||||
}
|
||||
|
||||
/// Default constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t() { }
|
||||
|
||||
/// Floating-point conversion - round toward nearest
|
||||
CUTLASS_HOST_DEVICE
|
||||
explicit bfloat16_t(float x) {
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) && (__CUDACC_VER_MAJOR__ >= 11)
|
||||
|
||||
asm("cvt.rn.bf16.f32 %0, %1;\n" : "=h"(storage) : "f"(x));
|
||||
|
||||
#else
|
||||
uint32_t bits = reinterpret_cast<uint32_t &>(x);
|
||||
|
||||
if ((bits & 0x7f800000) != 0x7f800000) {
|
||||
|
||||
bool mantissa_bit = ((bits & (1 << 16)) != 0);
|
||||
bool round_bit = ((bits & (1 << 15)) != 0);
|
||||
bool sticky_bit = ((bits & ((1 << 15) - 1)) != 0);
|
||||
|
||||
if ((round_bit && sticky_bit) || (round_bit && mantissa_bit)) {
|
||||
bits += uint32_t(1 << 16);
|
||||
}
|
||||
}
|
||||
else if (bits & ~0xff800000) {
|
||||
bits = 0x7fffffff;
|
||||
}
|
||||
|
||||
storage = uint16_t((bits >> 16) & 0xffff);
|
||||
#endif
|
||||
}
|
||||
|
||||
/// Floating-point conversion - round toward nearest
|
||||
CUTLASS_HOST_DEVICE
|
||||
explicit bfloat16_t(double x): bfloat16_t(float(x)) {
|
||||
|
||||
}
|
||||
|
||||
/// Integer conversion - round toward nearest
|
||||
CUTLASS_HOST_DEVICE
|
||||
explicit bfloat16_t(int x) {
|
||||
float flt = static_cast<float>(x);
|
||||
storage = uint16_t(reinterpret_cast<uint32_t const &>(flt) >> 16);
|
||||
}
|
||||
|
||||
/// Converts to float
|
||||
CUTLASS_HOST_DEVICE
|
||||
operator float() const {
|
||||
unsigned bits = (unsigned(storage) << 16);
|
||||
return reinterpret_cast<float const &>(bits);
|
||||
}
|
||||
|
||||
/// Converts to float
|
||||
CUTLASS_HOST_DEVICE
|
||||
operator double() const {
|
||||
return double(float(*this));
|
||||
}
|
||||
|
||||
/// Converts to int
|
||||
CUTLASS_HOST_DEVICE
|
||||
explicit operator int() const {
|
||||
return int(float(*this));
|
||||
}
|
||||
|
||||
/// Casts to bool
|
||||
CUTLASS_HOST_DEVICE
|
||||
operator bool() const {
|
||||
return (float(*this) != 0.0f);
|
||||
}
|
||||
|
||||
/// Obtains raw bits
|
||||
CUTLASS_HOST_DEVICE
|
||||
uint16_t raw() const {
|
||||
return storage;
|
||||
}
|
||||
/// Returns the sign bit
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool signbit() const {
|
||||
return ((raw() & 0x8000) != 0);
|
||||
}
|
||||
|
||||
/// Returns the biased exponent
|
||||
CUTLASS_HOST_DEVICE
|
||||
int exponent_biased() const {
|
||||
return int((raw() >> 7) & 0x0ff);
|
||||
}
|
||||
|
||||
/// Returns the unbiased exponent
|
||||
CUTLASS_HOST_DEVICE
|
||||
int exponent() const {
|
||||
return exponent_biased() - 127;
|
||||
}
|
||||
|
||||
/// Returns the mantissa
|
||||
CUTLASS_HOST_DEVICE
|
||||
int mantissa() const {
|
||||
return int(raw() & 0x7f);
|
||||
}
|
||||
};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool signbit(cutlass::bfloat16_t const& h) {
|
||||
return h.signbit();
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
cutlass::bfloat16_t abs(cutlass::bfloat16_t const& h) {
|
||||
return cutlass::bfloat16_t::bitcast(h.raw() & 0x7fffffff);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool isnan(cutlass::bfloat16_t const& h) {
|
||||
return (h.exponent_biased() == 0x0ff) && h.mantissa();
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool isfinite(cutlass::bfloat16_t const& h) {
|
||||
return (h.exponent_biased() != 0x0ff);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
cutlass::bfloat16_t nan_bf16(const char*) {
|
||||
// NVIDIA canonical NaN
|
||||
return cutlass::bfloat16_t::bitcast(0x7fff);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool isinf(cutlass::bfloat16_t const& h) {
|
||||
return (h.exponent_biased() == 0x0ff) && !h.mantissa();
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool isnormal(cutlass::bfloat16_t const& h) {
|
||||
return h.exponent_biased() && h.exponent_biased() != 0x0ff;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
int fpclassify(cutlass::bfloat16_t const& h) {
|
||||
int exp = h.exponent_biased();
|
||||
int mantissa = h.mantissa();
|
||||
if (exp == 0x0ff) {
|
||||
if (mantissa) {
|
||||
return FP_NAN;
|
||||
}
|
||||
else {
|
||||
return FP_INFINITE;
|
||||
}
|
||||
}
|
||||
else if (!exp) {
|
||||
if (mantissa) {
|
||||
return FP_SUBNORMAL;
|
||||
}
|
||||
else {
|
||||
return FP_ZERO;
|
||||
}
|
||||
}
|
||||
return FP_NORMAL;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
cutlass::bfloat16_t sqrt(cutlass::bfloat16_t const& h) {
|
||||
#if defined(__CUDACC_RTC__)
|
||||
return cutlass::bfloat16_t(sqrtf(float(h)));
|
||||
#else
|
||||
return cutlass::bfloat16_t(std::sqrt(float(h)));
|
||||
#endif
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t copysign(bfloat16_t const& a, bfloat16_t const& b) {
|
||||
|
||||
uint16_t a_mag = (reinterpret_cast<uint16_t const &>(a) & 0x7fff);
|
||||
uint16_t b_sign = (reinterpret_cast<uint16_t const &>(b) & 0x8000);
|
||||
uint16_t result = (a_mag | b_sign);
|
||||
|
||||
return reinterpret_cast<bfloat16_t const &>(result);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Standard Library operations and definitions
|
||||
//
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace std {
|
||||
|
||||
#if !defined(__CUDACC_RTC__)
|
||||
/// Numeric limits
|
||||
template <>
|
||||
struct numeric_limits<cutlass::bfloat16_t> {
|
||||
static bool const is_specialized = true;
|
||||
static bool const is_signed = true;
|
||||
static bool const is_integer = false;
|
||||
static bool const is_exact = false;
|
||||
static bool const has_infinity = true;
|
||||
static bool const has_quiet_NaN = true;
|
||||
static bool const has_signaling_NaN = false;
|
||||
static std::float_denorm_style const has_denorm = std::denorm_present;
|
||||
static bool const has_denorm_loss = true;
|
||||
static std::float_round_style const round_style = std::round_to_nearest;
|
||||
static bool const is_iec559 = false;
|
||||
static bool const is_bounded = true;
|
||||
static bool const is_modulo = false;
|
||||
static int const digits = 7;
|
||||
|
||||
/// Least positive value
|
||||
CUTLASS_HOST_DEVICE
|
||||
static cutlass::bfloat16_t min() { return cutlass::bfloat16_t::bitcast(0x01); }
|
||||
|
||||
/// Minimum finite value
|
||||
CUTLASS_HOST_DEVICE
|
||||
static cutlass::bfloat16_t lowest() { return cutlass::bfloat16_t::bitcast(0xff7f); }
|
||||
|
||||
/// Maximum finite value
|
||||
CUTLASS_HOST_DEVICE
|
||||
static cutlass::bfloat16_t max() { return cutlass::bfloat16_t::bitcast(0x7f7f); }
|
||||
|
||||
/// Returns smallest finite value
|
||||
CUTLASS_HOST_DEVICE
|
||||
static cutlass::bfloat16_t epsilon() { return cutlass::bfloat16_t::bitcast(0x1000); }
|
||||
|
||||
/// Returns smallest finite value
|
||||
CUTLASS_HOST_DEVICE
|
||||
static cutlass::bfloat16_t round_error() { return cutlass::bfloat16_t(0.5f); }
|
||||
|
||||
/// Returns smallest finite value
|
||||
CUTLASS_HOST_DEVICE
|
||||
static cutlass::bfloat16_t infinity() { return cutlass::bfloat16_t::bitcast(0x7f80); }
|
||||
|
||||
/// Returns smallest finite value
|
||||
CUTLASS_HOST_DEVICE
|
||||
static cutlass::bfloat16_t quiet_NaN() { return cutlass::bfloat16_t::bitcast(0x7fff); }
|
||||
|
||||
/// Returns smallest finite value
|
||||
CUTLASS_HOST_DEVICE
|
||||
static cutlass::bfloat16_t signaling_NaN() { return cutlass::bfloat16_t::bitcast(0x7fff); }
|
||||
|
||||
/// Returns smallest finite value
|
||||
CUTLASS_HOST_DEVICE
|
||||
static cutlass::bfloat16_t denorm_min() { return cutlass::bfloat16_t::bitcast(0x1); }
|
||||
};
|
||||
#endif
|
||||
|
||||
} // namespace std
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
//
|
||||
// Arithmetic operators
|
||||
//
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool operator==(bfloat16_t const& lhs, bfloat16_t const& rhs) {
|
||||
return float(lhs) == float(rhs);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool operator!=(bfloat16_t const& lhs, bfloat16_t const& rhs) {
|
||||
return float(lhs) != float(rhs);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool operator<(bfloat16_t const& lhs, bfloat16_t const& rhs) {
|
||||
return float(lhs) < float(rhs);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool operator<=(bfloat16_t const& lhs, bfloat16_t const& rhs) {
|
||||
return float(lhs) <= float(rhs);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool operator>(bfloat16_t const& lhs, bfloat16_t const& rhs) {
|
||||
return float(lhs) > float(rhs);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool operator>=(bfloat16_t const& lhs, bfloat16_t const& rhs) {
|
||||
return float(lhs) >= float(rhs);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t operator+(bfloat16_t const& lhs, bfloat16_t const& rhs) {
|
||||
return bfloat16_t(float(lhs) + float(rhs));
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t operator-(bfloat16_t const& lhs) {
|
||||
return bfloat16_t(-float(lhs));
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t operator-(bfloat16_t const& lhs, bfloat16_t const& rhs) {
|
||||
return bfloat16_t(float(lhs) - float(rhs));
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t operator*(bfloat16_t const& lhs, bfloat16_t const& rhs) {
|
||||
return bfloat16_t(float(lhs) * float(rhs));
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t operator/(bfloat16_t const& lhs, bfloat16_t const& rhs) {
|
||||
return bfloat16_t(float(lhs) / float(rhs));
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t& operator+=(bfloat16_t & lhs, bfloat16_t const& rhs) {
|
||||
lhs = bfloat16_t(float(lhs) + float(rhs));
|
||||
return lhs;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t& operator-=(bfloat16_t & lhs, bfloat16_t const& rhs) {
|
||||
lhs = bfloat16_t(float(lhs) - float(rhs));
|
||||
return lhs;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t& operator*=(bfloat16_t & lhs, bfloat16_t const& rhs) {
|
||||
lhs = bfloat16_t(float(lhs) * float(rhs));
|
||||
return lhs;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t& operator/=(bfloat16_t & lhs, bfloat16_t const& rhs) {
|
||||
lhs = bfloat16_t(float(lhs) / float(rhs));
|
||||
return lhs;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t& operator++(bfloat16_t & lhs) {
|
||||
float tmp(lhs);
|
||||
++tmp;
|
||||
lhs = bfloat16_t(tmp);
|
||||
return lhs;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t& operator--(bfloat16_t & lhs) {
|
||||
float tmp(lhs);
|
||||
--tmp;
|
||||
lhs = bfloat16_t(tmp);
|
||||
return lhs;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t operator++(bfloat16_t & lhs, int) {
|
||||
bfloat16_t ret(lhs);
|
||||
float tmp(lhs);
|
||||
tmp++;
|
||||
lhs = bfloat16_t(tmp);
|
||||
return ret;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bfloat16_t operator--(bfloat16_t & lhs, int) {
|
||||
bfloat16_t ret(lhs);
|
||||
float tmp(lhs);
|
||||
tmp--;
|
||||
lhs = bfloat16_t(tmp);
|
||||
return ret;
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
//
|
||||
// User-defined literals
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
cutlass::bfloat16_t operator "" _bf16(long double x) {
|
||||
return cutlass::bfloat16_t(float(x));
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
cutlass::bfloat16_t operator "" _bf16(unsigned long long int x) {
|
||||
return cutlass::bfloat16_t(int(x));
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -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:
|
||||
@@ -35,6 +35,9 @@
|
||||
#include "cutlass/half.h"
|
||||
#include "cutlass/real.h"
|
||||
|
||||
#include "cutlass/bfloat16.h"
|
||||
#include "cutlass/tfloat32.h"
|
||||
|
||||
#if !defined(__CUDACC_RTC__)
|
||||
#include <iosfwd>
|
||||
#endif
|
||||
@@ -370,6 +373,15 @@ template <typename T>
|
||||
CUTLASS_HOST_DEVICE complex<T> conj(complex<T> const &z) {
|
||||
return complex<T>(real(z), -imag(z));
|
||||
}
|
||||
/// Indentity transform for non-complex types
|
||||
template <typename T>
|
||||
CUTLASS_HOST_DEVICE T conj(T const &z) {
|
||||
static_assert( !std::is_same<T, cuComplex>::value &&
|
||||
!std::is_same<T, cuDoubleComplex>::value &&
|
||||
!std::is_same<T, cutlass::complex<double>>::value &&
|
||||
!std::is_same<T, cutlass::complex<float>>::value, "May not be a complex data type");
|
||||
return z;
|
||||
}
|
||||
|
||||
/// Projects the complex number z onto the Riemann sphere
|
||||
template <typename T>
|
||||
@@ -429,6 +441,7 @@ template <typename T>
|
||||
struct RealType< complex<T> > {
|
||||
using Type = T;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
static complex<T> from_real(double x) {
|
||||
return complex<T>(static_cast<T>(x));
|
||||
}
|
||||
|
||||
+25
-1
@@ -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:
|
||||
@@ -360,6 +360,29 @@ public:
|
||||
|
||||
namespace cutlass {
|
||||
|
||||
|
||||
/// Scalar multiplication
|
||||
template <int Rank, typename Index>
|
||||
CUTLASS_HOST_DEVICE
|
||||
Coord<Rank, Index> operator*(Index s, Coord<Rank, Index> coord) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < Rank; ++i) {
|
||||
coord[i] *= s;
|
||||
}
|
||||
return coord;
|
||||
}
|
||||
|
||||
/// Scalar multiplication
|
||||
template <int Rank, typename Index>
|
||||
CUTLASS_HOST_DEVICE
|
||||
Coord<Rank, Index> operator*(Coord<Rank, Index> coord, Index s) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < Rank; ++i) {
|
||||
coord[i] *= s;
|
||||
}
|
||||
return coord;
|
||||
}
|
||||
|
||||
/// Scalar division
|
||||
template <int Rank, typename Index>
|
||||
CUTLASS_HOST_DEVICE
|
||||
@@ -419,3 +442,4 @@ Coord<4> make_Coord(int _0, int _1, int _2, int _3) {
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass
|
||||
|
||||
|
||||
@@ -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:
|
||||
@@ -33,9 +33,14 @@
|
||||
|
||||
#include "cutlass/coord.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/layout/pitch_linear.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
namespace cutlass {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// stream operators for cutlass namespace //
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <int Rank>
|
||||
@@ -47,8 +52,6 @@ std::ostream& operator<<(std::ostream& out, Coord<Rank> const& coord) {
|
||||
return out;
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
inline
|
||||
std::istream & operator>>(std::istream &stream, half_t &x) {
|
||||
float tmp;
|
||||
@@ -62,6 +65,16 @@ std::ostream & operator<<(std::ostream &out, half_t const &x) {
|
||||
return out << float(x);
|
||||
}
|
||||
|
||||
inline
|
||||
std::ostream & operator<<(std::ostream &out, bfloat16_t const &x) {
|
||||
return out << float(x);
|
||||
}
|
||||
|
||||
inline
|
||||
std::ostream & operator<<(std::ostream &out, tfloat32_t const &x) {
|
||||
return out << float(x);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Helper to enable formatted printing of CUTLASS scalar types to an ostream
|
||||
@@ -98,7 +111,54 @@ inline std::ostream &operator<<(std::ostream &out, ScalarIO<uint8_t> const &scal
|
||||
return out << unsigned(scalar.value);
|
||||
}
|
||||
|
||||
|
||||
/// Default printing to ostream for MatrixShape
|
||||
template <int Row, int Column>
|
||||
inline
|
||||
std::ostream & operator<<(std::ostream &out, cutlass::MatrixShape<Row, Column> const &matrix_shape) {
|
||||
out << "cutlass::MatrixShape::(kRow, kColumn) {"
|
||||
<< cutlass::MatrixShape<Row,Column>::kRow <<","
|
||||
<< cutlass::MatrixShape<Row,Column>::kColumn <<"}";
|
||||
return out;
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// stream operators for cutlass::gemm namespace //
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
namespace gemm {
|
||||
|
||||
/// Default printing to ostream for GemmShape
|
||||
template <int M, int N, int K>
|
||||
inline
|
||||
std::ostream & operator<<(std::ostream &out, cutlass::gemm::GemmShape<M,N,K> const &gemm_shape) {
|
||||
out << "cutlass::GemmShape::(kM, kN, kK) {"
|
||||
<< cutlass::gemm::GemmShape<M,N,K>::kM <<","
|
||||
<< cutlass::gemm::GemmShape<M,N,K>::kN <<","
|
||||
<< cutlass::gemm::GemmShape<M,N,K>::kK << "}";
|
||||
return out;
|
||||
}
|
||||
|
||||
} //namespace gemm
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// stream operators for cutlass::layout namespace //
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
namespace layout {
|
||||
|
||||
/// Default printing to ostream for PitchLinearShape
|
||||
template < int Contiguous, int Strided>
|
||||
inline
|
||||
std::ostream & operator<<(std::ostream &out, cutlass::layout::PitchLinearShape<Contiguous, Strided> const &pitch_linear_shape) {
|
||||
out << "cutlass::layout::PitchLinearShape::(kContiguous, kStrided) {"
|
||||
<< cutlass::layout::PitchLinearShape<Contiguous,Strided>::kContiguous <<","
|
||||
<< cutlass::layout::PitchLinearShape<Contiguous,Strided>::kStrided <<"}";
|
||||
return out;
|
||||
}
|
||||
|
||||
} //namespace layout
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -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:
|
||||
@@ -55,6 +55,8 @@ enum class Status {
|
||||
kErrorNotSupported, ///< Operation is not supported on current device.
|
||||
kErrorWorkspaceNull, ///< The given workspace is null when it is required to be non-null.
|
||||
kErrorInternal, ///< An error within CUTLASS occurred.
|
||||
kErrorArchMismatch, ///< CUTLASS runs on a device that it was not compiled for.
|
||||
kErrorInsufficientDriver, ///< CUTLASS runs with a driver that is too old.
|
||||
kInvalid ///< Status is unspecified.
|
||||
};
|
||||
|
||||
@@ -78,6 +80,10 @@ static char const* cutlassGetStatusString(cutlass::Status status) {
|
||||
return "Error Workspace Null";
|
||||
case cutlass::Status::kErrorInternal:
|
||||
return "Error Internal";
|
||||
case cutlass::Status::kErrorInsufficientDriver:
|
||||
return "Error Insufficient Driver";
|
||||
case cutlass::Status::kErrorArchMismatch:
|
||||
return "Erroor Architecture Mismatch";
|
||||
case cutlass::Status::kInvalid: break;
|
||||
}
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief This extends the contents of cutlass/functional.h with frequently used activation functions.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
|
||||
#include "cutlass/complex.h"
|
||||
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/half.h"
|
||||
#include "cutlass/functional.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace thread {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// ReLu operator - propagates NaNs
|
||||
template <typename T>
|
||||
struct ReLu {
|
||||
CUTLASS_HOST_DEVICE
|
||||
T operator()(T const & threshold, T const &value) const {
|
||||
if (value < threshold) {
|
||||
value = threshold;
|
||||
}
|
||||
return value;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N>
|
||||
struct ReLu<Array<T, N>> {
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<T, N> operator()(T const & threshold, Array<T, N> const &frag) const {
|
||||
Array<T, N> result;
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N; ++i) {
|
||||
T value = frag[i];
|
||||
if (value < threshold) {
|
||||
value = threshold;
|
||||
}
|
||||
result[i] = value;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
// Sigmoid operator
|
||||
template <typename T>
|
||||
struct Sigmoid {
|
||||
CUTLASS_HOST_DEVICE
|
||||
T operator()(T const &scalar) const {
|
||||
return T(1) / (T(1) + exp(-scalar));
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Sigmoid<float> {
|
||||
CUTLASS_HOST_DEVICE
|
||||
float operator()(float const &scalar) const {
|
||||
return 1.0f / (1.0f + expf(-scalar));
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N>
|
||||
struct Sigmoid<Array<T, N> > {
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<T, N> operator()(Array<T, N> const &rhs) const {
|
||||
Array<T, N> y;
|
||||
Sigmoid<T> sigmoid_op;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < int(rhs.size()); ++i) {
|
||||
y[i] = sigmoid_op(rhs[i]);
|
||||
}
|
||||
|
||||
return y;
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace thread
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -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:
|
||||
@@ -101,7 +101,7 @@ public:
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentOutput const &source,
|
||||
FragmentOutput const &source = FragmentOutput(),
|
||||
ElementCompute uniform = ElementCompute(0)) const {
|
||||
|
||||
// Convert to destination numeric type
|
||||
|
||||
@@ -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:
|
||||
@@ -165,6 +165,28 @@ public:
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
FragmentAccumulator const &accumulator) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
ComputeFragment intermediate;
|
||||
multiplies<ComputeFragment> mul_accumulator;
|
||||
|
||||
intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -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:
|
||||
@@ -178,6 +178,40 @@ public:
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
FragmentAccumulator const &accumulator) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<ComputeFragment> mul_accumulator;
|
||||
|
||||
minimum<ComputeFragment> min_accumulator;
|
||||
maximum<ComputeFragment> max_accumulator;
|
||||
|
||||
intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
/// Clamping constant value
|
||||
ElementCompute const kClamp =
|
||||
ElementCompute((1U << (sizeof_bits<ElementOutput>::value - 1)) - 1);
|
||||
|
||||
intermediate = max_accumulator(intermediate, -kClamp - ElementCompute(1));
|
||||
intermediate = min_accumulator(intermediate, kClamp);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -278,7 +312,7 @@ public:
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
@@ -316,6 +350,37 @@ public:
|
||||
|
||||
return destination_converter(scaled_accumulator);
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(FragmentAccumulator const &accumulator) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Compute linear scaling in floating point
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<ComputeFragment> mul_add_accumulator;
|
||||
|
||||
// Float min-max
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
// Convert floats back to INT
|
||||
FragmentAccumulator scaled_accumulator;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
scaled_accumulator[i] = static_cast<int>(intermediate[i]);
|
||||
}
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, int, kCount, Round> destination_converter;
|
||||
|
||||
return destination_converter(scaled_accumulator);
|
||||
}
|
||||
};
|
||||
|
||||
#endif // Conditional guards to enable partial specialization for packed integers
|
||||
@@ -410,7 +475,7 @@ class FastLinearCombinationClamp {
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(FragmentAccumulator const &accumulator,
|
||||
@@ -453,6 +518,41 @@ class FastLinearCombinationClamp {
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(FragmentAccumulator const &accumulator) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
FastNumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
|
||||
accumulator_converter;
|
||||
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Compute linear scaling in floating point
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<ComputeFragment> mul_accumulator;
|
||||
|
||||
minimum<ComputeFragment> min_accumulator;
|
||||
maximum<ComputeFragment> max_accumulator;
|
||||
|
||||
// Float min-max
|
||||
intermediate = mul_accumulator(alpha_, converted_accumulator);
|
||||
|
||||
/// Clamping constant value
|
||||
ElementCompute const kClamp =
|
||||
ElementCompute(1 << (sizeof_bits<ElementOutput>::value - 1));
|
||||
|
||||
intermediate = max_accumulator(intermediate, -kClamp);
|
||||
intermediate = min_accumulator(intermediate, kClamp - ElementCompute(1));
|
||||
|
||||
// Convert to destination numeric type
|
||||
FastNumericArrayConverter<ElementOutput, ElementCompute, kCount, Round>
|
||||
destination_converter;
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -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:
|
||||
@@ -185,6 +185,39 @@ public:
|
||||
destination_converter(intermediate.real),
|
||||
destination_converter(intermediate.imag));
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
FragmentAccumulator const &accumulator) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
ComputeFragment converted_accumulator(
|
||||
accumulator_converter(accumulator.real),
|
||||
accumulator_converter(accumulator.imag));
|
||||
|
||||
// Perform binary operations
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<Array<ElementCompute, kCount> > mul_op;
|
||||
multiply_add<Array<ElementCompute, kCount> > mul_add_op;
|
||||
|
||||
// complex multiply-add: I = alpha * AB + I
|
||||
intermediate.real = mul_add_op(alpha_.real(), converted_accumulator.real);
|
||||
intermediate.imag = mul_add_op(alpha_.real(), converted_accumulator.imag);
|
||||
|
||||
intermediate.real = mul_add_op(-alpha_.imag(), converted_accumulator.imag, intermediate.real);
|
||||
intermediate.imag = mul_add_op( alpha_.imag(), converted_accumulator.real, intermediate.imag);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
|
||||
return FragmentOutput(
|
||||
destination_converter(intermediate.real),
|
||||
destination_converter(intermediate.imag));
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -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:
|
||||
@@ -23,8 +23,7 @@
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief Functor performing linear combination operations used by epilogues. Values are clamped before
|
||||
converting to the output element type.
|
||||
\brief Functor performing linear combination with a maximum operation used by epilogues.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
@@ -34,6 +33,7 @@
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/functional.h"
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -43,8 +43,7 @@ namespace thread {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Applies a linear combination operator to an array of elements then clamps the output before
|
||||
/// converting to the output element type.
|
||||
/// Applies a linear combination operator to an array of elements.
|
||||
///
|
||||
/// D = alpha * accumulator + beta * source + uniform
|
||||
///
|
||||
@@ -75,10 +74,10 @@ public:
|
||||
|
||||
ElementCompute alpha; ///< scales accumulators
|
||||
ElementCompute beta; ///< scales source tensor
|
||||
ElementCompute threshold; ///< Relu threshold
|
||||
ElementCompute threshold; ///< minimum value that is output
|
||||
ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
|
||||
ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
|
||||
|
||||
ElementCompute const *threshold_ptr; ///< pointer to threshold scalar - if not null, loads from memory
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
@@ -87,16 +86,17 @@ public:
|
||||
Params():
|
||||
alpha(ElementCompute(1)),
|
||||
beta(ElementCompute(0)),
|
||||
threshold(ElementCompute(0)),
|
||||
threshold(ElementCompute(0)),
|
||||
alpha_ptr(nullptr),
|
||||
beta_ptr(nullptr) { }
|
||||
beta_ptr(nullptr),
|
||||
threshold_ptr(nullptr) { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute alpha,
|
||||
ElementCompute beta,
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr) {
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr), threshold_ptr(nullptr) {
|
||||
|
||||
}
|
||||
|
||||
@@ -104,8 +104,8 @@ public:
|
||||
Params(
|
||||
ElementCompute const *alpha_ptr,
|
||||
ElementCompute const *beta_ptr,
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(0), beta(0), threshold(threshold), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
|
||||
ElementCompute const *threshold_ptr = nullptr
|
||||
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr), threshold_ptr(threshold_ptr) {
|
||||
|
||||
}
|
||||
};
|
||||
@@ -128,7 +128,7 @@ public:
|
||||
|
||||
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
|
||||
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
|
||||
threshold_ = params.threshold;
|
||||
threshold_ = (params.threshold_ptr ? *params.threshold_ptr : params.threshold);
|
||||
}
|
||||
|
||||
/// Returns true if source is needed
|
||||
@@ -144,13 +144,12 @@ public:
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentOutput const &source,
|
||||
ElementCompute uniform = ElementCompute(0)) const {
|
||||
FragmentOutput const &source) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementOutput, kCount, Round> source_converter;
|
||||
@@ -160,18 +159,44 @@ public:
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<ComputeFragment> mul_add_source;
|
||||
multiply_add<ComputeFragment> mul_add_accumulator;
|
||||
|
||||
maximum<ComputeFragment> max_accumulator;
|
||||
ReLu<ComputeFragment> relu;
|
||||
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
|
||||
|
||||
intermediate = max_accumulator(intermediate, threshold_);
|
||||
|
||||
// Compute threshold optionally
|
||||
intermediate = relu(threshold_, intermediate);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
FragmentAccumulator const &accumulator) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<ComputeFragment> mul_accumulator;
|
||||
ReLu<ComputeFragment> relu;
|
||||
|
||||
intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
// Compute threshold optionally
|
||||
intermediate = relu(threshold_, intermediate);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
@@ -180,24 +205,24 @@ public:
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Conditional guards to enable partial specialization for packed integers
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 720) && \
|
||||
((__CUDACC_VER_MAJOR__ > 10) || \
|
||||
((__CUDACC_VER_MAJOR__ >= 10) && (__CUDACC_VER_MINOR__ >= 2)))
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 720) && ((__CUDACC_VER_MAJOR__ > 10) || ((__CUDACC_VER_MAJOR__ >= 10) && (__CUDACC_VER_MINOR__ >= 2)))
|
||||
|
||||
/// Applies a linear combination operator to an array of elements then clamps the output before
|
||||
/// converting to the output element type.
|
||||
/// Applies a linear combination operator to an array of elements.
|
||||
///
|
||||
/// D = alpha * accumulator + beta * source + uniform
|
||||
///
|
||||
/// Special handling for int types
|
||||
|
||||
template <
|
||||
typename ElementOutput_, ///< Data type used to load and store tensors
|
||||
int Count, ///< Number of elements computed per operation
|
||||
FloatRoundStyle Round
|
||||
>
|
||||
class LinearCombinationRelu<ElementOutput_, Count, int, float, Round> {
|
||||
class LinearCombinationRelu <ElementOutput_, Count, int, float, Round> {
|
||||
public:
|
||||
|
||||
using ElementOutput = ElementOutput_;
|
||||
@@ -217,10 +242,10 @@ public:
|
||||
|
||||
ElementCompute alpha; ///< scales accumulators
|
||||
ElementCompute beta; ///< scales source tensor
|
||||
ElementCompute threshold; ///< Relu threshold
|
||||
ElementCompute threshold; ///< minimum value that is output
|
||||
ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
|
||||
ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
|
||||
|
||||
ElementCompute const *threshold_ptr; ///< pointer to threshold scalar - if not null, loads from memory
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
@@ -229,16 +254,17 @@ public:
|
||||
Params():
|
||||
alpha(ElementCompute(1)),
|
||||
beta(ElementCompute(0)),
|
||||
threshold(ElementCompute(0)),
|
||||
threshold(ElementCompute(0)),
|
||||
alpha_ptr(nullptr),
|
||||
beta_ptr(nullptr) { }
|
||||
beta_ptr(nullptr),
|
||||
threshold_ptr(nullptr) { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute alpha,
|
||||
ElementCompute beta,
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr) {
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr), threshold_ptr(nullptr) {
|
||||
|
||||
}
|
||||
|
||||
@@ -246,8 +272,8 @@ public:
|
||||
Params(
|
||||
ElementCompute const *alpha_ptr,
|
||||
ElementCompute const *beta_ptr,
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(0), beta(0), threshold(threshold), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
|
||||
ElementCompute const *threshold_ptr = nullptr
|
||||
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr), threshold_ptr(threshold_ptr) {
|
||||
|
||||
}
|
||||
};
|
||||
@@ -270,7 +296,7 @@ public:
|
||||
|
||||
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
|
||||
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
|
||||
threshold_ = params.threshold;
|
||||
threshold_ = (params.threshold_ptr ? *params.threshold_ptr : params.threshold);
|
||||
}
|
||||
|
||||
/// Returns true if source is needed
|
||||
@@ -286,13 +312,12 @@ public:
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentOutput const &source,
|
||||
ElementCompute uniform = ElementCompute(0)) const {
|
||||
FragmentOutput const &source) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementOutput, kCount, Round> source_converter;
|
||||
@@ -302,21 +327,16 @@ public:
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<ComputeFragment> mul_add_source;
|
||||
multiply_add<ComputeFragment> mul_add_accumulator;
|
||||
|
||||
maximum<ComputeFragment> max_accumulator;
|
||||
ReLu<FragmentAccumulator> relu;
|
||||
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
|
||||
|
||||
// Clamp to theshold
|
||||
intermediate = max_accumulator(intermediate, threshold_);
|
||||
|
||||
// Convert back to accumulator data type
|
||||
// Convert floats back to INT
|
||||
FragmentAccumulator scaled_accumulator;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
@@ -324,8 +344,46 @@ public:
|
||||
scaled_accumulator[i] = static_cast<int>(intermediate[i]);
|
||||
}
|
||||
|
||||
// Convert to destination numeric type and pack
|
||||
NumericArrayConverter<ElementOutput, ElementAccumulator, kCount, Round> destination_converter;
|
||||
// Compute threshold optionally
|
||||
scaled_accumulator = relu(threshold_, scaled_accumulator);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, int, kCount, Round> destination_converter;
|
||||
|
||||
return destination_converter(scaled_accumulator);
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
FragmentAccumulator const &accumulator) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<ComputeFragment> mul_accumulator;
|
||||
ReLu<FragmentAccumulator> relu;
|
||||
|
||||
intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
// Convert floats back to INT
|
||||
FragmentAccumulator scaled_accumulator;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
scaled_accumulator[i] = static_cast<int>(intermediate[i]);
|
||||
}
|
||||
|
||||
// Compute threshold optionally
|
||||
scaled_accumulator = relu(threshold_, scaled_accumulator);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, int, kCount, Round> destination_converter;
|
||||
|
||||
return destination_converter(scaled_accumulator);
|
||||
}
|
||||
@@ -338,3 +396,6 @@ public:
|
||||
} // namespace thread
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
@@ -0,0 +1,206 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief Functor performing linear combination operations used by epilogues.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/functional.h"
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace thread {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Applies a linear combination operator to an array of elements.
|
||||
///
|
||||
/// D = alpha * accumulator + beta * source + uniform
|
||||
///
|
||||
template <
|
||||
typename ElementOutput_, ///< Data type used to load and store tensors
|
||||
int Count, ///< Number of elements computed per operation
|
||||
typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type
|
||||
typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination
|
||||
FloatRoundStyle Round = FloatRoundStyle::round_to_nearest
|
||||
>
|
||||
class LinearCombinationSigmoid {
|
||||
public:
|
||||
|
||||
using ElementOutput = ElementOutput_;
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using ElementCompute = ElementCompute_;
|
||||
|
||||
static int const kCount = Count;
|
||||
|
||||
using FragmentOutput = Array<ElementOutput, kCount>;
|
||||
using FragmentAccumulator = Array<ElementAccumulator, kCount>;
|
||||
using ComputeFragment = Array<ElementCompute, kCount>;
|
||||
|
||||
static FloatRoundStyle const kRound = Round;
|
||||
|
||||
/// Host-constructable parameters structure
|
||||
struct Params {
|
||||
|
||||
ElementCompute alpha; ///< scales accumulators
|
||||
ElementCompute beta; ///< scales source tensor
|
||||
ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
|
||||
ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params():
|
||||
alpha(ElementCompute(1)),
|
||||
beta(ElementCompute(0)),
|
||||
alpha_ptr(nullptr),
|
||||
beta_ptr(nullptr) { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute alpha,
|
||||
ElementCompute beta
|
||||
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) {
|
||||
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute const *alpha_ptr,
|
||||
ElementCompute const *beta_ptr
|
||||
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
ElementCompute alpha_;
|
||||
ElementCompute beta_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs the function object, possibly loading from pointers in host memory
|
||||
CUTLASS_HOST_DEVICE
|
||||
LinearCombinationSigmoid(Params const ¶ms) {
|
||||
|
||||
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
|
||||
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
|
||||
}
|
||||
|
||||
/// Returns true if source is needed
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool is_source_needed() const {
|
||||
return beta_ != ElementCompute(0);
|
||||
}
|
||||
|
||||
/// Functionally required for serial reduction in the epilogue
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_k_partition(int k_partition) {
|
||||
if (k_partition) {
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentOutput const &source) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementOutput, kCount, Round> source_converter;
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
ComputeFragment converted_source = source_converter(source);
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<ComputeFragment> mul_add_source;
|
||||
multiply_add<ComputeFragment> mul_add_accumulator;
|
||||
Sigmoid<ComputeFragment> sigmoid;
|
||||
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
|
||||
|
||||
intermediate = sigmoid(intermediate);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentOutput operator()(
|
||||
FragmentAccumulator const &accumulator) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<ComputeFragment> mul_add_accumulator;
|
||||
Sigmoid<ComputeFragment> sigmoid;
|
||||
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
intermediate = sigmoid(intermediate);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace thread
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -45,6 +45,7 @@
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h"
|
||||
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_complex_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_gaussian_complex_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/tile_iterator_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/default_thread_map_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
@@ -76,6 +77,7 @@ template <
|
||||
/// Elements accessed by inner-most loop of AccumulatorFragmentIterator::load()
|
||||
int ElementsPerAccess,
|
||||
/// Multiply-add operator
|
||||
/// Selects between (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex)
|
||||
typename Operator_ = arch::OpMultiplyAddComplex>
|
||||
struct DefaultEpilogueComplexTensorOp {
|
||||
|
||||
@@ -146,6 +148,91 @@ struct DefaultEpilogueComplexTensorOp {
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/// Partial specialization and defines sensible defaults for epilogues for complex*complex case
|
||||
// 3 real-valued mma operations (Gaussian Complex)
|
||||
// A = (ar + j ai), B = (br +j bi), D = AB
|
||||
// P1 = (ar + ai) * br, P2 = - ar * (br - bi), P3 = ai * (br + bi)
|
||||
// D = dr + j di = (P1 - P3) + j (P1 + P2)
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
template <
|
||||
typename Shape_,
|
||||
typename WarpMmaTensorOp_,
|
||||
int PartitionsK,
|
||||
typename OutputOp_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultEpilogueComplexTensorOp <Shape_, WarpMmaTensorOp_, PartitionsK,
|
||||
OutputOp_, ElementsPerAccess,
|
||||
arch::OpMultiplyAddGaussianComplex> {
|
||||
|
||||
using Shape = Shape_;
|
||||
using WarpMmaTensorOp = WarpMmaTensorOp_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using OutputOp = OutputOp_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
using Operator = arch::OpMultiplyAddGaussianComplex;
|
||||
|
||||
using ElementOutput = typename OutputOp::ElementOutput;
|
||||
using LayoutC = typename WarpMmaTensorOp::LayoutC;
|
||||
using ElementAccumulator = typename WarpMmaTensorOp::ElementC;
|
||||
|
||||
//
|
||||
// Thread map
|
||||
//
|
||||
|
||||
using OutputTileThreadMap = typename cutlass::epilogue::threadblock::DefaultThreadMapTensorOp<
|
||||
Shape,
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
kPartitionsK,
|
||||
ElementOutput,
|
||||
kElementsPerAccess
|
||||
>::Type;
|
||||
|
||||
using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<
|
||||
OutputTileThreadMap,
|
||||
ElementOutput
|
||||
>;
|
||||
|
||||
using AccumulatorFragmentIterator = cutlass::epilogue::warp::FragmentIteratorGaussianComplexTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::ElementC,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC
|
||||
>;
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
ElementAccumulator,
|
||||
LayoutC
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<
|
||||
typename OutputTileThreadMap::CompactedThreadMap,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
/// Hard-coded padding elements added
|
||||
using Padding = cutlass::MatrixShape<0, 0>;
|
||||
|
||||
//
|
||||
// Define the epilogue
|
||||
//
|
||||
using Epilogue = cutlass::epilogue::threadblock::Epilogue<
|
||||
Shape,
|
||||
WarpMmaTensorOp,
|
||||
kPartitionsK,
|
||||
OutputTileIterator,
|
||||
AccumulatorFragmentIterator,
|
||||
WarpTileIterator,
|
||||
SharedLoadIterator,
|
||||
OutputOp,
|
||||
Padding
|
||||
>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
|
||||
@@ -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:
|
||||
@@ -148,6 +148,44 @@ struct DefaultEpiloguePlanarComplex<
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues.
|
||||
template <
|
||||
typename ThreadblockShape_,
|
||||
typename WarpMmaOperator_,
|
||||
int PartitionsK,
|
||||
typename OutputOp_,
|
||||
int ElementsPerAccess
|
||||
>
|
||||
struct DefaultEpiloguePlanarComplex<
|
||||
ThreadblockShape_,
|
||||
WarpMmaOperator_,
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
PartitionsK,
|
||||
OutputOp_,
|
||||
ElementsPerAccess> {
|
||||
|
||||
using RealEpilogue = DefaultEpilogueTensorOp<
|
||||
ThreadblockShape_,
|
||||
WarpMmaOperator_,
|
||||
PartitionsK,
|
||||
OutputOp_,
|
||||
ElementsPerAccess
|
||||
>;
|
||||
|
||||
using Epilogue = EpiloguePlanarComplex<
|
||||
ThreadblockShape_,
|
||||
WarpMmaOperator_,
|
||||
PartitionsK,
|
||||
typename RealEpilogue::OutputTileIterator,
|
||||
typename RealEpilogue::AccumulatorFragmentIterator,
|
||||
typename RealEpilogue::WarpTileIterator,
|
||||
typename RealEpilogue::SharedLoadIterator,
|
||||
OutputOp_,
|
||||
typename RealEpilogue::Padding
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues.
|
||||
|
||||
@@ -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:
|
||||
@@ -39,6 +39,7 @@
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
|
||||
#include "cutlass/epilogue/thread/conversion_op.h"
|
||||
#include "cutlass/epilogue/thread/reduction_op.h"
|
||||
|
||||
|
||||
@@ -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:
|
||||
@@ -39,16 +39,20 @@
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
|
||||
#include "cutlass/epilogue/thread/conversion_op.h"
|
||||
#include "cutlass/epilogue/thread/reduction_op.h"
|
||||
|
||||
#include "cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h"
|
||||
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/fragment_iterator_complex_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/tile_iterator_tensor_op.h"
|
||||
#include "cutlass/epilogue/warp/tile_iterator_tensor_op_mixed.h"
|
||||
#include "cutlass/epilogue/threadblock/default_thread_map_tensor_op.h"
|
||||
#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
|
||||
#include "cutlass/epilogue/threadblock/shared_load_iterator.h"
|
||||
#include "cutlass/epilogue/threadblock/shared_load_iterator_mixed.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/epilogue.h"
|
||||
#include "cutlass/epilogue/threadblock/interleaved_epilogue.h"
|
||||
@@ -61,6 +65,177 @@ namespace threadblock {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace detail {
|
||||
|
||||
template <
|
||||
typename ElementOutput,
|
||||
typename ElementAccumulator,
|
||||
int ElementsPerAccess,
|
||||
typename ThreadblockShape,
|
||||
typename WarpShape,
|
||||
typename InstructionShape,
|
||||
typename ThreadMap
|
||||
>
|
||||
struct DefaultIteratorsTensorOp {
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOp<
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
ElementAccumulator,
|
||||
layout::RowMajor
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<
|
||||
ThreadMap,
|
||||
ElementAccumulator
|
||||
>;
|
||||
};
|
||||
|
||||
/// Partial specialization for half <= float x 8 epilogues avoids shared memory bank conflicts.
|
||||
template <
|
||||
typename ThreadblockShape,
|
||||
typename WarpShape,
|
||||
typename InstructionShape,
|
||||
typename ThreadMap
|
||||
>
|
||||
struct DefaultIteratorsTensorOp<
|
||||
half_t,
|
||||
float,
|
||||
8,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
ThreadMap> {
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOpMixed<
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
float,
|
||||
32,
|
||||
16,
|
||||
8,
|
||||
8
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIteratorMixed<
|
||||
ThreadMap,
|
||||
float,
|
||||
32,
|
||||
16,
|
||||
8,
|
||||
8
|
||||
>;
|
||||
};
|
||||
|
||||
/// Partial specialization for int8_t x 16 <= int32_t x 16 epilogues avoids shared memory bank conflicts.
|
||||
template <
|
||||
int K,
|
||||
typename InstructionShape,
|
||||
typename ThreadMap
|
||||
>
|
||||
struct DefaultIteratorsTensorOp<
|
||||
int8_t,
|
||||
int32_t,
|
||||
16,
|
||||
gemm::GemmShape<128, 128, K>,
|
||||
gemm::GemmShape<64, 64, K>,
|
||||
InstructionShape,
|
||||
ThreadMap> {
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOpMixed<
|
||||
gemm::GemmShape<64, 64, K>,
|
||||
InstructionShape,
|
||||
int32_t,
|
||||
32,
|
||||
8,
|
||||
16,
|
||||
8
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIteratorMixed<
|
||||
ThreadMap,
|
||||
int32_t,
|
||||
32,
|
||||
8,
|
||||
16,
|
||||
8
|
||||
>;
|
||||
};
|
||||
|
||||
/// Partial specialization for int8_t x 8 <= int32_t x 8 epilogues avoids shared memory bank conflicts.
|
||||
template <
|
||||
int K,
|
||||
typename InstructionShape,
|
||||
typename ThreadMap
|
||||
>
|
||||
struct DefaultIteratorsTensorOp<
|
||||
int8_t,
|
||||
int32_t,
|
||||
8,
|
||||
gemm::GemmShape<128, 64, K>,
|
||||
gemm::GemmShape<64, 32, K>,
|
||||
InstructionShape,
|
||||
ThreadMap> {
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOpMixed<
|
||||
gemm::GemmShape<64, 32, K>,
|
||||
InstructionShape,
|
||||
int32_t,
|
||||
32,
|
||||
8,
|
||||
8,
|
||||
8
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIteratorMixed<
|
||||
ThreadMap,
|
||||
int32_t,
|
||||
32,
|
||||
8,
|
||||
8,
|
||||
8
|
||||
>;
|
||||
};
|
||||
|
||||
/// Partial specialization for int8_t x 8 <= int32_t x 8 epilogues avoids shared memory bank conflicts.
|
||||
template <
|
||||
int K,
|
||||
typename InstructionShape,
|
||||
typename ThreadMap
|
||||
>
|
||||
struct DefaultIteratorsTensorOp<
|
||||
int8_t,
|
||||
int32_t,
|
||||
8,
|
||||
gemm::GemmShape<64, 64, K>,
|
||||
gemm::GemmShape<32, 32, K>,
|
||||
InstructionShape,
|
||||
ThreadMap> {
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOpMixed<
|
||||
gemm::GemmShape<32, 32, K>,
|
||||
InstructionShape,
|
||||
int32_t,
|
||||
32,
|
||||
8,
|
||||
8,
|
||||
8
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIteratorMixed<
|
||||
ThreadMap,
|
||||
int32_t,
|
||||
32,
|
||||
8,
|
||||
8,
|
||||
8
|
||||
>;
|
||||
};
|
||||
|
||||
} // namespace detail
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines sensible defaults for epilogues for TensorOps.
|
||||
template <
|
||||
typename Shape_,
|
||||
@@ -98,25 +273,33 @@ struct DefaultEpilogueTensorOp {
|
||||
ElementOutput
|
||||
>;
|
||||
|
||||
using AccumulatorFragmentIterator = cutlass::epilogue::warp::FragmentIteratorTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::ElementC,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC
|
||||
>;
|
||||
using AccumulatorFragmentIterator = typename std::conditional<is_complex<ElementOutput>::value,
|
||||
cutlass::epilogue::warp::FragmentIteratorComplexTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::ElementC,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC>,
|
||||
cutlass::epilogue::warp::FragmentIteratorTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::ElementC,
|
||||
typename WarpMmaTensorOp::Policy::Operator::FragmentC,
|
||||
LayoutC> >::type;
|
||||
|
||||
using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOp<
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
/// Support several implementations depending on structure of epilogue
|
||||
using DefaultIterators = detail::DefaultIteratorsTensorOp<
|
||||
ElementOutput,
|
||||
ElementAccumulator,
|
||||
LayoutC
|
||||
kElementsPerAccess,
|
||||
Shape,
|
||||
typename WarpMmaTensorOp::Shape,
|
||||
typename WarpMmaTensorOp::Policy::Operator::Shape,
|
||||
typename OutputTileThreadMap::CompactedThreadMap
|
||||
>;
|
||||
|
||||
using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<
|
||||
typename OutputTileThreadMap::CompactedThreadMap,
|
||||
ElementAccumulator
|
||||
>;
|
||||
using WarpTileIterator = typename DefaultIterators::WarpTileIterator;
|
||||
using SharedLoadIterator = typename DefaultIterators::SharedLoadIterator;
|
||||
|
||||
/// Hard-coded padding elements added
|
||||
using Padding = cutlass::MatrixShape<0, 64 / sizeof_bits<ElementAccumulator>::value * 4>;
|
||||
@@ -184,6 +367,7 @@ struct DefaultInterleavedEpilogueTensorOp {
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
@@ -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:
|
||||
@@ -39,6 +39,7 @@
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
|
||||
#include "cutlass/epilogue/thread/conversion_op.h"
|
||||
#include "cutlass/epilogue/thread/reduction_op.h"
|
||||
|
||||
|
||||
@@ -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:
|
||||
@@ -39,6 +39,7 @@
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/linear_combination.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
|
||||
#include "cutlass/epilogue/thread/conversion_op.h"
|
||||
#include "cutlass/epilogue/thread/reduction_op.h"
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -146,54 +146,6 @@ struct DefaultInterleavedThreadMapTensorOp {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Defines the optimal thread map for TensorOp accumulator layouts
|
||||
template <typename ThreadblockShape_, typename WarpShape_, int PartitionsK,
|
||||
typename Element_, int ElementsPerAccess, int InterleavedK>
|
||||
struct DefaultInterleavedConvThreadMapTensorOp {
|
||||
using ThreadblockShape = ThreadblockShape_;
|
||||
using WarpShape = WarpShape_;
|
||||
static int const kPartitionsK = PartitionsK;
|
||||
using Element = Element_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kInterleavedK = InterleavedK;
|
||||
|
||||
//
|
||||
// Definitions
|
||||
//
|
||||
|
||||
struct Detail {
|
||||
/// Tensor Operations fundamentally perform operations on 8 rows
|
||||
static int const kTensorOpRows = 8;
|
||||
static int const kWarpSize = 32;
|
||||
|
||||
static_assert(!(ThreadblockShape::kM % WarpShape::kM) &&
|
||||
!(ThreadblockShape::kM % WarpShape::kM),
|
||||
"Divisibility");
|
||||
|
||||
/// Number of warps
|
||||
using WarpCount =
|
||||
gemm::GemmShape<ThreadblockShape::kM / WarpShape::kM,
|
||||
ThreadblockShape::kN / WarpShape::kN, kPartitionsK>;
|
||||
|
||||
/// Number of participating threads
|
||||
static int const kThreads = WarpCount::kCount * kWarpSize;
|
||||
};
|
||||
|
||||
//
|
||||
// ThreadMap
|
||||
//
|
||||
|
||||
/// ThreadMap to be used by epilogue::MaskedTileIterator satisfying concept
|
||||
/// InterleavedOutputTileThreadMap
|
||||
using Type = InterleavedConvOutputTileThreadMap<
|
||||
MatrixShape<Detail::WarpCount::kM, Detail::WarpCount::kN>,
|
||||
MatrixShape<WarpShape::kM / Detail::kTensorOpRows,
|
||||
WarpShape::kN / InterleavedK>,
|
||||
Detail::kThreads, kElementsPerAccess, sizeof_bits<Element>::value>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -175,15 +175,106 @@ public:
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
OutputTileIterator destination_iterator, ///< Tile iterator for destination
|
||||
AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile
|
||||
OutputTileIterator source_iterator, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
int64_t imag_stride_dest = 0, ///< Arguments required for planar complex case - not used in real-valued case
|
||||
int64_t imag_stride_src = 0) { ///<
|
||||
|
||||
typename OutputTileIterator::Fragment source_fragment;
|
||||
|
||||
OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
|
||||
if (!output_op.is_source_needed()) {
|
||||
source_iterator.clear_mask();
|
||||
compute_source_not_needed_(output_op, destination_iterator, accumulators);
|
||||
}
|
||||
else {
|
||||
compute_source_needed_(output_op, destination_iterator, accumulators, source_iterator);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
|
||||
/// Streams the result to global memory
|
||||
CUTLASS_DEVICE
|
||||
void compute_source_not_needed_(
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
OutputTileIterator destination_iterator, ///< Tile iterator for destination
|
||||
AccumulatorTile const &accumulators) { ///< Complete warp-level accumulator tile
|
||||
|
||||
//
|
||||
// Iterator over warp-level accumulator fragment
|
||||
//
|
||||
|
||||
AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
|
||||
|
||||
//
|
||||
// Iterate over accumulator tile
|
||||
//
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
|
||||
|
||||
//
|
||||
// Convert and store fragment
|
||||
//
|
||||
|
||||
__syncthreads();
|
||||
|
||||
typename AccumulatorFragmentIterator::Fragment accum_fragment;
|
||||
|
||||
accum_fragment_iterator.load(accum_fragment);
|
||||
++accum_fragment_iterator;
|
||||
|
||||
this->warp_tile_iterator_.store(accum_fragment);
|
||||
|
||||
__syncthreads();
|
||||
|
||||
//
|
||||
// Load fragments from shared memory
|
||||
//
|
||||
|
||||
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
|
||||
|
||||
shared_load_iterator_.load(aligned_accum_fragment[0]);
|
||||
|
||||
// If the number of k-slices is > 1 - perform a reduction amongst the k-slices
|
||||
if (kPartitionsK > 1)
|
||||
{
|
||||
plus <typename SharedLoadIterator::Fragment> add_fragments;
|
||||
const int tile_row_offset = Base::SharedStorage::StorageShape::kRow / PartitionsK;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for ( int i = 1; i < kPartitionsK; ++i) {
|
||||
shared_load_iterator_.add_tile_offset({tile_row_offset , 0});
|
||||
shared_load_iterator_.load(aligned_accum_fragment[i]);
|
||||
aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]);
|
||||
}
|
||||
|
||||
shared_load_iterator_.add_tile_offset({-1 * (kPartitionsK-1) * tile_row_offset, 0});
|
||||
}
|
||||
|
||||
//
|
||||
// Compute the output result
|
||||
//
|
||||
|
||||
typename OutputTileIterator::Fragment output_fragment;
|
||||
|
||||
apply_output_operator_source_not_needed_(output_fragment, output_op, aligned_accum_fragment[0]);
|
||||
|
||||
|
||||
//
|
||||
// Store the final result
|
||||
//
|
||||
|
||||
destination_iterator.store(output_fragment);
|
||||
++destination_iterator;
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Streams the result to global memory
|
||||
CUTLASS_DEVICE
|
||||
void compute_source_needed_(
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
OutputTileIterator destination_iterator, ///< Tile iterator for destination
|
||||
AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile
|
||||
OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
|
||||
|
||||
typename OutputTileIterator::Fragment source_fragment;
|
||||
|
||||
source_fragment.clear();
|
||||
|
||||
@@ -265,8 +356,6 @@ public:
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
|
||||
/// Helper to invoke the output functor over each vector of output
|
||||
CUTLASS_DEVICE
|
||||
void apply_output_operator_(
|
||||
@@ -294,6 +383,30 @@ private:
|
||||
output_frag_ptr[i] = output_op(compute_frag_ptr[i], source_frag_ptr[i]);
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper to invoke the output functor over each vector of output
|
||||
CUTLASS_DEVICE
|
||||
void apply_output_operator_source_not_needed_(
|
||||
typename OutputTileIterator::Fragment &output_fragment,
|
||||
OutputOp const &output_op, ///< Output operator
|
||||
typename SharedLoadIterator::Fragment const &aligned_accum_fragment) {
|
||||
|
||||
OutputAccessType *output_frag_ptr =
|
||||
reinterpret_cast<OutputAccessType *>(&output_fragment);
|
||||
|
||||
AccumulatorAccessType const *compute_frag_ptr =
|
||||
reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment);
|
||||
|
||||
int const kOutputOpIterations =
|
||||
OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kOutputOpIterations; ++i) {
|
||||
|
||||
// Call the output operator
|
||||
output_frag_ptr[i] = output_op(compute_frag_ptr[i]);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -490,67 +490,6 @@ struct InterleavedOutputTileThreadMap {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Template metaprogram for partitioning a 4D interleaved layout across warps
|
||||
/// to achieve several performance objectives:
|
||||
///
|
||||
/// - coalesced memory accesses in units of 64 Byte lines
|
||||
/// - minimal address arithmetic
|
||||
/// - minimal predicate calculations
|
||||
///
|
||||
template <typename WarpCount_, typename Iterations_, int Threads,
|
||||
int ElementsPerAccess, int ElementSize>
|
||||
struct InterleavedConvOutputTileThreadMap {
|
||||
using WarpCount = WarpCount_;
|
||||
|
||||
static int const kWarpSize = 32;
|
||||
static int const kThreads = Threads;
|
||||
static int const kWarpCount = kThreads / kWarpSize;
|
||||
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kElementSize = ElementSize;
|
||||
|
||||
//
|
||||
// Metaprogram computation
|
||||
//
|
||||
|
||||
struct Detail {};
|
||||
|
||||
//
|
||||
// Output
|
||||
//
|
||||
|
||||
using Iterations = Iterations_;
|
||||
|
||||
using Delta = MatrixShape<kWarpSize / 4, 4 * kElementsPerAccess>;
|
||||
|
||||
/// Initial offset function
|
||||
CUTLASS_HOST_DEVICE
|
||||
static MatrixCoord initial_offset(int thread_idx) {
|
||||
int warp_idx = thread_idx / kWarpSize;
|
||||
int lane_idx = thread_idx % kWarpSize;
|
||||
|
||||
// Compute warp location
|
||||
MatrixCoord warp_footprint{
|
||||
Delta::kRow * Iterations::kRow,
|
||||
Delta::kColumn * Iterations::kColumn,
|
||||
};
|
||||
|
||||
MatrixCoord warp_offset{warp_idx % WarpCount::kRow,
|
||||
warp_idx / WarpCount::kRow};
|
||||
|
||||
// Compute per-lane offset
|
||||
MatrixCoord thread_offset_in_warp{lane_idx / 4,
|
||||
(lane_idx % 4) * kElementsPerAccess};
|
||||
|
||||
MatrixCoord thread_offset_in_threadblock_tile =
|
||||
warp_footprint * warp_offset + thread_offset_in_warp;
|
||||
|
||||
return thread_offset_in_threadblock_tile;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
@@ -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:
|
||||
@@ -41,7 +41,7 @@
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "cutlass/transform/pitch_linear_thread_map.h"
|
||||
#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
|
||||
|
||||
#include "cutlass/arch/memory.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -306,10 +306,15 @@ public:
|
||||
|
||||
bool guard = row_guard && mask_.predicates[column];
|
||||
|
||||
if (guard) {
|
||||
frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column] =
|
||||
memory_pointer[column * ThreadMap::Delta::kColumn / kElementsPerAccess];
|
||||
}
|
||||
cutlass::arch::global_load<
|
||||
AccessType,
|
||||
sizeof(AccessType)
|
||||
>(
|
||||
frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn +
|
||||
column],
|
||||
(void *)&memory_pointer[column * ThreadMap::Delta::kColumn /
|
||||
kElementsPerAccess],
|
||||
guard);
|
||||
}
|
||||
|
||||
if (row + 1 < ThreadMap::Iterations::kRow) {
|
||||
@@ -365,11 +370,12 @@ public:
|
||||
|
||||
bool guard = row_guard && mask_.predicates[column];
|
||||
|
||||
if (guard) {
|
||||
|
||||
memory_pointer[column * ThreadMap::Delta::kColumn / kElementsPerAccess] =
|
||||
frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column];
|
||||
}
|
||||
cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
|
||||
frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn +
|
||||
column],
|
||||
(void *)&memory_pointer[column * ThreadMap::Delta::kColumn /
|
||||
kElementsPerAccess],
|
||||
guard);
|
||||
}
|
||||
|
||||
if (row + 1 < ThreadMap::Iterations::kRow) {
|
||||
@@ -660,9 +666,13 @@ public:
|
||||
|
||||
bool guard = col_guard && mask_.predicates[iteration_contiguous_];
|
||||
|
||||
if (guard) {
|
||||
*frag_ptr = *memory_pointer;
|
||||
}
|
||||
cutlass::arch::global_load<
|
||||
AccessType,
|
||||
sizeof(AccessType)
|
||||
>(
|
||||
*frag_ptr,
|
||||
(void *)memory_pointer,
|
||||
guard);
|
||||
}
|
||||
|
||||
/// Stores a fragment to memory
|
||||
@@ -678,9 +688,8 @@ public:
|
||||
|
||||
bool guard = col_guard && mask_.predicates[iteration_contiguous_];
|
||||
|
||||
if (guard) {
|
||||
*memory_pointer = *frag_ptr;
|
||||
}
|
||||
cutlass::arch::global_store<AccessType, sizeof(AccessType)>(
|
||||
*frag_ptr, (void *)memory_pointer, guard);
|
||||
}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
@@ -732,6 +741,7 @@ public:
|
||||
}
|
||||
};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
|
||||
@@ -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:
|
||||
@@ -96,6 +96,15 @@ public:
|
||||
ThreadMap::kElementsPerAccess,
|
||||
kAlignment>;
|
||||
|
||||
/// Vector type used for SMEM loads
|
||||
using LoadType = AlignedArray<
|
||||
Element,
|
||||
const_min(128 / sizeof_bits<Element>::value, ThreadMap::kElementsPerAccess),
|
||||
const_min(16, kAlignment)
|
||||
>;
|
||||
|
||||
static int const kLoadsPerAccess = AccessType::kElements / LoadType::kElements;
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
@@ -149,7 +158,6 @@ public:
|
||||
CUTLASS_DEVICE
|
||||
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
|
||||
|
||||
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
|
||||
@@ -169,15 +177,19 @@ public:
|
||||
int frag_row_idx =
|
||||
(row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
|
||||
|
||||
AccessType const *memory_pointer = reinterpret_cast<AccessType const *>(byte_pointer);
|
||||
LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
|
||||
LoadType const *memory_pointer = reinterpret_cast<LoadType const *>(byte_pointer);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
|
||||
|
||||
int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column;
|
||||
|
||||
frag_ptr[frag_idx] =
|
||||
memory_pointer[column * ThreadMap::Delta::kColumn / kElementsPerAccess];
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int v = 0; v < kLoadsPerAccess; ++v) {
|
||||
frag_ptr[frag_idx * kLoadsPerAccess + v] =
|
||||
memory_pointer[(column * ThreadMap::Delta::kColumn / kElementsPerAccess) * kLoadsPerAccess + v];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,559 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief Epilogue for threadblock scoped GEMMs using Tensor Ops optimized for mixed-precision.
|
||||
|
||||
This assumes the shared memory tile is in a permuted layout which avoids bank conflicts on loading.
|
||||
|
||||
When the fragment is loaded into registers, it matches the row-major thread map assumed by
|
||||
the predicated tile iterator writing to global memory.
|
||||
|
||||
The epilogue rearranges the result of a matrix product through shared memory to match canonical
|
||||
tensor layouts in global memory. Epilogues support conversion and reduction operations.
|
||||
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/matrix_shape.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
|
||||
#include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tile iterator used to load output tile from shared memory in epilogue.
|
||||
///
|
||||
/// Satisfies: ReadableTileIterator
|
||||
///
|
||||
template <
|
||||
typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap)
|
||||
typename Element_, ///< Accumulator data type
|
||||
int ElementSizeBits_, ///< Size of accumulator in bits
|
||||
int OutputSizeBits_, ///< Size of output element in bits
|
||||
int ElementsPerAccess, ///< Vector length of output vector
|
||||
int ContiguousLanes ///< Number of lanes in the warp writing to contiguous elements
|
||||
/// in the global memory tensor
|
||||
>
|
||||
class SharedLoadIteratorMixed;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tile iterator used to load output tile from shared memory in epilogue.
|
||||
///
|
||||
/// Satisfies: ReadableTileIterator
|
||||
///
|
||||
template <
|
||||
typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap)
|
||||
typename Element_ ///< Accumulator data type
|
||||
>
|
||||
class SharedLoadIteratorMixed<ThreadMap_, Element_, 32, 16, 8, 8> {
|
||||
public:
|
||||
using ThreadMap = ThreadMap_;
|
||||
using Shape = typename ThreadMap::Shape;
|
||||
|
||||
using Element = Element_;
|
||||
|
||||
using Layout = layout::RowMajor;
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using ConstTensorRef = typename TensorRef::ConstTensorRef;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
using TensorCoord = MatrixCoord;
|
||||
|
||||
static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
|
||||
|
||||
static int const kAlignment = ThreadMap::kElementsPerAccess * sizeof_bits<Element_>::value / 8;
|
||||
|
||||
static int const kThreads = ThreadMap::kThreads;
|
||||
|
||||
/// Fragment object
|
||||
using Fragment = Array<
|
||||
Element,
|
||||
ThreadMap::Iterations::kColumn *
|
||||
ThreadMap::Iterations::kRow *
|
||||
ThreadMap::Iterations::kGroup *
|
||||
ThreadMap::Iterations::kCluster *
|
||||
ThreadMap::kElementsPerAccess>;
|
||||
|
||||
/// Memory access size
|
||||
using AccessType = AlignedArray<
|
||||
Element,
|
||||
ThreadMap::kElementsPerAccess,
|
||||
kAlignment>;
|
||||
|
||||
/// Vector type used for SMEM loads
|
||||
using LoadType = AlignedArray<
|
||||
Element,
|
||||
const_min(128 / sizeof_bits<Element>::value, ThreadMap::kElementsPerAccess),
|
||||
const_min(16, kAlignment)
|
||||
>;
|
||||
|
||||
static int const kLoadsPerAccess = AccessType::kElements / LoadType::kElements;
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Byte-level pointer
|
||||
LoadType const *pointers_[kLoadsPerAccess];
|
||||
|
||||
/// Stride along adjacent rows in units of LoadType
|
||||
int stride_;
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
SharedLoadIteratorMixed(
|
||||
TensorRef ref,
|
||||
int thread_idx
|
||||
):
|
||||
stride_((ref.stride(0) / LoadType::kElements)) {
|
||||
|
||||
TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
|
||||
|
||||
// Initialize pointers
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kLoadsPerAccess; ++i) {
|
||||
pointers_[i] = reinterpret_cast<LoadType const *>(ref.data());
|
||||
|
||||
int col_idx = (thread_offset.column() / kElementsPerAccess) * kLoadsPerAccess;
|
||||
int bank_offset = (col_idx * sizeof(LoadType) / 128) % kLoadsPerAccess;
|
||||
|
||||
col_idx += (bank_offset + i) % kLoadsPerAccess;
|
||||
|
||||
pointers_[i] += thread_offset.row() * stride_ + col_idx;
|
||||
}
|
||||
}
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kLoadsPerAccess; ++i) {
|
||||
pointers_ += pointer_offset / LoadType::kElements;
|
||||
}
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void add_tile_offset(TensorCoord const &offset) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kLoadsPerAccess; ++i) {
|
||||
pointers_[i] += offset.row() * stride_ + offset.column() / LoadType::kElements;
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
|
||||
|
||||
int row_ptr_offset =
|
||||
row * ThreadMap::Delta::kRow * stride_ +
|
||||
group * ThreadMap::Delta::kGroup* stride_ +
|
||||
cluster * ThreadMap::Delta::kCluster * stride_ +
|
||||
pointer_offset / LoadType::kElements;
|
||||
|
||||
int frag_row_idx = (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
|
||||
|
||||
LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
|
||||
|
||||
int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int v = 0; v < kLoadsPerAccess; ++v) {
|
||||
|
||||
int vector_idx = (column * ThreadMap::Delta::kColumn / kElementsPerAccess * kLoadsPerAccess);
|
||||
|
||||
LoadType const *memory_pointer = pointers_[v] + row_ptr_offset;
|
||||
|
||||
frag_ptr[frag_idx * kLoadsPerAccess + v] = memory_pointer[vector_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads a fragment
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
|
||||
load_with_pointer_offset(frag, 0);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for int32_t x 16 => int8_t x 16
|
||||
template <
|
||||
typename ThreadMap_ ///< Thread map (conept: OutputTileThreadMap)
|
||||
>
|
||||
class SharedLoadIteratorMixed<ThreadMap_, int32_t, 32, 8, 16, 8> {
|
||||
public:
|
||||
using ThreadMap = ThreadMap_;
|
||||
using Shape = typename ThreadMap::Shape;
|
||||
|
||||
using Element = int32_t;
|
||||
|
||||
using Layout = layout::RowMajor;
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using ConstTensorRef = typename TensorRef::ConstTensorRef;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
using TensorCoord = MatrixCoord;
|
||||
|
||||
static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
|
||||
|
||||
static int const kAlignment = 16;
|
||||
|
||||
static int const kThreads = ThreadMap::kThreads;
|
||||
|
||||
/// Fragment object
|
||||
using Fragment = Array<
|
||||
Element,
|
||||
ThreadMap::Iterations::kColumn *
|
||||
ThreadMap::Iterations::kRow *
|
||||
ThreadMap::Iterations::kGroup *
|
||||
ThreadMap::Iterations::kCluster *
|
||||
ThreadMap::kElementsPerAccess>;
|
||||
|
||||
/// Memory access size
|
||||
using AccessType = AlignedArray<
|
||||
Element,
|
||||
16,
|
||||
kAlignment>;
|
||||
|
||||
/// Vector type used for SMEM loads
|
||||
using LoadType = AlignedArray<
|
||||
Element,
|
||||
4,
|
||||
16
|
||||
>;
|
||||
|
||||
static int const kLoadsPerAccess = 4;
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Byte-level pointer
|
||||
LoadType const *pointers_[kLoadsPerAccess];
|
||||
|
||||
/// Stride along adjacent rows in units of LoadType
|
||||
int stride_;
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
SharedLoadIteratorMixed(
|
||||
TensorRef ref,
|
||||
int thread_idx
|
||||
):
|
||||
stride_((ref.stride(0) / LoadType::kElements)) {
|
||||
|
||||
TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
|
||||
|
||||
// Initialize pointers
|
||||
LoadType const *base_ptr = reinterpret_cast<LoadType const *>(ref.data()) + thread_offset.row() * stride_;
|
||||
|
||||
int lane_col_idx = thread_offset.column() / 16;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kLoadsPerAccess; ++i) {
|
||||
int lane_offset = (lane_col_idx % 2) * 4 | ((lane_col_idx / 2) * 8) | ((lane_col_idx / 2) ^ i);
|
||||
|
||||
pointers_[i] = base_ptr + lane_offset;
|
||||
}
|
||||
}
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kLoadsPerAccess; ++i) {
|
||||
pointers_[i] += pointer_offset / LoadType::kElements;
|
||||
}
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void add_tile_offset(TensorCoord const &offset) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kLoadsPerAccess; ++i) {
|
||||
pointers_[i] += offset.row() * stride_ + offset.column() / LoadType::kElements;
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
|
||||
|
||||
int row_ptr_offset =
|
||||
row * ThreadMap::Delta::kRow * stride_ +
|
||||
group * ThreadMap::Delta::kGroup* stride_ +
|
||||
cluster * ThreadMap::Delta::kCluster * stride_ +
|
||||
pointer_offset / LoadType::kElements;
|
||||
|
||||
int frag_row_idx = (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
|
||||
|
||||
LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
|
||||
|
||||
int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int v = 0; v < kLoadsPerAccess; ++v) {
|
||||
|
||||
LoadType const *memory_pointer = pointers_[v];
|
||||
|
||||
frag_ptr[frag_idx * kLoadsPerAccess + v] = memory_pointer[row_ptr_offset];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads a fragment
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
|
||||
load_with_pointer_offset(frag, 0);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for int32_t x 8 => int8_t x 8
|
||||
template <
|
||||
typename ThreadMap_ ///< Thread map (conept: OutputTileThreadMap)
|
||||
>
|
||||
class SharedLoadIteratorMixed<ThreadMap_, int32_t, 32, 8, 8, 8> {
|
||||
public:
|
||||
using ThreadMap = ThreadMap_;
|
||||
using Shape = typename ThreadMap::Shape;
|
||||
|
||||
using Element = int32_t;
|
||||
|
||||
using Layout = layout::RowMajor;
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using ConstTensorRef = typename TensorRef::ConstTensorRef;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
using TensorCoord = MatrixCoord;
|
||||
|
||||
static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
|
||||
|
||||
static int const kAlignment = 8;
|
||||
|
||||
static int const kThreads = ThreadMap::kThreads;
|
||||
|
||||
/// Fragment object
|
||||
using Fragment = Array<
|
||||
Element,
|
||||
ThreadMap::Iterations::kColumn *
|
||||
ThreadMap::Iterations::kRow *
|
||||
ThreadMap::Iterations::kGroup *
|
||||
ThreadMap::Iterations::kCluster *
|
||||
ThreadMap::kElementsPerAccess>;
|
||||
|
||||
/// Memory access size
|
||||
using AccessType = AlignedArray<
|
||||
Element,
|
||||
8,
|
||||
kAlignment>;
|
||||
|
||||
/// Vector type used for SMEM loads
|
||||
using LoadType = AlignedArray<
|
||||
Element,
|
||||
4,
|
||||
16
|
||||
>;
|
||||
|
||||
static int const kLoadsPerAccess = 2;
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Byte-level pointer
|
||||
LoadType const *pointers_[kLoadsPerAccess];
|
||||
|
||||
/// Stride along adjacent rows in units of LoadType
|
||||
int stride_;
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_DEVICE
|
||||
SharedLoadIteratorMixed(
|
||||
TensorRef ref,
|
||||
int thread_idx
|
||||
):
|
||||
stride_((ref.stride(0) / LoadType::kElements)) {
|
||||
|
||||
TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx);
|
||||
|
||||
// Initialize pointers
|
||||
LoadType const *base_ptr = reinterpret_cast<LoadType const *>(ref.data()) + thread_offset.row() * stride_;
|
||||
|
||||
int lane_col_idx = thread_offset.column() / 8;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kLoadsPerAccess; ++i) {
|
||||
int lane_offset = (lane_col_idx % 8) * 2 | ((lane_col_idx / 4) ^ i);
|
||||
|
||||
pointers_[i] = base_ptr + lane_offset;
|
||||
}
|
||||
}
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kLoadsPerAccess; ++i) {
|
||||
pointers_[i] += pointer_offset / LoadType::kElements;
|
||||
}
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void add_tile_offset(TensorCoord const &offset) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kLoadsPerAccess; ++i) {
|
||||
pointers_[i] += offset.row() * stride_ + offset.column() / LoadType::kElements;
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
|
||||
|
||||
int row_ptr_offset =
|
||||
row * ThreadMap::Delta::kRow * stride_ +
|
||||
group * ThreadMap::Delta::kGroup* stride_ +
|
||||
cluster * ThreadMap::Delta::kCluster * stride_ +
|
||||
pointer_offset / LoadType::kElements;
|
||||
|
||||
int frag_row_idx = (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster));
|
||||
|
||||
LoadType *frag_ptr = reinterpret_cast<LoadType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) {
|
||||
|
||||
int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int v = 0; v < kLoadsPerAccess; ++v) {
|
||||
|
||||
LoadType const *memory_pointer = pointers_[v];
|
||||
|
||||
frag_ptr[frag_idx * kLoadsPerAccess + v] = memory_pointer[row_ptr_offset];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads a fragment
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
|
||||
load_with_pointer_offset(frag, 0);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -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:
|
||||
|
||||
@@ -0,0 +1,188 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile
|
||||
that participate in one warp-level store operation.
|
||||
|
||||
Typically, the accumulator tile is the largest single block of register-backed storage
|
||||
within the kernel. Storing it to memory is best accomplished by partitioning it into
|
||||
smaller tiles and storing these sequentially.
|
||||
|
||||
Round trips through shared memory during the Epilogue phase require partitioning, as
|
||||
shared memory capacity is typically insufficient for a threadblock's total accumulator
|
||||
size.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
|
||||
#include "cutlass/epilogue/warp/tensor_op_policy.h"
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace warp {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
///
|
||||
template <
|
||||
typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape)
|
||||
typename OperatorShape, ///< matrix multiply operation shape (concept: gemm::GemmShape)
|
||||
typename OperatorElementC, ///< matrix multiply operation data type (concept: data type)
|
||||
typename OperatorFragmentC, ///< matrix multiply operation fragment (concept: Array)
|
||||
typename Layout ///< target shared memory layout
|
||||
>
|
||||
class FragmentIteratorGaussianComplexTensorOp;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
/// Partial specialization for row-major shared memory
|
||||
template <
|
||||
typename WarpShape_, ///< shape of the warp-level GEMM tile
|
||||
typename OperatorShape_, ///< underlying real-valued matrix multiply operation shape (concept: gemm::GemmShape)
|
||||
typename OperatorElementC_, ///< underlying real-valued matrix multiply operation data type
|
||||
typename OperatorFragmentC_ ///< underlying real-valued matrix multiply operation fragment (concept: Array)
|
||||
>
|
||||
class FragmentIteratorGaussianComplexTensorOp<WarpShape_, OperatorShape_, OperatorElementC_, OperatorFragmentC_, layout::RowMajor> {
|
||||
public:
|
||||
|
||||
using WarpShape = WarpShape_;
|
||||
using OperatorShape = OperatorShape_;
|
||||
using OperatorElementC = OperatorElementC_;
|
||||
using OperatorFragmentC = OperatorFragmentC_;
|
||||
using Layout = layout::RowMajor;
|
||||
|
||||
using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
|
||||
|
||||
/// This is the fragment size produced by one access of the iterator.
|
||||
using Fragment = Array<
|
||||
complex<OperatorElementC>,
|
||||
Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
|
||||
|
||||
/// Size of one part of accumulator of 3-part accumulator in units of number of OperatorElementC
|
||||
static int const kElementsAccumulatorPerPart =
|
||||
OperatorFragmentC::kElements * Policy::OperatorCount::kRow * Policy::OperatorCount::kColumn;
|
||||
|
||||
/// Offset into the accumulator fragment part 1
|
||||
static int const kPart1Index = kElementsAccumulatorPerPart * 0;
|
||||
|
||||
/// Offset into the accumulator fragment part 2
|
||||
static int const kPart2Index = kElementsAccumulatorPerPart * 1;
|
||||
|
||||
/// Offset into the accumulator fragment part 3
|
||||
static int const kPart3Index = kElementsAccumulatorPerPart * 2;
|
||||
|
||||
/// This is the complete warp-level accumulator tile holding part1, part2, and part3
|
||||
using AccumulatorTile = Array<OperatorElementC, kElementsAccumulatorPerPart * 3>;
|
||||
|
||||
/// This is the complete warp-level accumulator tile holding final output of complex<T> type
|
||||
using OutputAccumulatorTile = Array<complex<OperatorElementC>, kElementsAccumulatorPerPart>;
|
||||
|
||||
/// Number of times this iterator can be incremented
|
||||
static int const kIterations = Policy::kIterations;
|
||||
|
||||
private:
|
||||
|
||||
/// Internal access type
|
||||
using AccessType = Array<OperatorElementC, Policy::kElementsPerAccess>;
|
||||
|
||||
using FragmentAccessType = Array<complex<OperatorElementC>, Policy::kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Accumulator tile
|
||||
AccessType const *accumulators_;
|
||||
|
||||
/// Internal index
|
||||
int index_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs an iterator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentIteratorGaussianComplexTensorOp(AccumulatorTile const &accum):
|
||||
accumulators_(reinterpret_cast<AccessType const *>(&accum)),
|
||||
index_(0) {
|
||||
}
|
||||
|
||||
/// Increments
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentIteratorGaussianComplexTensorOp &operator++() {
|
||||
++index_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Decrements
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentIteratorGaussianComplexTensorOp &operator--() {
|
||||
--index_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Loads a fragment from the referenced part of the accumulator tile
|
||||
CUTLASS_HOST_DEVICE
|
||||
void load(Fragment &frag, int index_offset = 0) const {
|
||||
|
||||
int index = index_ + index_offset;
|
||||
|
||||
FragmentAccessType *frag_ptr = reinterpret_cast<FragmentAccessType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
|
||||
|
||||
int accumulator_access_offset =
|
||||
index + n * Policy::kAccumulatorColumnStride / Policy::kElementsPerAccess;
|
||||
|
||||
auto const & part1_accum_array = accumulators_[accumulator_access_offset + kPart1Index];
|
||||
auto const & part2_accum_array = accumulators_[accumulator_access_offset + kPart2Index / Policy::kElementsPerAccess];
|
||||
auto const & part3_accum_array = accumulators_[accumulator_access_offset + kPart3Index / Policy::kElementsPerAccess];
|
||||
|
||||
// Pack parts 1, 2, and 3 into a structure. This is likely to result in MOVs
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < Policy::kElementsPerAccess; ++i) {
|
||||
|
||||
frag_ptr[n][i].real() = part1_accum_array[i] - part3_accum_array[i];
|
||||
frag_ptr[n][i].imag() = part1_accum_array[i] + part2_accum_array[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace warp
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -0,0 +1,675 @@
|
||||
/***************************************************************************************************
|
||||
* 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:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * Redistributions in binary form must reproduce the above copyright notice, this list of
|
||||
* conditions and the following disclaimer in the documentation and/or other materials
|
||||
* provided with the distribution.
|
||||
* * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
|
||||
* to endorse or promote products derived from this software without specific prior written
|
||||
* permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
|
||||
* IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
|
||||
* FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
|
||||
* OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
* STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/layout/pitch_linear.h"
|
||||
|
||||
#include "cutlass/arch/memory_sm75.h"
|
||||
#include "cutlass/epilogue/warp/tensor_op_policy.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace warp {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Template for reading and writing tiles of accumulators to shared memory. This is optimized
|
||||
/// for mixed-precision epilogues in which the accumulators are 32b in width, but the output
|
||||
/// data type is smaller.
|
||||
template <
|
||||
typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape)
|
||||
typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape)
|
||||
typename Element_, ///< data type of accumulator element
|
||||
int ElementSizeBits, ///< Size of accumulator element in bits
|
||||
int OutputSizeBits, ///< Size of output element in bits
|
||||
int OutputElementCount, ///< number of elements in output vector
|
||||
int ContiguousLanes ///< Number of consecutive lanes writing to contiguous memory
|
||||
>
|
||||
class TileIteratorTensorOpMixed {
|
||||
public:
|
||||
|
||||
using WarpShape = WarpShape_;
|
||||
using OperatorShape = OperatorShape_;
|
||||
using Element = Element_;
|
||||
using Layout = layout::RowMajor;
|
||||
static int const kOutputElementCount = OutputElementCount;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>; ///< Tensor Reference object
|
||||
using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor
|
||||
using Index = typename TensorRef::Index;
|
||||
using LongIndex = typename TensorRef::LongIndex;
|
||||
|
||||
using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
|
||||
|
||||
/// Shape of the tile in memory
|
||||
using Shape = MatrixShape<
|
||||
Policy::kRowsPerIteration,
|
||||
WarpShape::kN
|
||||
>;
|
||||
|
||||
/// This is the fragment size produced by one access of the iterator.
|
||||
using Fragment = Array<
|
||||
Element,
|
||||
Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
|
||||
|
||||
/// This is the complete warp-level accumulator tile.
|
||||
//using AccumulatorTile = typename Operator::FragmentC;
|
||||
|
||||
/// Number of times this iterator can be incremented
|
||||
static int const kIterations = Policy::kIterations;
|
||||
|
||||
// Internal constants
|
||||
struct Detail {
|
||||
static int const kLanesInQuad = 4;
|
||||
|
||||
/// Number of pointers needed to write accumulators
|
||||
static int const kPointerCount =
|
||||
(OutputElementCount * sizeof_bits<Element>::value) / (const_min(128, OutputElementCount * sizeof_bits<Element>::value));
|
||||
|
||||
static_assert(kPointerCount <= 4, "Can only accommodate four pointers at present.");
|
||||
static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32).");
|
||||
};
|
||||
|
||||
/// Padding quantity
|
||||
using Padding = MatrixShape<
|
||||
0,
|
||||
Detail::kLanesInQuad * Policy::kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
|
||||
/// Storage type for accessing memory
|
||||
using AccessType = AlignedArray<Element, Policy::kElementsPerAccess>;
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Internal pointer to memory
|
||||
AccessType *pointers_[Detail::kPointerCount];
|
||||
|
||||
/// Stride in units of AccessType
|
||||
int stride_;
|
||||
|
||||
/// Logical column in which warp tile is aligned
|
||||
int warp_column_;
|
||||
|
||||
public:
|
||||
|
||||
/// Default constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed() {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
|
||||
pointers_[i] = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
/// Constructor from TensorRef
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed(
|
||||
TensorRef const &ref,
|
||||
unsigned lane_id
|
||||
):
|
||||
stride_(ref.stride()[0] / Policy::kElementsPerAccess),
|
||||
warp_column_(0) {
|
||||
|
||||
int quad_id = (lane_id / Detail::kLanesInQuad);
|
||||
int lane_in_quad = (lane_id % Detail::kLanesInQuad);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
|
||||
AccessType *ptr = reinterpret_cast<AccessType *>(ref.data()) + quad_id * stride_;
|
||||
int column_idx = (lane_in_quad % 2) + (((lane_in_quad / 2) + i) % Detail::kPointerCount) * 2;
|
||||
|
||||
ptr += column_idx;
|
||||
|
||||
if (i == 0) {
|
||||
pointers_[0 % Detail::kPointerCount] = ptr;
|
||||
}
|
||||
else if (i == 1) {
|
||||
pointers_[1 % Detail::kPointerCount] = ptr;
|
||||
}
|
||||
else if (i == 2) {
|
||||
pointers_[2 % Detail::kPointerCount] = ptr;
|
||||
}
|
||||
else if (i == 3) {
|
||||
pointers_[3 % Detail::kPointerCount] = ptr;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Adds a pointer offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
|
||||
pointers_[i] += pointer_offset / Policy::kElementsPerAccess;
|
||||
}
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
///< advances in units of whole tiles along the logical coordinate space of the tensor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
|
||||
pointers_[i] += tile_offset.row() * Shape::kRow * stride_ +
|
||||
tile_offset.column() * Shape::kColumn / Policy::kElementsPerAccess;
|
||||
}
|
||||
|
||||
warp_column_ += tile_offset.column() * Shape::kColumn;
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
///< advances in units of whole tiles along the logical coordinate space of the tensor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) {
|
||||
return add_tile_offset(tile_offset);
|
||||
}
|
||||
|
||||
/// Store
|
||||
CUTLASS_DEVICE
|
||||
void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
|
||||
|
||||
AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int64_t n = 0; n < Policy::OperatorCount::kColumn; ++n) {
|
||||
|
||||
int column_idx = warp_column_ + n * Detail::kLanesInQuad * Policy::kElementsPerAccess;
|
||||
int ptr_idx = ((column_idx * sizeof_bits<Element>::value) / 1024) % Detail::kPointerCount;
|
||||
|
||||
AccessType *ptr;
|
||||
if (ptr_idx == 0) {
|
||||
ptr = pointers_[0 % Detail::kPointerCount];
|
||||
}
|
||||
else if (ptr_idx == 1) {
|
||||
ptr = pointers_[1 % Detail::kPointerCount];
|
||||
}
|
||||
else if (ptr_idx == 2) {
|
||||
ptr = pointers_[2 % Detail::kPointerCount];
|
||||
}
|
||||
else if (ptr_idx == 3) {
|
||||
ptr = pointers_[3 % Detail::kPointerCount];
|
||||
}
|
||||
|
||||
int offset = n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess;
|
||||
#if 0
|
||||
// Using inline PTX to avoid generic memory
|
||||
AccessType *smem_ptr = pointers_[ptr_idx];
|
||||
smem_ptr[offset] = frag_ptr[n];
|
||||
#else
|
||||
uint32_t smem_addr = arch::cutlass_get_smem_pointer(ptr);
|
||||
uint32_t const *data = reinterpret_cast<uint32_t const *>(frag_ptr + n);
|
||||
uint32_t offset_in_bytes = offset * sizeof(AccessType);
|
||||
|
||||
asm volatile(
|
||||
"{ .reg .u32 smem_ptr; add.u32 smem_ptr, %0, %1; st.shared.v2.u32 [smem_ptr], {%2, %3}; }\n"
|
||||
: : "r"(smem_addr), "r"(offset_in_bytes), "r"(data[0]), "r"(data[1])
|
||||
);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
/// Store
|
||||
CUTLASS_HOST_DEVICE
|
||||
void store(Fragment const &frag) {
|
||||
store_with_pointer_offset(frag, 0);
|
||||
}
|
||||
|
||||
/// Load
|
||||
CUTLASS_HOST_DEVICE
|
||||
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const {
|
||||
|
||||
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int64_t n = 0; n < Policy::OperatorCount::kColumn; ++n) {
|
||||
|
||||
int column_idx = warp_column_ + n * Detail::kLanesInQuad * Policy::kElementsPerAccess;
|
||||
int ptr_idx = ((column_idx * sizeof_bits<Element>::value) / 1024) % Detail::kPointerCount;
|
||||
|
||||
AccessType const *smem_ptr = pointers_[ptr_idx];
|
||||
frag_ptr[n] = smem_ptr[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess];
|
||||
}
|
||||
}
|
||||
|
||||
/// Load
|
||||
CUTLASS_HOST_DEVICE
|
||||
void load(Fragment &frag) const {
|
||||
load_with_pointer_offset(frag, 0);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for int32_t x 16 => int8_t x 16
|
||||
template <
|
||||
typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape)
|
||||
typename OperatorShape_ ///< matrix multiply operation shape (concept: gemm::GemmShape)
|
||||
>
|
||||
class TileIteratorTensorOpMixed<WarpShape_, OperatorShape_, int32_t, 32, 8, 16, 8> {
|
||||
public:
|
||||
|
||||
using WarpShape = WarpShape_;
|
||||
using OperatorShape = OperatorShape_;
|
||||
using Element = int32_t;
|
||||
using Layout = layout::RowMajor;
|
||||
static int const kOutputElementCount = 16;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>; ///< Tensor Reference object
|
||||
using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor
|
||||
using Index = typename TensorRef::Index;
|
||||
using LongIndex = typename TensorRef::LongIndex;
|
||||
|
||||
using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
|
||||
|
||||
/// Shape of the tile in memory
|
||||
using Shape = MatrixShape<
|
||||
Policy::kRowsPerIteration,
|
||||
WarpShape::kN
|
||||
>;
|
||||
|
||||
/// This is the fragment size produced by one access of the iterator.
|
||||
using Fragment = Array<
|
||||
Element,
|
||||
Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
|
||||
|
||||
/// This is the complete warp-level accumulator tile.
|
||||
//using AccumulatorTile = typename Operator::FragmentC;
|
||||
|
||||
/// Number of times this iterator can be incremented
|
||||
static int const kIterations = Policy::kIterations;
|
||||
|
||||
// Internal constants
|
||||
struct Detail {
|
||||
static int const kLanesInQuad = 4;
|
||||
|
||||
/// Number of pointers needed to write accumulators
|
||||
static int const kPointerCount = 2;
|
||||
|
||||
/// Offsets added
|
||||
static int const kOffsetCount = 4;
|
||||
|
||||
static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32).");
|
||||
};
|
||||
|
||||
/// Padding quantity
|
||||
using Padding = MatrixShape<0, Detail::kLanesInQuad * 2>;
|
||||
|
||||
private:
|
||||
|
||||
/// Storage type for accessing memory
|
||||
using AccessType = AlignedArray<Element, 2>;
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Internal pointer to memory
|
||||
AccessType *pointers_[Detail::kPointerCount];
|
||||
|
||||
/// Stride in units of AccessType
|
||||
int stride_;
|
||||
|
||||
/// Uniform offset in bytes added to warp tile iterator
|
||||
int uniform_offset_[Detail::kOffsetCount];
|
||||
|
||||
public:
|
||||
|
||||
/// Default constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed() {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
|
||||
pointers_[i] = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
/// Constructor from TensorRef
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed(
|
||||
TensorRef const &ref,
|
||||
unsigned lane_id
|
||||
):
|
||||
stride_(ref.stride()[0] / AccessType::kElements) {
|
||||
|
||||
int quad_id = (lane_id / Detail::kLanesInQuad);
|
||||
int lane_in_quad = (lane_id % Detail::kLanesInQuad);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < Detail::kPointerCount; ++i) {
|
||||
AccessType *ptr = reinterpret_cast<AccessType *>(ref.data()) + quad_id * stride_;
|
||||
int column_idx = lane_in_quad ^ (i * 2);
|
||||
|
||||
ptr += column_idx;
|
||||
|
||||
if (i == 0) {
|
||||
pointers_[0] = ptr;
|
||||
}
|
||||
else if (i == 1) {
|
||||
pointers_[1] = ptr;
|
||||
}
|
||||
}
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < Detail::kOffsetCount; ++i) {
|
||||
uniform_offset_[i] = (i ^ 0) * 4 * sizeof(AccessType);
|
||||
}
|
||||
}
|
||||
|
||||
/// Adds a pointer offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
|
||||
pointers_[i] += pointer_offset / AccessType::kElements;
|
||||
}
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
///< advances in units of whole tiles along the logical coordinate space of the tensor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) {
|
||||
|
||||
int ptr_offset = tile_offset.row() * Shape::kRow * stride_ +
|
||||
tile_offset.column() * Shape::kColumn / AccessType::kElements;
|
||||
|
||||
pointers_[0] += ptr_offset;
|
||||
pointers_[1] += ptr_offset;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < Detail::kOffsetCount; ++i) {
|
||||
uniform_offset_[i] = (i ^ tile_offset.column()) * 4 * sizeof(AccessType);
|
||||
}
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
///< advances in units of whole tiles along the logical coordinate space of the tensor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) {
|
||||
return add_tile_offset(tile_offset);
|
||||
}
|
||||
|
||||
/// Store
|
||||
CUTLASS_DEVICE
|
||||
void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
|
||||
|
||||
AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
|
||||
|
||||
int ptr_idx = (n / 4);
|
||||
int offset_idx = (n % 4);
|
||||
|
||||
AccessType *ptr;
|
||||
if (ptr_idx == 0) {
|
||||
ptr = pointers_[0];
|
||||
}
|
||||
else if (ptr_idx == 1) {
|
||||
ptr = pointers_[1];
|
||||
}
|
||||
|
||||
int offset = (n / 4) * 16 + pointer_offset / AccessType::kElements;
|
||||
|
||||
#if 0
|
||||
//
|
||||
// Using inline PTX to avoid generic memory
|
||||
//
|
||||
AccessType *smem_ptr = pointers_[ptr_idx];
|
||||
smem_ptr[offset] = frag_ptr[n];
|
||||
#else
|
||||
uint32_t smem_addr = arch::cutlass_get_smem_pointer(ptr);
|
||||
uint32_t const *data = reinterpret_cast<uint32_t const *>(frag_ptr + n);
|
||||
uint32_t offset_in_bytes = offset * sizeof(AccessType) + uniform_offset_[offset_idx];
|
||||
|
||||
asm volatile(
|
||||
"{ .reg .u32 smem_ptr; add.u32 smem_ptr, %0, %1; st.shared.v2.u32 [smem_ptr], {%2, %3}; }\n"
|
||||
: : "r"(smem_addr), "r"(offset_in_bytes), "r"(data[0]), "r"(data[1])
|
||||
);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
/// Store
|
||||
CUTLASS_HOST_DEVICE
|
||||
void store(Fragment const &frag) {
|
||||
store_with_pointer_offset(frag, 0);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for int32_t x 8 => int8_t x 8
|
||||
template <
|
||||
typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape)
|
||||
typename OperatorShape_ ///< matrix multiply operation shape (concept: gemm::GemmShape)
|
||||
>
|
||||
class TileIteratorTensorOpMixed<WarpShape_, OperatorShape_, int32_t, 32, 8, 8, 8> {
|
||||
public:
|
||||
|
||||
using WarpShape = WarpShape_;
|
||||
using OperatorShape = OperatorShape_;
|
||||
using Element = int32_t;
|
||||
using Layout = layout::RowMajor;
|
||||
static int const kOutputElementCount = 8;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>; ///< Tensor Reference object
|
||||
using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor
|
||||
using Index = typename TensorRef::Index;
|
||||
using LongIndex = typename TensorRef::LongIndex;
|
||||
|
||||
using Policy = TensorOpPolicy<WarpShape, OperatorShape, Layout>;
|
||||
|
||||
/// Shape of the tile in memory
|
||||
using Shape = MatrixShape<
|
||||
Policy::kRowsPerIteration,
|
||||
WarpShape::kN
|
||||
>;
|
||||
|
||||
/// This is the fragment size produced by one access of the iterator.
|
||||
using Fragment = Array<
|
||||
Element,
|
||||
Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>;
|
||||
|
||||
/// This is the complete warp-level accumulator tile.
|
||||
//using AccumulatorTile = typename Operator::FragmentC;
|
||||
|
||||
/// Number of times this iterator can be incremented
|
||||
static int const kIterations = Policy::kIterations;
|
||||
|
||||
// Internal constants
|
||||
struct Detail {
|
||||
static int const kLanesInQuad = 4;
|
||||
|
||||
/// Number of pointers needed to write accumulators
|
||||
static int const kPointerCount = 2;
|
||||
|
||||
static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32).");
|
||||
};
|
||||
|
||||
/// Padding quantity
|
||||
using Padding = MatrixShape<0, Detail::kLanesInQuad * 2>;
|
||||
|
||||
private:
|
||||
|
||||
/// Storage type for accessing memory
|
||||
using AccessType = AlignedArray<Element, 2>;
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Internal pointer to memory
|
||||
AccessType *pointers_[Detail::kPointerCount];
|
||||
|
||||
/// Stride in units of AccessType
|
||||
int stride_;
|
||||
|
||||
public:
|
||||
|
||||
/// Default constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed() {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
|
||||
pointers_[i] = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
/// Constructor from TensorRef
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed(
|
||||
TensorRef const &ref,
|
||||
unsigned lane_id
|
||||
):
|
||||
stride_(ref.stride()[0] / AccessType::kElements) {
|
||||
|
||||
int quad_id = (lane_id / Detail::kLanesInQuad);
|
||||
int lane_in_quad = (lane_id % Detail::kLanesInQuad);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < Detail::kPointerCount; ++i) {
|
||||
AccessType *ptr = reinterpret_cast<AccessType *>(ref.data()) + quad_id * stride_;
|
||||
int column_idx = lane_in_quad ^ (i * 2);
|
||||
|
||||
ptr += column_idx;
|
||||
|
||||
if (i == 0) {
|
||||
pointers_[0] = ptr;
|
||||
}
|
||||
else if (i == 1) {
|
||||
pointers_[1] = ptr;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Adds a pointer offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int64_t i = 0; i < Detail::kPointerCount; ++i) {
|
||||
pointers_[i] += pointer_offset / AccessType::kElements;
|
||||
}
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
///< advances in units of whole tiles along the logical coordinate space of the tensor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) {
|
||||
|
||||
int ptr_offset = tile_offset.row() * Shape::kRow * stride_ +
|
||||
tile_offset.column() * Shape::kColumn / AccessType::kElements;
|
||||
|
||||
pointers_[0] += ptr_offset;
|
||||
pointers_[1] += ptr_offset;
|
||||
|
||||
if (tile_offset.column() % 2) {
|
||||
auto tmp = pointers_[0];
|
||||
pointers_[0] = pointers_[1];
|
||||
pointers_[1] = tmp;
|
||||
}
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
///< advances in units of whole tiles along the logical coordinate space of the tensor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) {
|
||||
return add_tile_offset(tile_offset);
|
||||
}
|
||||
|
||||
/// Store
|
||||
CUTLASS_DEVICE
|
||||
void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
|
||||
|
||||
AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) {
|
||||
|
||||
int ptr_idx = (n / 4);
|
||||
|
||||
AccessType *ptr;
|
||||
if (ptr_idx == 0) {
|
||||
ptr = pointers_[0];
|
||||
}
|
||||
else if (ptr_idx == 1) {
|
||||
ptr = pointers_[1];
|
||||
}
|
||||
|
||||
int offset = (n / 4) * 16 + pointer_offset / AccessType::kElements + (n % 4) * 4;
|
||||
|
||||
#if 0
|
||||
//
|
||||
// Using inline PTX to avoid generic memory
|
||||
//
|
||||
AccessType *smem_ptr = pointers_[ptr_idx];
|
||||
smem_ptr[offset] = frag_ptr[n];
|
||||
#else
|
||||
uint32_t smem_addr = arch::cutlass_get_smem_pointer(ptr);
|
||||
uint32_t const *data = reinterpret_cast<uint32_t const *>(frag_ptr + n);
|
||||
uint32_t offset_in_bytes = offset * sizeof(AccessType);
|
||||
|
||||
asm volatile(
|
||||
"{ .reg .u32 smem_ptr; add.u32 smem_ptr, %0, %1; st.shared.v2.u32 [smem_ptr], {%2, %3}; }\n"
|
||||
: : "r"(smem_addr), "r"(offset_in_bytes), "r"(data[0]), "r"(data[1])
|
||||
);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
/// Store
|
||||
CUTLASS_HOST_DEVICE
|
||||
void store(Fragment const &frag) {
|
||||
store_with_pointer_offset(frag, 0);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace warp
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -96,6 +96,16 @@ struct multiply_add {
|
||||
}
|
||||
};
|
||||
|
||||
/// Fused multiply-add
|
||||
template <typename T>
|
||||
struct and_add {
|
||||
CUTLASS_HOST_DEVICE
|
||||
T operator()(T const &a, T const &b, T const &c) const {
|
||||
return ((a & b) + c);
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
/// Fused multiply-add
|
||||
template <typename T>
|
||||
struct xor_add {
|
||||
@@ -1207,6 +1217,212 @@ struct multiply_add<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>> {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Fused multiply-add
|
||||
template <int N>
|
||||
struct multiply_add<Array<bfloat16_t, N>, Array<bfloat16_t, N>, Array<bfloat16_t, N>> {
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<bfloat16_t, N> operator()(
|
||||
Array<bfloat16_t, N> const &a,
|
||||
Array<bfloat16_t, N> const &b,
|
||||
Array<bfloat16_t, N> const &c) const {
|
||||
|
||||
Array<bfloat16_t, N> result;
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
|
||||
unsigned *result_ptr = reinterpret_cast<unsigned *>(&result);
|
||||
unsigned const *a_ptr = reinterpret_cast<unsigned const *>(&a);
|
||||
unsigned const *b_ptr = reinterpret_cast<unsigned const *>(&b);
|
||||
unsigned const *c_ptr = reinterpret_cast<unsigned const *>(&c);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N / 2; ++i) {
|
||||
asm ("fma.rn.bf16x2 %0, %1, %2, %3;\n"
|
||||
: "=r"(result_ptr[i])
|
||||
: "r"(a_ptr[i]), "r"(b_ptr[i]), "r"(c_ptr[i])
|
||||
);
|
||||
}
|
||||
|
||||
if (N % 2) {
|
||||
|
||||
uint16_t *result_ptr = reinterpret_cast<uint16_t *>(&result);
|
||||
uint16_t const *a_residual_ptr = reinterpret_cast<uint16_t const *>(&a);
|
||||
uint16_t const *b_residual_ptr = reinterpret_cast<uint16_t const *>(&b);
|
||||
uint16_t const *c_residual_ptr = reinterpret_cast<uint16_t const *>(&c);
|
||||
|
||||
asm ("fma.rn.bf16 %0, %1, %2, %3;\n"
|
||||
: "=h"(result_ptr[N - 1])
|
||||
: "h"(a_residual_ptr[N - 1]), "h"(b_residual_ptr[N - 1]), "h"(c_residual_ptr[N - 1])
|
||||
);
|
||||
}
|
||||
|
||||
#else
|
||||
|
||||
multiply_add<bfloat16_t> op;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N; ++i) {
|
||||
result[i] = op(a[i], b[i], c[i]);
|
||||
}
|
||||
#endif
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<bfloat16_t, N> operator()(
|
||||
bfloat16_t const &a,
|
||||
Array<bfloat16_t, N> const &b,
|
||||
Array<bfloat16_t, N> const &c) const {
|
||||
|
||||
Array<bfloat16_t, N> result;
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
|
||||
unsigned *result_ptr = reinterpret_cast<unsigned *>(&result);
|
||||
|
||||
unsigned const *b_ptr = reinterpret_cast<unsigned const *>(&b);
|
||||
unsigned const *c_ptr = reinterpret_cast<unsigned const *>(&c);
|
||||
|
||||
unsigned a_packed = static_cast<unsigned>(a.raw());
|
||||
a_packed = (a_packed | (a_packed << 16));
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N / 2; ++i) {
|
||||
asm ("fma.rn.bf16x2 %0, %1, %2, %3;\n"
|
||||
: "=r"(result_ptr[i])
|
||||
: "r"(a_packed), "r"(b_ptr[i]), "r"(c_ptr[i])
|
||||
);
|
||||
}
|
||||
|
||||
if (N % 2) {
|
||||
|
||||
uint16_t *result_ptr = reinterpret_cast<uint16_t *>(&result);
|
||||
uint16_t const *a_residual_ptr = reinterpret_cast<uint16_t const *>(&a);
|
||||
uint16_t const *b_residual_ptr = reinterpret_cast<uint16_t const *>(&b);
|
||||
uint16_t const *c_residual_ptr = reinterpret_cast<uint16_t const *>(&c);
|
||||
|
||||
asm ("fma.rn.bf16 %0, %1, %2, %3;\n"
|
||||
: "=h"(result_ptr[N - 1])
|
||||
: "h"(a_residual_ptr[0]), "h"(b_residual_ptr[N - 1]), "h"(c_residual_ptr[N - 1])
|
||||
);
|
||||
}
|
||||
|
||||
#else
|
||||
|
||||
multiply_add<bfloat16_t> op;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N; ++i) {
|
||||
result[i] = op(a, b[i], c[i]);
|
||||
}
|
||||
#endif
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<bfloat16_t, N> operator()(
|
||||
Array<bfloat16_t, N> const &a,
|
||||
bfloat16_t const &b,
|
||||
Array<bfloat16_t, N> const &c) const {
|
||||
|
||||
Array<bfloat16_t, N> result;
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
|
||||
unsigned *result_ptr = reinterpret_cast<unsigned *>(&result);
|
||||
|
||||
unsigned const *a_ptr = reinterpret_cast<unsigned const *>(&a);
|
||||
unsigned const *c_ptr = reinterpret_cast<unsigned const *>(&c);
|
||||
|
||||
unsigned b_packed = static_cast<unsigned>(b.raw());
|
||||
b_packed = (b_packed | (b_packed << 16));
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N / 2; ++i) {
|
||||
asm ("fma.rn.bf16x2 %0, %1, %2, %3;\n"
|
||||
: "=r"(result_ptr[i])
|
||||
: "r"(a_ptr[i]), "r"(b_packed), "r"(c_ptr[i])
|
||||
);
|
||||
}
|
||||
|
||||
if (N % 2) {
|
||||
|
||||
uint16_t *result_ptr = reinterpret_cast<uint16_t *>(&result);
|
||||
uint16_t const *a_residual_ptr = reinterpret_cast<uint16_t const *>(&a);
|
||||
uint16_t const *b_residual_ptr = reinterpret_cast<uint16_t const *>(&b);
|
||||
uint16_t const *c_residual_ptr = reinterpret_cast<uint16_t const *>(&c);
|
||||
|
||||
asm ("fma.rn.bf16 %0, %1, %2, %3;\n"
|
||||
: "=h"(result_ptr[N - 1])
|
||||
: "h"(a_residual_ptr[N - 1]), "h"(b_residual_ptr[0]), "h"(c_residual_ptr[N - 1])
|
||||
);
|
||||
}
|
||||
|
||||
#else
|
||||
|
||||
multiply_add<bfloat16_t> op;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N; ++i) {
|
||||
result[i] = op(a[i], b, c[i]);
|
||||
}
|
||||
#endif
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<bfloat16_t, N> operator()(
|
||||
Array<bfloat16_t, N> const &a,
|
||||
Array<bfloat16_t, N> const &b,
|
||||
bfloat16_t const &c) const {
|
||||
|
||||
Array<bfloat16_t, N> result;
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
|
||||
unsigned *result_ptr = reinterpret_cast<unsigned *>(&result);
|
||||
|
||||
unsigned const *a_ptr = reinterpret_cast<unsigned const *>(&a);
|
||||
unsigned const *b_ptr = reinterpret_cast<unsigned const *>(&b);
|
||||
|
||||
unsigned c_packed = static_cast<unsigned>(c.raw());
|
||||
c_packed = (c_packed | (c_packed << 16));
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N / 2; ++i) {
|
||||
asm ("fma.rn.bf16x2 %0, %1, %2, %3;\n"
|
||||
: "=r"(result_ptr[i])
|
||||
: "r"(a_ptr[i]), "r"(b_ptr[i]), "r"(c_packed)
|
||||
);
|
||||
}
|
||||
|
||||
if (N % 2) {
|
||||
|
||||
uint16_t *result_ptr = reinterpret_cast<uint16_t *>(&result);
|
||||
uint16_t const *a_residual_ptr = reinterpret_cast<uint16_t const *>(&a);
|
||||
uint16_t const *b_residual_ptr = reinterpret_cast<uint16_t const *>(&b);
|
||||
uint16_t const *c_residual_ptr = reinterpret_cast<uint16_t const *>(&c);
|
||||
|
||||
asm ("fma.rn.bf16 %0, %1, %2, %3;\n"
|
||||
: "=h"(result_ptr[N - 1])
|
||||
: "h"(a_residual_ptr[N - 1]), "h"(b_residual_ptr[N - 1]), "h"(c_residual_ptr[0])
|
||||
);
|
||||
}
|
||||
|
||||
#else
|
||||
|
||||
multiply_add<bfloat16_t> op;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N; ++i) {
|
||||
result[i] = op(a[i], b[i], c);
|
||||
}
|
||||
#endif
|
||||
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass
|
||||
|
||||
@@ -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:
|
||||
@@ -422,6 +422,342 @@ struct DefaultGemmConfiguration<
|
||||
using Operator = arch::OpMultiplyAddSaturate;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm75,
|
||||
uint1b_t,
|
||||
uint1b_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<uint1b_t>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<uint1b_t>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 512>;
|
||||
using WarpShape = GemmShape<64, 64, 512>;
|
||||
using InstructionShape = GemmShape<8, 8, 128>;
|
||||
static int const kStages = 2;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
|
||||
|
||||
using Operator = arch::OpXorPopc;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename ElementA, typename ElementB, typename ElementC,
|
||||
typename ElementAccumulator>
|
||||
struct DefaultGemmConfiguration<arch::OpClassTensorOp, arch::Sm80, ElementA,
|
||||
ElementB, ElementC, ElementAccumulator> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<ElementA>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<ElementA>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 64>;
|
||||
using WarpShape = GemmShape<64, 64, 64>;
|
||||
using InstructionShape = GemmShape<16, 8, 16>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombination<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, ElementAccumulator,
|
||||
ElementAccumulator>;
|
||||
|
||||
using Operator = typename platform::conditional<
|
||||
(platform::is_same<ElementA, int8_t>::value ||
|
||||
platform::is_same<ElementA, int4b_t>::value ||
|
||||
platform::is_same<ElementA, uint8_t>::value ||
|
||||
platform::is_same<ElementA, uint4b_t>::value),
|
||||
arch::OpMultiplyAddSaturate, arch::OpMultiplyAdd>::type;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename ElementC,
|
||||
typename ElementAccumulator>
|
||||
struct DefaultGemmConfiguration<arch::OpClassTensorOp, arch::Sm80, double,
|
||||
double, ElementC, ElementAccumulator> {
|
||||
|
||||
static int const kAlignmentA = 1;
|
||||
static int const kAlignmentB = 1;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 64>;
|
||||
using WarpShape = GemmShape<64, 64, 64>;
|
||||
using InstructionShape = GemmShape<16, 8, 16>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombination<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, ElementAccumulator,
|
||||
ElementAccumulator>;
|
||||
|
||||
using Operator = arch::OpMultiplyAdd;
|
||||
};
|
||||
|
||||
|
||||
template <>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
complex<double>,
|
||||
complex<double>,
|
||||
complex<double>,
|
||||
complex<double>
|
||||
> {
|
||||
|
||||
static int const kAlignmentA = 1;
|
||||
static int const kAlignmentB = 1;
|
||||
|
||||
using ThreadblockShape = GemmShape<64, 64, 16>;
|
||||
using WarpShape = GemmShape<32, 32, 16>;
|
||||
using InstructionShape = GemmShape<8, 8, 4>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombination<
|
||||
complex<double>, 1, complex<double>,
|
||||
complex<double>>;
|
||||
|
||||
using Operator = arch::OpMultiplyAddComplex;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
int8_t,
|
||||
int8_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<int8_t>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<int8_t>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 64>;
|
||||
using WarpShape = GemmShape<64, 64, 64>;
|
||||
using InstructionShape = GemmShape<16, 8, 32>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
|
||||
|
||||
using Operator = arch::OpMultiplyAddSaturate;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
int8_t,
|
||||
uint8_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<int8_t>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<uint8_t>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 64>;
|
||||
using WarpShape = GemmShape<64, 64, 64>;
|
||||
using InstructionShape = GemmShape<16, 8, 32>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
|
||||
|
||||
using Operator = arch::OpMultiplyAddSaturate;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<uint8_t>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<int8_t>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 64>;
|
||||
using WarpShape = GemmShape<64, 64, 64>;
|
||||
using InstructionShape = GemmShape<16, 8, 32>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
|
||||
|
||||
using Operator = arch::OpMultiplyAddSaturate;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
uint8_t,
|
||||
uint8_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<uint8_t>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<uint8_t>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 64>;
|
||||
using WarpShape = GemmShape<64, 64, 64>;
|
||||
using InstructionShape = GemmShape<16, 8, 32>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
|
||||
|
||||
using Operator = arch::OpMultiplyAddSaturate;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
int4b_t,
|
||||
int4b_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<int4b_t>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<int4b_t>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 128>;
|
||||
using WarpShape = GemmShape<64, 64, 128>;
|
||||
using InstructionShape = GemmShape<16, 8, 64>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
|
||||
|
||||
using Operator = arch::OpMultiplyAddSaturate;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
int4b_t,
|
||||
uint4b_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<int4b_t>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<uint4b_t>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 128>;
|
||||
using WarpShape = GemmShape<64, 64, 128>;
|
||||
using InstructionShape = GemmShape<16, 8, 64>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
|
||||
|
||||
using Operator = arch::OpMultiplyAddSaturate;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
uint4b_t,
|
||||
int4b_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<uint4b_t>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<int4b_t>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 128>;
|
||||
using WarpShape = GemmShape<64, 64, 128>;
|
||||
using InstructionShape = GemmShape<16, 8, 64>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
|
||||
|
||||
using Operator = arch::OpMultiplyAddSaturate;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
uint4b_t,
|
||||
uint4b_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<uint4b_t>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<uint4b_t>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 128>;
|
||||
using WarpShape = GemmShape<64, 64, 128>;
|
||||
using InstructionShape = GemmShape<16, 8, 64>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
|
||||
|
||||
using Operator = arch::OpMultiplyAddSaturate;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename ElementC>
|
||||
struct DefaultGemmConfiguration<
|
||||
arch::OpClassTensorOp,
|
||||
arch::Sm80,
|
||||
uint1b_t,
|
||||
uint1b_t,
|
||||
ElementC,
|
||||
int32_t> {
|
||||
|
||||
static int const kAlignmentA = 128 / sizeof_bits<uint1b_t>::value;
|
||||
static int const kAlignmentB = 128 / sizeof_bits<uint1b_t>::value;
|
||||
|
||||
using ThreadblockShape = GemmShape<128, 256, 512>;
|
||||
using WarpShape = GemmShape<64, 64, 512>;
|
||||
using InstructionShape = GemmShape<16, 8, 256>;
|
||||
static int const kStages = 3;
|
||||
|
||||
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
|
||||
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
|
||||
|
||||
using Operator = arch::OpMultiplyAdd;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
} // namespace device
|
||||
} // namespace gemm
|
||||
|
||||
@@ -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:
|
||||
@@ -193,7 +193,7 @@ template <
|
||||
ElementAccumulator_>::EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle_ =
|
||||
typename threadblock::GemmCohortThreadblockSwizzle<LayoutA_, LayoutB_>,
|
||||
typename threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -192,7 +192,7 @@ template <
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle,
|
||||
typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
@@ -202,6 +202,7 @@ template <
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB = ComplexTransform::kNone,
|
||||
/// Multiply-add operator
|
||||
// (selects complex or gaussian complex)
|
||||
typename Operator_ = arch::OpMultiplyAddComplex,
|
||||
/// If true, kernel supports split-K with serial reduction
|
||||
bool SplitKSerial = false
|
||||
@@ -506,6 +507,7 @@ template <
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB,
|
||||
/// Multiply-add operator
|
||||
// (selects complex or gaussian complex)
|
||||
typename Operator_,
|
||||
/// If true, kernel supports split-K as a serial reduction
|
||||
bool SplitKSerial
|
||||
@@ -571,8 +573,8 @@ public:
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
TransformA,
|
||||
TransformB,
|
||||
TransformA,
|
||||
Operator,
|
||||
SplitKSerial
|
||||
>;
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -89,7 +89,7 @@ template <
|
||||
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
|
||||
ElementAccumulator_>::EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle,
|
||||
typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages =
|
||||
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
|
||||
|
||||
@@ -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:
|
||||
@@ -41,14 +41,78 @@ namespace device {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace detail {
|
||||
|
||||
template <
|
||||
typename ElementA_,
|
||||
typename LayoutA_,
|
||||
ComplexTransform TransformA,
|
||||
int AlignmentA,
|
||||
typename ElementB_,
|
||||
typename LayoutB_,
|
||||
ComplexTransform TransformB,
|
||||
int AlignmentB,
|
||||
typename LayoutC_,
|
||||
bool Transpose
|
||||
>
|
||||
struct MapArguments {
|
||||
using ElementA = ElementA_;
|
||||
using LayoutA = LayoutA_;
|
||||
static ComplexTransform const kTransformA = TransformA;
|
||||
static int const kAlignmentA = AlignmentA;
|
||||
using ElementB = ElementB_;
|
||||
using LayoutB = LayoutB_;
|
||||
static ComplexTransform const kTransformB = TransformB;
|
||||
static int const kAlignmentB = AlignmentB;
|
||||
using LayoutC = LayoutC_;
|
||||
};
|
||||
|
||||
template <
|
||||
typename ElementA_,
|
||||
typename LayoutA_,
|
||||
ComplexTransform TransformA,
|
||||
int AlignmentA,
|
||||
typename ElementB_,
|
||||
typename LayoutB_,
|
||||
ComplexTransform TransformB,
|
||||
int AlignmentB,
|
||||
typename LayoutC_
|
||||
>
|
||||
struct MapArguments<
|
||||
ElementA_,
|
||||
LayoutA_,
|
||||
TransformA,
|
||||
AlignmentA,
|
||||
ElementB_,
|
||||
LayoutB_,
|
||||
TransformB,
|
||||
AlignmentB,
|
||||
LayoutC_,
|
||||
true
|
||||
> {
|
||||
using ElementA = ElementB_;
|
||||
using LayoutA = typename layout::LayoutTranspose<LayoutB_>::type;
|
||||
static ComplexTransform const kTransformA = TransformB;
|
||||
static int const kAlignmentA = AlignmentB;
|
||||
using ElementB = ElementA_;
|
||||
using LayoutB = typename layout::LayoutTranspose<LayoutA_>::type;
|
||||
static ComplexTransform const kTransformB = TransformA;
|
||||
static int const kAlignmentB = AlignmentA;
|
||||
using LayoutC = typename layout::LayoutTranspose<LayoutC_>::type;
|
||||
};
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename GemmKernel_>
|
||||
class GemmUniversalAdapter {
|
||||
public:
|
||||
|
||||
using GemmKernel = GemmKernel_;
|
||||
|
||||
static_assert(std::is_same<typename GemmKernel::LayoutC, cutlass::layout::RowMajor>::value,
|
||||
"Universal adapter expects the kernel to be row-major and transposes its arguments.");
|
||||
static bool const kInternalTranspose =
|
||||
std::is_same<typename GemmKernel::LayoutC, cutlass::layout::RowMajor>::value;
|
||||
|
||||
using ThreadblockShape = typename GemmKernel::Mma::Shape;
|
||||
using WarpShape = typename GemmKernel::WarpShape;
|
||||
@@ -56,26 +120,39 @@ public:
|
||||
|
||||
using OperatorClass = typename GemmKernel::OperatorClass;
|
||||
using ArchTag = typename GemmKernel::ArchTag;
|
||||
|
||||
|
||||
// Type, layout, and complex transform deliberately exchanged with B
|
||||
using ElementA = typename GemmKernel::ElementB;
|
||||
using LayoutA = typename layout::LayoutTranspose<typename GemmKernel::LayoutB>::type;
|
||||
using TensorRefA = TensorRef<ElementA const, LayoutA>;
|
||||
static ComplexTransform const kTransformA = GemmKernel::kTransformB;
|
||||
using MapArguments = detail::MapArguments<
|
||||
typename GemmKernel::ElementA,
|
||||
typename GemmKernel::LayoutA,
|
||||
GemmKernel::kTransformA,
|
||||
GemmKernel::kAlignmentA,
|
||||
typename GemmKernel::ElementB,
|
||||
typename GemmKernel::LayoutB,
|
||||
GemmKernel::kTransformB,
|
||||
GemmKernel::kAlignmentB,
|
||||
typename GemmKernel::LayoutC,
|
||||
kInternalTranspose
|
||||
>;
|
||||
|
||||
using ElementA = typename MapArguments::ElementA;
|
||||
using LayoutA = typename MapArguments::LayoutA;
|
||||
static ComplexTransform const kTransformA = MapArguments::kTransformA;
|
||||
static int const kAlignmentA = GemmKernel::kAlignmentA;
|
||||
|
||||
// Type, layout, and complex transform deliberately exchanged with A
|
||||
using ElementB = typename GemmKernel::ElementA;
|
||||
using LayoutB = typename layout::LayoutTranspose<typename GemmKernel::LayoutA>::type;
|
||||
using TensorRefB = TensorRef<ElementB const, LayoutB>;
|
||||
static ComplexTransform const kTransformB = GemmKernel::kTransformA;
|
||||
using ElementB = typename MapArguments::ElementB;
|
||||
using LayoutB = typename MapArguments::LayoutB;
|
||||
static ComplexTransform const kTransformB = MapArguments::kTransformB;
|
||||
static int const kAlignmentB = GemmKernel::kAlignmentB;
|
||||
|
||||
|
||||
using ElementC = typename GemmKernel::ElementC;
|
||||
using LayoutC = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = typename MapArguments::LayoutC;
|
||||
static int const kAlignmentC = GemmKernel::kAlignmentC;
|
||||
|
||||
using TensorRefA = TensorRef<ElementA const, LayoutA>;
|
||||
using TensorRefB = TensorRef<ElementB const, LayoutB>;
|
||||
using TensorRefC = TensorRef<ElementC const, LayoutC>;
|
||||
using TensorRefD = TensorRef<ElementC, LayoutC>;
|
||||
static int const kAlignmentC = GemmKernel::kAlignmentC;
|
||||
|
||||
using ElementAccumulator = typename GemmKernel::Mma::Policy::Operator::ElementC;
|
||||
|
||||
@@ -99,7 +176,12 @@ public:
|
||||
|
||||
/// Helper to construct a transposed equivalent for the underying GEMM operator
|
||||
static Arguments to_underlying_arguments(Arguments const &args) {
|
||||
return args.transposed_problem();
|
||||
if (kInternalTranspose) {
|
||||
return args.transposed_problem();
|
||||
}
|
||||
else {
|
||||
return args;
|
||||
}
|
||||
}
|
||||
|
||||
/// Determines whether the GEMM can execute the given problem.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -400,7 +400,8 @@ enum class GemmUniversalMode {
|
||||
kGemm,
|
||||
kGemmSplitKParallel,
|
||||
kBatched,
|
||||
kArray
|
||||
kArray,
|
||||
kInvalid
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -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:
|
||||
@@ -49,6 +49,7 @@
|
||||
#include "cutlass/gemm/kernel/gemm_pipelined.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_core_sm75.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_core_sm70.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_core_sm80.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_core_simt.h"
|
||||
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
|
||||
@@ -116,6 +117,68 @@ template <
|
||||
struct DefaultGemm;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for Ampere Architecture
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// If true, kernel is configured to support serial reduction in the
|
||||
/// epilogue
|
||||
bool SplitKSerial,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator>
|
||||
struct DefaultGemm<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementC,
|
||||
layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp,
|
||||
arch::Sm80, ThreadblockShape, WarpShape, InstructionShape,
|
||||
EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial,
|
||||
Operator> {
|
||||
/// Define the threadblock-scoped matrix multiply-accumulate
|
||||
using Mma = typename cutlass::gemm::threadblock::DefaultMma<
|
||||
ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB,
|
||||
ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80,
|
||||
ThreadblockShape, WarpShape, InstructionShape, Stages,
|
||||
Operator>::ThreadblockMma;
|
||||
|
||||
static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
|
||||
|
||||
/// Define the epilogue
|
||||
using Epilogue =
|
||||
typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
|
||||
ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp,
|
||||
EpilogueOutputOp::kCount>::Epilogue;
|
||||
|
||||
/// Define the kernel-level GEMM operator.
|
||||
using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
|
||||
};
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for Turing Architecture
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
@@ -201,6 +264,75 @@ struct DefaultGemm<
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for Ampere Integer Matrix Multiply Interleaved layout
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// Number of Interleaved k
|
||||
int InterleavedK,
|
||||
/// If true, kernel is configured to support serial reduction in the
|
||||
/// epilogue
|
||||
bool SplitKSerial,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator,
|
||||
/// Is Beta zero or not
|
||||
bool IsBetaZero>
|
||||
struct DefaultGemm<
|
||||
ElementA, layout::ColumnMajorInterleaved<InterleavedK>, kAlignmentA,
|
||||
ElementB, layout::RowMajorInterleaved<InterleavedK>, kAlignmentB, ElementC,
|
||||
layout::ColumnMajorInterleaved<InterleavedK>, int32_t,
|
||||
arch::OpClassTensorOp, arch::Sm80, ThreadblockShape, WarpShape,
|
||||
InstructionShape, EpilogueOutputOp, ThreadblockSwizzle, Stages,
|
||||
SplitKSerial, Operator, IsBetaZero> {
|
||||
using LayoutA = layout::ColumnMajorInterleaved<InterleavedK>;
|
||||
using LayoutB = layout::RowMajorInterleaved<InterleavedK>;
|
||||
using LayoutC = layout::ColumnMajorInterleaved<InterleavedK>;
|
||||
|
||||
using ElementAccumulator = int32_t;
|
||||
|
||||
/// Define the threadblock-scoped matrix multiply-accumulate
|
||||
using Mma = typename cutlass::gemm::threadblock::DefaultMma<
|
||||
ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB,
|
||||
ElementAccumulator, LayoutC, arch::OpClassTensorOp, arch::Sm80,
|
||||
ThreadblockShape, WarpShape, InstructionShape, Stages, Operator,
|
||||
true>::ThreadblockMma;
|
||||
|
||||
static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
|
||||
|
||||
/// Define the epilogue
|
||||
using Epilogue = typename cutlass::epilogue::threadblock::
|
||||
DefaultInterleavedEpilogueTensorOp<
|
||||
ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp,
|
||||
64 / sizeof_bits<ElementC>::value, InterleavedK,
|
||||
IsBetaZero>::Epilogue;
|
||||
|
||||
/// Define the kernel-level GEMM operator.
|
||||
using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for Turing Integer Matrix Multiply Interleaved layout
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
@@ -439,6 +571,80 @@ struct DefaultGemm<
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for Ampere
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages
|
||||
int Stages,
|
||||
/// If true, kernel is configured to support serial reduction in the epilogue
|
||||
bool SplitKSerial,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator>
|
||||
struct DefaultGemm<ElementA,
|
||||
LayoutA,
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
kAlignmentB,
|
||||
ElementC,
|
||||
layout::RowMajor,
|
||||
ElementAccumulator,
|
||||
arch::OpClassSimt,
|
||||
arch::Sm80,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
GemmShape<1, 1, 1>,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
SplitKSerial,
|
||||
Operator> {
|
||||
|
||||
/// Define the threadblock-scoped matrix multiply-accumulate
|
||||
using Mma = typename cutlass::gemm::threadblock::DefaultMma<
|
||||
ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB,
|
||||
ElementAccumulator, layout::RowMajor, arch::OpClassSimt, arch::Sm80,
|
||||
ThreadblockShape, WarpShape, GemmShape<1, 1, 1>, Stages,
|
||||
Operator>::ThreadblockMma;
|
||||
|
||||
static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount;
|
||||
static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars");
|
||||
|
||||
/// Define the epilogue
|
||||
using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueSimt<
|
||||
ThreadblockShape,
|
||||
typename Mma::Operator,
|
||||
EpilogueOutputOp,
|
||||
kEpilogueElementsPerAccess
|
||||
>::Epilogue;
|
||||
|
||||
/// Define the kernel-level GEMM operator.
|
||||
using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
/// Partial specialization for SIMT DP4A
|
||||
|
||||
@@ -516,7 +722,6 @@ struct DefaultGemm<int8_t, LayoutA, kAlignmentA, int8_t, LayoutB, kAlignmentB,
|
||||
using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
|
||||
};
|
||||
|
||||
|
||||
#if defined(CUTLASS_ARCH_WMMA_ENABLED)
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
/// Partial specialization for Wmma Gemm Kernel
|
||||
|
||||
@@ -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:
|
||||
@@ -49,7 +49,9 @@
|
||||
#include "cutlass/gemm/kernel/gemm_pipelined.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_core_sm75.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_core_sm70.h"
|
||||
#include "cutlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma.h"
|
||||
#include "cutlass/gemm/threadblock/default_multistage_mma_complex.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_core_simt.h"
|
||||
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
|
||||
#include "cutlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h"
|
||||
@@ -101,6 +103,7 @@ template <
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB,
|
||||
/// Multiply-add operator
|
||||
// (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex)
|
||||
typename Operator,
|
||||
/// If true, kernel is configured to support serial reduction in the epilogue
|
||||
bool SplitKSerial
|
||||
@@ -109,6 +112,64 @@ struct DefaultGemmComplex;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for Ampere Architecture
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// Complex elementwise transformation on A operand
|
||||
ComplexTransform TransformA,
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB,
|
||||
/// Multiply-add operator
|
||||
// (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex)
|
||||
typename Operator,
|
||||
/// If true, kernel is configured to support serial reduction in the epilogue
|
||||
bool SplitKSerial
|
||||
>
|
||||
struct DefaultGemmComplex<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementC,
|
||||
layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp,
|
||||
arch::Sm80, ThreadblockShape, WarpShape, InstructionShape,
|
||||
EpilogueOutputOp, ThreadblockSwizzle, Stages, TransformA, TransformB, Operator, SplitKSerial> {
|
||||
|
||||
/// Define the threadblock-scoped matrix multiply-accumulate
|
||||
using Mma = typename cutlass::gemm::threadblock::DefaultMultistageMmaComplex<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementAccumulator,
|
||||
layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, ThreadblockShape,
|
||||
WarpShape, InstructionShape, Stages, TransformA, TransformB, Operator>::ThreadblockMma;
|
||||
|
||||
/// Define the epilogue
|
||||
using Epilogue =
|
||||
typename cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp<
|
||||
ThreadblockShape, typename Mma::Operator, 1, EpilogueOutputOp,
|
||||
EpilogueOutputOp::kCount, Operator>::Epilogue;
|
||||
|
||||
/// Define the kernel-level GEMM operator.
|
||||
using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace kernel
|
||||
|
||||
@@ -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:
|
||||
@@ -49,6 +49,7 @@
|
||||
|
||||
#include "cutlass/epilogue/threadblock/default_epilogue_planar_complex.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_planar_complex_pipelined.h"
|
||||
#include "cutlass/gemm/threadblock/default_mma_planar_complex_multistage.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -222,6 +223,122 @@ struct DefaultGemmPlanarComplexUniversal<
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Partial specialization for multiple pipeline stages.
|
||||
template <
|
||||
/// Element type for A matrix operand
|
||||
typename ElementA,
|
||||
/// Layout type for A matrix operand
|
||||
typename LayoutA,
|
||||
/// Complex elementwise transformation on A operand
|
||||
ComplexTransform TransformA,
|
||||
/// Access granularity of A matrix in units of elements
|
||||
int kAlignmentA,
|
||||
/// Element type for B matrix operand
|
||||
typename ElementB,
|
||||
/// Layout type for B matrix operand
|
||||
typename LayoutB,
|
||||
/// Complex elementwise transformation on B operand
|
||||
ComplexTransform TransformB,
|
||||
/// Access granularity of B matrix in units of elements
|
||||
int kAlignmentB,
|
||||
/// Element type for C and D matrix operands
|
||||
typename ElementC,
|
||||
/// Layout type for C and D matrix operands
|
||||
typename LayoutC,
|
||||
/// Element type for internal accumulation
|
||||
typename ElementAccumulator,
|
||||
/// Operator class tag
|
||||
typename OperatorClass,
|
||||
/// Tag indicating architecture to tune for
|
||||
typename ArchTag,
|
||||
/// Threadblock-level tile size (concept: GemmShape)
|
||||
typename ThreadblockShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename WarpShape,
|
||||
/// Warp-level tile size (concept: GemmShape)
|
||||
typename InstructionShape,
|
||||
/// Epilogue output operator
|
||||
typename EpilogueOutputOp,
|
||||
/// Threadblock-level swizzling operator
|
||||
typename ThreadblockSwizzle,
|
||||
/// Number of stages used in the pipelined mainloop
|
||||
int Stages,
|
||||
/// Operation performed by GEMM
|
||||
typename Operator
|
||||
>
|
||||
struct DefaultGemmPlanarComplexUniversal<
|
||||
ElementA,
|
||||
LayoutA,
|
||||
TransformA,
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
TransformB,
|
||||
kAlignmentB,
|
||||
ElementC,
|
||||
LayoutC,
|
||||
ElementAccumulator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
EpilogueOutputOp,
|
||||
ThreadblockSwizzle,
|
||||
Stages,
|
||||
Operator,
|
||||
typename std::enable_if<(Stages > 2)>::type
|
||||
> {
|
||||
|
||||
/// Define planar complex valued variants instead
|
||||
using Mma = typename gemm::threadblock::DefaultMmaPlanarComplexMultistage<
|
||||
ElementA,
|
||||
LayoutA,
|
||||
kAlignmentA,
|
||||
ElementB,
|
||||
LayoutB,
|
||||
kAlignmentB,
|
||||
ElementAccumulator,
|
||||
LayoutC,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape,
|
||||
WarpShape,
|
||||
InstructionShape,
|
||||
Stages,
|
||||
TransformA,
|
||||
TransformB,
|
||||
Operator
|
||||
>::ThreadblockMma;
|
||||
|
||||
/// Planar complex epilogue
|
||||
using Epilogue = typename epilogue::threadblock::DefaultEpiloguePlanarComplex<
|
||||
ThreadblockShape,
|
||||
typename Mma::Policy::Operator,
|
||||
OperatorClass,
|
||||
ArchTag,
|
||||
ThreadblockShape::kK / WarpShape::kK,
|
||||
EpilogueOutputOp,
|
||||
EpilogueOutputOp::kCount
|
||||
>::Epilogue;
|
||||
|
||||
/// Define the kernel in terms of the default kernel
|
||||
using GemmKernel = kernel::GemmPlanarComplex<
|
||||
Mma,
|
||||
Epilogue,
|
||||
ThreadblockSwizzle
|
||||
>;
|
||||
|
||||
// Array variant
|
||||
using GemmArrayKernel = kernel::GemmPlanarComplexArray<
|
||||
Mma,
|
||||
Epilogue,
|
||||
ThreadblockSwizzle
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace kernel
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -421,6 +421,13 @@ public:
|
||||
|
||||
cutlass::gemm::GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset();
|
||||
|
||||
// Early exit if CTA is out of range
|
||||
if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
|
||||
params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
int offset_k = 0;
|
||||
int problem_size_k = params.problem_size.k();
|
||||
|
||||
|
||||
@@ -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:
|
||||
@@ -377,6 +377,14 @@ public:
|
||||
ThreadblockSwizzle threadblock_swizzle;
|
||||
|
||||
cutlass::gemm::GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset();
|
||||
|
||||
// Early exit if CTA is out of range
|
||||
if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
|
||||
params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
int batch_idx = threadblock_tile_offset.k();
|
||||
|
||||
int problem_size_m = params.problem_size.m();
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
@@ -71,7 +71,7 @@ public:
|
||||
using OperatorClass = typename Mma::Operator::OperatorClass;
|
||||
using ThreadblockShape = typename Mma::Shape;
|
||||
using WarpShape = typename Mma::Operator::Shape;
|
||||
using InstructionShape = typename Mma::Policy::Operator::Shape;
|
||||
using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
|
||||
using ArchTag = typename Mma::ArchTag;
|
||||
|
||||
static int const kStages = Mma::kStages;
|
||||
@@ -259,9 +259,9 @@ public:
|
||||
Arguments const &args,
|
||||
void *workspace = nullptr) {
|
||||
|
||||
ptr_A = args.ptr_A;
|
||||
ptr_B = args.ptr_B;
|
||||
ptr_C = args.ptr_C;
|
||||
ptr_A = const_cast<void *>(args.ptr_A);
|
||||
ptr_B = const_cast<void *>(args.ptr_B);
|
||||
ptr_C = const_cast<void *>(args.ptr_C);
|
||||
ptr_D = args.ptr_D;
|
||||
|
||||
output_op = args.epilogue;
|
||||
@@ -303,6 +303,10 @@ public:
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
static Status can_implement(Arguments const &args) {
|
||||
return can_implement(args.problem_size);
|
||||
}
|
||||
|
||||
/// Executes one GEMM
|
||||
CUTLASS_DEVICE
|
||||
void operator()(Params const ¶ms, SharedStorage &shared_storage) {
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user