saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-15 21:00:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions
pytorch/torch/csrc/jit/codegen/cuda/executor.cpp
Richard Barnes 3ece9fb45d Check all CUDA API calls for errors in torch/ (#81560) 
Summary:
Original commit changeset: 0bb770d2cdb2

Original Phabricator Diff: D35194935 (79e5b053b6)

Differential Revision: D35291874

Pull Request resolved: https://github.com/pytorch/pytorch/pull/81560
Approved by: https://github.com/ezyang
2022-10-28 00:40:48 +00:00

1320 lines
49 KiB
C++
Raw Blame History

#include <torch/csrc/jit/codegen/cuda/executor.h>
#include <torch/csrc/jit/codegen/cuda/codegen.h>
#include <torch/csrc/jit/codegen/cuda/executor_kernel_arg.h>
#include <torch/csrc/jit/codegen/cuda/executor_utils.h>
#include <torch/csrc/jit/codegen/cuda/instrumentation.h>
#include <torch/csrc/jit/codegen/cuda/ir_all_nodes.h>
#include <torch/csrc/jit/codegen/cuda/ir_utils.h>
#include <torch/csrc/jit/codegen/cuda/iter_visitor.h>
#include <torch/csrc/jit/codegen/cuda/kernel_ir.h>
#include <torch/csrc/jit/codegen/cuda/utils.h>
#include <ATen/core/LegacyTypeDispatch.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/llvm_jit_strings.h>
#include <ATen/cuda/nvrtc_stub/ATenNVRTC.h>
#include <c10/core/DeviceGuard.h>
#include <c10/cuda/CUDAFunctions.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/util/irange.h>
#include <fstream>
namespace torch {
namespace jit {
namespace fuser {
namespace cuda {
int FusionExecutor::fusion_id_counter_ = 0; // NOLINT
namespace {
static const char* defineIndexMode(KernelIndexMode index_mode) {
switch (index_mode) {
case KernelIndexMode::INT32:
return "typedef int nvfuser_index_t;\n";
case KernelIndexMode::INT64:
return "typedef int64_t nvfuser_index_t;\n";
default:
break;
}
TORCH_INTERNAL_ASSERT(false, "unknow indexing mode");
return "";
}
static const char* defineIntegerTypes() {
return R"(
typedef signed char int8_t;
typedef unsigned char uint8_t;
typedef short int int16_t;
typedef unsigned short int uint16_t;
typedef int int32_t;
typedef unsigned int uint32_t;
typedef long long int int64_t;
typedef unsigned long long int uint64_t;
)";
}
static const std::string& defineComplexTypes() {
static std::string result = std::string(R"ESCAPE(
#define POS_INFINITY __int_as_float(0x7f800000)
#define INFINITY POS_INFINITY
#define NEG_INFINITY __int_as_float(0xff800000)
#define NAN __int_as_float(0x7fffffff)
)ESCAPE") +
at::cuda::get_traits_string() + at::cuda::get_complex_body_string() +
at::cuda::get_cmath_string() + at::cuda::get_complex_math_string();
return result;
}
} // namespace
std::string FusionExecutor::getStructuredCode(const std::string& kernel) {
// generating cuda code;
std::string code = "";
#ifdef USE_ROCM
#if ROCM_VERSION < 40200
code += std::string("#include <hip/hip_runtime.h>\n") +
std::string("#include <hip/hip_bf16.h>\n") +
std::string("#include <hip/hip_fp16.h>\n");
#endif
code += std::string("#pragma clang force_cuda_host_device begin\n");
#endif
code += std::string("namespace ") + FusionExecutor::kernelNamespace() +
" {\n" + defineIntegerTypes() + defineIndexMode(options_.index_mode) +
defineComplexTypes() + executor_utils::kernelPreamble() + kernel + "}\n";
#ifdef USE_ROCM
code += std::string("#pragma clang force_cuda_host_device end\n");
#endif
if (isDebugDumpEnabled(DebugDumpOption::CudaKernel)) {
std::cout << "\n======= Codegen output for kernel: " << kernelName()
<< " =======\n\n"
<< kernel << "\n======================================\n\n";
} else if (isDebugDumpEnabled(DebugDumpOption::CudaFull)) {
std::cout << "\n======= Codegen output for kernel: " << kernelName()
<< " =======\n\n"
<< code << "\n======================================\n\n";
}
if (isDebugDumpEnabled(DebugDumpOption::CudaToFile) ||
isDebugDumpEnabled(DebugDumpOption::DebugInfo)) {
std::stringstream file_name;
file_name << "__tmp_kernel" << fusion_id_ << ".cu";
std::cout << "PRINTING: " << file_name.str() << std::endl;
std::ofstream out(file_name.str());
out << code << std::endl;
out.close();
}
return code;
}
// TODO: come up with a more user friendly interface
void FusionExecutor::debugCompileFusionFromStr(
Fusion* fusion,
const std::string& code,
const std::string& name,
int id,
CompileOptions options) {
options_ = options;
if (isDebugDumpEnabled(DebugDumpOption::FusionIr)) {
fusion->print();
} else if (isDebugDumpEnabled(DebugDumpOption::FusionIrMath)) {
fusion->printMath();
}
if (isDebugDumpEnabled(DebugDumpOption::CudaFull)) {
std::cout << "\n==== codegen output for kernel: " << kernelName()
<< " ====" << std::endl
<< code << std::endl
<< "======================================\n"
<< std::endl;
}
lowered_ = std::make_unique<GpuLower>(fusion);
const auto kernel = lowered_->kernel();
fusion_ = lowered_->kernel();
fusion_id_ = id;
setUsedTVs();
if (isDebugDumpEnabled(DebugDumpOption::KernelIr)) {
kernel->print();
}
const auto& kernel_summary = kernel->summary();
if (!kernel_summary.static_smem_allocations.empty()) {
kir::ExpressionEvaluator static_evaluator;
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
const auto static_smem_size = computeSharedMemory(
static_evaluator, kernel_summary.static_smem_allocations);
TORCH_INTERNAL_ASSERT(
static_smem_size < max_static_smem_,
"The static shared memory allocation is larger than available memory.");
}
std::tie(compiled_kernel_, last_compiler_log_) =
executor_utils::nvrtcCompile(code, name, fusion_id_);
TORCH_INTERNAL_ASSERT(
fusion_id_ > 0, "assign a fusion_id_ <= 0 is not accepted.");
}
void FusionExecutor::compileFusion(
Fusion* fusion,
const KernelArgumentHolder& args,
const LaunchParams& launch_constraints) {
FUSER_PERF_SCOPE("compileFusion");
TORCH_INTERNAL_ASSERT(
!fusion->outputs().empty(), "No output found for this kernel, aborting.");
for (auto out : fusion->outputs()) {
TORCH_INTERNAL_ASSERT(
out->getValType() == ValType::TensorView,
"Output types from fusions that are not tensors are not supported at this point.");
const auto maybe_rfactor_domain =
out->as<TensorView>()->getMaybeRFactorDomain();
// walking through outputs to see if output shapes are dependent on
// non-tensor inputs. For which case, we should have disabled output
// allocation, since the caching id only looks at tensor shapes.
// See issue https://github.com/csarofeen/pytorch/issues/2002
std::vector<Val*> output_extents;
for (const auto id : maybe_rfactor_domain) {
Val* extent = nullptr;
if (id->isReduction() || id->isStride()) {
continue;
} else if (id->isBroadcast() && id->hasExpandedExtent()) {
extent = id->expandedExtent();
} else {
extent = id->extent();
}
output_extents.emplace_back(extent);
}
auto dependencies = InputsOf::outputs(fusion, output_extents);
if (std::any_of(dependencies.begin(), dependencies.end(), [](Val* val) {
return val->isFusionInput();
})) {
// TODO: parameter cache is too big a hammer here. We should consider
// separate the caching logic of output sizes & launch params. Since
// output size dependency should only invalidate the output sizes
disable_parameter_cache_ = true;
break;
}
}
if (isDebugDumpEnabled(DebugDumpOption::FusionIr)) {
fusion->print();
} else if (isDebugDumpEnabled(DebugDumpOption::FusionIrMath)) {
fusion->printMath();
}
// TODO: refactor the options_ passed through
options_.device = c10::Device(c10::DeviceType::CUDA, args.getDeviceIndex());
options_.index_mode = args.getIndexMode();
c10::DeviceGuard dg(options_.device);
TORCH_INTERNAL_ASSERT(
options_.device.is_cuda(), "Provided device to CUDA fuser is the CPU.");
auto properties = at::cuda::getDeviceProperties(options_.device.index());
configured_device_smem_ = properties->sharedMemPerBlock;
#ifndef USE_ROCM
device_smem_limit_ = properties->sharedMemPerBlockOptin;
#else
// don't know if rocm supports opt-in shared memroy reconfiguration
device_smem_limit_ = properties->sharedMemPerBlock;
#endif
warp_size_ = properties->warpSize;
lowered_ = std::make_unique<GpuLower>(
fusion,
options_.index_mode == KernelIndexMode::INT64 ? DataType::Int
: DataType::Int32);
const auto kernel = lowered_->kernel();
fusion_ = lowered_->kernel()->as<Fusion>();
fusion_id_ = ++fusion_id_counter_;
setUsedTVs();
if (isDebugDumpEnabled(DebugDumpOption::KernelIr)) {
kernel->print();
}
kernel_code_ = codegen::generateCudaKernel(kernel, kernelName());
const auto structured_code = getStructuredCode(kernel_code_);
const auto& kernel_summary = kernel->summary();
// We currently shouldn't allocate any more shared mem
// tensors statically but could keep this path if
// needed in later development.
if (!kernel_summary.static_smem_allocations.empty()) {
kir::ExpressionEvaluator static_evaluator;
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
const auto static_smem_size = computeSharedMemory(
static_evaluator, kernel_summary.static_smem_allocations);
TORCH_INTERNAL_ASSERT(
static_smem_size < max_static_smem_,
"The static shared memory allocation is larger than available memory.");
}
if (kernel_summary.has_dynamic_local_memory_allocations) {
std::stringstream ss;
ss << "Allocations must be based on constant integers for local memory. However, found: ";
for (auto alloc : kernel_summary.dynamic_lmem_allocations) {
ss << alloc->buffer()->toString() << ", ";
}
ss << " have dynamic allocations but are placed in local memory.";
TORCH_INTERNAL_ASSERT(false, ss.str());
}
// TODO: pass block_size here;
c10::optional<int> block_size = c10::nullopt;
if (!args.empty()) {
auto expr_eval = executor_utils::bindKernelInputs(args, kernel);
auto launch_params =
computeLaunchParams(launch_constraints, expr_eval, warp_size_);
block_size = launch_params.nThreads();
TORCH_INTERNAL_ASSERT(
block_size > 0, "launch param inferred block size < 0");
}
// TODO: high water mark should be computed via occupancy API after
// compilation.
// Basically setting high water martk as 1 when we don't provide args for
// compilation, it will just generate a kernel that gets ditched at the first
// run - not great. We should have better heuristics.
block_size_high_water_mark = std::max<int64_t>(
(block_size.has_value() ? block_size.value() : 1),
block_size_high_water_mark);
std::tie(compiled_kernel_, last_compiler_log_) = executor_utils::nvrtcCompile(
structured_code,
(kernelNamespace() + "::" + kernelName()).c_str(),
fusion_id_,
block_size);
TORCH_INTERNAL_ASSERT(
fusion_id_ > 0, "failed to assign a fusion_id_ after compilation.");
#ifndef USE_ROCM
// The driver API call requires an int argument.
int max_dynamic_smem = 0;
AT_CUDA_DRIVER_CHECK(at::globalContext().getNVRTC().cuFuncGetAttribute(
&max_dynamic_smem,
CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
compiled_kernel_.function));
maybe_available_dynamic_smem_ = max_dynamic_smem;
#endif
}
namespace {
at::Tensor inferAndAlloc(
const TensorView* tv,
const std::vector<Val*>& sizes,
kir::ExpressionEvaluator& expr_eval,
// Map from dim -> expanded size of TV if any expanded broadcast dimensions
// exist
std::unordered_map<int, Val*> expanded_map,
const CompileOptions& options,
bool zero_init = false) {
FUSER_PERF_SCOPE("inferAndAlloc");
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
// Going to infer all the sizes of the TensorView
std::vector<int64_t> inferred_sizes;
// Expanded sizes is at maximum the same size of inferred_sizes, as you could
// have a fully broadcasted tensor that's being expanded
std::vector<int64_t> expanded_sizes;
bool expanded_dim = false;
for (const auto size : sizes) {
const auto inferred_val = expr_eval.evaluate(size);
TORCH_INTERNAL_ASSERT(
inferred_val.has_value(),
"Could not launch kernel as program could not infer ",
size->toString(),
"(",
size->name(),
") for the buffer ",
tv->toString());
inferred_sizes.push_back(inferred_val->as<int64_t>());
if (expanded_map.count(expanded_sizes.size())) {
auto expanded_size = expanded_map.at(expanded_sizes.size());
const auto inferred_expanded_size = expr_eval.evaluate(expanded_size);
TORCH_INTERNAL_ASSERT(
inferred_expanded_size.has_value(),
"Could not launch kernel as program could not infer the expanded extent ",
expanded_size->toString(),
"(",
expanded_size->name(),
") for the buffer ",
tv->toString());
if (inferred_val.value() != 1) {
TORCH_INTERNAL_ASSERT(
inferred_val.value() == inferred_expanded_size.value(),
"Attempted an expand on a non-broadcasted dimension,",
" but the expand doesn't match the dimensions size.");
} else {
expanded_dim = true;
}
expanded_sizes.push_back(inferred_expanded_size->as<int64_t>());
} else {
expanded_sizes.push_back(inferred_val->as<int64_t>());
}
}
const auto at_type = data_type_to_aten(tv->dtype());
const auto tensor_options =
at::TensorOptions().dtype(at_type).device(options.device);
c10::IntArrayRef isizes(inferred_sizes);
if (zero_init) {
auto zeros = at::zeros(isizes, tensor_options);
if (expanded_dim) {
return zeros.expand(expanded_sizes);
}
return zeros;
} else {
// Non Variable type guard for empty_cuda call
at::AutoDispatchBelowADInplaceOrView non_variable_type_mode;
auto empty = at::empty(isizes, tensor_options);
if (expanded_dim) {
return empty.expand(expanded_sizes);
}
return empty;
}
}
at::Tensor inferAndAllocOutput(
const TensorView* tv,
kir::ExpressionEvaluator& expr_eval,
const CompileOptions& options,
bool zero_init = false) {
const auto domain = tv->domain();
const auto maybe_rfactor_domain = domain->hasRFactor()
? domain->getRFactorDomain()
: domain->getRootDomain();
std::vector<Val*> sizes;
std::unordered_map<int, Val*> expand_map;
for (const auto id : maybe_rfactor_domain) {
if (id->isReduction() || id->isStride()) {
continue;
}
sizes.push_back(id->extent());
if (id->isBroadcast() && id->hasExpandedExtent()) {
expand_map[sizes.size() - 1] = id->expandedExtent();
}
}
return inferAndAlloc(tv, sizes, expr_eval, expand_map, options, zero_init);
}
} // namespace
uint64_t FusionExecutor::computeSharedMemory(
kir::ExpressionEvaluator& expr_eval,
const std::vector<const kir::Allocate*>& buffers,
bool align_padding,
uint64_t total) {
FUSER_PERF_SCOPE("computeSharedMemory");
for (auto smem_alloc : buffers) {
// If this buffer aliases another buffer,
// then do not allocate memory for this buffer.
if (smem_alloc->alias() == nullptr) {
const auto inferred_val = expr_eval.evaluate(smem_alloc->size());
if (inferred_val.has_value()) {
const uint64_t data_size = dataTypeSize(smem_alloc->buffer()->dtype());
// Add padding to align dynamic shared memory
if (align_padding) {
#ifndef USE_ROCM
const int align_size = 16; // always align to 16B/128b.
#else
const int align_size = 8; // see codegen.cpp for HIP
#endif
total = ceilDiv(total, align_size) * align_size;
}
total += inferred_val->as<int64_t>() * data_size;
} else {
TORCH_INTERNAL_ASSERT(
false,
"Failed to evaluate the size ",
smem_alloc->size(),
" of shared memory buffer - T",
smem_alloc->buffer()->name());
}
}
}
return total;
}
LaunchParams FusionExecutor::computeLaunchParams(
const LaunchParams& launch_constraints,
kir::ExpressionEvaluator& expr_eval,
const int warp_size) {
FUSER_PERF_SCOPE("FusionExecutor::ComputeLaunchParams");
TORCH_INTERNAL_ASSERT(warp_size > 0, "WARP_SIZE should be larger than 0");
LaunchParams launch_params;
auto data_cache = compileTimeDataCache();
auto lower = lowered_.get();
auto& used_tvs = getUsedTVs();
auto parallel_binding_ids_entry =
executor_utils::caching::ExecutorCompileTimeEntry<
executor_utils::caching::ParallelBindingIterDomains>(
data_cache, [&used_tvs, &lower]() {
return std::make_unique<std::vector<IterDomain*>>(
executor_utils::getParallelBindingsIterDomains(
lower, used_tvs));
});
auto& parallel_binding_ids = parallel_binding_ids_entry.get();
auto parallel_iter_extent_entry =
executor_utils::caching::ExecutorCompileTimeEntry<
executor_utils::caching::ParallelIterExtentMap>(
data_cache, [&parallel_binding_ids]() {
return executor_utils::getParallelIterExtents(parallel_binding_ids);
});
auto& parallel_iter_extents = parallel_iter_extent_entry.get();
auto simplified_parallel_iter_extent_entry =
executor_utils::caching::ExecutorCompileTimeEntry<
executor_utils::caching::SimplifiedParallelIterExtentMap>(
data_cache, [&parallel_binding_ids, &lower]() {
return executor_utils::getSimplifiedParallelIterExtents(
lower, parallel_binding_ids);
});
auto& simplified_parallel_iter_extents =
simplified_parallel_iter_extent_entry.get();
auto warp_padded_parallel_entry =
executor_utils::caching::ExecutorCompileTimeEntry<
executor_utils::caching::WarpPaddedParallelExtents>(
data_cache, [&parallel_binding_ids, &lower]() {
return executor_utils::getWarpPaddedExtentsInfo(
lower->kernel(), parallel_binding_ids);
});
auto& warp_padded_extent_set =
warp_padded_parallel_entry.get().warp_padded_extent_set;
auto& warp_padded_constant =
warp_padded_parallel_entry.get().warp_padded_constant;
// TODO: Need to redesign this part a bit to
// find the right place to trigger evaluate
if (expr_eval.precomputedValues()) {
expr_eval.precomputedValues()->bindParallelExtents(
parallel_iter_extents, launch_constraints);
expr_eval.precomputedValues()->evaluate();
}
// If any dimension was set in launch constraints we need to run through
// IterDomains that have been parallelized, and bind those values. Or make
// sure if they could be inferred the inference matches what was set.
for (auto& entry : parallel_iter_extents) {
auto p_type = entry.first;
if (launch_constraints.hasDim(p_type)) {
auto parallel_extents = entry.second;
for (auto extent : parallel_extents) {
auto inferred_val = expr_eval.evaluate(extent);
if (inferred_val.has_value()) {
// This value could have been inferred, make sure it was set right.
bool valid =
inferred_val.value() == launch_constraints.getDim(p_type) ||
launch_constraints.getRawVal(p_type) == -1;
if (!useFallback() && !valid) {
TORCH_WARN_ONCE(
"Cannot validate parallelization scheme, "
"this may be due to mixed broadcast axes that are parallelized.");
}
} else if (!expr_eval.precomputedValues()) {
expr_eval.bind(extent, launch_constraints.getDim(p_type));
}
if (!launch_params.hasDim(p_type)) {
// Bind the launch constraint into our evaluation context
launch_params.bind(launch_constraints.getDim(p_type), p_type);
// Makes sure the p-types bound to evaluators are the
// final values that will become the actual launch
// param size to ensure accurate smem buffer size
// computation.
expr_eval.bind(p_type, launch_constraints.getDim(p_type));
}
}
}
}
// Run through the rest of the parallel IterDomains and infer their size
for (auto& entry : simplified_parallel_iter_extents) {
FUSER_PERF_SCOPE("FusionExecutor::ParallelBindingResolution");
auto p_type = entry.first;
auto parallel_extents = entry.second;
// Select the maxmimum value out of all the parallel extents
int64_t maximum_value = std::numeric_limits<int64_t>::min();
for (auto extent : parallel_extents) {
auto val = expr_eval.evaluate(extent);
TORCH_INTERNAL_ASSERT(
val.has_value(),
"Tried to evaluate the extent, ",
extent->toInlineString(),
" for the ptype: ",
p_type,
" to set launch bounds but could not.");
// apply padding to the extent if needed
if (warp_padded_extent_set.count(extent)) {
// Check if the extent has const value
auto padded_constant_it = warp_padded_constant.find(extent);
if (padded_constant_it != warp_padded_constant.end()) {
// If already specified padded to constant, need to check
// runtime value not over the constant bound
TORCH_INTERNAL_ASSERT(*val <= padded_constant_it->second);
*val = padded_constant_it->second;
} else {
// If no specified constant, pad to the smallest multiple of warp
// above the value.
auto padded_number_of_warps = (*val + warp_size - 1) / warp_size;
*val = warp_size * padded_number_of_warps;
}
TORCH_INTERNAL_ASSERT(
*val <= 1024, "padded dimension larger than max block size");
}
maximum_value = std::max(maximum_value, val->as<int64_t>());
}
// Protect for size-0 tensors, they still have a value so would prefer to
// bind nothing than 0
if (maximum_value > 0) {
expr_eval.bind(p_type, maximum_value);
launch_params.bind(maximum_value, p_type);
}
}
// Re-run the integer machine with all
// the thread sizes now determined.
if (expr_eval.precomputedValues()) {
expr_eval.precomputedValues()->evaluate();
}
const auto kernel = lowered_->kernel();
const auto& kernel_summary = kernel->summary();
// Calculate Dynamic Shared Memory Size
// Add workspace for reduction and broadcast
uint64_t reduction_broadcast_workspace = 0;
const bool has_workspace = kernel_summary.has_block_reductions ||
kernel_summary.has_grid_reductions ||
kernel_summary.has_block_broadcasts || kernel_summary.has_grid_broadcasts;
if (has_workspace &&
kernel_summary.largest_smem_data_type != DataType::Null) {
// Not using nThreads here since it does not handle uninitialized value
// TODO: here is an optimization opportunity since welford uses int64_t for
// N while the data type is not neccessarily double. But it may need more
// work on the alignment
const int welford_factor =
kernel_summary.has_block_welford || kernel_summary.has_grid_welford ? 3
: 1;
reduction_broadcast_workspace =
dataTypeSize(kernel_summary.largest_smem_data_type) * welford_factor *
launch_params.bdimx() * launch_params.bdimy() * launch_params.bdimz();
}
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
const uint64_t dynamic_smem_size = computeSharedMemory(
expr_eval,
kernel_summary.dynamic_smem_allocations,
true,
reduction_broadcast_workspace);
// Check that requested smem size can be dynamically allocated.
// This check is only done once a kernel has been compiled, since
// maybe_available_dynamic_smem_ needs to be evaluated on
// a compiled kernel.
if (maybe_available_dynamic_smem_.has_value()) {
// Dynamic shared memory space that we can allocate without
// carving more space from L1.
const uint64_t available_dynamic_smem_without_reconfiguration =
maybe_available_dynamic_smem_.value();
// Maximum additional shared memory size we could request
// if we do re-configuration.
const uint64_t additional_dynamic_smem_available_through_reconfiguration =
device_smem_limit_ - configured_device_smem_;
TORCH_INTERNAL_ASSERT(
(dynamic_smem_size) <
(available_dynamic_smem_without_reconfiguration +
additional_dynamic_smem_available_through_reconfiguration),
"The total shared memory allocation is larger than available memory.",
" Dynamic size: ",
dynamic_smem_size,
". Available size: ",
maybe_available_dynamic_smem_.value(),
". Configured smem size: ",
configured_device_smem_,
". Device limit size: ",
device_smem_limit_);
}
launch_params.setSmem(dynamic_smem_size);
return launch_params;
}
FusionExecutor::GlobalBuffers FusionExecutor::allocGlobalVals(
kir::ExpressionEvaluator& expr_eval) {
FUSER_PERF_SCOPE("FusionExecutor::AllocGlobalVals");
GlobalBuffers global_buffers;
const auto kernel = lowered_->kernel();
const auto& kernel_summary = kernel->summary();
for (auto alloc : kernel_summary.global_allocations) {
TORCH_INTERNAL_ASSERT(
alloc->buffer()->isA<TensorView>(),
"Cannot allocate global buffers that are not tensors.");
auto tv = alloc->buffer()->as<TensorView>();
if (tv->isFusionOutput()) {
continue;
}
if (alloc->zeroInit()) {
global_buffers.buffers.push_back(
inferAndAlloc(tv, alloc->shape(), expr_eval, {}, options_, true));
global_buffers.zero_init.push_back(true);
} else {
global_buffers.buffers.push_back(
inferAndAlloc(tv, alloc->shape(), expr_eval, {}, options_, false));
global_buffers.zero_init.push_back(false);
}
// Remember the tensor buffer used for storing kernel profile
if (isOptionEnabled(EnableOption::KernelProfile) &&
tv == kernel->profile().getBuffer()) {
global_buffers.profile_buffer = global_buffers.buffers.back();
}
}
return global_buffers;
}
std::vector<at::Tensor> FusionExecutor::allocOutputs(
kir::ExpressionEvaluator& expr_eval,
const std::unordered_set<int>& alias_indices) {
FUSER_PERF_SCOPE("FusionExecutor::AllocOutputs");
const auto kernel = lowered_->kernel();
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector<at::Tensor> outputs;
for (const auto out_i : c10::irange(kernel->outputs().size())) {
// TODO: FIX this short-cut where we trivially forward inputs to outputs
if (kernel->outputs()[out_i]->isFusionInput()) {
TORCH_INTERNAL_ASSERT(false, "trivial input forwarding NOT IMPLEMENTED");
// for (auto inp_i : c10::irange(kernel->inputs().size())) {
// if (kernel->inputs()[inp_i] == kernel->outputs()[out_i]) {
// TORCH_INTERNAL_ASSERT(
// inp_i < inputs.size(),
// "Issue with an input showing up as output, couldn't find
// input.");
// TORCH_INTERNAL_ASSERT(
// inputs[inp_i].isTensor(),
// "Cannot register a scalar as an output in a fusion.");
// outputs.push_back(inputs[inp_i].toTensor());
// break;
// }
// }
} else {
TORCH_INTERNAL_ASSERT(
kernel->outputs()[out_i]->isA<TensorView>(),
"Cannot allocate outputs that are not tensors.");
auto output = kernel->outputs()[out_i]->as<TensorView>();
if (alias_indices.count(out_i) != 0) {
// aliasing to inputs, no need to allocate real output, just push empty
// tensor here.
outputs.emplace_back();
} else {
outputs.push_back(
inferAndAllocOutput(output, expr_eval, options_, false));
}
}
}
return outputs;
}
void FusionExecutor::setUsedTVs() {
auto used_vals = fusion_->usedMathVals();
auto used_tvs = ir_utils::filterByType<TensorView>(used_vals);
used_tvs_.clear();
used_tvs_.insert(used_tvs_.begin(), used_tvs.begin(), used_tvs.end());
}
KernelArgumentHolder FusionExecutor::evaluateOutputSizes(
const KernelArgumentHolder& args,
kir::ExpressionEvaluator& expr_eval,
const std::unordered_set<int>& alias_indices) {
FUSER_PERF_SCOPE("FusionExecutor::AllocOutputs");
const auto kernel = lowered_->kernel();
KernelArgumentHolder ret(args.getIndexMode());
ret.setDeviceIndex(args.getDeviceIndex());
CompileOptions meta_options = options_;
meta_options.device = c10::Device(DeviceType::Meta, 0);
for (const auto out_i : c10::irange(kernel->outputs().size())) {
// If the output is just trivially the input, just "copy" it over.
if (kernel->outputs()[out_i]->isFusionInput()) {
for (auto inp_i : c10::irange(kernel->inputs().size())) {
if (kernel->inputs()[inp_i] == kernel->outputs()[out_i]) {
TORCH_INTERNAL_ASSERT(
inp_i < args.size(),
"Issue with an input showing up as output, couldn't find input.");
auto tensor_arg_abstract =
dynamic_cast<const TensorArgAbstract*>(args[inp_i]);
TORCH_INTERNAL_ASSERT(
tensor_arg_abstract,
"Cannot register a scalar as an output in a fusion.");
ret.push(tensor_arg_abstract);
break;
}
}
} else {
TORCH_INTERNAL_ASSERT(
kernel->outputs()[out_i]->isA<TensorView>(),
"Cannot allocate outputs that are not tensors.");
auto output = kernel->outputs()[out_i]->as<TensorView>();
if (alias_indices.count(out_i) != 0) {
// aliasing to inputs, no need to allocate real output
// but we still need to push an entry here.
ret.push(int64_t(0));
} else {
// TODO: we are using meta here, which is bad since it doesn't account
// for devices. Switch to fake tensor instead
ret.push(inferAndAllocOutput(output, expr_eval, meta_options, false));
}
}
}
return ret;
}
KernelArgumentHolder FusionExecutor::inferOutputSizes(
const KernelArgumentHolder& args,
const LaunchParams& launch_constraints) {
FUSER_PERF_SCOPE("FusionExecutor::RunFusion");
ExecutorEntry* executor_entry = nullptr;
c10::optional<size_t> opt_code = args.getCacheId();
if (opt_code.has_value()) {
executor_entry = &executor_entry_lookup_[*opt_code];
}
executor_utils::initializeCudaContext();
TORCH_INTERNAL_ASSERT(lowered_);
TORCH_INTERNAL_ASSERT(
!executor_entry || !executor_entry->init,
"compile kernel shouldn't hit a pre-existing cache");
FUSER_PERF_SCOPE("ExecutorRunFusion::ValidateAndInitialize");
// TODO: validate kernel inputs currently won't be happy, since our fusion
// args are mapped with `meta` tensor instead of `cuda` tensor, check if this
// would be resolved with FakeTensor
// executor_utils::validateKernelInputs(fusion_, args, options_.device);
if (!evaluator_precomputed_values_) {
evaluator_precomputed_values_ =
std::make_unique<KernelPrecomputedValues>(lowered_->kernel());
}
kir::ExpressionEvaluator expr_eval;
evaluator_precomputed_values_->bindKernelInputs(lowered_->kernel(), args);
expr_eval.precomputedValues() = evaluator_precomputed_values_.get();
// I think this binds something to expr_eval, so even though we are not using
// launch_params_, we still need this in order to infer output shapes.
launch_params_ =
computeLaunchParams(launch_constraints, expr_eval, warp_size_);
executor_utils::validateVectorizedTensors(
lowered_.get()->kernel(), args, {}, compileTimeDataCache(), expr_eval);
auto alias_indices_entry = executor_utils::caching::ExecutorCompileTimeEntry<
executor_utils::caching::InputAliasIndices>(
compileTimeDataCache(), [&]() {
return std::make_unique<std::vector<std::pair<int, int>>>(
fusion_->getInputAliasIndices());
});
auto& alias_indices = alias_indices_entry.get();
// NOLINTNEXTLINE(bugprone-branch-clone)
auto output_alias_indices_entry =
executor_utils::caching::ExecutorCompileTimeEntry<
executor_utils::caching::OutputAliasIndices>(
compileTimeDataCache(), [&]() {
return std::make_unique<std::unordered_set<int>>(
fusion_->getOutputAliasIndices());
});
auto& output_alias_indices = output_alias_indices_entry.get();
auto ret = evaluateOutputSizes(args, expr_eval, output_alias_indices);
for (const auto& entry : alias_indices) {
auto aliased_output_index = entry.first;
auto aliased_input_index = entry.second;
TORCH_INTERNAL_ASSERT(
args[aliased_input_index]->isType(ArgType::Tensor),
"alias io only supports tensor");
ret.swap(aliased_output_index, args[aliased_input_index]);
}
return ret;
}
std::vector<at::Tensor> FusionExecutor::runFusion(
KernelArgumentHolder& args,
const LaunchParams& launch_constraints,
const std::vector<at::Tensor>& outputs) {
FUSER_PERF_SCOPE("FusionExecutor::RunFusion");
TORCH_INTERNAL_ASSERT(compiled());
TORCH_INTERNAL_ASSERT(
fusion_id_ > 0, "Cannot run fusion, it was not compiled.");
TORCH_INTERNAL_ASSERT(
!args.getCacheId().has_value() || outputs.empty(),
"short cut input cache is not compatible with pre-allocated output");
if (isDebugDumpEnabled(DebugDumpOption::FusionArgs)) {
std::cout << "Arguments for fusion" << fusion_id_ << ":" << std::endl
<< "Inputs:" << std::endl;
for (auto i : c10::irange(args.size())) {
args[i]->print();
}
std::cout << "Outputs:" << std::endl;
for (const auto& output : outputs) {
std::cout << " " << output.scalar_type() << " " << output.sizes()
<< " (strides = " << output.strides() << ")" << std::endl;
}
std::cout << launch_constraints.toString();
}
ExecutorEntry* executor_entry = nullptr;
if (args.getCacheId().has_value()) {
executor_entry = &executor_entry_lookup_[*args.getCacheId()];
}
c10::DeviceGuard dg(options_.device);
auto stream = at::cuda::getCurrentCUDAStream();
executor_utils::initializeCudaContext();
TORCH_INTERNAL_ASSERT(lowered_);
launch_params_ = LaunchParams();
// NOLINTNEXTLINE(cppcoreguidelines-init-variables)
std::vector<at::Tensor> allocated_outputs;
GlobalBuffers global_buffers;
uint64_t rand_offset = 0;
if (executor_entry && executor_entry->init && !disable_parameter_cache_) {
{
// context manager to disable auto grad for `empty_cuda` calls later
at::AutoDispatchBelowADInplaceOrView non_variable_type_mode;
// take the short-cut for launch if we see a recorded input set again
launch_params_ = executor_entry->launch_params;
// only allocate outputs when not given
if (outputs.empty()) {
FUSER_PERF_SCOPE("ExecutorRunFusion::OutputAlloc");
for (const auto i : c10::irange(executor_entry->output_sizes.size())) {
allocated_outputs.push_back(at::native::empty_strided_cuda(
executor_entry->output_sizes[i],
executor_entry->output_strides[i],
executor_entry->output_types[i],
c10::nullopt,
options_.device,
c10::nullopt));
}
// Note: aliased output is not returned as output. But we still need it
// for kernel execution, so would need to push them to args
for (const auto& entry : executor_entry->io_alias_indices) {
auto aliased_output_index = entry.first;
auto aliased_input_index = entry.second;
auto tensor_arg_abstract =
dynamic_cast<const TensorArgAbstract*>(args[aliased_input_index]);
TORCH_INTERNAL_ASSERT(
tensor_arg_abstract, "alias io only supports tensor");
allocated_outputs[aliased_output_index] =
tensor_arg_abstract->getTensor();
}
args.push(allocated_outputs);
} else {
TORCH_INTERNAL_ASSERT(
outputs.size() == fusion_->outputs().size(),
__func__,
" provided number of outputs does match fusion output");
allocated_outputs = outputs;
args.push(outputs);
}
{
FUSER_PERF_SCOPE("ExecutorRunFusion::IntermediateBufferAlloc");
for (const auto i : c10::irange(executor_entry->buffer_sizes.size())) {
if (executor_entry->buffer_zero_init[i]) {
global_buffers.buffers.push_back(at::zeros(
executor_entry->buffer_sizes[i],
at::TensorOptions()
.dtype(executor_entry->buffer_types[i])
.device(options_.device)));
global_buffers.zero_init.push_back(true);
} else {
global_buffers.buffers.push_back(at::native::empty_cuda(
executor_entry->buffer_sizes[i],
executor_entry->buffer_types[i],
c10::nullopt,
options_.device,
c10::nullopt));
global_buffers.zero_init.push_back(false);
}
}
}
}
rand_offset = executor_entry->rand_offset;
} else {
FUSER_PERF_SCOPE("ExecutorRunFusion::ValidateAndInitialize");
// code path to take when either:
// 1. no opt_code is provided or
// 2. `executor_entry` is not initialized
executor_utils::validateKernelInputs(fusion_, args, options_.device);
if (!evaluator_precomputed_values_) {
evaluator_precomputed_values_ =
std::make_unique<KernelPrecomputedValues>(lowered_->kernel());
}
kir::ExpressionEvaluator expr_eval;
evaluator_precomputed_values_->bindKernelInputs(lowered_->kernel(), args);
expr_eval.precomputedValues() = evaluator_precomputed_values_.get();
launch_params_ =
computeLaunchParams(launch_constraints, expr_eval, warp_size_);
// Recompile the kernel if the number of threads in the block has increased
if (launch_params_.nThreads() > block_size_high_water_mark) {
const auto kernel = lowered_->kernel();
kernel_code_ = codegen::generateCudaKernel(kernel, kernelName());
const auto structured_code = getStructuredCode(kernel_code_);
block_size_high_water_mark = launch_params_.nThreads();
std::tie(compiled_kernel_, last_compiler_log_) =
executor_utils::nvrtcCompile(
structured_code,
(kernelNamespace() + "::" + kernelName()).c_str(),
fusion_id_,
block_size_high_water_mark);
}
if (kernel()->summary().has_cooperative_grid_reduction) {
#ifndef USE_ROCM
int num_blocks_per_SM = -1;
at::globalContext().getNVRTC().cuOccupancyMaxActiveBlocksPerMultiprocessor(
&num_blocks_per_SM,
compiled_kernel_.function,
(int)(launch_params_.bdimx() * launch_params_.bdimy() * launch_params_.bdimz()),
(size_t)launch_params_.smem());
TORCH_INTERNAL_ASSERT(
(int64_t)(
num_blocks_per_SM *
at::cuda::getDeviceProperties(options_.device.index())
->multiProcessorCount) >= launch_params_.gdimx() *
launch_params_.gdimy() * launch_params_.gdimz(),
"Wanted to launch a cooperative kernel, however the number of blocks is greater than ",
"what can be resident on the GPU at once. Need: ",
launch_params_.gdimx() * launch_params_.gdimy() *
launch_params_.gdimz(),
" (",
launch_params_.gdimx(),
" * ",
launch_params_.gdimy(),
" * ",
launch_params_.gdimz(),
") but limited to ",
num_blocks_per_SM,
" * ",
at::cuda::getDeviceProperties(options_.device.index())
->multiProcessorCount);
#else
TORCH_INTERNAL_ASSERT(
false, "Cross grid communication not supported with HIP.");
#endif
}
executor_utils::validateVectorizedTensors(
lowered_.get()->kernel(),
args,
outputs,
compileTimeDataCache(),
expr_eval);
auto alias_indices_entry =
executor_utils::caching::ExecutorCompileTimeEntry<
executor_utils::caching::InputAliasIndices>(
compileTimeDataCache(), [&]() {
return std::make_unique<std::vector<std::pair<int, int>>>(
fusion_->getInputAliasIndices());
});
auto& alias_indices = alias_indices_entry.get();
// NOLINTNEXTLINE(bugprone-branch-clone)
if (outputs.empty()) {
auto output_alias_indices_entry =
executor_utils::caching::ExecutorCompileTimeEntry<
executor_utils::caching::OutputAliasIndices>(
compileTimeDataCache(), [&]() {
return std::make_unique<std::unordered_set<int>>(
fusion_->getOutputAliasIndices());
});
auto& output_alias_indices = output_alias_indices_entry.get();
allocated_outputs = allocOutputs(expr_eval, output_alias_indices);
for (const auto& entry : alias_indices) {
auto aliased_output_index = entry.first;
auto aliased_input_index = entry.second;
auto tensor_arg_abstract =
dynamic_cast<const TensorArgAbstract*>(args[aliased_input_index]);
TORCH_INTERNAL_ASSERT(
tensor_arg_abstract, "alias io only supports tensor");
allocated_outputs[aliased_output_index] =
tensor_arg_abstract->getTensor();
}
args.push(allocated_outputs);
} else {
allocated_outputs = outputs;
args.push(outputs);
executor_utils::validateKernelOutputs(
fusion_, allocated_outputs, options_.device);
}
global_buffers = allocGlobalVals(expr_eval);
if (kernel()->summary().max_rng_offsets >= 0) {
// NOTE: this is how we map offset to PW kernels in order to have
// identical random number generator to match native PyTorch results.
// But it doesn't really work as it takes assumption how threads are
// binded but is not generally how we handle that in scheduler.
// Refer to `Philox` in generated kernel to understand how the mapping
// works.
rand_offset = (kernel()->summary().max_rng_offsets + 1) * 4;
}
// This is the entry when we have provided `opt_code` but the entry has not
// been initialized yet.
if (executor_entry) {
FUSER_PERF_SCOPE("ExecutorRunFusion::FillCacheEntry");
// record the the short-cut executor entry for the given input set;
executor_entry->launch_params = launch_params_;
executor_entry->io_alias_indices = alias_indices;
for (const auto& output : allocated_outputs) {
executor_entry->output_sizes.push_back(output.sizes().vec());
executor_entry->output_strides.push_back(output.strides().vec());
executor_entry->output_types.push_back(output.scalar_type());
}
for (const auto& i : c10::irange(global_buffers.buffers.size())) {
executor_entry->buffer_sizes.push_back(
global_buffers.buffers[i].sizes().vec());
executor_entry->buffer_types.push_back(
global_buffers.buffers[i].scalar_type());
executor_entry->buffer_zero_init.push_back(global_buffers.zero_init[i]);
}
executor_entry->rand_offset = rand_offset;
executor_entry->init = true;
}
}
// push back global buffers
args.push(global_buffers.buffers);
// push back RNG state if needed
if (lowered_->kernel()->summary().max_rng_offsets >= 0) {
args.appendPhiloxRNGSeed(rand_offset);
}
if (isDebugDumpEnabled(DebugDumpOption::LaunchParam)) {
launch_params_.print();
}
if (isDebugDumpEnabled(DebugDumpOption::KernelArgs)) {
std::cout << "Arguments for kernel" << fusion_id_ << ":" << std::endl
<< "Inputs:" << std::endl;
for (auto i : c10::irange(args.size())) {
args[i]->print();
}
std::cout << "Outputs:" << std::endl;
// note: add aliased outputs here.
for (const auto& output : allocated_outputs) {
std::cout << " " << output.scalar_type() << " " << output.sizes()
<< " (strides = " << output.strides()
<< ", address = " << output.data_ptr() << ")" << std::endl;
}
std::cout << "Reduction and semaphore buffers:" << std::endl;
TORCH_INTERNAL_ASSERT(
global_buffers.buffers.size() == global_buffers.zero_init.size(),
"global_buffer buffer & zero_init container should have identical sizes");
for (const auto i : c10::irange(global_buffers.buffers.size())) {
const auto& buffer = global_buffers.buffers[i];
const auto& zero_init = global_buffers.zero_init[i];
std::cout << " " << buffer.scalar_type() << " " << buffer.sizes()
<< " is_zero_initialized: " << zero_init << std::endl;
}
}
cudaEvent_t start_event = {};
cudaEvent_t finish_event = {};
if (measure_kernel_time_ ||
isDebugDumpEnabled(DebugDumpOption::EffectiveBandwidth) ||
isDebugDumpEnabled(DebugDumpOption::PerfDebugVerbose)) {
C10_CUDA_CHECK(cudaEventCreate(&start_event));
C10_CUDA_CHECK(cudaEventCreate(&finish_event));
C10_CUDA_CHECK(cudaEventRecord(start_event));
}
if (execute_kernel_) {
if (maybe_available_dynamic_smem_.has_value() &&
launch_params_.smem() > maybe_available_dynamic_smem_.value()) {
#ifndef USE_ROCM
// Increase limit of dynamic shared memory if needed.
AT_CUDA_DRIVER_CHECK(at::globalContext().getNVRTC().cuFuncSetAttribute(
compiled_kernel_.function,
CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
launch_params_.smem()));
#else
TORCH_INTERNAL_ASSERT(
false, "cuFuncSetAttribute not supported with HIP.");
#endif
}
if (!kernel()->summary().has_cooperative_grid_reduction) {
FUSER_PERF_SCOPE("ExecutorRunFusion::cuLaunchKernel");
AT_CUDA_DRIVER_CHECK(at::globalContext().getNVRTC().cuLaunchKernel(
compiled_kernel_.function,
launch_params_.gdimx(),
launch_params_.gdimy(),
launch_params_.gdimz(),
launch_params_.bdimx(),
launch_params_.bdimy(),
launch_params_.bdimz(),
launch_params_.smem(),
stream,
args.getBuffer(),
nullptr));
} else {
#ifndef USE_ROCM
FUSER_PERF_SCOPE("ExecutorRunFusion::cuLaunchCooperativeKernel");
AT_CUDA_DRIVER_CHECK(
at::globalContext().getNVRTC().cuLaunchCooperativeKernel(
compiled_kernel_.function,
launch_params_.gdimx(),
launch_params_.gdimy(),
launch_params_.gdimz(),
launch_params_.bdimx(),
launch_params_.bdimy(),
launch_params_.bdimz(),
launch_params_.smem(),
stream,
args.getBuffer()));
#else
TORCH_INTERNAL_ASSERT(
false, "Cross grid communication not supported with HIP.");
#endif
}
}
if (measure_kernel_time_ ||
isDebugDumpEnabled(DebugDumpOption::EffectiveBandwidth) ||
isDebugDumpEnabled(DebugDumpOption::PerfDebugVerbose)) {
C10_CUDA_CHECK(cudaEventRecord(finish_event));
C10_CUDA_CHECK(cudaEventSynchronize(start_event));
C10_CUDA_CHECK(cudaEventSynchronize(finish_event));
C10_CUDA_CHECK(
cudaEventElapsedTime(&kernel_time_ms_, start_event, finish_event));
C10_CUDA_CHECK(cudaEventDestroy(start_event));
C10_CUDA_CHECK(cudaEventDestroy(finish_event));
bytes_processed_ = 0;
// Figure how many bytes are inputs, outputs, and temporary buffers
for (auto i : c10::irange(args.size())) {
if (auto tensor_arg_abstract =
dynamic_cast<const TensorArgAbstract*>(args[i])) {
bytes_processed_ += tensor_arg_abstract->numel() *
dataTypeSize(tensor_arg_abstract->getDataType());
}
}
for (const auto& output : allocated_outputs) {
bytes_processed_ += output.numel() *
dataTypeSize(aten_to_data_type(output.scalar_type()));
}
if (isDebugDumpEnabled(DebugDumpOption::EffectiveBandwidth)) {
double gb_per_s =
((double)bytes_processed_ / ((double)kernel_time_ms_ / 1000)) /
(double)1.0e9;
std::cout << "kernel" << fusion_id_ << " run in " << kernel_time_ms_
<< " ms, achieved: " << gb_per_s << " GB/s" << std::endl;
}
}
if (isOptionEnabled(EnableOption::KernelProfile)) {
std::cout << kernel()->profile().toString(global_buffers.profile_buffer);
}
return allocated_outputs;
}
void FusionExecutor::compileRtc(
const std::string& code,
const std::string& name,
bool structured,
CompileOptions options) {
FUSER_PERF_SCOPE("ExecutorRunFusion::compileRtc");
std::string scode;
if (!structured) {
scode = getStructuredCode(code);
} else {
scode = code;
}
fusion_id_ = 1;
options_ = options;
std::tie(compiled_kernel_, last_compiler_log_) =
executor_utils::nvrtcCompile(scode, name, fusion_id_);
}
void FusionExecutor::runRtc(
const LaunchParams& launch_params,
const std::vector<at::Tensor>& args) {
FUSER_PERF_SCOPE("runFusion");
c10::DeviceGuard dg(options_.device);
auto stream = at::cuda::getCurrentCUDAStream();
KernelArgumentHolder kernel_arguments(options_.index_mode);
kernel_arguments.push(args);
AT_CUDA_DRIVER_CHECK(at::globalContext().getNVRTC().cuLaunchKernel(
compiled_kernel_.function,
launch_params.gdimx(),
launch_params.gdimy(),
launch_params.gdimz(),
launch_params.bdimx(),
launch_params.bdimy(),
launch_params.bdimz(),
launch_params.smem(),
stream,
kernel_arguments.getBuffer(),
nullptr));
}
} // namespace cuda
} // namespace fuser
} // namespace jit
} // namespace torch
Reference in a new issue View git blame Copy permalink