diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h b/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h index 8c9c31c33d5..216f48b2991 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h @@ -51,6 +51,7 @@ struct TensorRTBlobHeader { bool hardware_compatible = false; ///< Engine built hardware-compatible; false when unstated. int device_id = 0; ///< Device the engine was built for; 0 when unstated. std::string engine_key; ///< Named-data key of an engine stored outside the blob; empty when it is inline. + std::string tensorrt_version; ///< TensorRT version that exported the engine; empty when the blob does not say. /** * @brief Start of the serialized engine inside a blob. diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index f607cc7f1e8..fda72dbe62f 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -194,10 +194,17 @@ but corrupted engine bytes can load and return wrong answers. ### Load errors +- With no usable CUDA device, for example under `CUDA_VISIBLE_DEVICES=""`, the delegate reports + itself unavailable, so loading returns `Error::NotFound`. The log says CUDA is unavailable and + gives CUDA's reason. - Loading returns `Error::MemoryAllocationFailed` when TensorRT runs out of memory while it deserializes the engine, applies the weight streaming budget, or creates the execution context. The log then shows the free and total GPU memory. Other TensorRT load failures return `Error::InvalidProgram`, with TensorRT's reason in the log. +- An engine loads only on the TensorRT version that built it, unless it was built version + compatible. When TensorRT cannot read an engine, the log names the TensorRT version that exported + it and the one the delegate runs on. A program exported before the export side recorded that + version says so instead. - An engine saved without its weights, for example compiled with `immutable_weights=False, strip_engine_weights=True`, is refused with `Error::InvalidProgram`, because this delegate cannot refit it. To check, the delegate asks TensorRT how many weight bytes diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index dca7601b46e..4be9d71a48c 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -130,6 +130,32 @@ Error report_load_error(const char* operation) { return out_of_memory ? Error::MemoryAllocationFailed : Error::InvalidProgram; } +// TensorRT reads an engine only with the version that built it, unless it was built version +// compatible, and its own error then names internal serialization tags rather than versions. +void log_tensorrt_versions(const std::string& exported_version) { + // TensorRT 10 and later encode the version as major * 10000 + minor * 100 + patch. + const int32_t version = getInferLibVersion(); + const int major = version / 10000; + const int minor = version / 100 % 100; + const int patch = version % 100; + if (exported_version.empty()) { + ET_LOG( + Error, + "TensorRTBackend::init: this runtime is TensorRT %d.%d.%d, and the program does not record which TensorRT exported the engine", + major, + minor, + patch); + } else { + ET_LOG( + Error, + "TensorRTBackend::init: the engine was exported with TensorRT %s, and this runtime is TensorRT %d.%d.%d", + exported_version.c_str(), + major, + minor, + patch); + } +} + Error check_registration() { if (kRegistrationResult != Error::Ok) { ET_LOG(Error, "TensorRTBackend registration failed: %s", ::executorch::runtime::to_string(kRegistrationResult)); @@ -253,6 +279,18 @@ nvinfer1::IRuntime* TensorRTBackend::shared_runtime() { // Retried while null, so one transient failure does not disable the backend for the process. if (runtime == nullptr) { const std::shared_lock capture_lock(cuda_graph_capture_mutex()); + // Without a usable CUDA device, createInferRuntime logs an error and then returns a pointer that + // is not null and is not a runtime either (seen with TensorRT 11.4), so check CUDA first. + int device_count = 0; + const cudaError_t cuda_err = cudaGetDeviceCount(&device_count); + if (cuda_err != cudaSuccess) { + cudaGetLastError(); + ET_LOG( + Error, + "TensorRTBackend: cannot create the TensorRT runtime, because CUDA is unavailable: %s", + cudaGetErrorString(cuda_err)); + return nullptr; + } if (logger == nullptr) { logger = new TRTLogger(); } @@ -398,11 +436,16 @@ Error load_engine( const void* data, uint64_t size, const WsBudget& request, + const std::string& exported_version, std::shared_ptr& out) { last_load_error = {}; std::shared_ptr engine(runtime.deserializeCudaEngine(data, size), TRTDeleter{}); if (engine == nullptr) { - return report_load_error("deserialize TensorRT engine"); + const Error err = report_load_error("deserialize TensorRT engine"); + if (err != Error::MemoryAllocationFailed) { + log_tensorrt_versions(exported_version); + } + return err; } // A plan serialized without its weights (built with STRIP_PLAN, or saved with EXCLUDE_WEIGHTS) // deserializes and runs without an error, but computes with placeholder weights, and this backend @@ -452,6 +495,7 @@ Error acquire_shared_engine( uint64_t size, int device_id, const WsBudget& request, + const std::string& exported_version, std::shared_ptr& out) { const SharedEngineKey key = shared_engine_key(data, size, device_id, request.valid ? request.bytes : -1); SharedEngines& engines = shared_engines(); @@ -464,7 +508,7 @@ Error acquire_shared_engine( // Declared out here so a copy that lost the race is freed after the lock is released. std::shared_ptr loaded; if (out == nullptr) { - const Error err = load_engine(runtime, data, size, request, loaded); + const Error err = load_engine(runtime, data, size, request, exported_version, loaded); if (err != Error::Ok) { // A racing load of the same bytes may have published its engine meanwhile, and this load may // have failed only because the device had no room for a second copy. @@ -1410,9 +1454,10 @@ Result TensorRTBackend::init( { const std::shared_lock capture_lock(cuda_graph_capture_mutex()); if (share) { - err = acquire_shared_engine(*runtime, engine_data, engine_size, handle->device_id, ws_request, handle->engine); + err = acquire_shared_engine( + *runtime, engine_data, engine_size, handle->device_id, ws_request, header.tensorrt_version, handle->engine); } else { - err = load_engine(*runtime, engine_data, engine_size, ws_request, handle->engine); + err = load_engine(*runtime, engine_data, engine_size, ws_request, header.tensorrt_version, handle->engine); } // The engine now lives on the GPU, so release the host copy before building the context. named_engine.reset(); diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index a5e2aff3a27..6a9484e52c6 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -487,6 +487,13 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso return false; } } + // Only an error message reads it, so a value that is not a string is ignored rather than refused. + out.tensorrt_version.clear(); + const std::size_t version_key = find_top_level_key(json, "\"tensorrt_version\""); + if (version_key != std::string::npos && + parse_string(json, skip_ws(json, json.find(':', version_key) + 1), out.tensorrt_version) == std::string::npos) { + out.tensorrt_version.clear(); + } const std::size_t hw_key = find_top_level_key(json, "\"hardware_compatible\""); const std::size_t device_key = find_top_level_key(json, "\"device_id\""); return parse_bool_after_key(json, hw_key, "\"hardware_compatible\"", out.hardware_compatible) && diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index a32301aa910..8e3bdaa5cd3 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -225,8 +225,9 @@ def cross_compile_for_windows( disabled_constant_fold_exclusions (Collection[str]): IDs of predefined Torch-TensorRT rules to disable. Rules are registered by lowering implementations and enabled by default; compilation users do not - need to register them. Disabling a rule makes its matching nodes - foldable. Supported IDs: ``"attention_mask_arange"``. + need to register them. Disabling a rule makes the nodes that only + it matches foldable. Supported IDs: ``"attention_mask_arange"``, + ``"scatter_add_positions"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False. @@ -428,12 +429,16 @@ def cross_compile_for_windows( # Apply lowering on the graph module. Note: constant_fold runs inside post_lowering and requires # module parameters to still be on GPU, so we must not deallocate before this call. + unlowered_gm = gm gm = post_lowering(gm, settings) logger.debug(f"CPU memory usage after post_lowering: {get_cpu_memory_usage()} MB") logger.debug("Lowered Input graph: " + str(gm.graph)) # Move the weights in the state_dict to CPU if offload_module_to_cpu: + # Complex lowering can return a new module without the weights that only + # constant folding read, so offload the module lowering started from too. + deallocate_module(unlowered_gm) deallocate_module(gm) logger.info( "The PyTorch model was moved to the CPU to allocate all GPU memory to TensorRT. To retain the model on the GPU, set offload_module_to_cpu=False" @@ -637,8 +642,9 @@ def compile( disabled_constant_fold_exclusions (Collection[str]): IDs of predefined Torch-TensorRT rules to disable. Rules are registered by lowering implementations and enabled by default; compilation users do not - need to register them. Disabling a rule makes its matching nodes - foldable. Supported IDs: ``"attention_mask_arange"``. + need to register them. Disabling a rule makes the nodes that only + it matches foldable. Supported IDs: ``"attention_mask_arange"``, + ``"scatter_add_positions"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False. @@ -909,6 +915,7 @@ def compile( # Apply lowering on the graph module. Note: constant_fold runs inside post_lowering and requires # module parameters to still be on GPU, so we must not deallocate before this call. + unlowered_gm = gm gm = post_lowering(gm, settings) logger.debug(f"CPU memory usage after post_lowering: {get_cpu_memory_usage()} MB") logger.debug("Lowered Input graph: " + str(gm.graph)) @@ -921,6 +928,9 @@ def compile( # Move the weights in the state_dict to CPU if offload_module_to_cpu: + # Complex lowering can return a new module without the weights that only + # constant folding read, so offload the module lowering started from too. + deallocate_module(unlowered_gm) deallocate_module(gm) logger.info( "The PyTorch model was moved to the CPU to allocate all GPU memory to TensorRT. To retain the model on the GPU, set offload_module_to_cpu=False" @@ -2102,8 +2112,9 @@ def convert_exported_program_to_serialized_trt_engine( disabled_constant_fold_exclusions (Collection[str]): IDs of predefined Torch-TensorRT rules to disable. Rules are registered by lowering implementations and enabled by default; compilation users do not - need to register them. Disabling a rule makes its matching nodes - foldable. Supported IDs: ``"attention_mask_arange"``. + need to register them. Disabling a rule makes the nodes that only + it matches foldable. Supported IDs: ``"attention_mask_arange"``, + ``"scatter_add_positions"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False. diff --git a/py/torch_tensorrt/dynamo/_refit.py b/py/torch_tensorrt/dynamo/_refit.py index 3762b3aed0c..79b3ab4b1bc 100644 --- a/py/torch_tensorrt/dynamo/_refit.py +++ b/py/torch_tensorrt/dynamo/_refit.py @@ -460,9 +460,7 @@ def refit_module_weights( if verify_output and arg_inputs is not None: new_gm.to(to_torch_device(settings.device)) # move to device for inference - # A complex placeholder stays complex in the lowered reference module; the - # unpacking to real happens inside the graph, so both modules take the same - # inputs the caller passed. + _restore_complex_io(new_gm) outputs_match = check_module_output( new_module=new_gm, refitted_module=compiled_module, @@ -479,6 +477,42 @@ def refit_module_weights( return compiled_module +def _restore_complex_io(gm: torch.fx.GraphModule) -> None: + """Make the lowered reference module take and return complex tensors, as the + compiled module does, so the output check can compare the two. + + The legacy complex rewriter replaces a complex input with a real + "_unpacked_complex" one, and both complex lowerings return a complex + output as a real tensor with a trailing dimension of 2. + """ + graph = gm.graph + unpacked = {f"{n}_unpacked_complex" for n in gm.meta.get("complex_input_names", [])} + for node in list(graph.nodes): + if node.op == "placeholder" and node.name in unpacked: + with graph.inserting_after(node): + real = graph.call_function(torch.ops.aten.view_as_real.default, (node,)) + node.replace_all_uses_with(real) + real.args = (node,) + output_node = list(graph.nodes)[-1] + output_arg = output_node.args[0] + outputs = ( + list(output_arg) if isinstance(output_arg, (list, tuple)) else [output_arg] + ) + for idx in gm.meta.get("complex_output_indices", []): + with graph.inserting_before(output_node): + outputs[idx] = graph.call_function( + torch.ops.aten.view_as_complex.default, (outputs[idx],) + ) + output_node.args = ( + ( + type(output_arg)(outputs) + if isinstance(output_arg, (list, tuple)) + else outputs[0] + ), + ) + gm.recompile() + + # Util functions ----------- import base64 diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py index 206f5354cee..1f7fabfbd0d 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -159,8 +159,9 @@ class CompilationSettings: disabled_constant_fold_exclusions (Collection[str]): IDs of predefined Torch-TensorRT rules to disable. Rules are registered by lowering implementations and enabled by default; compilation users do not - need to register them. Disabling a rule makes its matching nodes - foldable. Supported IDs: ``"attention_mask_arange"``. + need to register them. Disabling a rule makes the nodes that only + it matches foldable. Supported IDs: ``"attention_mask_arange"``, + ``"scatter_add_positions"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False. diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 7cda4e4dc7a..dec408c2fb4 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -1753,7 +1753,27 @@ def aten_ops_scatter( ) -@dynamo_tensorrt_converter(torch.ops.aten.select.int, supports_dynamic_shapes=True) +def select_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Keep a select of a uint8 tensor on the PyTorch path. + + TensorRT allows uint8 only at engine inputs and at outputs made by a cast or + identity layer, so the gather's uint8 result fails the engine build, as it does + for aten.gather. + """ + meta = getattr(args_bounds_check(node.args, 0), "meta", {}) + if getattr(meta.get("val", meta.get("tensor_meta")), "dtype", None) == torch.uint8: + _LOGGER.debug("select of a uint8 tensor is not supported, falling back") + return False + return True + + +@dynamo_tensorrt_converter( + torch.ops.aten.select.int, + capability_validator=select_validator, + supports_dynamic_shapes=True, +) def aten_ops_select( ctx: ConversionContext, target: Target, @@ -2142,12 +2162,42 @@ def aten_ops_permute( ) +def _truncatable_float64_copy(node: Node, source: Argument) -> bool: + """Whether a float64 copy can run as a float32 copy and give PyTorch's result. + + The source has to be a type float32 holds exactly, the same rule impl.cast.to_copy + applies, and the PyTorch result has to have the contiguous layout a TensorRT copy + produces. The check does not depend on where the copy sits, because partitioning + can turn any copy into an engine output and the converter is validated again then. + """ + input_meta = source.meta.get("val") if isinstance(source, Node) else None + if ( + not isinstance(input_meta, torch.Tensor) + or input_meta.dtype not in impl.cast.FLOAT64_TRUNCATABLE_SOURCES + ): + return False + # Alone on a placeholder after partitioning, a dynamic copy must also pass + # _dynamic_placeholder_copy_supported, which refuses any explicit memory_format. + if node.kwargs.get("memory_format") is not None and node_has_dynamic_shapes(node): + return False + output_meta = node.meta.get("val") + if not isinstance(output_meta, torch.Tensor): + return False + # A data-dependent stride cannot be decided here, so leave the copy to PyTorch. + try: + return bool(output_meta.is_contiguous()) + except Exception: + return False + + def to_copy_dtype_validator( placeholder_only: bool, settings: Optional[CompilationSettings] = None ) -> Callable[[Node, CompilationSettings], bool]: """Return validator for to_copy node with placeholder restrictions""" - def validate_dtype(to_copy_node: Node) -> bool: + def validate_dtype( + to_copy_node: Node, settings: Optional[CompilationSettings] + ) -> bool: """Returns true if the to_copy node can be converted to TRT Based on data type being casted to @@ -2187,6 +2237,15 @@ def validate_dtype(to_copy_node: Node) -> bool: if "dtype" in to_copy_node.kwargs: if to_copy_node.kwargs["dtype"] in allowed_casts: return True + # truncate_double means the caller accepts float32 for float64, and + # the converter maps it. + elif ( + to_copy_node.kwargs["dtype"] is torch.float64 + and settings is not None + and settings.truncate_double + and _truncatable_float64_copy(to_copy_node, to_copy_node.args[0]) + ): + return True else: _LOGGER.debug( f"_to_copy converter rejected node {to_copy_node} with dtype {to_copy_node.kwargs['dtype']}" @@ -2206,7 +2265,7 @@ def validator( """ # The placeholder restriction is satsfied if placeholder_only is the same # truth value as is_only_operator_on_placeholder(to_copy_node) - return validate_dtype(to_copy_node) and ( + return validate_dtype(to_copy_node, settings) and ( (not placeholder_only) ^ is_only_operator_on_placeholder(to_copy_node) ) @@ -4750,7 +4809,22 @@ def aten_ops_trunc( ) -@dynamo_tensorrt_converter(torch.ops.aten.copy.default, supports_dynamic_shapes=True) +def _copy_validator(node: Node, settings: Optional[CompilationSettings] = None) -> bool: + output_meta = node.meta.get("val") + if not isinstance(output_meta, torch.Tensor) or output_meta.dtype != torch.float64: + return True + return ( + settings is not None + and settings.truncate_double + and _truncatable_float64_copy(node, node.args[1]) + ) + + +@dynamo_tensorrt_converter( + torch.ops.aten.copy.default, + capability_validator=_copy_validator, + supports_dynamic_shapes=True, +) @enforce_tensor_types( { 1: (TRTTensor,), @@ -4764,6 +4838,25 @@ def aten_ops_copy( name: str, ) -> Union[TRTTensor, Sequence[TRTTensor]]: dest, src = args[0], args[1] + # copy broadcasts src, so the result always has the shape of dest. + if tuple(src.shape) != tuple(dest.shape): + src = impl.slice.expand( + ctx, + target, + SourceIR.ATEN, + f"{name}_expand", + src, + [ + ( + impl.shape.shape( + ctx, target, SourceIR.ATEN, f"{name}_dest_dim{i}", dest, i + ) + if size == DYNAMIC_DIM + else size + ) + for i, size in enumerate(dest.shape) + ], + ) return impl.cast.to_copy( ctx, target, diff --git a/py/torch_tensorrt/dynamo/conversion/impl/cast.py b/py/torch_tensorrt/dynamo/conversion/impl/cast.py index f4953f8ce91..478f274c19f 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/cast.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/cast.py @@ -21,6 +21,20 @@ LOGGER: logging.Logger = logging.getLogger(__name__) +# With truncate_double, a float64 result from these sources runs in float32, which +# holds their values exactly. A float64 source is already float32 inside TensorRT. +FLOAT64_TRUNCATABLE_SOURCES = frozenset( + { + torch.bool, + torch.int8, + torch.uint8, + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + } +) + def to_copy( ctx: ConversionContext, @@ -34,6 +48,15 @@ def to_copy( if not isinstance(input, TRTTensor): input = get_trt_tensor(ctx, input, f"{name}_copy_tensor") + # TensorRT has no float64. With truncate_double the caller accepts float32. + if ( + ctx.compilation_settings.truncate_double + and _enums.dtype._from(dtype) == _enums.dtype.f64 + and _enums.dtype._from(input.dtype).to(torch.dtype) + in FLOAT64_TRUNCATABLE_SOURCES + ): + dtype = torch.float32 + # If cast is forced, insert identity layer regardless of whether the dtype # doesn't change if force_layer: diff --git a/py/torch_tensorrt/dynamo/conversion/impl/select.py b/py/torch_tensorrt/dynamo/conversion/impl/select.py index d473842cc81..e0b5c4f6e82 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/select.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/select.py @@ -789,13 +789,12 @@ def index_put_converter( 0, ) ) - # When any index has a dynamic size, use the first dynamic value - # (all valid indices are guaranteed to have the same N after broadcasting). - # Python's max() cannot compare ITensors, so we avoid it when dynamic. - if any(isinstance(s, ITensor) for s in index_shapes): - N = next(s for s in index_shapes if isinstance(s, ITensor)) - else: - N = max(index_shapes) if index_shapes else 1 + # Index lengths broadcast together. A dynamic length must be 1 or equal a + # static length larger than 1, so that static length is N. Otherwise use + # the first dynamic length, since Python's max() cannot compare ITensors. + static_n = max((s for s in index_shapes if isinstance(s, int)), default=1) + dynamic_n = [s for s in index_shapes if isinstance(s, ITensor)] + N = dynamic_n[0] if dynamic_n and static_n == 1 else static_n else: N = 1 @@ -999,8 +998,6 @@ def index_put_converter( # Define the expected shape based on (N,) + F_shape_values. # N may be a ITensor when it comes from a dynamic source (e.g. nonzero). - # Pass it directly — impl.slice.expand handles ITensor shape elements and - # will emit a stride-0 broadcast when expanding a size-1 dim to a dynamic N. expected_shape = (N,) + tuple(F_shape_values) # Broadcast 'values' to match the expected shape @@ -1117,17 +1114,15 @@ def index_put_converter( values_shape_padded = [1] * ( len(expected_shape) - len(values.shape) ) + list(values.shape) - broadcast_shape = [] for exp_dim, val_dim in zip(expected_shape, values_shape_padded): - if val_dim == DYNAMIC_DIM or exp_dim == DYNAMIC_DIM: - broadcast_shape.append(-1) - elif val_dim == 1 or exp_dim == val_dim: - broadcast_shape.append(exp_dim) - else: + if not ( + isinstance(exp_dim, ITensor) + or DYNAMIC_DIM in (exp_dim, val_dim) + or val_dim in (1, exp_dim) + ): raise ValueError( f"Cannot broadcast {values.shape} to {expected_shape}" ) - values_expanded = impl.slice.expand( ctx, target, diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py index a89113b6c3f..eab48b9cd00 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py @@ -157,8 +157,11 @@ def slice_op( # TODO: This should be slice not whatever is in base output_shape = list(input.shape) if input.shape[dim] != -1 and isinstance(start, int) and isinstance(stop, int): - start = get_positive_dim(start, input.shape[dim]) - stop = get_positive_dim(stop, input.shape[dim]) + # Clamp out-of-range bounds like Python slicing instead of wrapping them + # like a dimension index: x[-15:] on a size-10 dim is the whole dim. + dim_size = input.shape[dim] + start = max(start + dim_size, 0) if start < 0 else min(start, dim_size) + stop = max(stop + dim_size, 0) if stop < 0 else min(stop, dim_size) start_slice[dim] = start else: # the start and stop or None is dynamic along dim or or start or stop is an ITensor @@ -272,7 +275,7 @@ def slice_op( # TODO: This should be slice not whatever is in base layer.set_input(3, stride_slice_tensor) return layer.get_output(0) - output_shape[dim] = math.ceil((stop - start) / step) + output_shape[dim] = max(math.ceil((stop - start) / step), 0) return slice( ctx, target, source_ir, name, input, start_slice, output_shape, stride_slice ) @@ -307,6 +310,14 @@ def expand( # After the above padding, the shape and tensor rank must be equal assert len(input_t.shape) == shape_rank + # PyTorch only expands a dim larger than 1 to its own size, so a runtime + # target there must equal the static input dim. Keep the dim instead of + # broadcasting its first element. + shape = [ + i if isinstance(o, TRTTensor) and i not in (1, DYNAMIC_DIM) else o + for i, o in zip(input_t.shape, shape) + ] + # Configure the start, strides and output shape tensors start = tuple([0] * shape_rank) @@ -314,7 +325,7 @@ def expand( # stride should be 1 for all non-broadcasted dims stride = [] input_tensor_shape = tuple(input_t.shape) - for i, o in zip(input_tensor_shape, shape): + for k, (i, o) in enumerate(zip(input_tensor_shape, shape)): # If input dim and target shape dim are static, broadcast if they are not equal # If input dim is known and target shape dim is dynamic we treat it as a broadcasted dim if ( @@ -326,6 +337,15 @@ def expand( stride.append(int(i == o)) elif isinstance(i, int) and i != DYNAMIC_DIM and isinstance(o, TRTTensor): stride.append(0) + elif i == DYNAMIC_DIM and isinstance(o, int) and o != DYNAMIC_DIM: + # A runtime dim is either 1, and broadcasts, or already the target size. + size = get_shape( + ctx, target, source_ir, f"{name}_stride_dim{k}", input_t, k + ) + is_full = impl.elementwise.gt( + ctx, target, source_ir, f"{name}_stride_gt{k}", size, 1 + ) + stride.append(cast_trt_tensor(ctx, is_full, trt.int32, f"{name}_stride{k}")) else: # No broadcasting is happening. The output should have the same size as input at this dimension. stride.append(1) @@ -343,7 +363,7 @@ def expand( target_shape_t = target_shape # Handle dynamic shapes case where shape has dynamic dimension - if any(isinstance(ele, TRTTensor) for ele in target_shape_t): + if any(isinstance(ele, TRTTensor) for ele in (*target_shape_t, *stride)): target_shape_t = cat( ctx, target, diff --git a/py/torch_tensorrt/dynamo/lowering/_decomposition_filter.py b/py/torch_tensorrt/dynamo/lowering/_decomposition_filter.py index 191c905b6a9..415311d3f2c 100644 --- a/py/torch_tensorrt/dynamo/lowering/_decomposition_filter.py +++ b/py/torch_tensorrt/dynamo/lowering/_decomposition_filter.py @@ -16,13 +16,16 @@ def _has_symbolic_scatter_add_extent(graph_module: torch.fx.GraphModule) -> bool if node.target != torch.ops.aten.scatter_add.default or len(node.args) < 4: continue dim = node.args[1] - src_node = node.args[3] - if not isinstance(dim, int) or not isinstance(src_node, torch.fx.Node): + # The decomposition loops over the extent of index along dim. + index_node = node.args[2] + if not isinstance(dim, int) or not isinstance(index_node, torch.fx.Node): continue - src_val = src_node.meta.get("val", src_node.meta.get("example_value")) - if not isinstance(src_val, torch.Tensor) or not src_val.ndim: + index_val = index_node.meta.get("val", index_node.meta.get("example_value")) + if not isinstance(index_val, torch.Tensor) or not index_val.ndim: continue - if isinstance(src_val.shape[get_positive_dim(dim, src_val.ndim)], torch.SymInt): + if isinstance( + index_val.shape[get_positive_dim(dim, index_val.ndim)], torch.SymInt + ): return True return False diff --git a/py/torch_tensorrt/dynamo/lowering/_decompositions.py b/py/torch_tensorrt/dynamo/lowering/_decompositions.py index 56cf0ca052f..1bfcfa5d2c4 100644 --- a/py/torch_tensorrt/dynamo/lowering/_decompositions.py +++ b/py/torch_tensorrt/dynamo/lowering/_decompositions.py @@ -13,6 +13,7 @@ _get_decomp_for_cia, ) from torch._ops import OpOverload +from torch.fx.experimental.symbolic_shapes import statically_known_true from torch_tensorrt.dynamo._defaults import default_device from torch_tensorrt.dynamo.conversion.converter_utils import get_positive_dim from torch_tensorrt.dynamo.utils import to_torch_device @@ -28,6 +29,9 @@ from .constant_fold_exclusions.attention_mask import ( exclude_attn_mask_aranges_from_constant_fold, ) +from .constant_fold_exclusions.scatter_add import ( + exclude_scatter_add_positions_from_constant_fold, +) logger = logging.getLogger(__name__) @@ -292,10 +296,12 @@ def scatter_add_decomposition( index: torch.Tensor, src_tensor: torch.Tensor, ) -> torch.Tensor: + if _use_integer_scatter_add(input_tensor, dim, index, src_tensor): + return _scatter_add_integer_decomposition(input_tensor, dim, index, src_tensor) + scatter_add_tensor = input_tensor - src_shape = list(src_tensor.shape) - src_dim = src_shape[dim] - for i in range(0, src_dim): + # PyTorch reads src only where index has values, so loop over index. + for i in range(0, index.shape[dim]): to_scatter_tensor = torch.zeros(input_tensor.shape, dtype=input_tensor.dtype) # index and src slice src_slice = torch.select(src_tensor, dim, i) @@ -320,6 +326,75 @@ def scatter_add_decomposition( return scatter_add_tensor +# The comparison holds index.numel() * input.shape[dim] elements. Keep it well +# under TensorRT's tensor volume limit and small enough (512 MiB as int64) for +# constant folding or a PyTorch fallback to materialize. +_MAX_INTEGER_SCATTER_ADD_COMPARISONS = 1 << 26 + + +def _use_integer_scatter_add( + input_tensor: torch.Tensor, + dim: int, + index: torch.Tensor, + src_tensor: torch.Tensor, +) -> bool: + # Integer reductions can regroup updates without changing their values. + if src_tensor.dtype not in (torch.int32, torch.int64): + return False + # The comparison adds one axis, and TensorRT tensors have at most 8. + if index.dim() >= 8: + return False + if input_tensor.dim() == 0: + extent = 1 + elif index.dim() and isinstance(index.shape[dim], int): + extent = input_tensor.shape[dim] + else: + # Otherwise keep the loop, as before. On a dynamic extent the loop would + # specialize it, which is why filter_decomposition_table keeps such an op + # out of decomposition on the compile paths. + return False + comparisons = index.numel() * extent + return statically_known_true(comparisons <= _MAX_INTEGER_SCATTER_ADD_COMPARISONS) + + +def _scatter_add_integer_decomposition( + input_tensor: torch.Tensor, + dim: int, + index: torch.Tensor, + src_tensor: torch.Tensor, +) -> torch.Tensor: + # TensorRT scatter overwrites colliding indices, so compare each index with + # every output position along dim and sum the matching src values. + if statically_known_true(index.numel() == 0): + return input_tensor.clone() + if input_tensor.dim() == 0: + # PyTorch treats a scalar destination as one element along dim 0. + return _scatter_add_integer_decomposition( + input_tensor.reshape(1), 0, index.reshape(-1), src_tensor.reshape(-1) + ).reshape(()) + dim = get_positive_dim(dim, input_tensor.dim()) + for k, size in enumerate(index.shape): + if not statically_known_true(src_tensor.shape[k] == size): + # A slice view would add guards on dynamic sizes, a gather does not. + src_tensor = src_tensor.index_select( + k, torch.arange(size, device=src_tensor.device) + ) + index = index.movedim(dim, -1).unsqueeze(-2) + src_tensor = src_tensor.movedim(dim, -1).unsqueeze(-2) + positions = torch.arange(input_tensor.shape[dim], device=index.device) + exclude_scatter_add_positions_from_constant_fold(positions) + hits = index == positions.view(-1, 1) + # The where converter narrows int64 to int32, multiplying keeps the dtype. + added = (hits.to(src_tensor.dtype) * src_tensor).sum(-1, dtype=src_tensor.dtype) + added = added.movedim(-1, dim) + pad = [] + for k in reversed(range(input_tensor.dim())): + pad += [0, input_tensor.shape[k] - added.shape[k]] + if not all(statically_known_true(p == 0) for p in pad): + added = torch.constant_pad_nd(added, pad, 0) + return input_tensor + added + + # enum class for reduce operation of scatter_reduce class ReduceOperation(Enum): SUM = ("Sum reduce operation", lambda x, y: torch.add(x, y)) @@ -509,7 +584,7 @@ def scaled_dot_product_attention_decomposition( attn_bias = torch.where(temp_mask.logical_not(), negative, zero) if attn_mask is not None: - # Unconditional: mark_constant_fold_exclusions revokes these marks when the + # Unconditional: mark_constant_fold_exclusions ignores the tags when the # rule is disabled, so the decompositions do not need the setting. exclude_attn_mask_aranges_from_constant_fold(attn_mask) if attn_mask.dtype == torch.bool: diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py index 8e88946240c..855faaf12d9 100644 --- a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py +++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/__init__.py @@ -9,11 +9,13 @@ validate_disabled_constant_fold_exclusions, ) from .attention_mask import ATTENTION_MASK_ARANGE_RULE_ID +from .scatter_add import SCATTER_ADD_POSITIONS_RULE_ID __all__ = [ "ATTENTION_MASK_ARANGE_RULE_ID", "CONSTANT_FOLD_EXCLUSION_META_KEY", "ConstantFoldExclusionRule", + "SCATTER_ADD_POSITIONS_RULE_ID", "register_constant_fold_exclusion_rule", "unregister_constant_fold_exclusion_rule", "validate_disabled_constant_fold_exclusions", diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py index c9a3c4fba2a..9b3a96a69d7 100644 --- a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py +++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/_core.py @@ -15,8 +15,7 @@ def _mark_constant_fold_exclusion(nodes: Iterable[torch.fx.Node], rule_id: str) """Record which rule wants ``nodes`` kept out of constant folding. The marks carry their rule ID rather than a bare flag so the exclusion pass - can revoke the ones belonging to disabled rules, whichever marking path - produced them. + can revoke the ones belonging to disabled rules. """ for node in nodes: node.meta.setdefault(CONSTANT_FOLD_EXCLUSION_META_KEY, set()).add(rule_id) diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py index 7abe4c753ca..d60234d431b 100644 --- a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py +++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/attention_mask.py @@ -7,12 +7,10 @@ from torch._subclasses.functional_tensor import mb_unwrap_functional_tensor from torch.fx.experimental.proxy_tensor import get_proxy_mode, get_proxy_slot -from ._core import ( - _mark_constant_fold_exclusion, - register_constant_fold_exclusion_rule, -) +from ._core import register_constant_fold_exclusion_rule ATTENTION_MASK_ARANGE_RULE_ID = "attention_mask_arange" +_ATTENTION_MASK_ARANGE_TAG = "_torch_tensorrt_attention_mask_arange" def _is_arange_node(node: torch.fx.Node) -> bool: @@ -62,7 +60,7 @@ def _find_attention_mask_aranges( def exclude_attn_mask_aranges_from_constant_fold(attn_mask: torch.Tensor) -> None: - """Mark aranges behind an attention mask while tracing a decomposition. + """Tag aranges behind an attention mask while tracing a decomposition. ``run_decompositions`` invokes decompositions once for functionalization and again while proxy tracing. Only the proxy-tracing invocation has an FX graph @@ -82,15 +80,18 @@ def exclude_attn_mask_aranges_from_constant_fold(attn_mask: torch.Tensor) -> Non mask_node = getattr(proxy, "node", None) if isinstance(mask_node, torch.fx.Node): - _mark_constant_fold_exclusion( - _find_attention_mask_aranges(mask_node), - ATTENTION_MASK_ARANGE_RULE_ID, - ) + # Retraces keep node.meta["custom"], so the rule below still finds + # these aranges after the attention op has been decomposed away. + for node in _find_attention_mask_aranges(mask_node): + node.meta.setdefault("custom", {})[_ATTENTION_MASK_ARANGE_TAG] = True @register_constant_fold_exclusion_rule(ATTENTION_MASK_ARANGE_RULE_ID) def _attention_mask_arange_rule(node: torch.fx.Node) -> Iterable[torch.fx.Node]: """Select aranges feeding an attention mask.""" + if node.meta.get("custom", {}).get(_ATTENTION_MASK_ARANGE_TAG): + return (node,) + # Every SDPA overload that takes a mask keeps it at positional index 3. # _scaled_dot_product_flash_attention is absent because it has no mask. attention_mask_args: dict[Any, tuple[int, str]] = { diff --git a/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py new file mode 100644 index 00000000000..ab6e902fef8 --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/constant_fold_exclusions/scatter_add.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +from typing import Iterable + +import torch +from torch._subclasses.functional_tensor import mb_unwrap_functional_tensor +from torch.fx.experimental.proxy_tensor import get_proxy_mode, get_proxy_slot + +from ._core import ( + CONSTANT_FOLD_EXCLUSION_META_KEY, + register_constant_fold_exclusion_rule, +) + +SCATTER_ADD_POSITIONS_RULE_ID = "scatter_add_positions" +_POSITIONS_TAG = "_torch_tensorrt_scatter_add_positions" + + +def exclude_scatter_add_positions_from_constant_fold(positions: torch.Tensor) -> None: + """Tag the positions arange of the scatter_add decomposition. + + With a constant index, folding it would store the comparison of every index + with every position, the 0/1 mask that src is multiplied by, as an engine + constant, even when src varies at run time. Retraces keep + ``node.meta["custom"]``, so the rule below finds the tag after the + decomposed op is gone, for example after the complex decomposition retrace + or a second ``run_decompositions``. + """ + proxy_mode = get_proxy_mode() + if proxy_mode is None: + return + + tracked = get_proxy_slot( + mb_unwrap_functional_tensor(positions), proxy_mode.tracer, default=None + ) + node = getattr(getattr(tracked, "proxy", tracked), "node", None) + if isinstance(node, torch.fx.Node): + node.meta.setdefault("custom", {})[_POSITIONS_TAG] = True + + +def _varies_at_runtime(node: torch.fx.Node) -> bool: + """Whether this node reads a value that constant folding leaves in the graph. + + That is an input, a random op, a node an earlier rule keeps out of folding, or + a view whose storage the graph returns or mutates. + """ + from torch_tensorrt.dynamo.lowering.passes.reset_folded_constructors import ( + _aliases_input, + _storage_escapes, + ) + + nodes_to_visit, visited = [node], set() + while nodes_to_visit: + current = nodes_to_visit.pop() + target_tags = getattr(current.target, "tags", ()) + if ( + current.op == "placeholder" + or torch.Tag.nondeterministic_seeded in target_tags + or current.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY) + or (_aliases_input(current) and _storage_escapes(current)) + ): + return True + if current not in visited: + visited.add(current) + nodes_to_visit.extend(current.all_input_nodes) + return False + + +@register_constant_fold_exclusion_rule(SCATTER_ADD_POSITIONS_RULE_ID) +def _scatter_add_positions_rule(node: torch.fx.Node) -> Iterable[torch.fx.Node]: + """Select the tagged positions arange when src varies at runtime. + + The decomposition multiplies ``(index == positions.view(-1, 1)).to(dtype)`` + by src. With constant src, a constant index makes the whole sum constant, + and it folds to a result the size of the output, so the positions stay + foldable. A runtime index keeps the comparison live either way. + """ + if not node.meta.get("custom", {}).get(_POSITIONS_TAG): + return () + # Several scatter_adds can share the positions, for example after CSE. + frontier, comparison, sources = [node], {node}, [] + for _ in range(4): # view, eq, cast, mul + frontier = [user for current in frontier for user in current.users] + for user in frontier: + if user.target == torch.ops.aten.mul.Tensor: + sources += [ + arg for arg in user.all_input_nodes if arg not in comparison + ] + comparison.update(frontier) + if sources and not any(map(_varies_at_runtime, sources)): + return () + return (node,) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py index e58c0fef06d..132b2c191fd 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py @@ -3,9 +3,11 @@ import logging import operator -from typing import Any, Callable, Optional, Sequence, Union +from typing import Any, Callable, List, Optional, Sequence, Union import torch +import torch.utils._pytree as pytree +from torch.fx.experimental.symbolic_shapes import statically_known_true from torch_tensorrt._utils import is_tegra_platform from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo.lowering.passes._FakeTensorUpdater import FakeTensorUpdater @@ -36,6 +38,7 @@ from .lower_associative_scan import lower_associative_scan from .normalize_negative_slice_stop import normalize_negative_slice_stop # noqa: F401 from .pass_manager import DynamoPassManager +from .preserve_attention_output_layout import preserve_attention_output_layout from .remove_alias import remove_alias from .remove_assert_nodes import remove_assert_nodes from .remove_detach import remove_detach @@ -51,6 +54,7 @@ remove_detach, remove_assert_nodes, rule_based_autocast, + preserve_attention_output_layout, ] @@ -193,12 +197,61 @@ def post_lowering( if fake_mode is not None: fake_tensor_updater.incremental_update(fake_mode) - for sub in iter_cond_subgraphs(gm): - post_lowering(sub, settings) + # A branch shared by several cond nodes is lowered once. + for branch in dict.fromkeys(iter_cond_subgraphs(gm)): + lowered = post_lowering(branch, settings) + # A pass can return a new module instead of editing the branch in place. + # The cond node was traced against the old one, so take the new module only + # if its inputs and outputs keep their dtypes and memory layout. + if lowered is branch or not _same_io(lowered, branch): + continue + for name, module in list(gm.named_modules(remove_duplicate=False)): + if module is branch: + gm.set_submodule(name, lowered) return gm +def _io_values(gm: torch.fx.GraphModule) -> List[Any]: + """The meta value of each graph input and output, in order.""" + + def value(node: Any) -> Any: + return node.meta.get("val") if isinstance(node, torch.fx.Node) else None + + placeholders = [value(n) for n in gm.graph.nodes if n.op == "placeholder"] + outputs = [value(n) for n in pytree.tree_leaves(gm.graph.output_node().args)] + return placeholders + outputs + + +def _same_io(a: torch.fx.GraphModule, b: torch.fx.GraphModule) -> bool: + """Whether two modules take and return tensors of the same kind and layout. + + Tensors match when their dtype, device, sizes, strides, storage offset and + storage size match. + """ + + def same(x: Any, y: Any) -> bool: + if not isinstance(x, torch.Tensor) or not isinstance(y, torch.Tensor): + return not isinstance(x, torch.Tensor) and not isinstance(y, torch.Tensor) + # A symbolic value that cannot be proven equal counts as different. + return ( + x.dtype == y.dtype + and x.device == y.device + and x.dim() == y.dim() + and all(statically_known_true(p == q) for p, q in zip(x.shape, y.shape)) + and statically_known_true( + x.untyped_storage().nbytes() == y.untyped_storage().nbytes() + ) + and statically_known_true(x.storage_offset() == y.storage_offset()) + and all( + statically_known_true(p == q) for p, q in zip(x.stride(), y.stride()) + ) + ) + + a_values, b_values = _io_values(a), _io_values(b) + return len(a_values) == len(b_values) and all(map(same, a_values, b_values)) + + def pre_export_lowering( ep: torch.export.ExportedProgram, settings: CompilationSettings = CompilationSettings(), diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py index e35265461e0..d598dcce680 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py @@ -27,8 +27,10 @@ such as ``aten.log2``/``aten.log10`` as of torch 2.14.dev). """ +import gc import logging -from typing import Any +from contextlib import contextmanager +from typing import Any, Iterator, Optional import torch from torch._export.utils import _detect_fake_mode_from_gm @@ -108,6 +110,7 @@ def complex_decomposition_adapter( # codegen then unflattens again. gm.graph._codegen = torch.fx.graph.CodeGen() gm.recompile() + decomposed_gm: Optional[GraphModule] = None try: # Both retraces return a new graph. Keep the original for the existing # unsupported-op fallback rather than replacing it before success. @@ -116,11 +119,12 @@ def complex_decomposition_adapter( if any(free_symbols(arg) for arg in flat_args) else decompose_complex_in_graph ) - decomposed_gm = decompose( - gm, - flat_args, - decompositions=_trt_decomposition_table(settings, gm), - ) + with _offloaded_constants_on_device(gm) as offloaded: + decomposed_gm = decompose( + gm, + flat_args, + decompositions=_trt_decomposition_table(settings, gm), + ) except Exception as e: # decompose_complex_in_graph is upstream, experimental PyTorch code # (torch._functorch._aot_autograd.complex_decomposition) with @@ -133,11 +137,17 @@ def complex_decomposition_adapter( f"complex decomposition failed ({type(e).__name__}: {e}); " "falling back to the legacy complex_graph_detection pass." ) - gm.graph._codegen = original_codegen - gm.recompile() - return complex_graph_detection(gm, settings) gm.graph._codegen = original_codegen # the caller's module, as it was gm.recompile() + if decomposed_gm is None: + # Not inside the except block: there the failed retrace's frames are still + # alive, and with them its device copies of the offloaded constants. + return complex_graph_detection(gm, settings) + _restore_offloaded_constants(decomposed_gm, offloaded) + if offloaded: + # The retrace leaves the device copies in reference cycles. Free them now + # rather than holding them through the engine build. + gc.collect() gm = decomposed_gm # The retrace brings back alias nodes. aten.alias has no converter, so the @@ -239,6 +249,54 @@ def _graph_has_complex(gm: GraphModule) -> bool: return False +@contextmanager +def _offloaded_constants_on_device( + gm: GraphModule, +) -> Iterator[dict[int, torch.Tensor]]: + """Hand the retrace each CPU-offloaded constant on the device it was computed on. + + With offload_module_to_cpu, constant folding keeps the folded values on the CPU + while their metadata still records the device that computed them. The retrace + runs the graph, so a CPU constant next to device activations fails with a device + mismatch. Copy each such constant to its metadata device for the retrace only; + yields a map from the id of each copy to its CPU tensor, for + _restore_offloaded_constants. + """ + swapped = {} + for node in gm.graph.nodes: + target = str(node.target) + if node.op != "get_attr" or not hasattr(gm, target): + continue + real, val = getattr(gm, target), node.meta.get("val", None) + if ( + isinstance(real, torch.Tensor) + and isinstance(val, torch.Tensor) + and real.device != val.device + ): + swapped[target] = (real, real.to(val.device)) + for target, (_, copy) in swapped.items(): + delattr(gm, target) + setattr(gm, target, copy) + try: + yield {id(copy): real for real, copy in swapped.values()} + finally: + for target, (real, _) in swapped.items(): + delattr(gm, target) + setattr(gm, target, real) + + +def _restore_offloaded_constants( + gm: GraphModule, offloaded: dict[int, torch.Tensor] +) -> None: + """Put the CPU tensors back wherever the retraced module kept a device copy.""" + for node in gm.graph.nodes: + if node.op != "get_attr": + continue + real = offloaded.get(id(getattr(gm, str(node.target), None))) + if real is not None: + setattr(gm, str(node.target), real) + + def _fake_flat_args(gm: GraphModule) -> list[torch.Tensor]: """Build the example inputs the upstream retrace (make_fx) needs. diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py index be9134e33d4..dd365cfcb4e 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py @@ -33,8 +33,6 @@ torch.ops.aten.add.Tensor, torch.ops.aten.sub.Tensor, torch.ops.aten.neg.default, - torch.ops.aten.mul.Scalar, # scalar*(re,im) — both parts scaled equally - torch.ops.aten.div.Scalar, # (re,im)/scalar — both parts divided equally # Structural / copy — operate on the whole tensor without touching content. # Note: permute.default is NOT here; it needs an explicit rewrite to append # the trailing real/imag dimension index to the dims list. @@ -463,7 +461,12 @@ def _rewrite_permute(self, node: Node) -> bool: self.gm.graph.erase_node(node) return True - @_complex_unpacker(torch.ops.aten.mul.Tensor, torch.ops.aten.div.Tensor) + @_complex_unpacker( + torch.ops.aten.mul.Tensor, + torch.ops.aten.div.Tensor, + torch.ops.aten.mul.Scalar, + torch.ops.aten.div.Scalar, + ) def _rewrite_mul_div_tensor(self, node: Node) -> bool: arg0_is_node = isinstance(node.args[0], torch.fx.Node) arg1_is_node = isinstance(node.args[1], torch.fx.Node) @@ -471,13 +474,29 @@ def _rewrite_mul_div_tensor(self, node: Node) -> bool: if not arg0_is_node and not arg1_is_node: return False # both scalars - if node.target == torch.ops.aten.mul.Tensor and ( - not arg0_is_node or not arg1_is_node + if node.target in (torch.ops.aten.mul.Scalar, torch.ops.aten.div.Scalar) and ( + arg1_is_node ): - return False # scalar * complex — elementwise-safe - - if node.target == torch.ops.aten.div.Tensor and not arg1_is_node: - return False # complex / scalar — elementwise-safe + return False # symbolic real scalar: elementwise-safe + + is_mul = node.target in (torch.ops.aten.mul.Tensor, torch.ops.aten.mul.Scalar) + if is_mul and (not arg0_is_node or not arg1_is_node): + z_node, scalar = ( + (node.args[0], node.args[1]) + if arg0_is_node + else (node.args[1], node.args[0]) + ) + if isinstance(scalar, complex): + return self._rewrite_mul_complex_scalar(node, z_node, scalar) + return False # real scalar * complex: elementwise-safe + + if not is_mul and not arg1_is_node: + c = node.args[1] + if isinstance(c, complex): + # Python raises on 1 / 0j; nan is what CUDA eager returns for it. + inverse = 1 / c if c else complex("nan+nanj") + return self._rewrite_mul_complex_scalar(node, node.args[0], inverse) + return False # complex / real scalar: elementwise-safe if node.target == torch.ops.aten.div.Tensor and not arg0_is_node: # scalar / complex: s/(a+bi) = (s*a/(a²+b²)) + i*(-s*b/(a²+b²)) @@ -625,21 +644,63 @@ def match_complex_mul( return True return False - @_complex_unpacker(torch.ops.aten.add.Tensor, torch.ops.aten.sub.Tensor) - def _rewrite_add_sub_tensor_scalar(self, node: Node) -> bool: + def _rewrite_mul_complex_scalar(self, node: Node, z_node: Node, c: complex) -> bool: + with SubgraphBuilder(self.gm.graph, node) as b: + if self._is_complex_layout_node(z_node): + # (a+bi)(x+yi) = (ax-by) + (ay+bx)i + re, im = self._inline_select_re_im(b, z_node) + re_x = b(torch.ops.aten.mul.Tensor, re, c.real) + im_y = b(torch.ops.aten.mul.Tensor, im, c.imag) + out_re = b(torch.ops.aten.sub.Tensor, re_x, im_y) + re_y = b(torch.ops.aten.mul.Tensor, re, c.imag) + im_x = b(torch.ops.aten.mul.Tensor, im, c.real) + out_im = b(torch.ops.aten.add.Tensor, re_y, im_x) + else: + # a(x+yi) = ax + ayi for a real tensor a + out_re = b(torch.ops.aten.mul.Tensor, z_node, c.real) + out_im = b(torch.ops.aten.mul.Tensor, z_node, c.imag) + out = self._inline_cat_re_im(b, out_re, out_im) + node.replace_all_uses_with(out) + self.gm.graph.erase_node(node) + return True + + @_complex_unpacker( + torch.ops.aten.add.Tensor, + torch.ops.aten.sub.Tensor, + torch.ops.aten.add.Scalar, + torch.ops.aten.sub.Scalar, + ) + def _rewrite_add_sub_scalar(self, node: Node) -> bool: # add.Tensor(z1, z2): tensor+tensor is componentwise on the [..., 2] # real layout — real and imag parts are handled independently — so no # rewrite is needed. Return True (handled, no modification). - # add.Tensor(z, scalar): scalar applies to real part only, so we must - # explicitly rewrite. - # tensor+tensor: componentwise on [..., 2] layout — no rewrite needed + # (a+bi) ± (x+yi) = (a±x) + (b±y)i, so a real scalar touches only the + # real part and a complex scalar adds its parts separately. if len(node.args) < 2 or isinstance(node.args[1], torch.fx.Node): return True inp, scalar = node.args[0], node.args[1] + scalar = scalar * node.kwargs.get("alpha", 1) + s_re, s_im = ( + (scalar.real, scalar.imag) if isinstance(scalar, complex) else (scalar, 0) + ) with SubgraphBuilder(self.gm.graph, node) as b: - re = b(torch.ops.aten.select.int, inp, -1, 0) - im = b(torch.ops.aten.select.int, inp, -1, 1) - new_re = b(node.target, re, scalar) + if self._is_complex_layout_node(inp): + re, im = self._inline_select_re_im(b, inp) + if s_im: + im = b(node.target, im, s_im) + else: + re, im = inp, None + new_re = b(node.target, re, s_re) + if im is None: + # A real tensor has no imaginary part, so the result's is ±y. Fill + # it like new_re, which already has the promoted float dtype. + is_add = node.target in ( + torch.ops.aten.add.Tensor, + torch.ops.aten.add.Scalar, + ) + im = b( + torch.ops.aten.full_like.default, new_re, s_im if is_add else -s_im + ) out = self._inline_cat_re_im(b, new_re, im) node.replace_all_uses_with(out) out.meta["is_complex_layout"] = True @@ -711,23 +772,6 @@ def _rewrite_mean_dim(self, node: Node) -> bool: self.gm.graph.erase_node(node) return True - @_complex_unpacker(torch.ops.aten.prod.dim_int) - def _rewrite_prod_dim(self, node: Node) -> bool: - # prod.dim_int(inp, dim, keepdim=False, dtype=None) - inp = node.args[0] - dim = node.args[1] - if dim >= 0: - return False # positive dim — pass-through is correct - new_dim = dim - 1 - keepdim = node.args[2] if len(node.args) > 2 else False - extra = list(node.args[3:]) - with SubgraphBuilder(self.gm.graph, node) as b: - out = b(torch.ops.aten.prod.dim_int, inp, new_dim, keepdim, *extra) - node.replace_all_uses_with(out) - out.meta["is_complex_layout"] = True - self.gm.graph.erase_node(node) - return True - @_complex_unpacker(torch.ops.aten.narrow.default) def _rewrite_narrow(self, node: Node) -> bool: # narrow(inp, dim, start, length) — shift negative dim by -1 @@ -915,6 +959,15 @@ def _rewrite_angle(self, node: Node) -> bool: self.gm.graph.erase_node(node) return True + @_complex_unpacker(torch.ops.aten.complex.default) + def _rewrite_complex(self, node: Node) -> bool: + # complex(re, im) is the [..., 2] layout of its two real inputs. + with SubgraphBuilder(self.gm.graph, node) as b: + out = self._inline_cat_re_im(b, node.args[0], node.args[1]) + node.replace_all_uses_with(out) + self.gm.graph.erase_node(node) + return True + @_complex_unpacker(torch.ops.aten.polar.default) def _rewrite_polar(self, node: Node) -> bool: # polar(r, theta) = r*cos(theta) + i*r*sin(theta) @@ -929,19 +982,6 @@ def _rewrite_polar(self, node: Node) -> bool: self.gm.graph.erase_node(node) return True - @_complex_unpacker(torch.ops.aten.add.Scalar, torch.ops.aten.sub.Scalar) - def _rewrite_add_sub_scalar(self, node: Node) -> bool: - # (a+bi) ± s = (a±s) + bi — scalar applies to real part only - inp, scalar = node.args[0], node.args[1] - with SubgraphBuilder(self.gm.graph, node) as b: - re = b(torch.ops.aten.select.int, inp, -1, 0) - im = b(torch.ops.aten.select.int, inp, -1, 1) - new_re = b(node.target, re, scalar) - out = self._inline_cat_re_im(b, new_re, im) - node.replace_all_uses_with(out) - self.gm.graph.erase_node(node) - return True - @_complex_unpacker(torch.ops.aten.log2.default, torch.ops.aten.log10.default) def _rewrite_log2_log10(self, node: Node) -> bool: # log_b(z) = log(z) / log(b) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py b/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py index 5c1150afdf3..7feb16fd4e5 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/mark_constant_fold_exclusions.py @@ -18,20 +18,22 @@ def mark_constant_fold_exclusions( """Apply the registered rules that exclude FX nodes from constant folding. This pass is the single authority on which rules are in effect. It runs - immediately before ``constant_fold`` and is the only marking path that sees - ``settings``: rules that mark nodes while a decomposition is traced run - during ``run_decompositions``, long before a settings object is reachable. - Those marks are therefore revoked here rather than suppressed where they are - made, so a caller only has to communicate the disabled rules once. + immediately before ``constant_fold`` and is the only code that marks nodes: + decompositions run during ``run_decompositions``, long before a settings + object is reachable, so they only tag nodes for a rule to select here. As a + guard, it also revokes marks of disabled rules that a node already carries, + so a caller only has to communicate the disabled rules once. """ disabled_rule_ids = validate_disabled_constant_fold_exclusions( settings.disabled_constant_fold_exclusions if settings is not None else () ) - for node in gm.graph.nodes: - for rule_id, rule in _CONSTANT_FOLD_EXCLUSION_RULES.items(): - if rule_id in disabled_rule_ids: - continue + # One rule covers the whole graph before the next starts, so a rule can read + # the marks of the rules registered before it. + for rule_id, rule in _CONSTANT_FOLD_EXCLUSION_RULES.items(): + if rule_id in disabled_rule_ids: + continue + for node in gm.graph.nodes: _mark_constant_fold_exclusion(rule(node), rule_id) if disabled_rule_ids: diff --git a/py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py b/py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py new file mode 100644 index 00000000000..1fa9cd702a2 --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/passes/preserve_attention_output_layout.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +import logging + +import torch +from torch._prims_common import is_contiguous_or_false +from torch._subclasses.fake_tensor import FakeTensor +from torch_tensorrt.dynamo._settings import CompilationSettings + +logger = logging.getLogger(__name__) + + +def _preserve_layouts(graph: torch.fx.Graph) -> int: + count = 0 + for node in list(graph.nodes): + if node.target != torch.ops.aten.scaled_dot_product_attention.default: + continue + val = node.meta.get("val") + # Under torch.vmap the recorded value is a batched tensor, not a FakeTensor, + # and those nodes lower without this pass. + if not isinstance(val, FakeTensor) or is_contiguous_or_false(val): + continue + dim_order = val.dim_order() + if dim_order == tuple(range(val.dim())): + continue + + inverse = [dim_order.index(dim) for dim in range(val.dim())] + with graph.inserting_before(node): + attention = graph.call_function(node.target, node.args, node.kwargs) + permuted = graph.call_function( + torch.ops.aten.permute.default, (attention, list(dim_order)) + ) + copied = graph.call_function( + torch.ops.aten.clone.default, + (permuted,), + {"memory_format": torch.contiguous_format}, + ) + restored = graph.call_function( + torch.ops.aten.permute.default, (copied, inverse) + ) + for new_node in (attention, permuted, copied, restored): + new_node.meta.update( + {key: value for key, value in node.meta.items() if key != "tensor_meta"} + ) + # Recorded as the decomposition returns it, so a second run leaves it alone. + attention.meta["val"] = val.contiguous() + permuted.meta["val"] = attention.meta["val"].permute(dim_order) + copied.meta["val"] = permuted.meta["val"].contiguous() + restored.meta["val"] = copied.meta["val"].permute(inverse) + # The original node becomes a copy that keeps the recorded layout, so every + # user, graph outputs and in-place updates included, keeps reading the same + # node. A copy rather than a view, so in-place ops such as detach_() work. + node.target = torch.ops.aten.clone.default + node.args = (restored,) + node.kwargs = {} + count += 1 + return count + + +def preserve_attention_output_layout( + gm: torch.fx.GraphModule, settings: CompilationSettings +) -> torch.fx.GraphModule: + """Keep the exported memory layout of attention outputs that will be decomposed. + + The fused attention kernels return their output in a layout that depends on the + backend and the query, and export drops a ``.contiguous()`` that is a no-op on it. + The attention decomposition returns a contiguous tensor instead, so a later + ``view`` fails when the graph is retraced. Copying each output into the layout + recorded at export keeps those views valid. + """ + if not settings.decompose_attention: + return gm + + count = 0 + # Submodules too: attention in torch.cond branches and in torch.no_grad() or + # autocast regions lives in subgraphs, nested ones included. + for module in gm.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + changed = _preserve_layouts(module.graph) + if changed: + # No dead code elimination: the graph is not functionalized yet, and it + # would drop nodes with no users such as the vmap nesting calls. + module.recompile() + count += changed + + if count: + logger.debug("Preserved the output layout of %d attention nodes", count) + return gm diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py index 0b46c420286..4d9ae72b0dc 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: BSD-3-Clause import logging +from collections import defaultdict, deque from typing import Collection, Dict, List, Optional, Set, Tuple import torch @@ -10,6 +11,7 @@ from torch.fx.passes.splitter_base import ( FxNetAccFusionsFinder, FxNetAccNodesFinder, + FxNetSplitterInternalError, Subgraph, _SplitterBase, _SplitterSettingBase, @@ -412,6 +414,67 @@ def starter_nodes(self) -> Tuple[NodeSet, NodeSet]: return starter_non_acc_nodes, starter_acc_nodes + def put_nodes_into_subgraphs(self) -> List[Subgraph]: + """Builds the same subgraphs as _SplitterBase.put_nodes_into_subgraphs + + starter_nodes puts every callable node in the frontier, and the base + class can rescan that whole frontier for a ready node at every step, + which is quadratic on large graphs. Counting the unplaced dependencies + of each node finds ready nodes in time proportional to the nodes plus + their dependency entries. Each subgraph still takes every node of its + kind whose dependencies are placed, so the node sets match. Fusion + groups need no extra step: every callable node is counted from the + start, and their constraints are already in deps. + """ + pending = { + node: len(self.deps[node]) + for node in self.module.graph.nodes + if node.op in CALLABLE_NODE_OPS + } + dependents: Dict[torch.fx.Node, List[torch.fx.Node]] = defaultdict(list) + for node in pending: + for dep in self.deps[node]: + dependents[dep].append(node) + + ready = { + is_acc: deque( + node + for node, count in pending.items() + if count == 0 and (node in self.acc_nodes) == is_acc + ) + for is_acc in (False, True) + } + acc_subgraph = not ready[False] + current_subgraph_nodes: List[torch.fx.Node] = [] + subgraphs: List[Subgraph] = [] + remaining = len(pending) + while remaining: + if not ready[acc_subgraph]: + if not current_subgraph_nodes: + raise FxNetSplitterInternalError("Subgraph can't be empty") + subgraphs.append( + Subgraph(is_acc=acc_subgraph, nodes=current_subgraph_nodes) + ) + acc_subgraph = not acc_subgraph + current_subgraph_nodes = [] + continue + + node = ready[acc_subgraph].popleft() + current_subgraph_nodes.append(node) + remaining -= 1 + for user in dependents[node]: + pending[user] -= 1 + if pending[user] == 0: + ready[user in self.acc_nodes].append(user) + + if current_subgraph_nodes: + subgraphs.append( + Subgraph(is_acc=acc_subgraph, nodes=current_subgraph_nodes) + ) + if not subgraphs: + raise FxNetSplitterInternalError("Couldn't create subgraphs") + return subgraphs + def partition( gm: torch.fx.GraphModule, diff --git a/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py b/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py index e66b75a3a7c..64354b0cbef 100644 --- a/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py +++ b/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py @@ -266,6 +266,7 @@ def __init__( self._input_buffers: List[torch.Tensor] = [] self._output_buffers: List[torch.Tensor] = [] self._bound_inputs_by_name: Dict[str, torch.Tensor] = {} + self._shape_input_buffers: List[Any] = [] self._caller_stream: Optional[torch.cuda.Stream] = None self._engine_stream: Optional[torch.cuda.Stream] = None self._owned_pool_stream: Optional[torch.cuda.Stream] = None @@ -407,6 +408,7 @@ def __setstate__(self, state: Any) -> None: self._input_buffers = [] self._output_buffers = [] self._bound_inputs_by_name = {} + self._shape_input_buffers = [] self._caller_stream = None self._engine_stream = None self._owned_pool_stream = None @@ -1192,6 +1194,7 @@ def setup_input_tensors( # create_output_tensors / the output-binding loop to alias outputs to # their source-input storage. self._bound_inputs_by_name = {} + self._shape_input_buffers = [] for i, binding in enumerate(self._input_binding_infos): input_name = binding.name @@ -1208,6 +1211,9 @@ def setup_input_tensors( if binding.is_shape_tensor: inputs_cpu = contiguous_inputs[i].cpu().to(torch.int64).numpy().copy() + # TensorRT reads this host memory when the engine runs, so the + # array must outlive this loop. + self._shape_input_buffers.append(inputs_cpu) self.context.set_tensor_address(input_name, inputs_cpu.ctypes.data) else: self.context.set_input_shape( diff --git a/py/torch_tensorrt/executorch/_export_utils.py b/py/torch_tensorrt/executorch/_export_utils.py index f4730fae149..7467f5c6548 100644 --- a/py/torch_tensorrt/executorch/_export_utils.py +++ b/py/torch_tensorrt/executorch/_export_utils.py @@ -12,6 +12,10 @@ from torch._library.fake_class_registry import FakeScriptObject from torch._subclasses.fake_tensor import is_fake from torch.export.graph_signature import InputKind +from torch.fx.experimental.symbolic_shapes import ( + guarding_hint_or_throw, + has_guarding_hint, +) logger = logging.getLogger(__name__) @@ -92,8 +96,39 @@ def get_engine_info_from_state( return list(state[0] if isinstance(state, tuple) else state) -def validate_engine_info(engine_info: Sequence[Any], *, node_name: str = "") -> None: - """Reject engine configurations unsupported by ExecuTorch export.""" +def _output_shapes_are_known(engine_obj: Any, node: Any) -> bool: + """Whether the engine knows every output shape for this node's input shapes. + + Torch-TensorRT marks an engine as needing an output allocator when any of its + operations has a data-dependent shape, which also keeps CUDA graphs off for it. + Only an output with such a shape needs the allocator ExecuTorch lacks. An input + without a concrete example shape makes the answer False, since nothing proves + the outputs known. + """ + input_shapes = [] + for arg in node.args[0]: + val = arg.meta.get("val") if isinstance(arg, torch.fx.Node) else None + if not isinstance(val, torch.Tensor): + return False + if not all(has_guarding_hint(dim) for dim in val.shape): + return False + input_shapes.append([guarding_hint_or_throw(dim) for dim in val.shape]) + output_shapes = engine_obj.infer_outputs(input_shapes) + return all( + -1 not in tuple(getattr(shape, "shape", shape)) for shape in output_shapes + ) + + +def validate_engine_info( + engine_info: Sequence[Any], + *, + node_name: str = "", + cause: BaseException | None = None, +) -> None: + """Reject engine configurations unsupported by ExecuTorch export. + + ``cause`` is chained to the refusal, so the error behind it stays visible. + """ from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( REQUIRES_OUTPUT_ALLOCATOR_IDX, ) @@ -107,7 +142,7 @@ def validate_engine_info(engine_info: Sequence[Any], *, node_name: str = "") -> "ExecuTorch export does not support TensorRT engines that require " "an output allocator (data-dependent output shapes)" f"{node_suffix}." - ) + ) from cause def _schema_name(node: Any) -> str: @@ -231,12 +266,19 @@ def validate_engine_program( Resolving an engine serializes it, so pass a dict as ``resolved`` to keep what was resolved here and let the rewrite reuse it instead of serializing a second time. - Validation reads only metadata, so the records kept in ``resolved`` are - metadata-only: their ENGINE_IDX slot is empty. The rewrite reuses them for the + Validation reads each engine's metadata. For an engine marked as needing an + output allocator it also asks the live engine for its output shapes, which sets + the node's input shapes on the engine's execution context and creates the context + if no run has yet. The records kept in ``resolved`` are metadata-only: their + ENGINE_IDX slot is empty. The rewrite reuses them for the plain-string slots and takes the engine itself from :func:`_resolve_engine_tensor`, which never reads a cached record's ENGINE_IDX -- so a metadata-only record is never cross-served as the engine bytes. """ + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + REQUIRES_OUTPUT_ALLOCATOR_IDX, + ) + count = 0 # One engine can feed several execute_engine nodes, and resolving it again would # serialize it again and keep a second copy of the bytes. @@ -263,7 +305,28 @@ def validate_engine_program( ) if engine_node is not None: by_engine_node[engine_node] = engine_info - validate_engine_info(engine_info, node_name=node.name) + inference_error = None + if ( + node.target is torch.ops.tensorrt.execute_engine.default + and len(engine_info) > REQUIRES_OUTPUT_ALLOCATOR_IDX + and str(engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX]) == "1" + ): + try: + shapes_known = _output_shapes_are_known( + _resolve_engine_object(exported_program, node), node + ) + except RuntimeError as error: + # The C++ engine raises when it cannot create an output for an + # unknown dimension, but also when anything else fails, so the + # refusal keeps the error as its cause. + shapes_known, inference_error = False, error + if shapes_known: + # ExecuTorch needs no output allocator then. This record becomes + # the ExecuTorch engine's arguments, and the backend refuses a + # marked one. + engine_info = list(engine_info) + engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = "0" + validate_engine_info(engine_info, node_name=node.name, cause=inference_error) if resolved is not None: resolved[node.name] = engine_info count += 1 diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index 6f08b8bf047..ac8bfc67a19 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -10,6 +10,7 @@ import operator from typing import Any, Container, Iterable, List, Optional, Set, final +import tensorrt import torch import torch.fx from executorch.exir._serialize._named_data_store import NamedDataStore @@ -662,6 +663,7 @@ def preprocess( device_id=_parse_device_id(engine_info[DEVICE_IDX]), serialized_metadata=_get_str(engine_info, SERIALIZED_METADATA_IDX), target_platform=_get_str(engine_info, TARGET_PLATFORM_IDX), + tensorrt_version=tensorrt.__version__, # Content-keyed, so identical engines in different methods share one # entry and different engines can never collide on a key. engine_key="tensorrt_engine_" + hashlib.sha256(engine_bytes).hexdigest(), diff --git a/py/torch_tensorrt/executorch/serialization.py b/py/torch_tensorrt/executorch/serialization.py index dc7b0a56814..d0214fd0cd3 100644 --- a/py/torch_tensorrt/executorch/serialization.py +++ b/py/torch_tensorrt/executorch/serialization.py @@ -64,6 +64,9 @@ class TensorRTBlobMetadata: # Named-data key of the engine when it is stored outside the blob, which then # carries an empty engine. Empty means the engine follows the metadata inline. engine_key: str = "" + # TensorRT version that exported the engine, so a runtime whose TensorRT cannot + # read the engine can name both versions. Empty when unknown. + tensorrt_version: str = "" def to_json(self) -> bytes: # Keep field order stable because the C++ parser is intentionally small. @@ -92,6 +95,7 @@ def to_json(self) -> bytes: "device_id": self.device_id, "serialized_metadata": self.serialized_metadata, "target_platform": self.target_platform, + "tensorrt_version": self.tensorrt_version, } if self.engine_key: data["engine_key"] = self.engine_key @@ -119,6 +123,7 @@ def from_json(cls, data: bytes) -> "TensorRTBlobMetadata": serialized_metadata=parsed.get("serialized_metadata", ""), target_platform=parsed.get("target_platform", ""), engine_key=parsed.get("engine_key", ""), + tensorrt_version=parsed.get("tensorrt_version", ""), ) diff --git a/tests/py/dynamo/conversion/test_casts.py b/tests/py/dynamo/conversion/test_casts.py index 90d3e6d678a..d27d2aaa753 100644 --- a/tests/py/dynamo/conversion/test_casts.py +++ b/tests/py/dynamo/conversion/test_casts.py @@ -5,7 +5,9 @@ import torch import torch.nn as nn +import torch_tensorrt from torch.testing._internal.common_utils import run_tests +from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo.conversion.aten_ops_converters import ( to_copy_dtype_validator, ) @@ -94,6 +96,115 @@ def forward(self, x): inputs, ) + def test_to_copy_float64_with_truncate_double(self): + class ToCopyDouble(nn.Module): + def forward(self, x): + y = x * 2 + return y, torch.ops.aten._to_copy.default(y, dtype=torch.float64) + + # aten.copy reaches the same cast without the _to_copy validator, and + # broadcasts its source to the buffer shape. + class CopyIntoDouble(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("dest", torch.zeros(2, 3, 10, dtype=torch.double)) + + def forward(self, x): + return torch.ops.aten.copy.default(self.dest, x * 2) + + inputs = [torch.rand((1, 3, 10), dtype=torch.float)] + for model in (ToCopyDouble(), CopyIntoDouble()): + with self.subTest(type(model).__name__): + self.run_test( + model, + inputs, + check_dtype=False, + use_dynamo_tracer=True, + require_full_compilation=True, + ) + + def test_to_copy_float64_dynamic_with_memory_format_after_pytorch_op(self): + # lgamma stays in PyTorch, so the cast ends up alone on an engine input. + class LgammaDouble(nn.Module): + def __init__(self, memory_format): + super().__init__() + self.memory_format = memory_format + + def forward(self, x): + return torch.lgamma(x).to( + torch.float64, memory_format=self.memory_format + ) + + for memory_format in (None, torch.contiguous_format, torch.preserve_format): + with self.subTest(memory_format=memory_format): + model = LgammaDouble(memory_format).eval().cuda() + x = torch.rand(3, 5, device="cuda") + 0.5 + batch = torch.export.Dim("batch", min=1, max=8) + exported = torch.export.export( + model, (x,), dynamic_shapes=({0: batch},) + ) + compiled = torch_tensorrt.dynamo.compile( + exported, + arg_inputs=[ + torch_tensorrt.Input( + min_shape=(1, 5), + opt_shape=(3, 5), + max_shape=(8, 5), + dtype=torch.float32, + ) + ], + min_block_size=1, + truncate_double=True, + ) + # Only the cast without a memory_format goes to TensorRT. + self.assertEqual( + sum("_run_on_acc" in name for name, _ in compiled.named_children()), + int(memory_format is None), + ) + for n in (1, 3, 8): + x = torch.rand(n, 5, device="cuda") + 0.5 + torch.testing.assert_close(compiled(x), model(x)) + + def test_float64_copies_share_one_truncation_rule(self): + # The float64 rule for _to_copy applies to aten.copy too: a source float32 + # cannot hold exactly, or a result layout TensorRT cannot produce, stays in + # PyTorch, and the other cases convert. + class CopyIntoDouble(nn.Module): + def __init__(self, memory_format): + super().__init__() + dest = torch.zeros(2, 3, 4, 5, dtype=torch.double) + self.register_buffer("dest", dest.to(memory_format=memory_format)) + + def forward(self, x): + return torch.ops.aten.copy.default(self.dest, x * 2) + + class BoolToDouble(nn.Module): + def forward(self, x): + return (x > 0.5).to(torch.float64) * 3 + + cases = ( + (CopyIntoDouble(torch.channels_last), torch.rand(2, 3, 4, 5)), + (CopyIntoDouble(torch.contiguous_format), torch.randint(9, (2, 3, 4, 5))), + (BoolToDouble(), torch.rand(2, 3, 4, 5)), + ) + for model, x in cases: + with self.subTest(type(model).__name__, dtype=x.dtype): + model = model.eval().cuda() + x = x.cuda() + compiled = torch_tensorrt.dynamo.compile( + torch.export.export(model, (x,)), + arg_inputs=[x], + min_block_size=1, + truncate_double=True, + ) + self.assertEqual( + sum("_run_on_acc" in name for name, _ in compiled.named_children()), + 1, + ) + out, expected = compiled(x), model(x) + torch.testing.assert_close(out, expected) + self.assertEqual(out.stride(), expected.stride()) + def test_to_copy_multiple_returns(self): class ToCopyReturns(nn.Module): def forward(self, x): @@ -133,6 +244,58 @@ def to_copy_node(device): self.assertFalse(validator(to_copy_node("cpu"))) self.assertTrue(validator(to_copy_node("cuda"))) + def test_to_copy_validator_float64_keeps_layout_in_pytorch(self): + class DoubleCopy(nn.Module): + def __init__(self, memory_format, view): + super().__init__() + self.memory_format = memory_format + self.view = view + + def forward(self, x): + y = self.view(x + 1) + return torch.ops.aten._to_copy.default( + y, dtype=torch.float64, memory_format=self.memory_format + ) + + def to_copy_node(memory_format, dtype=torch.float32, view=lambda y: y): + x = torch.ones(2, 3, 4, 5, dtype=dtype) + exported = torch.export.export(DoubleCopy(memory_format, view), (x,)) + return next( + node + for node in exported.graph.nodes + if node.target is torch.ops.aten._to_copy.default + ) + + validator = to_copy_dtype_validator(placeholder_only=False) + settings = CompilationSettings(truncate_double=True) + self.assertTrue(validator(to_copy_node(torch.contiguous_format), settings)) + self.assertFalse(validator(to_copy_node(torch.channels_last), settings)) + # Copying a strided input gives a contiguous result, which TensorRT matches. + self.assertTrue( + validator( + to_copy_node(torch.preserve_format, view=lambda y: y[..., ::2]), + settings, + ) + ) + self.assertFalse( + validator( + to_copy_node(torch.preserve_format, view=lambda y: y.transpose(1, 2)), + settings, + ) + ) + # TensorRT cannot cast float8 to float32. + self.assertFalse( + validator( + to_copy_node(torch.contiguous_format, torch.float8_e4m3fn), settings + ) + ) + self.assertFalse( + validator( + to_copy_node(torch.contiguous_format), + CompilationSettings(truncate_double=False), + ) + ) + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/conversion/test_copy_aten.py b/tests/py/dynamo/conversion/test_copy_aten.py index 2752e1faf14..5ac56d31d12 100644 --- a/tests/py/dynamo/conversion/test_copy_aten.py +++ b/tests/py/dynamo/conversion/test_copy_aten.py @@ -17,6 +17,8 @@ class TestCopyConverter(DispatchTestCase): ((1, 10), (1, 10), False), ((2, 3, 4), (2, 3, 4), True), ((2, 3, 4, 5), (2, 3, 4, 5), True), + ((2, 3), (3,), False), + ((2, 3, 4), (2, 1, 4), False), ] ) def test_copy_float(self, input_shape, src_shape, non_blocking): @@ -91,6 +93,23 @@ def forward(self, input, src): input_specs, ) + @parameterized.expand([((5,),), ((1, 5),)]) + def test_dynamic_shape_copy_broadcast(self, src_shape): + class Copy(nn.Module): + def forward(self, input, src): + return torch.ops.aten.copy.default(input, src) + + input_specs = [ + Input( + min_shape=(1, 5), + opt_shape=(3, 5), + max_shape=(8, 5), + dtype=torch.float32, + ), + Input(shape=src_shape, dtype=torch.float32), + ] + self.run_test_with_dynamic_shape(Copy(), input_specs) + def test_copy_mixed_dtype_dest_fp32_src_fp16(self): # Regression for pytorch/TensorRT#4265: aten.copy(self, src) returns a # tensor with self's dtype (src is implicitly cast). The previous diff --git a/tests/py/dynamo/conversion/test_expand_aten.py b/tests/py/dynamo/conversion/test_expand_aten.py index fed09e29671..fe0b5a6e7b0 100644 --- a/tests/py/dynamo/conversion/test_expand_aten.py +++ b/tests/py/dynamo/conversion/test_expand_aten.py @@ -3,6 +3,7 @@ import torch import torch.nn as nn +import torch_tensorrt from parameterized import parameterized from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input @@ -84,6 +85,23 @@ def forward(self, x): ExpandTargetDynamic(), input_specs, use_dynamo_tracer=True ) + def test_expand_static_dim_to_data_dependent_size(self): + # The True count is a runtime target that must equal the static dim of 4. + class ExpandToCount(nn.Module): + def forward(self, x, mask): + return x.expand(mask.nonzero().shape[0], 8) + + mask = torch.zeros(10, dtype=torch.bool, device="cuda") + mask[3:7] = True + inputs = (torch.randn(4, 8, device="cuda"), mask) + compiled = torch_tensorrt.dynamo.compile( + torch.export.export(ExpandToCount(), inputs), + arg_inputs=inputs, + min_block_size=1, + pass_through_build_failures=True, + ) + torch.testing.assert_close(compiled(*inputs), ExpandToCount()(*inputs)) + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/conversion/test_index_put_aten.py b/tests/py/dynamo/conversion/test_index_put_aten.py index 1460e5657c9..5864edb8905 100644 --- a/tests/py/dynamo/conversion/test_index_put_aten.py +++ b/tests/py/dynamo/conversion/test_index_put_aten.py @@ -518,6 +518,62 @@ def forward(self, source_tensor, indices_tensor, value_tensor): torch.allclose(result, torch_output, atol=1e-4, rtol=1e-4) + def test_bool_mask_with_data_dependent_count(self): + # The True count is only known at runtime, while values has a static row count. + class AddRows(torch.nn.Module): + def forward(self, data, mask, values): + data = data.clone() + data[mask, :] += values # Qwen3-VL adds visual embeddings this way + return data + + class SetRows(torch.nn.Module): + def forward(self, data, mask, values): + data = data.clone() + data[mask] = values + return data + + class SetRowsAtColumns(torch.nn.Module): + def forward(self, data, mask, values): + data = data.clone() + data[mask, torch.tensor([0, 2, 4, 5], device=data.device)] = values + return data + + values = torch.randn(4, 8, device="cuda") + mask = torch.zeros(1, 10, dtype=torch.bool, device="cuda") + mask[0, 3:7] = True + one_true = torch.zeros(10, dtype=torch.bool, device="cuda") + one_true[5] = True + cases = [ + (AddRows(), torch.randn(1, 10, 8, device="cuda"), mask, ()), + (SetRows(), torch.randn(10, 8, device="cuda"), mask[0], ()), + # One True entry broadcasts across the four column indices. + ( + SetRowsAtColumns(), + torch.randn(10, 6, 8, device="cuda"), + mask[0], + (one_true,), + ), + ] + for model, data, case_mask, other_masks in cases: + with self.subTest(type(model).__name__): + inputs = (data, case_mask, values) + compiled = torchtrt.dynamo.compile( + torch.export.export(model, inputs), + arg_inputs=inputs, + min_block_size=1, + pass_through_build_failures=True, + ) + for run_mask in (case_mask, *other_masks): + torch.testing.assert_close( + compiled(data, run_mask, values), model(data, run_mask, values) + ) + # A True count that differs from the value rows fails, as in PyTorch, + # which raises IndexError when the op runs outside TensorRT. + fewer = case_mask.clone() + fewer[..., 6] = False + with self.assertRaises((RuntimeError, IndexError)): + compiled(data, fewer, values) + def test_updates_are_cast_to_int64_destination_dtype(self): class IndexPutArgmax(torch.nn.Module): def forward(self, data, idx, scores): diff --git a/tests/py/dynamo/conversion/test_select_aten.py b/tests/py/dynamo/conversion/test_select_aten.py index 15eb1d82a69..260e6f6ab8d 100644 --- a/tests/py/dynamo/conversion/test_select_aten.py +++ b/tests/py/dynamo/conversion/test_select_aten.py @@ -3,9 +3,12 @@ import torch import torch.nn as nn +import torch_tensorrt from parameterized import parameterized +from torch.fx.passes.shape_prop import ShapeProp from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input +from torch_tensorrt.dynamo.conversion.aten_ops_converters import select_validator from .harness import DispatchTestCase @@ -163,6 +166,33 @@ def forward(self, values, index): check_dtype=False, ) + def test_select_uint8_input_compiles(self): + # The uint8 select would fail the engine build, so it must stay in PyTorch. + class SelectThenCast(nn.Module): + def forward(self, frames): + return frames[1].float() / 255 + + frames = torch.randint(0, 256, (2, 8, 8, 3), dtype=torch.uint8, device="cuda") + model = SelectThenCast() + compiled = torch_tensorrt.dynamo.compile( + torch.export.export(model, (frames,)), + arg_inputs=[frames], + min_block_size=1, + pass_through_build_failures=True, + ) + torch.testing.assert_close(compiled(frames), model(frames)) + + def test_select_validator_reads_tensor_meta(self): + class Select(nn.Module): + def forward(self, frames): + return torch.ops.aten.select.int(frames, 0, 1) + + graph_module = torch.fx.symbolic_trace(Select()) + ShapeProp(graph_module).propagate(torch.zeros(2, 4, dtype=torch.uint8)) + node = next(n for n in graph_module.graph.nodes if n.op == "call_function") + self.assertNotIn("val", node.args[0].meta) + self.assertFalse(select_validator(node)) + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/conversion/test_slice_aten.py b/tests/py/dynamo/conversion/test_slice_aten.py index bdf171f5fbb..7a8f2beb82c 100644 --- a/tests/py/dynamo/conversion/test_slice_aten.py +++ b/tests/py/dynamo/conversion/test_slice_aten.py @@ -19,6 +19,9 @@ class TestSliceConverter(DispatchTestCase): ("slice_dim_start_stop_step_max_int", 2, 0, 2**63 - 1, 1), ("slice_dim_start_stop_step_past_end", 2, 0, 2048, 1), ("slice_dim_start_stop_step_none", 2, None, None, 1), + ("slice_dim_start_before_begin", 1, -15, None, 1), + ("slice_dim_stop_before_begin", 1, 0, -15, 1), + ("slice_dim_start_after_stop", 1, 7, 3, 1), ] ) def test_slice(self, _, dim, start, stop, step): diff --git a/tests/py/dynamo/executorch/test_backend.py b/tests/py/dynamo/executorch/test_backend.py index 4cafbdb1ffe..580cbfdb411 100644 --- a/tests/py/dynamo/executorch/test_backend.py +++ b/tests/py/dynamo/executorch/test_backend.py @@ -11,6 +11,7 @@ executorch = pytest.importorskip("executorch.exir") +import tensorrt # noqa: E402 import torch # noqa: E402 from executorch.exir.backend.compile_spec_schema import CompileSpec # noqa: E402 from torch.export.graph_signature import InputKind # noqa: E402 @@ -201,6 +202,7 @@ def test_preprocess_serializes_engine_blob(): assert _named_engine(result) == b"engine-bytes" _, metadata = deserialize_engine(result.processed_bytes) assert metadata.device_id == 2 + assert metadata.tensorrt_version == tensorrt.__version__ assert [binding.name for binding in metadata.io_bindings] == ["x", "y"] assert [binding.is_input for binding in metadata.io_bindings] == [True, False] diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index 8e72788e332..2d26ef22675 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -489,6 +489,22 @@ def test_validate_engine_program_rejects_output_allocator_without_mutation(): assert program.state_dict == {} +@pytest.mark.unit +def test_validate_engine_program_keeps_the_output_shape_error_as_the_cause(): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + program, _, input_node, engine = _engine_program( + lifted=False, requires_output_allocator=True + ) + input_node.meta["val"] = torch.empty(2, 3) + error = RuntimeError("Unable to create TensorRT execution context") + engine.infer_outputs = MagicMock(side_effect=error) + + with pytest.raises(RuntimeError, match="output allocator") as refusal: + export_utils.validate_engine_program(program) + + assert refusal.value.__cause__ is error + + @pytest.mark.unit def test_export_returns_edge_and_forwards_all_options(monkeypatch): export_module, lower = _patch_lowering(monkeypatch) @@ -857,6 +873,45 @@ def forward(self, value): assert output.stat().st_size > 0 +@pytest.mark.unit +def test_engine_with_data_dependent_shape_only_inside_exports_a_pte(tmp_path): + """The boolean mask selection has a data-dependent shape, the engine output does not.""" + + class MaskedSum(torch.nn.Module): + def forward(self, value, mask): + return value + value[mask].sum() + + class MaskedRows(torch.nn.Module): + def forward(self, value, mask): + return value[mask] * 2 + + value = torch.randn(10, 4, device="cuda") + mask = torch.arange(10, device="cuda") % 3 == 0 + + def save(model): + exported = torch.export.export(model.eval().cuda(), (value, mask)) + compiled = torch_tensorrt.dynamo.compile( + exported, + arg_inputs=[value, mask], + min_block_size=1, + pass_through_build_failures=True, + ) + output = tmp_path / f"{type(model).__name__}.pte" + torch_tensorrt.save( + compiled, + str(output), + output_format="executorch", + arg_inputs=(value, mask), + retrace=False, + ) + return output + + assert save(MaskedSum()).stat().st_size > 0 + # An output shape that depends on the data still needs an output allocator. + with pytest.raises(RuntimeError, match="output allocator"): + save(MaskedRows()) + + @pytest.mark.unit def test_lowering_stores_the_staged_engine_bytes_as_named_data(monkeypatch, caplog): """The rewrite stages each engine as ``bytes``, and lowering must hand that same diff --git a/tests/py/dynamo/executorch/test_load_errors.py b/tests/py/dynamo/executorch/test_load_errors.py new file mode 100644 index 00000000000..3e19982536f --- /dev/null +++ b/tests/py/dynamo/executorch/test_load_errors.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""Load real TensorRT programs in ways that have to fail, and check each fails with a clear error.""" + +import importlib.util +import os +import subprocess +import sys + +import pytest + +pytest.importorskip("executorch.exir") + +import torch # noqa: E402 + + +@pytest.fixture(autouse=True) +def requires_runtime(): + # Decided at run time, not collection time, for the reason test_cuda_partitioner_composition.py + # gives: remote GPU runners collect off the GPU host. + if not torch.cuda.is_available(): + pytest.skip("needs a CUDA device to build a TensorRT engine") + if importlib.util.find_spec("torch_tensorrt_executorch_runtime") is None: + pytest.skip( + "needs the torch-tensorrt-executorch-runtime package to load the program" + ) + + +def _save(path, **kwargs): + import torch_tensorrt + + torch.manual_seed(0) + inputs = (torch.randn(4, 16, device="cuda"),) + model = torch.nn.Linear(16, 16).eval().cuda() + trt_gm = torch_tensorrt.dynamo.compile( + torch.export.export(model, inputs), inputs=list(inputs), min_block_size=1 + ) + torch_tensorrt.save( + trt_gm, + str(path), + output_format="executorch", + retrace=False, + arg_inputs=list(inputs), + **kwargs, + ) + + +def test_loading_with_no_visible_gpu_fails_with_a_clear_error(tmp_path): + from executorch.exir import ExecutorchBackendConfig + + pte = tmp_path / "model.pte" + # Planned in host memory, so loading reaches the delegate. A device-planned buffer would stop + # the load earlier, in ExecuTorch's CUDA allocator, which reports a missing GPU on its own. + _save( + pte, + backend_config=ExecutorchBackendConfig(enable_non_cpu_memory_planning=False), + ) + script = ( + "import sys\n" + "from pathlib import Path\n" + "import torch_tensorrt_executorch_runtime\n" + "from executorch.runtime import Runtime\n" + "program = Runtime.get().load_program(Path(sys.argv[1]))\n" + "try:\n" + " program.load_method('forward')\n" + "except RuntimeError as error:\n" + " print(error)\n" + " sys.exit(3)\n" + ) + result = subprocess.run( + [sys.executable, "-c", script, str(pte)], + env={**os.environ, "CUDA_VISIBLE_DEVICES": ""}, + capture_output=True, + text=True, + timeout=600, + ) + assert result.returncode == 3, (result.returncode, result.stderr[-2000:]) + output = result.stdout + result.stderr + assert "CUDA is unavailable" in output + # The backend reports itself unavailable, which ExecuTorch turns into NotFound (0x20). + assert "0x20" in output or "0x:20" in output, output[-2000:] + + +def test_an_engine_tensorrt_cannot_read_names_both_tensorrt_versions(tmp_path, capfd): + """A real version mismatch needs a second TensorRT install, so this corrupts the stored engine. + + TensorRT then refuses it on the same path it takes for an engine from another version, which + is the path that has to name both versions. + """ + import tensorrt + import torch_tensorrt_executorch_runtime # noqa: F401 + from executorch.exir._serialize._program import deserialize_pte_binary + from executorch.runtime import Runtime + + pte = tmp_path / "model.pte" + _save(pte) + data = bytearray(pte.read_bytes()) + named = deserialize_pte_binary(bytes(data)).named_data + (engine,) = [bytes(buffer) for buffer in named.buffers] + start = data.find(engine) + assert start > 0 and data.find(engine, start + 1) == -1 + data[start : start + 64] = bytes(64) + # A stand-in export version of the same length, so the log can only have read it from the + # program and not from the TensorRT that is loading it. + exported = tensorrt.__version__.encode() + stand_in = b".".join(b"9" * len(part) for part in exported.split(b".")) + assert data.count(exported) == 1 + data[data.find(exported) : data.find(exported) + len(exported)] = stand_in + pte.write_bytes(data) + + program = Runtime.get().load_program(pte) + with pytest.raises(RuntimeError, match="0x:?23") as raised: + program.load_method("forward") + log = str(raised.value) + capfd.readouterr().err + runtime_version = ".".join(tensorrt.__version__.split(".")[:3]) + assert f"exported with TensorRT {stand_in.decode()}" in log, log + assert f"this runtime is TensorRT {runtime_version}" in log, log diff --git a/tests/py/dynamo/executorch/test_serialization.py b/tests/py/dynamo/executorch/test_serialization.py index f2f8ef9514f..c933d6f6161 100644 --- a/tests/py/dynamo/executorch/test_serialization.py +++ b/tests/py/dynamo/executorch/test_serialization.py @@ -25,6 +25,7 @@ def test_serialize_engine_writes_tr01_blob(): ], device_id=1, hardware_compatible=True, + tensorrt_version="11.4.0.106", ) blob = serialize_engine(b"engine-bytes", metadata) @@ -36,6 +37,7 @@ def test_serialize_engine_writes_tr01_blob(): assert engine == b"engine-bytes" assert parsed.device_id == 1 assert parsed.hardware_compatible is True + assert parsed.tensorrt_version == "11.4.0.106" assert [b.name for b in parsed.io_bindings] == ["x", "y"] assert [b.is_input for b in parsed.io_bindings] == [True, False] diff --git a/tests/py/dynamo/hlo/test_complex_ops.py b/tests/py/dynamo/hlo/test_complex_ops.py index 75e1f04262c..9f2351afb26 100644 --- a/tests/py/dynamo/hlo/test_complex_ops.py +++ b/tests/py/dynamo/hlo/test_complex_ops.py @@ -905,6 +905,92 @@ def test_complex_output_dynamic_batch(): torch._dynamo.reset() +class ComplexOp(nn.Module): + """Apply one complex op to a real (..., 2) input viewed as complex.""" + + def __init__(self, op): + super().__init__() + self.op = op + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.view_as_real(self.op(torch.view_as_complex(x))) + + +@pytest.mark.unit +@pytest.mark.parametrize("dynamic", [False, True], ids=["static", "dynamic"]) +@pytest.mark.parametrize( + "op", + [ + lambda z: torch.prod(z, dim=1), + lambda z: torch.prod(z, dim=-1, keepdim=True), + lambda z: z + (1 + 2j), + lambda z: torch.ops.aten.sub.Scalar(z, 1 + 2j), + lambda z: (2 - 1j) * z, + lambda z: z / (2 - 1j), + lambda z: torch.ops.aten.mul.Scalar(z, 2 - 1j), + lambda z: z.real * (2 - 1j), + lambda z: z.imag - (1 + 2j), + lambda z: (z.real * 4).to(torch.int64) + (1.25 + 2.75j), + lambda z: torch.add(z, 1 + 2j, alpha=2), + lambda z: torch.complex(z.real * 2, z.imag), + ], + ids=[ + "prod_dim", + "prod_neg_dim_keepdim", + "add_complex_scalar", + "sub_complex_scalar", + "mul_complex_scalar", + "div_complex_scalar", + "mul_scalar_overload_complex_scalar", + "real_tensor_mul_complex_scalar", + "real_tensor_sub_complex_scalar", + "int_tensor_add_complex_scalar", + "add_complex_scalar_alpha", + "complex_from_parts", + ], +) +def test_complex_op_matches_eager(op, dynamic): + model = ComplexOp(op).eval().cuda() + x = torch.randn(2, 4, 2).cuda() + dynamic_shapes = ({0: Dim("batch", min=2, max=8)},) if dynamic else None + ep = torch.export.export(model, (x,), dynamic_shapes=dynamic_shapes) + trt_model = torchtrt.dynamo.compile( + ep, + inputs=[x], + min_block_size=1, + pass_through_build_failures=True, + **_DECOMP, + ) + for batch in (2, 5) if dynamic else (2,): + xb = torch.randn(batch, 4, 2).cuda() + torch.testing.assert_close(trt_model(xb), model(xb), rtol=1e-3, atol=1e-3) + torch._dynamo.reset() + + +class ComplexDivBySize(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + z = torch.view_as_complex(x) + size = torch.sym_float(x.shape[0]) + return torch.view_as_real(torch.ops.aten.div.Scalar(z, size)) + + +@pytest.mark.unit +def test_complex_divided_by_symbolic_size(): + # Under torch.compile with dynamic shapes the divisor stays a graph value. + model = ComplexDivBySize().eval().cuda() + compiled = torch.compile( + model, + backend="torch_tensorrt", + fullgraph=True, + dynamic=True, + options={"min_block_size": 1, "pass_through_build_failures": True, **_DECOMP}, + ) + for batch in (3, 5): + x = torch.randn(batch, 4, 2).cuda() + torch.testing.assert_close(compiled(x), model(x), rtol=1e-3, atol=1e-3) + torch._dynamo.reset() + + # =========================================================================== # 11. Numerical precision: complex64 vs truncated complex128 # =========================================================================== diff --git a/tests/py/dynamo/lowering/test_aten_lowering_passes.py b/tests/py/dynamo/lowering/test_aten_lowering_passes.py index 80d531f29b4..110a5e644d0 100644 --- a/tests/py/dynamo/lowering/test_aten_lowering_passes.py +++ b/tests/py/dynamo/lowering/test_aten_lowering_passes.py @@ -454,6 +454,98 @@ def test_eliminates_noop_sym_min_int64_max(self): self.assertEqual(output_node.args[0], (x, x)) +class TestPreserveAttentionOutputLayout(TestCase): + def test_view_after_decomposed_attention(self): + # Fused attention returns this transposed query's layout, so export drops the + # .contiguous() below and records a view that only works on that layout. The + # attention output must stay a tensor of its own (detach_), the in-place + # update through the view must reach it, and the no_grad region puts all of + # it in a subgraph. + class ViewAfterAttention(torch.nn.Module): + def forward(self, query, key, value): + with torch.no_grad(): + query, key, value = (t.transpose(1, 2) for t in (query, key, value)) + out = torch.nn.functional.scaled_dot_product_attention( + query, key, value + ) + out.detach_() + flat = out.transpose(1, 2).contiguous() + flat = flat.view(flat.shape[0], flat.shape[1], -1) + flat.add_(1) + return out, flat + + inputs = [torch.randn(2, 16, 4, 32).half().cuda() for _ in range(3)] + model = ViewAfterAttention().cuda() + ep = torch.export.export(model, tuple(inputs)) + # With a contiguous attention output, export keeps the .contiguous() and + # this test passes without the pass. + (attention,) = ( + node + for module in ep.graph_module.modules() + if isinstance(module, torch.fx.GraphModule) + for node in module.graph.nodes + if node.target == torch.ops.aten.scaled_dot_product_attention.default + ) + self.assertFalse(attention.meta["val"].is_contiguous()) + trt_module = torch_tensorrt.dynamo.compile( + ep, inputs, min_block_size=1, decompose_attention=True + ) + torch.testing.assert_close( + model(*inputs), trt_module(*inputs), rtol=1e-2, atol=1e-2 + ) + + def test_view_after_attention_in_cond_branch(self): + class CondViewAfterAttention(torch.nn.Module): + def forward(self, pred, query, key, value): + def attend(query, key, value): + query, key, value = (t.transpose(1, 2) for t in (query, key, value)) + out = torch.nn.functional.scaled_dot_product_attention( + query, key, value + ) + out = out.transpose(1, 2).contiguous() + return out.view(out.shape[0], out.shape[1], -1) + + def flatten(query, key, value): + return query.flatten(2).clone() + + return torch.cond(pred, attend, flatten, (query, key, value)) + + inputs = [torch.randn(2, 16, 4, 32).half().cuda() for _ in range(3)] + model = CondViewAfterAttention().cuda() + pred = torch.tensor(True).cuda() + ep = torch.export.export(model, (pred, *inputs)) + trt_module = torch_tensorrt.dynamo.compile( + ep, [pred, *inputs], min_block_size=1, decompose_attention=True + ) + for pred in (True, False): + pred = torch.tensor(pred).cuda() + torch.testing.assert_close( + model(pred, *inputs), trt_module(pred, *inputs), rtol=1e-2, atol=1e-2 + ) + + def test_attention_on_data_dependent_rows(self): + # Boolean indexing gives attention a batch size known only at runtime. + class MaskedAttention(torch.nn.Module): + def forward(self, x, mask): + query = x[mask].transpose(1, 2) + return torch.nn.functional.scaled_dot_product_attention( + query, query, query + ) + + model = MaskedAttention().cuda() + x = torch.randn(3, 16, 4, 32).half().cuda() + mask = torch.tensor([True, False, True]).cuda() + ep = torch.export.export(model, (x, mask)) + trt_module = torch_tensorrt.dynamo.compile( + ep, [x, mask], min_block_size=1, decompose_attention=True + ) + for mask in ([True, False, False], [True, True, True]): + mask = torch.tensor(mask).cuda() + torch.testing.assert_close( + model(x, mask), trt_module(x, mask), rtol=1e-2, atol=1e-2 + ) + + class TestRewriteEfficientAttention(TestCase): def test_force_causal_efficient_attention(self): class RewriteEfficientAttention(torch.nn.Module): diff --git a/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py b/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py index 66b0d6a1f0e..6b19c929ba9 100644 --- a/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py +++ b/tests/py/dynamo/lowering/test_attention_mask_constant_fold_exclusion.py @@ -46,11 +46,16 @@ def _export(self): return torch.export.export(self.AttentionWithCausalMask(), inputs) def _assert_only_attention_aranges_survive( - self, decompose_attention, disabled_constant_fold_exclusions=() + self, + decompose_attention, + disabled_constant_fold_exclusions=(), + decompositions_runs=1, ): - exported_program = self._export().run_decompositions( - get_decompositions(decompose_attention=decompose_attention) - ) + exported_program = self._export() + for _ in range(decompositions_runs): + exported_program = exported_program.run_decompositions( + get_decompositions(decompose_attention=decompose_attention) + ) gm = post_lowering( exported_program.module(), CompilationSettings( @@ -78,6 +83,48 @@ def test_decomposed_attention_mask_aranges_are_not_folded(self): def test_native_attention_mask_aranges_are_not_folded(self): self._assert_only_attention_aranges_survive(decompose_attention=False) + def test_decomposed_attention_mask_aranges_survive_a_second_decomposition(self): + self._assert_only_attention_aranges_survive( + decompose_attention=True, decompositions_runs=2 + ) + + def test_decomposed_attention_mask_aranges_survive_complex_retrace(self): + class AttentionWithComplexOp(self.AttentionWithCausalMask): + def forward(self, query, key, value, attention_mask, rotation): + attention, unrelated_arange = super().forward( + query, key, value, attention_mask + ) + rotation = torch.view_as_complex(rotation) + return attention, unrelated_arange, torch.view_as_real(rotation**2) + + inputs = ( + torch.randn(1, 2, 8, 16, device="cuda"), + torch.randn(1, 2, 8, 16, device="cuda"), + torch.randn(1, 2, 8, 16, device="cuda"), + torch.ones(8, 8, dtype=torch.bool, device="cuda"), + torch.randn(4, 2, device="cuda"), + ) + exported_program = torch.export.export( + AttentionWithComplexOp(), inputs + ).run_decompositions(get_decompositions(decompose_attention=True)) + gm = post_lowering( + exported_program.module(), + CompilationSettings(use_complex_decomposition=True), + ) + arange_nodes = [ + node + for node in gm.graph.nodes + if node.op == "call_function" + and getattr(node.target, "overloadpacket", None) is torch.ops.aten.arange + ] + self.assertEqual(len(arange_nodes), 2) + self.assertTrue( + all( + node.meta.get(CONSTANT_FOLD_EXCLUSION_META_KEY, False) + for node in arange_nodes + ) + ) + def test_native_attention_rules_can_be_disabled(self): exported_program = self._export().run_decompositions( get_decompositions(decompose_attention=False) @@ -100,9 +147,9 @@ def test_native_attention_rules_can_be_disabled(self): def test_decomposed_attention_rules_can_be_disabled(self): """post_lowering is the only place a rule has to be disabled. - The decompositions mark unconditionally while tracing, well before any - settings object is reachable, so post_lowering revokes those marks - instead of the caller having to communicate the setting twice. + The decompositions tag unconditionally while tracing, well before any + settings object is reachable, and post_lowering skips the disabled rule + that would turn the tags into marks, so the caller sets it only once. """ exported_program = self._export().run_decompositions( get_decompositions(decompose_attention=True) diff --git a/tests/py/dynamo/lowering/test_constant_folding.py b/tests/py/dynamo/lowering/test_constant_folding.py index 5648932eea9..6df25a789a0 100644 --- a/tests/py/dynamo/lowering/test_constant_folding.py +++ b/tests/py/dynamo/lowering/test_constant_folding.py @@ -3,11 +3,13 @@ import copy import unittest +from unittest import mock import torch import torch_tensorrt from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.lowering.passes import complex_decomposition_adapter from torch_tensorrt.dynamo.lowering.passes.constant_folding import constant_fold @@ -154,6 +156,107 @@ def test_aliased_view_models_compile_end_to_end(self): torch.testing.assert_close(actual, expected) + @unittest.skipIf( + not torch.cuda.is_available() + or not complex_decomposition_adapter.has_complex_decomposition(), + "needs CUDA and the upstream complex decomposition", + ) + def test_offloaded_constants_reach_the_complex_decomposition(self): + # offload_module_to_cpu leaves folded constants on the CPU; the complex + # decomposition retraces the graph and must not fall back over them. + class Rotary(torch.nn.Module): + def __init__(self): + super().__init__() + angles = torch.outer(torch.arange(16.0), torch.ones(4)) + self.register_buffer( + "freqs", torch.polar(torch.ones_like(angles), angles) + ) + # Read only by constant folding, so the retraced module drops it. + self.scale = torch.nn.Parameter(torch.ones(4)) + + def forward(self, x): + freqs = self.freqs[: x.shape[0]] * self.scale.mean() + z = torch.view_as_complex(x.reshape(x.shape[0], 4, 2)) + return torch.view_as_real(z * freqs).flatten(1) + + model = Rotary().eval().cuda() + example_input = torch.randn(8, 8, device="cuda") + with ( + torch.no_grad(), + mock.patch.object( + complex_decomposition_adapter, + "complex_graph_detection", + side_effect=AssertionError("fell back to the legacy complex pass"), + ), + ): + expected = model(example_input) + exported_program = torch.export.export(model, (example_input,)) + compiled_model = torch_tensorrt.dynamo.compile( + exported_program, + inputs=[example_input], + min_block_size=1, + offload_module_to_cpu=True, + ) + actual = compiled_model(example_input) + + torch.testing.assert_close(actual, expected) + self.assertEqual(model.scale.device.type, "cpu") + + @unittest.skipIf( + not torch.cuda.is_available() + or not complex_decomposition_adapter.has_complex_decomposition(), + "needs CUDA and the upstream complex decomposition", + ) + def test_complex_fallback_does_not_hold_device_copies_of_offloaded_constants( + self, + ): + from torch._functorch._aot_autograd import complex_decomposition + + class Rotary(torch.nn.Module): + def __init__(self): + super().__init__() + angles = torch.outer(torch.arange(256.0), torch.ones(512)) + self.register_buffer( + "freqs", torch.polar(torch.ones_like(angles), angles) + ) + self.scale = torch.nn.Parameter(torch.ones(512)) + + def forward(self, x): + z = torch.view_as_complex(x.reshape(256, 512, 2)) + return torch.view_as_real(z * (self.freqs * self.scale.mean())) + + settings = CompilationSettings(offload_module_to_cpu=True) + example_input = torch.randn(256, 1024, device="cuda") + exported_program = torch.export.export(Rotary().cuda(), (example_input,)) + # Folds freqs * scale into a 1 MiB constant kept on the CPU. + gm = constant_fold(exported_program.module(), settings) + + def failing_retrace(gm, flat_args, decompositions): + # Like a real retrace, its frame holds the constants it was handed. + constants = [ + getattr(gm, n.target) for n in gm.graph.nodes if n.op == "get_attr" + ] + raise NotImplementedError(f"unsupported op, {len(constants)} constants") + + allocated = {} + + def legacy_pass(gm, settings): + allocated["fallback"] = torch.cuda.memory_allocated() + return gm + + with ( + mock.patch.object( + complex_decomposition, "decompose_complex_in_graph", failing_retrace + ), + mock.patch.object( + complex_decomposition_adapter, "complex_graph_detection", legacy_pass + ), + ): + allocated["before"] = torch.cuda.memory_allocated() + complex_decomposition_adapter.complex_decomposition_adapter(gm, settings) + + self.assertEqual(allocated["fallback"], allocated["before"]) + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/lowering/test_decompositions.py b/tests/py/dynamo/lowering/test_decompositions.py index 629f88f0501..df3fbde79f7 100644 --- a/tests/py/dynamo/lowering/test_decompositions.py +++ b/tests/py/dynamo/lowering/test_decompositions.py @@ -1362,17 +1362,14 @@ def forward(self, base, index, src): model = ScatterAdd().eval() base = torch.zeros((16, 8)) index = torch.zeros((4, 8), dtype=torch.int64) - src = torch.ones((4, 8)) - src_rows = torch.export.Dim("src_rows", min=1, max=16) + src = torch.ones((16, 8)) + index_rows = torch.export.Dim("index_rows", min=1, max=16) + # The decomposition loops over index, so only its extent matters. dynamic_export = torch.export.export( model, (base, index, src), - dynamic_shapes={ - "base": {}, - "index": {0: src_rows}, - "src": {0: src_rows}, - }, + dynamic_shapes={"base": {}, "index": {0: index_rows}, "src": {}}, ) base_decompositions = get_decompositions() dynamic_decompositions = filter_decomposition_table( @@ -1395,9 +1392,9 @@ def forward(self, base, index, src): for node in dynamic_export.graph_module.graph.nodes if node.target == torch.ops.aten.scatter_add.default ) - dynamic_src = dynamic_scatter.args[3] - self.assertIsInstance(dynamic_src, torch.fx.Node) - dynamic_src.meta["example_value"] = dynamic_src.meta.pop("val") + dynamic_index = dynamic_scatter.args[2] + self.assertIsInstance(dynamic_index, torch.fx.Node) + dynamic_index.meta["example_value"] = dynamic_index.meta.pop("val") dynamo_decompositions = filter_decomposition_table( get_decompositions(), dynamic_export.graph_module ) @@ -1409,6 +1406,95 @@ def forward(self, base, index, src): ) self.assertIn(torch.ops.aten.scatter_add.default, static_decompositions) + @parameterized.expand( + [ + ("rule_enabled", (), 1, "input", True), + ("rule_disabled", {"scatter_add_positions"}, 1, "input", False), + ("decomposed_twice", (), 2, "input", True), + ("constant_src_folds", (), 1, "ones", False), + ("random_src", (), 1, "randint", True), + # Constant src that constant folding still leaves in the graph. + ("attention_mask_src", (), 1, "attention_mask_arange", True), + ("returned_src_view", (), 1, "returned_view", True), + ] + ) + def test_scatter_add_constant_index_positions_not_folded( + self, + _, + disabled_constant_fold_exclusions, + decompositions_runs, + src_kind, + positions_kept, + ): + class ConstantIndexScatterAdd(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("index", torch.randint(0, 64, (16, 8))) + + def forward(self, input, src): + if src_kind == "ones": + src = torch.ones_like(self.index) + elif src_kind == "randint": + src = torch.randint_like(self.index, 5) + elif src_kind == "attention_mask_arange": + sequence = torch.arange(8, device=input.device) + src = sequence.expand_as(self.index) + elif src_kind == "returned_view": + src = self.index.t().t() + result = input.scatter_add(0, self.index, src) + if src_kind == "attention_mask_arange": + # The attention comes after the scatter_add positions in the graph. + mask = sequence.unsqueeze(0) <= sequence.unsqueeze(1) + query = input[:8].float().view(1, 1, 8, 8) + attention = torch.ops.aten.scaled_dot_product_attention.default( + query, query, query, mask + ) + return result, attention + if src_kind == "returned_view": + return result, src + return result + + exported = torch.export.export( + ConstantIndexScatterAdd().cuda(), + ( + torch.zeros(64, 8, dtype=torch.int64).cuda(), + torch.ones(16, 8, dtype=torch.int64).cuda(), + ), + ) + for _ in range(decompositions_runs): + exported = exported.run_decompositions(get_decompositions()) + gm = post_lowering( + exported.module(), + CompilationSettings( + disabled_constant_fold_exclusions=disabled_constant_fold_exclusions + ), + ) + has_positions = any( + getattr(node.target, "overloadpacket", None) is torch.ops.aten.arange + and node.meta["val"].shape == (64,) + for node in gm.graph.nodes + ) + self.assertEqual(has_positions, positions_kept) + + def test_scatter_add_large_comparison_keeps_loop(self): + class ScatterAdd(torch.nn.Module): + def forward(self, input, index, src): + return input.scatter_add(0, index, src) + + def lowered_targets(updates): + inputs = ( + torch.zeros(2**20, dtype=torch.int64), + torch.zeros(updates, dtype=torch.int64), + # src longer than index also exercises the loop above the bound. + torch.ones(updates + 1, dtype=torch.int64), + ) + exported = torch.export.export(ScatterAdd(), inputs) + lowered = exported.run_decompositions(get_decompositions()) + return {node.target for node in lowered.graph.nodes} + + self.assertNotIn(torch.ops.aten.scatter.src, lowered_targets(64)) + self.assertIn(torch.ops.aten.scatter.src, lowered_targets(128)) + @parameterized.expand( [ ( @@ -1448,12 +1534,22 @@ def forward(self, base, index, src): 1, torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1], [3, 2, 1, 2]]).cuda(), torch.tensor( - [[1, 2, 3, 1], [5, 6, 5, 5], [2, 4, 3, 2]], dtype=torch.int32 + [[1, 2, 3, 1], [5, 6, 5, 5], [2, 4, 3, 2]], dtype=torch.float32 ).cuda(), { torch.ops.aten.add.Tensor, }, ), + ( + "scatter_add_zero_dim_src_larger_than_index", + 0, + torch.tensor([[0, 1, 2, 0], [2, 2, 0, 1]]).cuda(), + torch.tensor( + [[1, 2, 3, 4, 5], [6, 2**40, 8, 9, 10], [11, 12, 13, 14, 15]], + dtype=torch.int64, + ).cuda(), + {torch.ops.aten.add.Tensor}, + ), ] ) def test_scatter_add(self, _, dim, index, src, expected_ops_param): @@ -1467,8 +1563,11 @@ def forward(self, input): # Operations expected to be included in the traced graph after decompositions expected_ops = expected_ops_param unexpected_ops = {torch.ops.aten.scatter_add.default} + if not src.is_floating_point(): + unexpected_ops.add(torch.ops.aten.scatter.src) - input = torch.zeros(3, 5, dtype=torch.int32).cuda() + # Nonzero, so cells without updates must keep their values. + input = torch.arange(15).reshape(3, 5).to(src.dtype).cuda() inputs = [input] fx_graph = torch.fx.symbolic_trace(TestModule()) diff --git a/tests/py/dynamo/models/test_cond.py b/tests/py/dynamo/models/test_cond.py index ee5ce8c5dfb..4fa93d21a02 100644 --- a/tests/py/dynamo/models/test_cond.py +++ b/tests/py/dynamo/models/test_cond.py @@ -12,6 +12,7 @@ post_lowering, pre_export_lowering, ) +from torch_tensorrt.dynamo.lowering.passes._aten_lowering_pass import _same_io def _cond_node(gm: torch.fx.GraphModule) -> torch.fx.Node: @@ -57,6 +58,71 @@ def forward(self, x, predicate): ) +class _ComplexCond(nn.Module): + def forward(self, x, predicate): + def rotate(value): + z = torch.view_as_complex(value.reshape(2, 2, 2)) + return torch.view_as_real(z * z).reshape(2, 4) + + return torch.cond(predicate, rotate, lambda value: value - 1, (x,)) + + +class _TransposedComplexCond(nn.Module): + def forward(self, x, predicate): + def rotate(value): + z = torch.view_as_complex(value) + return torch.view_as_real((z * z).transpose(0, 1)) + + result = torch.cond( + predicate, rotate, lambda value: (value - 1).transpose(0, 1), (x,) + ) + return torch.as_strided(result, (24,), (1,)) + + +class _SlicedComplexCond(nn.Module): + def __init__(self, rows): + super().__init__() + self.rows = rows + + def forward(self, x, predicate): + def rotate(value): + z = torch.view_as_complex(value) + return torch.view_as_real((z * z)[self.rows]) + + result = torch.cond( + predicate, rotate, lambda value: (value - 1)[self.rows], (x,) + ) + return torch.as_strided(result, (24,), (1,), 0) + + +class _DataDependentComplexCond(nn.Module): + def forward(self, x, n, predicate): + k = n.item() + torch._check(k >= 1) + torch._check(k <= x.shape[0]) + + def rotate(value): + z = torch.view_as_complex(value) + return torch.view_as_real((z * z).transpose(0, 1)) + + return torch.cond( + predicate, + rotate, + lambda value: (value - 1).transpose(0, 1), + (x[:k].contiguous(),), + ) + + +class _ComplexOutputCond(nn.Module): + def forward(self, x, predicate): + return torch.cond( + predicate, + lambda value: torch.view_as_complex(value + 1), + lambda value: torch.view_as_complex(value - 1), + (x,), + ) + + def _export_lowered(mod: nn.Module, inputs: tuple) -> torch.fx.GraphModule: settings = CompilationSettings() exported = torch.export.export(mod.eval(), inputs) @@ -73,6 +139,61 @@ def test_cond_is_supported_when_branches_are_convertible(self): gm = _export_lowered(_AddSubCond(), (x, pred)) self.assertTrue(_cond_node(gm) in CONVERTERS) + def test_cond_keeps_branches_returned_by_lowering(self): + # The complex lowering returns a new branch module rather than editing it. + x = torch.randn(2, 4) + model = _ComplexCond() + gm = _export_lowered(model, (x, torch.tensor(True))) + self.assertTrue(_cond_node(gm) in CONVERTERS) + for pred in (torch.tensor(True), torch.tensor(False)): + torch.testing.assert_close(gm(x, pred), model(x, pred)) + + def test_cond_keeps_branch_when_lowering_changes_output_layout(self): + pred = torch.tensor(True) + cases = [ + # The lowered branch is contiguous, so as_strided reads another order. + (_TransposedComplexCond(), (torch.randn(3, 4, 2), pred)), + # The lowered branches own smaller storages, one also at another offset. + (_SlicedComplexCond(slice(1, None)), (torch.randn(4, 3, 2), pred)), + (_SlicedComplexCond(slice(None, -1)), (torch.randn(4, 3, 2), pred)), + # Data-dependent strides cannot be proven equal. + ( + _DataDependentComplexCond(), + (torch.randn(8, 4, 2), torch.tensor(5), pred), + ), + ] + for model, inputs in cases: + with self.subTest(type(model).__name__): + gm = _export_lowered(model, inputs) + torch.testing.assert_close(gm(*inputs), model(*inputs)) + + def test_cond_keeps_branch_when_lowering_changes_output_dtype(self): + # The lowered branch returns the real [..., 2] layout instead of complex. + x = torch.randn(2, 2, 2) + pred = torch.tensor(True) + model = _ComplexOutputCond() + gm = _export_lowered(model, (x, pred)) + torch.testing.assert_close(gm(x, pred), model(x, pred)) + + def test_cond_branch_check_compares_sizes_and_device(self): + def identity(value): + graph = torch.fx.Graph() + node = graph.placeholder("x") + node.meta["val"] = value + graph.output(node) + return torch.fx.GraphModule(torch.nn.Module(), graph) + + storage = torch.empty(12) + cases = { + # Same dtype, strides, storage size and offset. + "sizes": (storage[:8].view(2, 4), storage.view(3, 4)), + "device": (torch.empty(2, 4), torch.empty(2, 4, device="meta")), + } + for name, (a, b) in cases.items(): + with self.subTest(name): + self.assertTrue(_same_io(identity(a), identity(a))) + self.assertFalse(_same_io(identity(a), identity(b))) + def test_cond_falls_back_when_branch_has_unsupported_op(self): x = torch.randn(1, 4).abs() + 0.1 pred = torch.tensor(True) diff --git a/tests/py/dynamo/models/test_model_refit.py b/tests/py/dynamo/models/test_model_refit.py index 491283088fb..a6b1b19eb64 100644 --- a/tests/py/dynamo/models/test_model_refit.py +++ b/tests/py/dynamo/models/test_model_refit.py @@ -14,6 +14,7 @@ import torch_tensorrt as torchtrt import torch_tensorrt as torch_trt from torch import nn +from torch_tensorrt._features import has_complex_decomposition from torch_tensorrt._utils import is_tensorrt_rtx_version_supported from torch_tensorrt.dynamo import partitioning, refit_module_weights from torch_tensorrt.dynamo._refit import ( @@ -1037,8 +1038,16 @@ def forward(self, x): "Refit feature is not supported in Python 3.13 or higher", ) @pytest.mark.unit +@pytest.mark.parametrize( + "complex_output", [False, True], ids=["real_out", "complex_out"] +) @pytest.mark.parametrize("dynamic_batch", [False, True]) -def test_complex_buffer_refit(dynamic_batch, caplog): +@pytest.mark.parametrize( + "use_complex_decomposition", [False, True], ids=["legacy", "decomp"] +) +def test_complex_buffer_refit( + use_complex_decomposition, dynamic_batch, complex_output, caplog +): """Refit a model whose weights include a complex-valued buffer (e.g. RoPE freqs). Exercises the combined complex_graph_detection + refit_module_weights path: @@ -1046,6 +1055,8 @@ def test_complex_buffer_refit(dynamic_batch, caplog): - complex placeholder input goes through view_as_real at the TRT boundary - after refitting with new frequencies the output matches the new model """ + if use_complex_decomposition and not has_complex_decomposition(): + pytest.skip("decompose_complex_in_graph requires torch>=2.14.dev") class ComplexFreqModel(nn.Module): def __init__(self, freqs: torch.Tensor): @@ -1053,8 +1064,8 @@ def __init__(self, freqs: torch.Tensor): self.register_buffer("freqs", freqs.cuda()) def forward(self, z: torch.Tensor) -> torch.Tensor: - # complex mul then expose as real so TRT can produce a real output - return torch.view_as_real(z * self.freqs) + out = z * self.freqs + return out if complex_output else torch.view_as_real(out) SEQ, DIM = 8, 32 @@ -1087,6 +1098,7 @@ def make_freqs() -> torch.Tensor: tuple(inputs), min_block_size=1, immutable_weights=False, + use_complex_decomposition=use_complex_decomposition, ) with caplog.at_level(logging.INFO, logger="torch_tensorrt.dynamo._refit"): diff --git a/tests/py/dynamo/partitioning/test_000_fast_partitioning.py b/tests/py/dynamo/partitioning/test_000_fast_partitioning.py index 6a666ca3dba..76dc3e8e337 100644 --- a/tests/py/dynamo/partitioning/test_000_fast_partitioning.py +++ b/tests/py/dynamo/partitioning/test_000_fast_partitioning.py @@ -1,18 +1,33 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause +from collections import defaultdict from copy import deepcopy import numpy as np import torch import torch_tensorrt from parameterized import parameterized +from torch.fx.passes.splitter_base import _SplitterBase +from torch.fx.passes.tools_common import CALLABLE_NODE_OPS from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt.dynamo import partitioning +from torch_tensorrt.dynamo.partitioning._adjacency_partitioner import ( + OpSupportTester, + TRTPartitioner, +) from ..testing_utilities import lower_graph_testing +class _CountingDeps(defaultdict): + reads = 0 + + def __getitem__(self, node): + self.reads += 1 + return super().__getitem__(node) + + class TestFastPartitioning(TestCase): def test_partition_fully_supported_one_op(self): class FullySupportedOneOp(torch.nn.Module): @@ -183,6 +198,39 @@ def forward(self, x, y): "Certain operators are set to run in Torch, expected 1 segment", ) + def test_put_nodes_into_subgraphs_matches_base_without_rescans(self): + class Interleaved(torch.nn.Module): + def forward(self, x, y): + for i in range(60): + a = torch.ops.aten.add.Tensor(x, y) + b = torch.ops.aten.mul.Tensor(a, y) + if i % 6 == 0: + b = torch.ops.aten.abs.default(b) + x = torch.ops.aten.sub.Tensor(b, a) + return x + + partitioner = TRTPartitioner( + torch.fx.symbolic_trace(Interleaved()), + OpSupportTester(torch_executed_ops={"torch.ops.aten.abs.default"}), + min_block_size=1, + ) + partitioner.deps = _CountingDeps(set, partitioner.deps) + subgraphs = partitioner.put_nodes_into_subgraphs() + reads = partitioner.deps.reads + + expected = _SplitterBase.put_nodes_into_subgraphs(partitioner) + self.assertEqual( + [(s.is_acc, set(s.nodes)) for s in subgraphs], + [(s.is_acc, set(s.nodes)) for s in expected], + ) + self.assertEqual(len(subgraphs), 21) + + # The base loop can reread the dependencies of all waiting nodes at every step. + num_nodes = len( + [n for n in partitioner.module.graph.nodes if n.op in CALLABLE_NODE_OPS] + ) + self.assertLessEqual(reads, 4 * num_nodes) + @parameterized.expand( [ (["torch.nn.modules.conv.Conv2d"], 2), diff --git a/tests/py/dynamo/runtime/test_host_shape_inputs.py b/tests/py/dynamo/runtime/test_host_shape_inputs.py index f607ebba86f..49d312a18fb 100644 --- a/tests/py/dynamo/runtime/test_host_shape_inputs.py +++ b/tests/py/dynamo/runtime/test_host_shape_inputs.py @@ -3,6 +3,7 @@ import torch_tensorrt from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt.dynamo.runtime import TorchTensorRTModule +from torch_tensorrt.dynamo.runtime._TRTEngine import TRTEngine class ShapeInputModel(nn.Module): @@ -69,6 +70,40 @@ def test_host_and_device_shape_input_agree(self): device_out = self._run_engine(trt_mod, x, 64, "cuda") torch.testing.assert_close(host_out, device_out, rtol=1e-5, atol=1e-5) + def test_python_runtime_keeps_every_shape_input(self): + class ThreeShapeInputModel(nn.Module): + def forward(self, x, a, b, c): + ranges = [torch.arange(0, n, 1, device=x.device) for n in (a, b, c)] + return torch.cat(ranges).to(torch.float32) + x.sum() + + model = ThreeShapeInputModel().eval().cuda() + inputs = ( + torch.randn(8, device="cuda"), + *(torch.tensor(n, dtype=torch.int64, device="cuda") for n in (3, 5, 7)), + ) + gm = torch_tensorrt.dynamo.compile( + torch.export.export(model, inputs), + inputs=list(inputs), + min_block_size=1, + pass_through_build_failures=True, + assume_dynamic_shape_support=True, + ) + trt_mod = next( + m for _, m in gm.named_children() if isinstance(m, TorchTensorRTModule) + ) + # A build with the C++ runtime backs the module with the C++ engine, so + # drive the Python engine directly from the same serialized info. + engine = TRTEngine(trt_mod._pack_engine_info()) + shape_names = sorted( + (name for name in engine.in_binding_names if name != "x"), + key=lambda name: (len(name), name), + ) + args = {"x": inputs[0], **dict(zip(shape_names, inputs[1:]))} + output = engine.execute([args[name] for name in engine.in_binding_names]) + if isinstance(output, (list, tuple)): + (output,) = output + torch.testing.assert_close(output, model(*inputs)) + if __name__ == "__main__": run_tests()