Memory layout optimization

Author: DoeringChristian

src/python/freeze.cpp

@@ -23,6 +23,9 @@
 #include <sstream>
 #include <vector>
 
+#define likely(x)   DRJIT_LIKELY(x)
+#define unlikely(x) DRJIT_UNLIKELY(x)
+
 /**
  * \brief Helper struct to profile and log frozen functions.
  */
@@ -229,10 +232,10 @@ bool FlatVariables::fill_opaque_mask(FlatVariables &prev,
         Layout &layout      = this->layout[i];
         Layout &prev_layout = prev.layout[i];
 
-        if (layout.flags & (uint32_t) LayoutFlag::Literal &&
-            prev_layout.flags & (uint32_t) LayoutFlag::Literal &&
-            (layout.literal != prev_layout.literal ||
-             layout.index != prev_layout.index)) {
+        if (unlikely(layout.flags & (uint32_t) LayoutFlag::Literal &&
+                     prev_layout.flags & (uint32_t) LayoutFlag::Literal &&
+                     (layout.literal != prev_layout.literal ||
+                      layout.index != prev_layout.index))) {
             opaque_mask[i] = true;
             new_opaques = true;
         }
@@ -288,14 +291,14 @@ void FlatVariables::schedule_jit_variables(bool schedule_force,
             // is nice to have a fallback.
             layout.literal = info.literal;
             // Store size in index variable, as this is not used for literals
-            layout.index = info.size;
-            layout.vt    = info.type;
+            layout.index         = info.size;
+            layout.vt            = (uint32_t) info.type;
             layout.literal_index = index;
 
             layout.flags |= (uint32_t) LayoutFlag::Literal;
         } else {
             layout.index = this->add_jit_index(index);
-            layout.vt    = info.type;
+            layout.vt    = (uint32_t) info.type;
             jit_var_dec_ref(index);
         }
     }
@@ -389,7 +392,7 @@ void FlatVariables::traverse_jit_index(uint32_t index, TraverseContext &ctx,
 
     layout.flags |= (uint32_t) LayoutFlag::JitIndex;
     layout.index = index;
-    layout.vt = jit_var_type(index);
+    layout.vt    = (uint32_t) jit_var_type(index);
 }
 
 /**
@@ -408,7 +411,7 @@ uint32_t FlatVariables::construct_jit_index(uint32_t prev_index) {
     if (layout.flags & (uint32_t) LayoutFlag::Literal) {
         index = layout.literal_index;
         jit_var_inc_ref(index);
-        vt    = layout.vt;
+        vt = (VarType) layout.vt;
     } else {
         VarLayout &var_layout = this->var_layout[layout.index];
         index = this->variables[layout.index];
@@ -1706,6 +1709,7 @@ nb::object FunctionRecording::replay(nb::callable func,
 }
 
 nb::object FrozenFunction::operator()(nb::args args, nb::kwargs kwargs) {
+    ProfilerPhase profiler("frozen function");
     nb::object result;
     {
         // Enter Isolate grad scope, so that gradients are not propagated

src/python/freeze.h

@@ -45,23 +45,27 @@ enum class LayoutFlag : uint32_t {
 /// sub-elements or their field keys. This can be used to reconstruct a PyTree
 /// from a flattened variable array.
 struct Layout {
-    /// Number of members in this container.
-    /// Can be used to traverse the layout without knowing the type.
-    uint32_t num = 0;
+
+    /// The literal data
+    uint64_t literal = 0;
+
     /// Optional field identifiers of the container
     /// for example: keys in dictionary
     drjit::vector<nb::object> fields;
+
+    /// Number of members in this container.
+    /// Can be used to traverse the layout without knowing the type.
+    uint32_t num = 0;
     /// The index in the flat_variables array of this variable.
     /// This can be used to determine aliasing.
     uint32_t index = 0;
 
     /// Flags, storing information about variables and literals.
-    uint32_t flags = 0;
+    uint32_t flags : 6; // LayoutFlag
 
-    /// The literal data
-    uint64_t literal = 0;
     /// Optional drjit type of the variable
-    VarType vt = VarType::Void;
+    uint32_t vt: 4; // VarType
+
     /// Variable index of literal. Instead of constructing a literal every time,
     /// we keep a reference to it.
     uint32_t literal_index = 0;
@@ -77,7 +81,9 @@ struct Layout {
     bool operator==(const Layout &rhs) const;
     bool operator!=(const Layout &rhs) const { return !(*this == rhs); }
 
-    Layout() = default;
+    Layout()
+        : literal(0), fields(), num(0), index(0), flags(0), vt(0),
+          literal_index(0), py_object(), type() {};
 
     Layout(const Layout &)            = delete;
     Layout &operator=(const Layout &) = delete;