Skip to content
Merged

Aaa #230

Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion modules/ocarina/ast/function.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
//

#include "function.h"
#include "core/util/switch_profile.h"
#include "core/type_system/type_desc.h"
#include "core/util/util.h"
#include "generator/source_emitter.h"
Expand Down Expand Up @@ -45,6 +46,7 @@ ocarina::list<CallExpr::Template> Function::resolve_ast_templates(ocarina::list<
}

void Function::correct() noexcept {
switch_profile::Scope profile{"Function::correct", "ast_correct"};
TIMER(FunctionCorrect)
FunctionCorrector().apply(this);
}
Expand Down Expand Up @@ -622,4 +624,4 @@ uint64_t Function::compute_hash() const noexcept {
return ret;
}

}// namespace ocarina
}// namespace ocarina
4 changes: 3 additions & 1 deletion modules/ocarina/ast/function.h
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include "core/header.h"
#include "core/type.h"
#include "core/type_system/precision_policy.h"
#include "core/util/switch_profile.h"
#include "expression.h"
#include "statement.h"
#include "ast_node.h"
Expand Down Expand Up @@ -257,6 +258,7 @@ class OC_AST_API Function : public concepts::Noncopyable, public Hashable {
}
template<typename Func>
static shared_ptr<Function> define_kernel(Func &&func) noexcept {
switch_profile::Scope profile{"Function::define_kernel", "dsl"};
shared_ptr<Function> function = _define(Tag::KERNEL, std::forward<Func>(func));
function->correct();
return function;
Expand Down Expand Up @@ -339,4 +341,4 @@ class OC_AST_API Function : public concepts::Noncopyable, public Hashable {
[[nodiscard]] constexpr const Type *return_type() const noexcept { return ret_; }
constexpr void set_raytracing(bool val) const noexcept { raytracing_ = val; }
};
}// namespace ocarina
}// namespace ocarina
170 changes: 94 additions & 76 deletions modules/ocarina/backends/cuda/cuda_compiler.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
//

#include "cuda_compiler.h"
#include "core/util/switch_profile.h"
#include "cuda_device.h"
#include "ast/function.h"
#include "ast_to_cuda_source.h"
Expand Down Expand Up @@ -30,6 +31,7 @@ constexpr unsigned ast_to_cuda_source_cache_version = 1u;
}

