@@ -0,0 +1,404 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2011-2019, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without modification, are not permit-
|
||||
* ted.
|
||||
*
|
||||
* 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 TORT (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 Barrier Operations on SM90+
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cutlass/arch/memory_sm75.h>
|
||||
#include <cute/arch/cluster_sm90.hpp>
|
||||
|
||||
namespace cutlass {
|
||||
/// @brief
|
||||
namespace arch {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && (__CUDACC_VER_MAJOR__ >= 12)
|
||||
#define CUDA_BARRIER_ENABLED 1
|
||||
#else
|
||||
#define CUDA_BARRIER_ENABLED 0
|
||||
#endif
|
||||
|
||||
class NamedBarrier {
|
||||
|
||||
// Data Members:
|
||||
|
||||
// Range = [1 , NUM_THREADS_PER_CTA]
|
||||
// Range % warp-size (i.e 32) == 0
|
||||
uint32_t const num_threads_;
|
||||
|
||||
// Range : [0, 15]
|
||||
uint32_t const id_;
|
||||
|
||||
public:
|
||||
|
||||
CUTLASS_DEVICE
|
||||
NamedBarrier(uint32_t num_threads, uint32_t id = 0)
|
||||
: num_threads_(num_threads), id_(id) {}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void arrive_and_wait() const {
|
||||
NamedBarrier::arrive_and_wait(num_threads_, id_);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void arrive() const {
|
||||
NamedBarrier::arrive(num_threads_, id_);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void sync() const {
|
||||
NamedBarrier::arrive_and_wait();
|
||||
}
|
||||
|
||||
// Static variants
|
||||
CUTLASS_DEVICE
|
||||
static void arrive_and_wait(uint32_t num_threads, uint32_t barrier_id) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
asm volatile("bar.sync %0, %1;" : : "r"(barrier_id), "r"(num_threads));
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
static void arrive(uint32_t num_threads, uint32_t barrier_id) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
asm volatile("bar.arrive %0, %1;" : : "r"(barrier_id), "r"(num_threads));
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
static void sync(uint32_t num_threads, uint32_t barrier_id) {
|
||||
NamedBarrier::arrive_and_wait(num_threads, barrier_id);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Hopper introduces a new cluster-wide barrier which handle with Cluster-wide AW behaviour.
|
||||
// This is an extension to the Ampere AW barriers
|
||||
// Note : Ampere AW Barriers have a larger max-arrive count (2^30) than Hopper AW Barriers (2^20).
|
||||
struct ClusterBarrier {
|
||||
|
||||
using ValueType = uint64_t;
|
||||
|
||||
protected:
|
||||
// Can never be initializated - can only be aliased to smem
|
||||
ValueType barrier_;
|
||||
|
||||
public:
|
||||
|
||||
CUTLASS_DEVICE
|
||||
ClusterBarrier() = delete;
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void init(uint32_t arrive_count) const {
|
||||
ClusterBarrier::init(&this->barrier_, arrive_count);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
uint32_t test_wait(uint32_t phase, uint32_t pred=true) const {
|
||||
return ClusterBarrier::test_wait(&this->barrier_, phase, pred);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void wait(uint32_t phase) const {
|
||||
ClusterBarrier::wait(&this->barrier_, phase);
|
||||
}
|
||||
|
||||
// Barrier arrive on local smem
|
||||
CUTLASS_DEVICE
|
||||
void arrive() const {
|
||||
ClusterBarrier::arrive(&this->barrier_);
|
||||
}
|
||||
|
||||
// Remote SMEM arrive with a perdicate (usually done to pick the thread doing the arrive)
|
||||
CUTLASS_DEVICE
|
||||
void arrive(uint32_t cta_id, uint32_t pred = true ) const {
|
||||
ClusterBarrier::arrive(&this->barrier_, cta_id, pred);
|
||||
}
|
||||
|
||||
//
|
||||
// Static Versions
|
||||
//
|
||||
CUTLASS_DEVICE
|
||||
static void init(ValueType const* smem_ptr, uint32_t arrive_count) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
"mbarrier.init.shared.b64 [%1], %0; \n"
|
||||
"}"
|
||||
:
|
||||
: "r"(arrive_count), "r"(smem_addr));
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Static version of wait - in case we don't want to burn a register
|
||||
CUTLASS_DEVICE
|
||||
static void wait(ValueType const* smem_ptr, uint32_t phase) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
// Arbitrarily large timer value after which try-wait expires and re-tries.
|
||||
uint32_t ticks = 0x989680;
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred P1; \n\t"
|
||||
"LAB_WAIT: \n\t"
|
||||
"mbarrier.try_wait.parity.shared.b64 P1, [%0], %1, %2; \n\t"
|
||||
"@P1 bra.uni DONE; \n\t"
|
||||
"bra.uni LAB_WAIT; \n\t"
|
||||
"DONE: \n\t"
|
||||
"}"
|
||||
:
|
||||
: "r"(smem_addr), "r"(phase), "r"(ticks));
|
||||
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
static uint32_t test_wait(ValueType const* smem_ptr, uint32_t phase, uint32_t pred) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
uint32_t waitComplete;
|
||||
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred P1; \n\t"
|
||||
".reg .pred P2; \n\t"
|
||||
"setp.eq.u32 P2, %3, 1;\n\t"
|
||||
"@P2 mbarrier.test_wait.parity.shared.b64 P1, [%1], %2; \n\t"
|
||||
"selp.b32 %0, 1, 0, P1; \n\t"
|
||||
"}"
|
||||
: "=r"(waitComplete)
|
||||
: "r"(smem_addr), "r"(phase), "r"(pred));
|
||||
|
||||
return waitComplete;
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
return 0;
|
||||
}
|
||||
|
||||
// Static Predicated version of the above - in case we know the address.
|
||||
CUTLASS_DEVICE
|
||||
static void arrive(ValueType const* smem_ptr, uint32_t cta_id, uint32_t pred) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred p;\n\t"
|
||||
".reg .b32 remAddr32;\n\t"
|
||||
"setp.eq.u32 p, %2, 1;\n\t"
|
||||
"@p mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t"
|
||||
"@p mbarrier.arrive.shared::cluster.b64 _, [remAddr32];\n\t"
|
||||
"}"
|
||||
:
|
||||
: "r"(smem_addr), "r"(cta_id), "r"(pred));
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Barrier arrive on local smem
|
||||
CUTLASS_DEVICE
|
||||
static void arrive(ValueType const* smem_ptr) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
uint64_t state = 0;
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
"mbarrier.arrive.shared.b64 %1, [%0];\n\t"
|
||||
"}"
|
||||
:
|
||||
: "r"(smem_addr), "l"(state));
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
static void invalidate(ValueType const* smem_ptr) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
"mbarrier.ival.shared.b64 [%0]; \n\t"
|
||||
"}"
|
||||
:
|
||||
: "r"(smem_addr));
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// SM90 also introduces a new type of cluster-barrier which supports sync.
|
||||
// not just based on Arrive Count, but also transaction count (in bytes)
|
||||
struct ClusterTransactionBarrier : public ClusterBarrier {
|
||||
|
||||
CUTLASS_DEVICE
|
||||
ClusterTransactionBarrier() = delete;
|
||||
|
||||
// Performs an arrive operation + bytes reset
|
||||
CUTLASS_DEVICE
|
||||
void arrive_and_reset_bytes(uint32_t transaction_bytes) const {
|
||||
ClusterTransactionBarrier::arrive_and_reset_bytes(&this->barrier_, transaction_bytes);
|
||||
}
|
||||
|
||||
// Performs an arrive operation + bytes reset
|
||||
CUTLASS_DEVICE
|
||||
void arrive_and_reset_bytes(uint32_t transaction_bytes, uint32_t cta_id) const {
|
||||
ClusterTransactionBarrier::arrive_and_reset_bytes(&this->barrier_, transaction_bytes , cta_id, true);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void commit(uint32_t transaction_bytes, uint32_t pred = 1) const {
|
||||
uint32_t cta_rank = cute::block_rank_in_cluster();
|
||||
ClusterTransactionBarrier::commit(&this->barrier_, cta_rank, transaction_bytes, pred);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void commit(uint32_t dst_cta_id, uint32_t transaction_bytes, uint32_t pred) const {
|
||||
ClusterTransactionBarrier::commit(&this->barrier_, dst_cta_id, transaction_bytes, pred);
|
||||
}
|
||||
|
||||
//
|
||||
// Static Versions
|
||||
//
|
||||
|
||||
// Performs an arrive operation + bytes reset
|
||||
CUTLASS_DEVICE
|
||||
static void arrive_and_reset_bytes(ValueType const* smem_ptr, uint32_t transaction_bytes) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
"mbarrier.arrive.expect_tx.shared.b64 _, [%1], %0; \n\t"
|
||||
"}"
|
||||
:
|
||||
: "r"(transaction_bytes), "r"(smem_addr));
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Performs an arrive operation + bytes reset for a remote cta_id in a Cluster
|
||||
CUTLASS_DEVICE
|
||||
static void arrive_and_reset_bytes(
|
||||
ValueType const* smem_ptr, uint32_t transaction_bytes, uint32_t cta_id, uint32_t pred) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred p;\n\t"
|
||||
".reg .b32 remAddr32;\n\t"
|
||||
"setp.eq.u32 p, %2, 1;\n\t"
|
||||
"@p mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t"
|
||||
"@p mbarrier.arrive.expect_tx.shared::cluster.b64 _, [remAddr32], %3;\n\t"
|
||||
"}"
|
||||
:
|
||||
: "r"(smem_addr), "r"(cta_id), "r"(pred), "r"(transaction_bytes));
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Performs an bytes reset without doing an arrive operation
|
||||
CUTLASS_DEVICE
|
||||
static void reset_bytes(ValueType const* smem_ptr, uint32_t transaction_bytes) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
"mbarrier.expect_tx.shared.b64 [%1], %0; \n\t"
|
||||
"}"
|
||||
:
|
||||
: "r"(transaction_bytes), "r"(smem_addr));
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Increments transaction bytes in the barrier
|
||||
CUTLASS_DEVICE
|
||||
static void commit(
|
||||
ValueType const* smem_ptr, uint32_t dst_cta_id, uint32_t transaction_bytes, uint32_t pred = 1) {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
smem_addr = cute::set_block_rank(smem_addr, dst_cta_id);
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred p;\n\t"
|
||||
"setp.eq.u32 p, %2, 1;\n\t"
|
||||
"@p mbarrier.complete_tx.shared::cluster.relaxed.cluster.b64 [%1], %0;"
|
||||
"}"
|
||||
:
|
||||
: "r"(transaction_bytes), "r"(smem_addr), "r"(pred));
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
// Helps with visibility of barrier init operations across warps / cta / cluster
|
||||
// Available as a separate function so as to batch inits across barriers and fence once
|
||||
// Note : It must be composed with an appropriate sync instruction with the right scope
|
||||
// to ensure visibility eg. __syncthreads() or a cluster_arrive() + cluster_wait()
|
||||
CUTLASS_DEVICE
|
||||
void fence_barrier_init() {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
"fence.mbarrier_init.release.cluster; \n"
|
||||
"}"
|
||||
::);
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Issue a shared memory fence for async operations
|
||||
CUTLASS_DEVICE
|
||||
void fence_view_async_shared() {
|
||||
#if CUDA_BARRIER_ENABLED
|
||||
asm volatile (
|
||||
"{\n\t"
|
||||
"fence.proxy.async.shared::cta; \n"
|
||||
"}"
|
||||
::);
|
||||
#else
|
||||
asm volatile ("brkpt;\n" ::);
|
||||
#endif
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
} // end namespace arch
|
||||
} // end namespace cutlass
|
||||
@@ -36,6 +36,7 @@
|
||||
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cute/arch/util.hpp"
|
||||
|
||||
namespace cutlass {
|
||||
namespace arch {
|
||||
@@ -65,74 +66,13 @@ inline __device__ void ldsm(Array<unsigned, MatrixCount> & D, void const* ptr);
|
||||
#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
|
||||
|
||||
#if ! defined(CUDA_NVVM_GET_SMEM_POINTER_ENABLED)
|
||||
#define CUDA_NVVM_GET_SMEM_POINTER_ENABLED CUDA_NVVM_GET_SMEM_POINTER_SUPPORTED
|
||||
#endif
|
||||
*/
|
||||
|
||||
#if (! defined (__clang__) && __CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2)
|
||||
extern "C" {
|
||||
//
|
||||
// This NVVM intrinsic is subject to change in future versions of CUDA.
|
||||
// Clients should not call it directly. Rather, they should use the
|
||||
// cutlass::arch::ldsm<>() template.
|
||||
//
|
||||
__device__ uint32_t __nvvm_get_smem_pointer(void *);
|
||||
}
|
||||
#endif
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// 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 (__clang__) && 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
|
||||
return static_cast<unsigned>(__cvta_generic_to_shared(ptr));
|
||||
|
||||
#elif (! defined (__clang__) && 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
|
||||
|
||||
CUTLASS_UNUSED(ptr);
|
||||
CUTLASS_NOT_IMPLEMENTED();
|
||||
return 0;
|
||||
|
||||
#endif
|
||||
return cute::cast_smem_ptr_to_uint(ptr);
|
||||
}
|
||||
|
||||
|
||||
/// 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));
|
||||
|
||||
@@ -224,5 +224,4 @@ struct SparseMma;
|
||||
#include "cutlass/arch/mma_sm80.h"
|
||||
#include "cutlass/arch/mma_sparse_sm80.h"
|
||||
#include "cutlass/arch/mma_sm90.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -2166,7 +2166,7 @@ struct Mma<
|
||||
"r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]));
|
||||
|
||||
#else
|
||||
|
||||
|
||||
CUTLASS_UNUSED(a);
|
||||
CUTLASS_UNUSED(b);
|
||||
CUTLASS_UNUSED(c);
|
||||
|
||||
@@ -47,10 +47,21 @@
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#if ((__CUDACC_VER_MAJOR__ > 11) || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 8))
|
||||
#define CUTLASS_ARCH_MMA_SM90_SUPPORTED 1
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
#define CUTLASS_ARCH_MMA_SM90_ENABLED
|
||||
#define CUTLASS_ARCH_MMA_SM90_F64_MMA_SUPPORTED
|
||||
#if (!defined(CUTLASS_ARCH_MMA_SM90_F64_MMA_ENABLED))
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
#define CUTLASS_ARCH_MMA_SM90_F64_MMA_ENABLED
|
||||
#endif
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#if (__CUDACC_VER_MAJOR__ >= 12)
|
||||
#define CUTLASS_ARCH_MMA_SM90_SUPPORTED
|
||||
#if (!defined(CUTLASS_ARCH_MMA_SM90_ENABLED))
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
#define CUTLASS_ARCH_MMA_SM90_ENABLED
|
||||
#endif
|
||||
#endif
|
||||
#endif
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -97,7 +108,7 @@ struct Mma<
|
||||
void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b,
|
||||
FragmentC const &c) const {
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM90_ENABLED)
|
||||
#if defined(CUTLASS_ARCH_MMA_SM90_F64_MMA_ENABLED)
|
||||
|
||||
double const *A = reinterpret_cast<double const *>(&a);
|
||||
double const *B = reinterpret_cast<double const *>(&b);
|
||||
@@ -105,10 +116,73 @@ struct Mma<
|
||||
double const *C = reinterpret_cast<double const *>(&c);
|
||||
double *D = reinterpret_cast<double *>(&d);
|
||||
|
||||
asm volatile("mma.sync.aligned.m16n8k4.row.col.f64.f64.f64.f64 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%7, %8, %9, %10};\n"
|
||||
asm volatile("mma.sync.aligned.m16n8k4.row.col.f64.f64.f64.f64.rn {%0, %1, %2, %3}, {%4, %5}, {%6}, {%7, %8, %9, %10};\n"
|
||||
: "=d"(D[0]), "=d"(D[1]), "=d"(D[2]), "=d"(D[3])
|
||||
: "d"(A[0]), "d"(A[1]),
|
||||
"d"(B[0]),
|
||||
: "d"(A[0]), "d"(A[1]),
|
||||
"d"(B[0]),
|
||||
"d"(C[0]), "d"(C[1]), "d"(C[2]), "d"(C[3]));
|
||||
|
||||
#else
|
||||
CUTLASS_UNUSED(d);
|
||||
CUTLASS_UNUSED(a);
|
||||
CUTLASS_UNUSED(b);
|
||||
CUTLASS_UNUSED(c);
|
||||
CUTLASS_NOT_IMPLEMENTED();
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
/// Matrix Multiply-Add 16x8x8 fp64
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Matrix multiply-add operation: F64 = F64 * F64 + F64
|
||||
template <>
|
||||
struct Mma<
|
||||
gemm::GemmShape<16,8,8>,
|
||||
32,
|
||||
double,
|
||||
layout::RowMajor,
|
||||
double,
|
||||
layout::ColumnMajor,
|
||||
double,
|
||||
layout::RowMajor,
|
||||
OpMultiplyAdd> {
|
||||
|
||||
using Shape = gemm::GemmShape<16,8,8>;
|
||||
|
||||
using ElementA = double;
|
||||
using LayoutA = layout::RowMajor;
|
||||
using FragmentA = Array<double, 4>;
|
||||
|
||||
using ElementB = double;
|
||||
using LayoutB = layout::ColumnMajor;
|
||||
using FragmentB = Array<double, 2>;
|
||||
|
||||
using ElementC = double;
|
||||
using LayoutC = layout::RowMajor;
|
||||
using FragmentC = Array<double, 4>;
|
||||
|
||||
using Operator = OpMultiplyAdd;
|
||||
|
||||
using ArchTag = arch::Sm90;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b,
|
||||
FragmentC const &c) const {
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM90_F64_MMA_ENABLED)
|
||||
|
||||
double const *A = reinterpret_cast<double const *>(&a);
|
||||
double const *B = reinterpret_cast<double const *>(&b);
|
||||
|
||||
double const *C = reinterpret_cast<double const *>(&c);
|
||||
double *D = reinterpret_cast<double *>(&d);
|
||||
|
||||
asm volatile("mma.sync.aligned.m16n8k8.row.col.f64.f64.f64.f64 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n"
|
||||
: "=d"(D[0]), "=d"(d[1]), "=d"(d[2]), "=d"(d[3])
|
||||
: "d"(A[0]), "d"(A[1]), "d"(A[2]), "d"(A[3]),
|
||||
"d"(B[0]), "d"(B[1]),
|
||||
"d"(C[0]), "d"(C[1]), "d"(C[2]), "d"(C[3]));
|
||||
|
||||
#else
|
||||
@@ -118,7 +192,65 @@ struct Mma<
|
||||
CUTLASS_UNUSED(b);
|
||||
CUTLASS_UNUSED(c);
|
||||
CUTLASS_NOT_IMPLEMENTED();
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
/// Matrix Multiply-Add 16x8x16 fp64
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Matrix multiply-add operation: F64 = F64 * F64 + F64
|
||||
template <>
|
||||
struct Mma<
|
||||
gemm::GemmShape<16,8,16>,
|
||||
32,
|
||||
double,
|
||||
layout::RowMajor,
|
||||
double,
|
||||
layout::ColumnMajor,
|
||||
double,
|
||||
layout::RowMajor,
|
||||
OpMultiplyAdd> {
|
||||
|
||||
using Shape = gemm::GemmShape<16,8,16>;
|
||||
|
||||
using ElementA = double;
|
||||
using LayoutA = layout::RowMajor;
|
||||
using FragmentA = Array<double, 8>;
|
||||
|
||||
using ElementB = double;
|
||||
using LayoutB = layout::ColumnMajor;
|
||||
using FragmentB = Array<double, 4>;
|
||||
|
||||
using ElementC = double;
|
||||
using LayoutC = layout::RowMajor;
|
||||
using FragmentC = Array<double, 4>;
|
||||
|
||||
using Operator = OpMultiplyAdd;
|
||||
|
||||
using ArchTag = arch::Sm90;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b,
|
||||
FragmentC const &c) const {
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM90_F64_MMA_ENABLED)
|
||||
|
||||
double const *A = reinterpret_cast<double const *>(&a);
|
||||
double const *B = reinterpret_cast<double const *>(&b);
|
||||
|
||||
double const *C = reinterpret_cast<double const *>(&c);
|
||||
double *D = reinterpret_cast<double *>(&d);
|
||||
|
||||
asm volatile("mma.sync.aligned.m16n8k16.row.col.f64.f64.f64.f64 {%0, %1, %2, %3}, {%4, %5, %6, %7, %8, %9, %10, %11}, {%12, %13, %14, %15}, {%16, %17, %18, %19};\n"
|
||||
: "=d"(D[0]), "=d"(D[1]), "=d"(D[2]), "=d"(D[3])
|
||||
: "d"(A[0]), "d"(A[2]), "d"(A[2]), "d"(A[3]), "d"(A[4]), "d"(A[5]), "d"(A[6]), "d"(A[7])
|
||||
"d"(B[0]), "d"(B[1]), "d"(B[2]), "d"(B[3]),
|
||||
"d"(C[0]), "d"(C[1]), "d"(C[2]), "d"(C[3]));
|
||||
|
||||
#else
|
||||
CUTLASS_NOT_IMPLEMENTED();
|
||||
#endif
|
||||
}
|
||||
};
|
||||
@@ -129,3 +261,4 @@ struct Mma<
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. 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.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 PTX for CTA Reconfiguration
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDACC_VER_MAJOR__ >= 12))
|
||||
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL))
|
||||
#define CUDA_CTA_RECONFIG_ACTIVATED 1
|
||||
#endif
|
||||
#else
|
||||
#define CUDA_CTA_RECONFIG_ACTIVATED 0
|
||||
#endif
|
||||
|
||||
namespace cutlass {
|
||||
namespace arch {
|
||||
|
||||
template<uint32_t RegCount>
|
||||
CUTLASS_DEVICE
|
||||
void warpgroup_reg_alloc(){
|
||||
#if CUDA_CTA_RECONFIG_ACTIVATED
|
||||
asm volatile( "setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount) );
|
||||
#endif
|
||||
}
|
||||
|
||||
template<uint32_t RegCount>
|
||||
CUTLASS_DEVICE
|
||||
void warpgroup_reg_dealloc(){
|
||||
#if CUDA_CTA_RECONFIG_ACTIVATED
|
||||
asm volatile( "setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount) );
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace arch
|
||||
} // namespace cutlass
|
||||
Reference in New Issue
Block a user