tasks example

Author: Felix Igelbrink

examples/compile_with_tasks.py

@@ -0,0 +1,78 @@
+import concurrent.futures
+
+import direct.task.Task
+
+import optix as ox
+import argparse
+import logging
+import sys
+from concurrent.futures import ThreadPoolExecutor
+import time
+
+logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
+log = logging.getLogger()
+
+
+if __name__ == "__main__":
+    if ox.optix_version()[1] < 4:
+        raise NotImplementedError("Parallel tasks are not implemented in optix versions < 7.3.")
+
+    parser = argparse.ArgumentParser("Compile OptiX modules using parallel tasks")
+    parser.add_argument('file', nargs=1, help="The input file (.ptx or .cu) to compile")
+    parser.add_argument('-na', '--num-attributes', type=int, default=2, required=False,
+                        help="Number of attribute values (up to 8, default 2)")
+    parser.add_argument('-npv', '--num-payload-values', type=int, default=2, required=False,
+                        help=f"Number of payload values (up to {ox.PipelineCompileOptions.DEFAULT_MAX_PAYLOAD_VALUE_COUNT}, default 2)")
+    parser.add_argument('-npt', '--num-payload-types', type=int, default=1, required=False,
+                        help=f"Number of payload types (up to {ox.ModuleCompileOptions.DEFAULT_MAX_PAYLOAD_TYPE_COUNT}, default 1)")
+    parser.add_argument('-ni', '--num-iters', type=int, default=1, required=False,
+                        help="Number of iterations to compile. > 1 disables disk cache (default 1)")
+    parser.add_argument('-dt', '--disable-tasks', action='store_true', required=False,
+                        help="Disable compilation with tasks (default enabled)")
+    parser.add_argument('-nt', '--num-threads', type=int, default=1, required=False,
+                        help="Number of threads (default 1)")
+    parser.add_argument('-mt', '--max-num-tasks', type=int, default=2, required=False,
+                        help="Maximum number of additional tasks (default 2)")
+
+    args = parser.parse_args()
+
+    logger = ox.Logger(log)
+    ctx = ox.DeviceContext(validation_mode=True, log_callback_function=logger, log_callback_level=3)
+
+    if args.num_iters > 1:
+        ctx.cache_enabled = False
+
+    # compile the file content to ptx in case a .cu file is given
+    ptx = ox.Module.compile_cuda_ptx(args.file[0])
+
+    pipeline_options = ox.PipelineCompileOptions(num_payload_values=0,
+                                                 num_attribute_values=args.num_attributes)
+
+    payload_semantics = [ox.PayloadSemantics.DEFAULT] * args.num_payload_values
+    payload_types = [payload_semantics] * args.num_payload_types
+
+    compile_opts = ox.ModuleCompileOptions(payload_types=payload_types)
+
+    use_tasks = not args.disable_tasks
+
+    if use_tasks:
+        tic = time.time()
+        with ThreadPoolExecutor(max_workers=args.num_threads) as executor:
+            for i in range(args.num_iters):
+                module, task = ox.Module.create_as_task(ctx, ptx, module_compile_options=compile_opts, pipeline_compile_options=pipeline_options)
+                task_futures = {executor.submit(task.execute, args.max_num_tasks)}
+            while task_futures:
+                done, not_done = concurrent.futures.wait(task_futures, timeout=0.25, return_when=concurrent.futures.FIRST_COMPLETED)
+                for future in done:
+                    new_tasks = future.result()
+                    if len(new_tasks) > 0:
+                        task_futures.update({executor.submit(t.execute, args.max_num_tasks) for t in new_tasks})
+                    task_futures.remove(future)
+
+        # wait for the executor to finish here
+        print("Overall run time with tasks", time.time()-tic)
+    else:
+        tic = time.time()
+        for i in range(args.num_iters):
+            module = ox.Module(ctx, ptx, module_compile_options=compile_opts, pipeline_compile_options=pipeline_options)
+        print("Overall run time without tasks", time.time()-tic)

optix/__init__.py

@@ -1,6 +1,6 @@
 from .context import DeviceContext, optix_version
 from .build import *
-from .module import Module, ModuleCompileOptions, CompileOptimizationLevel, CompileDebugLevel
+from .module import Module, ModuleCompileOptions, CompileOptimizationLevel, CompileDebugLevel, PayloadSemantics, Task
 from .program_group import ProgramGroup
 from .struct import SbtRecord, LaunchParamsRecord
 from .shader_binding_table import ShaderBindingTable

optix/context.pyx

@@ -12,6 +12,8 @@ def optix_version():
     return _OPTIX_VERSION_MAJOR, _OPTIX_VERSION_MINOR, _OPTIX_VERSION_MICRO
 
 
+
+
 cdef class _LogWrapper:
     def __init__(self, log_function):
         self.log_function = log_function

optix/module.pxd

