Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/bvh.py (Python)

Code before:
        result = -1
        if j >= 0 and j < self.n_aabbs:
            result = 64
            x = self.morton_codes[i_b, ti.i32(i)] ^ self.morton_codes[i_b, ti.i32(j)]
            for b in range(64):
                if x & (ti.u64(1) << (63 - b)):
                    result = b
                    break
        return result

    def compute_bounds(self):
        """
        Compute the bounds of the BVH nodes.

        Starts from the leaf nodes and works upwards layer by layer.
        """
        self._kernel_compute_bounds_init()
        while self.updated[None]:
            self._kernel_compute_bounds_one_layer()

    @ti.kernel
    def _kernel_compute_bounds_init(self):
        self.updated[None] = True

        for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs - 1):
            self.internal_node_active[i_b, i] = ti.u8(0)
            self.internal_node_ready[i_b, i] = ti.u8(0)

        for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs):
            idx = ti.i32(self.morton_codes[i_b, i])
            self.nodes[i_b, i + self.n_aabbs - 1].bound.min = self.aabbs[i_b, idx].min
            self.nodes[i_b, i + self.n_aabbs - 1].bound.max = self.aabbs[i_b, idx].max
            parent_idx = self.nodes[i_b, i + self.n_aabbs - 1].parent
            if parent_idx != -1:
                self.internal_node_active[i_b, parent_idx] = ti.u8(1)

    @ti.kernel
    def _kernel_compute_bounds_one_layer(self):
        self.updated[None] = False
        for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs - 1):
            if self.internal_node_active[i_b, i] == ti.uint8(0):
                continue
            left_bound = self.nodes[i_b, self.nodes[i_b, i].left].bound
            right_bound = self.nodes[i_b, self.nodes[i_b, i].right].bound
            self.nodes[i_b, i].bound.min = ti.min(left_bound.min, right_bound.min)
            self.nodes[i_b, i].bound.max = ti.max(left_bound.max, right_bound.max)
            parent_idx = self.nodes[i_b, i].parent
            if parent_idx != -1:
                self.internal_node_ready[i_b, parent_idx] = ti.u8(1)
            self.internal_node_active[i_b, i] = ti.u8(0)
            self.updated[None] = True

Review comment (suggestion):
`= 0` should be fine, the compiler will do the job, since the variable is already allocated with the right type.

Code after:
        result = -1
        if j >= 0 and j < self.n_aabbs:
            result = 64
            x = self.morton_codes[i_b, ti.i32(i)] ^ self.morton_codes[i_b, ti.i32(j)]
            for b in range(64):
                if x & (ti.u64(1) << (63 - b)):
                    result = b
                    break
        return result

    def compute_bounds(self):
        """
        Compute the bounds of the BVH nodes.

        Starts from the leaf nodes and works upwards layer by layer.
        """
        self._kernel_compute_bounds_init()
        while self.updated[None]:
            self._kernel_compute_bounds_one_layer()

    @ti.kernel
    def _kernel_compute_bounds_init(self):
        self.updated[None] = True
        self.internal_node_active.fill(0)
        self.internal_node_ready.fill(0)

        for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs):
            idx = ti.i32(self.morton_codes[i_b, i])
            self.nodes[i_b, i + self.n_aabbs - 1].bound.min = self.aabbs[i_b, idx].min
            self.nodes[i_b, i + self.n_aabbs - 1].bound.max = self.aabbs[i_b, idx].max
            parent_idx = self.nodes[i_b, i + self.n_aabbs - 1].parent
            if parent_idx != -1:
                self.internal_node_active[i_b, parent_idx] = 1

    @ti.kernel
    def _kernel_compute_bounds_one_layer(self):
        self.updated[None] = False
        for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs - 1):
            if self.internal_node_active[i_b, i] == 0:
                continue
            left_bound = self.nodes[i_b, self.nodes[i_b, i].left].bound
            right_bound = self.nodes[i_b, self.nodes[i_b, i].right].bound
            self.nodes[i_b, i].bound.min = ti.min(left_bound.min, right_bound.min)
            self.nodes[i_b, i].bound.max = ti.max(left_bound.max, right_bound.max)
            parent_idx = self.nodes[i_b, i].parent
            if parent_idx != -1:
                self.internal_node_ready[i_b, parent_idx] = 1
            self.internal_node_active[i_b, i] = 0
            self.updated[None] = True

        for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs - 1):