[[nodiscard]] string emit_cuda_source(const Function &function) noexcept {
switch_profile::Scope profile{"emit_cuda_source", "codegen"};
switch (Env::shader_codegen_path()) {
case ShaderCodegenPath::EAstToSource: {
AstToCudaSource emitter{Env::code_obfuscation()};
Expand Down Expand Up @@ -66,89 +68,107 @@ CUDACompiler::CUDACompiler(CUDADevice *device)
: device_(device) {}

ocarina::string CUDACompiler::compile(const Function &function, int sm) const noexcept {

fs::path cuda_path = get_cuda_path();
cuda_path = cuda_path / "include";

string header_path = "-I" + cuda_path.string();
int ver_major = 0;
int ver_minor = 0;
OC_NVRTC_CHECK(nvrtcVersion(&ver_major, &ver_minor));
int nvrtc_version = ver_major * 10000 + ver_minor * 100;
const auto storage_policy = function.storage_policy();
const int use_float32 = storage_policy.policy == PrecisionPolicy::force_f16 ? 0 : 1;
auto nvrtc_option = fmt::format("-DLC_NVRTC_VERSION={}", nvrtc_version);
auto real_option = fmt::format("-DOCARINA_CUDA_USE_FLOAT32={}", use_float32);
std::vector header_names{"cuda_device_std.h", "cuda_device_scalar.h", "cuda_device_vector.h",
"cuda_device_matrix.h", "cuda_device_builtin.h", "cuda_tensor_builtin.h", "cuda_vector_func.h",
"cuda_device_math.h", "cuda_matrix_func.h", "cuda_device_resource.h"};
switch_profile::Scope profile{"CUDACompiler::compile", "compiler_setup"};
// These strings own the storage referenced by NVRTC's pointer arrays.
// Keep them alive after the options phase ends, through the compile call.
string header_path, nvrtc_option, real_option, compute_sm, rt_option, const_option;
std::vector<const char *> header_names;
std::vector<string> header_sources;
std::vector<const char *> header_sources_ptr;

for (auto fn : header_names) {
string source = RHIContext::read_file(string("cuda/") + fn);
header_sources_ptr.push_back(source.c_str());
header_sources.push_back(ocarina::move(source));
}

// NVRTC will embed the chosen architecture into the PTX (e.g. .target sm_120).
// OptiX's PTX frontend can lag behind the newest GPU targets; if we hand it PTX
// with an unsupported .target, optixModuleCreate() fails with "Invalid PTX input".
//
// PTX is ABI-compatible by the driver for the real GPU anyway, so we can OptiX modules
// It's safer to cap the PTX target to a widely supported SM.
int ptx_sm = sm;
if (function.is_raytracing()) {
// keep this conservative(for compatibility) given OptiX releases.
ptx_sm = std::min(ptx_sm, 89);
}
auto compute_sm = ocarina::format("-arch=compute_{}", ptx_sm);
auto rt_option = fmt::format("-DLC_OPTIX_VERSION={}", OPTIX_VERSION);
auto const_option = fmt::format("-Dlc_constant={}", nvrtc_version <= 110200 ? "const" : "constexpr");
ocarina::vector<const char *> compile_option = {
"--std=c++17",
compute_sm.c_str(),
const_option.c_str(),
rt_option.c_str(),
nvrtc_option.c_str(),
real_option.c_str(),
"-default-device",
"--use_fast_math",
"-restrict",
header_path.c_str(),
//#ifndef NDEBUG
"-lineinfo",
//#endif
"-extra-device-vectorization",
"-dw",
"-w"};
ocarina::vector<const char *> compile_option;
ocarina::list<string> includes;
if (function.is_raytracing()) {
static string inc_path = (fs::current_path() / "cuda" / "optix").string();
includes.push_back(ocarina::format("-I {}", inc_path));
compile_option.push_back(includes.back().c_str());
header_names.push_back("optix_device_header.h");
string source = RHIContext::read_file(string("cuda/optix_device_header.h"));
header_sources_ptr.push_back(source.c_str());
header_sources.push_back(ocarina::move(source));
int ptx_sm = sm;
{
switch_profile::Scope options_profile{"CUDACompiler::compile.headers_options", "compiler_setup"};
fs::path cuda_path = get_cuda_path();
cuda_path = cuda_path / "include";

header_path = "-I" + cuda_path.string();
int ver_major = 0;
int ver_minor = 0;
OC_NVRTC_CHECK(nvrtcVersion(&ver_major, &ver_minor));
int nvrtc_version = ver_major * 10000 + ver_minor * 100;
const auto storage_policy = function.storage_policy();
const int use_float32 = storage_policy.policy == PrecisionPolicy::force_f16 ? 0 : 1;
nvrtc_option = fmt::format("-DLC_NVRTC_VERSION={}", nvrtc_version);
real_option = fmt::format("-DOCARINA_CUDA_USE_FLOAT32={}", use_float32);
header_names = {"cuda_device_std.h", "cuda_device_scalar.h", "cuda_device_vector.h",
"cuda_device_matrix.h", "cuda_device_builtin.h", "cuda_tensor_builtin.h", "cuda_vector_func.h",
"cuda_device_math.h", "cuda_matrix_func.h", "cuda_device_resource.h"};

for (auto fn : header_names) {
string source = RHIContext::read_file(string("cuda/") + fn);
header_sources_ptr.push_back(source.c_str());
header_sources.push_back(ocarina::move(source));
}

// NVRTC will embed the chosen architecture into the PTX (e.g. .target sm_120).
// OptiX's PTX frontend can lag behind the newest GPU targets; if we hand it PTX
// with an unsupported .target, optixModuleCreate() fails with "Invalid PTX input".
//
// PTX is ABI-compatible by the driver for the real GPU anyway, so we can OptiX modules
// It's safer to cap the PTX target to a widely supported SM.
if (function.is_raytracing()) {
// keep this conservative(for compatibility) given OptiX releases.
ptx_sm = std::min(ptx_sm, 89);
}
compute_sm = ocarina::format("-arch=compute_{}", ptx_sm);
rt_option = fmt::format("-DLC_OPTIX_VERSION={}", OPTIX_VERSION);
const_option = fmt::format("-Dlc_constant={}", nvrtc_version <= 110200 ? "const" : "constexpr");
compile_option = {
"--std=c++17",
compute_sm.c_str(),
const_option.c_str(),
rt_option.c_str(),
nvrtc_option.c_str(),
real_option.c_str(),
"-default-device",
"--use_fast_math",
"-restrict",
header_path.c_str(),
//#ifndef NDEBUG
"-lineinfo",
//#endif
"-extra-device-vectorization",
"-dw",
"-w"};
if (function.is_raytracing()) {
static string inc_path = (fs::current_path() / "cuda" / "optix").string();
includes.push_back(ocarina::format("-I {}", inc_path));
compile_option.push_back(includes.back().c_str());
header_names.push_back("optix_device_header.h");
string source = RHIContext::read_file(string("cuda/optix_device_header.h"));
header_sources_ptr.push_back(source.c_str());
header_sources.push_back(ocarina::move(source));
}
for (const auto &header_name : header_names) {
includes.push_back(ocarina::format("-include={}", header_name));
compile_option.push_back(includes.back().c_str());
}
}
for (const auto &header_name : header_names) {
includes.push_back(ocarina::format("-include={}", header_name));
compile_option.push_back(includes.back().c_str());
string fn;
{
switch_profile::Scope hash_profile{"CUDACompiler::compile.cache_key", "cache_key"};
uint64_t ext_hash = hash64(hash64_list(compile_option),
hash64_list(header_sources_ptr),
ast_to_cuda_source_cache_version,
shader_codegen_path_version(Env::shader_codegen_path()));
fn = function.func_name(ext_hash, function.description());
}

uint64_t ext_hash = hash64(hash64_list(compile_option),
hash64_list(header_sources_ptr),
ast_to_cuda_source_cache_version,
shader_codegen_path_version(Env::shader_codegen_path()));

auto compile = [&](const string &cu, const string &fn, int sm) -> string {
TIMER_TAG(compile, "compile " + fn);
OC_INFO_FORMAT("compiling shader with NVRTC: {}", fn);
TIMER_TAG(compile, "NVRTC compile " + fn);
switch_profile::Scope nvrtc_profile{"CUDACompiler::compile.nvrtc", "nvrtc_setup"};
nvrtcProgram program{};
OC_NVRTC_CHECK(nvrtcCreateProgram(&program, cu.c_str(), fn.c_str(),
header_names.size(), header_sources_ptr.data(),
header_names.data()));
const nvrtcResult compile_res = nvrtcCompileProgram(program, compile_option.size(), compile_option.data());
nvrtcResult compile_res;
{
switch_profile::Scope compile_profile{"nvrtcCompileProgram", "nvrtc"};
compile_res = nvrtcCompileProgram(program, compile_option.size(), compile_option.data());
}
size_t log_size = 0;
OC_NVRTC_CHECK(nvrtcGetProgramLogSize(program, &log_size));
string log;
Expand All @@ -169,14 +189,12 @@ ocarina::string CUDACompiler::compile(const Function &function, int sm) const no
return ptx;
};

string fn = function.func_name(ext_hash, function.description());

ocarina::string ptx_fn = fn + ".ptx";
string cu_fn = fn + ".cu";
ocarina::string ptx;
RHIContext *context = device_->context();
if (!context->is_exist_cache(ptx_fn)) {
OC_INFO_FORMAT("miss ptx file {}", ptx_fn);
OC_INFO_FORMAT("shader PTX cache miss: {}", ptx_fn);
if (!context->is_exist_cache(cu_fn)) {
const ocarina::string cu = emit_cuda_source(function);
context->write_global_cache(cu_fn, cu);
Expand All @@ -188,7 +206,7 @@ ocarina::string CUDACompiler::compile(const Function &function, int sm) const no
context->write_global_cache(ptx_fn, ptx);
}
} else {
OC_INFO_FORMAT("find ptx file {}", ptx_fn);
OC_INFO_FORMAT("shader PTX cache hit: {}", ptx_fn);
ptx = context->read_global_cache(ptx_fn);
}
return ptx;
Expand Down
4 changes: 3 additions & 1 deletion modules/ocarina/backends/cuda/cuda_device.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
//

#include "cuda_device.h"
#include "core/util/switch_profile.h"
#include "cuda_stream.h"
#include "cuda_texture.h"
#include "cuda_shader.h"
Expand Down Expand Up @@ -239,6 +240,7 @@ handle_ty CUDADevice::create_texture_from_external(ocarina::uint tex_handle) noe
}

handle_ty CUDADevice::create_shader(const Function &function) noexcept {
switch_profile::Scope profile{"CUDADevice::create_shader", "backend"};
CUDACompiler compiler(this);
ocarina::string ptx = compiler.compile(function, compute_capability_);

Expand Down Expand Up @@ -502,4 +504,4 @@ OC_EXPORT_API ocarina::CUDADevice *create_device(ocarina::RHIContext *context) {
OC_EXPORT_API void destroy(ocarina::CUDADevice *device) {
ocarina::delete_with_allocator(device);
OC_INFO("cuda device is destroy!");
}
}
Loading
Loading