@@ -10,6 +10,8 @@ cdef extern from "optix_includes.h" nogil:
     cdef OptixResult optixInit()
 
     cdef size_t OPTIX_COMPILE_DEFAULT_MAX_REGISTER_COUNT
+    cdef size_t OPTIX_COMPILE_DEFAULT_MAX_PAYLOAD_VALUE_COUNT
+
 
     cdef enum OptixCompileOptimizationLevel:
         OPTIX_COMPILE_OPTIMIZATION_DEFAULT,
@@ -27,6 +29,8 @@ cdef extern from "optix_includes.h" nogil:
 
 
     IF _OPTIX_VERSION > 70300:  # switch to new version
+        cdef size_t OPTIX_COMPILE_DEFAULT_MAX_PAYLOAD_TYPE_COUNT
+
         cdef enum OptixPayloadSemantics:
             OPTIX_PAYLOAD_SEMANTICS_TRACE_CALLER_NONE,
             OPTIX_PAYLOAD_SEMANTICS_TRACE_CALLER_READ,

optix/module.pyx

@@ -112,18 +112,17 @@ IF _OPTIX_VERSION > 70300:
                 tasks.append(t)
             return tasks
 
-ELSE:
-    class PayloadType(IntFlag):
-        DEFAULT = 0 # only for interface. Ignored in Optix versions < 7.4
-
 
 cdef class ModuleCompileOptions(OptixObject):
     """
     Wraps the OptixModuleCompileOptions struct.
     """
     DEFAULT_MAX_REGISTER_COUNT = OPTIX_COMPILE_DEFAULT_MAX_REGISTER_COUNT
+    DEFAULT_MAX_PAYLOAD_TYPE_COUNT = OPTIX_COMPILE_DEFAULT_MAX_PAYLOAD_TYPE_COUNT
+    DEFAULT_MAX_PAYLOAD_VALUE_COUNT = OPTIX_COMPILE_DEFAULT_MAX_PAYLOAD_VALUE_COUNT
+
     def __init__(self,
-                 max_register_count=OPTIX_COMPILE_DEFAULT_MAX_REGISTER_COUNT,
+                 max_register_count=DEFAULT_MAX_REGISTER_COUNT,
                  opt_level=CompileOptimizationLevel.DEFAULT,
                  debug_level= CompileDebugLevel.DEFAULT,
                  payload_types=None): #TODO add bound values
@@ -185,7 +184,6 @@ cdef _is_ptx(src):
     if not isinstance(src, (bytes, bytearray)):
         return False
     for line in src.splitlines():
-        print(line)
         if len(line) == 0 or line.startswith(b'//') or line.startswith(b'\n'):
             continue
         return line.startswith(b'.version')
@@ -241,13 +239,10 @@ cdef class Module(OptixContextObject):
         self._compile_flags = list(compile_flags)
 
         if src is not None:
-            if not _is_ptx(src):
-                ptx = self._compile_cuda_ptx(src, compile_flags, name=program_name)
-            else:
-                ptx = src
+            ptx = self.compile_cuda_ptx(src, compile_flags, name=program_name)
             c_ptx = ptx
