CUTLASS 3.5.0 (#1411)
This commit is contained in:
Vijay Thakkar
@@ -260,7 +260,7 @@ private:
|
||||
if (options.vectorize <= 2) return std::make_pair(false, -1);
|
||||
|
||||
// Boundary check.
|
||||
if (i > elements.size() || (i + options.vectorize - 1) > elements.size())
|
||||
if (i > int(elements.size()) || (i + options.vectorize - 1) > int(elements.size()))
|
||||
return std::make_pair(false, -1);
|
||||
|
||||
// Check if either all elements are valid or invalid.
|
||||
|
||||
@@ -94,7 +94,7 @@ __global__ void copy(
|
||||
|
||||
typename Iterator::Fragment fragment;
|
||||
|
||||
for(int i = 0; i < fragment.size(); ++i) {
|
||||
for(size_t i = 0; i < fragment.size(); ++i) {
|
||||
fragment[i] = 0;
|
||||
}
|
||||
|
||||
|
||||
@@ -207,15 +207,15 @@ cudaError_t strided_batched_gemm_nn_reference(
|
||||
|
||||
cudaError_t result = cudaSuccess;
|
||||
|
||||
if (A.size() < lda * k * batch_count) {
|
||||
if (A.size() < size_t(lda * k * batch_count)) {
|
||||
std::cout << "the size of A is too small" << std::endl;
|
||||
return cudaErrorInvalidValue;
|
||||
}
|
||||
if (B.size() < ldb * n) {
|
||||
if (B.size() < size_t(ldb * n)) {
|
||||
std::cout << "the size of B is too small" << std::endl;
|
||||
return cudaErrorInvalidValue;
|
||||
}
|
||||
if (C.size() < ldc * n * batch_count) {
|
||||
if (C.size() < size_t(ldc * n * batch_count)) {
|
||||
std::cout << "the size of C is too small" << std::endl;
|
||||
return cudaErrorInvalidValue;
|
||||
}
|
||||
|
||||
@@ -102,7 +102,7 @@ struct B2bFusedGroupedGemmRun
|
||||
if (dist_kind == cutlass::Distribution::Uniform) {
|
||||
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
view, seed, 2, -2, 0);
|
||||
view, seed, 1, -1, 0);
|
||||
}
|
||||
else if (dist_kind == cutlass::Distribution::Identity) {
|
||||
|
||||
|
||||
@@ -157,35 +157,34 @@ struct B2bGemm {
|
||||
// Data members
|
||||
//
|
||||
|
||||
GemmUniversalMode mode;
|
||||
GemmCoord problem_size_0;
|
||||
GemmCoord problem_size_1;
|
||||
typename B2bMma::IteratorA0::TensorRef ref_A0;
|
||||
typename B2bMma::IteratorB0::TensorRef ref_B0;
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_C0;
|
||||
typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Scale0;
|
||||
typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Bias0;
|
||||
typename B2bMma::IteratorB1::TensorRef ref_B1;
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_C1;
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_D1;
|
||||
int64_t batch_stride_A0;
|
||||
int64_t batch_stride_B0;
|
||||
int64_t batch_stride_B1;
|
||||
int64_t batch_stride_C1;
|
||||
int64_t batch_stride_D1;
|
||||
int64_t batch_stride_Bias0;
|
||||
int64_t batch_stride_Scale0;
|
||||
typename OutputOp0::Params epilogue0;
|
||||
typename OutputOp1::Params epilogue1;
|
||||
int batch_count;
|
||||
GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm;
|
||||
GemmCoord problem_size_0{0,0,0};
|
||||
GemmCoord problem_size_1{0,0,0};
|
||||
typename B2bMma::IteratorA0::TensorRef ref_A0{};
|
||||
typename B2bMma::IteratorB0::TensorRef ref_B0{};
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_C0{};
|
||||
typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Scale0{};
|
||||
typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Bias0{};
|
||||
typename B2bMma::IteratorB1::TensorRef ref_B1{};
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_C1{};
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_D1{};
|
||||
int64_t batch_stride_A0{0};
|
||||
int64_t batch_stride_B0{0};
|
||||
int64_t batch_stride_B1{0};
|
||||
int64_t batch_stride_C1{0};
|
||||
int64_t batch_stride_D1{0};
|
||||
int64_t batch_stride_Bias0{0};
|
||||
int64_t batch_stride_Scale0{0};
|
||||
typename OutputOp0::Params epilogue0 {};
|
||||
typename OutputOp1::Params epilogue1 {};
|
||||
int batch_count{1};
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Default ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Arguments() : mode(mode), problem_size_0(0, 0, 0), problem_size_1(0, 0, 0), batch_count(1) {}
|
||||
Arguments() = default;
|
||||
|
||||
/// Constructs an Arguments structure
|
||||
CUTLASS_HOST_DEVICE
|
||||
@@ -285,47 +284,45 @@ struct B2bGemm {
|
||||
|
||||
/// Parameters structure
|
||||
struct Params {
|
||||
cutlass::gemm::GemmUniversalMode mode;
|
||||
cutlass::gemm::GemmCoord problem_size_0;
|
||||
cutlass::gemm::GemmCoord problem_size_1;
|
||||
cutlass::gemm::GemmCoord grid_tiled_shape;
|
||||
int swizzle_log_tile;
|
||||
typename B2bMma::IteratorA0::Params params_A0;
|
||||
typename B2bMma::IteratorA0::TensorRef ref_A0;
|
||||
typename B2bMma::IteratorB0::Params params_B0;
|
||||
typename B2bMma::IteratorB0::TensorRef ref_B0;
|
||||
typename Epilogue::OutputTileIterator::Params params_C0;
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_C0;
|
||||
typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Scale0;
|
||||
typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Bias0;
|
||||
typename B2bMma::IteratorB1::Params params_B1;
|
||||
typename B2bMma::IteratorB1::TensorRef ref_B1;
|
||||
typename Epilogue::OutputTileIterator::Params params_C1;
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_C1;
|
||||
typename Epilogue::OutputTileIterator::Params params_D1;
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_D1;
|
||||
typename OutputOp0::Params output_op_0;
|
||||
typename OutputOp1::Params output_op_1;
|
||||
int64_t batch_stride_A0;
|
||||
int64_t batch_stride_B0;
|
||||
int64_t batch_stride_B1;
|
||||
int64_t batch_stride_C1;
|
||||
int64_t batch_stride_D1;
|
||||
int64_t batch_stride_Bias0;
|
||||
int64_t batch_stride_Scale0;
|
||||
int *semaphore;
|
||||
int gemm_k_iterations_0;
|
||||
int gemm_k_size_0;
|
||||
int gemm_k_iterations_1;
|
||||
int gemm_k_size_1;
|
||||
cutlass::gemm::GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm;
|
||||
cutlass::gemm::GemmCoord problem_size_0{};
|
||||
cutlass::gemm::GemmCoord problem_size_1{};
|
||||
cutlass::gemm::GemmCoord grid_tiled_shape{};
|
||||
int swizzle_log_tile{0};
|
||||
typename B2bMma::IteratorA0::Params params_A0{};
|
||||
typename B2bMma::IteratorA0::TensorRef ref_A0{};
|
||||
typename B2bMma::IteratorB0::Params params_B0{};
|
||||
typename B2bMma::IteratorB0::TensorRef ref_B0{};
|
||||
typename Epilogue::OutputTileIterator::Params params_C0{};
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_C0{};
|
||||
typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Scale0{};
|
||||
typename B2bMma::IteratorAccumulatorScaleBias::TensorRef ref_Bias0{};
|
||||
typename B2bMma::IteratorB1::Params params_B1{};
|
||||
typename B2bMma::IteratorB1::TensorRef ref_B1{};
|
||||
typename Epilogue::OutputTileIterator::Params params_C1{};
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_C1{};
|
||||
typename Epilogue::OutputTileIterator::Params params_D1{};
|
||||
typename Epilogue::OutputTileIterator::TensorRef ref_D1{};
|
||||
typename OutputOp0::Params output_op_0{};
|
||||
typename OutputOp1::Params output_op_1{};
|
||||
int64_t batch_stride_A0{0};
|
||||
int64_t batch_stride_B0{0};
|
||||
int64_t batch_stride_B1{0};
|
||||
int64_t batch_stride_C1{0};
|
||||
int64_t batch_stride_D1{0};
|
||||
int64_t batch_stride_Bias0{0};
|
||||
int64_t batch_stride_Scale0{0};
|
||||
int *semaphore = nullptr;
|
||||
int gemm_k_iterations_0{0};
|
||||
int gemm_k_size_0{0};
|
||||
int gemm_k_iterations_1{0};
|
||||
int gemm_k_size_1{0};
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(): mode(mode), swizzle_log_tile(0), semaphore(0), gemm_k_iterations_0(0), gemm_k_size_0(0),
|
||||
gemm_k_iterations_1(0), gemm_k_size_1(0) { }
|
||||
Params() = default;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
|
||||
@@ -27,10 +27,14 @@
|
||||
# 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.
|
||||
|
||||
|
||||
set(TEST_STANDARD --m=1024 --n=1024 --k=1024)
|
||||
set(TEST_LARGE_PERFCHECK --m=4096 --n=3456 --k=4096 --perf-check)
|
||||
|
||||
cutlass_example_add_executable(
|
||||
23_ampere_gemm_operand_reduction_fusion
|
||||
ampere_gemm_operand_reduction_fusion.cu
|
||||
TEST_COMMAND_OPTIONS
|
||||
TEST_STANDARD
|
||||
TEST_LARGE_PERFCHECK
|
||||
)
|
||||
|
||||
|
||||
+6
-6
@@ -377,22 +377,22 @@ Result profile(Options const &options) {
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_a.host_view(),
|
||||
1997,
|
||||
ElementInputA(2),
|
||||
ElementInputA(-2),
|
||||
ElementInputA(1),
|
||||
ElementInputA(-1),
|
||||
0); // <- Fill tensor A on host with uniform-distribution random data
|
||||
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_b.host_view(),
|
||||
2003,
|
||||
ElementInputB(2),
|
||||
ElementInputB(-2),
|
||||
ElementInputB(1),
|
||||
ElementInputB(-1),
|
||||
0); // <- Fill tensor B on host with uniform-distribution random data
|
||||
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_c.host_view(),
|
||||
2017,
|
||||
ElementOutput(2),
|
||||
ElementOutput(-2),
|
||||
ElementOutput(1),
|
||||
ElementOutput(-1),
|
||||
0); // <- Fill matrix C on host with uniform-distribution random data
|
||||
cutlass::reference::host::TensorFill(
|
||||
tensor_d.host_view()); // <- fill matrix D on host with zeros
|
||||
|
||||
@@ -789,7 +789,7 @@ public:
|
||||
problem_count_check += bin.second.size();
|
||||
}
|
||||
|
||||
if (problem_count_check != this->problem_count()) {
|
||||
if (problem_count_check != size_t(this->problem_count())) {
|
||||
std::cout << "\n***\nERROR in BINNING LOGIC!\n***\n" << std::endl;
|
||||
}
|
||||
|
||||
|
||||
@@ -113,10 +113,10 @@ cudaError_t CutlassSsyrkNN(
|
||||
>,
|
||||
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>,
|
||||
5, // Stages
|
||||
1, // AligmentA
|
||||
1, // AlignmentA
|
||||
false, // SplitKSerail
|
||||
cutlass::arch::OpMultiplyAdd,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
cutlass::arch::OpMultiplyAdd,
|
||||
cutlass::ComplexTransform::kNone,
|
||||
cutlass::BlasMode::kSymmetric
|
||||
>;
|
||||
|
||||
@@ -149,7 +149,7 @@ cudaError_t CutlassSsyrkNN(
|
||||
//
|
||||
// Launch the CUTLASS SYRK kernel.
|
||||
//
|
||||
|
||||
|
||||
cutlass::Status status = syrk_operator(args);
|
||||
|
||||
//
|
||||
|
||||
@@ -456,7 +456,7 @@ struct Testbed {
|
||||
bool verify_tensor(std::vector<Element> vector_Input, \
|
||||
std::vector<Element> vector_Input_Ref) {
|
||||
|
||||
int64_t size = (vector_Input.size() < vector_Input_Ref.size()) ? vector_Input.size() : vector_Input_Ref.size();
|
||||
auto size = int64_t((vector_Input.size() < vector_Input_Ref.size()) ? vector_Input.size() : vector_Input_Ref.size());
|
||||
float abs_tol = options.tolerance;
|
||||
float rel_tol = options.tolerance;
|
||||
|
||||
|
||||
@@ -454,48 +454,48 @@ struct Testbed {
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_A0.host_view(),
|
||||
options.seed,
|
||||
ElementInputA0(5),
|
||||
ElementInputA0(-5),
|
||||
ElementInputA0(4),
|
||||
ElementInputA0(-4),
|
||||
0
|
||||
);
|
||||
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_B0.host_view(),
|
||||
options.seed + 1,
|
||||
ElementInputB0(5),
|
||||
ElementInputB0(-5),
|
||||
ElementInputB0(4),
|
||||
ElementInputB0(-4),
|
||||
0
|
||||
);
|
||||
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_A1.host_view(),
|
||||
options.seed + 2,
|
||||
ElementInputA1(5),
|
||||
ElementInputA1(-5),
|
||||
ElementInputA1(4),
|
||||
ElementInputA1(-4),
|
||||
0
|
||||
);
|
||||
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_Beta.host_view(),
|
||||
options.seed + 3,
|
||||
ElementInputScaleBias(5),
|
||||
ElementInputScaleBias(-5),
|
||||
ElementInputScaleBias(4),
|
||||
ElementInputScaleBias(-4),
|
||||
0
|
||||
);
|
||||
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_Gamma.host_view(),
|
||||
options.seed + 4,
|
||||
ElementInputScaleBias(5),
|
||||
ElementInputScaleBias(-5),
|
||||
ElementInputScaleBias(4),
|
||||
ElementInputScaleBias(-4),
|
||||
0
|
||||
);
|
||||
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
tensor_Shifted_K.host_view(),
|
||||
options.seed + 5,
|
||||
ElementOutput(5),
|
||||
ElementOutput(-6),
|
||||
ElementOutput(4),
|
||||
ElementOutput(-5),
|
||||
0
|
||||
);
|
||||
|
||||
|
||||
@@ -803,7 +803,7 @@ public:
|
||||
// Use 'D' for the in/out workspace
|
||||
this->block_D.copy_from_device(this->block_C.get());
|
||||
|
||||
for (int i = 0; i < this->options.problem_sizes.size(); ++i) {
|
||||
for (size_t i = 0; i < this->options.problem_sizes.size(); ++i) {
|
||||
cutlass::gemm::GemmCoord const & problem = this->options.problem_sizes[i];
|
||||
int32_t batch_count = 1;
|
||||
int64_t lda = this->lda_host.at(i);
|
||||
@@ -904,10 +904,10 @@ public:
|
||||
// Run profiling loop
|
||||
//
|
||||
|
||||
int last_stream_idx = 0;
|
||||
size_t last_stream_idx = 0;
|
||||
|
||||
for (int iter = 0; iter < this->options.iterations; ++iter) {
|
||||
for (int i = 0; i < this->options.problem_sizes.size(); ++i) {
|
||||
for (size_t i = 0; i < this->options.problem_sizes.size(); ++i) {
|
||||
cutlass::gemm::GemmCoord const & problem = this->options.problem_sizes[i];
|
||||
int32_t batch_count = 1;
|
||||
int64_t lda = this->lda_host.at(i);
|
||||
@@ -1146,7 +1146,7 @@ public:
|
||||
);
|
||||
|
||||
// Initialize the Rank2K object
|
||||
Rank2K rank2k;
|
||||
Rank2K rank2k{};
|
||||
size_t workspace_size = rank2k.get_workspace_size(args);
|
||||
cutlass::DeviceAllocation<uint8_t> workspace(workspace_size);
|
||||
|
||||
|
||||
@@ -40,7 +40,7 @@
|
||||
// Nans & inf detection
|
||||
#define NANCHECK(frag) \
|
||||
{ \
|
||||
for (int _i = 0; _i < frag.size(); ++_i) { \
|
||||
for (size_t _i = 0; _i < frag.size(); ++_i) { \
|
||||
assert(std::isfinite(float(frag[_i]))); \
|
||||
assert(!std::isnan(float(frag[_i]))); \
|
||||
} \
|
||||
@@ -147,7 +147,7 @@ constexpr __string_view __get_type_name() {
|
||||
{ \
|
||||
auto typeStr = __get_type_name<decltype(frag)>(); \
|
||||
PRINT_B0_T0("printing %s (%s)", name, typeStr.data); \
|
||||
for (int _start = 0; _start < frag.size(); _start += 8) { \
|
||||
for (size_t _start = 0; _start < frag.size(); _start += 8) { \
|
||||
PRINT_ACCUM8_T0_L0_START(" ", frag, _start); \
|
||||
} \
|
||||
/*__syncthreads(); \
|
||||
|
||||
@@ -167,58 +167,39 @@ public:
|
||||
// Data members
|
||||
//
|
||||
|
||||
GemmCoord *problem_sizes0;
|
||||
GemmCoord *problem_sizes1;
|
||||
GemmCoord *problem_sizes0{nullptr};
|
||||
GemmCoord *problem_sizes1{nullptr};
|
||||
|
||||
int problem_count;
|
||||
int threadblock_count;
|
||||
int problem_count{0};
|
||||
int threadblock_count{0};
|
||||
|
||||
ElementQ ** ptr_Q;
|
||||
ElementK ** ptr_K;
|
||||
ElementP ** ptr_P;
|
||||
ElementV ** ptr_V;
|
||||
ElementO ** ptr_O;
|
||||
ElementOAccum ** ptr_O_accum;
|
||||
ElementQ ** ptr_Q{nullptr};
|
||||
ElementK ** ptr_K{nullptr};
|
||||
ElementP ** ptr_P{nullptr};
|
||||
ElementV ** ptr_V{nullptr};
|
||||
ElementO ** ptr_O{nullptr};
|
||||
ElementOAccum ** ptr_O_accum{nullptr};
|
||||
|
||||
typename LayoutQ::Stride::LongIndex *ldq;
|
||||
typename LayoutK::Stride::LongIndex *ldk;
|
||||
typename LayoutP::Stride::LongIndex *ldv;
|
||||
typename LayoutO::Stride::LongIndex *ldo;
|
||||
|
||||
// Scale
|
||||
ElementAccumulator scale;
|
||||
typename LayoutQ::Stride::LongIndex *ldq{nullptr};
|
||||
typename LayoutK::Stride::LongIndex *ldk{nullptr};
|
||||
typename LayoutP::Stride::LongIndex *ldv{nullptr};
|
||||
typename LayoutO::Stride::LongIndex *ldo{nullptr};
|
||||
|
||||
// Whether causal masking is to be performed
|
||||
bool causal;
|
||||
bool causal{false};
|
||||
|
||||
// Scale
|
||||
ElementAccumulator scale{0};
|
||||
|
||||
// Only used by device-level operator
|
||||
GemmCoord *host_problem_sizes;
|
||||
GemmCoord *host_problem_sizes{nullptr};
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Default ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
Arguments():
|
||||
problem_count(0),
|
||||
threadblock_count(0),
|
||||
ptr_Q(nullptr),
|
||||
ptr_K(nullptr),
|
||||
ptr_P(nullptr),
|
||||
ptr_V(nullptr),
|
||||
ptr_O(nullptr),
|
||||
ptr_O_accum(nullptr),
|
||||
ldq(nullptr),
|
||||
ldk(nullptr),
|
||||
ldv(nullptr),
|
||||
ldo(nullptr),
|
||||
scale(0),
|
||||
causal(false),
|
||||
host_problem_sizes(nullptr)
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
/// Default ctor
|
||||
Arguments() = default;
|
||||
|
||||
/// Ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
|
||||
@@ -286,7 +286,7 @@ struct Options {
|
||||
// Number of real-valued multiply-adds
|
||||
int64_t fops = int64_t();
|
||||
|
||||
for (int i = 0; i < problem_sizes0.size(); ++i) {
|
||||
for (size_t i = 0; i < problem_sizes0.size(); ++i) {
|
||||
auto const& problem0 = problem_sizes0[i];
|
||||
auto const& problem1 = problem_sizes1[i];
|
||||
for (int row = 0; row < problem0.m(); ++row) {
|
||||
|
||||
@@ -340,7 +340,7 @@ struct Options {
|
||||
// Number of real-valued multiply-adds
|
||||
int64_t fops = int64_t();
|
||||
|
||||
for (int i = 0; i < problem_sizes0.size(); ++i) {
|
||||
for (size_t i = 0; i < problem_sizes0.size(); ++i) {
|
||||
auto const& problem0 = problem_sizes0[i];
|
||||
auto const& problem1 = problem_sizes1[i];
|
||||
|
||||
|
||||
@@ -244,11 +244,13 @@ class CustomMmaMultistage : public CustomMmaBase<Shape_, Policy_, Stages> {
|
||||
CUTLASS_DEVICE
|
||||
bool set_prologue_done(bool value) {
|
||||
prologue_done_ = value;
|
||||
return true;
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
bool set_zero_outside_bounds(bool value) {
|
||||
zero_outside_bounds_ = value;
|
||||
return true;
|
||||
}
|
||||
|
||||
template <bool kLoadA = true, bool kLoadB = true>
|
||||
|
||||
@@ -1446,7 +1446,7 @@ struct AttentionBackwardKernel {
|
||||
uint8_t lane_id) {
|
||||
cutlass::Array<cutlass::uint1b_t, MatmulDOIVJ::Mma::FragmentC::kElements>
|
||||
dropout_keep_mask_doivj;
|
||||
dropout_keep_mask_doivj.fill(1);
|
||||
dropout_keep_mask_doivj.fill(cutlass::uint1b_t{1});
|
||||
const float dropout_scale =
|
||||
kApplyDropout ? 1.0 / (1.0 - p.dropout_prob) : 1.0f;
|
||||
|
||||
@@ -1744,7 +1744,7 @@ struct AttentionBackwardKernel {
|
||||
[&](int accum_m) {},
|
||||
[&](int accum_m /*q*/, int accum_n /*k*/, int idx) {
|
||||
if (zij.at({accum_n, accum_m}) == scalar_t(0)) {
|
||||
dropout_keep_mask_doivj[idx] = cutlass::uint1b_t(0);
|
||||
dropout_keep_mask_doivj[idx] = cutlass::uint1b_t{0};
|
||||
}
|
||||
},
|
||||
[&](int accum_m) {});
|
||||
|
||||
@@ -40,7 +40,6 @@
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
|
||||
#include "cutlass/bfloat16.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
|
||||
@@ -452,7 +452,7 @@ public:
|
||||
// Determine SMEM requirements and waive if not satisfied
|
||||
//
|
||||
|
||||
int smem_size = int(sizeof(typename Gemm::GemmKernel::SharedStorage));
|
||||
size_t smem_size = sizeof(typename Gemm::GemmKernel::SharedStorage);
|
||||
|
||||
cudaDeviceProp properties;
|
||||
int device_idx;
|
||||
@@ -509,7 +509,7 @@ public:
|
||||
);
|
||||
|
||||
// Initialize the GEMM object
|
||||
Gemm gemm;
|
||||
Gemm gemm{};
|
||||
|
||||
result.status = gemm.initialize(args);
|
||||
|
||||
|
||||
@@ -102,7 +102,8 @@ gett_kernel(
|
||||
ElementB, StrideB, 128 / cutlass::sizeof_bits<ElementB>::value,
|
||||
ElementAccumulator,
|
||||
TileShape, Shape<_1,_2,_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
|
||||
@@ -289,7 +289,8 @@ struct ExampleRunner
|
||||
ElementAccumulator,
|
||||
Shape<_128,_128,_64>,
|
||||
Shape<_2,_2,_1>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename EpilogueOpt::SharedStorage)>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename EpilogueOpt::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
|
||||
@@ -39,6 +39,11 @@
|
||||
|
||||
#include "gather_tensor.hpp"
|
||||
|
||||
namespace cutlass {
|
||||
///Forward declaration
|
||||
struct CudaHostAdapter;
|
||||
}
|
||||
|
||||
namespace cutlass::gemm::kernel {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
@@ -143,10 +148,10 @@ public:
|
||||
|
||||
// Kernel entry point API
|
||||
struct Params {
|
||||
GemmUniversalMode mode;
|
||||
ProblemShape problem_shape;
|
||||
MainloopParams mainloop;
|
||||
EpilogueParams epilogue;
|
||||
GemmUniversalMode mode{};
|
||||
ProblemShape problem_shape{};
|
||||
MainloopParams mainloop{};
|
||||
EpilogueParams epilogue{};
|
||||
GatherA gather_A{};
|
||||
GatherB gather_B{};
|
||||
};
|
||||
@@ -191,14 +196,15 @@ public:
|
||||
}
|
||||
|
||||
static
|
||||
int
|
||||
size_t
|
||||
get_workspace_size(Arguments const& args) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
static
|
||||
cutlass::Status
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
|
||||
CudaHostAdapter* cuda_adapter = nullptr) {
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
|
||||
@@ -39,7 +39,7 @@
|
||||
#include "cutlass/epilogue/collective/detail.hpp"
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/numeric/int.hpp"
|
||||
#include "cute/numeric/numeric_types.hpp"
|
||||
|
||||
#include "gather_tensor.hpp"
|
||||
|
||||
|
||||
@@ -393,7 +393,8 @@ private:
|
||||
ElementB, StrideB, 128 / cutlass::sizeof_bits<ElementB>::value,
|
||||
ElementAccumulator,
|
||||
TileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpilogue::SharedStorage)>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
@@ -403,7 +404,8 @@ private:
|
||||
ElementB, StrideBPermute, 128 / cutlass::sizeof_bits<ElementB>::value,
|
||||
ElementAccumulator,
|
||||
TileShapePermute, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename CollectiveEpiloguePermute::SharedStorage)>,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(sizeof(typename CollectiveEpiloguePermute::SharedStorage))>,
|
||||
cutlass::gemm::collective::KernelScheduleAuto
|
||||
>::CollectiveOp;
|
||||
|
||||
|
||||
@@ -37,7 +37,7 @@
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/tensor_view.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cute/numeric/uint128.hpp"
|
||||
#include "cute/numeric/numeric_types.hpp"
|
||||
|
||||
namespace example
|
||||
{
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
|
||||
cutlass_example_add_executable(
|
||||
58_ada_fp8_gemm
|
||||
ada_fp8_gemm.cu
|
||||
)
|
||||
@@ -0,0 +1,826 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 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 Example of running an Ada FP8 GEMM.
|
||||
|
||||
In addition to using FP8 Tensor Core instructions, the Ada FP8 GEMM uses a distinct epilogue
|
||||
that enables additional scaling of operands/outputs, storing a pre-activation-function output
|
||||
tensor (called the "auxiliary" output), and computing the absolute maximum value of the
|
||||
outputs.
|
||||
|
||||
Pseudocode for this epilogue is as follows:
|
||||
|
||||
Aux = ((alpha * scale_a * scale_b) * accumulator) + ((beta * scale_c) * source) + bias
|
||||
D = activation(Aux)
|
||||
|
||||
if Aux is fp8 type:
|
||||
abs_max_output = max( abs(aux) | (for every aux in Aux))
|
||||
Aux = scale_aux * Aux
|
||||
endif
|
||||
|
||||
if D is fp8 type:
|
||||
abs_max_output = max( abs(d) | (for every d in D))
|
||||
D = scale_d * D
|
||||
endif
|
||||
|
||||
Parameter Aux is optionally stored to global memory
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
#include <sstream>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
#include "cutlass/util/command_line.h"
|
||||
#include "cutlass/util/host_tensor.h"
|
||||
#include "cutlass/util/reference/host/gemm_complex.h"
|
||||
#include "cutlass/util/tensor_view_io.h"
|
||||
#include "cutlass/util/distribution.h"
|
||||
#include "cutlass/util/reference/host/tensor_fill.h"
|
||||
#include "cutlass/util/reference/host/tensor_copy.h"
|
||||
#include "cutlass/util/reference/host/tensor_compare.h"
|
||||
#include "cutlass/util/reference/host/tensor_norm.h"
|
||||
#include "cutlass/util/reference/host/gemm.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
#include "cutlass/epilogue/thread/linear_combination_generic_with_scaling.h"
|
||||
#include "cutlass/gemm/device/gemm_universal_with_absmax.h"
|
||||
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/matrix_coord.h"
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
|
||||
|
||||
using ElementA = cutlass::float_e4m3_t;
|
||||
using ElementB = cutlass::float_e4m3_t;
|
||||
using ElementOutput = cutlass::float_e4m3_t;
|
||||
using ElementAuxOutput = ElementOutput;
|
||||
using ElementAccumulator = float;
|
||||
using LayoutA = cutlass::layout::RowMajor;
|
||||
using LayoutB = cutlass::layout::ColumnMajor;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
static int const kStages = 3;
|
||||
static int const kAlignmentA = 16;
|
||||
static int const kAlignmentB = 16;
|
||||
|
||||
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationGenericWithScalingAndAbsMax<
|
||||
cutlass::epilogue::thread::ReLu,
|
||||
ElementOutput,
|
||||
ElementAuxOutput,
|
||||
128 / cutlass::sizeof_bits<ElementOutput>::value,
|
||||
ElementAccumulator,
|
||||
ElementAccumulator
|
||||
>;
|
||||
|
||||
template <typename MathOperator>
|
||||
using Gemm_ = cutlass::gemm::device::GemmUniversalWithAbsMax<
|
||||
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
|
||||
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
|
||||
cutlass::gemm::GemmShape<128, 256, 64>, cutlass::gemm::GemmShape<64, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
|
||||
EpilogueOutputOp, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, kStages,
|
||||
kAlignmentA, kAlignmentB, MathOperator
|
||||
>;
|
||||
|
||||
using ElementAbsmax = typename EpilogueOutputOp::ElementAbsmax;
|
||||
|
||||
|
||||
// Command line options parsing
|
||||
struct Options {
|
||||
|
||||
bool help;
|
||||
bool error;
|
||||
bool reference_check;
|
||||
cutlass::gemm::GemmCoord problem_size;
|
||||
|
||||
int iterations;
|
||||
int warmup_iterations;
|
||||
|
||||
bool scale_A;
|
||||
bool scale_B;
|
||||
bool scale_C;
|
||||
|
||||
float alpha;
|
||||
float beta;
|
||||
|
||||
Options():
|
||||
help(false),
|
||||
error(false),
|
||||
reference_check(false),
|
||||
iterations(20),
|
||||
warmup_iterations(5),
|
||||
scale_A(true),
|
||||
scale_B(true),
|
||||
scale_C(true),
|
||||
alpha(1.f),
|
||||
beta(0.f)
|
||||
{ }
|
||||
|
||||
// Parses the command line
|
||||
void parse(int argc, char const **args) {
|
||||
cutlass::CommandLine cmd(argc, args);
|
||||
|
||||
if (cmd.check_cmd_line_flag("help")) {
|
||||
help = true;
|
||||
return;
|
||||
}
|
||||
|
||||
cmd.get_cmd_line_argument("iterations", iterations, 20);
|
||||
cmd.get_cmd_line_argument("warmup_iterations", warmup_iterations, 5);
|
||||
cmd.get_cmd_line_argument("reference-check", reference_check, false);
|
||||
cmd.get_cmd_line_argument("scale-A", scale_A, true);
|
||||
cmd.get_cmd_line_argument("scale-B", scale_B, true);
|
||||
cmd.get_cmd_line_argument("scale-C", scale_C, true);
|
||||
cmd.get_cmd_line_argument("alpha", alpha, 1.f);
|
||||
cmd.get_cmd_line_argument("beta", beta, 0.f);
|
||||
|
||||
int m, n, k;
|
||||
cmd.get_cmd_line_argument("m", m, 1024);
|
||||
cmd.get_cmd_line_argument("n", n, 1024);
|
||||
cmd.get_cmd_line_argument("k", k, 1024);
|
||||
|
||||
problem_size = cutlass::gemm::GemmCoord{m, n, k};
|
||||
}
|
||||
|
||||
/// Prints the usage statement.
|
||||
std::ostream & print_usage(std::ostream &out) const {
|
||||
|
||||
out << "58_ada_fp8_gemm\n\n"
|
||||
<< " This example executes a GEMM using Ada FP8 Tensor Core operations. In addition to performing\n"
|
||||
<< " a normal GEMM, the kernel performs the following operations:\n"
|
||||
<< " Aux = ((alpha * scale_a * scale_b) * accumulator) + ((beta * scale_c) * source) + bias\n"
|
||||
<< " D = activation(Aux)\n\n"
|
||||
<< " if Aux is fp8:\n"
|
||||
<< " abs_max_output = max( abs(aux) | (for every aux in Aux) )\n"
|
||||
<< " Aux = scale_aux * Aux\n\n"
|
||||
<< " if D is fp8 type:\n"
|
||||
<< " abs_max_output = max( abs(d) | (for every d in D) )\n"
|
||||
<< " D = scale_d * D\n\n"
|
||||
<< "Options:\n\n"
|
||||
<< " --help If specified, displays this usage statement\n\n"
|
||||
<< " --m=<int> Sets the M dimension of the GEMM\n"
|
||||
<< " --n=<int> Sets the N dimension of the GEMM\n"
|
||||
<< " --k=<int> Sets the K dimension of the GEMM\n"
|
||||
<< " --scale-A=<bool> Whether to apply a scaling factor to operand A (default: true)\n"
|
||||
<< " --scale-B=<bool> Whether to apply a scaling factor to operand B (default: true)\n"
|
||||
<< " --scale-C=<bool> Whether to apply a scaling factor to operand C (default: true)\n"
|
||||
<< " --iterations=<int> Number of profiling iterations to perform\n"
|
||||
<< " --warmup-iterations=<int> Number of warmup iterations to perform\n"
|
||||
<< " --reference-check=<bool> If true, performs reference check\n";
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
/// Compute performance in GFLOP/s
|
||||
float gflops(float runtime_s) const {
|
||||
// Two flops per multiply-add
|
||||
return 2.0f * float(problem_size.product()) / float(1.0e9) / runtime_s;
|
||||
}
|
||||
};
|
||||
|
||||
/// Helper class to run the kernel
|
||||
template <typename Gemm>
|
||||
struct TestbedRunner {
|
||||
|
||||
using ElementAccumulator = typename Gemm::ElementAccumulator;
|
||||
using ElementCompute = typename Gemm::GemmKernel::Epilogue::OutputOp::ElementCompute;
|
||||
using ElementScalingFactor = typename Gemm::EpilogueOutputOp::ElementScalingFactor;
|
||||
|
||||
static bool const kScaleAux = Gemm::EpilogueOutputOp::kIsScalingAndAmaxAuxOutputNeeded;
|
||||
static bool const kScaleOutput = Gemm::EpilogueOutputOp::kIsScalingAndAmaxOutputNeeded;
|
||||
|
||||
/// Initialization
|
||||
cutlass::Distribution::Kind init_A;
|
||||
cutlass::Distribution::Kind init_B;
|
||||
cutlass::Distribution::Kind init_C;
|
||||
uint64_t seed;
|
||||
|
||||
cutlass::HostTensor<typename Gemm::ElementA, typename Gemm::LayoutA> tensor_A;
|
||||
cutlass::HostTensor<typename Gemm::ElementB, typename Gemm::LayoutB> tensor_B;
|
||||
cutlass::HostTensor<typename Gemm::ElementC, typename Gemm::LayoutC> tensor_C;
|
||||
cutlass::HostTensor<typename Gemm::EpilogueOutputOp::ElementAuxOutput, typename Gemm::LayoutC> tensor_Aux;
|
||||
cutlass::HostTensor<typename Gemm::EpilogueOutputOp::ElementOutput, typename Gemm::LayoutC> tensor_D;
|
||||
cutlass::HostTensor<typename Gemm::ElementC, typename Gemm::LayoutC> tensor_Vector;
|
||||
cutlass::HostTensor<ElementAccumulator, typename Gemm::LayoutC> tmp_D;
|
||||
cutlass::HostTensor<typename Gemm::EpilogueOutputOp::ElementOutput, typename Gemm::LayoutC> reference_D;
|
||||
cutlass::HostTensor<typename Gemm::EpilogueOutputOp::ElementAuxOutput, typename Gemm::LayoutC> reference_Aux;
|
||||
cutlass::HostTensor<ElementScalingFactor, typename Gemm::LayoutC> scale_A;
|
||||
cutlass::HostTensor<ElementScalingFactor, typename Gemm::LayoutC> scale_B;
|
||||
cutlass::HostTensor<ElementScalingFactor, typename Gemm::LayoutC> scale_C;
|
||||
cutlass::HostTensor<ElementScalingFactor, typename Gemm::LayoutC> scale_D;
|
||||
cutlass::HostTensor<ElementScalingFactor, typename Gemm::LayoutC> scale_Aux;
|
||||
cutlass::HostTensor<ElementAbsmax, typename Gemm::LayoutC> abs_max_Aux;
|
||||
cutlass::HostTensor<ElementAbsmax, typename Gemm::LayoutC> abs_max_D;
|
||||
cutlass::HostTensor<ElementAbsmax, typename Gemm::LayoutC> reference_abs_max_Aux;
|
||||
cutlass::HostTensor<ElementAbsmax, typename Gemm::LayoutC> reference_abs_max_D;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
TestbedRunner(
|
||||
bool scaleA = true,
|
||||
bool scaleB = true,
|
||||
bool scaleC = true,
|
||||
cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform,
|
||||
cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform,
|
||||
cutlass::Distribution::Kind init_C_ = cutlass::Distribution::Uniform,
|
||||
uint64_t seed_ = 2080
|
||||
):
|
||||
init_A(init_A_), init_B(init_B_), init_C(init_C_), seed(seed_) { }
|
||||
|
||||
/// Helper to initialize scaling factors
|
||||
template <typename Element, typename Layout>
|
||||
bool initialize_scale_factor(cutlass::TensorView<Element, Layout> view, uint64_t seed, int bits=0) {
|
||||
cutlass::reference::host::TensorFillRandomUniform(view, seed, double(1.), double(0.), bits);
|
||||
return true;
|
||||
}
|
||||
|
||||
/// Helper to initialize a tensor view
|
||||
template <typename Element, typename Layout>
|
||||
bool initialize_tensor(
|
||||
cutlass::TensorView<Element, Layout> view,
|
||||
cutlass::Distribution::Kind dist_kind,
|
||||
uint64_t seed) {
|
||||
|
||||
if (dist_kind == cutlass::Distribution::Uniform) {
|
||||
|
||||
double scope_max, scope_min;
|
||||
int bits_input = cutlass::sizeof_bits<Element>::value;
|
||||
int bits_output = cutlass::sizeof_bits<typename Gemm::ElementC>::value;
|
||||
|
||||
if (bits_input == 1) {
|
||||
scope_max = 2;
|
||||
scope_min = 0;
|
||||
} else if (bits_input <= 8) {
|
||||
scope_max = 2;
|
||||
scope_min = -2;
|
||||
} else if (bits_output == 16) {
|
||||
scope_max = 5;
|
||||
scope_min = -5;
|
||||
} else {
|
||||
scope_max = 8;
|
||||
scope_min = -8;
|
||||
}
|
||||
|
||||
cutlass::reference::host::TensorFillRandomUniform(
|
||||
view, seed, scope_max, scope_min, 0);
|
||||
}
|
||||
else if (dist_kind == cutlass::Distribution::Identity) {
|
||||
|
||||
cutlass::reference::host::TensorFillIdentity(view);
|
||||
}
|
||||
else if (dist_kind == cutlass::Distribution::Gaussian) {
|
||||
|
||||
cutlass::reference::host::TensorFillRandomGaussian(view, seed, 0, 0.5);
|
||||
}
|
||||
else if (dist_kind == cutlass::Distribution::Sequential) {
|
||||
|
||||
cutlass::reference::host::BlockFillSequential(
|
||||
view.data(), view.capacity());
|
||||
}
|
||||
else {
|
||||
std::cerr << "Not implemented";
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/// Initializes data structures
|
||||
void initialize(const Options& options) {
|
||||
//
|
||||
// Allocate the GEMM workspace
|
||||
//
|
||||
|
||||
tensor_A.resize(options.problem_size.mk());
|
||||
tensor_B.resize(options.problem_size.kn());
|
||||
tensor_C.resize(options.problem_size.mn());
|
||||
tensor_D.resize(options.problem_size.mn());
|
||||
tensor_Vector.resize({1, options.problem_size.n()});
|
||||
reference_D.resize(options.problem_size.mn(), false);
|
||||
tmp_D.resize(options.problem_size.mn(), false);
|
||||
|
||||
initialize_tensor(tensor_A.host_view(), init_A, seed + 2019);
|
||||
initialize_tensor(tensor_B.host_view(), init_B, seed + 2018);
|
||||
initialize_tensor(tensor_C.host_view(), init_C, seed + 2017);
|
||||
initialize_tensor(tensor_Vector.host_view(), init_C, seed + 2020);
|
||||
|
||||
// It is possible to randomly initialize to all zeros, so override this with non-zeros
|
||||
// in the upper left corner of each operand.
|
||||
cutlass::Coord<2> origin(0);
|
||||
tensor_A.host_view().at(origin) = typename Gemm::ElementA(1);
|
||||
tensor_B.host_view().at(origin) = typename Gemm::ElementB(1);
|
||||
tensor_C.host_view().at(origin) = typename Gemm::ElementC(1);
|
||||
tensor_Vector.host_view().at(origin) = typename Gemm::ElementC(1);
|
||||
|
||||
cutlass::reference::host::TensorFill(tensor_D.host_view());
|
||||
cutlass::reference::host::TensorCopy(reference_D.host_view(), tensor_C.host_view());
|
||||
|
||||
tensor_A.sync_device();
|
||||
tensor_B.sync_device();
|
||||
tensor_C.sync_device();
|
||||
tensor_D.sync_device();
|
||||
tensor_Vector.sync_device();
|
||||
|
||||
int scale_bits = 2;
|
||||
if (options.scale_A) {
|
||||
scale_A.resize({1, 1});
|
||||
initialize_scale_factor(scale_A.host_view(), seed + 2021, scale_bits);
|
||||
scale_A.sync_device();
|
||||
}
|
||||
|
||||
if (options.scale_B) {
|
||||
scale_B.resize({1, 1});
|
||||
initialize_scale_factor(scale_B.host_view(), seed + 2022, scale_bits);
|
||||
scale_B.sync_device();
|
||||
}
|
||||
|
||||
if (options.scale_C) {
|
||||
scale_C.resize({1, 1});
|
||||
initialize_scale_factor(scale_C.host_view(), seed + 2023, scale_bits);
|
||||
scale_C.sync_device();
|
||||
}
|
||||
|
||||
if (kScaleOutput) {
|
||||
scale_D.resize({1, 1});
|
||||
initialize_scale_factor(scale_D.host_view(), seed + 2024, scale_bits);
|
||||
scale_D.sync_device();
|
||||
|
||||
abs_max_D.resize({1, 1});
|
||||
cutlass::reference::host::TensorFill(abs_max_D.host_view());
|
||||
abs_max_D.sync_device();
|
||||
|
||||
reference_abs_max_D.resize({1, 1});
|
||||
}
|
||||
|
||||
if (kScaleAux) {
|
||||
tensor_Aux.resize(options.problem_size.mn());
|
||||
cutlass::reference::host::TensorFill(tensor_Aux.host_view());
|
||||
tensor_Aux.sync_device();
|
||||
|
||||
scale_Aux.resize({1, 1});
|
||||
initialize_scale_factor(scale_Aux.host_view(), seed + 2025, scale_bits);
|
||||
scale_Aux.sync_device();
|
||||
|
||||
abs_max_Aux.resize({1, 1});
|
||||
cutlass::reference::host::TensorFill(abs_max_Aux.host_view());
|
||||
abs_max_Aux.sync_device();
|
||||
|
||||
reference_Aux.resize(options.problem_size.mn(), false);
|
||||
reference_abs_max_Aux.resize({1, 1});
|
||||
}
|
||||
}
|
||||
|
||||
/// Compares computed reference with device reference and outputs to a file if incorrect
|
||||
bool compare_reference(const Options& options) {
|
||||
|
||||
tensor_D.sync_host();
|
||||
|
||||
bool passed = cutlass::reference::host::TensorEquals(reference_D.host_view(), tensor_D.host_view());
|
||||
|
||||
if (kScaleAux) {
|
||||
tensor_Aux.sync_host();
|
||||
abs_max_Aux.sync_host();
|
||||
passed &= cutlass::reference::host::TensorEquals(reference_Aux.host_view(), tensor_Aux.host_view());
|
||||
passed &= cutlass::reference::host::TensorEquals(abs_max_Aux.host_view(), reference_abs_max_Aux.host_view());
|
||||
}
|
||||
|
||||
if (kScaleOutput) {
|
||||
abs_max_D.sync_host();
|
||||
passed &= cutlass::reference::host::TensorEquals(abs_max_D.host_view(), reference_abs_max_D.host_view());
|
||||
}
|
||||
|
||||
if (!passed) {
|
||||
std::cerr << "Reference check failed" << std::endl;
|
||||
|
||||
std::string output_file = "testbed_with_amax_errors.txt";
|
||||
std::ofstream file(output_file);
|
||||
|
||||
file
|
||||
<< "problem: " << options.problem_size
|
||||
<< ", alpha: " << options.alpha << ", beta: " << options.beta << "\n\n";
|
||||
|
||||
file
|
||||
<< "A =\n" << tensor_A.host_view()
|
||||
<< "\nB =\n" << tensor_B.host_view()
|
||||
<< "\nC =\n" << tensor_C.host_view()
|
||||
<< "\nVector =\n" << tensor_Vector.host_view()
|
||||
<< "\nScaleA = " << scale_A.host_view()
|
||||
<< "\nScaleB = " << scale_B.host_view()
|
||||
<< "\nScaleC = " << scale_C.host_view()
|
||||
<< "\nScaleD = " << scale_D.host_view()
|
||||
<< "\nScaleAux = " << scale_Aux.host_view()
|
||||
<< "\n\nReference D =\n" << reference_D.host_view()
|
||||
<< "\nComputed D =\n" << tensor_D.host_view();
|
||||
if (kScaleAux) {
|
||||
file
|
||||
<< "\n\nReference Aux =\n" << reference_Aux.host_view()
|
||||
<< "\nComputed Aux =\n" << tensor_Aux.host_view()
|
||||
<< "\n\nReference Absmax Aux = " << reference_abs_max_Aux.host_view()
|
||||
<< "\nComputed Absmax Aux = " << abs_max_Aux.host_view();
|
||||
}
|
||||
if (kScaleOutput) {
|
||||
file
|
||||
<< "\n\nReference Absmax D = " << reference_abs_max_D.host_view()
|
||||
<< "\nComputed Absmax D = " << abs_max_D.host_view();
|
||||
}
|
||||
|
||||
std::cerr << "Dumped results to " << output_file << std::endl;
|
||||
|
||||
}
|
||||
|
||||
return passed;
|
||||
}
|
||||
|
||||
/// Verifies the result is a GEMM
|
||||
bool verify(const Options& options) {
|
||||
|
||||
cutlass::Coord<2> origin(0);
|
||||
ElementCompute scaled_alpha = options.alpha;
|
||||
if (options.scale_A) {
|
||||
scaled_alpha *= scale_A.host_view().at(origin);
|
||||
}
|
||||
if (options.scale_B) {
|
||||
scaled_alpha *= scale_B.host_view().at(origin);
|
||||
}
|
||||
|
||||
ElementCompute scaled_beta = options.beta;
|
||||
if (options.scale_C) {
|
||||
scaled_beta *= scale_C.host_view().at(origin);
|
||||
}
|
||||
|
||||
//
|
||||
// Verify
|
||||
//
|
||||
|
||||
cutlass::reference::host::GemmComplex<
|
||||
typename Gemm::ElementA, typename Gemm::LayoutA,
|
||||
typename Gemm::ElementB, typename Gemm::LayoutB,
|
||||
typename Gemm::ElementC, typename Gemm::LayoutC,
|
||||
ElementCompute, ElementAccumulator, ElementAccumulator
|
||||
>(
|
||||
options.problem_size,
|
||||
scaled_alpha,
|
||||
tensor_A.host_ref(),
|
||||
Gemm::kTransformA,
|
||||
tensor_B.host_ref(),
|
||||
Gemm::kTransformB,
|
||||
scaled_beta,
|
||||
tensor_C.host_ref(),
|
||||
tmp_D.host_ref(),
|
||||
ElementAccumulator(0)
|
||||
);
|
||||
|
||||
ElementCompute tmp_abs_max_Aux(0.);
|
||||
ElementCompute tmp_abs_max_D(0.);
|
||||
|
||||
cutlass::NumericConverter<ElementCompute, typename Gemm::ElementC> cvt_c_to_compute;
|
||||
cutlass::NumericConverter<ElementCompute, ElementAccumulator> cvt_accum_to_compute;
|
||||
cutlass::NumericConverter<ElementAccumulator, ElementCompute> cvt_compute_to_accum;
|
||||
cutlass::NumericConverter<typename Gemm::EpilogueOutputOp::ElementOutput, ElementCompute> cvt_compute_to_d;
|
||||
cutlass::NumericConverter<typename Gemm::EpilogueOutputOp::ElementAuxOutput, ElementCompute> cvt_compute_to_aux;
|
||||
|
||||
cutlass::absolute_value_op<ElementCompute> abs;
|
||||
cutlass::maximum_with_nan_propogation<ElementCompute> max;
|
||||
cutlass::epilogue::thread::ReLu<ElementCompute> act;
|
||||
|
||||
ElementScalingFactor d_scale = kScaleOutput ? scale_D.host_view().at(origin) : ElementScalingFactor(1.);
|
||||
|
||||
for (int m = 0; m < options.problem_size.m(); ++m) {
|
||||
for (int n = 0; n < options.problem_size.n(); ++n) {
|
||||
ElementCompute intermediate = cvt_accum_to_compute(tmp_D.host_view().at({m, n}));
|
||||
ElementCompute bias = cvt_c_to_compute(tensor_Vector.host_view().at({0, n}));
|
||||
ElementCompute aux = intermediate + bias;
|
||||
ElementCompute d = act(aux);
|
||||
tmp_abs_max_Aux = max(abs(aux), tmp_abs_max_Aux);
|
||||
tmp_abs_max_D = max(abs(d), tmp_abs_max_D);
|
||||
reference_D.host_view().at({m, n}) = cvt_compute_to_d(d * d_scale);
|
||||
|
||||
if (kScaleAux) {
|
||||
reference_Aux.host_view().at({m, n}) = cvt_compute_to_aux(aux * scale_Aux.host_view().at(origin));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (kScaleAux) {
|
||||
reference_abs_max_Aux.host_view().at(origin) = cvt_compute_to_accum(tmp_abs_max_Aux);
|
||||
}
|
||||
|
||||
if (kScaleOutput) {
|
||||
reference_abs_max_D.host_view().at(origin) = cvt_compute_to_accum(tmp_abs_max_D);
|
||||
}
|
||||
|
||||
return compare_reference(options);
|
||||
}
|
||||
|
||||
/// Returns true if the CUDA device is sufficient to execute the kernel.
|
||||
bool sufficient() const {
|
||||
|
||||
if (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 4)) {
|
||||
std::cerr << "This example requires CUDA 12.4 or greater." << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t smem_size = sizeof(typename Gemm::GemmKernel::SharedStorage);
|
||||
|
||||
cudaDeviceProp properties;
|
||||
int device_idx;
|
||||
cudaError_t result = cudaGetDevice(&device_idx);
|
||||
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaGetDevice() failed with error: " << cudaGetErrorString(result) << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
result = cudaGetDeviceProperties(&properties, device_idx);
|
||||
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "cudaGetDeviceProperties() failed with error: " << cudaGetErrorString(result) << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
if (properties.major < 8 || (properties.major == 8 && properties.minor < 9)) {
|
||||
std::cerr << "CUTLASS's Ada FP8 GEMM example requires a device of compute capability 89 or higher.\n" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
if (properties.sharedMemPerBlockOptin < smem_size) {
|
||||
std::cerr << "Insufficient shared memory. Need " << smem_size
|
||||
<< ", but device only has " << properties.sharedMemPerBlockOptin << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/// Executes one test
|
||||
bool run(Options& options)
|
||||
{
|
||||
|
||||
// Waive test if insufficient CUDA device
|
||||
if (!sufficient()) {
|
||||
std::cerr << "Insufficient resources to run the kernel." << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
this->initialize(options);
|
||||
|
||||
//
|
||||
// Initialize the GEMM operator
|
||||
//
|
||||
|
||||
typename Gemm::EpilogueOutputOp::Params::ActivationParams activation_params{
|
||||
ElementCompute(options.alpha),
|
||||
ElementCompute(options.beta)
|
||||
};
|
||||
typename Gemm::EpilogueOutputOp::Params epilogue_params{
|
||||
activation_params,
|
||||
scale_A.device_data(),
|
||||
scale_B.device_data(),
|
||||
scale_C.device_data(),
|
||||
scale_D.device_data(),
|
||||
scale_Aux.device_data(),
|
||||
abs_max_Aux.device_data(),
|
||||
abs_max_D.device_data()
|
||||
};
|
||||
|
||||
typename Gemm::Arguments arguments{
|
||||
cutlass::gemm::GemmUniversalMode::kGemm,
|
||||
options.problem_size,
|
||||
/* batch_count = */ 1,
|
||||
epilogue_params,
|
||||
tensor_A.device_data(),
|
||||
tensor_B.device_data(),
|
||||
tensor_C.device_data(),
|
||||
tensor_D.device_data(),
|
||||
tensor_Aux.device_data(),
|
||||
tensor_Vector.device_data(),
|
||||
options.problem_size.m() * options.problem_size.k(),
|
||||
options.problem_size.n() * options.problem_size.k(),
|
||||
options.problem_size.m() * options.problem_size.n(),
|
||||
options.problem_size.m() * options.problem_size.n(),
|
||||
(int)options.problem_size.m(), // Batch stride vector
|
||||
tensor_A.layout().stride(0),
|
||||
tensor_B.layout().stride(0),
|
||||
tensor_C.layout().stride(0),
|
||||
tensor_D.layout().stride(0),
|
||||
(int64_t)0 // Leading dimension of vector. This must be 0
|
||||
};
|
||||
|
||||
Gemm gemm_op;
|
||||
|
||||
cutlass::Status status = gemm_op.can_implement(arguments);
|
||||
if (status != cutlass::Status::kSuccess) {
|
||||
std::cerr << "Gemm::can_implement() failed" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t workspace_size = Gemm::get_workspace_size(arguments);
|
||||
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
|
||||
|
||||
status = gemm_op.initialize(arguments, workspace.get());
|
||||
if (status != cutlass::Status::kSuccess) {
|
||||
std::cerr << "Gemm::initialize() failed" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
//
|
||||
// Run the GEMM
|
||||
//
|
||||
|
||||
status = gemm_op();
|
||||
|
||||
if (status != cutlass::Status::kSuccess) {
|
||||
std::cerr << "Gemm::run() failed" << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
cudaError_t cuda_error = cudaDeviceSynchronize();
|
||||
if (cuda_error != cudaSuccess) {
|
||||
std::cerr << "CUDA error: " << cudaGetErrorString(cuda_error) << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
//
|
||||
// Verify
|
||||
//
|
||||
|
||||
bool passed = true;
|
||||
if (options.reference_check) {
|
||||
passed &= this->verify(options);
|
||||
} else {
|
||||
std::cout << "Skipped reference check" << std::endl;
|
||||
}
|
||||
|
||||
//
|
||||
// Warm up
|
||||
//
|
||||
|
||||
for (int i = 0; i < options.warmup_iterations; ++i) {
|
||||
gemm_op();
|
||||
}
|
||||
|
||||
//
|
||||
// Profile
|
||||
//
|
||||
|
||||
cudaEvent_t events[2];
|
||||
cudaError_t error;
|
||||
for (auto & event : events) {
|
||||
error = cudaEventCreate(&event);
|
||||
if (error != cudaSuccess) {
|
||||
std::cerr << "cudaEventCreate() failed: " << cudaGetErrorString(error) << std::endl;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Record an event at the start of a series of GEMM operations
|
||||
error = cudaEventRecord(events[0]);
|
||||
if (error != cudaSuccess) {
|
||||
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(error) << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
// Run profiling loop
|
||||
for (int iter = 0; iter < options.iterations; ++iter) {
|
||||
gemm_op();
|
||||
}
|
||||
|
||||
// Record an event when the GEMM operations have been launched.
|
||||
error = cudaEventRecord(events[1]);
|
||||
if (error != cudaSuccess) {
|
||||
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(error) << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
// Wait for work on the device to complete.
|
||||
error = cudaEventSynchronize(events[1]);
|
||||
if (error != cudaSuccess) {
|
||||
std::cerr << "cudaEventSynchronize() failed: " << cudaGetErrorString(error) << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
// Measure elapsed runtime
|
||||
float runtime_ms = 0;
|
||||
error = cudaEventElapsedTime(&runtime_ms, events[0], events[1]);
|
||||
if (error != cudaSuccess) {
|
||||
std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(error) << std::endl;
|
||||
return false;
|
||||
}
|
||||
|
||||
// Compute average runtime and GFLOPs.
|
||||
runtime_ms = runtime_ms / float(options.iterations);
|
||||
float gflops = options.gflops(runtime_ms / 1000.0f);
|
||||
|
||||
std::cout << "Problem size: " << options.problem_size.m() << 'x' << options.problem_size.n() << 'x' << options.problem_size.k() << std::endl;
|
||||
std::cout << "Runtime (ms): " << runtime_ms << std::endl;
|
||||
std::cout << "GFLOPs/sec: " << gflops << std::endl;
|
||||
|
||||
// Cleanup
|
||||
for (auto event : events) {
|
||||
(void)cudaEventDestroy(event);
|
||||
}
|
||||
|
||||
return passed;
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
int main(int argc, char const** argv) {
|
||||
|
||||
cudaDeviceProp props;
|
||||
|
||||
cudaError_t error = cudaGetDeviceProperties(&props, 0);
|
||||
if (error != cudaSuccess) {
|
||||
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
if (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 4) ||
|
||||
(props.major != 8 && props.minor != 9)) {
|
||||
|
||||
//
|
||||
// This example requires an NVIDIA Ada-architecture GPU.
|
||||
//
|
||||
|
||||
std::cout
|
||||
<< "CUTLASS's FP8 SM89 example requires a GPU of NVIDIA's Ada architecture "
|
||||
<< "and CUDA toolkit version 12.4 or later.\n";
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
//
|
||||
// Parse options
|
||||
//
|
||||
|
||||
Options options;
|
||||
|
||||
options.parse(argc, argv);
|
||||
|
||||
if (options.help) {
|
||||
options.print_usage(std::cout) << std::endl;
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (options.error) {
|
||||
std::cerr << "Aborting execution." << std::endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
std::cout << "Running GEMM with staged accumulation (OpMultiplyAdd)" << std::endl;
|
||||
std::cout << "=====================================================" << std::endl;
|
||||
TestbedRunner<Gemm_<cutlass::arch::OpMultiplyAdd>> testbed_staged_accum;
|
||||
bool passed = testbed_staged_accum.run(options);
|
||||
|
||||
if (passed) {
|
||||
std::cout << "Passed" << std::endl;
|
||||
} else {
|
||||
std::cout << "Failed" << std::endl;
|
||||
}
|
||||
|
||||
std::cout << "\nRunning GEMM with fast accumulation (OpMultiplyAddFastAccum)" << std::endl;
|
||||
std::cout << "============================================================" << std::endl;
|
||||
TestbedRunner<Gemm_<cutlass::arch::OpMultiplyAddFastAccum>> testbed_fast_accum;
|
||||
passed = testbed_fast_accum.run(options);
|
||||
|
||||
if (passed) {
|
||||
std::cout << "Passed" << std::endl;
|
||||
} else {
|
||||
std::cout << "Failed" << std::endl;
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
# Copyright (c) 2024 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.
|
||||
|
||||
cutlass_example_add_executable(
|
||||
59_ampere_gather_scatter_conv
|
||||
ampere_gather_scatter_conv.cu
|
||||
)
|
||||
|
||||
if (CUTLASS_ENABLE_OPENMP_TESTS AND OpenMP_CXX_FOUND)
|
||||
target_link_libraries(59_ampere_gather_scatter_conv PRIVATE OpenMP::OpenMP_CXX)
|
||||
endif()
|
||||
@@ -0,0 +1,209 @@
|
||||
# Example 59: Ampere gather/scatter convolution
|
||||
|
||||
CuTe and CUTLASS 3.x based Ampere convolution forward propagation kernel capable of operating on both affine and gather/scatter tensors.
|
||||
|
||||
Example executions:
|
||||
```sh
|
||||
./59_ampere_gather_scatter_conv
|
||||
./59_ampere_gather_scatter_conv --n=108
|
||||
./59_ampere_gather_scatter_conv --n=4096 --i=1
|
||||
./59_ampere_gather_scatter_conv --n=1080 --i=1000
|
||||
./59_ampere_gather_scatter_conv --n=131072 --i=1000 --no-check
|
||||
```
|
||||
|
||||
This example demonstrates a few super cool features of CUTLASS and CuTe. It shows off
|
||||
1. A dense conv 3D fprop kernel written as a single file ...
|
||||
2. ... that leverages off-the-shelf CUTLASS collectives to show how custom kernels can use collectives ...
|
||||
3. ... and uses the exact same templated kernel to also stamp out a gather/scatter 3D fprop conv ...
|
||||
4. ... while getting near peak performance of the Ampere class tensor core on Ampere and Ada GPUs ...
|
||||
5. ... by using static cute shapes and strides in case problem shapes are known at compile time.
|
||||
|
||||
## A dense conv 3D fprop kernel written in CUTLASS 3.x and CuTe
|
||||
|
||||
The most common strategy for implementing high performance convolution kernels on the GPU is to transform
|
||||
the activation tensor in such a way that we can perform the computation as a GEMM. This is called the
|
||||
image to column (im2col) transformation. [CUTLASS 2.x implementation of im2col based convolutions is
|
||||
documented separately](../../media/docs/implicit_gemm_convolution.md), and here we consider a fresh approach for CuTe.
|
||||
|
||||
A 3D convolution has the following input tensors:
|
||||
- Activation tensor (Act): `((N,(D,H,W)), (C,(1,1,1)))`
|
||||
- Filter tensor (Flt): `( K, (C,(T,R,S)))`
|
||||
- Output tensor (Out): `((N,(Z,P,Q)), K )`
|
||||
|
||||
Where
|
||||
- N := number of images
|
||||
- DHW := spatial dimensions of the activation tensor
|
||||
- C := channel dimension of the activation tensor
|
||||
- K := channel dimension of the filter and output tensor
|
||||
- TRS := spoke dimensions of the filter tensor
|
||||
- ZPQ := spatial dimensions of the output tensor
|
||||
|
||||
As is evident in the tensor shapes, these cannot be issued to a GEMM just yet, since there is no
|
||||
logical M, N, and K modes we can group the tensor modes into.
|
||||
|
||||
Notice that every spoke of the filter tensor (TRS) will be applied to some (offset) view of the
|
||||
activation tensor, thus expanding the logical size of the activation tensor.
|
||||
Additionally, a similar logical transform of the spatial dimensions can be encoded as a function of the
|
||||
padding, dilations, traversal strides, and filter spokes. This gets us to our im2col transform:
|
||||
|
||||
im2col transform affects the component shapes/strides of the activation tensor in the following way:
|
||||
- ZPQ Shape : changes DHW domain with formula `(1 + (DHW + pad - (((TRS-1) * dilation) + 1)) / traversal_stride)`
|
||||
- TRS Shape : TRS domain instead of `(1,1,1)`
|
||||
- ZPQ Strides : Original DHW strides get `elem_scale()`-ed by traversal strides DHW
|
||||
- TRS Strides : Original DHW strides get `elem_scale()`-ed by dilation DHW
|
||||
|
||||
With this transform applied, we end up with a set of input and output tensors that
|
||||
are logically consistent in their MNK dimensions, thus allowing us to dispatch to a GEMM.
|
||||
im2col activation layout: ((N,(Z,P,Q)), (C,(T,R,S))) // logical (M,K)
|
||||
filter layout : ( K, (C,(T,R,S))) // logical (N,K)
|
||||
output layout : ((N,(Z,P,Q)), K ) // logical (M,N)
|
||||
|
||||
CuTe's layout representation and algebra make these folded tensors easy to represent and manipulate.
|
||||
This is most evident in the reference check code used in this example:
|
||||
|
||||
```cpp
|
||||
for (size_t logical_m = 0; logical_m < size<0>(mOutputRef); ++logical_m) {
|
||||
for (size_t logical_n = 0; logical_n < size<1>(mOutputRef); ++logical_n) {
|
||||
auto accumulator = float(0);
|
||||
for (size_t logical_k = 0; logical_k < size<1>(mStencil); ++logical_k) {
|
||||
accumulator += mStencil(logical_m, logical_k) * mActivation(logical_n, logical_k);
|
||||
}
|
||||
mOutputRef(logical_m, logical_n) = accumulator;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Which succinctly demonstrates how im2col transform allows us to implement convolutions
|
||||
as GEMMs with special layout transformations on the input tensor.
|
||||
|
||||
Note: in the example kernel's implementation we treat activations as the B tensor
|
||||
and filter as the A tensor, thus making their logical dimensions NK and MK respectively.
|
||||
|
||||
## Leveraging CUTLASS collectives off the shelf in a custom kernel
|
||||
|
||||
Now that we have transformed our problem in such a way that allows us to dispatch to a GEMM,
|
||||
we can reuse much of the machinery CUTLASS offers to implement this forward pass convolution
|
||||
operator. CUTLASS decomposes these "moving parts" of GPU linear algebra into reusable,
|
||||
modular software components abstracted by C++ template classes. This example
|
||||
demonstrates how some of the lower layers of the hierarchy can be re-used for custom kernels
|
||||
by writing a custom kernel for convolution that re-uses the Ampere/Ada GEMM collectives
|
||||
from CUTLASS 3.
|
||||
|
||||
A kernel author is free to compose their custom components with any of the existing templates
|
||||
in the CUTLASS hierarchy to leverage existing high performance implementations from the CUTLASS
|
||||
team. In this example, we write a custom kernel layer and compose with an existing collective.
|
||||
However, any of the CUTLASS kernels can be composed with bespoke collectives if the desired
|
||||
customization is a mainloop or epilogue fusion without changes to the grid planning,
|
||||
tile scheduling, load balancing, or thread marshalling.
|
||||
|
||||
## Implementing gather/scatter and dense convolution with the same kernel
|
||||
|
||||
Functionality and correctness of the implemented kernel, as a virtue of using
|
||||
CuTe and off the shelf CUTLASS collectives, only relies on the logical consistency of
|
||||
the layouts of input and output tensors. This means that we can freely change how
|
||||
the logical coordinates of the tensors map into the index space, and even how these dereferences
|
||||
happen. [CUTLASS example 52](../52_hopper_gather_scatter_fusion/) demonstrates this by implementing a custom stride that
|
||||
supports indexed indirection for tensor data accesses. This allows for example 52
|
||||
to implement a GEMM where inputs are gathered and output is scattered based on an index buffer.
|
||||
|
||||
We re-use the same custom stride utilities in this example to implement a convolution kernel
|
||||
that gathers along the NDHW dimensions of the activation tensor and scatters the output along the
|
||||
NZPQ dimensions of the output tensor, treating the channel dimensions as the dense vectors.
|
||||
|
||||
Our dense affine im2col transformed activation tensor:
|
||||
|
||||
```cpp
|
||||
// im2col transformed activation layout: ((nzpq), (ctrs)) => idx
|
||||
auto xformed_act_layout = make_layout(
|
||||
make_shape (make_shape ( N, Z, P, Q), make_shape ( C, T, R, S)),
|
||||
make_stride(make_stride(D*H*W*C, H*W*C, W*C, C), make_stride(_1{}, H*W*C, W*C, C)));
|
||||
```
|
||||
|
||||
now becomes a composed layout that uses `IndexedGather`:
|
||||
|
||||
```cpp
|
||||
// Inner layout of the composition:
|
||||
// ((nzpq), (csrt)) => (idx_buffer_idx, dense_offset)
|
||||
auto EG = E<0>{}; // Gather basis (1,0) (idx_buffer_idx)
|
||||
auto EC = E<1>{}; // Contiguous basis (0,1) (dense_offset)
|
||||
auto xformed_act_logical_inner = make_layout(
|
||||
make_shape (make_shape ( N, Z, P, Q), make_shape ( C, T, R, S)),
|
||||
make_stride(make_stride(D*H*W*EG, H*W*EG, W*EG, EG), make_stride(EC, H*W*EG, W*EG, EG)));
|
||||
|
||||
// Outer layout of the composition:
|
||||
// (idx_buffer_idx, dense_offset) => idx
|
||||
// IndexedGather obtains idx by applying (gmem_base_ptr + gather_idx_buf[idx_buffer_idx] + dense_offset)
|
||||
auto xformed_act_gather_outer = make_layout(
|
||||
make_shape(_1{},_1{}),
|
||||
make_stride(CustomStride{IndexedGather{gather_idx_buf}, C}, _1{}));
|
||||
|
||||
// Compose the inner and outer layouts
|
||||
// ((nzpq), (ctrs)) => idx
|
||||
auto xformed_act_composed_layout = composition(
|
||||
xformed_act_gather_outer,
|
||||
make_arithmetic_tuple(_0{}, _0{}),
|
||||
xformed_act_logical_inner);
|
||||
```
|
||||
|
||||
Here, we create a composed layout whose inner layout has the same logical MK shape as earlier,
|
||||
but with an outer layout that uses the custom strides with an index buffer to access memory with
|
||||
indirections. A custom stride requires two inputs to compute the index that a certain coordinate maps to:
|
||||
the index buffer offset and the dense offset into the vector. This entails that our inner layout
|
||||
(the one with the logical MK shape) has a rank-2 codomain `(idx_buffer_idx, dense_offset)`.
|
||||
We can set up such a layout with scaled basis strides, which allow us to map a domain onto a
|
||||
codomain with multiple orthogonal bases. The two codomain basis are the
|
||||
index buffer offsets (rank 0 basis), and the dense vector offsets (rank 1 basis).
|
||||
A similar composed layout is set up for the output scatter tensor.
|
||||
|
||||
This tensor still has a logical MK shape and is backed by a CuTe layout, which means we can still
|
||||
tile, partition, and otherwise manipulate it with CuTe's layout algebra in the same way we would any
|
||||
other tensor. Substituting the activation tensor's affine layout for this gather layout requires
|
||||
no changes to the implementation of the kernel whatsoever. Everything composes. This example
|
||||
stamps out a dense 3D convolution as well as gather/scatter 3D convolution using the same kernel template,
|
||||
with the only difference between them being the layouts of the input and output tensors.
|
||||
|
||||
Convolutions are just a special case of tensor contractions, and as [example 51](../51_hopper_gett)
|
||||
demonstrates, the exact same collective used in this example can also be used to implement arbitrary GETTs.
|
||||
Of course, this also means that the same kernel can implement gather/scatter GETTs as well!
|
||||
|
||||
This demonstrates the composition power of not just CuTe, but also CUTLASS 3's two level
|
||||
micro kernel abstraction. A single highly tuned temporal micro-kernel (collective) can be implemented once
|
||||
and applied to compute dense GETTs, gather/scatter GETTs, dense convolutions, and gather/scatter convolutions.
|
||||
|
||||
## Peak performance on Ampere and Ada GPUs by leveraging domain specific knowledge
|
||||
|
||||
Often, when implementing custom kernels, a user has more knowledge of the problem domain that can be
|
||||
exploited to deliver higher performance than otherwise could be through general kernels. In this example
|
||||
we presume that the shape of each of the images (DHWC dimensions) as well as the filter (TRS) are available
|
||||
a-priori and that the tile shape evenly divides the problem. Number of images (N) is still left as a runtime
|
||||
parameter.
|
||||
|
||||
Knowing the extents of our tensors at compile time allows us to encode them as static cute shapes rather than
|
||||
a dynamic problem shape, resulting in the elimination of most of the index computation instructions such as
|
||||
expensive div/mods. Knowing that the problem shape is divisible by the tile shape allows us to use the
|
||||
Ampere collective that does not perform predication on global memory loads, further reducing overheads
|
||||
and allowing us to achieve near peak performance on RTX Ampere and Ada GPUs.
|
||||
|
||||
Running this example on an RTX 3080Ti prints the following performance numbers (some output culled for brevity):
|
||||
|
||||
```
|
||||
$> ./examples/59_ampere_gather_scatter_conv/59_ampere_gather_scatter_conv --n=131072 --i=128 --no-check
|
||||
Ampere convolution forward propogation kernel supporting both affine and gather/scatter tensors.
|
||||
|
||||
Allocating tensors ... done.
|
||||
Initializing data ... done.
|
||||
Initializing gather/scatter index buffers ... done.
|
||||
|
||||
Running dense fprop kernel
|
||||
Conv TFLOP count = 0.927713
|
||||
Conv dense perf: 31.027376ms | TFLOP/s = 29.899819
|
||||
|
||||
Running gather/scatter fprop kernel
|
||||
Conv TFLOP count = 0.927713
|
||||
Conv gather/scatter perf: 28.973721ms | TFLOP/s = 32.019117
|
||||
```
|
||||
|
||||
With this in mind, this example kernel has the following limitations:
|
||||
- This example kernel only supports dynamic image count, all other conv problem shape must be defined as `cute::Constant<>`s
|
||||
- Problem shapes (including dynamic image count `N`) must be evenly divisible by the tile shape
|
||||
- It does not perform fp32->tf32 numeric conversion, gmem inputs must be rounded to tf32 already
|
||||
@@ -0,0 +1,320 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#pragma once
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
#include "cute/atom/copy_atom.hpp"
|
||||
#include <random>
|
||||
|
||||
#include "cutlass/util/print_error.hpp"
|
||||
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/gemm/collective/collective_mma.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
struct AmpereUnpredicatedFprop {
|
||||
//
|
||||
// Static config for conv problem shape
|
||||
//
|
||||
using D = _6;
|
||||
using H = _4;
|
||||
using W = _4;
|
||||
|
||||
using T = _3;
|
||||
using R = _3;
|
||||
using S = _3;
|
||||
|
||||
using Z = _4;
|
||||
using P = _2;
|
||||
using Q = _2;
|
||||
|
||||
using C = _64;
|
||||
using K = _128;
|
||||
|
||||
// Tiler config
|
||||
using Tiler_K = decltype(cute::min(K{}, _128{}));
|
||||
using Tiler_C = decltype(cute::min(C{}, _32{}));
|
||||
using Tiler_N = _4;
|
||||
using TileM = Tiler_K;
|
||||
using TileN = Shape<Tiler_N, Z, P, Q>;
|
||||
using TileK = Shape<Tiler_C,_1,_1,_1>;
|
||||
using PIPE = _3;
|
||||
using TilerFlt = Shape<TileM, TileK>;
|
||||
using TilerAct = Shape<TileN, TileK>;
|
||||
using TilerOut = Shape<TileM, TileN>;
|
||||
|
||||
using TileSizeM = Int<size(TileM{})>;
|
||||
using TileSizeN = Int<size(TileN{})>;
|
||||
using TileSizeK = Int<size(TileK{})>;
|
||||
static constexpr int Stages = PIPE::value;
|
||||
|
||||
using ElementFlt = tfloat32_t;
|
||||
using ElementAct = tfloat32_t;
|
||||
using ElementOut = float;
|
||||
|
||||
using TiledMma = TiledMMA<
|
||||
MMA_Atom<SM80_16x8x8_F32TF32TF32F32_TN>,
|
||||
Layout<Shape<_2,_2,_1>>,
|
||||
Tile<_32,_32,Underscore>>;
|
||||
|
||||
static constexpr int MaxThreadsPerBlock = size(TiledMma{});
|
||||
static constexpr int MinBlocksPerMultiprocessor = 1;
|
||||
|
||||
union SharedStorage {
|
||||
struct {
|
||||
ElementFlt sAMatrix[size(TileM{}) * size(TileK{}) * size(PIPE{})];
|
||||
ElementAct sBMatrix[size(TileN{}) * size(TileK{}) * size(PIPE{})];
|
||||
} mainloop;
|
||||
|
||||
struct {
|
||||
ElementOut sCMatrix[size(TileM{}) * size(TileN{})];
|
||||
} epilogue;
|
||||
};
|
||||
|
||||
//
|
||||
// Stencil tensor
|
||||
//
|
||||
|
||||
using GmemLayoutFlt = decltype(make_ordered_layout(
|
||||
Shape< K, Shape< C, T, R, S>>{},
|
||||
tuple<_4, tuple<_0,_3,_2,_1>>{}));
|
||||
|
||||
// We have 64 elements * 32b each in the major mode that we can vectorize
|
||||
// Max vector size is 128b, so lay 16 threads along the major mode with a vector size of 4
|
||||
// Rest along the minor mode
|
||||
using GmemTiledCopyFlt = decltype(make_tiled_copy(
|
||||
Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, ElementFlt>{},
|
||||
Layout<Shape <_16, _8>,
|
||||
Stride< _8, _1>>{},
|
||||
Layout<Shape < _1, _4>>{}));
|
||||
|
||||
// Following layout is also correct, but trades off dynamic strides in the slice for bank conflict free accesses
|
||||
// using SmemLayoutFlt = decltype(
|
||||
// composition(Swizzle<3,2,3>{},
|
||||
// make_ordered_layout(
|
||||
// Shape<TileSizeM,TileSizeK,PIPE>{},
|
||||
// tuple< _1, _0, _2>{})));
|
||||
|
||||
using SmemLayoutAtomFlt = decltype(
|
||||
composition(Swizzle<1,2,3>{},
|
||||
Layout<Shape <_8,Shape <_4, _2>>,
|
||||
Stride<_4,Stride<_1,_32>>>{}));
|
||||
|
||||
using SmemCopyAtomFlt = Copy_Atom<SM75_U32x4_LDSM_N, ElementFlt>;
|
||||
|
||||
//
|
||||
// Activation tensor
|
||||
//
|
||||
|
||||
// Activation tensor is major in the contraction mode, so vectorize that mode first
|
||||
// Then lay out the rest of the threads along the other mode
|
||||
using GmemTiledCopyAct = decltype(make_tiled_copy(
|
||||
Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, ElementAct>{},
|
||||
Layout<Shape <_16, _8>,
|
||||
Stride< _8, _1>>{},
|
||||
Layout<Shape < _1, _4>>{}));
|
||||
|
||||
// Following layout is also correct, but trades off dynamic strides in the slice for bank conflict free accesses
|
||||
// using SmemLayoutAct = decltype(
|
||||
// composition(Swizzle<3,2,3>{},
|
||||
// make_ordered_layout(
|
||||
// Shape<TileSizeN,TileSizeK,PIPE>{},
|
||||
// tuple< _1, _0, _2>{})));
|
||||
|
||||
using SmemLayoutAtomAct = decltype(
|
||||
composition(Swizzle<1,2,3>{},
|
||||
Layout<Shape <_8,Shape <_4, _2>>,
|
||||
Stride<_4,Stride<_1,_32>>>{}));
|
||||
|
||||
using SmemCopyAtomAct = Copy_Atom<SM75_U32x4_LDSM_N, ElementAct>;
|
||||
|
||||
//
|
||||
// Output tensor
|
||||
//
|
||||
|
||||
using GmemTiledCopyOut = decltype(make_tiled_copy(
|
||||
Copy_Atom<UniversalCopy<uint128_t>, ElementAct>{},
|
||||
Layout<Shape <_8, _16>,
|
||||
Stride<_1, _8>>{},
|
||||
Layout<Shape <_4, _1>>{}));
|
||||
|
||||
using SmemCopyAtomOut = Copy_Atom<UniversalCopy<uint32_t>, ElementOut>;
|
||||
|
||||
// This can be optimized to make accesses BCF, but we use a col-major layout here to show off composability
|
||||
using SmemLayoutOut = Layout<Shape<TileSizeM, TileSizeN>>;
|
||||
|
||||
//
|
||||
// Conv functor
|
||||
//
|
||||
template <class EngineFlt, class TensorActivation, class TensorOutput>
|
||||
void __device__
|
||||
operator()(cute::Tensor<EngineFlt, GmemLayoutFlt> mFlt, // ( K, (C,T,R,S))
|
||||
TensorActivation mAct, // ((N,Z,P,Q), (C,T,R,S))
|
||||
TensorOutput mOut, // ( K, (N,Z,P,Q))
|
||||
char* smem_buf) const {
|
||||
using namespace cute;
|
||||
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveMma<
|
||||
cutlass::gemm::MainloopSm80CpAsyncUnpredicated<PIPE::value>,
|
||||
Shape<TileM,TileN,TileK>,
|
||||
ElementFlt,
|
||||
Underscore, // Ignore the stride, we are passing full cute::Tensor to operator()
|
||||
ElementAct,
|
||||
Underscore, // Ignore the stride, we are passing full cute::Tensor to operator()
|
||||
TiledMma,
|
||||
GmemTiledCopyFlt,
|
||||
SmemLayoutAtomFlt,
|
||||
SmemCopyAtomFlt,
|
||||
cute::identity,
|
||||
GmemTiledCopyAct,
|
||||
SmemLayoutAtomAct,
|
||||
SmemCopyAtomAct,
|
||||
cute::identity>;
|
||||
|
||||
TiledMma tiled_mma;
|
||||
Tensor accum = partition_fragment_C(tiled_mma, TilerOut{});
|
||||
clear(accum);
|
||||
|
||||
// Set up tensors
|
||||
// NOTE: blockIdx.x projects onto act-NDHW mode, y along the flt-K mode for the sake of higher dynamic range in NDHW
|
||||
Tensor gA_mk = local_tile(mFlt, TilerFlt{}, make_coord(_,_)); // (BLK_M,BLK_K,m',k')
|
||||
Tensor gB_nk = local_tile(mAct, TilerAct{}, make_coord(_,_)); // (BLK_N,BLK_K,n',_1)
|
||||
Tensor gC_mn = local_tile(mOut, TilerOut{}, make_coord(_,_)); // (BLK_M,BLK_N,m',n')
|
||||
|
||||
// Compute m_coord and n_coord with their post-tiled shapes
|
||||
auto m_coord = idx2crd(int(blockIdx.y), shape<2>(gA_mk));
|
||||
auto n_coord = idx2crd(int(blockIdx.x), shape<2>(gB_nk));
|
||||
Tensor gA = gA_mk(_,_,m_coord,_); // (BLK_M,BLK_K,k')
|
||||
Tensor gB = gB_nk(_,_,n_coord,_); // (BLK_N,BLK_K,_1)
|
||||
Tensor gC = gC_mn(_,_,m_coord,n_coord); // (BLK_M,BLK_N)
|
||||
|
||||
auto k_tile_iter = cute::make_coord_iterator(size<2>(gA));
|
||||
int k_tile_count = size<2>(gA);
|
||||
|
||||
CollectiveMainloop collective_mma;
|
||||
collective_mma(
|
||||
accum,
|
||||
gA,
|
||||
gB,
|
||||
accum,
|
||||
k_tile_iter, k_tile_count,
|
||||
Underscore{}, // no residue since we do not support predication
|
||||
threadIdx.x,
|
||||
smem_buf);
|
||||
|
||||
//
|
||||
// Epilogue
|
||||
//
|
||||
SharedStorage& storage = *reinterpret_cast<SharedStorage*>(smem_buf);
|
||||
Tensor sC = make_tensor(make_smem_ptr(&storage.epilogue.sCMatrix[0]), SmemLayoutOut{});
|
||||
|
||||
auto smem_tiled_copy_C = make_tiled_copy_C(SmemCopyAtomOut{}, tiled_mma);
|
||||
auto smem_thr_copy_C = smem_tiled_copy_C.get_slice(threadIdx.x);
|
||||
auto tCrC = smem_thr_copy_C.retile_S(accum);
|
||||
auto tCsC = smem_thr_copy_C.partition_D(sC);
|
||||
copy(smem_tiled_copy_C, tCrC, tCsC);
|
||||
|
||||
__syncthreads();
|
||||
|
||||
GmemTiledCopyOut gmem_tiled_copy_C;
|
||||
auto gmem_thr_copy_C = gmem_tiled_copy_C.get_slice(threadIdx.x);
|
||||
auto tDsC = gmem_thr_copy_C.partition_S(sC);
|
||||
auto tDgC = gmem_thr_copy_C.partition_D(gC);
|
||||
copy(gmem_tiled_copy_C, tDsC, tDgC);
|
||||
|
||||
#if 0
|
||||
if (thread0()) {
|
||||
print("mAct = "); print(mAct); print('\n');
|
||||
print("mFlt = "); print(mFlt); print('\n');
|
||||
print("mOut = "); print(mOut); print('\n');
|
||||
print("gA = "); print(gA); print('\n');
|
||||
print("gB = "); print(gB); print('\n');
|
||||
print("gC = "); print(gC); print('\n');
|
||||
print("sA = "); print(sA.layout()); print('\n');
|
||||
print("sB = "); print(sB.layout()); print('\n');
|
||||
print("sC = "); print(sC.layout()); print('\n');
|
||||
print("tAgA = "); print(tAgA.layout()); print('\n');
|
||||
print("tBgB = "); print(tBgB.layout()); print('\n');
|
||||
print("tAsA = "); print(tAsA.layout()); print('\n');
|
||||
print("tBsB = "); print(tBsB.layout()); print('\n');
|
||||
print("tCsA = "); print(tCsA.layout()); print('\n');
|
||||
print("tCsB = "); print(tCsB.layout()); print('\n');
|
||||
print("tCrC = "); print(tCrC.layout()); print('\n');
|
||||
print("tCsC = "); print(tCsC.layout()); print('\n');
|
||||
print("tDsC = "); print(tDsC.layout()); print('\n');
|
||||
print("tDgC = "); print(tDgC.layout()); print('\n');
|
||||
print("gmem tiled copy A = "); print(gmem_tiled_copy_A); print('\n');
|
||||
print("gmem tiled copy B = "); print(gmem_tiled_copy_B); print('\n');
|
||||
print("gmem tiled copy C = "); print(gmem_tiled_copy_C); print('\n');
|
||||
print("k_tile_count = "); print(size<2>(gA)); print('\n');
|
||||
print("k_tile_iter = "); print(*k_tile_iter); print('\n');
|
||||
print("K_BLOCK_MAX = "); print(K_BLOCK_MAX); print('\n');
|
||||
}
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
template <class TensorFlt, class TensorAct, class TensorOut>
|
||||
inline int
|
||||
fprop_reference(
|
||||
TensorFlt mStencil, // Logical MK: ( K, (C,T,R,S))
|
||||
TensorAct mActivation, // Logical NK: ((N,Z,P,Q), (C,T,R,S))
|
||||
TensorOut mOutput, // Logical MN: ( K, (N,Z,P,Q))
|
||||
TensorOut mOutputRef) {
|
||||
int32_t N = size<1,0>(mOutputRef);
|
||||
int32_t Z = size<1,1>(mOutputRef);
|
||||
int32_t P = size<1,2>(mOutputRef);
|
||||
int32_t Q = size<1,3>(mOutputRef);
|
||||
int32_t T = size<1,3>(mStencil);
|
||||
int32_t R = size<1,2>(mStencil);
|
||||
int32_t S = size<1,1>(mStencil);
|
||||
int32_t C = size<1,0>(mStencil);
|
||||
|
||||
size_t K = static_cast<size_t>(size<0>(mOutputRef));
|
||||
size_t NZPQ = static_cast<size_t>(size<1>(mOutputRef));
|
||||
size_t CTRS = static_cast<size_t>(size<1>(mStencil));
|
||||
|
||||
#if defined(_OPENMP)
|
||||
#pragma omp parallel for
|
||||
#endif
|
||||
for (size_t logical_m = 0; logical_m < K; ++logical_m) {
|
||||
for (size_t logical_n = 0; logical_n < NZPQ; ++logical_n) {
|
||||
auto accumulator = float(0);
|
||||
for (size_t logical_k = 0; logical_k < CTRS; ++logical_k) {
|
||||
accumulator += mStencil(logical_m, logical_k) * mActivation(logical_n, logical_k);
|
||||
}
|
||||
mOutputRef(logical_m, logical_n) = accumulator;
|
||||
}
|
||||
}
|
||||
|
||||
return print_relative_error(mOutput, mOutputRef, /*print_verbose*/ false, /*print_error*/ true, /*error_margin*/ 0.01);
|
||||
}
|
||||
@@ -0,0 +1,392 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2024 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 Example demonstrating CuTe and CUTLASS 3.x based Ampere convolution forward propogation kernel
|
||||
capable of operating on both affine and gather/scatter tensors.
|
||||
|
||||
This example demonstartes a few super cool features of CUTLASS and CuTe. It shows off
|
||||
1. A dense conv 3D fprop kernel written as a single file ...
|
||||
2. ... that leverages off the shelf CUTLASS collectives to show how custom kernels can use collectives ...
|
||||
3. ... and uses the exact same templated kernel to also stamp out a gather/scatter 3D fprop conv ...
|
||||
4. ... while getting near peak performance of the Ampere class tensor core on Ampere and Ada GPUs ...
|
||||
5. ... by using static cute shapes and strides in case problem shapes are known at compile time.
|
||||
|
||||
Full documentation for this example can be found within the README.md file in this directory.
|
||||
|
||||
Example executions:
|
||||
./59_ampere_gather_scatter_conv
|
||||
./59_ampere_gather_scatter_conv --n=108
|
||||
./59_ampere_gather_scatter_conv --n=4096 --i=1
|
||||
./59_ampere_gather_scatter_conv --n=1080 --i=1000
|
||||
./59_ampere_gather_scatter_conv --n=131072 --i=1000 --no-check
|
||||
*/
|
||||
|
||||
#include <thrust/sequence.h>
|
||||
#include <thrust/universal_vector.h>
|
||||
|
||||
#include "ampere_conv_kernel.h"
|
||||
#include "gather_tensor.hpp"
|
||||
|
||||
#include "cutlass/util/command_line.h"
|
||||
|
||||
bool check_cuda_result(cudaError_t code, const char* file, int line) {
|
||||
if (code == cudaSuccess) {
|
||||
return true;
|
||||
}
|
||||
|
||||
std::cerr << "CUDA error at (" << file << "," << line << ")\n\t" << unsigned(code) << " -- " << cudaGetErrorString(code) << "\n";
|
||||
return false;
|
||||
}
|
||||
|
||||
#define CHECK_CUDA(code) (check_cuda_result(code, __FILE__, __LINE__))
|
||||
|
||||
using namespace cute;
|
||||
using example::IndexedGather;
|
||||
using example::CustomStride;
|
||||
|
||||
template<class Operator, class FilterTensor, class ActivationTensor, class OutputTensor>
|
||||
__global__
|
||||
__launch_bounds__(Operator::MaxThreadsPerBlock, Operator::MinBlocksPerMultiprocessor)
|
||||
void kernel_entrypoint(FilterTensor mFlt, ActivationTensor mAct, OutputTensor mOut) {
|
||||
extern __shared__ char smem_buf[];
|
||||
Operator op;
|
||||
op(mFlt, mAct, mOut, smem_buf);
|
||||
}
|
||||
|
||||
int ampere_dense_conv_fprop(
|
||||
int num_images,
|
||||
float* activations,
|
||||
float* filter,
|
||||
float* output,
|
||||
float* output_ref,
|
||||
int num_iterations = 1,
|
||||
bool do_ref_check = true) {
|
||||
auto D = typename AmpereUnpredicatedFprop::D{};
|
||||
auto H = typename AmpereUnpredicatedFprop::H{};
|
||||
auto W = typename AmpereUnpredicatedFprop::W{};
|
||||
auto Z = typename AmpereUnpredicatedFprop::Z{};
|
||||
auto P = typename AmpereUnpredicatedFprop::P{};
|
||||
auto Q = typename AmpereUnpredicatedFprop::Q{};
|
||||
auto C = typename AmpereUnpredicatedFprop::C{};
|
||||
auto K = typename AmpereUnpredicatedFprop::K{};
|
||||
auto S = typename AmpereUnpredicatedFprop::S{};
|
||||
auto R = typename AmpereUnpredicatedFprop::R{};
|
||||
auto T = typename AmpereUnpredicatedFprop::T{};
|
||||
|
||||
int N = num_images; // dynamic
|
||||
if (num_images % int(typename AmpereUnpredicatedFprop::Tiler_N{}) != 0) {
|
||||
printf("ERROR: Input image count must be evenly divisible by CTA tiler N.\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Tensor Activation: (n,d,h,w,c)::(?,6,4,4,64):(6144,1536,384,64,1)
|
||||
auto activation_layout = make_layout(
|
||||
make_shape (make_shape ( N, D, H, W), make_shape ( C, _1{},_1{},_1{})),
|
||||
make_stride(make_stride(D*H*W*C, H*W*C, W*C, C), make_stride(_1{}, _0{},_0{},_0{})));
|
||||
|
||||
auto xformed_act_layout = make_layout(
|
||||
make_shape (make_shape(N, Z, P, Q), make_shape ( C, T, R, S)),
|
||||
make_stride(stride<0>(activation_layout), make_stride(_1{}, H*W*C, W*C, C)));
|
||||
|
||||
// Tensor Filter : (k,c,s,r,t)::(128,3,3,3,64):(1728,576,192,64,1)
|
||||
auto filter_layout = AmpereUnpredicatedFprop::GmemLayoutFlt{};
|
||||
|
||||
// Tensor Output : (n,z,p,q,k)::(?,4,2,2,128):(2048,1024,512,128,1)
|
||||
auto output_layout = make_ordered_layout(
|
||||
make_shape( K, make_shape( N, Z, P, Q)),
|
||||
make_tuple(_0{}, make_tuple(_4{},_3{},_2{},_1{})));
|
||||
|
||||
Tensor mActivation = make_tensor(make_gmem_ptr(activations), activation_layout);
|
||||
Tensor mXformedAct = make_tensor(make_gmem_ptr(activations), xformed_act_layout);
|
||||
Tensor mFilter = make_tensor(make_gmem_ptr(filter), filter_layout);
|
||||
Tensor mOutput = make_tensor(make_gmem_ptr(output), output_layout); // (K, (N,Z,P,Q))
|
||||
Tensor mOutputRef = make_tensor(make_gmem_ptr(output_ref), output_layout);
|
||||
|
||||
print("xformed act layout ((N,Z,P,Q), (C,T,R,S)) = "); print(xformed_act_layout); print("\n");
|
||||
|
||||
cudaEvent_t start, stop;
|
||||
CHECK_CUDA(cudaEventCreate(&start));
|
||||
CHECK_CUDA(cudaEventCreate(&stop));
|
||||
|
||||
constexpr size_t smem_size = sizeof(typename AmpereUnpredicatedFprop::SharedStorage);
|
||||
Tensor gOutput_mn = zipped_divide(mOutput, typename AmpereUnpredicatedFprop::TilerOut{}); // ((BLK_M, BLK_N), (m', n'))
|
||||
dim3 lauch_grid {static_cast<uint32_t>(size<1,1>(gOutput_mn)), static_cast<uint32_t>(size<1,0>(gOutput_mn)), 1};
|
||||
|
||||
CHECK_CUDA(cudaFuncSetAttribute(
|
||||
kernel_entrypoint<AmpereUnpredicatedFprop, decltype(mFilter), decltype(mXformedAct), decltype(mOutput)>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
smem_size));
|
||||
|
||||
CHECK_CUDA(cudaEventRecord(start));
|
||||
for (int i = 0; i < num_iterations; ++i) {
|
||||
kernel_entrypoint<AmpereUnpredicatedFprop, decltype(mFilter), decltype(mXformedAct), decltype(mOutput)>
|
||||
<<<lauch_grid, AmpereUnpredicatedFprop::MaxThreadsPerBlock, smem_size>>>(
|
||||
mFilter, mXformedAct, mOutput);
|
||||
}
|
||||
CHECK_CUDA(cudaEventRecord(stop));
|
||||
CHECK_CUDA(cudaEventSynchronize(stop));
|
||||
|
||||
float milliseconds = 0;
|
||||
cudaEventElapsedTime(&milliseconds, start, stop);
|
||||
milliseconds /= float(num_iterations);
|
||||
|
||||
double tflop_count = (2 * double(size<0>(xformed_act_layout)) * double(size(filter_layout))) / double(1e12);
|
||||
double tflops = tflop_count / (double(milliseconds) / double(1e3));
|
||||
|
||||
printf("Conv TFLOP count = %f\n", tflop_count);
|
||||
printf("Conv dense perf: %fms | TFLOP/s = %f\n", milliseconds, tflops);
|
||||
|
||||
if (do_ref_check) {
|
||||
printf("Running host reference check ...\n");
|
||||
return fprop_reference(mFilter, mXformedAct, mOutput, mOutputRef);
|
||||
}
|
||||
else {
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
int ampere_gather_scatter_conv_fprop(
|
||||
int num_images,
|
||||
float* activations,
|
||||
uint32_t *gather_idx_buf,
|
||||
float* filter,
|
||||
float* output,
|
||||
uint32_t *scatter_idx_buf,
|
||||
int num_iterations = 1) {
|
||||
auto D = typename AmpereUnpredicatedFprop::D{};
|
||||
auto H = typename AmpereUnpredicatedFprop::H{};
|
||||
auto W = typename AmpereUnpredicatedFprop::W{};
|
||||
auto Z = typename AmpereUnpredicatedFprop::Z{};
|
||||
auto P = typename AmpereUnpredicatedFprop::P{};
|
||||
auto Q = typename AmpereUnpredicatedFprop::Q{};
|
||||
auto C = typename AmpereUnpredicatedFprop::C{};
|
||||
auto K = typename AmpereUnpredicatedFprop::K{};
|
||||
auto S = typename AmpereUnpredicatedFprop::S{};
|
||||
auto R = typename AmpereUnpredicatedFprop::R{};
|
||||
auto T = typename AmpereUnpredicatedFprop::T{};
|
||||
|
||||
int N = num_images; // dynamic
|
||||
if (N % int(typename AmpereUnpredicatedFprop::Tiler_N{}) != 0) {
|
||||
printf("ERROR: Input image count must be evenly divisible by CTA tiler N. Got num_images = %d\n", N);
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Tensor Filter : (k,c,s,r,t)::(128,3,3,3,64):(1728,576,192,64,1)
|
||||
auto filter_layout = AmpereUnpredicatedFprop::GmemLayoutFlt{};
|
||||
|
||||
// Tensor Output : (n,z,p,q,k)::(?,4,2,2,128):(2048,1024,512,128,1)
|
||||
auto output_layout = make_ordered_layout(
|
||||
make_shape( K, make_shape( N, Z, P, Q)),
|
||||
make_tuple(_0{}, make_tuple(_4{},_3{},_2{},_1{})));
|
||||
|
||||
// Input gather layout
|
||||
// inner_layout(make_coord((nzpq), (csrt))) => (idx_buffer_idx, dense_c_idx)
|
||||
auto EG = E<0>{}; // Gather basis (1,0) (idx_buffer_idx)
|
||||
auto EC = E<1>{}; // Contiguous basis (0,1) (dense_offset)
|
||||
auto xformed_act_logical_inner = make_layout(
|
||||
make_shape (make_shape ( N, Z, P, Q), make_shape ( C, T, R, S)),
|
||||
make_stride(make_stride(D*H*W*EG, H*W*EG, W*EG, EG), make_stride(EC, H*W*EG, W*EG, EG)));
|
||||
|
||||
// outer_layout(make_coord(idx_buffer_idx, dense_c_idx)) => idx
|
||||
// IndexedGather obtains idx by applying (gmem_base_ptr + gather_idx_buf[idx_buffer_idx] + dense_offset)
|
||||
auto xformed_act_gather_outer = make_layout(
|
||||
make_shape(_1{},_1{}),
|
||||
make_stride(CustomStride{IndexedGather{gather_idx_buf}, C}, _1{}));
|
||||
|
||||
// Compose the inner and outer layouts
|
||||
// gather_composed(make_coord((nzpq), (csrt))) => idx
|
||||
auto xformed_act_composed_layout = composition(
|
||||
xformed_act_gather_outer,
|
||||
make_arithmetic_tuple(_0{}, _0{}),
|
||||
xformed_act_logical_inner);
|
||||
|
||||
// Output scatter layout
|
||||
auto out_basis_stride = make_stride(
|
||||
E<1>{},
|
||||
make_stride(Z*P*Q*E<0>{}, P*Q*E<0>{}, Q*E<0>{}, _1{}*E<0>{})); // -> (crd0, crd1)
|
||||
auto out_basis_layout = make_layout(shape(output_layout), out_basis_stride);
|
||||
auto out_scatter_layout = make_layout(
|
||||
make_shape(_1{},_1{}),
|
||||
make_stride(CustomStride{IndexedGather{scatter_idx_buf}, K}, _1{}));
|
||||
auto out_composed_layout = composition(
|
||||
out_scatter_layout,
|
||||
make_arithmetic_tuple(_0{},_0{}),
|
||||
out_basis_layout);
|
||||
|
||||
Tensor mXformedActGather = make_tensor(make_gmem_ptr(activations), xformed_act_composed_layout);
|
||||
Tensor mFilter = make_tensor(make_gmem_ptr(filter), filter_layout);
|
||||
Tensor mOutputScatter = make_tensor(make_gmem_ptr(output), out_composed_layout); // (K, (N,Z,P,Q))
|
||||
|
||||
Tensor gOutput_mn = zipped_divide(mOutputScatter, typename AmpereUnpredicatedFprop::TilerOut{}); // ((BLK_M, BLK_N), (m', n'))
|
||||
dim3 lauch_grid {static_cast<uint32_t>(size<1,1>(gOutput_mn)), static_cast<uint32_t>(size<1,0>(gOutput_mn)), 1};
|
||||
constexpr size_t smem_size = sizeof(typename AmpereUnpredicatedFprop::SharedStorage);
|
||||
|
||||
print("xforemed gather layout ((N,Z,P,Q), (C,T,R,S)) = "); print(xformed_act_composed_layout); print("\n");
|
||||
print("Output scatter layout ( K, (N,Z,P,Q)) = "); print(out_composed_layout); print("\n");
|
||||
print("Filter layout ( K, (C,T,R,S)) = "); print(filter_layout); print("\n");
|
||||
|
||||
CHECK_CUDA(cudaFuncSetAttribute(
|
||||
kernel_entrypoint<AmpereUnpredicatedFprop, decltype(mFilter), decltype(mXformedActGather), decltype(mOutputScatter)>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
smem_size));
|
||||
|
||||
cudaEvent_t start, stop;
|
||||
CHECK_CUDA(cudaEventCreate(&start));
|
||||
CHECK_CUDA(cudaEventCreate(&stop));
|
||||
CHECK_CUDA(cudaEventRecord(start));
|
||||
for (int i = 0; i < num_iterations; ++i) {
|
||||
kernel_entrypoint<AmpereUnpredicatedFprop, decltype(mFilter), decltype(mXformedActGather), decltype(mOutputScatter)>
|
||||
<<<lauch_grid, AmpereUnpredicatedFprop::MaxThreadsPerBlock, smem_size>>>(
|
||||
mFilter, mXformedActGather, mOutputScatter);
|
||||
}
|
||||
CHECK_CUDA(cudaEventRecord(stop));
|
||||
CHECK_CUDA(cudaEventSynchronize(stop));
|
||||
float milliseconds = 0;
|
||||
cudaEventElapsedTime(&milliseconds, start, stop);
|
||||
milliseconds /= float(num_iterations);
|
||||
|
||||
double tflop_count = (2 * double(size<0>(xformed_act_logical_inner)) * double(size(filter_layout))) / double(1e12);
|
||||
double tflops = tflop_count / (double(milliseconds) / double(1e3));
|
||||
printf("Conv TFLOP count = %f\n", tflop_count);
|
||||
printf("Conv gather/scatter perf: %fms | TFLOP/s = %f\n", milliseconds, tflops);
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int
|
||||
main(int argc, char const** argv) {
|
||||
cutlass::CommandLine cmd(argc, argv);
|
||||
std::cout << "Ampere convolution forward propogation kernel supporting both affine and gather/scatter tensors.\n\n";
|
||||
if (cmd.check_cmd_line_flag("help")) {
|
||||
std::cout
|
||||
<< "Options:\n"
|
||||
"\t--n=<int> Sets the number of images for the input activation tensor (dataset size). Default = 131072.\n"
|
||||
"\t--i=<int> Sets the benchmarking repetitions. Default = 128.\n"
|
||||
"\t--nocheck If specified, skips the reference check for dense kernel.\n"
|
||||
"\t--help Displays this help message and exits.\n";
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
cudaDeviceProp props;
|
||||
cudaError_t error = cudaGetDeviceProperties(&props, 0);
|
||||
if (error != cudaSuccess) {
|
||||
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
|
||||
return -1;
|
||||
}
|
||||
if (props.major < 8) {
|
||||
std::cerr << "This example requires an Ampere GPU or newer.\n";
|
||||
return 0;
|
||||
}
|
||||
|
||||
int num_images = 4320;
|
||||
cmd.get_cmd_line_argument("n", num_images, 4320);
|
||||
int num_iterations = 128;
|
||||
cmd.get_cmd_line_argument("i", num_iterations, 128);
|
||||
bool do_host_ref_check = not cmd.check_cmd_line_flag("no-check");
|
||||
|
||||
auto D = typename AmpereUnpredicatedFprop::D{};
|
||||
auto H = typename AmpereUnpredicatedFprop::H{};
|
||||
auto W = typename AmpereUnpredicatedFprop::W{};
|
||||
auto Z = typename AmpereUnpredicatedFprop::Z{};
|
||||
auto P = typename AmpereUnpredicatedFprop::P{};
|
||||
auto Q = typename AmpereUnpredicatedFprop::Q{};
|
||||
auto C = typename AmpereUnpredicatedFprop::C{};
|
||||
auto K = typename AmpereUnpredicatedFprop::K{};
|
||||
|
||||
auto activation_layout = make_layout(
|
||||
make_shape (make_shape (num_images, D, H, W), make_shape ( C, _1{},_1{},_1{})),
|
||||
make_stride(make_stride( D*H*W*C, H*W*C, W*C, C), make_stride(_1{}, _0{},_0{},_0{})));
|
||||
|
||||
auto filter_layout = typename AmpereUnpredicatedFprop::GmemLayoutFlt{};
|
||||
|
||||
auto output_layout = make_ordered_layout(
|
||||
make_shape( K, make_shape(num_images, Z, P, Q)),
|
||||
make_step (_0{}, make_step ( _4{},_3{},_2{},_1{})));
|
||||
|
||||
print("Filter layout ( K, (C,T,R,S)) = "); print(filter_layout); print("\n");
|
||||
print("Activation layout ((N,D,H,W), (C,1,1,1)) = "); print(activation_layout); print("\n");
|
||||
print("Output layout ( K, (N,Z,P,Q)) = "); print(output_layout); print("\n");
|
||||
|
||||
// allocate tensors
|
||||
std::cout << "Allocating tensors ... ";
|
||||
thrust::universal_vector<float> activation_data(size_t(cute::size(activation_layout)), float(0));
|
||||
thrust::universal_vector<float> filter_data(size_t(cute::size(filter_layout)), float(0));
|
||||
thrust::universal_vector<float> output_data(size_t(cute::size(output_layout)), float(0));
|
||||
thrust::universal_vector<float> output_data_ref(size_t(cute::size(output_layout)), float(0));
|
||||
std::cout << "done.\n";
|
||||
|
||||
// init tensors
|
||||
std::cout << "Initializing data ... " << std::flush;
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> uniform_dist(-1.0, 1.0);
|
||||
for (std::size_t i = 0; i < size_t(cute::size(activation_layout)); ++i) {
|
||||
activation_data[i] = uniform_dist(gen);
|
||||
}
|
||||
|
||||
for (std::size_t i = 0; i < size_t(cute::size(filter_layout)); ++i) {
|
||||
filter_data[i] = uniform_dist(gen);
|
||||
}
|
||||
std::cout << "done.\n";
|
||||
|
||||
// set up index buffers for gather/scatter, fill with indireciton indices in reversed order
|
||||
std::cout << "Initializing gather/scatter index buffers ... ";
|
||||
thrust::universal_vector<uint32_t> gather_idx_buf(size_t(size<0>(activation_layout)));
|
||||
thrust::universal_vector<uint32_t> scatter_idx_buf(size_t(size<1>(output_layout)));
|
||||
thrust::sequence(gather_idx_buf.rbegin(), gather_idx_buf.rend());
|
||||
thrust::sequence(scatter_idx_buf.rbegin(), scatter_idx_buf.rend());
|
||||
std::cout << "done.\n";
|
||||
|
||||
// launch dense
|
||||
std::cout << "\nRunning dense fprop kernel\n";
|
||||
int passed = ampere_dense_conv_fprop(
|
||||
num_images,
|
||||
activation_data.data().get(),
|
||||
filter_data.data().get(),
|
||||
output_data.data().get(),
|
||||
output_data_ref.data().get(),
|
||||
num_iterations,
|
||||
do_host_ref_check);
|
||||
|
||||
// launch gather/scatter
|
||||
std::cout << "\nRunning gather/scatter fprop kernel\n";
|
||||
ampere_gather_scatter_conv_fprop(
|
||||
num_images,
|
||||
activation_data.data().get(),
|
||||
gather_idx_buf.data().get(),
|
||||
filter_data.data().get(),
|
||||
output_data.data().get(),
|
||||
scatter_idx_buf.data().get(),
|
||||
num_iterations);
|
||||
|
||||
return passed;
|
||||
}
|
||||
@@ -138,6 +138,8 @@ foreach(EXAMPLE
|
||||
55_hopper_mixed_dtype_gemm
|
||||
56_hopper_ptr_array_batched_gemm
|
||||
57_hopper_grouped_gemm
|
||||
58_ada_fp8_gemm
|
||||
59_ampere_gather_scatter_conv
|
||||
)
|
||||
|
||||
add_subdirectory(${EXAMPLE})
|
||||
|
||||
+7
-6
@@ -32,6 +32,7 @@
|
||||
|
||||
#include "cute/layout.hpp"
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cute/util/print.hpp"
|
||||
|
||||
namespace example {
|
||||
|
||||
@@ -59,7 +60,7 @@ struct IndexedGather
|
||||
CUTE_HOST_DEVICE friend
|
||||
void
|
||||
print(IndexedGather const &s) {
|
||||
print("Indexed");
|
||||
cute::print("Indexed");
|
||||
}
|
||||
|
||||
Index const *indices_;
|
||||
@@ -81,9 +82,9 @@ struct StridedGather
|
||||
CUTE_HOST_DEVICE friend
|
||||
void
|
||||
print(StridedGather const &s) {
|
||||
print("Strided{");
|
||||
cute::print("Strided{");
|
||||
print(s.stride_);
|
||||
print("}");
|
||||
cute::print("}");
|
||||
}
|
||||
|
||||
Stride stride_;
|
||||
@@ -109,11 +110,11 @@ struct CustomStride
|
||||
CUTE_HOST_DEVICE friend
|
||||
void
|
||||
print(CustomStride const & s) {
|
||||
print("Custom{");
|
||||
cute::print("Custom{");
|
||||
print(s.func_);
|
||||
print(",");
|
||||
cute::print(",");
|
||||
print(s.stride_);
|
||||
print("}");
|
||||
cute::print("}");
|
||||
}
|
||||
|
||||
template<class Div>
|
||||
@@ -29,8 +29,23 @@
|
||||
|
||||
|
||||
cutlass_example_add_executable(
|
||||
sgemm_nt_1
|
||||
sgemm_nt_1.cu
|
||||
sgemm_1
|
||||
sgemm_1.cu
|
||||
)
|
||||
|
||||
cutlass_example_add_executable(
|
||||
sgemm_2
|
||||
sgemm_2.cu
|
||||
)
|
||||
|
||||
cutlass_example_add_executable(
|
||||
sgemm_sm70
|
||||
sgemm_sm70.cu
|
||||
)
|
||||
|
||||
cutlass_example_add_executable(
|
||||
sgemm_sm80
|
||||
sgemm_sm80.cu
|
||||
)
|
||||
|
||||
cutlass_example_add_executable(
|
||||
|
||||
@@ -0,0 +1,469 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2023 - 2024 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#include <cstdlib>
|
||||
#include <cstdio>
|
||||
#include <cassert>
|
||||
|
||||
#include <thrust/host_vector.h>
|
||||
#include <thrust/device_vector.h>
|
||||
|
||||
#include <cute/tensor.hpp>
|
||||
|
||||
#include "cutlass/util/print_error.hpp"
|
||||
#include "cutlass/util/GPU_Clock.hpp"
|
||||
#include "cutlass/util/helper_cuda.hpp"
|
||||
|
||||
template <class ProblemShape, class CtaTiler,
|
||||
class TA, class AStride, class ASmemLayout, class AThreadLayout,
|
||||
class TB, class BStride, class BSmemLayout, class BThreadLayout,
|
||||
class TC, class CStride, class CSmemLayout, class CThreadLayout,
|
||||
class Alpha, class Beta>
|
||||
__global__ static
|
||||
__launch_bounds__(decltype(size(CThreadLayout{}))::value)
|
||||
void
|
||||
gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler,
|
||||
TA const* A, AStride dA, ASmemLayout sA_layout, AThreadLayout tA,
|
||||
TB const* B, BStride dB, BSmemLayout sB_layout, BThreadLayout tB,
|
||||
TC * C, CStride dC, CSmemLayout , CThreadLayout tC,
|
||||
Alpha alpha, Beta beta)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Preconditions
|
||||
CUTE_STATIC_ASSERT_V(rank(shape_MNK) == Int<3>{}); // (M, N, K)
|
||||
CUTE_STATIC_ASSERT_V(rank(cta_tiler) == Int<3>{}); // (BLK_M, BLK_N, BLK_K)
|
||||
|
||||
static_assert(is_static<AThreadLayout>::value);
|
||||
static_assert(is_static<BThreadLayout>::value);
|
||||
static_assert(is_static<CThreadLayout>::value);
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size(tA) == size(tB)); // NumThreads
|
||||
CUTE_STATIC_ASSERT_V(size(tC) == size(tA)); // NumThreads
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<0>(cta_tiler) % size<0>(tA) == Int<0>{}); // BLK_M / THR_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(cta_tiler) % size<1>(tA) == Int<0>{}); // BLK_K / THR_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(cta_tiler) % size<0>(tB) == Int<0>{}); // BLK_N / THR_N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(cta_tiler) % size<1>(tB) == Int<0>{}); // BLK_K / THR_K
|
||||
CUTE_STATIC_ASSERT_V(size<0>(cta_tiler) % size<0>(tC) == Int<0>{}); // BLK_M / THR_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(cta_tiler) % size<1>(tC) == Int<0>{}); // BLK_N / THR_N
|
||||
|
||||
static_assert(is_static<ASmemLayout>::value);
|
||||
static_assert(is_static<BSmemLayout>::value);
|
||||
static_assert(is_static<CSmemLayout>::value);
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<0>(ASmemLayout{}) == size<0>(cta_tiler)); // BLK_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<0>(cta_tiler)); // BLK_M
|
||||
CUTE_STATIC_ASSERT_V(size<0>(BSmemLayout{}) == size<1>(cta_tiler)); // BLK_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<1>(cta_tiler)); // BLK_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(ASmemLayout{}) == size<2>(cta_tiler)); // BLK_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(BSmemLayout{}) == size<2>(cta_tiler)); // BLK_K
|
||||
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<0,2>(shape_MNK), dA)); // dA strides for shape MK
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<1,2>(shape_MNK), dB)); // dB strides for shape NK
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<0,1>(shape_MNK), dC)); // dC strides for shape MN
|
||||
|
||||
//
|
||||
// Full and Tiled Tensors
|
||||
//
|
||||
|
||||
// Represent the full tensors
|
||||
Tensor mA = make_tensor(make_gmem_ptr(A), select<0,2>(shape_MNK), dA); // (M,K)
|
||||
Tensor mB = make_tensor(make_gmem_ptr(B), select<1,2>(shape_MNK), dB); // (N,K)
|
||||
Tensor mC = make_tensor(make_gmem_ptr(C), select<0,1>(shape_MNK), dC); // (M,N)
|
||||
|
||||
// Get the appropriate blocks for this thread block
|
||||
auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k)
|
||||
Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
|
||||
Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N)
|
||||
|
||||
// Shared memory buffers
|
||||
__shared__ TA smemA[cosize_v<ASmemLayout>];
|
||||
__shared__ TB smemB[cosize_v<BSmemLayout>];
|
||||
Tensor sA = make_tensor(make_smem_ptr(smemA), sA_layout); // (BLK_M,BLK_K)
|
||||
Tensor sB = make_tensor(make_smem_ptr(smemB), sB_layout); // (BLK_N,BLK_K)
|
||||
|
||||
//
|
||||
// Partition the copying of A and B tiles across the threads
|
||||
//
|
||||
|
||||
// TUTORIAL: Example of simple raked partitioning of ThreadLayouts tA|tB over data A|B tiles
|
||||
|
||||
Tensor tAgA = local_partition(gA, tA, threadIdx.x); // (THR_M,THR_K,k)
|
||||
Tensor tAsA = local_partition(sA, tA, threadIdx.x); // (THR_M,THR_K)
|
||||
|
||||
Tensor tBgB = local_partition(gB, tB, threadIdx.x); // (THR_N,THR_K,k)
|
||||
Tensor tBsB = local_partition(sB, tB, threadIdx.x); // (THR_N,THR_K)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<0>(tAgA) == size<0>(tAsA)); // THR_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tAsA)); // THR_K
|
||||
CUTE_STATIC_ASSERT_V(size<0>(tBgB) == size<0>(tBsB)); // THR_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBsB)); // THR_K
|
||||
|
||||
//
|
||||
// Define A/B partitioning and C accumulators
|
||||
//
|
||||
|
||||
// TUTORIAL: Example of partitioning via projections of a ThreadLayout tC
|
||||
|
||||
// Partition sA (M,K) by the rows of tC
|
||||
Tensor tCsA = local_partition(sA, tC, threadIdx.x, Step<_1, X>{}); // (THR_M,BLK_K)
|
||||
// Partition sB (N,K) by the cols of tC
|
||||
Tensor tCsB = local_partition(sB, tC, threadIdx.x, Step< X,_1>{}); // (THR_N,BLK_K)
|
||||
// Partition gC (M,N) by the tile of tC
|
||||
Tensor tCgC = local_partition(gC, tC, threadIdx.x, Step<_1,_1>{}); // (THR_M,THR_N)
|
||||
|
||||
// Allocate the accumulators -- same shape/layout as the partitioned data
|
||||
Tensor tCrC = make_tensor_like(tCgC); // (THR_M,THR_N)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<0>(tCrC) == size<0>(tCgC)); // THR_M
|
||||
CUTE_STATIC_ASSERT_V(size<0>(tCrC) == size<0>(tCsA)); // THR_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCrC) == size<1>(tCgC)); // THR_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCrC) == size<0>(tCsB)); // THR_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCsB)); // BLK_K
|
||||
|
||||
// Clear the accumulators
|
||||
clear(tCrC);
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mA : "); print( mA); print("\n");
|
||||
print(" gA : "); print( gA); print("\n");
|
||||
print(" sA : "); print( sA); print("\n");
|
||||
print("tAgA : "); print(tAgA); print("\n");
|
||||
print("tAsA : "); print(tAsA); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mB : "); print( mB); print("\n");
|
||||
print(" gB : "); print( gB); print("\n");
|
||||
print(" sB : "); print( sB); print("\n");
|
||||
print("tBgB : "); print(tBgB); print("\n");
|
||||
print("tBsB : "); print(tBsB); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mC : "); print( mC); print("\n");
|
||||
print(" gC : "); print( gC); print("\n");
|
||||
print("tCsA : "); print(tCsA); print("\n");
|
||||
print("tCsB : "); print(tCsB); print("\n");
|
||||
print("tCgC : "); print(tCgC); print("\n");
|
||||
print("tCrC : "); print(tCrC); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 1
|
||||
|
||||
// TUTORIAL: Example of a simple mainloop that read tiles of data into shared memory,
|
||||
// and then computes on those tiles.
|
||||
// copy(.) operates on the global and shared memory via the tA|tB partitioning
|
||||
// gemm(.) operates on the shared and register memory via the tC partitioning
|
||||
|
||||
auto K_TILE_MAX = size<2>(tAgA);
|
||||
|
||||
for (int k_tile = 0; k_tile < K_TILE_MAX; ++k_tile)
|
||||
{
|
||||
// Copy gmem to smem with tA|tB thread-partitioned tensors
|
||||
copy(tAgA(_,_,k_tile), tAsA); // A (THR_M,THR_K) -> (THR_M,THR_K)
|
||||
copy(tBgB(_,_,k_tile), tBsB); // B (THR_N,THR_K) -> (THR_N,THR_K)
|
||||
|
||||
// TUTORIAL: The above call to copy(tAgA(_,_,k_tile), tAsA) is equivalent to
|
||||
// Tensor tAgAk = tAgA(_,_,k_tile);
|
||||
// CUTE_UNROLL
|
||||
// for (int i = 0; i < size(tAsA); ++i) {
|
||||
// tAsA(i) = tAgAk(i);
|
||||
// }
|
||||
|
||||
cp_async_fence(); // Label the end of (potential) cp.async instructions
|
||||
cp_async_wait<0>(); // Sync on all (potential) cp.async instructions
|
||||
__syncthreads(); // Wait for all threads to write to smem
|
||||
|
||||
// Compute gemm on tC thread-partitioned smem
|
||||
gemm(tCsA, tCsB, tCrC); // (THR_M,THR_N) += (THR_M,BLK_K) * (THR_N,BLK_K)
|
||||
|
||||
// TUTORIAL: The above call to gemm(tCsA, tCsB, tCrC) is equivalent to
|
||||
// CUTE_UNROLL
|
||||
// for (int k = 0; k < size<1>(tCsA); ++k) {
|
||||
// CUTE_UNROLL
|
||||
// for (int m = 0; m < size<0>(tCrC); ++m) {
|
||||
// CUTE_UNROLL
|
||||
// for (int n = 0; n < size<1>(tCrC); ++n) {
|
||||
// tCrC(m,n) += tCsA(m,k) * tCsB(n,k);
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
|
||||
__syncthreads(); // Wait for all threads to read from smem
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
//
|
||||
// Epilogue
|
||||
//
|
||||
|
||||
axpby(alpha, tCrC, beta, tCgC);
|
||||
|
||||
// TUTORIAL: The above call to axpby(alpha, tCrC, beta, tCgC) is equivalent to
|
||||
// CUTE_UNROLL
|
||||
// for (int i = 0; i < size(tCsA); ++i) {
|
||||
// tCgC(i) = alpha * tCrC(i) + beta * tCgC(i);
|
||||
// }
|
||||
}
|
||||
|
||||
// Setup params for an NT GEMM
|
||||
// Use m-major smem sA, n-major smem sB, and mn-major threads tA|tB
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm_nt(int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Define shapes (dynamic)
|
||||
auto M = int(m);
|
||||
auto N = int(n);
|
||||
auto K = int(k);
|
||||
auto prob_shape = make_shape(M, N, K); // (M, N, K)
|
||||
|
||||
// Define NT strides (mixed)
|
||||
auto dA = make_stride(Int<1>{}, ldA); // (dM, dK)
|
||||
auto dB = make_stride(Int<1>{}, ldB); // (dN, dK)
|
||||
auto dC = make_stride(Int<1>{}, ldC); // (dM, dN)
|
||||
|
||||
// Define CTA tile sizes (static)
|
||||
auto bM = Int<128>{};
|
||||
auto bN = Int<128>{};
|
||||
auto bK = Int< 8>{};
|
||||
auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K)
|
||||
|
||||
// Define the smem layouts (static)
|
||||
auto sA = make_layout(make_shape(bM, bK)); // (m,k) -> smem_idx; m-major
|
||||
auto sB = make_layout(make_shape(bN, bK)); // (n,k) -> smem_idx; n-major
|
||||
auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx; m-major
|
||||
|
||||
// Define the thread layouts (static)
|
||||
auto tA = make_layout(make_shape(Int<32>{}, Int< 8>{})); // (m,k) -> thr_idx
|
||||
auto tB = make_layout(make_shape(Int<32>{}, Int< 8>{})); // (n,k) -> thr_idx
|
||||
auto tC = make_layout(make_shape(Int<16>{}, Int<16>{})); // (m,n) -> thr_idx
|
||||
|
||||
dim3 dimBlock(size(tC));
|
||||
dim3 dimGrid(size(ceil_div(M, bM)),
|
||||
size(ceil_div(N, bN)));
|
||||
gemm_device<<<dimGrid, dimBlock, 0, stream>>>
|
||||
(prob_shape, cta_tiler,
|
||||
A, dA, sA, tA,
|
||||
B, dB, sB, tB,
|
||||
C, dC, sC, tC,
|
||||
alpha, beta);
|
||||
}
|
||||
|
||||
// Setup params for a TN GEMM
|
||||
// Use padded m-major smem sA, padded n-major smem sB, and k-major threads tA|tB
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm_tn(int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Define shapes (dynamic)
|
||||
auto M = int(m);
|
||||
auto N = int(n);
|
||||
auto K = int(k);
|
||||
auto prob_shape = make_shape(M, N, K); // (M, N, K)
|
||||
|
||||
// Define TN strides (mixed)
|
||||
auto dA = make_stride(ldA, Int<1>{}); // (dM, dK)
|
||||
auto dB = make_stride(ldB, Int<1>{}); // (dN, dK)
|
||||
auto dC = make_stride(Int<1>{}, ldC); // (dM, dN)
|
||||
|
||||
// Define CTA tile sizes (static)
|
||||
auto bM = Int<128>{};
|
||||
auto bN = Int<128>{};
|
||||
auto bK = Int< 8>{};
|
||||
auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K)
|
||||
|
||||
// Define the smem layouts (static)
|
||||
auto sA = make_layout(make_shape(bM,bK), LayoutRight{}); // (m,k) -> smem_idx; k-major
|
||||
auto sB = make_layout(make_shape(bN,bK), LayoutRight{}); // (n,k) -> smem_idx; k-major
|
||||
auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx; m-major
|
||||
|
||||
// Define the thread layouts (static)
|
||||
auto tA = make_layout(make_shape(Int<32>{}, Int< 8>{}), LayoutRight{}); // (m,k) -> thr_idx; k-major
|
||||
auto tB = make_layout(make_shape(Int<32>{}, Int< 8>{}), LayoutRight{}); // (n,k) -> thr_idx; k-major
|
||||
auto tC = make_layout(make_shape(Int<16>{}, Int<16>{})); // (m,n) -> thr_idx; m-major
|
||||
|
||||
dim3 dimBlock(size(tC));
|
||||
dim3 dimGrid(size(ceil_div(M, bM)),
|
||||
size(ceil_div(N, bN)));
|
||||
gemm_device<<<dimGrid, dimBlock, 0, stream>>>
|
||||
(prob_shape, cta_tiler,
|
||||
A, dA, sA, tA,
|
||||
B, dB, sB, tB,
|
||||
C, dC, sC, tC,
|
||||
alpha, beta);
|
||||
}
|
||||
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm(char transA, char transB, int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
if (transA == 'N' && transB == 'T') {
|
||||
return gemm_nt(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream);
|
||||
} else
|
||||
if (transA == 'T' && transB == 'N') {
|
||||
return gemm_tn(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream);
|
||||
}
|
||||
assert(false && "Not implemented");
|
||||
}
|
||||
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
int m = 5120;
|
||||
if (argc >= 2)
|
||||
sscanf(argv[1], "%d", &m);
|
||||
|
||||
int n = 5120;
|
||||
if (argc >= 3)
|
||||
sscanf(argv[2], "%d", &n);
|
||||
|
||||
int k = 4096;
|
||||
if (argc >= 4)
|
||||
sscanf(argv[3], "%d", &k);
|
||||
|
||||
char transA = 'N';
|
||||
if (argc >= 5)
|
||||
sscanf(argv[4], "%c", &transA);
|
||||
|
||||
char transB = 'T';
|
||||
if (argc >= 6)
|
||||
sscanf(argv[5], "%c", &transB);
|
||||
|
||||
using TA = float;
|
||||
using TB = float;
|
||||
using TC = float;
|
||||
using TI = float;
|
||||
|
||||
TI alpha = 1.0;
|
||||
TI beta = 0.0;
|
||||
|
||||
std::cout << "M = " << m << std::endl;
|
||||
std::cout << "N = " << n << std::endl;
|
||||
std::cout << "K = " << k << std::endl;
|
||||
std::cout << "C = A^" << transA << " B^" << transB << std::endl;
|
||||
|
||||
cute::device_init(0);
|
||||
|
||||
thrust::host_vector<TA> h_A(m*k);
|
||||
thrust::host_vector<TB> h_B(n*k);
|
||||
thrust::host_vector<TC> h_C(m*n);
|
||||
|
||||
for (int j = 0; j < m*k; ++j) h_A[j] = static_cast<TA>( 2*(rand() / double(RAND_MAX)) - 1 );
|
||||
for (int j = 0; j < n*k; ++j) h_B[j] = static_cast<TB>( 2*(rand() / double(RAND_MAX)) - 1 );
|
||||
for (int j = 0; j < m*n; ++j) h_C[j] = static_cast<TC>(-1);
|
||||
|
||||
thrust::device_vector<TA> d_A = h_A;
|
||||
thrust::device_vector<TB> d_B = h_B;
|
||||
thrust::device_vector<TC> d_C = h_C;
|
||||
|
||||
double gflops = (2.0*m*n*k) * 1e-9;
|
||||
|
||||
const int timing_iterations = 100;
|
||||
GPU_Clock timer;
|
||||
|
||||
int ldA = 0, ldB = 0, ldC = m;
|
||||
|
||||
if (transA == 'N') {
|
||||
ldA = m;
|
||||
} else if (transA == 'T') {
|
||||
ldA = k;
|
||||
} else {
|
||||
assert(false);
|
||||
}
|
||||
|
||||
if (transB == 'N') {
|
||||
ldB = k;
|
||||
} else if (transB == 'T') {
|
||||
ldB = n;
|
||||
} else {
|
||||
assert(false);
|
||||
}
|
||||
// Run once
|
||||
d_C = h_C;
|
||||
gemm(transA, transB, m, n, k,
|
||||
alpha,
|
||||
d_A.data().get(), ldA,
|
||||
d_B.data().get(), ldB,
|
||||
beta,
|
||||
d_C.data().get(), ldC);
|
||||
CUTE_CHECK_LAST();
|
||||
thrust::host_vector<TC> cute_result = d_C;
|
||||
|
||||
// Timing iterations
|
||||
timer.start();
|
||||
for (int i = 0; i < timing_iterations; ++i) {
|
||||
gemm(transA, transB, m, n, k,
|
||||
alpha,
|
||||
d_A.data().get(), ldA,
|
||||
d_B.data().get(), ldB,
|
||||
beta,
|
||||
d_C.data().get(), ldC);
|
||||
}
|
||||
double cute_time = timer.seconds() / timing_iterations;
|
||||
CUTE_CHECK_LAST();
|
||||
printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000);
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,523 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2023 - 2024 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#include <cstdlib>
|
||||
#include <cstdio>
|
||||
#include <cassert>
|
||||
|
||||
#include <thrust/host_vector.h>
|
||||
#include <thrust/device_vector.h>
|
||||
|
||||
#include <cute/tensor.hpp>
|
||||
|
||||
#include "cutlass/util/print_error.hpp"
|
||||
#include "cutlass/util/GPU_Clock.hpp"
|
||||
#include "cutlass/util/helper_cuda.hpp"
|
||||
|
||||
template <class ProblemShape, class CtaTiler,
|
||||
class TA, class AStride, class ASmemLayout, class TiledCopyA,
|
||||
class TB, class BStride, class BSmemLayout, class TiledCopyB,
|
||||
class TC, class CStride, class CSmemLayout, class TiledMma,
|
||||
class Alpha, class Beta>
|
||||
__global__ static
|
||||
__launch_bounds__(decltype(size(TiledMma{}))::value)
|
||||
void
|
||||
gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler,
|
||||
TA const* A, AStride dA, ASmemLayout sA_layout, TiledCopyA copy_a,
|
||||
TB const* B, BStride dB, BSmemLayout sB_layout, TiledCopyB copy_b,
|
||||
TC * C, CStride dC, CSmemLayout , TiledMma mma,
|
||||
Alpha alpha, Beta beta)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Preconditions
|
||||
CUTE_STATIC_ASSERT_V(rank(shape_MNK) == Int<3>{}); // (M, N, K)
|
||||
CUTE_STATIC_ASSERT_V(rank(cta_tiler) == Int<3>{}); // (BLK_M, BLK_N, BLK_K)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size(copy_a) == size(mma)); // NumThreads
|
||||
CUTE_STATIC_ASSERT_V(size(copy_b) == size(mma)); // NumThreads
|
||||
|
||||
static_assert(is_static<ASmemLayout>::value);
|
||||
static_assert(is_static<BSmemLayout>::value);
|
||||
static_assert(is_static<CSmemLayout>::value);
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<0>(ASmemLayout{}) == size<0>(cta_tiler)); // BLK_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<0>(cta_tiler)); // BLK_M
|
||||
CUTE_STATIC_ASSERT_V(size<0>(BSmemLayout{}) == size<1>(cta_tiler)); // BLK_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<1>(cta_tiler)); // BLK_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(ASmemLayout{}) == size<2>(cta_tiler)); // BLK_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(BSmemLayout{}) == size<2>(cta_tiler)); // BLK_K
|
||||
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<0,2>(shape_MNK), dA)); // dA strides for shape MK
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<1,2>(shape_MNK), dB)); // dB strides for shape NK
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<0,1>(shape_MNK), dC)); // dC strides for shape MN
|
||||
|
||||
//
|
||||
// Full and Tiled Tensors
|
||||
//
|
||||
|
||||
// Represent the full tensors
|
||||
Tensor mA = make_tensor(make_gmem_ptr(A), select<0,2>(shape_MNK), dA); // (M,K)
|
||||
Tensor mB = make_tensor(make_gmem_ptr(B), select<1,2>(shape_MNK), dB); // (N,K)
|
||||
Tensor mC = make_tensor(make_gmem_ptr(C), select<0,1>(shape_MNK), dC); // (M,N)
|
||||
|
||||
// Get the appropriate blocks for this thread block
|
||||
auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k)
|
||||
Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
|
||||
Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N)
|
||||
|
||||
// Shared memory buffers
|
||||
__shared__ TA smemA[cosize_v<ASmemLayout>];
|
||||
__shared__ TB smemB[cosize_v<BSmemLayout>];
|
||||
Tensor sA = make_tensor(make_smem_ptr(smemA), sA_layout); // (BLK_M,BLK_K)
|
||||
Tensor sB = make_tensor(make_smem_ptr(smemB), sB_layout); // (BLK_N,BLK_K)
|
||||
|
||||
//
|
||||
// Partition the copying of A and B tiles across the threads
|
||||
//
|
||||
|
||||
// TUTORIAL: Example of partitioning via a TiledCopy
|
||||
|
||||
ThrCopy thr_copy_a = copy_a.get_slice(threadIdx.x);
|
||||
Tensor tAgA = thr_copy_a.partition_S(gA); // (CPY,CPY_M,CPY_K,k)
|
||||
Tensor tAsA = thr_copy_a.partition_D(sA); // (CPY,CPY_M,CPY_K)
|
||||
// Allocate registers same shape/layout as partitioned data
|
||||
Tensor tArA = make_fragment_like(tAsA); // (CPY,CPY_M,CPY_K)
|
||||
|
||||
ThrCopy thr_copy_b = copy_b.get_slice(threadIdx.x);
|
||||
Tensor tBgB = thr_copy_b.partition_S(gB); // (CPY,CPY_N,CPY_K,k)
|
||||
Tensor tBsB = thr_copy_b.partition_D(sB); // (CPY,CPY_N,CPY_K)
|
||||
// Allocate registers same shape/layout as partitioned data
|
||||
Tensor tBrB = make_fragment_like(tBsB); // (CPY,CPY_N,CPY_K)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tAsA)); // CPY_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tArA)); // CPY_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tAgA) == size<2>(tAsA)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tAgA) == size<2>(tArA)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBsB)); // CPY_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBrB)); // CPY_N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tBgB) == size<2>(tBsB)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tBgB) == size<2>(tBrB)); // CPY_K
|
||||
|
||||
// Copy gmem to rmem for k_tile=0
|
||||
copy(copy_a, tAgA(_,_,_,0), tArA);
|
||||
copy(copy_b, tBgB(_,_,_,0), tBrB);
|
||||
//
|
||||
// Define A/B partitioning and C accumulators
|
||||
//
|
||||
|
||||
// TUTORIAL: Example of partitioning via a TiledMMA
|
||||
|
||||
ThrMMA thr_mma = mma.get_slice(threadIdx.x);
|
||||
Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K)
|
||||
Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K)
|
||||
Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N)
|
||||
|
||||
// Allocate the accumulators -- same size as the projected data
|
||||
Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N)
|
||||
|
||||
CUTE_STATIC_ASSERT_V( shape(tCrC) == shape(tCgC)); // (MMA,MMA_M,MMA_N)
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCgC) == size<1>(tCsA)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCgC) == size<1>(tCsB)); // MMA_N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // MMA_K
|
||||
|
||||
// Clear the accumulators
|
||||
clear(tCrC);
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mA : "); print( mA); print("\n");
|
||||
print(" gA : "); print( gA); print("\n");
|
||||
print(" sA : "); print( sA); print("\n");
|
||||
print("tAgA : "); print(tAgA); print("\n");
|
||||
print("tAsA : "); print(tAsA); print("\n");
|
||||
print("tArA : "); print(tArA); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mB : "); print( mB); print("\n");
|
||||
print(" gB : "); print( gB); print("\n");
|
||||
print(" sB : "); print( sB); print("\n");
|
||||
print("tBgB : "); print(tBgB); print("\n");
|
||||
print("tBsB : "); print(tBsB); print("\n");
|
||||
print("tArA : "); print(tArA); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mC : "); print( mC); print("\n");
|
||||
print(" gC : "); print( gC); print("\n");
|
||||
print("tCsA : "); print(tCsA); print("\n");
|
||||
print("tCsB : "); print(tCsB); print("\n");
|
||||
print("tCgC : "); print(tCgC); print("\n");
|
||||
print("tCrC : "); print(tCrC); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 1
|
||||
|
||||
// TUTORIAL: Example of an inner loop that pipelines compute with reads
|
||||
// from global memory by staging through register and shared memory.
|
||||
// Data is read from global to registers, then to shared via the TiledCopy partitions
|
||||
// gemm(.) operates on the shared memory directly via the TiledMMA partitions
|
||||
|
||||
auto K_TILE_MAX = size<3>(tAgA);
|
||||
|
||||
for (int k_tile = 0; k_tile < K_TILE_MAX; ++k_tile)
|
||||
{
|
||||
// Copy rmem to smem with tA|tB thread-partitioned tensors
|
||||
__syncthreads(); // Wait for all threads to consume smem
|
||||
copy(tArA, tAsA);
|
||||
copy(tBrB, tBsB);
|
||||
__syncthreads(); // Wait for all threads to consume smem
|
||||
|
||||
// Copy gmem to rmem for k_tile+1 with tA|tB thread-partitioned tensors
|
||||
int k_tile_next = (k_tile + 1 < K_TILE_MAX) ? k_tile + 1 : k_tile;
|
||||
copy(copy_a, tAgA(_,_,_,k_tile_next), tArA);
|
||||
copy(copy_b, tBgB(_,_,_,k_tile_next), tBrB);
|
||||
// TUTORIAL: The above call to copy(copy_a, tAgA(_,_,_,k_tile_next), tArA) is equivalent to
|
||||
// CUTE_UNROLL
|
||||
// for (int k = 0; k < size<1>(tCsA); ++k) {
|
||||
// CUTE_UNROLL
|
||||
// for (int m = 0; m < size<0>(tCrC); ++m) {
|
||||
// copy_a.call(tAgA(_,m,k), tArA(_,m,k);
|
||||
// }
|
||||
// }
|
||||
|
||||
// Compute gemm on mma-partitioned smem
|
||||
gemm(mma, tCsA, tCsB, tCrC);
|
||||
// TUTORIAL: The above call to gemm(tCsA, tCsB, tCrC) is equivalent to
|
||||
// CUTE_UNROLL
|
||||
// for (int k = 0; k < size<1>(tCsA); ++k) {
|
||||
// CUTE_UNROLL
|
||||
// for (int m = 0; m < size<0>(tCrC); ++m) {
|
||||
// CUTE_UNROLL
|
||||
// for (int n = 0; n < size<1>(tCrC); ++n) {
|
||||
// mma.call(tCsA(_,m,k), tCsB(_,n,k), tCrC(_,m,n);
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
//
|
||||
// Epilogue
|
||||
//
|
||||
|
||||
axpby(alpha, tCrC, beta, tCgC);
|
||||
}
|
||||
|
||||
// Setup params for a NT GEMM
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm_nt(int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Define shapes (dynamic)
|
||||
auto M = int(m);
|
||||
auto N = int(n);
|
||||
auto K = int(k);
|
||||
auto prob_shape = make_shape(M, N, K); // (M, N, K)
|
||||
|
||||
// Define NT strides (mixed)
|
||||
auto dA = make_stride(Int<1>{}, ldA); // (dM, dK)
|
||||
auto dB = make_stride(Int<1>{}, ldB); // (dN, dK)
|
||||
auto dC = make_stride(Int<1>{}, ldC); // (dM, dN)
|
||||
|
||||
// Define CTA tile sizes (static)
|
||||
auto bM = Int<128>{};
|
||||
auto bN = Int<128>{};
|
||||
auto bK = Int< 8>{};
|
||||
auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K)
|
||||
|
||||
// Define the smem layouts (static)
|
||||
auto sA = make_layout(make_shape(bM, bK)); // (m,k) -> smem_idx; m-major
|
||||
auto sB = make_layout(make_shape(bN, bK)); // (n,k) -> smem_idx; n-major
|
||||
auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx; m-major
|
||||
|
||||
// Define the thread layouts (static)
|
||||
|
||||
// TUTORIAL: Construct TiledCopy with a particular Copy_Atom to use and
|
||||
// define the partitioning pattern to apply.
|
||||
// Each thread will (try to) copy 4x1 elements of type TA using 128-bit copy.
|
||||
// Use 32x8 of these threads.
|
||||
|
||||
TiledCopy copyA = make_tiled_copy(Copy_Atom<UniversalCopy<uint128_t>, TA>{},
|
||||
Layout<Shape<_32,_8>>{}, // Thr layout 32x8 m-major
|
||||
Layout<Shape< _4,_1>>{}); // Val layout 4x1 m-major
|
||||
TiledCopy copyB = make_tiled_copy(Copy_Atom<UniversalCopy<uint128_t>, TB>{},
|
||||
Layout<Shape<_32,_8>>{}, // Thr layout 32x8 n-major
|
||||
Layout<Shape< _4,_1>>{}); // Val layout 4x1 n-major
|
||||
|
||||
// TUTORIAL: Construct TiledMMA with a particular MMA_Atom to use and
|
||||
// define the partitioning pattern to apply.
|
||||
// Use a 1x1x1 FMA on the types TC += TA * TB. Each atom requires a single thread.
|
||||
// Reproduce that atom 16x16x1 times (m-major) across threads so that we use 256 threads.
|
||||
|
||||
TiledMMA mmaC = make_tiled_mma(UniversalFMA<TC,TA,TB>{},
|
||||
Layout<Shape<_16,_16,_1>>{}); // 16x16x1 UniversalFMA
|
||||
|
||||
#if 0
|
||||
print(copyA);
|
||||
print(copyB);
|
||||
print(mmaC);
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
print_latex(copyA);
|
||||
print_latex(copyB);
|
||||
print_latex(mmaC);
|
||||
#endif
|
||||
|
||||
dim3 dimBlock(size(mmaC));
|
||||
dim3 dimGrid(size(ceil_div(M, bM)),
|
||||
size(ceil_div(N, bN)));
|
||||
gemm_device<<<dimGrid, dimBlock, 0, stream>>>
|
||||
(prob_shape, cta_tiler,
|
||||
A, dA, sA, copyA,
|
||||
B, dB, sB, copyB,
|
||||
C, dC, sC, mmaC,
|
||||
alpha, beta);
|
||||
}
|
||||
|
||||
// Setup params for a TN GEMM
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm_tn(int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Define shapes (dynamic)
|
||||
auto M = int(m);
|
||||
auto N = int(n);
|
||||
auto K = int(k);
|
||||
auto prob_shape = make_shape(M, N, K); // (M, N, K)
|
||||
|
||||
// Define TN strides (mixed)
|
||||
auto dA = make_stride(ldA, Int<1>{}); // (dM, dK)
|
||||
auto dB = make_stride(ldB, Int<1>{}); // (dN, dK)
|
||||
auto dC = make_stride(Int<1>{}, ldC); // (dM, dN)
|
||||
|
||||
// Define CTA tile sizes (static)
|
||||
auto bM = Int<128>{};
|
||||
auto bN = Int<128>{};
|
||||
auto bK = Int< 8>{};
|
||||
auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K)
|
||||
|
||||
// Define the smem layouts (static)
|
||||
auto sA = make_layout(make_shape ( bM, bK),
|
||||
make_stride(Int<1>{}, bM+Int<1>{})); // (m,k) -> smem_idx; padded m-major
|
||||
auto sB = make_layout(make_shape ( bN, bK),
|
||||
make_stride(Int<1>{}, bN+Int<1>{})); // (n,k) -> smem_idx; padded n-major
|
||||
auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx
|
||||
|
||||
// TUTORIAL: Construct TiledCopy to define the Copy_Atom to use and the
|
||||
// partitioning pattern to apply.
|
||||
// Each thread will copy 1x1 elements of type TA.
|
||||
// Use 32x8 of these threads arranged in k-major.
|
||||
|
||||
TiledCopy copyA = make_tiled_copy(Copy_Atom<UniversalCopy<TA>, TA>{},
|
||||
Layout<Shape<_32,_8>,Stride<_8,_1>>{}, // Thr layout 32x8 k-major
|
||||
Layout<Shape< _1,_1>>{}); // Val layout 1x1
|
||||
TiledCopy copyB = make_tiled_copy(Copy_Atom<UniversalCopy<TB>, TB>{},
|
||||
Layout<Shape<_32,_8>,Stride<_8,_1>>{}, // Thr layout 32x8 k-major
|
||||
Layout<Shape< _1,_1>>{}); // Val layout 1x1
|
||||
|
||||
// TUTORIAL: Construct TiledMMA to define the MMA_Atom to use and the
|
||||
// partitioning pattern to apply.
|
||||
// Use a 1x1x1 FMA on the types TC += TA * TB. Each atom requires a single thread.
|
||||
// Reproduce that atom 16x16x1 times (m-major) across threads so that we use 256 threads.
|
||||
|
||||
TiledMMA mmaC = make_tiled_mma(UniversalFMA<TC,TA,TB>{},
|
||||
Layout<Shape<_16,_16,_1>>{}); // 16x16x1 TiledMMA
|
||||
|
||||
#if 0
|
||||
print(copyA);
|
||||
print(copyB);
|
||||
print(mmaC);
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
print_latex(copyA);
|
||||
print_latex(copyB);
|
||||
print_latex(mmaC);
|
||||
#endif
|
||||
|
||||
dim3 dimBlock(size(mmaC));
|
||||
dim3 dimGrid(size(ceil_div(M, bM)),
|
||||
size(ceil_div(N, bN)));
|
||||
gemm_device<<<dimGrid, dimBlock, 0, stream>>>
|
||||
(prob_shape, cta_tiler,
|
||||
A, dA, sA, copyA,
|
||||
B, dB, sB, copyB,
|
||||
C, dC, sC, mmaC,
|
||||
alpha, beta);
|
||||
}
|
||||
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm(char transA, char transB, int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
if (transA == 'N' && transB == 'T') {
|
||||
return gemm_nt(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream);
|
||||
} else
|
||||
if (transA == 'T' && transB == 'N') {
|
||||
return gemm_tn(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream);
|
||||
}
|
||||
assert(false && "Not implemented");
|
||||
}
|
||||
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
int m = 5120;
|
||||
if (argc >= 2)
|
||||
sscanf(argv[1], "%d", &m);
|
||||
|
||||
int n = 5120;
|
||||
if (argc >= 3)
|
||||
sscanf(argv[2], "%d", &n);
|
||||
|
||||
int k = 4096;
|
||||
if (argc >= 4)
|
||||
sscanf(argv[3], "%d", &k);
|
||||
|
||||
char transA = 'N';
|
||||
if (argc >= 5)
|
||||
sscanf(argv[4], "%c", &transA);
|
||||
|
||||
char transB = 'T';
|
||||
if (argc >= 6)
|
||||
sscanf(argv[5], "%c", &transB);
|
||||
|
||||
using TA = float;
|
||||
using TB = float;
|
||||
using TC = float;
|
||||
using TI = float;
|
||||
|
||||
TI alpha = 1.0;
|
||||
TI beta = 0.0;
|
||||
|
||||
std::cout << "M = " << m << std::endl;
|
||||
std::cout << "N = " << n << std::endl;
|
||||
std::cout << "K = " << k << std::endl;
|
||||
std::cout << "C = A^" << transA << " B^" << transB << std::endl;
|
||||
|
||||
cute::device_init(0);
|
||||
|
||||
thrust::host_vector<TA> h_A(m*k);
|
||||
thrust::host_vector<TB> h_B(n*k);
|
||||
thrust::host_vector<TC> h_C(m*n);
|
||||
|
||||
for (int j = 0; j < m*k; ++j) h_A[j] = static_cast<TA>( 2*(rand() / double(RAND_MAX)) - 1 );
|
||||
for (int j = 0; j < n*k; ++j) h_B[j] = static_cast<TB>( 2*(rand() / double(RAND_MAX)) - 1 );
|
||||
for (int j = 0; j < m*n; ++j) h_C[j] = static_cast<TC>(-1);
|
||||
|
||||
thrust::device_vector<TA> d_A = h_A;
|
||||
thrust::device_vector<TB> d_B = h_B;
|
||||
thrust::device_vector<TC> d_C = h_C;
|
||||
|
||||
double gflops = (2.0*m*n*k) * 1e-9;
|
||||
|
||||
const int timing_iterations = 100;
|
||||
GPU_Clock timer;
|
||||
|
||||
int ldA = 0, ldB = 0, ldC = m;
|
||||
|
||||
if (transA == 'N') {
|
||||
ldA = m;
|
||||
} else if (transA == 'T') {
|
||||
ldA = k;
|
||||
} else {
|
||||
assert(false);
|
||||
}
|
||||
|
||||
if (transB == 'N') {
|
||||
ldB = k;
|
||||
} else if (transB == 'T') {
|
||||
ldB = n;
|
||||
} else {
|
||||
assert(false);
|
||||
}
|
||||
|
||||
// Run once
|
||||
d_C = h_C;
|
||||
gemm(transA, transB, m, n, k,
|
||||
alpha,
|
||||
d_A.data().get(), ldA,
|
||||
d_B.data().get(), ldB,
|
||||
beta,
|
||||
d_C.data().get(), ldC);
|
||||
CUTE_CHECK_LAST();
|
||||
thrust::host_vector<TC> cute_result = d_C;
|
||||
|
||||
// Timing iterations
|
||||
timer.start();
|
||||
for (int i = 0; i < timing_iterations; ++i) {
|
||||
gemm(transA, transB, m, n, k,
|
||||
alpha,
|
||||
d_A.data().get(), ldA,
|
||||
d_B.data().get(), ldB,
|
||||
beta,
|
||||
d_C.data().get(), ldC);
|
||||
}
|
||||
double cute_time = timer.seconds() / timing_iterations;
|
||||
CUTE_CHECK_LAST();
|
||||
printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000);
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -1,426 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2023 - 2024 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#include <thrust/host_vector.h>
|
||||
#include <thrust/device_vector.h>
|
||||
|
||||
#include <cute/tensor.hpp>
|
||||
|
||||
#include "cutlass/util/print_error.hpp"
|
||||
#include "cutlass/util/GPU_Clock.hpp"
|
||||
#if defined(CUTLASS_ENABLE_CUBLAS) && CUTLASS_ENABLE_CUBLAS != 0
|
||||
# include "cutlass/util/cublas_wrappers.hpp"
|
||||
#endif
|
||||
#include "cutlass/util/helper_cuda.hpp"
|
||||
|
||||
template <class MShape, class NShape, class KShape,
|
||||
class TA, class AStride, class ABlockLayout, class AThreadLayout,
|
||||
class TB, class BStride, class BBlockLayout, class BThreadLayout,
|
||||
class TC, class CStride, class CBlockLayout, class CThreadLayout,
|
||||
class Alpha, class Beta>
|
||||
__global__ static
|
||||
__launch_bounds__(decltype(size(CThreadLayout{}))::value)
|
||||
void
|
||||
gemm_device(MShape M, NShape N, KShape K,
|
||||
TA const* A, AStride dA, ABlockLayout blockA, AThreadLayout tA,
|
||||
TB const* B, BStride dB, BBlockLayout blockB, BThreadLayout tB,
|
||||
TC * C, CStride dC, CBlockLayout , CThreadLayout tC,
|
||||
Alpha alpha, Beta beta)
|
||||
{
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
|
||||
// Preconditions
|
||||
CUTE_STATIC_ASSERT(is_static<ABlockLayout>::value);
|
||||
CUTE_STATIC_ASSERT(is_static<BBlockLayout>::value);
|
||||
CUTE_STATIC_ASSERT(is_static<CBlockLayout>::value);
|
||||
|
||||
CUTE_STATIC_ASSERT(is_static<AThreadLayout>::value);
|
||||
CUTE_STATIC_ASSERT(is_static<BThreadLayout>::value);
|
||||
CUTE_STATIC_ASSERT(is_static<CThreadLayout>::value);
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size(tA) == size(tC));
|
||||
CUTE_STATIC_ASSERT_V(size(tB) == size(tC));
|
||||
|
||||
//CUTE_STATIC_ASSERT_V(shape<0>(blockA) == shape<0>(blockC)); // BLK_M
|
||||
//CUTE_STATIC_ASSERT_V(shape<0>(blockB) == shape<1>(blockC)); // BLK_N
|
||||
CUTE_STATIC_ASSERT_V(shape<1>(blockA) == shape<1>(blockB)); // BLK_K
|
||||
|
||||
// Shared memory buffers
|
||||
__shared__ TA smemA[cosize_v<ABlockLayout>];
|
||||
__shared__ TB smemB[cosize_v<BBlockLayout>];
|
||||
auto sA = make_tensor(make_smem_ptr(smemA), blockA); // (BLK_M,BLK_K)
|
||||
auto sB = make_tensor(make_smem_ptr(smemB), blockB); // (BLK_N,BLK_K)
|
||||
|
||||
// Represent the full tensors
|
||||
auto mA = make_tensor(make_gmem_ptr(A), make_shape(M,K), dA); // (M,K)
|
||||
auto mB = make_tensor(make_gmem_ptr(B), make_shape(N,K), dB); // (N,K)
|
||||
auto mC = make_tensor(make_gmem_ptr(C), make_shape(M,N), dC); // (M,N)
|
||||
|
||||
// Get the appropriate blocks for this thread block --
|
||||
// potential for thread block locality
|
||||
auto blk_shape = make_shape(size<0>(sA), size<0>(sB), size<1>(sB));// (BLK_M,BLK_N,BLK_K)
|
||||
auto blk_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k)
|
||||
|
||||
auto gA = local_tile(mA, blk_shape, blk_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k)
|
||||
auto gB = local_tile(mB, blk_shape, blk_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
|
||||
auto gC = local_tile(mC, blk_shape, blk_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N)
|
||||
|
||||
//
|
||||
// Partition the copying of A and B tiles across the threads
|
||||
//
|
||||
|
||||
// TUTORIAL: Example of simple partitioning of A|B tiles over tA|tB
|
||||
// Default is a raked partition, but can be changed with Step<X,Y> parameter
|
||||
|
||||
auto tAgA = local_partition(gA, tA, threadIdx.x); // (THR_M,THR_K,k)
|
||||
auto tAsA = local_partition(sA, tA, threadIdx.x); // (THR_M,THR_K)
|
||||
|
||||
auto tBgB = local_partition(gB, tB, threadIdx.x); // (THR_N,THR_K,k)
|
||||
auto tBsB = local_partition(sB, tB, threadIdx.x); // (THR_N,THR_K)
|
||||
|
||||
//
|
||||
// Define C accumulators and A/B partitioning
|
||||
//
|
||||
|
||||
// TUTORIAL: Example of partitioning via projections of tC
|
||||
|
||||
// Partition sA (M,K) by the rows of tC
|
||||
auto tCsA = local_partition(sA, tC, threadIdx.x, Step<_1, X>{}); // (THR_M,BLK_K)
|
||||
// Partition sB (N,K) by the cols of tC
|
||||
auto tCsB = local_partition(sB, tC, threadIdx.x, Step< X,_1>{}); // (THR_N,BLK_K)
|
||||
// Partition gC (M,N) by the tile of tC
|
||||
auto tCgC = local_partition(gC, tC, threadIdx.x, Step<_1,_1>{}); // (THR_M,THR_N)
|
||||
|
||||
// Allocate the accumulators -- same size as the projected data
|
||||
auto tCrC = make_fragment_like(tCgC); // (THR_M,THR_N)
|
||||
|
||||
// Clear the accumulators
|
||||
clear(tCrC);
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print("mA\n");
|
||||
print(mA.shape()); print("\n"); print(mA.stride());
|
||||
print("\n\ngA\n");
|
||||
print(gA.shape()); print("\n"); print(gA.stride());
|
||||
print("\n\ntAgA\n");
|
||||
print(tAgA.shape()); print("\n"); print(tAgA.stride());
|
||||
print("\n\nsA\n");
|
||||
print(sA.shape()); print("\n"); print(sA.stride());
|
||||
print("\n\ntAsA\n");
|
||||
print(tAsA.shape()); print("\n"); print(tAsA.stride());
|
||||
print("\n\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print("mB\n");
|
||||
print(mB.shape()); print("\n"); print(mB.stride());
|
||||
print("\n\ngB\n");
|
||||
print(gB.shape()); print("\n"); print(gB.stride());
|
||||
print("\n\ntBgB\n");
|
||||
print(tBgB.shape()); print("\n"); print(tBgB.stride());
|
||||
print("\n\nsB\n");
|
||||
print(sB.shape()); print("\n"); print(sB.stride());
|
||||
print("\n\ntBsB\n");
|
||||
print(tBsB.shape()); print("\n"); print(tBsB.stride());
|
||||
print("\n\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print("mC\n");
|
||||
print(mC.shape()); print("\n"); print(mC.stride());
|
||||
print("\n\ngC\n");
|
||||
print(gC.shape()); print("\n"); print(gC.stride());
|
||||
print("\n\ntCsA\n");
|
||||
print(tCsA.shape()); print("\n"); print(tCsA.stride());
|
||||
print("\n\ntCsB\n");
|
||||
print(tCsB.shape()); print("\n"); print(tCsB.stride());
|
||||
print("\n\ntCgC\n");
|
||||
print(tCgC.shape()); print("\n"); print(tCgC.stride());
|
||||
print("\n\ntCrC\n");
|
||||
print(tCrC.shape()); print("\n"); print(tCrC.stride());
|
||||
print("\n\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 1
|
||||
|
||||
// TUTORIAL: Example of a very simple compute loop
|
||||
// Data is read from global to shared memory via the tA|tB partitioning
|
||||
// gemm(.) operates on the shared memory directly via the tC partitioning
|
||||
|
||||
auto k_max = size<2>(tAgA);
|
||||
|
||||
for (int k = 0; k < k_max; ++k)
|
||||
{
|
||||
// Copy gmem to smem
|
||||
copy(tAgA(_,_,k), tAsA);
|
||||
copy(tBgB(_,_,k), tBsB);
|
||||
|
||||
// In case copy uses cp.async, make sure that the cp.async
|
||||
// instructions are ordered with respect to other cp.async
|
||||
// instructions (fence), then wait on all the outstanding copy
|
||||
// operations (wait<0>()). __syncthreads() alone does not do
|
||||
// this.
|
||||
//
|
||||
// NOTE: cp_async_wait<0>() currently issues cp.async.wait_all.
|
||||
// This is equivalent to cp.async.commit_group followed by
|
||||
// cp.async_wait_group 0. This should make the first
|
||||
// cp_async_fence() (which also issues cp.async.commit_group)
|
||||
// redundant. The tutorial works as-is, so we'll leave the
|
||||
// redundant fence in for now and study its removal later.
|
||||
cp_async_fence();
|
||||
cp_async_wait<0>();
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Compute gemm on smem
|
||||
gemm(tCsA, tCsB, tCrC);
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
//
|
||||
// Epilogue
|
||||
//
|
||||
|
||||
axpby(alpha, tCrC, beta, tCgC);
|
||||
}
|
||||
|
||||
|
||||
template <typename TA, typename TB, typename TC,
|
||||
typename Alpha, typename Beta>
|
||||
void
|
||||
gemm(int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Define shapes (dynamic)
|
||||
auto M = int(m);
|
||||
auto N = int(n);
|
||||
auto K = int(k);
|
||||
|
||||
// Define strides (mixed)
|
||||
auto dA = make_stride(Int<1>{}, ldA);
|
||||
auto dB = make_stride(Int<1>{}, ldB);
|
||||
auto dC = make_stride(Int<1>{}, ldC);
|
||||
|
||||
// Define block sizes (static)
|
||||
auto bM = Int<128>{};
|
||||
auto bN = Int<128>{};
|
||||
auto bK = Int< 8>{};
|
||||
|
||||
// Define the block layouts (static)
|
||||
auto sA = make_layout(make_shape(bM,bK));
|
||||
auto sB = make_layout(make_shape(bN,bK));
|
||||
auto sC = make_layout(make_shape(bM,bN));
|
||||
|
||||
// Define the thread layouts (static)
|
||||
auto tA = make_layout(make_shape(Int<32>{}, Int< 8>{}));
|
||||
auto tB = make_layout(make_shape(Int<32>{}, Int< 8>{}));
|
||||
auto tC = make_layout(make_shape(Int<16>{}, Int<16>{}));
|
||||
|
||||
dim3 dimBlock(size(tC));
|
||||
dim3 dimGrid(ceil_div(size(M), size(bM)),
|
||||
ceil_div(size(N), size(bN)));
|
||||
gemm_device
|
||||
<<< dimGrid, dimBlock, 0, stream >>>
|
||||
(M, N, K,
|
||||
A, dA, sA, tA,
|
||||
B, dB, sB, tB,
|
||||
C, dC, sC, tC,
|
||||
alpha, beta);
|
||||
}
|
||||
|
||||
#include <cstdlib>
|
||||
#include <cstdio>
|
||||
#include <cassert>
|
||||
|
||||
void test_gemm(int m, int n, int k)
|
||||
{
|
||||
cute::device_init(0);
|
||||
|
||||
std::cout << "M = " << m << std::endl;
|
||||
std::cout << "N = " << n << std::endl;
|
||||
std::cout << "K = " << k << std::endl;
|
||||
|
||||
using TA = float;
|
||||
using TB = float;
|
||||
using TC = float;
|
||||
using TI = float;
|
||||
|
||||
thrust::host_vector<TA> h_A(m*k);
|
||||
thrust::host_vector<TB> h_B(n*k);
|
||||
thrust::host_vector<TC> h_C(m*n);
|
||||
|
||||
for (int j = 0; j < m*k; ++j) h_A[j] = static_cast<TA>( 2*(rand() / double(RAND_MAX)) - 1 );
|
||||
for (int j = 0; j < n*k; ++j) h_B[j] = static_cast<TB>( 2*(rand() / double(RAND_MAX)) - 1 );
|
||||
for (int j = 0; j < m*n; ++j) h_C[j] = static_cast<TC>(-1);
|
||||
|
||||
thrust::device_vector<TA> d_A = h_A;
|
||||
thrust::device_vector<TB> d_B = h_B;
|
||||
thrust::device_vector<TC> d_C = h_C;
|
||||
|
||||
TI alpha = 1.0;
|
||||
TI beta = 0.0;
|
||||
|
||||
double gflops = (2.0*m*n*k) * 1e-9;
|
||||
|
||||
const int timing_iterations = 100;
|
||||
GPU_Clock timer;
|
||||
|
||||
#if defined(CUTLASS_ENABLE_CUBLAS) && CUTLASS_ENABLE_CUBLAS != 0
|
||||
//
|
||||
// cuBLas
|
||||
//
|
||||
|
||||
cublasHandle_t handle;
|
||||
cublasCreate(&handle);
|
||||
|
||||
// Run once
|
||||
d_C = h_C;
|
||||
blam::cublas::gemm(handle, CUBLAS_OP_N, CUBLAS_OP_T,
|
||||
m, n, k,
|
||||
&alpha,
|
||||
d_A.data().get(), m,
|
||||
d_B.data().get(), n,
|
||||
&beta,
|
||||
d_C.data().get(), m);
|
||||
CUTE_CHECK_LAST();
|
||||
|
||||
thrust::host_vector<TC> cublas_result = d_C;
|
||||
|
||||
// Timing iterations
|
||||
timer.start();
|
||||
for (int i = 0; i < timing_iterations; ++i) {
|
||||
blam::cublas::gemm(handle, CUBLAS_OP_N, CUBLAS_OP_T,
|
||||
m, n, k,
|
||||
&alpha,
|
||||
d_A.data().get(), m,
|
||||
d_B.data().get(), n,
|
||||
&beta,
|
||||
d_C.data().get(), m);
|
||||
}
|
||||
double cublas_time = timer.seconds() / timing_iterations;
|
||||
CUTE_CHECK_LAST();
|
||||
printf("CUBLAS_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cublas_time, cublas_time*1000);
|
||||
|
||||
#else
|
||||
|
||||
std::cout << "Verification by comparison with cuBLAS is disabled, "
|
||||
"either because the CMake option CUTLASS_ENABLE_CUBLAS "
|
||||
"was explicitly set to OFF, or because CMake could not find cuBLAS. "
|
||||
"If you would like to enable verification with cuBLAS, "
|
||||
"please set the CMake option CUTLASS_ENABLE_CUBLAS to ON, "
|
||||
"rerun CMake, and recompile this example.\n";
|
||||
|
||||
#endif // CUTLASS_ENABLE_CUBLAS
|
||||
|
||||
//
|
||||
// CuTe
|
||||
//
|
||||
|
||||
// Run once (and check)
|
||||
d_C = h_C;
|
||||
gemm(m, n, k,
|
||||
alpha,
|
||||
d_A.data().get(), m,
|
||||
d_B.data().get(), n,
|
||||
beta,
|
||||
d_C.data().get(), m);
|
||||
CUTE_CHECK_LAST();
|
||||
thrust::host_vector<TC> cute_result = d_C;
|
||||
|
||||
// Timing iterations
|
||||
timer.start();
|
||||
for (int i = 0; i < timing_iterations; ++i) {
|
||||
gemm(m, n, k,
|
||||
alpha,
|
||||
d_A.data().get(), m,
|
||||
d_B.data().get(), n,
|
||||
beta,
|
||||
d_C.data().get(), m);
|
||||
}
|
||||
double cute_time = timer.seconds() / timing_iterations;
|
||||
CUTE_CHECK_LAST();
|
||||
printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000);
|
||||
|
||||
#if defined(CUTLASS_ENABLE_CUBLAS) && CUTLASS_ENABLE_CUBLAS != 0
|
||||
printf("Empirical Perf: %.1f%%\n", (cublas_time / cute_time) * 100);
|
||||
|
||||
auto host_matrix_to_const_column_major_cute_tensor =
|
||||
[](const auto& X, int num_rows, int num_cols, int LDX) {
|
||||
const auto shape = cute::Shape<int, int>{num_rows, num_cols};
|
||||
const auto strides = cute::Stride<int, int>{1, LDX};
|
||||
return cute::make_tensor(X.data(), cute::make_layout(shape, strides));
|
||||
};
|
||||
|
||||
const auto A_view = host_matrix_to_const_column_major_cute_tensor(h_A, m, k, m);
|
||||
// B^T is k x n, so B is n x k.
|
||||
const auto B_view = host_matrix_to_const_column_major_cute_tensor(h_B, n, k, n);
|
||||
const auto C_computed_view = host_matrix_to_const_column_major_cute_tensor(cute_result, m, n, m);
|
||||
const auto C_expected_view = host_matrix_to_const_column_major_cute_tensor(cublas_result, m, n, m);
|
||||
print_matrix_multiply_mollified_relative_error("float", A_view, B_view, C_computed_view, C_expected_view);
|
||||
|
||||
#endif // CUTLASS_ENABLE_CUBLAS
|
||||
}
|
||||
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
int m = 5120;
|
||||
if (argc >= 2)
|
||||
sscanf(argv[1], "%d", &m);
|
||||
|
||||
int n = 5120;
|
||||
if (argc >= 3)
|
||||
sscanf(argv[2], "%d", &n);
|
||||
|
||||
int k = 4096;
|
||||
if (argc >= 4)
|
||||
sscanf(argv[3], "%d", &k);
|
||||
|
||||
test_gemm(m, n, k);
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,526 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2023 - 2024 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#include <cstdlib>
|
||||
#include <cstdio>
|
||||
#include <cassert>
|
||||
|
||||
#include <thrust/host_vector.h>
|
||||
#include <thrust/device_vector.h>
|
||||
|
||||
#include <cute/tensor.hpp>
|
||||
|
||||
#include "cutlass/util/print_error.hpp"
|
||||
#include "cutlass/util/GPU_Clock.hpp"
|
||||
#include "cutlass/util/helper_cuda.hpp"
|
||||
|
||||
template <class ProblemShape, class CtaTiler,
|
||||
class TA, class AStride, class ASmemLayout, class TiledCopyA,
|
||||
class TB, class BStride, class BSmemLayout, class TiledCopyB,
|
||||
class TC, class CStride, class CSmemLayout, class TiledMma,
|
||||
class Alpha, class Beta>
|
||||
__global__ static
|
||||
__launch_bounds__(decltype(size(TiledMma{}))::value)
|
||||
void
|
||||
gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler,
|
||||
TA const* A, AStride dA, ASmemLayout sA_layout, TiledCopyA copy_a,
|
||||
TB const* B, BStride dB, BSmemLayout sB_layout, TiledCopyB copy_b,
|
||||
TC * C, CStride dC, CSmemLayout , TiledMma mma,
|
||||
Alpha alpha, Beta beta)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Preconditions
|
||||
CUTE_STATIC_ASSERT_V(rank(shape_MNK) == Int<3>{}); // (M, N, K)
|
||||
CUTE_STATIC_ASSERT_V(rank(cta_tiler) == Int<3>{}); // (BLK_M, BLK_N, BLK_K)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size(copy_a) == size(mma)); // NumThreads
|
||||
CUTE_STATIC_ASSERT_V(size(copy_b) == size(mma)); // NumThreads
|
||||
|
||||
static_assert(is_static<ASmemLayout>::value);
|
||||
static_assert(is_static<BSmemLayout>::value);
|
||||
static_assert(is_static<CSmemLayout>::value);
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<0>(ASmemLayout{}) == size<0>(cta_tiler)); // BLK_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<0>(cta_tiler)); // BLK_M
|
||||
CUTE_STATIC_ASSERT_V(size<0>(BSmemLayout{}) == size<1>(cta_tiler)); // BLK_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<1>(cta_tiler)); // BLK_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(ASmemLayout{}) == size<2>(cta_tiler)); // BLK_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(BSmemLayout{}) == size<2>(cta_tiler)); // BLK_K
|
||||
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<0,2>(shape_MNK), dA)); // dA strides for shape MK
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<1,2>(shape_MNK), dB)); // dB strides for shape NK
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<0,1>(shape_MNK), dC)); // dC strides for shape MN
|
||||
|
||||
//
|
||||
// Full and Tiled Tensors
|
||||
//
|
||||
|
||||
// Represent the full tensors
|
||||
Tensor mA = make_tensor(make_gmem_ptr(A), select<0,2>(shape_MNK), dA); // (M,K)
|
||||
Tensor mB = make_tensor(make_gmem_ptr(B), select<1,2>(shape_MNK), dB); // (N,K)
|
||||
Tensor mC = make_tensor(make_gmem_ptr(C), select<0,1>(shape_MNK), dC); // (M,N)
|
||||
|
||||
// Get the appropriate blocks for this thread block
|
||||
auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k)
|
||||
Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
|
||||
Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N)
|
||||
|
||||
// Shared memory buffers
|
||||
__shared__ TA smemA[cosize_v<ASmemLayout>];
|
||||
__shared__ TB smemB[cosize_v<BSmemLayout>];
|
||||
Tensor sA = make_tensor(make_smem_ptr(smemA), sA_layout); // (BLK_M,BLK_K)
|
||||
Tensor sB = make_tensor(make_smem_ptr(smemB), sB_layout); // (BLK_N,BLK_K)
|
||||
|
||||
//
|
||||
// Partition the copying of A and B tiles across the threads
|
||||
//
|
||||
|
||||
// TUTORIAL: Example of partitioning via a TiledCopy
|
||||
|
||||
ThrCopy thr_copy_a = copy_a.get_slice(threadIdx.x);
|
||||
Tensor tAgA = thr_copy_a.partition_S(gA); // (CPY,CPY_M,CPY_K,k)
|
||||
Tensor tAsA = thr_copy_a.partition_D(sA); // (CPY,CPY_M,CPY_K)
|
||||
Tensor tArA = make_fragment_like(tAsA); // (CPY,CPY_M,CPY_K)
|
||||
|
||||
ThrCopy thr_copy_b = copy_b.get_slice(threadIdx.x);
|
||||
Tensor tBgB = thr_copy_b.partition_S(gB); // (CPY,CPY_N,CPY_K,k)
|
||||
Tensor tBsB = thr_copy_b.partition_D(sB); // (CPY,CPY_N,CPY_K)
|
||||
Tensor tBrB = make_fragment_like(tBsB); // (CPY,CPY_N,CPY_K)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tAsA)); // CPY_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tArA)); // CPY_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tAgA) == size<2>(tAsA)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tAgA) == size<2>(tArA)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBsB)); // CPY_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBrB)); // CPY_N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tBgB) == size<2>(tBsB)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tBgB) == size<2>(tBrB)); // CPY_K
|
||||
|
||||
// Copy gmem to rmem for k_tile=0
|
||||
copy(copy_a, tAgA(_,_,_,0), tArA);
|
||||
copy(copy_b, tBgB(_,_,_,0), tBrB);
|
||||
//
|
||||
// Define A/B partitioning and C accumulators
|
||||
//
|
||||
|
||||
// TUTORIAL: Example of partitioning via a TiledMMA
|
||||
|
||||
ThrMMA thr_mma = mma.get_slice(threadIdx.x);
|
||||
Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K)
|
||||
Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K)
|
||||
Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N)
|
||||
|
||||
// Allocate registers for pipelining
|
||||
Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K)
|
||||
Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K)
|
||||
// Allocate the accumulators -- same size as the projected data
|
||||
Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N)
|
||||
|
||||
CUTE_STATIC_ASSERT_V( shape(tCrA) == shape(tCsA)); // (MMA,MMA_M,MMA_K)
|
||||
CUTE_STATIC_ASSERT_V( shape(tCrB) == shape(tCsB)); // (MMA,MMA_N,MMA_K)
|
||||
CUTE_STATIC_ASSERT_V( shape(tCrC) == shape(tCgC)); // (MMA,MMA_M,MMA_N)
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCgC) == size<1>(tCsA)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCgC) == size<1>(tCsB)); // MMA_N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // MMA_K
|
||||
|
||||
// Clear the accumulators
|
||||
clear(tCrC);
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mA : "); print( mA); print("\n");
|
||||
print(" gA : "); print( gA); print("\n");
|
||||
print(" sA : "); print( sA); print("\n");
|
||||
print("tAgA : "); print(tAgA); print("\n");
|
||||
print("tAsA : "); print(tAsA); print("\n");
|
||||
print("tArA : "); print(tArA); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mB : "); print( mB); print("\n");
|
||||
print(" gB : "); print( gB); print("\n");
|
||||
print(" sB : "); print( sB); print("\n");
|
||||
print("tBgB : "); print(tBgB); print("\n");
|
||||
print("tBsB : "); print(tBsB); print("\n");
|
||||
print("tArA : "); print(tArA); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mC : "); print( mC); print("\n");
|
||||
print(" gC : "); print( gC); print("\n");
|
||||
print("tCsA : "); print(tCsA); print("\n");
|
||||
print("tCsB : "); print(tCsB); print("\n");
|
||||
print("tCgC : "); print(tCgC); print("\n");
|
||||
print("tCrC : "); print(tCrC); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 1
|
||||
|
||||
// Copy rmem to smem
|
||||
copy(tArA, tAsA);
|
||||
copy(tBrB, tBsB);
|
||||
__syncthreads();
|
||||
|
||||
//
|
||||
// PIPELINED MAIN LOOP
|
||||
// TUTORIAL: Example of a gemm loop that pipelines shared memory AND register memory
|
||||
// Data is read from global to registers, then to shared via the tA|tB partitions
|
||||
// Data is then copied from shared to registers in multiple waves via the tC partitions
|
||||
// and gemm(.) operates on the current register wave
|
||||
//
|
||||
|
||||
// Load A, B shmem->regs for k_block=0
|
||||
copy(tCsA(_,_,0), tCrA(_,_,0));
|
||||
copy(tCsB(_,_,0), tCrB(_,_,0));
|
||||
auto K_TILE_MAX = size<3>(tAgA);
|
||||
auto K_BLOCK_MAX = size<2>(tCrA);
|
||||
|
||||
CUTE_NO_UNROLL
|
||||
for (int k_tile = 0; k_tile < K_TILE_MAX; ++k_tile)
|
||||
{
|
||||
// Pipeline the k-mode of the block registers
|
||||
CUTE_UNROLL
|
||||
for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block)
|
||||
{
|
||||
if (k_block == K_BLOCK_MAX - 1)
|
||||
{
|
||||
// Copy rmem to smem
|
||||
__syncthreads();
|
||||
copy(tArA, tAsA);
|
||||
copy(tBrB, tBsB);
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Copy smem to rmem for k_block+1
|
||||
int k_block_next = (k_block + 1) % K_BLOCK_MAX;
|
||||
copy(tCsA(_,_,k_block_next), tCrA(_,_,k_block_next));
|
||||
copy(tCsB(_,_,k_block_next), tCrB(_,_,k_block_next));
|
||||
if (k_block == 0)
|
||||
{
|
||||
// Copy gmem to rmem for k_tile+1
|
||||
int k_tile_next = (k_tile + 1 < K_TILE_MAX) ? k_tile + 1 : k_tile;
|
||||
copy(copy_a, tAgA(_,_,_,k_tile_next), tArA);
|
||||
copy(copy_b, tBgB(_,_,_,k_tile_next), tBrB);
|
||||
}
|
||||
// Thread-level register gemm for k_block
|
||||
gemm(mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);
|
||||
} // k_block
|
||||
} // k_tile
|
||||
|
||||
#endif
|
||||
|
||||
//
|
||||
// Epilogue
|
||||
//
|
||||
|
||||
axpby(alpha, tCrC, beta, tCgC);
|
||||
}
|
||||
|
||||
// Setup params for a NT GEMM
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm_nt(int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Define shapes (dynamic)
|
||||
auto M = int(m);
|
||||
auto N = int(n);
|
||||
auto K = int(k);
|
||||
auto prob_shape = make_shape(M, N, K); // (M, N, K)
|
||||
|
||||
// Define NT strides (mixed)
|
||||
auto dA = make_stride(Int<1>{}, ldA); // (dM, dK)
|
||||
auto dB = make_stride(Int<1>{}, ldB); // (dN, dK)
|
||||
auto dC = make_stride(Int<1>{}, ldC); // (dM, dN)
|
||||
|
||||
// Define CTA tile sizes (static)
|
||||
auto bM = Int<128>{};
|
||||
auto bN = Int<128>{};
|
||||
auto bK = Int< 8>{};
|
||||
auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K)
|
||||
|
||||
// Define the smem layouts (static)
|
||||
auto sA = make_layout(make_shape(bM, bK)); // (m,k) -> smem_idx; m-major
|
||||
auto sB = make_layout(make_shape(bN, bK)); // (n,k) -> smem_idx; n-major
|
||||
auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx; m-major
|
||||
|
||||
// Define the thread layouts (static)
|
||||
TiledCopy copyA = make_tiled_copy(Copy_Atom<UniversalCopy<uint128_t>, TA>{},
|
||||
Layout<Shape<_32,_8>>{}, // Thr layout 32x8 m-major
|
||||
Layout<Shape< _4,_1>>{}); // Val layout 4x1 m-major
|
||||
TiledCopy copyB = make_tiled_copy(Copy_Atom<UniversalCopy<uint128_t>, TB>{},
|
||||
Layout<Shape<_32,_8>>{}, // Thr layout 32x8 n-major
|
||||
Layout<Shape< _4,_1>>{}); // Val layout 4x1 n-major
|
||||
|
||||
TiledMMA mmaC = make_tiled_mma(UniversalFMA<TC,TA,TB>{},
|
||||
Layout<Shape<_16,_16,_1>>{}); // 16x16x1 TiledMMA
|
||||
|
||||
#if 0
|
||||
print(copyA);
|
||||
print(copyB);
|
||||
print(mmaC);
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
print_latex(copyA);
|
||||
print_latex(copyB);
|
||||
print_latex(mmaC);
|
||||
#endif
|
||||
|
||||
dim3 dimBlock(size(mmaC));
|
||||
dim3 dimGrid(size(ceil_div(M, bM)),
|
||||
size(ceil_div(N, bN)));
|
||||
gemm_device<<<dimGrid, dimBlock, 0, stream>>>
|
||||
(prob_shape, cta_tiler,
|
||||
A, dA, sA, copyA,
|
||||
B, dB, sB, copyB,
|
||||
C, dC, sC, mmaC,
|
||||
alpha, beta);
|
||||
}
|
||||
|
||||
// Setup params for a TN GEMM
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm_tn(int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Define shapes (dynamic)
|
||||
auto M = int(m);
|
||||
auto N = int(n);
|
||||
auto K = int(k);
|
||||
auto prob_shape = make_shape(M, N, K); // (M, N, K)
|
||||
|
||||
// Define TN strides (mixed)
|
||||
auto dA = make_stride(ldA, Int<1>{}); // (dM, dK)
|
||||
auto dB = make_stride(ldB, Int<1>{}); // (dN, dK)
|
||||
auto dC = make_stride(Int<1>{}, ldC); // (dM, dN)
|
||||
|
||||
// Define CTA tile sizes (static)
|
||||
auto bM = Int<128>{};
|
||||
auto bN = Int<128>{};
|
||||
auto bK = Int< 8>{};
|
||||
auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K)
|
||||
|
||||
// Define the smem layouts (static)
|
||||
auto sA = make_layout(make_shape ( bM, bK),
|
||||
make_stride(Int<1>{}, bM+Int<1>{})); // (m,k) -> smem_idx; padded m-major
|
||||
auto sB = make_layout(make_shape ( bN, bK),
|
||||
make_stride(Int<1>{}, bN+Int<1>{})); // (n,k) -> smem_idx; padded n-major
|
||||
auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx
|
||||
|
||||
// Define the thread layouts (static)
|
||||
|
||||
TiledCopy copyA = make_tiled_copy(Copy_Atom<UniversalCopy<TA>, TA>{},
|
||||
Layout<Shape<_32,_8>,Stride<_8,_1>>{}, // Thr layout 32x8 k-major
|
||||
Layout<Shape< _1,_1>>{}); // Val layout 1x1
|
||||
TiledCopy copyB = make_tiled_copy(Copy_Atom<UniversalCopy<TB>, TB>{},
|
||||
Layout<Shape<_32,_8>,Stride<_8,_1>>{}, // Thr layout 32x8 k-major
|
||||
Layout<Shape< _1,_1>>{}); // Val layout 1x1
|
||||
|
||||
TiledMMA mmaC = make_tiled_mma(UniversalFMA<TC,TA,TB>{},
|
||||
Layout<Shape<_16,_16,_1>>{}); // 16x16x1 TiledMMA
|
||||
|
||||
#if 0
|
||||
print(copyA);
|
||||
print(copyB);
|
||||
print(mmaC);
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
print_latex(copyA);
|
||||
print_latex(copyB);
|
||||
print_latex(mmaC);
|
||||
#endif
|
||||
|
||||
dim3 dimBlock(size(mmaC));
|
||||
dim3 dimGrid(size(ceil_div(M, bM)),
|
||||
size(ceil_div(N, bN)));
|
||||
gemm_device<<<dimGrid, dimBlock, 0, stream>>>
|
||||
(prob_shape, cta_tiler,
|
||||
A, dA, sA, copyA,
|
||||
B, dB, sB, copyB,
|
||||
C, dC, sC, mmaC,
|
||||
alpha, beta);
|
||||
}
|
||||
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm(char transA, char transB, int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
if (transA == 'N' && transB == 'T') {
|
||||
return gemm_nt(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream);
|
||||
} else
|
||||
if (transA == 'T' && transB == 'N') {
|
||||
return gemm_tn(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream);
|
||||
}
|
||||
assert(false && "Not implemented");
|
||||
}
|
||||
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
cudaDeviceProp props;
|
||||
cudaError_t error = cudaGetDeviceProperties(&props, 0);
|
||||
if (error != cudaSuccess) {
|
||||
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
if (props.major < 7) {
|
||||
std::cout << "This example requires an Volta GPU or newer (CC >= 70)" << std::endl;
|
||||
// Return 0 so tests pass if run on unsupported architectures or CUDA Toolkits.
|
||||
return 0;
|
||||
}
|
||||
|
||||
int m = 5120;
|
||||
if (argc >= 2)
|
||||
sscanf(argv[1], "%d", &m);
|
||||
|
||||
int n = 5120;
|
||||
if (argc >= 3)
|
||||
sscanf(argv[2], "%d", &n);
|
||||
|
||||
int k = 4096;
|
||||
if (argc >= 4)
|
||||
sscanf(argv[3], "%d", &k);
|
||||
|
||||
char transA = 'N';
|
||||
if (argc >= 5)
|
||||
sscanf(argv[4], "%c", &transA);
|
||||
|
||||
char transB = 'T';
|
||||
if (argc >= 6)
|
||||
sscanf(argv[5], "%c", &transB);
|
||||
|
||||
using TA = float;
|
||||
using TB = float;
|
||||
using TC = float;
|
||||
using TI = float;
|
||||
|
||||
TI alpha = 1.0;
|
||||
TI beta = 0.0;
|
||||
|
||||
std::cout << "M = " << m << std::endl;
|
||||
std::cout << "N = " << n << std::endl;
|
||||
std::cout << "K = " << k << std::endl;
|
||||
std::cout << "C = A^" << transA << " B^" << transB << std::endl;
|
||||
|
||||
thrust::host_vector<TA> h_A(m*k);
|
||||
thrust::host_vector<TB> h_B(n*k);
|
||||
thrust::host_vector<TC> h_C(m*n);
|
||||
|
||||
for (int j = 0; j < m*k; ++j) h_A[j] = static_cast<TA>( 2*(rand() / double(RAND_MAX)) - 1 );
|
||||
for (int j = 0; j < n*k; ++j) h_B[j] = static_cast<TB>( 2*(rand() / double(RAND_MAX)) - 1 );
|
||||
for (int j = 0; j < m*n; ++j) h_C[j] = static_cast<TC>(-1);
|
||||
|
||||
thrust::device_vector<TA> d_A = h_A;
|
||||
thrust::device_vector<TB> d_B = h_B;
|
||||
thrust::device_vector<TC> d_C = h_C;
|
||||
|
||||
double gflops = (2.0*m*n*k) * 1e-9;
|
||||
|
||||
const int timing_iterations = 100;
|
||||
GPU_Clock timer;
|
||||
|
||||
int ldA = 0, ldB = 0, ldC = m;
|
||||
|
||||
if (transA == 'N') {
|
||||
ldA = m;
|
||||
} else if (transA == 'T') {
|
||||
ldA = k;
|
||||
} else {
|
||||
assert(false);
|
||||
}
|
||||
|
||||
if (transB == 'N') {
|
||||
ldB = k;
|
||||
} else if (transB == 'T') {
|
||||
ldB = n;
|
||||
} else {
|
||||
assert(false);
|
||||
}
|
||||
|
||||
// Run once
|
||||
d_C = h_C;
|
||||
gemm(transA, transB, m, n, k,
|
||||
alpha,
|
||||
d_A.data().get(), ldA,
|
||||
d_B.data().get(), ldB,
|
||||
beta,
|
||||
d_C.data().get(), ldC);
|
||||
CUTE_CHECK_LAST();
|
||||
thrust::host_vector<TC> cute_result = d_C;
|
||||
|
||||
// Timing iterations
|
||||
timer.start();
|
||||
for (int i = 0; i < timing_iterations; ++i) {
|
||||
gemm(transA, transB, m, n, k,
|
||||
alpha,
|
||||
d_A.data().get(), ldA,
|
||||
d_B.data().get(), ldB,
|
||||
beta,
|
||||
d_C.data().get(), ldC);
|
||||
}
|
||||
double cute_time = timer.seconds() / timing_iterations;
|
||||
CUTE_CHECK_LAST();
|
||||
printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000);
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,567 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2023 - 2024 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#include <cstdlib>
|
||||
#include <cstdio>
|
||||
#include <cassert>
|
||||
|
||||
#include <thrust/host_vector.h>
|
||||
#include <thrust/device_vector.h>
|
||||
|
||||
#include <cute/tensor.hpp>
|
||||
|
||||
#include "cutlass/util/print_error.hpp"
|
||||
#include "cutlass/util/GPU_Clock.hpp"
|
||||
#include "cutlass/util/helper_cuda.hpp"
|
||||
|
||||
template <class ProblemShape, class CtaTiler,
|
||||
class TA, class AStride, class ASmemLayout, class TiledCopyA,
|
||||
class TB, class BStride, class BSmemLayout, class TiledCopyB,
|
||||
class TC, class CStride, class CSmemLayout, class TiledMma,
|
||||
class Alpha, class Beta>
|
||||
__global__ static
|
||||
__launch_bounds__(decltype(size(TiledMma{}))::value)
|
||||
void
|
||||
gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler,
|
||||
TA const* A, AStride dA, ASmemLayout sA_layout, TiledCopyA copy_a,
|
||||
TB const* B, BStride dB, BSmemLayout sB_layout, TiledCopyB copy_b,
|
||||
TC * C, CStride dC, CSmemLayout , TiledMma mma,
|
||||
Alpha alpha, Beta beta)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Preconditions
|
||||
CUTE_STATIC_ASSERT_V(rank(shape_MNK) == Int<3>{}); // (M, N, K)
|
||||
CUTE_STATIC_ASSERT_V(rank(cta_tiler) == Int<3>{}); // (BLK_M, BLK_N, BLK_K)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size(copy_a) == size(mma)); // NumThreads
|
||||
CUTE_STATIC_ASSERT_V(size(copy_b) == size(mma)); // NumThreads
|
||||
|
||||
static_assert(is_static<ASmemLayout>::value);
|
||||
static_assert(is_static<BSmemLayout>::value);
|
||||
static_assert(is_static<CSmemLayout>::value);
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<0>(ASmemLayout{}) == size<0>(cta_tiler)); // BLK_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<0>(cta_tiler)); // BLK_M
|
||||
CUTE_STATIC_ASSERT_V(size<0>(BSmemLayout{}) == size<1>(cta_tiler)); // BLK_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(CSmemLayout{}) == size<1>(cta_tiler)); // BLK_N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(ASmemLayout{}) == size<2>(cta_tiler)); // BLK_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(BSmemLayout{}) == size<2>(cta_tiler)); // BLK_K
|
||||
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<0,2>(shape_MNK), dA)); // dA strides for shape MK
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<1,2>(shape_MNK), dB)); // dB strides for shape NK
|
||||
CUTE_STATIC_ASSERT_V(congruent(select<0,1>(shape_MNK), dC)); // dC strides for shape MN
|
||||
|
||||
//
|
||||
// Full and Tiled Tensors
|
||||
//
|
||||
|
||||
// Represent the full tensors
|
||||
Tensor mA = make_tensor(make_gmem_ptr(A), select<0,2>(shape_MNK), dA); // (M,K)
|
||||
Tensor mB = make_tensor(make_gmem_ptr(B), select<1,2>(shape_MNK), dB); // (N,K)
|
||||
Tensor mC = make_tensor(make_gmem_ptr(C), select<0,1>(shape_MNK), dC); // (M,N)
|
||||
|
||||
// Get the appropriate blocks for this thread block
|
||||
auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _); // (m,n,k)
|
||||
Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
|
||||
Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N)
|
||||
|
||||
// Shared memory buffers
|
||||
__shared__ TA smemA[cosize_v<ASmemLayout>];
|
||||
__shared__ TB smemB[cosize_v<BSmemLayout>];
|
||||
Tensor sA = make_tensor(make_smem_ptr(smemA), sA_layout); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sB = make_tensor(make_smem_ptr(smemB), sB_layout); // (BLK_N,BLK_K,PIPE)
|
||||
|
||||
//
|
||||
// Partition the copying of A and B tiles across the threads
|
||||
//
|
||||
|
||||
ThrCopy thr_copy_a = copy_a.get_slice(threadIdx.x);
|
||||
Tensor tAgA = thr_copy_a.partition_S(gA); // (CPY,CPY_M,CPY_K,k)
|
||||
Tensor tAsA = thr_copy_a.partition_D(sA); // (CPY,CPY_M,CPY_K,PIPE)
|
||||
|
||||
ThrCopy thr_copy_b = copy_b.get_slice(threadIdx.x);
|
||||
Tensor tBgB = thr_copy_b.partition_S(gB); // (CPY,CPY_N,CPY_K,k)
|
||||
Tensor tBsB = thr_copy_b.partition_D(sB); // (CPY,CPY_N,CPY_K,PIPE)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tAgA) == size<1>(tAsA)); // CPY_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tAgA) == size<2>(tAsA)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tBgB) == size<1>(tBsB)); // CPY_N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tBgB) == size<2>(tBsB)); // CPY_K
|
||||
|
||||
//
|
||||
// PREFETCH
|
||||
//
|
||||
|
||||
auto K_PIPE_MAX = size<3>(tAsA);
|
||||
|
||||
// Total count of tiles
|
||||
int k_tile_count = size<3>(tAgA);
|
||||
// Current tile index in gmem to read from
|
||||
int k_tile_next = 0;
|
||||
|
||||
// Start async loads for all pipes but the last
|
||||
CUTE_UNROLL
|
||||
for (int k_pipe = 0; k_pipe < K_PIPE_MAX-1; ++k_pipe) {
|
||||
copy(copy_a, tAgA(_,_,_,k_tile_next), tAsA(_,_,_,k_pipe));
|
||||
copy(copy_b, tBgB(_,_,_,k_tile_next), tBsB(_,_,_,k_pipe));
|
||||
cp_async_fence();
|
||||
--k_tile_count;
|
||||
if (k_tile_count > 0) { ++k_tile_next; }
|
||||
}
|
||||
|
||||
//
|
||||
// Define A/B partitioning and C accumulators
|
||||
//
|
||||
|
||||
ThrMMA thr_mma = mma.get_slice(threadIdx.x);
|
||||
Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
|
||||
Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N)
|
||||
|
||||
// Allocate registers for pipelining
|
||||
Tensor tCrA = thr_mma.make_fragment_A(tCsA(_,_,_,0)); // (MMA,MMA_M,MMA_K)
|
||||
Tensor tCrB = thr_mma.make_fragment_B(tCsB(_,_,_,0)); // (MMA,MMA_N,MMA_K)
|
||||
// Allocate the accumulators -- same size as the projected data
|
||||
Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N)
|
||||
|
||||
CUTE_STATIC_ASSERT_V( shape(tCrA) == shape(tCsA)); // (MMA,MMA_M,MMA_K)
|
||||
CUTE_STATIC_ASSERT_V( shape(tCrB) == shape(tCsB)); // (MMA,MMA_N,MMA_K)
|
||||
CUTE_STATIC_ASSERT_V( shape(tCrC) == shape(tCgC)); // (MMA,MMA_M,MMA_N)
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCgC) == size<1>(tCsA)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCgC) == size<1>(tCsB)); // MMA_N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // MMA_K
|
||||
|
||||
// Clear the accumulators
|
||||
clear(tCrC);
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mA : "); print( mA); print("\n");
|
||||
print(" gA : "); print( gA); print("\n");
|
||||
print(" sA : "); print( sA); print("\n");
|
||||
print("tAgA : "); print(tAgA); print("\n");
|
||||
print("tAsA : "); print(tAsA); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mB : "); print( mB); print("\n");
|
||||
print(" gB : "); print( gB); print("\n");
|
||||
print(" sB : "); print( sB); print("\n");
|
||||
print("tBgB : "); print(tBgB); print("\n");
|
||||
print("tBsB : "); print(tBsB); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
if(thread0()) {
|
||||
print(" mC : "); print( mC); print("\n");
|
||||
print(" gC : "); print( gC); print("\n");
|
||||
print("tCsA : "); print(tCsA); print("\n");
|
||||
print("tCsB : "); print(tCsB); print("\n");
|
||||
print("tCgC : "); print(tCgC); print("\n");
|
||||
print("tCrA : "); print(tCrA); print("\n");
|
||||
print("tCrB : "); print(tCrB); print("\n");
|
||||
print("tCrC : "); print(tCrC); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
#if 1
|
||||
|
||||
// Current pipe index in smem to read from
|
||||
int smem_pipe_read = 0;
|
||||
// Current pipe index in smem to write to
|
||||
int smem_pipe_write = K_PIPE_MAX-1;
|
||||
|
||||
// Pipe slice
|
||||
Tensor tCsA_p = tCsA(_,_,_,smem_pipe_read);
|
||||
Tensor tCsB_p = tCsB(_,_,_,smem_pipe_read);
|
||||
|
||||
// Size of the register pipeline
|
||||
auto K_BLOCK_MAX = size<2>(tCrA);
|
||||
|
||||
// PREFETCH register pipeline
|
||||
if (K_BLOCK_MAX > 1) {
|
||||
// Wait until our first prefetched tile is loaded in
|
||||
cp_async_wait<K_PIPE_MAX-2>();
|
||||
__syncthreads();
|
||||
|
||||
// Prefetch the first rmem from the first k-tile
|
||||
copy(tCsA_p(_,_,Int<0>{}), tCrA(_,_,Int<0>{}));
|
||||
copy(tCsB_p(_,_,Int<0>{}), tCrB(_,_,Int<0>{}));
|
||||
}
|
||||
|
||||
//
|
||||
// PIPELINED MAIN LOOP
|
||||
// TUTORIAL: Example of a gemm loop that pipelines shared memory using SM80's cp.async instructions
|
||||
// and explicit pipelines in shared memory.
|
||||
// Data is read from global(k_tile_next) to shared(smem_pipe_write).
|
||||
// Data is read from shared(smem_pipe_read) to registers(k_block_next).
|
||||
// Data is computed on registers(b_block).
|
||||
//
|
||||
// This allows all copies and compute to overlap:
|
||||
// Copy from gmem->smem can overlap with copies from smem->rmem and compute on rmem.
|
||||
// Copy from smem->rmem can overlap with compute on rmem.
|
||||
//
|
||||
|
||||
CUTE_NO_UNROLL
|
||||
while (k_tile_count > -(K_PIPE_MAX-1))
|
||||
{
|
||||
CUTE_UNROLL
|
||||
for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block)
|
||||
{
|
||||
if (k_block == K_BLOCK_MAX - 1)
|
||||
{
|
||||
// Slice the smem_pipe_read smem
|
||||
tCsA_p = tCsA(_,_,_,smem_pipe_read);
|
||||
tCsB_p = tCsB(_,_,_,smem_pipe_read);
|
||||
|
||||
// Commit the smem for smem_pipe_read
|
||||
cp_async_wait<K_PIPE_MAX-2>();
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Load A, B shmem->regs for k_block+1
|
||||
auto k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX; // static
|
||||
copy(tCsA_p(_,_,k_block_next), tCrA(_,_,k_block_next));
|
||||
copy(tCsB_p(_,_,k_block_next), tCrB(_,_,k_block_next));
|
||||
// Copy gmem to smem before computing gemm on each k-pipe
|
||||
if (k_block == 0)
|
||||
{
|
||||
copy(copy_a, tAgA(_,_,_,k_tile_next), tAsA(_,_,_,smem_pipe_write));
|
||||
copy(copy_b, tBgB(_,_,_,k_tile_next), tBsB(_,_,_,smem_pipe_write));
|
||||
cp_async_fence();
|
||||
|
||||
// Advance the gmem tile
|
||||
--k_tile_count;
|
||||
if (k_tile_count > 0) { ++k_tile_next; }
|
||||
|
||||
// Advance the smem pipe
|
||||
smem_pipe_write = smem_pipe_read;
|
||||
++smem_pipe_read;
|
||||
smem_pipe_read = (smem_pipe_read == K_PIPE_MAX) ? 0 : smem_pipe_read;
|
||||
}
|
||||
// Thread-level register gemm for k_block
|
||||
gemm(mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
//
|
||||
// Epilogue
|
||||
//
|
||||
|
||||
axpby(alpha, tCrC, beta, tCgC);
|
||||
}
|
||||
|
||||
// Setup params for a NT GEMM
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm_nt(int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Define shapes (dynamic)
|
||||
auto M = int(m);
|
||||
auto N = int(n);
|
||||
auto K = int(k);
|
||||
auto prob_shape = make_shape(M, N, K); // (M, N, K)
|
||||
|
||||
// Define NT strides (mixed)
|
||||
auto dA = make_stride(Int<1>{}, ldA); // (dM, dK)
|
||||
auto dB = make_stride(Int<1>{}, ldB); // (dN, dK)
|
||||
auto dC = make_stride(Int<1>{}, ldC); // (dM, dN)
|
||||
|
||||
// Define CTA tile sizes (static)
|
||||
auto bM = Int<128>{};
|
||||
auto bN = Int<128>{};
|
||||
auto bK = Int< 8>{};
|
||||
auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K)
|
||||
auto bP = Int<3>{}; // Pipeline
|
||||
|
||||
// Define the smem layouts (static)
|
||||
auto sA = make_layout(make_shape(bM, bK, bP)); // (m,k,p) -> smem_idx; m-major
|
||||
auto sB = make_layout(make_shape(bN, bK, bP)); // (n,k,p) -> smem_idx; n-major
|
||||
auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx; m-major
|
||||
|
||||
// Define the thread layouts (static)
|
||||
|
||||
TiledCopy copyA = make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, TA>{},
|
||||
Layout<Shape<_32,_8>>{}, // Thr layout 32x8 m-major
|
||||
Layout<Shape< _4,_1>>{});// Val layout 4x1 m-major
|
||||
TiledCopy copyB = make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, TB>{},
|
||||
Layout<Shape<_32,_8>>{}, // Thr layout 32x8 n-major
|
||||
Layout<Shape< _4,_1>>{});// Val layout 4x1 n-major
|
||||
|
||||
TiledMMA mmaC = make_tiled_mma(UniversalFMA<TC,TA,TB>{},
|
||||
Layout<Shape<_16,_16,_1>>{}); // 16x16x1 TiledMMA
|
||||
|
||||
#if 0
|
||||
print(copyA);
|
||||
print(copyB);
|
||||
print(mmaC);
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
print_latex(copyA);
|
||||
print_latex(copyB);
|
||||
print_latex(mmaC);
|
||||
#endif
|
||||
|
||||
dim3 dimBlock(size(mmaC));
|
||||
dim3 dimGrid(size(ceil_div(M, bM)),
|
||||
size(ceil_div(N, bN)));
|
||||
gemm_device<<<dimGrid, dimBlock, 0, stream>>>
|
||||
(prob_shape, cta_tiler,
|
||||
A, dA, sA, copyA,
|
||||
B, dB, sB, copyB,
|
||||
C, dC, sC, mmaC,
|
||||
alpha, beta);
|
||||
}
|
||||
|
||||
// Setup params for a NT GEMM
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm_tn(int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
// Define shapes (dynamic)
|
||||
auto M = int(m);
|
||||
auto N = int(n);
|
||||
auto K = int(k);
|
||||
auto prob_shape = make_shape(M, N, K); // (M, N, K)
|
||||
|
||||
// Define TN strides (mixed)
|
||||
auto dA = make_stride(ldA, Int<1>{}); // (dM, dK)
|
||||
auto dB = make_stride(ldB, Int<1>{}); // (dN, dK)
|
||||
auto dC = make_stride(Int<1>{}, ldC); // (dM, dN)
|
||||
|
||||
// Define CTA tile sizes (static)
|
||||
auto bM = Int<128>{};
|
||||
auto bN = Int<128>{};
|
||||
auto bK = Int< 8>{};
|
||||
auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M, BLK_N, BLK_K)
|
||||
auto bP = Int<3>{}; // Pipeline
|
||||
|
||||
// Define the smem layouts (static)
|
||||
auto sA_atom = make_layout(make_shape ( bM, bK),
|
||||
make_stride(Int<1>{}, bM+Int<1>{})); // (m,k) -> smem_idx; padded m-major
|
||||
auto sB_atom = make_layout(make_shape ( bN, bK),
|
||||
make_stride(Int<1>{}, bN+Int<1>{})); // (n,k) -> smem_idx; padded n-major
|
||||
auto sA = tile_to_shape(sA_atom, make_shape(bM, bK, bP));
|
||||
auto sB = tile_to_shape(sA_atom, make_shape(bN, bK, bP));
|
||||
auto sC = make_layout(make_shape(bM, bN)); // (m,n) -> smem_idx
|
||||
|
||||
// Define the thread layouts (static)
|
||||
|
||||
TiledCopy copyA = make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<TA>, TA>{},
|
||||
Layout<Shape<_32,_8>,Stride<_8,_1>>{}, // Thr layout 32x8 k-major
|
||||
Layout<Shape< _1,_1>>{}); // Val layout 1x1
|
||||
TiledCopy copyB = make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<TB>, TB>{},
|
||||
Layout<Shape<_32,_8>,Stride<_8,_1>>{}, // Thr layout 32x8 k-major
|
||||
Layout<Shape< _1,_1>>{}); // Val layout 1x1
|
||||
|
||||
TiledMMA mmaC = make_tiled_mma(UniversalFMA<TC,TA,TB>{},
|
||||
Layout<Shape<_16,_16,_1>>{}); // 16x16x1 TiledMMA
|
||||
|
||||
#if 0
|
||||
print(copyA);
|
||||
print(copyB);
|
||||
print(mmaC);
|
||||
#endif
|
||||
|
||||
#if 0
|
||||
print_latex(copyA);
|
||||
print_latex(copyB);
|
||||
print_latex(mmaC);
|
||||
#endif
|
||||
|
||||
dim3 dimBlock(size(mmaC));
|
||||
dim3 dimGrid(size(ceil_div(M, bM)),
|
||||
size(ceil_div(N, bN)));
|
||||
gemm_device<<<dimGrid, dimBlock, 0, stream>>>
|
||||
(prob_shape, cta_tiler,
|
||||
A, dA, sA, copyA,
|
||||
B, dB, sB, copyB,
|
||||
C, dC, sC, mmaC,
|
||||
alpha, beta);
|
||||
}
|
||||
|
||||
template <class TA, class TB, class TC,
|
||||
class Alpha, class Beta>
|
||||
void
|
||||
gemm(char transA, char transB, int m, int n, int k,
|
||||
Alpha alpha,
|
||||
TA const* A, int ldA,
|
||||
TB const* B, int ldB,
|
||||
Beta beta,
|
||||
TC * C, int ldC,
|
||||
cudaStream_t stream = 0)
|
||||
{
|
||||
if (transA == 'N' && transB == 'T') {
|
||||
return gemm_nt(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream);
|
||||
} else
|
||||
if (transA == 'T' && transB == 'N') {
|
||||
return gemm_tn(m, n, k, alpha, A, ldA, B, ldB, beta, C, ldC, stream);
|
||||
}
|
||||
assert(false && "Not implemented");
|
||||
}
|
||||
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
cudaDeviceProp props;
|
||||
cudaError_t error = cudaGetDeviceProperties(&props, 0);
|
||||
if (error != cudaSuccess) {
|
||||
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
if (props.major < 8) {
|
||||
std::cout << "This example requires an Ampere GPU or newer (CC >= 80)" << std::endl;
|
||||
// Return 0 so tests pass if run on unsupported architectures or CUDA Toolkits.
|
||||
return 0;
|
||||
}
|
||||
|
||||
int m = 5120;
|
||||
if (argc >= 2)
|
||||
sscanf(argv[1], "%d", &m);
|
||||
|
||||
int n = 5120;
|
||||
if (argc >= 3)
|
||||
sscanf(argv[2], "%d", &n);
|
||||
|
||||
int k = 4096;
|
||||
if (argc >= 4)
|
||||
sscanf(argv[3], "%d", &k);
|
||||
|
||||
char transA = 'N';
|
||||
if (argc >= 5)
|
||||
sscanf(argv[4], "%c", &transA);
|
||||
|
||||
char transB = 'T';
|
||||
if (argc >= 6)
|
||||
sscanf(argv[5], "%c", &transB);
|
||||
|
||||
using TA = float;
|
||||
using TB = float;
|
||||
using TC = float;
|
||||
using TI = float;
|
||||
|
||||
TI alpha = 1.0;
|
||||
TI beta = 0.0;
|
||||
|
||||
std::cout << "M = " << m << std::endl;
|
||||
std::cout << "N = " << n << std::endl;
|
||||
std::cout << "K = " << k << std::endl;
|
||||
std::cout << "C = A^" << transA << " B^" << transB << std::endl;
|
||||
|
||||
thrust::host_vector<TA> h_A(m*k);
|
||||
thrust::host_vector<TB> h_B(n*k);
|
||||
thrust::host_vector<TC> h_C(m*n);
|
||||
|
||||
for (int j = 0; j < m*k; ++j) h_A[j] = static_cast<TA>( 2*(rand() / double(RAND_MAX)) - 1 );
|
||||
for (int j = 0; j < n*k; ++j) h_B[j] = static_cast<TB>( 2*(rand() / double(RAND_MAX)) - 1 );
|
||||
for (int j = 0; j < m*n; ++j) h_C[j] = static_cast<TC>(-1);
|
||||
|
||||
thrust::device_vector<TA> d_A = h_A;
|
||||
thrust::device_vector<TB> d_B = h_B;
|
||||
thrust::device_vector<TC> d_C = h_C;
|
||||
|
||||
double gflops = (2.0*m*n*k) * 1e-9;
|
||||
|
||||
const int timing_iterations = 100;
|
||||
GPU_Clock timer;
|
||||
|
||||
int ldA = 0, ldB = 0, ldC = m;
|
||||
|
||||
if (transA == 'N') {
|
||||
ldA = m;
|
||||
} else if (transA == 'T') {
|
||||
ldA = k;
|
||||
} else {
|
||||
assert(false);
|
||||
}
|
||||
|
||||
if (transB == 'N') {
|
||||
ldB = k;
|
||||
} else if (transB == 'T') {
|
||||
ldB = n;
|
||||
} else {
|
||||
assert(false);
|
||||
}
|
||||
|
||||
// Run once
|
||||
d_C = h_C;
|
||||
gemm(transA, transB, m, n, k,
|
||||
alpha,
|
||||
d_A.data().get(), ldA,
|
||||
d_B.data().get(), ldB,
|
||||
beta,
|
||||
d_C.data().get(), ldC);
|
||||
CUTE_CHECK_LAST();
|
||||
thrust::host_vector<TC> cute_result = d_C;
|
||||
|
||||
// Timing iterations
|
||||
timer.start();
|
||||
for (int i = 0; i < timing_iterations; ++i) {
|
||||
gemm(transA, transB, m, n, k,
|
||||
alpha,
|
||||
d_A.data().get(), ldA,
|
||||
d_B.data().get(), ldB,
|
||||
beta,
|
||||
d_C.data().get(), ldC);
|
||||
}
|
||||
double cute_time = timer.seconds() / timing_iterations;
|
||||
CUTE_CHECK_LAST();
|
||||
printf("CUTE_GEMM: [%6.1f]GFlop/s (%6.4f)ms\n", gflops / cute_time, cute_time*1000);
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -67,7 +67,7 @@
|
||||
//
|
||||
// Uses local_partition() to partition a tile among threads arranged as (THR_M, THR_N).
|
||||
template <class TensorS, class TensorD, class ThreadLayout>
|
||||
__global__ void copy_kernel(TensorS S, TensorD D, ThreadLayout)
|
||||
__global__ void copy_kernel(TensorS S, TensorD D, ThreadLayout)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
@@ -77,12 +77,13 @@ __global__ void copy_kernel(TensorS S, TensorD D, ThreadLayout)
|
||||
|
||||
// Construct a partitioning of the tile among threads with the given thread arrangement.
|
||||
|
||||
// Concept: Tensor Layout Index
|
||||
Tensor thr_tile_S = local_partition(tile_S, ThreadLayout{}, threadIdx.x);
|
||||
Tensor thr_tile_D = local_partition(tile_D, ThreadLayout{}, threadIdx.x);
|
||||
// Concept: Tensor ThrLayout ThrIndex
|
||||
Tensor thr_tile_S = local_partition(tile_S, ThreadLayout{}, threadIdx.x); // (ThrValM, ThrValN)
|
||||
Tensor thr_tile_D = local_partition(tile_D, ThreadLayout{}, threadIdx.x); // (ThrValM, ThrValN)
|
||||
|
||||
// Construct a register-backed Tensor with the same shape as each thread's partition
|
||||
auto fragment = make_fragment_like(thr_tile_S);
|
||||
// Use make_tensor to try to match the layout of thr_tile_S
|
||||
Tensor fragment = make_tensor_like(thr_tile_S); // (ThrValM, ThrValN)
|
||||
|
||||
// Copy from GMEM to RMEM and from RMEM to GMEM
|
||||
copy(thr_tile_S, fragment);
|
||||
@@ -95,17 +96,17 @@ __global__ void copy_kernel(TensorS S, TensorD D, ThreadLayout)
|
||||
/// has the precondition that pointers are aligned to the vector size.
|
||||
///
|
||||
template <class TensorS, class TensorD, class ThreadLayout, class VecLayout>
|
||||
__global__ void copy_kernel_vectorized(TensorS S, TensorD D, ThreadLayout, VecLayout)
|
||||
__global__ void copy_kernel_vectorized(TensorS S, TensorD D, ThreadLayout, VecLayout)
|
||||
{
|
||||
using namespace cute;
|
||||
using Element = typename TensorS::value_type;
|
||||
|
||||
// Slice the tensors to obtain a view into each tile.
|
||||
Tensor tile_S = S(make_coord(_, _), blockIdx.x, blockIdx.y); // (BlockShape_M, BlockShape_N)
|
||||
Tensor tile_D = D(make_coord(_, _), blockIdx.x, blockIdx.y); // (BlockShape_M, BlockShape_N)
|
||||
Tensor tile_S = S(make_coord(_, _), blockIdx.x, blockIdx.y); // (BlockShape_M, BlockShape_N)
|
||||
Tensor tile_D = D(make_coord(_, _), blockIdx.x, blockIdx.y); // (BlockShape_M, BlockShape_N)
|
||||
|
||||
// Define `AccessType` which controls the size of the actual memory access.
|
||||
using AccessType = cutlass::AlignedArray<Element, size(shape(VecLayout{}))>;
|
||||
using AccessType = cutlass::AlignedArray<Element, size(VecLayout{})>;
|
||||
|
||||
// A copy atom corresponds to one hardware memory access.
|
||||
using Atom = Copy_Atom<UniversalCopy<AccessType>, Element>;
|
||||
@@ -125,29 +126,18 @@ __global__ void copy_kernel_vectorized(TensorS S, TensorD D, ThreadLayout, VecLa
|
||||
// Construct a Tensor corresponding to each thread's slice.
|
||||
auto thr_copy = tiled_copy.get_thread_slice(threadIdx.x);
|
||||
|
||||
Tensor thr_tile_S = thr_copy.partition_S(tile_S);
|
||||
Tensor thr_tile_D = thr_copy.partition_D(tile_D);
|
||||
Tensor thr_tile_S = thr_copy.partition_S(tile_S); // (CopyOp, CopyM, CopyN)
|
||||
Tensor thr_tile_D = thr_copy.partition_D(tile_D); // (CopyOp, CopyM, CopyN)
|
||||
|
||||
// Construct a register-backed Tensor with the same shape as each thread's partition
|
||||
auto fragment = make_fragment_like(thr_tile_D);
|
||||
// Use make_fragment because the first mode is the instruction-local mode
|
||||
Tensor fragment = make_fragment_like(thr_tile_D); // (CopyOp, CopyM, CopyN)
|
||||
|
||||
// Copy from GMEM to RMEM and from RMEM to GMEM
|
||||
copy(tiled_copy, thr_tile_S, fragment);
|
||||
copy(tiled_copy, fragment, thr_tile_D);
|
||||
}
|
||||
|
||||
/// Helper to convert a shape to a dim3
|
||||
template <class Shape>
|
||||
dim3 shape_to_dim3(Shape shape)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
CUTE_STATIC_ASSERT_V(rank(shape) <= Int<3>{});
|
||||
auto result = append<3>(product_each(shape), 1u);
|
||||
|
||||
return dim3(get<0>(result), get<1>(result), get<2>(result));
|
||||
}
|
||||
|
||||
/// Main function
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
@@ -161,13 +151,13 @@ int main(int argc, char** argv)
|
||||
// Define a tensor shape with dynamic extents (m, n)
|
||||
auto tensor_shape = make_shape(256, 512);
|
||||
|
||||
//
|
||||
// Allocate and initialize
|
||||
//
|
||||
|
||||
thrust::host_vector<Element> h_S(size(tensor_shape));
|
||||
thrust::host_vector<Element> h_D(size(tensor_shape));
|
||||
|
||||
//
|
||||
// Initialize
|
||||
//
|
||||
|
||||
for (size_t i = 0; i < h_S.size(); ++i) {
|
||||
h_S[i] = static_cast<Element>(i);
|
||||
h_D[i] = Element{};
|
||||
@@ -180,33 +170,36 @@ int main(int argc, char** argv)
|
||||
// Make tensors
|
||||
//
|
||||
|
||||
Tensor tensor_S = make_tensor(make_gmem_ptr(d_S.data().get()), make_layout(tensor_shape));
|
||||
Tensor tensor_D = make_tensor(make_gmem_ptr(d_D.data().get()), make_layout(tensor_shape));
|
||||
Tensor tensor_S = make_tensor(make_gmem_ptr(thrust::raw_pointer_cast(d_S.data())), make_layout(tensor_shape));
|
||||
Tensor tensor_D = make_tensor(make_gmem_ptr(thrust::raw_pointer_cast(d_D.data())), make_layout(tensor_shape));
|
||||
|
||||
//
|
||||
// Partition
|
||||
// Tile tensors
|
||||
//
|
||||
|
||||
|
||||
// Define a statically sized block (M, N).
|
||||
//
|
||||
// Note, by convention, capital letters are used to represent static modes.
|
||||
auto block_shape = make_shape(Int<128>{}, Int<64>{});
|
||||
|
||||
if ((get<0>(tensor_shape) % get<0>(block_shape)) || (get<1>(tensor_shape) % get<1>(block_shape))) {
|
||||
if ((size<0>(tensor_shape) % size<0>(block_shape)) || (size<1>(tensor_shape) % size<1>(block_shape))) {
|
||||
std::cerr << "The tensor shape must be divisible by the block shape." << std::endl;
|
||||
return -1;
|
||||
}
|
||||
// Equivalent check to the above
|
||||
if (not weakly_compatible(block_shape, tensor_shape)) {
|
||||
std::cerr << "Expected the tensors to be weakly compatible with the block_shape." << std::endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Tile the tensor (m, m) ==> ((M, N), m', n') where (M, N) is the static tile
|
||||
// Tile the tensor (m, n) ==> ((M, N), m', n') where (M, N) is the static tile
|
||||
// shape, and modes (m', n') correspond to the number of tiles.
|
||||
//
|
||||
// These will be used to determine the CUDA kernel grid dimensinos.
|
||||
Tensor tiled_tensor_S = tiled_divide(tensor_S, block_shape);
|
||||
Tensor tiled_tensor_D = tiled_divide(tensor_D, block_shape);
|
||||
//
|
||||
// These will be used to determine the CUDA kernel grid dimensions.
|
||||
Tensor tiled_tensor_S = tiled_divide(tensor_S, block_shape); // ((M, N), m', n')
|
||||
Tensor tiled_tensor_D = tiled_divide(tensor_D, block_shape); // ((M, N), m', n')
|
||||
|
||||
// Thread arrangement
|
||||
Layout thr_layout = make_layout(make_shape(Int<32>{}, Int< 8>{}));
|
||||
Layout thr_layout = make_layout(make_shape(Int<32>{}, Int<8>{}));
|
||||
|
||||
// Vector dimensions
|
||||
Layout vec_layout = make_layout(make_shape(Int<4>{}, Int<1>{}));
|
||||
@@ -215,16 +208,16 @@ int main(int argc, char** argv)
|
||||
// Determine grid and block dimensions
|
||||
//
|
||||
|
||||
dim3 gridDim = shape_to_dim3(select<1,2>(shape(tiled_tensor_D))); // Grid shape corresponds to modes m' and n'
|
||||
dim3 blockDim(size(shape(thr_layout)));
|
||||
dim3 gridDim (size<1>(tiled_tensor_D), size<2>(tiled_tensor_D)); // Grid shape corresponds to modes m' and n'
|
||||
dim3 blockDim(size(thr_layout));
|
||||
|
||||
//
|
||||
// Launch the kernel
|
||||
//
|
||||
copy_kernel_vectorized<<< gridDim, blockDim >>>(
|
||||
tiled_tensor_S,
|
||||
tiled_tensor_D,
|
||||
thr_layout,
|
||||
tiled_tensor_S,
|
||||
tiled_tensor_D,
|
||||
thr_layout,
|
||||
vec_layout);
|
||||
|
||||
cudaError result = cudaDeviceSynchronize();
|
||||
|
||||
Reference in New Issue
Block a user