# Copyright (c) 2023 NVIDIA CORPORATION. All rights reserved. # NVIDIA CORPORATION and its licensors retain all intellectual property # and proprietary rights in and to this software, related documentation # and any modifications thereto. Any use, reproduction, disclosure or # distribution of this software and related documentation without an express # license agreement from NVIDIA CORPORATION is strictly prohibited. """Backend implementation for kernel node(s).""" from __future__ import annotations from enum import IntFlag import functools import hashlib import importlib.util import json import operator import os import tempfile from typing import ( Any, Callable, Mapping, NamedTuple, Optional, Sequence, Tuple, Union, ) import omni.graph.core as og import warp as wp from omni.warp.nodes._impl.common import ( IntEnum, get_warp_type_from_data_type_name, type_convert_og_to_warp, ) from omni.warp.nodes._impl.attributes import ( ATTR_BUNDLE_TYPE, attr_cast_array_to_warp, attr_get_base_name, attr_get_name, attr_join_name, ) _ATTR_PORT_TYPE_INPUT = og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT _ATTR_PORT_TYPE_OUTPUT = og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT EXPLICIT_SOURCE = "explicit" # Enumerators # ------------------------------------------------------------------------------ class UserAttributesEvent(IntFlag): """User attributes event.""" NONE = 0 CREATED = 1 << 0 REMOVED = 1 << 1 class OutputArrayShapeSource(IntEnum): """Method to infer the shape of output attribute arrays.""" AS_INPUT_OR_AS_KERNEL = (0, "as input if any, or as kernel") AS_KERNEL = (1, "as kernel") class OutputBundleTypeSource(IntEnum): """Method to infer the type of output attribute bundles.""" AS_INPUT = (0, "as input if any") AS_INPUT_OR_EXPLICIT = (1, "as input if any, or explicit") EXPLICIT = (2, "explicit") class ArrayAttributeFormat(IntEnum): """Format describing how attribute arrays are defined on the node.""" RAW = (0, "raw") BUNDLE = (1, "bundle") # User Attributes Description # ------------------------------------------------------------------------------ class UserAttributeDesc(NamedTuple): """Description of an attribute added dynamically by users through the UI. This struct is what the Attribute Editor UI passes to the node in order to communicate any attribute metadata. """ port_type: og.AttributePortType base_name: str data_type_name: str is_array: bool array_format: ArrayAttributeFormat array_shape_source: Union[None, OutputArrayShapeSource] optional: bool @classmethod def deserialize( cls, data: Mapping[str:Any], ) -> Optional[UserAttributeDesc]: """Creates a new instance based on a serialized representation.""" # Retrieve the port type. It's invalid not to have any set. port_type = data.get("port_type") if port_type is None: return None port_type = og.AttributePortType(port_type) # Define sensible default values. # Although this class requires all of its member values to be explicitly # defined upon initialization, it's possible that the incoming data was # serialized with an older version of this class, in which case we might # want to try filling any gap. values = { "array_format": ArrayAttributeFormat.RAW, "array_shape_source": ( OutputArrayShapeSource.AS_INPUT_OR_AS_KERNEL if port_type == _ATTR_PORT_TYPE_OUTPUT else None ), "optional": False, } # Override the default values with the incoming data. values.update({k: v for k, v in data.items() if k in cls._fields}) # Ensure that the member values are set using their rightful types. values.update( { "port_type": port_type, "array_format": ArrayAttributeFormat(values["array_format"]), "array_shape_source": ( None if values["array_shape_source"] is None else OutputArrayShapeSource(values["array_shape_source"]) ), } ) try: # This might error in case some members are still missing. return cls(**values) except TypeError: return None @property def name(self) -> str: """Retrieves the attribute's name prefixed with its port type.""" return attr_join_name(self.port_type, self.base_name) @property def type(self) -> og.Attribute: """Retrieves OmniGraph's attribute type.""" return og.AttributeType.type_from_sdf_type_name(self.type_name) @property def type_name(self) -> str: """Retrieves OmniGraph's attribute type name.""" if self.is_array: return "{}[]".format(self.data_type_name) return self.data_type_name def serialize(self) -> Mapping[str:Any]: """Converts this instance into a serialized representation.""" return self._replace( port_type=int(self.port_type), )._asdict() def deserialize_user_attribute_descs( data: str, ) -> Mapping[str, UserAttributeDesc]: """Deserializes a string into a mapping of (name, desc).""" descs = {attr_join_name(x["port_type"], x["base_name"]): UserAttributeDesc.deserialize(x) for x in json.loads(data)} # Filter out any invalid description. return {k: v for k, v in descs.items() if v is not None} def serialize_user_attribute_descs( descs: Mapping[str, UserAttributeDesc], ) -> str: """Serializes a mapping of (name, desc) into a string.""" return json.dumps(tuple(x.serialize() for x in descs.values())) # User Attributes Information # ------------------------------------------------------------------------------ class OutputAttributeInfo(NamedTuple): """Information relating to an output node attribute.""" array_shape_source: Optional[OutputArrayShapeSource] bundle_type_source: Optional[OutputBundleTypeSource] bundle_type_explicit: Optional[str] = None class AttributeInfo(NamedTuple): """Information relating to a node attribute. This struct contains all the metadata required by the node to initialize and evaluate. This includes compiling the kernel and initializing the Inputs and Outputs structs that are then passed to the kernel as parameters. We don't directly store the array shape, if any, since it is possible that it might vary between each evaluation of the node's compute. Instead, we store which method to use to infer the array's shape and let the node determine the actual shape during each compute step. Note ---- The `warp_type` member represents the type of the kernel parameter corresdonding to that attribute. If the attribute is a bundle, then it is expected to be a `wp.struct` holding the values of the bundle, unless the bundle is of type :class:`Array`, in which case `warp_type` should be a standard `wp.array`. """ port_type: og.AttributePortType base_name: str og_type: og.Type warp_type: type output: Optional[OutputAttributeInfo] = None @property def name(self) -> str: return attr_join_name(self.port_type, self.base_name) @property def og_data_type(self) -> og.Type: return og.Type( self.og_type.base_type, tuple_count=self.og_type.tuple_count, array_depth=0, role=self.og_type.role, ) @property def is_array(self) -> bool: return self.og_type.array_depth > 0 @property def is_bundle(self) -> bool: return self.og_type == ATTR_BUNDLE_TYPE @property def dim_count(self) -> int: if self.is_array: return self.warp_type.ndim return 0 @property def warp_data_type(self) -> type: if self.is_array: return self.warp_type.dtype return self.warp_type @property def warp_type_name(self) -> str: if self.is_bundle: return self.warp_type.cls.__name__ return get_warp_type_from_data_type_name( self.warp_data_type.__name__, dim_count=self.dim_count, as_str=True, ) @property def warp_data_type_name(self) -> str: if self.is_bundle: return self.warp_type.cls.__name__ return get_warp_type_from_data_type_name( self.warp_data_type.__name__, dim_count=0, as_str=True, ) def gather_attribute_infos( node: og.Node, db_inputs: Any, db_outputs: Any, attr_descs: Mapping[str, UserAttributeDesc], kernel_dim_count: int, ) -> Mapping[og.AttributePortType, Tuple[AttributeInfo, ...]]: """Gathers the information for each user attribute. See also: :class:`AttributeInfo`. """ def extract_partial_info_from_attr(attr: og.Attribute) -> Tuple[Any, ...]: """Extract a partial information set from an attribute.""" name = attr_get_name(attr) base_name = attr_get_base_name(attr) og_type = attr.get_resolved_type() is_array = og_type.array_depth > 0 return (name, base_name, og_type, is_array) # Retrieve the user attributes defined on the node. attrs = tuple(x for x in node.get_attributes() if x.is_dynamic()) # Gather the information for the input attributes. input_attr_infos = [] for attr in attrs: if attr.get_port_type() != _ATTR_PORT_TYPE_INPUT: continue (name, base_name, og_type, is_array) = extract_partial_info_from_attr(attr) og_data_type = og.Type( og_type.base_type, tuple_count=og_type.tuple_count, array_depth=0, role=og_type.role, ) input_attr_infos.append( AttributeInfo( port_type=_ATTR_PORT_TYPE_INPUT, base_name=base_name, og_type=og_type, warp_type=type_convert_og_to_warp( og_data_type, dim_count=int(is_array), ), ) ) # Gather the information for the output attributes. output_attr_infos = [] for attr in attrs: if attr.get_port_type() != _ATTR_PORT_TYPE_OUTPUT: continue (name, base_name, og_type, is_array) = extract_partial_info_from_attr(attr) desc = attr_descs.get(name) if desc is None: # Fallback for nodes created before the attribute description # feature was implemented. array_shape_source = OutputArrayShapeSource.AS_INPUT_OR_AS_KERNEL else: array_shape_source = desc.array_shape_source if array_shape_source == OutputArrayShapeSource.AS_INPUT_OR_AS_KERNEL: # Check if we have an input attribute with a matching name, # in which case we use its array dimension count. try: dim_count = next(x.dim_count for x in input_attr_infos if x.base_name == base_name) except StopIteration: # Fallback to using the kernel's dimension count. dim_count = kernel_dim_count elif array_shape_source == OutputArrayShapeSource.AS_KERNEL: dim_count = kernel_dim_count else: assert False, "Unexpected array shape source method '{}'.".format(array_shape_source) og_data_type = og.Type( og_type.base_type, tuple_count=og_type.tuple_count, array_depth=0, role=og_type.role, ) output_attr_infos.append( AttributeInfo( port_type=_ATTR_PORT_TYPE_OUTPUT, base_name=base_name, og_type=og_type, warp_type=type_convert_og_to_warp( og_data_type, dim_count=dim_count, ), output=OutputAttributeInfo( array_shape_source=array_shape_source, bundle_type_source=OutputBundleTypeSource.AS_INPUT, ), ) ) return { _ATTR_PORT_TYPE_INPUT: tuple(input_attr_infos), _ATTR_PORT_TYPE_OUTPUT: tuple(output_attr_infos), } # Kernel Code # ------------------------------------------------------------------------------ _STRUCT_DECLARATION_CODE_TEMPLATE = """@wp.struct class {name}: {members} """ def _generate_struct_declaration_code(warp_struct: wp.struct) -> str: """Generates the code declaring a Warp struct.""" lines = [] for label, var in warp_struct.vars.items(): warp_type = var.type if isinstance(warp_type, wp.array): warp_data_type = warp_type.dtype dim_count = warp_type.ndim else: warp_data_type = warp_type dim_count = 0 warp_type_name = get_warp_type_from_data_type_name( warp_data_type.__name__, dim_count=dim_count, as_str=True, ) lines.append(" {}: {}".format(label, warp_type_name)) return _STRUCT_DECLARATION_CODE_TEMPLATE.format( name=warp_struct.cls.__name__, members="\n".join(lines), ) _HEADER_CODE_TEMPLATE = """import warp as wp {declarations} @wp.struct class Inputs: {inputs} pass @wp.struct class Outputs: {outputs} pass """ def _generate_header_code( attr_infos: Mapping[og.AttributePortType, Tuple[AttributeInfo, ...]], ) -> str: """Generates the code header based on the node's attributes.""" # Retrieve all the Warp struct types corresponding to bundle attributes. struct_types = {x.warp_type_name: x.warp_type for _, v in attr_infos.items() for x in v if x.is_bundle} # Generate the code that declares the Warp structs found. declarations = [""] declarations.extend(_generate_struct_declaration_code(x) for _, x in struct_types.items()) # Generate the lines of code declaring the members for each port type. lines = {k: tuple(" {}: {}".format(x.base_name, x.warp_type_name) for x in v) for k, v in attr_infos.items()} # Return the template code populated with the members. return _HEADER_CODE_TEMPLATE.format( declarations="\n".join(declarations), inputs="\n".join(lines.get(_ATTR_PORT_TYPE_INPUT, ())), outputs="\n".join(lines.get(_ATTR_PORT_TYPE_OUTPUT, ())), ) def _get_user_code(code_provider: str, code_str: str, code_file: str) -> str: """Retrieves the code provided by the user.""" if code_provider == "embedded": return code_str if code_provider == "file": with open(code_file, "r") as f: return f.read() assert False, "Unexpected code provider '{}'.".format(code_provider) # Kernel Module # ------------------------------------------------------------------------------ def _load_code_as_module(code: str, name: str) -> Any: """Loads a Python module from the given source code.""" # It's possible to use the `exec()` built-in function to create and # populate a Python module with the source code defined in a string, # however warp requires access to the source code of the kernel's # function, which is only available when the original source file # pointed by the function attribute `__code__.co_filename` can # be opened to read the lines corresponding to that function. # As such, we must write the source code into a temporary file # on disk before importing it as a module and having the function # turned into a kernel by warp's mechanism. # Create a temporary file. file, file_path = tempfile.mkstemp(suffix=".py") try: # Save the embedded code into the temporary file. with os.fdopen(file, "w") as f: f.write(code) # Import the temporary file as a Python module. spec = importlib.util.spec_from_file_location(name, file_path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) finally: # The resulting Python module is stored into memory as a bytcode # object and the kernel function has already been parsed by warp # as long as it was correctly decorated, so it's now safe to # clean-up the temporary file. os.remove(file_path) return module def initialize_kernel_module( attr_infos: Mapping[og.AttributePortType, Tuple[AttributeInfo, ...]], code_provider: str, code_str: str, code_file: str, ) -> wp.context.Module: # Ensure that all output parameters are arrays. Writing to non-array # types is not supported as per CUDA's design. invalid_attrs = tuple(x.name for x in attr_infos[_ATTR_PORT_TYPE_OUTPUT] if not x.is_array and not x.is_bundle) if invalid_attrs: raise RuntimeError( "Output attributes are required to be arrays or bundles but " "the following attributes are not: {}.".format(", ".join(invalid_attrs)) ) # Retrieve the kernel code to evaluate. code_header = _generate_header_code(attr_infos) user_code = _get_user_code(code_provider, code_str, code_file) code = "{}\n{}".format(code_header, user_code) # Create a Python module made of the kernel code. # We try to keep its name unique to ensure that it's not clashing with # other kernel modules from the same session. uid = hashlib.blake2b(bytes(code, encoding="utf-8"), digest_size=8) module_name = "warp-kernelnode-{}".format(uid.hexdigest()) kernel_module = _load_code_as_module(code, module_name) # Validate the module's contents. if not hasattr(kernel_module, "compute"): raise RuntimeError("The code must define a kernel function named 'compute'.") if not isinstance(kernel_module.compute, wp.context.Kernel): raise RuntimeError("The 'compute' function must be decorated with '@wp.kernel'.") # Configure warp to only compute the forward pass. wp.set_module_options({"enable_backward": False}, module=kernel_module) return kernel_module # Data I/O # ------------------------------------------------------------------------------ def _infer_output_array_shape( attr_info: AttributeInfo, input_attr_infos: Tuple[AttributeInfo, ...], kernel_inputs: Any, kernel_shape: Sequence[int], ) -> Tuple[int, ...]: if attr_info.output.array_shape_source == OutputArrayShapeSource.AS_INPUT_OR_AS_KERNEL: # Check if we have an input attribute with a matching name, # in which case we use its array shape. try: ref_attr_base_name = next( x.base_name for x in input_attr_infos if (x.base_name == attr_info.base_name and x.is_array and x.dim_count == attr_info.dim_count) ) return getattr(kernel_inputs, ref_attr_base_name).shape except StopIteration: # Fallback to using the kernel's shape. return tuple(kernel_shape) if attr_info.output.array_shape_source == OutputArrayShapeSource.AS_KERNEL: return tuple(kernel_shape) assert False, "Unexpected array shape source method '{}'.".format(attr_info.output.array_shape_source) class KernelArgsConfig(NamedTuple): """Configuration for resolving kernel arguments.""" input_bundle_handlers: Optional[Mapping[str, Callable]] = None output_bundle_handlers: Optional[Mapping[str, Callable]] = None def get_kernel_args( db_inputs: Any, db_outputs: Any, attr_infos: Mapping[og.AttributePortType, Tuple[AttributeInfo, ...]], kernel_module: Any, kernel_shape: Sequence[int], device: Optional[wp.context.Device] = None, config: Optional[KernelArgsConfig] = None, ) -> Tuple[Any, Any]: """Retrieves the in/out argument values to pass to the kernel.""" if device is None: device = wp.get_device() if config is None: config = KernelArgsConfig() # Initialize the kernel's input data. inputs = kernel_module.Inputs() for info in attr_infos[_ATTR_PORT_TYPE_INPUT]: # Retrieve the input attribute value and cast it to # the corresponding warp type. if info.is_array: value = getattr(db_inputs, info.base_name) # The array value might define 2 dimensions when tuples such as # wp.vec3 are used as data type, so we preserve only the first # dimension to retrieve the actual shape since OmniGraph only # supports 1D arrays anyways. shape = value.shape[:1] value = attr_cast_array_to_warp( value, info.warp_data_type, shape, device, ) elif info.is_bundle: raise NotImplementedError("Bundle attributes are not yet supported.") else: value = getattr(db_inputs, info.base_name) # Store the result in the inputs struct. setattr(inputs, info.base_name, value) # Initialize the kernel's output data. outputs = kernel_module.Outputs() for info in attr_infos[_ATTR_PORT_TYPE_OUTPUT]: # Retrieve the output attribute value and cast it to the corresponding # warp type. if info.is_array: shape = _infer_output_array_shape( info, attr_infos[_ATTR_PORT_TYPE_INPUT], inputs, kernel_shape, ) # Allocate a buffer for the array. size = functools.reduce(operator.mul, shape) setattr(db_outputs, "{}_size".format(info.base_name), size) value = getattr(db_outputs, info.base_name) value = attr_cast_array_to_warp( value, info.warp_data_type, shape, device, ) elif info.is_bundle: raise NotImplementedError("Bundle attributes are not yet supported.") else: assert False, "Output attributes are expected to be arrays or bundles." # Store the result in the outputs struct. setattr(outputs, info.base_name, value) return (inputs, outputs) def write_output_attrs( db_outputs: Any, attr_infos: Mapping[og.AttributePortType, Tuple[AttributeInfo, ...]], kernel_outputs: Any, device: Optional[wp.context.Device] = None, ) -> None: """Writes the output values to the node's attributes.""" if device is None: device = wp.get_device() if device.is_cuda: # CUDA attribute arrays are directly being written to by Warp. return for info in attr_infos[_ATTR_PORT_TYPE_OUTPUT]: value = getattr(kernel_outputs, info.base_name) setattr(db_outputs, info.base_name, value) # Validation # ------------------------------------------------------------------------------ def validate_input_arrays( node: og.Node, attr_infos: Mapping[og.AttributePortType, Tuple[AttributeInfo, ...]], kernel_inputs: Any, ) -> None: """Validates array input attributes.""" for info in attr_infos[_ATTR_PORT_TYPE_INPUT]: value = getattr(kernel_inputs, info.base_name) if not isinstance(value, wp.array): continue # Ensure that all array input attributes are not NULL, # unless they are set as being optional. attr = og.Controller.attribute(info.name, node) if not attr.is_optional_for_compute and not value.ptr: raise RuntimeError("Empty value for non-optional attribute '{}'.".format(info.name)) # Node's Internal State # ------------------------------------------------------------------------------ class InternalStateBase: """Base class for the node's internal state.""" def __init__(self) -> None: self._code_provider = None self._code_str = None self._code_file = None self._code_file_timestamp = None self.attr_infos = None self.kernel_module = None self.is_valid = False def needs_initialization( self, db: Any, check_file_modified_time: bool, ) -> bool: """Checks if the internal state needs to be (re)initialized.""" if self.is_valid: # If everything is in order, we only need to recompile the kernel # when attributes are removed, since adding new attributes is not # a breaking change. if self.kernel_module is None or UserAttributesEvent.REMOVED & db.state.userAttrsEvent: return True else: # If something previously went wrong, we always recompile the kernel # when attributes are edited, in case it might fix code that # errored out due to referencing a non-existing attribute. if db.state.userAttrsEvent != UserAttributesEvent.NONE: return True if self.attr_infos is None: return True if self._code_provider != db.inputs.codeProvider: return True if self._code_provider == "embedded": if self._code_str != db.inputs.codeStr: return True elif self._code_provider == "file": if self._code_file != db.inputs.codeFile or ( check_file_modified_time and (self._code_file_timestamp != os.path.getmtime(self._code_file)) ): return True else: assert False, ("Unexpected code provider '{}'.".format(self._code_provider),) return False def initialize(self, db: Any) -> bool: """Initialize the internal state and recompile the kernel.""" # Cache the node attribute values relevant to this internal state. # They're the ones used to check whether this state is outdated or not. self._code_provider = db.inputs.codeProvider self._code_str = db.inputs.codeStr self._code_file = db.inputs.codeFile return True