Skip to content

Commit 2f3a144

Browse files
Using jit_var_info instead of jit_set_backend
1 parent 1d76fdd commit 2f3a144

File tree

3 files changed

+31
-21
lines changed

3 files changed

+31
-21
lines changed

src/python/freeze.cpp

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,7 @@ bool FlatVariables::fill_opaque_mask(FlatVariables &prev,
289289

290290
void FlatVariables::schedule_jit_variables(bool schedule_force,
291291
drjit::vector<bool> *opaque_mask) {
292+
ProfilerPhase profiler("schedule_jit_variables");
292293
for (uint32_t i = layout_index; i < layout.size(); i++) {
293294
Layout &layout = this->layout[i];
294295

@@ -311,7 +312,7 @@ void FlatVariables::schedule_jit_variables(bool schedule_force,
311312
jit_var_inc_ref(index);
312313
}
313314

314-
VarInfo info = jit_set_backend(index);
315+
VarInfo info = jit_var_info(index);
315316
if (backend == info.backend || this->backend == JitBackend::None) {
316317
backend = info.backend;
317318
} else {
@@ -352,12 +353,13 @@ void FlatVariables::schedule_jit_variables(bool schedule_force,
352353
* over the collected indices and collects that information.
353354
*/
354355
void FlatVariables::record_jit_variables() {
356+
ProfilerPhase profiler("record_jit_variables");
355357
assert(variables.size() == var_layout.size());
356358
for (uint32_t i = 0; i < var_layout.size(); i++){
357359
uint32_t index = variables[i];
358360
VarLayout &layout = var_layout[i];
359361

360-
VarInfo info = jit_set_backend(index);
362+
VarInfo info = jit_var_info(index);
361363
if (info.type == VarType::Pointer) {
362364
// We do not support pointers as inputs. It might be possible with
363365
// some extra handling, but they are never used directly.
@@ -444,8 +446,6 @@ uint32_t FlatVariables::construct_jit_index(uint32_t prev_index) {
444446
uint32_t index;
445447
VarType vt;
446448
if (layout.flags & (uint32_t) LayoutFlag::Literal) {
447-
// index = jit_var_literal(this->backend, layout.vt, &layout.literal,
448-
// layout.index);
449449
index = layout.literal_index;
450450
jit_var_inc_ref(index);
451451
vt = layout.vt;
@@ -1288,6 +1288,28 @@ void FlatVariables::assign_with_registry(nb::handle dst, TraverseContext &ctx) {
12881288
}
12891289
}
12901290

1291+
FlatVariables::~FlatVariables() {
1292+
state_lock_guard guard;
1293+
for (uint32_t i = 0; i < layout.size(); ++i) {
1294+
Layout &l = layout[i];
1295+
if (l.flags & (uint32_t) LayoutFlag::Literal && l.literal_index) {
1296+
jit_var_dec_ref(l.literal_index);
1297+
}
1298+
}
1299+
}
1300+
1301+
void FlatVariables::borrow() {
1302+
state_lock_guard guard;
1303+
for (uint32_t &index : this->variables)
1304+
jit_var_inc_ref(index);
1305+
}
1306+
1307+
void FlatVariables::release() {
1308+
state_lock_guard guard;
1309+
for (uint32_t &index : this->variables)
1310+
jit_var_dec_ref(index);
1311+
}
1312+
12911313
bool log_diff_variable(LogLevel level, const FlatVariables &curr,
12921314
const FlatVariables &prev, std::string &path,
12931315
uint32_t slot) {
@@ -1777,6 +1799,7 @@ nb::object FrozenFunction::operator()(nb::args args, nb::kwargs kwargs) {
17771799
in_variables->layout_index = 0;
17781800

17791801
{ // Evaluate the variables, scheduled when traversing
1802+
ProfilerPhase profiler("eval");
17801803
nb::gil_scoped_release guard;
17811804
jit_eval();
17821805
}

src/python/freeze.h

Lines changed: 3 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -226,14 +226,7 @@ struct FlatVariables {
226226
FlatVariables(FlatVariables &&) = default;
227227
FlatVariables &operator=(FlatVariables &&) = default;
228228

229-
~FlatVariables() {
230-
for (uint32_t i = 0; i < layout.size(); ++i) {
231-
Layout &l = layout[i];
232-
if (l.flags & (uint32_t) LayoutFlag::JitIndex && l.literal_index) {
233-
jit_var_dec_ref(l.literal_index);
234-
}
235-
}
236-
}
229+
~FlatVariables();
237230

238231
void clear() {
239232
this->layout_index = 0;
@@ -243,15 +236,9 @@ struct FlatVariables {
243236
this->backend = JitBackend::None;
244237
}
245238
/// Borrow all variables held by this struct.
246-
void borrow() {
247-
for (uint32_t &index : this->variables)
248-
jit_var_inc_ref(index);
249-
}
239+
void borrow();
250240
/// Release all variables held by this struct.
251-
void release() {
252-
for (uint32_t &index : this->variables)
253-
jit_var_dec_ref(index);
254-
}
241+
void release();
255242

256243
/**
257244
* Generates a mask of variables that should be made opaque in the next

0 commit comments

Comments
 (0)