@@ -278,6 +278,7 @@ execute_process(
|
||||
--architectures "${CUTLASS_NVCC_ARCHS_ENABLED}"
|
||||
--kernels "${CUTLASS_LIBRARY_KERNELS}"
|
||||
--ignore-kernels "${CUTLASS_LIBRARY_IGNORE_KERNELS}"
|
||||
--exclude-kernels "${CUTLASS_LIBRARY_EXCLUDE_KERNELS}"
|
||||
--kernel-filter-file "${KERNEL_FILTER_FILE}"
|
||||
--selected-kernel-list "${CUTLASS_LIBRARY_GENERATED_KERNEL_LIST_FILE}"
|
||||
--cuda-version "${CUTLASS_GENERATOR_CUDA_COMPILER_VERSION}"
|
||||
|
||||
@@ -292,6 +292,7 @@ struct GemmUniversalArguments {
|
||||
int sm_count{0};
|
||||
|
||||
library::RasterOrder raster_order{};
|
||||
int swizzle_size{1};
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -595,12 +595,10 @@ 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,
|
||||
// which just sets D's strides to C's strides.
|
||||
const TensorStride stride_D = stride_C;
|
||||
|
||||
const int num_groups = config.problem_size.groups;
|
||||
if (num_groups != 1) {
|
||||
@@ -773,9 +771,7 @@ 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 TensorStride stride_D = 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.");
|
||||
@@ -852,6 +848,12 @@ private:
|
||||
std::cerr << "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) {
|
||||
return status;
|
||||
}
|
||||
|
||||
out_args.mainloop.ptr_A = reinterpret_cast<ElementA const*>(in_args.A);
|
||||
out_args.mainloop.ptr_B = reinterpret_cast<ElementB const*>(in_args.B);
|
||||
|
||||
|
||||
@@ -250,6 +250,10 @@ protected:
|
||||
/* Query device SM count to pass onto the kernel as an argument, where needed */
|
||||
operator_args.hw_info.sm_count = arguments->sm_count;
|
||||
|
||||
if constexpr (!std::is_const_v<decltype(operator_args.scheduler.max_swizzle_size)>) {
|
||||
operator_args.scheduler.max_swizzle_size = arguments->swizzle_size;
|
||||
}
|
||||
|
||||
if constexpr (!std::is_const_v<decltype(operator_args.scheduler.raster_order)>) {
|
||||
using Enum_t = decltype(operator_args.scheduler.raster_order);
|
||||
switch (arguments->raster_order) {
|
||||
|
||||
@@ -489,7 +489,7 @@ template <
|
||||
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
|
||||
>
|
||||
void make_conv_fprop(Manifest &manifest) {
|
||||
|
||||
#if !defined(CUTLASS_PROFILER_DISABLE_REFERENCE)
|
||||
manifest.append(new ConvReferenceOperation<
|
||||
Provider::kReferenceHost,
|
||||
cutlass::conv::Operator::kFprop,
|
||||
@@ -515,6 +515,7 @@ void make_conv_fprop(Manifest &manifest) {
|
||||
ConvertOp_,
|
||||
InnerProductOp_
|
||||
>);
|
||||
#endif // !defined(CUTLASS_PROFILER_DISABLE_REFERENCE)
|
||||
}
|
||||
|
||||
/// Constructs Dgrad and Wgrad reference operators.
|
||||
@@ -532,7 +533,7 @@ template <
|
||||
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
|
||||
>
|
||||
void make_conv_backwards(Manifest &manifest) {
|
||||
|
||||
#if !defined(CUTLASS_PROFILER_DISABLE_REFERENCE)
|
||||
manifest.append(new ConvReferenceOperation<
|
||||
Provider::kReferenceHost,
|
||||
cutlass::conv::Operator::kDgrad,
|
||||
@@ -584,6 +585,7 @@ void make_conv_backwards(Manifest &manifest) {
|
||||
ConvertOp_,
|
||||
InnerProductOp_
|
||||
>);
|
||||
#endif // !defined(CUTLASS_PROFILER_DISABLE_REFERENCE)
|
||||
}
|
||||
|
||||
/// Six operators for the price of one.
|
||||
|
||||
@@ -293,7 +293,7 @@ template <
|
||||
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
|
||||
>
|
||||
void make_gemm(Manifest &manifest) {
|
||||
|
||||
#if !defined(CUTLASS_PROFILER_DISABLE_REFERENCE)
|
||||
manifest.append(new GemmReferenceOperation<
|
||||
Provider::kReferenceHost,
|
||||
ElementA_, LayoutA_, TransformA,
|
||||
@@ -317,6 +317,7 @@ void make_gemm(Manifest &manifest) {
|
||||
ConvertOp_,
|
||||
InnerProductOp_
|
||||
>);
|
||||
#endif
|
||||
}
|
||||
|
||||
/// Helper to create NN, NT, TN, and TT GEMM layouts.
|
||||
|
||||
@@ -84,6 +84,8 @@ public:
|
||||
int batch_count{1};
|
||||
|
||||
cutlass::library::RasterOrder raster_order{cutlass::library::RasterOrder::kHeuristic};
|
||||
int swizzle_size{1};
|
||||
|
||||
// gemm with parallel interleaved reduction
|
||||
// gemm epilogue (alpha, beta) = (1.0, 0.0)
|
||||
// reduction epilogue (alpha, beta) = (GemmProblem::alpha, GemmProblem::beta)
|
||||
|
||||
@@ -72,7 +72,7 @@ public:
|
||||
// Methods
|
||||
//
|
||||
|
||||
Library(CommandLine const &cmdline);
|
||||
explicit Library(CommandLine const &cmdline);
|
||||
|
||||
void print_usage(std::ostream &out) const;
|
||||
void print_options(std::ostream &out, int indent = 0) const;
|
||||
@@ -94,7 +94,7 @@ public:
|
||||
// Methods
|
||||
//
|
||||
|
||||
Device(CommandLine const &cmdline);
|
||||
explicit Device(CommandLine const &cmdline);
|
||||
|
||||
void print_usage(std::ostream &out) const;
|
||||
void print_options(std::ostream &out, int indent = 0) const;
|
||||
@@ -128,7 +128,7 @@ public:
|
||||
// Methods
|
||||
//
|
||||
|
||||
Initialization(CommandLine const &cmdline);
|
||||
explicit Initialization(CommandLine const &cmdline);
|
||||
|
||||
void print_usage(std::ostream &out) const;
|
||||
void print_options(std::ostream &out, int indent = 0) const;
|
||||
@@ -170,7 +170,7 @@ public:
|
||||
// Methods
|
||||
//
|
||||
|
||||
Verification(CommandLine const &cmdline);
|
||||
explicit Verification(CommandLine const &cmdline);
|
||||
|
||||
void print_usage(std::ostream &out) const;
|
||||
void print_options(std::ostream &out, int indent = 0) const;
|
||||
@@ -186,22 +186,26 @@ public:
|
||||
struct Profiling {
|
||||
|
||||
/// Number of workspaces to rotate through to avoid cache-resident working sets
|
||||
int workspace_count;
|
||||
int workspace_count{0};
|
||||
|
||||
/// Number of iterations to warmup each kernel prior to profiling
|
||||
int warmup_iterations;
|
||||
int warmup_iterations{10};
|
||||
|
||||
/// Number of iterations to profile each kernel - if 0, kernels are launched up to the profiling duration
|
||||
int iterations;
|
||||
int iterations{100};
|
||||
|
||||
/// Number of ms to sleep between profiling periods (ms)
|
||||
int sleep_duration;
|
||||
int sleep_duration{50};
|
||||
|
||||
/// If true, profiling is actually conducted.
|
||||
bool enabled;
|
||||
bool enabled{true};
|
||||
|
||||
/// If true, profiling returns an error code if no kernels are found to match the filters.
|
||||
bool error_on_no_match = false;
|
||||
bool error_on_no_match{false};
|
||||
|
||||
/// If true, profiling returns an error code if no kernel are profiled
|
||||
// Sometimes the kernel matches but failed to profile (e.g. can_implement() error)
|
||||
bool error_if_nothing_is_profiled{false};
|
||||
|
||||
/// List of providers of each functionality to be profiled
|
||||
ProviderVector providers;
|
||||
@@ -210,7 +214,7 @@ public:
|
||||
// Methods
|
||||
//
|
||||
|
||||
Profiling(CommandLine const &cmdline);
|
||||
explicit Profiling(CommandLine const &cmdline);
|
||||
|
||||
void print_usage(std::ostream &out) const;
|
||||
void print_options(std::ostream &out, int indent = 0) const;
|
||||
@@ -255,7 +259,7 @@ public:
|
||||
// Methods
|
||||
//
|
||||
|
||||
Report(CommandLine const &cmdline);
|
||||
explicit Report(CommandLine const &cmdline);
|
||||
|
||||
void print_usage(std::ostream &out) const;
|
||||
void print_options(std::ostream &out, int indent = 0) const;
|
||||
@@ -277,7 +281,7 @@ public:
|
||||
// Methods
|
||||
//
|
||||
|
||||
About(CommandLine const &cmdline);
|
||||
explicit About(CommandLine const &cmdline);
|
||||
|
||||
void print_usage(std::ostream &out) const;
|
||||
void print_options(std::ostream &out, int indent = 0) const;
|
||||
@@ -320,7 +324,7 @@ public:
|
||||
|
||||
public:
|
||||
|
||||
Options(CommandLine const &cmdline);
|
||||
explicit Options(CommandLine const &cmdline);
|
||||
|
||||
void print_usage(std::ostream &out) const;
|
||||
void print_options(std::ostream &out) const;
|
||||
|
||||
@@ -256,7 +256,7 @@ struct ScalarArgument : public KernelArgument {
|
||||
virtual std::ostream &print(std::ostream &out) const;
|
||||
};
|
||||
|
||||
using ValueCollection = std::vector<std::string>;
|
||||
using ValueCollection = std::vector<decltype(ScalarValue::value)>;
|
||||
|
||||
/// Abstract base class to iterate over values within arguments
|
||||
struct ScalarValueIterator : public KernelArgument::ValueIterator {
|
||||
@@ -271,7 +271,7 @@ struct ScalarArgument : public KernelArgument {
|
||||
// Methods
|
||||
//
|
||||
|
||||
ScalarValueIterator(ScalarArgument const *argument = nullptr);
|
||||
explicit ScalarValueIterator(ScalarArgument const *argument = nullptr);
|
||||
|
||||
virtual void operator++();
|
||||
virtual bool operator==(ValueIterator const &it) const;
|
||||
@@ -292,7 +292,7 @@ struct ScalarArgument : public KernelArgument {
|
||||
//
|
||||
|
||||
/// Default ctor
|
||||
ScalarArgument(
|
||||
explicit ScalarArgument(
|
||||
ArgumentDescription const *description
|
||||
):
|
||||
KernelArgument(description) { }
|
||||
@@ -632,7 +632,7 @@ struct TensorArgument : public KernelArgument {
|
||||
// Methods
|
||||
//
|
||||
|
||||
TensorValueIterator(TensorArgument const *argument_);
|
||||
explicit TensorValueIterator(TensorArgument const *argument_);
|
||||
|
||||
virtual void operator++();
|
||||
virtual bool operator==(ValueIterator const &it) const;
|
||||
@@ -649,7 +649,7 @@ struct TensorArgument : public KernelArgument {
|
||||
//
|
||||
|
||||
/// Default ctor
|
||||
TensorArgument(
|
||||
explicit TensorArgument(
|
||||
ArgumentDescription const *description
|
||||
):
|
||||
KernelArgument(description) { }
|
||||
@@ -690,7 +690,7 @@ struct EnumeratedTypeArgument : public KernelArgument {
|
||||
virtual std::ostream &print(std::ostream &out) const;
|
||||
};
|
||||
|
||||
using ValueCollection = std::vector<std::string>;
|
||||
using ValueCollection = std::vector<decltype(EnumeratedTypeValue::element)>;
|
||||
|
||||
/// Abstract base class to iterate over values within arguments
|
||||
struct EnumeratedTypeValueIterator : public KernelArgument::ValueIterator {
|
||||
@@ -705,7 +705,7 @@ struct EnumeratedTypeArgument : public KernelArgument {
|
||||
// Methods
|
||||
//
|
||||
|
||||
EnumeratedTypeValueIterator(EnumeratedTypeArgument const *argument_ = nullptr);
|
||||
explicit EnumeratedTypeValueIterator(EnumeratedTypeArgument const *argument_ = nullptr);
|
||||
|
||||
virtual void operator++();
|
||||
virtual bool operator==(ValueIterator const &it) const;
|
||||
@@ -725,7 +725,7 @@ struct EnumeratedTypeArgument : public KernelArgument {
|
||||
//
|
||||
|
||||
/// Default ctor
|
||||
EnumeratedTypeArgument(ArgumentDescription const *description):
|
||||
explicit EnumeratedTypeArgument(ArgumentDescription const *description):
|
||||
KernelArgument(description) {}
|
||||
|
||||
virtual bool not_null() const {
|
||||
@@ -819,7 +819,7 @@ public:
|
||||
//
|
||||
|
||||
/// Default ctor
|
||||
ProblemSpace() {}
|
||||
ProblemSpace() = default;
|
||||
|
||||
/// Constructs a problem space from a vector of arguments. This vector must outlive
|
||||
/// the ProblemSpace object, which stores pointers to objects within the
|
||||
|
||||
@@ -143,16 +143,19 @@ void CutlassProfiler::enumerate_() {
|
||||
/// Profiles all operations
|
||||
int CutlassProfiler::profile_() {
|
||||
|
||||
int result = 0;
|
||||
// Keep track of all device memory tensor in map
|
||||
DeviceContext device_context;
|
||||
// For all profilers
|
||||
|
||||
int result = 0;
|
||||
// For all profilers (e.g. gemm/sparse_gemm/conv2d...)
|
||||
for (auto & profiler : operation_profilers_) {
|
||||
|
||||
if (options_.operation_kind == library::OperationKind::kInvalid ||
|
||||
options_.operation_kind == profiler->kind()) {
|
||||
options_.operation_kind == profiler->kind()) {
|
||||
|
||||
result = profiler->profile_all(options_, library::Singleton::get().manifest, device_context);
|
||||
|
||||
// If some profile failed, terminate immediately
|
||||
if (result) {
|
||||
return result;
|
||||
}
|
||||
|
||||
@@ -75,6 +75,7 @@ GemmOperationProfiler::GemmOperationProfiler(Options const &options):
|
||||
{ArgumentTypeID::kInteger, {"split_k_slices", "split-k-slices"}, "Number of partitions of K dimension"},
|
||||
{ArgumentTypeID::kInteger, {"batch_count", "batch-count"}, "Number of GEMMs computed in one batch"},
|
||||
{ArgumentTypeID::kEnumerated, {"raster_order", "raster-order"}, "Raster order (heuristic, along_n, along_m)"},
|
||||
{ArgumentTypeID::kInteger, {"swizzle_size", "swizzle-size"}, "Size to swizzle"},
|
||||
},
|
||||
{ library::Provider::kCUBLAS}
|
||||
) {
|
||||
@@ -191,6 +192,14 @@ Status GemmOperationProfiler::GemmProblem::parse(
|
||||
this->mode = library::GemmUniversalMode::kBatched;
|
||||
}
|
||||
|
||||
if (!arg_as_int(this->swizzle_size, "swizzle_size", problem_space, problem)) {
|
||||
// default value
|
||||
this->swizzle_size = 1;
|
||||
if (this->swizzle_size <= 0) {
|
||||
return Status::kErrorInvalidProblem;
|
||||
}
|
||||
}
|
||||
|
||||
if (!arg_as_RasterOrder(this->raster_order, "raster_order", problem_space, problem)) {
|
||||
// default value
|
||||
this->raster_order = library::RasterOrder::kHeuristic;
|
||||
@@ -329,6 +338,8 @@ void GemmOperationProfiler::GemmProblem::initialize_result(
|
||||
set_argument(result, "split_k_slices", problem_space, split_k_slices);
|
||||
set_argument(result, "batch_count", problem_space, batch_count);
|
||||
set_argument(result, "raster_order", problem_space, library::to_string(raster_order));
|
||||
set_argument(result, "swizzle_size", problem_space, swizzle_size);
|
||||
|
||||
set_argument(result, "alpha", problem_space,
|
||||
library::lexical_cast(alpha, operation_desc.element_epilogue));
|
||||
|
||||
@@ -383,6 +394,7 @@ Status GemmOperationProfiler::initialize_configuration(
|
||||
gemm_workspace_.arguments.alpha = problem_.alpha.data();
|
||||
gemm_workspace_.arguments.beta = problem_.beta.data();
|
||||
gemm_workspace_.arguments.pointer_mode = library::ScalarPointerMode::kHost;
|
||||
gemm_workspace_.arguments.swizzle_size = problem_.swizzle_size;
|
||||
gemm_workspace_.arguments.raster_order = problem_.raster_order;
|
||||
// initialize reduction operation for parallel splitKMode
|
||||
if (problem_.split_k_mode == library::SplitKMode::kParallel) {
|
||||
|
||||
@@ -345,6 +345,7 @@ int OperationProfiler::profile_all(
|
||||
|
||||
// For each operation in manifest
|
||||
int matched_operation_count = 0;
|
||||
int profiled_operation_count = 0;
|
||||
for (auto const& operation_ptr : manifest) {
|
||||
|
||||
library::Operation const *operation = operation_ptr.get();
|
||||
@@ -434,7 +435,7 @@ int OperationProfiler::profile_all(
|
||||
// If there was an internal error, consume the CUDA error and move to the next operation.
|
||||
(void)cudaGetLastError();
|
||||
|
||||
report.append_results(results_);
|
||||
report.append_result(model_result_);
|
||||
continue;
|
||||
}
|
||||
else if (status != Status::kSuccess) {
|
||||
@@ -522,25 +523,42 @@ int OperationProfiler::profile_all(
|
||||
operation,
|
||||
problem_space,
|
||||
problem);
|
||||
|
||||
// Count op as profiled, even it failed to profile
|
||||
profiled_operation_count++;
|
||||
}
|
||||
|
||||
report.append_results(results_);
|
||||
results_.clear();
|
||||
}
|
||||
} // if op satisfied compute capacity
|
||||
|
||||
if (!continue_profiling) {
|
||||
// break out of `for op in manifest` loop and move to next problem
|
||||
// `for each problem in problem space` conditional check on not continue profiling
|
||||
break;
|
||||
}
|
||||
}
|
||||
} // for op in manifest
|
||||
|
||||
// If we did not find any kernels that match our filters and error_on_no_match was set, report an error
|
||||
if (options.profiling.error_on_no_match && matched_operation_count <= 0) {
|
||||
#if !NDEBUG
|
||||
std::cout << "Error: No matching kernels found with kernel selection filters [--error_on_no_match]" << std::endl;
|
||||
std::cerr << "Error: No matching kernels found with kernel selection filters [--error_on_no_match]" << std::endl;
|
||||
#endif
|
||||
retval = 1;
|
||||
retval |= 1;
|
||||
// Stop profiling on error no match
|
||||
continue_profiling = false;
|
||||
}
|
||||
}
|
||||
|
||||
if (options.profiling.error_if_nothing_is_profiled && options.profiling.enabled && profiled_operation_count <= 0) {
|
||||
#if !NDEBUG
|
||||
std::cerr << "Error: No kernels profiled found with kernel selection filters [--error_if_nothing_is_profiled]" << std::endl;
|
||||
#endif
|
||||
retval |= 1;
|
||||
// Stop profiling on error no match
|
||||
continue_profiling = false;
|
||||
}
|
||||
|
||||
} // for each problem in problem space
|
||||
|
||||
return retval;
|
||||
}
|
||||
|
||||
@@ -726,11 +726,13 @@ Options::Options(cutlass::CommandLine const &cmdline):
|
||||
else if (cmdline.check_cmd_line_flag("kernels")) {
|
||||
cmdline.get_cmd_line_arguments("kernels", operation_names);
|
||||
profiling.error_on_no_match = cmdline.check_cmd_line_flag("error-on-no-match");
|
||||
profiling.error_if_nothing_is_profiled = cmdline.check_cmd_line_flag("error-if-nothing-is-profiled");
|
||||
}
|
||||
|
||||
if (cmdline.check_cmd_line_flag("ignore-kernels")) {
|
||||
cmdline.get_cmd_line_arguments("ignore-kernels", excluded_operation_names);
|
||||
profiling.error_on_no_match = cmdline.check_cmd_line_flag("error-on-no-match");
|
||||
profiling.error_if_nothing_is_profiled = cmdline.check_cmd_line_flag("error-if-nothing-is-profiled");
|
||||
}
|
||||
|
||||
// Prevent launches on the device for anything other than CUTLASS operation
|
||||
|
||||
@@ -395,10 +395,6 @@ std::unique_ptr<KernelArgument::ValueIterator> EnumeratedTypeArgument::end() con
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
ProblemSpace::Iterator::Iterator() {
|
||||
|
||||
}
|
||||
|
||||
ProblemSpace::Iterator::Iterator(ProblemSpace const &problem_space) {
|
||||
for (auto const & arg_ptr : problem_space.arguments) {
|
||||
construct_(arg_ptr.get());
|
||||
|
||||
@@ -104,23 +104,15 @@ public:
|
||||
/// Constant reference to element in tensor
|
||||
using ConstReference = typename ConstTensorRef::Reference;
|
||||
|
||||
/// Note: Below is used to handle packing of subbyte elements
|
||||
/// kBitsStoredVec : The bits of store vec that could be divisiable by the element
|
||||
/// kElementsPerStoredVec : The number of elements could be stored in per store vec
|
||||
/// kNumStoragePerStoredVec : How much storage(i.e. sizeof(element storage)) the store vec needs to consume.
|
||||
/// Usually the element storage of subbyte is uint8_t.
|
||||
/// Example
|
||||
/// int2: kBitsStoredVec = 8; kElementsPerStoredVec = 4; kNumStoragePerStoredVec = 1 uint8_t;
|
||||
/// int4: kBitsStoredVec = 8; kElementsPerStoredVec = 2; kNumStoragePerStoredVec = 1 uint8_t;
|
||||
static constexpr int kBitsStoredVec = (sizeof_bits<Element>::value < 8) ? cutlass::lcm(sizeof_bits<Element>::value, 8) : sizeof_bits<Element>::value;
|
||||
static constexpr int kElementsPerStoredVec = kBitsStoredVec / sizeof_bits<Element>::value;
|
||||
static constexpr int kNumStoragePerStoredVec = kBitsStoredVec / (sizeof(Element) * 8);
|
||||
|
||||
static_assert(kBitsStoredVec != 0, "kBitsStoredVec can not be zero");
|
||||
static_assert(kElementsPerStoredVec != 0, "kElementsPerStoredVec can not be zero");
|
||||
static_assert(kNumStoragePerStoredVec != 0, "kNumStoragePerStoredVec can not be zero");
|
||||
|
||||
private:
|
||||
private:
|
||||
using StorageUnit = typename platform::conditional_t<std::is_same_v<Element, bool>, uint8_t, // Avoid the std::vector<bool> specialization
|
||||
typename platform::conditional_t<sizeof_bits<Element>::value % 8 == 0, // Handle subbyte types
|
||||
Element, uint8_t>>;
|
||||
using StorageContainerCalculator = cutlass::detail::StorageContainerCalculator<Element, StorageUnit>;
|
||||
static constexpr int kContainerTypeNumBits = StorageContainerCalculator::kContainerTypeNumBits;
|
||||
static constexpr int kContainerTypeNumLogicalElements = StorageContainerCalculator::kContainerTypeNumLogicalElements;
|
||||
static constexpr int kContainerTypeNumBytes = StorageContainerCalculator::kContainerTypeNumBytes;
|
||||
static constexpr int kContainerTypeNumStorageUnit = StorageContainerCalculator::kContainerTypeNumStorageUnit;
|
||||
|
||||
//
|
||||
// Data members
|
||||
@@ -133,13 +125,17 @@ public:
|
||||
Layout layout_;
|
||||
|
||||
/// Host-side memory allocation
|
||||
/// avoid the std::vector<bool> specialization
|
||||
std::vector<std::conditional_t<std::is_same_v<Element,bool>, uint8_t, Element>> host_;
|
||||
std::vector<StorageUnit> host_;
|
||||
|
||||
/// Device-side memory
|
||||
device_memory::allocation<Element> device_;
|
||||
device_memory::allocation<StorageUnit> device_;
|
||||
|
||||
public:
|
||||
/// number of containers
|
||||
size_t count_to_container_storage_unit_count(size_t count) {
|
||||
return (count + kContainerTypeNumLogicalElements - 1) / kContainerTypeNumLogicalElements * kContainerTypeNumStorageUnit;
|
||||
}
|
||||
|
||||
public:
|
||||
//
|
||||
// Device and Host Methods
|
||||
//
|
||||
@@ -185,15 +181,15 @@ public:
|
||||
device_.reset();
|
||||
host_.clear();
|
||||
|
||||
count = (count + kElementsPerStoredVec - 1) / kElementsPerStoredVec * kNumStoragePerStoredVec;
|
||||
host_.resize(count);
|
||||
size_t count_container = count_to_container_storage_unit_count(count);
|
||||
host_.resize(count_container);
|
||||
|
||||
// Allocate memory
|
||||
Element* device_memory = nullptr;
|
||||
StorageUnit* device_memory = nullptr;
|
||||
if (device_backed_) {
|
||||
device_memory = device_memory::allocate<Element>(count);
|
||||
device_memory = device_memory::allocate<StorageUnit>(count_container);
|
||||
}
|
||||
device_.reset(device_memory, device_backed_ ? count : 0);
|
||||
device_.reset(device_memory, device_backed_ ? count_container : 0);
|
||||
}
|
||||
|
||||
/// Updates the extent and layout of the HostTensor. Allocates memory according to the new
|
||||
@@ -229,8 +225,9 @@ public:
|
||||
layout_ = layout;
|
||||
|
||||
LongIndex new_size = size_t(layout_.capacity(extent_));
|
||||
LongIndex new_size_container = count_to_container_storage_unit_count((layout_.capacity(extent_)));
|
||||
|
||||
if (static_cast<decltype(host_.size())>(new_size) > host_.size()) {
|
||||
if (static_cast<decltype(host_.size())>(new_size_container) > host_.size()) {
|
||||
reserve(new_size, device_backed_);
|
||||
}
|
||||
}
|
||||
@@ -244,14 +241,14 @@ public:
|
||||
resize(extent, Layout::packed(extent), device_backed_);
|
||||
}
|
||||
|
||||
/// Returns the number of elements stored in the host tensor
|
||||
/// Returns the logical number of elements stored in the host tensor
|
||||
size_t size() const {
|
||||
return host_.size() / kNumStoragePerStoredVec * kElementsPerStoredVec;
|
||||
return layout_.capacity(extent_);
|
||||
}
|
||||
|
||||
/// Returns the logical capacity based on extent and layout. May differ from size().
|
||||
/// Returns the logical capacity in terms of number of elements. May be larger than the size().
|
||||
LongIndex capacity() const {
|
||||
return layout_.capacity(extent_);
|
||||
return host_.size() / kContainerTypeNumStorageUnit * kContainerTypeNumLogicalElements;
|
||||
}
|
||||
|
||||
/// Gets pointer to host data
|
||||
@@ -277,10 +274,10 @@ public:
|
||||
}
|
||||
|
||||
/// Gets pointer to device data
|
||||
Element * device_data() { return device_.get(); }
|
||||
Element * device_data() { return reinterpret_cast<Element *>(device_.get()); }
|
||||
|
||||
/// Gets pointer to device data
|
||||
Element const * device_data() const { return device_.get(); }
|
||||
Element const * device_data() const { return reinterpret_cast<Element const *>(device_.get()); }
|
||||
|
||||
/// Gets pointer to device data with a pointer offset
|
||||
Element * device_data_ptr_offset(LongIndex ptr_element_offset) { return &ReferenceFactory<Element>::get(device_data(), ptr_element_offset); }
|
||||
@@ -389,7 +386,7 @@ public:
|
||||
void sync_host() {
|
||||
if (device_backed()) {
|
||||
device_memory::copy_to_host(
|
||||
host_data(), device_data(), size());
|
||||
host_.data(), device_.get(), device_.size());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -397,7 +394,7 @@ public:
|
||||
void sync_device() {
|
||||
if (device_backed()) {
|
||||
device_memory::copy_to_device(
|
||||
device_data(), host_data(), size());
|
||||
device_.get(), host_.data(), host_.capacity());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -412,8 +409,9 @@ public:
|
||||
else {
|
||||
count = __NV_STD_MIN(capacity(), count);
|
||||
}
|
||||
size_t container_count = count_to_container_storage_unit_count(count);
|
||||
device_memory::copy_to_host(
|
||||
host_data(), ptr_device, count);
|
||||
host_.data(), reinterpret_cast<StorageUnit const *>(ptr_device), container_count);
|
||||
}
|
||||
|
||||
/// Copy data from a caller-supplied device pointer into host memory.
|
||||
@@ -427,8 +425,9 @@ public:
|
||||
else {
|
||||
count = __NV_STD_MIN(capacity(), count);
|
||||
}
|
||||
size_t container_count = count_to_container_storage_unit_count(count);
|
||||
device_memory::copy_device_to_device(
|
||||
device_data(), ptr_device, count);
|
||||
device_.get(), reinterpret_cast<StorageUnit const *>(ptr_device), container_count);
|
||||
}
|
||||
|
||||
/// Copy data from a caller-supplied device pointer into host memory.
|
||||
@@ -442,8 +441,9 @@ public:
|
||||
else {
|
||||
count = __NV_STD_MIN(capacity(), count);
|
||||
}
|
||||
size_t container_count = count_to_container_storage_unit_count(count);
|
||||
device_memory::copy_to_device(
|
||||
device_data(), ptr_host, count);
|
||||
device_.get(), reinterpret_cast<StorageUnit const *>(ptr_host), container_count);
|
||||
}
|
||||
|
||||
/// Copy data from a caller-supplied device pointer into host memory.
|
||||
@@ -457,8 +457,9 @@ public:
|
||||
else {
|
||||
count = __NV_STD_MIN(capacity(), count);
|
||||
}
|
||||
size_t container_count = count_to_container_storage_unit_count(count);
|
||||
device_memory::copy_host_to_host(
|
||||
host_data(), ptr_host, count);
|
||||
host_.data(), reinterpret_cast<StorageUnit const *>(ptr_host), container_count);
|
||||
}
|
||||
|
||||
/// Copy data from a caller-supplied device pointer into host memory.
|
||||
@@ -472,8 +473,9 @@ public:
|
||||
else {
|
||||
count = __NV_STD_MIN(capacity(), count);
|
||||
}
|
||||
size_t container_count = count_to_container_storage_unit_count(count);
|
||||
device_memory::copy_to_host(
|
||||
ptr_host, device_data(), count);
|
||||
reinterpret_cast<StorageUnit *>(ptr_host), device_.get(), container_count);
|
||||
}
|
||||
|
||||
/// Copy data from a caller-supplied device pointer into host memory.
|
||||
@@ -487,8 +489,9 @@ public:
|
||||
else {
|
||||
count = __NV_STD_MIN(capacity(), count);
|
||||
}
|
||||
size_t container_count = count_to_container_storage_unit_count(count);
|
||||
device_memory::copy_device_to_device(
|
||||
ptr_device, device_data(), count);
|
||||
reinterpret_cast<StorageUnit *>(ptr_device), device_.get(), container_count);
|
||||
}
|
||||
|
||||
/// Copy data from a caller-supplied device pointer into host memory.
|
||||
@@ -502,8 +505,9 @@ public:
|
||||
else {
|
||||
count = __NV_STD_MIN(capacity(), count);
|
||||
}
|
||||
size_t container_count = count_to_container_storage_unit_count(count);
|
||||
device_memory::copy_to_device(
|
||||
ptr_device, host_data(), count);
|
||||
reinterpret_cast<StorageUnit *>(ptr_device), host_.data(), container_count);
|
||||
}
|
||||
|
||||
/// Copy data from a caller-supplied device pointer into host memory.
|
||||
@@ -517,8 +521,9 @@ public:
|
||||
else {
|
||||
count = __NV_STD_MIN(capacity(), count);
|
||||
}
|
||||
size_t container_count = count_to_container_storage_unit_count(count);
|
||||
device_memory::copy_host_to_host(
|
||||
ptr_host, host_data(), count);
|
||||
reinterpret_cast<StorageUnit *>(ptr_host), host_.data(), container_count);
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -458,6 +458,7 @@ make_cute_packed_stride(
|
||||
|
||||
// Filter cutlass::layout::TensorKCSR -> rank-3 stride (k, (_1, s, r), _0)
|
||||
template <class IntT>
|
||||
CUTLASS_HOST_DEVICE
|
||||
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>, cute::Int<0>>
|
||||
make_cute_packed_stride(
|
||||
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>, cute::Int<0>> s,
|
||||
@@ -497,6 +498,71 @@ make_cute_packed_stride(
|
||||
return s_copy;
|
||||
}
|
||||
|
||||
|
||||
//
|
||||
// Wgrad output tensor ((_1, s, r, t), k, _0)
|
||||
//
|
||||
|
||||
// Filter cutlass::layout::TensorCSK -> rank-3 stride ((_1, s), k, _0)
|
||||
template <class IntT>
|
||||
CUTLASS_HOST_DEVICE
|
||||
cute::Stride<cute::Stride<cute::Int<1>, IntT>, IntT, cute::Int<0>>
|
||||
make_cute_packed_stride(
|
||||
cute::Stride<cute::Stride<cute::Int<1>, IntT>, IntT, cute::Int<0>> s,
|
||||
[[maybe_unused]] cute::array<int32_t, 3> shape_output,
|
||||
cute::array<IntT, 3> stride_ksc,
|
||||
conv::Operator ConvOp) {
|
||||
static_assert(std::is_integral_v<IntT>,
|
||||
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
|
||||
|
||||
assert(stride_ksc[2] == 1);
|
||||
auto s_copy = s;
|
||||
cute::get<1,0>(s_copy) = stride_ksc[0];
|
||||
cute::get<0,1>(s_copy) = stride_ksc[1];
|
||||
return s_copy;
|
||||
}
|
||||
|
||||
// Filter cutlass::layout::TensorCSRK -> rank-3 stride ((_1, s, r), k, _0)
|
||||
template <class IntT>
|
||||
CUTLASS_HOST_DEVICE
|
||||
cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT>, IntT, cute::Int<0>>
|
||||
make_cute_packed_stride(
|
||||
cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT>, IntT, cute::Int<0>> s,
|
||||
[[maybe_unused]] cute::array<int32_t, 4> shape_output,
|
||||
cute::array<IntT, 4> stride_krsc,
|
||||
conv::Operator ConvOp) {
|
||||
static_assert(std::is_integral_v<IntT>,
|
||||
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
|
||||
|
||||
assert(stride_krsc[3] == 1);
|
||||
auto s_copy = s;
|
||||
cute::get<1,0>(s_copy) = stride_krsc[0];
|
||||
cute::for_each(cute::make_seq<2>{}, [&](auto i) {
|
||||
cute::get<0,2-i>(s_copy) = stride_krsc[i+1];
|
||||
});
|
||||
return s_copy;
|
||||
}
|
||||
|
||||
// Filter cutlass::layout::TensorCSRTK -> rank-3 stride ((_1, s, r, t), k, _0)
|
||||
template <class IntT>
|
||||
CUTLASS_HOST_DEVICE
|
||||
cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT, IntT>, IntT, cute::Int<0>>
|
||||
make_cute_packed_stride(
|
||||
cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT, IntT>, IntT, cute::Int<0>> s,
|
||||
[[maybe_unused]] cute::array<int32_t, 5> shape_output,
|
||||
cute::array<IntT, 5> stride_ktrsc,
|
||||
conv::Operator ConvOp) {
|
||||
static_assert(std::is_integral_v<IntT>,
|
||||
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
|
||||
|
||||
assert(stride_ktrsc[4] == 1);
|
||||
auto s_copy = s;
|
||||
cute::get<1,0>(s_copy) = stride_ktrsc[0];
|
||||
cute::for_each(cute::make_seq<3>{}, [&](auto i) {
|
||||
cute::get<0,3-i>(s_copy) = stride_ktrsc[i+1];
|
||||
});
|
||||
return s_copy;
|
||||
}
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass
|
||||
|
||||
@@ -443,10 +443,10 @@ struct RandomUniformFunc {
|
||||
int int_scale_ = -1
|
||||
):
|
||||
seed(seed_),
|
||||
range(static_cast<FloatType>(max_ - min)),
|
||||
range(static_cast<FloatType>(max_) - static_cast<FloatType>(min)),
|
||||
max(static_cast<FloatType>(max_)),
|
||||
int_scale(int_scale_) {
|
||||
|
||||
|
||||
float_scale_up = FloatType(IntType(2) << int_scale); // scale up to clamp low order bits
|
||||
float_scale_down = FloatType(1) / FloatType(IntType(2) << int_scale);
|
||||
}
|
||||
|
||||
@@ -125,7 +125,8 @@ template<
|
||||
class ShapePadding,
|
||||
class StrideTraversal,
|
||||
class ShapeDilation,
|
||||
class EpilogueFusionParams>
|
||||
class EpilogueFusionParams
|
||||
>
|
||||
struct ConvReferenceImpl {
|
||||
using ElementAcc = typename EpilogueFusionParams::ElementAcc;
|
||||
using ElementC = typename EpilogueFusionParams::ElementC;
|
||||
@@ -145,7 +146,6 @@ struct ConvReferenceImpl {
|
||||
NumericConverter<ElementOut, ElementCompute> output_converter;
|
||||
|
||||
EpilogueFusionParams& epi_fusion_params_;
|
||||
|
||||
TensorA const& tensor_a_;
|
||||
TensorB const& tensor_b_;
|
||||
TensorC const& tensor_c_;
|
||||
@@ -174,7 +174,8 @@ struct ConvReferenceImpl {
|
||||
padding_(padding),
|
||||
tstride_(tstride),
|
||||
dilation_(dilation),
|
||||
epi_fusion_params_(epi_fusion_params) {
|
||||
epi_fusion_params_(epi_fusion_params)
|
||||
{
|
||||
static_assert(rank(ShapePadding{}) == rank(ShapeDilation{}));
|
||||
static_assert(rank(ShapePadding{}) == rank(StrideTraversal{}));
|
||||
}
|
||||
@@ -211,7 +212,9 @@ private:
|
||||
for (int32_t c = 0; c < C; ++c) {
|
||||
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
|
||||
if (detail::is_activation_in_bounds(tensor_a_, n, w, c)) {
|
||||
accumulator += ElementAcc(tensor_a_(c, w, n) * tensor_b_(c, s, k));
|
||||
auto a = tensor_a_(c, w, n);
|
||||
auto b = tensor_b_(c, s, k);
|
||||
accumulator += ElementAcc(a * b);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -256,7 +259,9 @@ private:
|
||||
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
|
||||
int32_t h = p * cute::get<1>(tstride_) - cute::get<1>(padding_) + r * cute::get<1>(dilation_);
|
||||
if (detail::is_activation_in_bounds(tensor_a_, n, h, w, c)) {
|
||||
accumulator += ElementAcc(tensor_a_(c, w, h, n) * tensor_b_(c, s, r, k));
|
||||
auto a = tensor_a_(c, w, h, n);
|
||||
auto b = tensor_b_(c, s, r, k);
|
||||
accumulator += ElementAcc(a * b);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -308,7 +313,9 @@ private:
|
||||
int32_t h = p * cute::get<1>(tstride_) - cute::get<1>(padding_) + r * cute::get<1>(dilation_);
|
||||
int32_t d = z * cute::get<2>(tstride_) - cute::get<2>(padding_) + t * cute::get<2>(dilation_);
|
||||
if (detail::is_activation_in_bounds(tensor_a_, n, d, h, w, c)) {
|
||||
accumulator += ElementAcc(tensor_a_(c, w, h, d, n) * tensor_b_(c, s, r, t, k));
|
||||
auto a = tensor_a_(c, w, h, d, n);
|
||||
auto b = tensor_b_(c, s, r, t, k);
|
||||
accumulator += ElementAcc(a * b);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -516,9 +523,12 @@ private:
|
||||
|
||||
// Specialization for 1D wgrad kernel
|
||||
void wgrad_reference(cute::Int<1> spatial_dims) {
|
||||
int32_t N = size<2>(tensor_a_);
|
||||
int32_t Q = size<1>(tensor_a_);
|
||||
int32_t K = size<0>(tensor_a_);
|
||||
int32_t N =
|
||||
size<2>(tensor_a_);
|
||||
int32_t Q =
|
||||
size<1>(tensor_a_);
|
||||
int32_t K =
|
||||
size<0>(tensor_a_);
|
||||
int32_t S = size<1>(tensor_d_);
|
||||
int32_t C = size<0>(tensor_d_);
|
||||
|
||||
@@ -536,8 +546,14 @@ private:
|
||||
for (int32_t n = 0; n < N; ++n) {
|
||||
for (int32_t q = 0; q < Q; ++q) {
|
||||
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
|
||||
if (detail::is_activation_in_bounds(tensor_b_, n, w, c)) {
|
||||
accumulator += ElementAcc(tensor_b_(c, w, n) * tensor_a_(k, q, n));
|
||||
bool is_in_bounds =
|
||||
detail::is_activation_in_bounds(tensor_b_, n, w, c);
|
||||
if (is_in_bounds) {
|
||||
auto act =
|
||||
tensor_b_(c, w, n);
|
||||
auto xformed_act =
|
||||
tensor_a_(k, q, n);
|
||||
accumulator += ElementAcc(act * xformed_act);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -555,10 +571,14 @@ private:
|
||||
|
||||
// Specialization for 2D wgrad kernel
|
||||
void wgrad_reference(cute::Int<2> spatial_dims) {
|
||||
int32_t N = size<3>(tensor_a_);
|
||||
int32_t P = size<2>(tensor_a_);
|
||||
int32_t Q = size<1>(tensor_a_);
|
||||
int32_t K = size<0>(tensor_a_);
|
||||
int32_t N =
|
||||
size<3>(tensor_a_);
|
||||
int32_t P =
|
||||
size<2>(tensor_a_);
|
||||
int32_t Q =
|
||||
size<1>(tensor_a_);
|
||||
int32_t K =
|
||||
size<0>(tensor_a_);
|
||||
int32_t R = size<2>(tensor_d_);
|
||||
int32_t S = size<1>(tensor_d_);
|
||||
int32_t C = size<0>(tensor_d_);
|
||||
@@ -580,8 +600,14 @@ private:
|
||||
for (int32_t q = 0; q < Q; ++q) {
|
||||
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
|
||||
int32_t h = p * cute::get<1>(tstride_) - cute::get<1>(padding_) + r * cute::get<1>(dilation_);
|
||||
if (detail::is_activation_in_bounds(tensor_b_, n, h, w, c)) {
|
||||
accumulator += ElementAcc(tensor_b_(c, w, h, n) * tensor_a_(k, q, p, n));
|
||||
bool is_in_bounds =
|
||||
detail::is_activation_in_bounds(tensor_b_, n, h, w, c);
|
||||
if (is_in_bounds) {
|
||||
auto act =
|
||||
tensor_b_(c, w, h, n);
|
||||
auto xformed_act =
|
||||
tensor_a_(k, q, p, n);
|
||||
accumulator += ElementAcc(act * xformed_act);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -601,11 +627,16 @@ private:
|
||||
|
||||
// Specialization for 3D wgrad kernel
|
||||
void wgrad_reference(cute::Int<3> spatial_dims) {
|
||||
int32_t N = size<4>(tensor_a_);
|
||||
int32_t Z = size<3>(tensor_a_);
|
||||
int32_t P = size<2>(tensor_a_);
|
||||
int32_t Q = size<1>(tensor_a_);
|
||||
int32_t K = size<0>(tensor_a_);
|
||||
int32_t N =
|
||||
size<4>(tensor_a_);
|
||||
int32_t Z =
|
||||
size<3>(tensor_a_);
|
||||
int32_t P =
|
||||
size<2>(tensor_a_);
|
||||
int32_t Q =
|
||||
size<1>(tensor_a_);
|
||||
int32_t K =
|
||||
size<0>(tensor_a_);
|
||||
int32_t T = size<3>(tensor_d_);
|
||||
int32_t R = size<2>(tensor_d_);
|
||||
int32_t S = size<1>(tensor_d_);
|
||||
@@ -631,8 +662,14 @@ private:
|
||||
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
|
||||
int32_t h = p * cute::get<1>(tstride_) - cute::get<1>(padding_) + r * cute::get<1>(dilation_);
|
||||
int32_t d = z * cute::get<2>(tstride_) - cute::get<2>(padding_) + t * cute::get<2>(dilation_);
|
||||
if (detail::is_activation_in_bounds(tensor_b_, n, d, h, w, c)) {
|
||||
accumulator += ElementAcc(tensor_b_(c, w, h, d, n) * tensor_a_(k, q, p, z, n));
|
||||
bool is_in_bounds =
|
||||
detail::is_activation_in_bounds(tensor_b_, n, d, h, w, c);
|
||||
if (is_in_bounds) {
|
||||
auto act =
|
||||
tensor_b_(c, w, h, d, n);
|
||||
auto xformed_act =
|
||||
tensor_a_(k, q, p, z, n);
|
||||
accumulator += ElementAcc(act * xformed_act);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -82,7 +82,6 @@ struct GettMainloopParams {
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template<
|
||||
class ElementScalar_,
|
||||
class ElementScalingFactor_,
|
||||
@@ -117,7 +116,6 @@ struct GettEpilogueParams {
|
||||
using EngineD = typename TensorD::engine_type;
|
||||
using LayoutD = typename TensorD::layout_type;
|
||||
static constexpr bool PerColumnBias = PerColumnBias_;
|
||||
|
||||
ElementScalar alpha = ElementScalar(1);
|
||||
ElementScalar beta = ElementScalar(0);
|
||||
|
||||
@@ -184,6 +182,8 @@ void gett_mainloop(
|
||||
|
||||
static_assert(cute::rank(typename MainloopParams::LayoutA{}) == 3, "M, K, B");
|
||||
static_assert(cute::rank(typename MainloopParams::LayoutB{}) == 3, "N, K, B");
|
||||
|
||||
using cute::raw_pointer_cast;
|
||||
|
||||
using ElementA = typename ElementTraits<typename MainloopParams::EngineA::value_type>::type;
|
||||
using ElementB = typename ElementTraits<typename MainloopParams::EngineB::value_type>::type;
|
||||
@@ -254,6 +254,8 @@ void gett_epilogue(
|
||||
static_assert(cute::rank(typename EpilogueParams::LayoutC{}) == 3, "M, K, B");
|
||||
static_assert(cute::rank(typename EpilogueParams::LayoutD{}) == 3, "N, K, B");
|
||||
|
||||
using cute::raw_pointer_cast;
|
||||
|
||||
using ElementCompute = typename EpilogueParams::ElementCompute;
|
||||
using ElementC = typename EpilogueParams::TensorC::value_type;
|
||||
using ElementD = typename EpilogueParams::TensorD::value_type;
|
||||
@@ -265,7 +267,6 @@ void gett_epilogue(
|
||||
using BiasBinaryOp = typename EpilogueParams::BiasBinaryOp;
|
||||
|
||||
constexpr bool PerColBias = EpilogueParams::PerColumnBias;
|
||||
|
||||
constexpr bool IsScalingAndAmaxOutputNeeded =
|
||||
cute::is_same_v<ElementD, cutlass::float_e4m3_t> or
|
||||
cute::is_same_v<ElementD, cutlass::float_e5m2_t>;
|
||||
@@ -300,7 +301,7 @@ void gett_epilogue(
|
||||
|
||||
// Output related converter
|
||||
NumericConverter<ElementD, ElementCompute> destination_converter;
|
||||
NumericConverter<ElementAux, ElementCompute> aux_destination_converter;
|
||||
[[maybe_unused]] NumericConverter<ElementAux, ElementCompute> aux_destination_converter;
|
||||
NumericConverter<ElementBias, ElementCompute> dBias_converter;
|
||||
|
||||
// Epilogue operations
|
||||
@@ -417,6 +418,7 @@ void gett_epilogue(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(_OPENMP)
|
||||
#pragma omp critical(Abs_Max_Data_Update)
|
||||
#endif
|
||||
|
||||
Reference in New Issue
Block a user