Skip to content

added GAS compaction fixes #3 #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
brendancol merged 1 commit into from
Jan 28, 2026
Merged

added GAS compaction fixes #3 #32

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 48 additions & 10 deletions rtxpy/rtx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,19 +5,19 @@
NVIDIA's OptiX ray tracing engine via the otk-pyoptix Python bindings.
"""

import os
import atexit
import struct
from dataclasses import dataclass, field
from typing import Dict, List, Optional

# CRITICAL: cupy must be imported before optix for proper CUDA context sharing
import cupy

Check failure on line 15 in rtxpy/rtx.py

View workflow job for this annotation

GitHub Actions / Lint & Import Check

Ruff (I001) 

rtxpy/rtx.py:8:1: I001 Import block is un-sorted or un-formatted
has_cupy = True

import optix

Check failure on line 18 in rtxpy/rtx.py

View workflow job for this annotation

GitHub Actions / Lint & Import Check

Ruff (E402) 

rtxpy/rtx.py:18:1: E402 Module level import not at top of file

import numpy as np

Check failure on line 20 in rtxpy/rtx.py

View workflow job for this annotation

GitHub Actions / Lint & Import Check

Ruff (E402) 

rtxpy/rtx.py:20:1: E402 Module level import not at top of file

Check failure on line 20 in rtxpy/rtx.py

View workflow job for this annotation

GitHub Actions / Lint & Import Check

Ruff (I001) 

rtxpy/rtx.py:18:1: I001 Import block is un-sorted or un-formatted


# -----------------------------------------------------------------------------
Expand Down Expand Up @@ -362,9 +362,9 @@
)
build_input.numVertices = num_vertices

# Acceleration structure options
# Acceleration structure options - enable compaction for memory savings
accel_options = optix.AccelBuildOptions(
buildFlags=optix.BUILD_FLAG_ALLOW_RANDOM_VERTEX_ACCESS,
buildFlags=optix.BUILD_FLAG_ALLOW_RANDOM_VERTEX_ACCESS | optix.BUILD_FLAG_ALLOW_COMPACTION,
operation=optix.BUILD_OPERATION_BUILD,
)

Expand All @@ -378,7 +378,10 @@
d_temp = cupy.zeros(buffer_sizes.tempSizeInBytes, dtype=cupy.uint8)
gas_buffer = cupy.zeros(buffer_sizes.outputSizeInBytes, dtype=cupy.uint8)

# Build acceleration structure
# Allocate buffer to receive compacted size
compacted_size_buffer = cupy.zeros(1, dtype=cupy.uint64)

# Build acceleration structure with compacted size emission
gas_handle = _state.context.accelBuild(
0, # stream
[accel_options],
Expand All @@ -387,9 +390,24 @@
buffer_sizes.tempSizeInBytes,
gas_buffer.data.ptr,
buffer_sizes.outputSizeInBytes,
[], # emitted properties
[optix.AccelEmitDesc(compacted_size_buffer.data.ptr, optix.PROPERTY_TYPE_COMPACTED_SIZE)],
)

# Synchronize to ensure compacted size is available
cupy.cuda.Stream.null.synchronize()

# Compact if it saves memory
compacted_size = int(compacted_size_buffer[0])
if compacted_size < gas_buffer.nbytes:
compacted_buffer = cupy.zeros(compacted_size, dtype=cupy.uint8)
gas_handle = _state.context.accelCompact(
0, # stream
gas_handle,
compacted_buffer.data.ptr,
compacted_size,
)
gas_buffer = compacted_buffer

return gas_handle, gas_buffer


Expand Down Expand Up @@ -559,9 +577,9 @@
)
build_input.numVertices = num_vertices

# Acceleration structure options
# Acceleration structure options - enable compaction for memory savings
accel_options = optix.AccelBuildOptions(
buildFlags=optix.BUILD_FLAG_ALLOW_RANDOM_VERTEX_ACCESS,
buildFlags=optix.BUILD_FLAG_ALLOW_RANDOM_VERTEX_ACCESS | optix.BUILD_FLAG_ALLOW_COMPACTION,
operation=optix.BUILD_OPERATION_BUILD,
)

Expand All @@ -573,20 +591,40 @@

# Allocate buffers
d_temp = cupy.zeros(buffer_sizes.tempSizeInBytes, dtype=cupy.uint8)
_state.gas_buffer = cupy.zeros(buffer_sizes.outputSizeInBytes, dtype=cupy.uint8)
gas_buffer = cupy.zeros(buffer_sizes.outputSizeInBytes, dtype=cupy.uint8)

# Build acceleration structure
# Allocate buffer to receive compacted size
compacted_size_buffer = cupy.zeros(1, dtype=cupy.uint64)

# Build acceleration structure with compacted size emission
_state.gas_handle = _state.context.accelBuild(
0, # stream
[accel_options],
[build_input],
d_temp.data.ptr,
buffer_sizes.tempSizeInBytes,
_state.gas_buffer.data.ptr,
gas_buffer.data.ptr,
buffer_sizes.outputSizeInBytes,
[], # emitted properties
[optix.AccelEmitDesc(compacted_size_buffer.data.ptr, optix.PROPERTY_TYPE_COMPACTED_SIZE)],
)

# Synchronize to ensure compacted size is available
cupy.cuda.Stream.null.synchronize()

# Compact if it saves memory
compacted_size = int(compacted_size_buffer[0])
if compacted_size < gas_buffer.nbytes:
compacted_buffer = cupy.zeros(compacted_size, dtype=cupy.uint8)
_state.gas_handle = _state.context.accelCompact(
0, # stream
_state.gas_handle,
compacted_buffer.data.ptr,
compacted_size,
)
_state.gas_buffer = compacted_buffer
else:
_state.gas_buffer = gas_buffer

_state.current_hash = hash_value
return 0

Expand Down Expand Up @@ -644,7 +682,7 @@
_state.d_rays_size = rays_size
_state.d_rays[:] = cupy.asarray(rays, dtype=cupy.float32)
d_rays = _state.d_rays
rays_on_host = True

Check failure on line 685 in rtxpy/rtx.py

View workflow job for this annotation

GitHub Actions / Lint & Import Check

Ruff (F841) 

rtxpy/rtx.py:685:9: F841 Local variable `rays_on_host` is assigned to but never used

# Ensure hits buffer is on GPU
if isinstance(hits, cupy.ndarray):
Expand Down
Loading