Skip to content

Commit

Permalink
refactor! : Update default workspace size based on platforms.
Browse files Browse the repository at this point in the history
BREAKING CHANGE: This commit sets the default workspace size to 1GB for GPU platforms and 256MB for Jetson Nano/TX1 platforms whose compute capability is < 6.

Signed-off-by: Dheeraj Peri <[email protected]>

Signed-off-by: Dheeraj Peri <[email protected]>

Signed-off-by: Dheeraj Peri <[email protected]>

Signed-off-by: Dheeraj Peri <[email protected]>

Signed-off-by: Dheeraj Peri <[email protected]>
  • Loading branch information
peri044 committed Oct 18, 2021
1 parent a1180ce commit 391a4c0
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 2 deletions.
16 changes: 15 additions & 1 deletion core/compiler.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -347,6 +347,21 @@ torch::jit::script::Module CompileGraph(const torch::jit::script::Module& mod, C
if (cfg.partition_info.enabled) {
return CompileGraphWithFallback(mod, cfg);
}
auto device_spec = cfg.convert_info.engine_settings.device;

// GPU default WS size : 1 GB
// Set WS = 256 Mb for Jetson nano/TX1 like platforms whose compute capability is 5.X.
auto workspace_size = cfg.convert_info.engine_settings.workspace_size;
cudaDeviceProp device_prop;
cudaGetDeviceProperties(&device_prop, device_spec.gpu_id);
if (workspace_size == 0) {
if (device_prop.major < 6) {
cfg.convert_info.engine_settings.workspace_size = 256 * (1 << 20);
} else {
cfg.convert_info.engine_settings.workspace_size = 1 << 30;
}
}

// TODO: Should be doing a functional transform but need PR #31978
// [jit] More robust mangling
// torch::jit::script::Module new_mod = mod.clone();
Expand All @@ -357,7 +372,6 @@ torch::jit::script::Module CompileGraph(const torch::jit::script::Module& mod, C
if (method.name().compare("forward") == 0) {
auto engine = ConvertGraphToTRTEngine(mod, method.name(), cfg);
auto new_g = std::make_shared<torch::jit::Graph>();
auto device_spec = cfg.convert_info.engine_settings.device;
auto cuda_device = runtime::CudaDevice(device_spec.gpu_id, device_spec.device_type);
AddEngineToGraph(new_mod, new_g, engine, cuda_device);
auto new_method = new_mod._ivalue()->compilation_unit()->create_function(method.name(), new_g);
Expand Down
2 changes: 1 addition & 1 deletion core/conversion/conversionctx/ConversionCtx.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
net = make_trt(
builder->createNetworkV2(1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));

LOG_DEBUG(build_settings);
LOG_INFO(settings);
cfg = make_trt(builder->createBuilderConfig());

for (auto p = settings.enabled_precisions.begin(); p != settings.enabled_precisions.end(); ++p) {
Expand Down

0 comments on commit 391a4c0

Please sign in to comment.