Skip to content

Commit c74b13e

Browse files
authored
Merge pull request #1247 from pytorch/anuragd/remove_hard_codings
refactor: Changed the hardcoded values to macros for DLA memory sizes
2 parents f814415 + ad7c081 commit c74b13e

File tree

3 files changed

+11
-6
lines changed

3 files changed

+11
-6
lines changed

core/conversion/conversionctx/ConversionCtx.cpp

+3-3
Original file line numberDiff line numberDiff line change
@@ -124,13 +124,13 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
124124
settings.enabled_precisions.find(nvinfer1::DataType::kFLOAT) == settings.enabled_precisions.end(),
125125
"DLA supports only fp16 or int8 precision");
126126
cfg->setDLACore(settings.device.dla_core);
127-
if (settings.dla_sram_size != 1048576) {
127+
if (settings.dla_sram_size != DLA_SRAM_SIZE) {
128128
cfg->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kDLA_MANAGED_SRAM, settings.dla_sram_size);
129129
}
130-
if (settings.dla_local_dram_size != 1073741824) {
130+
if (settings.dla_local_dram_size != DLA_LOCAL_DRAM_SIZE) {
131131
cfg->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kDLA_LOCAL_DRAM, settings.dla_local_dram_size);
132132
}
133-
if (settings.dla_global_dram_size != 536870912) {
133+
if (settings.dla_global_dram_size != DLA_GLOBAL_DRAM_SIZE) {
134134
cfg->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kDLA_GLOBAL_DRAM, settings.dla_global_dram_size);
135135
}
136136
}

core/conversion/conversionctx/ConversionCtx.h

+3-3
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,9 @@ struct BuilderSettings {
3535
nvinfer1::IInt8Calibrator* calibrator = nullptr;
3636
uint64_t num_avg_timing_iters = 1;
3737
uint64_t workspace_size = 0;
38-
uint64_t dla_sram_size = 1048576;
39-
uint64_t dla_local_dram_size = 1073741824;
40-
uint64_t dla_global_dram_size = 536870912;
38+
uint64_t dla_sram_size = DLA_SRAM_SIZE;
39+
uint64_t dla_local_dram_size = DLA_LOCAL_DRAM_SIZE;
40+
uint64_t dla_global_dram_size = DLA_GLOBAL_DRAM_SIZE;
4141

4242
BuilderSettings() = default;
4343
BuilderSettings(const BuilderSettings& other) = default;

core/util/macros.h

+5
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,11 @@
44

55
#define GET_MACRO(_1, _2, NAME, ...) NAME
66

7+
// DLA Memory related macros
8+
#define DLA_SRAM_SIZE 1048576
9+
#define DLA_LOCAL_DRAM_SIZE 1073741824
10+
#define DLA_GLOBAL_DRAM_SIZE 536870912
11+
712
#define TORCHTRT_LOG(l, sev, msg) \
813
do { \
914
std::stringstream ss{}; \

0 commit comments

Comments
 (0)