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

3.6.0 update (#2005)

Browse Source 
* 3.6.0 update

* doc and swap stuff

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
2024-12-25 01:34:40 -05:00
committed by GitHub
co-authored by yuzhai Haicheng Wu
parent e1cd8c7866
commit 3d261a5974
258 changed files with 10863 additions and 3883 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/library/src/conv2d_operation.h
+12 -2
View File
@@ -326,7 +326,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
@@ -578,7 +583,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
tools/library/src/conv3d_operation.h
+6 -1
View File
@@ -317,7 +317,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
tools/library/src/conv_operation_3x.hpp
+136 -22
View File
@@ -236,12 +236,14 @@ public:
typename Operator::Arguments out_args{};
status = update_operator_arguments_from_configuration_2d_or_3d(out_args, configuration);
if (status != Status::kSuccess) {
CUTLASS_TRACE_HOST("*** can_implement: update_operator_arguments_from_configuration_2d_or_3d failed");
return status;
}
auto* in_args_ptr = reinterpret_cast<ConvArguments const*>(arguments);
status = update_operator_arguments_from_arguments(out_args, *in_args_ptr);
if (status != Status::kSuccess) {
CUTLASS_TRACE_HOST("*** can_implement: update_operator_arguments_from_arguments failed");
return status;
}
@@ -332,7 +334,8 @@ public:
void const* arguments,
void* host_workspace,
void* device_workspace = nullptr,
cudaStream_t stream = nullptr) const override
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const override
{
auto status = Status::kInvalid;
@@ -358,7 +361,7 @@ public:
}
auto* op = reinterpret_cast<Operator*>(host_workspace);
return op->run(out_args, device_workspace, stream);
return op->run(out_args, device_workspace, stream, nullptr, launch_with_pdl);
}
private:
@@ -482,6 +485,11 @@ private:
typename Operator::Arguments& out_args,
Conv2dConfiguration const& config)
{
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
CUTLASS_TRACE_HOST("ConvOperator3x::"
"update_operator_arguments_from_configuration"
"(Conv2dConfiguration)\n");
#endif
using detail::vector_to_array_strides;
constexpr int num_spatial_dims = Operator::NumSpatialDimensions;
@@ -595,6 +603,7 @@ private:
const TensorStride stride_A = vector_to_array_strides(config.stride_a, the_stride_size);
const TensorStride stride_B = vector_to_array_strides(config.stride_b, the_stride_size);
const TensorStride stride_C = vector_to_array_strides(config.stride_c, the_stride_size);
// cutlass::library::Conv2dConfiguration has no member stride_d.
// The code below imitates the testbed,
@@ -605,12 +614,57 @@ private:
CUTLASS_TRACE_HOST("CUTLASS 3 kernels currently only support groups = 1.");
return Status::kInvalid;
}
// ConvProblemShape is how CUTLASS 3 kernels represent
// convolution problems. ConvProblemShape's constructors take
// shape_act, stride_act, shape_flt, and stride_flt, and set
// shape_A, stride_A, shape_B, stride_B, shape_C, and stride_C
// according to Fprop / Dgrad / Wgrad.
//
// This means that stride_act isn't always config.stride_A,
// depending on Fprop / Dgrad / Wgrad. The code here "undoes"
// the logic in Conv2dWorkspace::set_stride_vector so that we
// can recover the strides of the activation and filter tensors.
// It doesn't need to worry about the so-called "output" tensor
// (which might not be C), as ConvProblemShape's constructor
// figures out its shapes and strides.
using TensorExtent = typename problem_shape_type::TensorExtent;
TensorExtent shape_act{N, H, W, C};
auto stride_act = [&] () {
// Some compilers consider conv_op (defined above), as
// captured by this lambda, as "not a constant expression."
constexpr auto conv_kind = Operator::DispatchPolicy::ConvOp;
if constexpr (conv_kind == cutlass::conv::Operator::kFprop) {
return stride_A;
}
else if constexpr (conv_kind == cutlass::conv::Operator::kDgrad) {
return stride_C;
}
else { // conv_kind == cutlass::conv::Operator::kWgrad
return stride_B;
}
} ();
TensorExtent shape_flt{K, R, S, C};
auto stride_flt = [&] () {
// Some compilers consider conv_op (defined above), as
// captured by this lambda, as "not a constant expression."
constexpr auto conv_kind = Operator::DispatchPolicy::ConvOp;
if constexpr (conv_kind == cutlass::conv::Operator::kFprop) {
return stride_B;
}
else if constexpr (conv_kind == cutlass::conv::Operator::kDgrad) {
return stride_B;
}
else { // conv_kind == cutlass::conv::Operator::kWgrad
return stride_C;
}
} ();
problem_shape_type problem_shape(
/* mode = */ mode,
/* shape_act = */ {N, H, W, C},
/* stride_act = */ stride_A,
/* shape_flt = */ {K, R, S, C},
/* stride_flt = */ stride_B,
/* shape_act = */ shape_act,
/* stride_act = */ stride_act,
/* shape_flt = */ shape_flt,
/* stride_flt = */ stride_flt,
/* lower_padding = */ {pad_h, pad_w},
/* upper_padding = */ {pad_h, pad_w},
/* traversal_stride = */ {traversal_stride_h, traversal_stride_w},
@@ -620,9 +674,11 @@ private:
// ConvProblemShape's constructor sets its shape_C member.
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
std::cerr << " problem_shape:\n"
<< " shape_C: " << problem_shape.shape_C << "\n";
std::cerr << " stride_C: " << problem_shape.stride_C << "\n";
printf("\n problem_shape.shape_C: ");
print(problem_shape.shape_C);
printf("\n problem_shape.stride_C: ");
print(problem_shape.stride_C);
printf("\n");
#endif
// Initialization of C's and D's strides follows the CUTLASS 3
// convolutions testbed (test/unit/conv/device_3x/testbed_conv.hpp).
@@ -670,6 +726,11 @@ private:
typename Operator::Arguments& out_args,
Conv3dConfiguration const& config)
{
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
CUTLASS_TRACE_HOST("ConvOperator3x::"
"update_operator_arguments_from_configuration"
"(Conv3dConfiguration)\n");
#endif
using detail::coord_to_array_strides;
constexpr int num_spatial_dims = Operator::NumSpatialDimensions;
@@ -762,6 +823,10 @@ private:
print_stride(input_stride_b, "input_stride_b");
print_stride(input_stride_c, "input_stride_c");
#endif
// Conv3dConfiguration stores the strides as Coord (with
// compile-time size), so there's no need to check sizes here
// (unlike Conv2dConfiguration, which stores strides as
// std::vector).
constexpr cutlass::conv::Operator conv_op = Operator::DispatchPolicy::ConvOp;
using problem_shape_type =
@@ -771,18 +836,68 @@ private:
const TensorStride stride_A = coord_to_array_strides(input_stride_a);
const TensorStride stride_B = coord_to_array_strides(input_stride_b);
const TensorStride stride_C = coord_to_array_strides(input_stride_c);
const int num_groups = config.problem_size.groups;
if (num_groups != 1) {
CUTLASS_TRACE_HOST("CUTLASS 3 kernels currently only support groups = 1.");
return Status::kInvalid;
}
// ConvProblemShape is how CUTLASS 3 kernels represent
// convolution problems. ConvProblemShape's constructors take
// shape_act, stride_act, shape_flt, and stride_flt, and set
// shape_A, stride_A, shape_B, stride_B, shape_C, and stride_C
// according to Fprop / Dgrad / Wgrad.
//
// Conv3dConfiguration differs a bit from Conv2dConfiguration,
// but the idea is the same: the "input_stride_a" from config
// depends on conv_kind (Fprop, Dgrad, or Wgrad), so stride_act
// isn't always input_stride_a. Analogously, stride_flt isn't
// always input_stride_b. The code here "undoes" the logic in
// config.layout_a(conv_kind) and config.layout_b(conv_kind)
// (analogous to Conv2dWorkspace::set_stride_vector) so that we
// can recover the strides of the activation and filter tensors.
// It doesn't need to worry about the so-called "output" tensor
// (which might not be C), as ConvProblemShape's constructor
// figures out its shapes and strides.
using TensorExtent = typename problem_shape_type::TensorExtent;
TensorExtent shape_act{N, D, H, W, C};
auto stride_act = [&] () {
// Some compilers consider conv_op (defined above), as
// captured by this lambda, as "not a constant expression."
constexpr auto conv_kind = Operator::DispatchPolicy::ConvOp;
if constexpr (conv_kind == cutlass::conv::Operator::kFprop) {
return stride_A;
}
else if constexpr (conv_kind == cutlass::conv::Operator::kDgrad) {
return stride_C;
}
else { // conv_kind == cutlass::conv::Operator::kWgrad
return stride_B;
}
} ();
TensorExtent shape_flt{K, T, R, S, C};
auto stride_flt = [&] () {
// Some compilers consider conv_op (defined above), as
// captured by this lambda, as "not a constant expression."
constexpr auto conv_kind = Operator::DispatchPolicy::ConvOp;
if constexpr (conv_kind == cutlass::conv::Operator::kFprop) {
return stride_B;
}
else if constexpr (conv_kind == cutlass::conv::Operator::kDgrad) {
return stride_B;
}
else { // conv_kind == cutlass::conv::Operator::kWgrad
return stride_C;
}
} ();
problem_shape_type problem_shape(
/* mode = */ mode,
/* shape_act = */ {N, D, H, W, C},
/* stride_act = */ stride_A,
/* shape_flt = */ {K, T, R, S, C},
/* stride_flt = */ stride_B,
/* shape_act = */ shape_act,
/* stride_act = */ stride_act,
/* shape_flt = */ shape_flt,
/* stride_flt = */ stride_flt,
/* lower_padding = */ {pad_d, pad_h, pad_w},
/* upper_padding = */ {pad_d, pad_h, pad_w},
/* traversal_stride = */ {traversal_stride_d, traversal_stride_h, traversal_stride_w},
@@ -792,15 +907,15 @@ private:
// ConvProblemShape's constructor sets its shape_C member.
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
std::cerr << " problem_shape:\n"
<< " shape_C: " << problem_shape.shape_C << "\n";
std::cerr << " stride_C: " << problem_shape.stride_C << "\n";
printf("\n problem_shape.shape_C: ");
print(problem_shape.shape_C);
printf("\n problem_shape.stride_C: ");
print(problem_shape.stride_C);
printf("\n");
#endif
// Initialization of C's and D's strides follows the CUTLASS 3
// convolutions testbed (test/unit/conv/device_3x/testbed_conv.hpp).
{
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
std::cerr << " Compute stride_C and stride_D\n";
#endif
using StrideC = typename Operator::ConvKernel::StrideC;
using StrideD = typename Operator::ConvKernel::StrideD;
auto stride_C = StrideC{};
@@ -845,9 +960,8 @@ private:
ConvArguments const& in_args) const
{
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
std::cerr << "ConvOperation3x::update_operator_arguments_from_arguments\n";
CUTLASS_TRACE_HOST("ConvOperation3x::update_operator_arguments_from_arguments\n");
#endif
auto status = UpdateFusionArgs<decltype(out_args.epilogue.thread)>::update_(
out_args.epilogue.thread, in_args);
if (status != Status::kSuccess) {
tools/library/src/gemm_operation.h
+36 -6
View File
@@ -296,7 +296,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
@@ -500,7 +505,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
@@ -721,7 +731,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
@@ -930,7 +945,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
@@ -1133,7 +1153,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
@@ -1337,7 +1362,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
tools/library/src/gemm_operation_3x.hpp
+4 -2
View File
@@ -35,6 +35,7 @@
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/detail/collective.hpp"
#include "cutlass/library/library.h"
#include "library_internal.h"
#include "cutlass/gemm/dispatch_policy.hpp"
@@ -331,7 +332,8 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const override {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const override {
OperatorArguments args;
Status status = update_arguments_(args, static_cast<GemmUniversalArguments const *>(arguments_ptr));
@@ -341,7 +343,7 @@ public:
Operator *op = static_cast<Operator *>(host_workspace);
// We need to call initialize() since we have to rebuild TMA desc for every new set of args
status = op->run(args, device_workspace, stream);
status = op->run(args, device_workspace, stream, nullptr, launch_with_pdl);
return status;
}
};
tools/library/src/handle.cu
+25 -6
View File
@@ -57,14 +57,12 @@ Handle::Handle(
scalar_pointer_mode_(ScalarPointerMode::kHost),
last_operation_(nullptr) {
int device_idx = -1;
cudaError_t error = cudaGetDevice(&device_idx);
cudaError_t error = cudaGetDevice(&device_idx_);
if (error != cudaSuccess) {
throw std::runtime_error("cudaGetDevice() failed");
}
error = cudaGetDeviceProperties(&device_, device_idx);
error = cudaGetDeviceProperties(&device_, device_idx_);
if (error != cudaSuccess) {
throw std::runtime_error("cudaGetDeviceProperties() failed");
}
@@ -78,8 +76,14 @@ Handle::Handle(
Handle::~Handle() {
if (workspace_) {
if (workspace_) {
cudaFree(workspace_);
int device_before;
cudaGetDevice(&device_before);
if (device_before != device_idx_) {
cudaSetDevice(device_idx_);
}
cudaFree(workspace_);
if (device_before != device_idx_) {
cudaSetDevice(device_before);
}
workspace_ = nullptr;
@@ -89,6 +93,10 @@ Handle::~Handle() {
/// Move constructor
Handle::Handle(Handle && handle) {
cudaError_t error = cudaGetDevice(&device_idx_);
if (error != cudaSuccess) {
throw std::runtime_error("cudaGetDevice() failed");
}
device_ = handle.device_;
workspace_size_ = handle.workspace_size_;
workspace_ = handle.workspace_;
@@ -112,6 +120,8 @@ Handle & Handle::operator=(Handle && handle) {
handle.workspace_ = nullptr;
handle.workspace_size_ = 0;
device_idx_ = handle.device_idx_;
return *this;
}
@@ -151,6 +161,12 @@ void *Handle::get_workspace() const {
/// Sets the size of device workspace, invalidating previous calls to get_device_workspace()
void Handle::set_workspace_size(size_t bytes) {
int device_before;
cudaGetDevice(&device_before);
if (device_before != device_idx_) {
cudaSetDevice(device_idx_);
}
if (bytes != workspace_size_) {
if (workspace_) {
@@ -177,6 +193,9 @@ void Handle::set_workspace_size(size_t bytes) {
throw std::runtime_error("Failed to clear workspace");
}
}
if (device_before != device_idx_) {
cudaSetDevice(device_before);
}
}
/// Gets the scalar pointer mode
tools/library/src/library_internal.h
+8
View File
@@ -72,6 +72,10 @@ template <> struct NumericTypeMap<cutlass::uint1b_t> {
static NumericTypeID const kId = NumericTypeID::kB1;
};
template <> struct NumericTypeMap<cutlass::int2b_t> {
static NumericTypeID const kId = NumericTypeID::kS2;
};
template <> struct NumericTypeMap<cutlass::int4b_t> {
static NumericTypeID const kId = NumericTypeID::kS4;
};
@@ -92,6 +96,10 @@ template <> struct NumericTypeMap<int64_t> {
static NumericTypeID const kId = NumericTypeID::kS64;
};
template <> struct NumericTypeMap<cutlass::uint2b_t> {
static NumericTypeID const kId = NumericTypeID::kU2;
};
template <> struct NumericTypeMap<cutlass::uint4b_t> {
static NumericTypeID const kId = NumericTypeID::kU4;
};
tools/library/src/rank_2k_operation.h
+6 -1
View File
@@ -314,7 +314,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
tools/library/src/rank_k_operation.h
+6 -1
View File
@@ -310,7 +310,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
tools/library/src/reduction/reduction_operation.h
+6 -1
View File
@@ -231,7 +231,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
tools/library/src/reference/conv_reference_operation.h
+6 -1
View File
@@ -432,7 +432,12 @@ public:
void const *arguments,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
ConvArguments const &args = *static_cast<ConvArguments const *>(arguments);
tools/library/src/reference/gemm_reference_operation.h
+6 -1
View File
@@ -192,7 +192,12 @@ public:
void const *arguments,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
GemmUniversalConfiguration const &config = *static_cast<GemmUniversalConfiguration const *>(host_workspace);
GemmUniversalArguments const &args = *static_cast<GemmUniversalArguments const *>(arguments);
tools/library/src/reference/gemm_s8_s8_s32.cu
+1 -1
View File
@@ -77,7 +77,7 @@ void initialize_gemm_reference_operations_s8_s8_s32(Manifest &manifest) {
int8_t, // ElementA
int8_t, // ElementB
int32_t, // ElementC
int32_t, // ElementScalar / ElementCompute
float, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
int32_t // ElementD
>(manifest);
tools/library/src/sparse_gemm_operation_3x.hpp
+26 -25
View File
@@ -35,6 +35,7 @@
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/detail/collective.hpp"
#include "cutlass/library/library.h"
#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp" // StructuredSparseCompressor
#include "cutlass/transform/device/transform_universal_adapter.hpp" // TransformUniversalAdapter
@@ -56,7 +57,7 @@ namespace cutlass::library {
///////////////////////////////////////////////////////////////////////////////////////////////////
// Limitation & Assumptions:
// Limitation & Assumptions:
// 1. The tensor must be densely packed. That is, lda is k if the tensor is k-major,
// and lda is m if the tensor is m-major.
// 2. Circular buffer for tensorA and tensorE may have a less count compared to tensorB and others.
@@ -169,7 +170,6 @@ protected:
return status;
}
// TODO: type erase Arguments structure in 3.0 GEMM
operator_args.problem_shape = cute::make_shape(
arguments->problem_size.m(),
arguments->problem_size.n(),
@@ -302,13 +302,15 @@ public:
}
Status initialize_with_profiler_workspace(
void const *configuration,
void *host_workspace,
void *device_workspace,
void const *configuration,
void *host_workspace,
void *device_workspace,
uint8_t **profiler_workspaces,
int problem_count_from_profiler,
cudaStream_t stream = nullptr) {
iter_idx.resize(static_cast<GemmUniversalConfiguration const*>(configuration)->device_count, 0);
// Set problem_count.
problem_count = problem_count_from_profiler;
@@ -319,13 +321,10 @@ public:
// * Construct Op
Operator *op = new (host_op_workspace_ptr) Operator;
// * Device Full Ptr
device_full_ptr = reinterpret_cast<uint8_t*>(device_workspace);
// * Device Ptr (1st iteration)
// Device workspace : | iter1 | iter2 | iter3 | .. | iterx |
// iteri : op_workspace | tensor_ac | tensor_e
auto* device_ptr_iter1 = device_full_ptr;
auto* device_ptr_iter1 = static_cast<uint8_t*>(device_workspace);
auto* device_op_workspace_ptr_iter1 = device_ptr_iter1;
auto* device_compressor_workspace_ptr_iter1 = device_op_workspace_ptr_iter1 + device_op_workspace_size;
auto* device_a_compressed_ptr_iter1 = device_compressor_workspace_ptr_iter1 + device_compress_workspace_size;
@@ -335,15 +334,15 @@ public:
auto* device_a_raw_ptr = profiler_workspaces[0];
// * Random fill 50% of TensorA w/ zero following the structured sparse requirement
cudaMemcpy(host_a_raw_ptr, device_a_raw_ptr, tensor_a_size, cudaMemcpyDeviceToHost);
CUDA_CHECK(cudaMemcpyAsync(host_a_raw_ptr, device_a_raw_ptr, tensor_a_size, cudaMemcpyDeviceToHost, stream));
compressor_utility.structure_sparse_zero_mask_fill(host_a_raw_ptr, 2000);
cudaMemcpy(device_a_raw_ptr, host_a_raw_ptr, tensor_a_size, cudaMemcpyHostToDevice);
CUDA_CHECK(cudaMemcpyAsync(device_a_raw_ptr, host_a_raw_ptr, tensor_a_size, cudaMemcpyHostToDevice, stream));
CUDA_CHECK(cudaGetLastError());
// * Compress DTensorA and get DTensorAC & DTensorE
cutlass::KernelHardwareInfo hw_info;
hw_info.device_id = 0;
CUDA_CHECK(cudaGetDevice(&hw_info.device_id));
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
typename Compressor::Arguments arguments{
{compressor_utility.M, 0, compressor_utility.K, compressor_utility.L},
@@ -372,23 +371,23 @@ public:
return status;
}
CUDA_CHECK(cudaStreamSynchronize(stream));
// * Copy Iter1's DTensorAC DTensorE to each iteration's DTensorAC DTensorE
for (int iter_i = 1; iter_i < problem_count; iter_i++) {
// * Device AC E Ptr per iteration
// Device workspace : | iter1 | iter2 | iter3 | .. | iterx |
// iteri : op_workspace | tensor_ac | tensor_e
auto* device_ptr_iteri = device_full_ptr + device_per_iter_workspace_size * iter_i;
auto* device_ptr_iteri = static_cast<uint8_t*>(device_workspace) + device_per_iter_workspace_size * iter_i;
auto* device_op_workspace_ptr = device_ptr_iteri;
auto* device_compressor_workspace_ptr = device_op_workspace_ptr + device_op_workspace_size;
auto* device_a_compressed_ptr = device_compressor_workspace_ptr + device_compress_workspace_size;
auto* device_e_ptr = device_a_compressed_ptr + tensor_ac_size;
cudaMemcpy(device_a_compressed_ptr, device_a_compressed_ptr_iter1, tensor_ac_size, cudaMemcpyDeviceToDevice);
cudaMemcpy(device_e_ptr, device_e_ptr_iter1, tensor_e_size, cudaMemcpyDeviceToDevice);
CUDA_CHECK(cudaMemcpyAsync(device_a_compressed_ptr, device_a_compressed_ptr_iter1, tensor_ac_size, cudaMemcpyDeviceToDevice, stream));
CUDA_CHECK(cudaMemcpyAsync(device_e_ptr, device_e_ptr_iter1, tensor_e_size, cudaMemcpyDeviceToDevice, stream));
}
CUDA_CHECK(cudaStreamSynchronize(stream));
CUDA_CHECK(cudaGetLastError());
return Status::kSuccess;
@@ -398,17 +397,21 @@ public:
Status run(
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const override {
void *device_workspace,
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const override {
OperatorArguments operator_args;
auto* device_ptr_iteri = device_full_ptr + device_per_iter_workspace_size * iter_idx;
const auto device_index = static_cast<GemmUniversalArguments const *>(arguments_ptr)->device_index;
auto* device_ptr_iteri = static_cast<uint8_t*>(device_workspace) + device_per_iter_workspace_size * iter_idx[device_index];
auto* device_op_workspace_ptr = device_ptr_iteri;
auto* device_compressor_workspace_ptr = device_op_workspace_ptr + device_op_workspace_size;
auto* device_a_compressed_ptr = device_compressor_workspace_ptr + device_compress_workspace_size;
auto* device_e_ptr = device_a_compressed_ptr + tensor_ac_size;
iter_idx = (iter_idx + 1) % problem_count;
iter_idx[device_index] = (iter_idx[device_index] + 1) % problem_count;
Status status = update_arguments_(operator_args, static_cast<GemmUniversalArguments const *>(arguments_ptr), compressor_utility, device_a_compressed_ptr, device_e_ptr );
@@ -418,7 +421,7 @@ public:
Operator *op = static_cast<Operator *>(host_workspace);
// We need to call initialize() since we have to rebuild TMA desc for every new set of args
status = op->run(operator_args, device_op_workspace_ptr, stream);
status = op->run(operator_args, device_op_workspace_ptr, stream, nullptr, launch_with_pdl);
return status;
}
@@ -426,9 +429,7 @@ private:
// Variables that must change in the const functions.
mutable CompressorUtility compressor_utility;
mutable int problem_count = 1;
mutable int iter_idx = 0;
uint8_t* device_full_ptr = nullptr;
mutable std::vector<int> iter_idx;
mutable uint64_t tensor_ac_size = 0;
mutable uint64_t tensor_e_size = 0;
tools/library/src/symm_operation.h
+6 -1
View File
@@ -312,7 +312,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;
tools/library/src/trmm_operation.h
+6 -1
View File
@@ -304,7 +304,12 @@ public:
void const *arguments_ptr,
void *host_workspace,
void *device_workspace = nullptr,
cudaStream_t stream = nullptr) const {
cudaStream_t stream = nullptr,
bool launch_with_pdl = false) const {
if (launch_with_pdl) {
return Status::kErrorNotSupported;
}
OperatorArguments args;