-            IF _OPTIX_VERSION > 70300:
-                self._check_payload_values(module_compile_options, pipeline_compile_options)
+            #IF _OPTIX_VERSION > 70300:
+            #    self._check_payload_values(module_compile_options, pipeline_compile_options)
 
             optix_check_return(optixModuleCreateFromPTX(self.context.c_context,
                                      &module_compile_options.compile_options,
@@ -270,18 +265,18 @@ cdef class Module(OptixContextObject):
                 optix_check_return(optixModuleGetCompilationState(self.module, &state))
             return ModuleCompileState(state)
 
-        @staticmethod
-        def _check_payload_values(ModuleCompileOptions module_compile_options, PipelineCompileOptions pipeline_compile_options):
-            IF _OPTIX_VERSION > 70300:
-                # check if the payload values match between the module and pipeline compile options
-                pipeline_payload_values = <unsigned int> pipeline_compile_options.compile_options.numPayloadValues
-                if module_compile_options.payload_types.size() > 0:
-                    for i in range(module_compile_options.compile_options.numPayloadTypes):
-                        if pipeline_payload_values != module_compile_options.compile_options.payloadTypes[
-                            i].numPayloadValues:
-                            raise ValueError(
-                                f"number of payload values in module compile options at index {i} does not match the num_payload_values in the pipeline_compile_options.")
-            return
+        # @staticmethod
+        # def _check_payload_values(ModuleCompileOptions module_compile_options, PipelineCompileOptions pipeline_compile_options):
+        #     IF _OPTIX_VERSION > 70300:
+        #         # check if the payload values match between the module and pipeline compile options
+        #         pipeline_payload_values = <unsigned int> pipeline_compile_options.compile_options.numPayloadValues
+        #         if module_compile_options.payload_types.size() > 0:
+        #             for i in range(module_compile_options.compile_options.numPayloadTypes):
+        #                 if pipeline_payload_values != module_compile_options.compile_options.payloadTypes[
+        #                     i].numPayloadValues:
+        #                     raise ValueError(
+        #                         f"number of payload values in module compile options at index {i} does not match the num_payload_values in the pipeline_compile_options.")
+        #     return
 
         @classmethod
         def create_as_task(cls,
@@ -323,12 +318,9 @@ cdef class Module(OptixContextObject):
             cdef Module module = Module(context, None, compile_flags=compile_flags)
             cdef const char * c_ptx
             cdef unsigned int pipeline_payload_values, i
-            cls._check_payload_values(module_compile_options, pipeline_compile_options)
+            #cls._check_payload_values(module_compile_options, pipeline_compile_options)
 
-            if not _is_ptx(src):
-                ptx = cls._compile_cuda_ptx(src, compile_flags, name=program_name)
-            else:
-                ptx = src
+            ptx = cls.compile_cuda_ptx(src, compile_flags, name=program_name)
             c_ptx = ptx
 
             cdef Task task = Task(module)
@@ -381,11 +373,13 @@ cdef class Module(OptixContextObject):
         return module
 
     @staticmethod
-    def _compile_cuda_ptx(src, compile_flags, name=None, **kwargs):
+    def compile_cuda_ptx(src, compile_flags=_nvrtc_compile_flags_default, name=None, **kwargs):
         if os.path.exists(src):
             name = src
             with open(src, 'r') as f:
                 src = f.read()
+        if _is_ptx(src):
+            return src
 
         elif name is None:
             name = "default_program"

optix/pipeline.pxd

@@ -5,6 +5,7 @@ from .program_group cimport ProgramGroup, OptixStackSizes
 from .shader_binding_table cimport OptixShaderBindingTable
 
 cdef extern from "optix_includes.h" nogil:
+    cdef size_t OPTIX_COMPILE_DEFAULT_MAX_PAYLOAD_VALUE_COUNT
 
     # pipeline functions and structs
     ctypedef struct OptixPipeline:
@@ -23,7 +24,7 @@ cdef extern from "optix_includes.h" nogil:
         OPTIX_TRAVERSABLE_GRAPH_FLAG_ALLOW_SINGLE_LEVEL_INSTANCING
 
 
-    IF _OPTIX_VERSION_MAJOR == 7 and _OPTIX_VERSION_MINOR > 3:  # switch to new instance flags
+    IF _OPTIX_VERSION > 70300:  # switch to new instance flags
         cdef enum OptixCompileDebugLevel:
             OPTIX_COMPILE_DEBUG_LEVEL_DEFAULT,
             OPTIX_COMPILE_DEBUG_LEVEL_NONE,

optix/pipeline.pyx

@@ -84,6 +84,8 @@ cdef class PipelineCompileOptions(OptixObject):
     """
     Class wrapping the OptixPipelineCompileOptions struct.
     """
+    DEFAULT_MAX_PAYLOAD_VALUE_COUNT = OPTIX_COMPILE_DEFAULT_MAX_PAYLOAD_VALUE_COUNT
+
     def __init__(self,
                  uses_motion_blur=False,
                  traversable_graph_flags = TraversableGraphFlags.ALLOW_ANY,
@@ -92,13 +94,13 @@ cdef class PipelineCompileOptions(OptixObject):
                  exception_flags = ExceptionFlags.NONE,
                  pipeline_launch_params_variable_name = "params",
                  uses_primitive_type_flags = PrimitiveTypeFlags.DEFAULT):
-        self.compile_options.usesMotionBlur = uses_motion_blur
-        self.compile_options.traversableGraphFlags = traversable_graph_flags.value
-        self.compile_options.numPayloadValues = num_payload_values
-        self.compile_options.numAttributeValues = num_attribute_values
-        self.compile_options.exceptionFlags = exception_flags.value
+        self.uses_motion_blur = uses_motion_blur
+        self.traversable_graph_flags = traversable_graph_flags
+        self.num_payload_values = num_payload_values
+        self.num_attribute_values = num_attribute_values
+        self.exception_flags = exception_flags
         self.pipeline_launch_params_variable_name = pipeline_launch_params_variable_name
-        self.compile_options.usesPrimitiveTypeFlags = uses_primitive_type_flags.value
+        self.uses_primitive_type_flags = uses_primitive_type_flags
 
     @property
     def uses_motion_blur(self):
@@ -122,6 +124,8 @@ cdef class PipelineCompileOptions(OptixObject):
 
     @num_payload_values.setter
     def num_payload_values(self, num_payload_values):
+        if num_payload_values > self.DEFAULT_MAX_PAYLOAD_VALUE_COUNT:
+            raise ValueError(f"A maximum of {self.DEFAULT_MAX_PAYLOAD_VALUE_COUNT} payload values is allowed.")
         self.compile_options.numPayloadValues = num_payload_values
 
     @property