leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 29 Packages Projects Releases Wiki Activity

Support parallel split K mode for porfiling (#277)

Browse Source 
* Support parallel split K mode for porfiling

Signed-off-by: Peter Han <fujun.han@iluvatar.ai>

* Parallel Split K support

  1. find gemm kernel by preference key
  2. switch m n for redution kernel

Signed-off-by: Peter Han <fujun.han@iluvatar.ai>

* parallel splitk for fp16 gemm

* add one missing file

Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Fujun Han
2022-01-27 10:37:37 -05:00
committed by GitHub
co-authored by Haicheng Wu
parent c3353add63
commit 1e4703cbab
13 changed files with 332 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/library/src/conv2d_operation.h
+3 -2
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
@@ -272,7 +272,8 @@ public:
/// Gets the device-side workspace
virtual uint64_t get_device_workspace_size(
void const *configuration_ptr) const {
void const *configuration_ptr,
void const *arguments_ptr = nullptr) const {
OperatorArguments args;
tools/library/src/conv3d_operation.h
+3 -2
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
@@ -266,7 +266,8 @@ public:
/// Gets the device-side workspace
virtual uint64_t get_device_workspace_size(
void const *configuration_ptr) const {
void const *configuration_ptr,
void const *arguments_ptr = nullptr) const {
OperatorArguments args;
tools/library/src/gemm_operation.h
+20 -7
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
@@ -242,7 +242,8 @@ public:
/// Gets the device-side workspace
virtual uint64_t get_device_workspace_size(
void const *configuration_ptr) const {
void const *configuration_ptr,
void const *arguments_ptr = nullptr) const {
OperatorArguments args;
@@ -443,7 +444,8 @@ public:
/// Gets the device-side workspace
virtual uint64_t get_device_workspace_size(
void const *configuration_ptr) const {
void const *configuration_ptr,
void const *arguments_ptr = nullptr) const {
OperatorArguments args;
@@ -569,7 +571,7 @@ protected:
operator_args.ldb = (configuration->ldb);
operator_args.ldc = (configuration->ldc);
operator_args.ldd = (configuration->ldd);
return Status::kSuccess;
}
@@ -649,7 +651,8 @@ public:
/// Gets the device-side workspace
virtual uint64_t get_device_workspace_size(
void const *configuration_ptr) const {
void const *configuration_ptr,
void const *arguments_ptr) const {
OperatorArguments args;
@@ -661,6 +664,14 @@ public:
return 0;
}
status = update_arguments_(
args,
static_cast<GemmUniversalArguments const *>(arguments_ptr));
if (status != Status::kSuccess) {
return 0;
}
uint64_t size = Operator::get_workspace_size(args);
return size;
@@ -855,7 +866,8 @@ public:
/// Gets the device-side workspace
virtual uint64_t get_device_workspace_size(
void const *configuration_ptr) const {
void const *configuration_ptr,
void const *arguments_ptr = nullptr) const {
OperatorArguments args;
@@ -1055,7 +1067,8 @@ public:
/// Gets the device-side workspace
virtual uint64_t get_device_workspace_size(
void const *configuration_ptr) const {
void const *configuration_ptr,
void const *arguments_ptr = nullptr) const {
OperatorArguments args;
tools/library/src/handle.cu
+54 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
@@ -1098,6 +1098,59 @@ Operation const* find_conv_operation_for_parallel_reduction(Operation const *ope
return nullptr;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Finds gemm operation instances with Gemm::ElementC = Reduction::ElementWorkspace
Operation const* find_gemm_operation_for_parallel_reduction(Operation const *operation) {
GemmDescription const &gemm_desc =
static_cast<GemmDescription const &>(operation->description());
// if the curren gemm operation accumulator and output data type match return operation
if(gemm_desc.tile_description.math_instruction.element_accumulator == gemm_desc.C.element) {
return operation;
}
// find gemm operation to match gemm output and reduction workspace data type
GemmFunctionalKey key(
library::Provider::kCUTLASS,
gemm_desc.gemm_kind,
gemm_desc.tile_description.math_instruction.element_accumulator,
gemm_desc.element_epilogue,
gemm_desc.A.element,
gemm_desc.A.layout,
gemm_desc.transform_A,
gemm_desc.B.element,
gemm_desc.B.layout,
gemm_desc.transform_B,
gemm_desc.tile_description.math_instruction.element_accumulator);
// gemm operation table
auto gemm_operations = Singleton::get().operation_table.gemm_operations;
// find ConvFunctionalKey in gemm operation table
auto operators_it = gemm_operations.find(key);
if (operators_it == gemm_operations.end()) {
return nullptr;
}
if (operators_it->second.empty()) {
return nullptr;
}
// A and B uses the same alignment in the generator.py
int alignment = gemm_desc.A.alignment;
// gemm operation for same compute capability and iterator algorithm
GemmPreferenceKey preference_key(
gemm_desc.tile_description.minimum_compute_capability,
alignment);
return find_gemm_operation(operators_it, preference_key);
}
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace library
tools/library/src/reduction/init_reduction_operations.cu
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
tools/library/src/reduction/reduction_device.cu
+4 -5
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
@@ -46,7 +46,7 @@ void initialize_reduce_add_linear_combination_f32_f32_f16(Manifest &manifest) {
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombination<
ElementOutput,
128 / cutlass::sizeof_bits<ElementOutput>::value,
128 / cutlass::sizeof_bits<ElementWorkspace>::value,
ElementAccumulator,
ElementCompute
>;
@@ -81,7 +81,7 @@ void initialize_reduce_add_linear_combination_f32_f32_f32(Manifest &manifest) {
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombination<
ElementOutput,
128 / cutlass::sizeof_bits<ElementOutput>::value,
128 / cutlass::sizeof_bits<ElementWorkspace>::value,
ElementAccumulator,
ElementCompute
>;
@@ -115,7 +115,7 @@ void initialize_reduce_add_linear_combination_cf32_cf32_cf32(Manifest &manifest)
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombination<
ElementOutput,
128 / cutlass::sizeof_bits<ElementOutput>::value,
128 / cutlass::sizeof_bits<ElementWorkspace>::value,
ElementAccumulator,
ElementCompute
>;
@@ -140,6 +140,5 @@ void initialize_reduce_add_linear_combination_cf32_cf32_cf32(Manifest &manifest)
));
}
}
}
tools/library/src/reduction/reduction_operation.h
+3 -1
View File
@@ -30,6 +30,7 @@
#include <iostream>
#include "cutlass/cutlass.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/epilogue/thread/linear_combination_clamp.h"
#include "cutlass/reduction/thread/reduction_operators.h"
#include "cutlass/reduction/device/reduce_split_k.h"
@@ -180,7 +181,8 @@ public:
/// Gets the device-side workspace
virtual uint64_t get_device_workspace_size(
void const *configuration_ptr) const {
void const *configuration_ptr,
void const *arguments_ptr = nullptr) const {
OperatorArguments args;
tools/library/src/reference/conv_reference_operation.h
+3 -2
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
@@ -403,7 +403,8 @@ public:
}
virtual uint64_t get_device_workspace_size(
void const *configuration) const {
void const *configuration,
void const *arguments = nullptr) const {
return 0;
}
tools/library/src/reference/gemm_reference_operation.h
+3 -2
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted
* provided that the following conditions are met:
@@ -161,7 +161,8 @@ public:
}
virtual uint64_t get_device_workspace_size(
void const *configuration) const {
void const *configuration,
void const *arguments = nullptr) const {
return 0;
}