@@ -289,6 +289,7 @@ bool FlatVariables::fill_opaque_mask(FlatVariables &prev,
289
289
290
290
void FlatVariables::schedule_jit_variables (bool schedule_force,
291
291
drjit::vector<bool > *opaque_mask) {
292
+ ProfilerPhase profiler (" schedule_jit_variables" );
292
293
for (uint32_t i = layout_index; i < layout.size (); i++) {
293
294
Layout &layout = this ->layout [i];
294
295
@@ -311,7 +312,7 @@ void FlatVariables::schedule_jit_variables(bool schedule_force,
311
312
jit_var_inc_ref (index );
312
313
}
313
314
314
- VarInfo info = jit_set_backend (index );
315
+ VarInfo info = jit_var_info (index );
315
316
if (backend == info.backend || this ->backend == JitBackend::None) {
316
317
backend = info.backend ;
317
318
} else {
@@ -352,12 +353,13 @@ void FlatVariables::schedule_jit_variables(bool schedule_force,
352
353
* over the collected indices and collects that information.
353
354
*/
354
355
void FlatVariables::record_jit_variables () {
356
+ ProfilerPhase profiler (" record_jit_variables" );
355
357
assert (variables.size () == var_layout.size ());
356
358
for (uint32_t i = 0 ; i < var_layout.size (); i++){
357
359
uint32_t index = variables[i];
358
360
VarLayout &layout = var_layout[i];
359
361
360
- VarInfo info = jit_set_backend (index );
362
+ VarInfo info = jit_var_info (index );
361
363
if (info.type == VarType::Pointer) {
362
364
// We do not support pointers as inputs. It might be possible with
363
365
// some extra handling, but they are never used directly.
@@ -444,8 +446,6 @@ uint32_t FlatVariables::construct_jit_index(uint32_t prev_index) {
444
446
uint32_t index ;
445
447
VarType vt;
446
448
if (layout.flags & (uint32_t ) LayoutFlag::Literal) {
447
- // index = jit_var_literal(this->backend, layout.vt, &layout.literal,
448
- // layout.index);
449
449
index = layout.literal_index ;
450
450
jit_var_inc_ref (index );
451
451
vt = layout.vt ;
@@ -1288,6 +1288,28 @@ void FlatVariables::assign_with_registry(nb::handle dst, TraverseContext &ctx) {
1288
1288
}
1289
1289
}
1290
1290
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
+
1291
1313
bool log_diff_variable (LogLevel level, const FlatVariables &curr,
1292
1314
const FlatVariables &prev, std::string &path,
1293
1315
uint32_t slot) {
@@ -1777,6 +1799,7 @@ nb::object FrozenFunction::operator()(nb::args args, nb::kwargs kwargs) {
1777
1799
in_variables->layout_index = 0 ;
1778
1800
1779
1801
{ // Evaluate the variables, scheduled when traversing
1802
+ ProfilerPhase profiler (" eval" );
1780
1803
nb::gil_scoped_release guard;
1781
1804
jit_eval ();
1782
1805
}
0 commit comments