diff --git a/python/tensorrt_model_connect/families/moge/model.py b/python/tensorrt_model_connect/families/moge/model.py index 08e49ecfb9..12b2d535cd 100644 --- a/python/tensorrt_model_connect/families/moge/model.py +++ b/python/tensorrt_model_connect/families/moge/model.py @@ -29,9 +29,66 @@ _PATCH = 14 _POSITION_GRID = 37 _NUM_TOKENS = 1800 +_FOCAL_RECOVERY_SIZE = 64 _MIN_IMAGE_SIZE = 64 _OPT_IMAGE_SIZE = 518 _MAX_IMAGE_SIZE = 2048 +_FAST_MIN_IMAGE_HEIGHT = 540 +_FAST_MIN_IMAGE_WIDTH = 608 +_FAST_OPT_IMAGE_HEIGHT = 1080 +_FAST_OPT_IMAGE_WIDTH = 1920 +_FAST_MAX_IMAGE_HEIGHT = 2160 +_FAST_MAX_IMAGE_WIDTH = 3840 +_ZERO_PAD_SELECTION = frozenset( + { + "mask_head.res_blocks.1.0.layers.5", + "mask_head.res_blocks.2.0.layers.2", + "mask_head.res_blocks.2.0.layers.5", + "mask_head.res_blocks.3.0.layers.2", + "mask_head.res_blocks.3.0.layers.5", + "mask_head.resamplers.1.1", + "mask_head.resamplers.2.1", + "neck.res_blocks.2.0.layers.2", + "neck.res_blocks.2.0.layers.5", + "neck.res_blocks.2.1.layers.2", + "neck.res_blocks.2.1.layers.5", + "neck.res_blocks.3.0.layers.2", + "neck.res_blocks.3.0.layers.5", + "neck.res_blocks.3.1.layers.2", + "neck.res_blocks.3.1.layers.5", + "points_head.res_blocks.1.0.layers.2", + "points_head.res_blocks.1.0.layers.5", + "points_head.res_blocks.2.0.layers.2", + "points_head.res_blocks.2.0.layers.5", + "points_head.res_blocks.3.0.layers.2", + "points_head.res_blocks.3.0.layers.5", + "points_head.resamplers.0.1", + "points_head.resamplers.1.1", + "points_head.resamplers.2.1", + } +) + + +def _fuse_half_pixel_x2_conv_weight(weight: np.ndarray) -> np.ndarray: + """Compose HALF_PIXEL bilinear x2 followed by a 3x3 cross-correlation.""" + + if weight.ndim != 4 or tuple(weight.shape[2:]) != (3, 3): + raise ValueError(f"MoGe fused resample requires OI33 weights, got {weight.shape}") + coefficients = np.asarray((0.25, 0.75, 0.75, 0.25), dtype=weight.dtype) + fused = np.zeros((weight.shape[1], weight.shape[0], 6, 6), dtype=weight.dtype) + transposed = weight.transpose(1, 0, 2, 3) + for resize_y, coefficient_y in enumerate(coefficients): + for resize_x, coefficient_x in enumerate(coefficients): + coefficient = coefficient_y * coefficient_x + for kernel_y in range(3): + for kernel_x in range(3): + fused[ + :, + :, + resize_y - kernel_y + 2, + resize_x - kernel_x + 2, + ] += coefficient * transposed[:, :, kernel_y, kernel_x] + return np.ascontiguousarray(fused) def _require_torch(): @@ -56,10 +113,13 @@ def _checkpoint_digest(path: Path) -> str: class _NativeMogeGraph: """Small family-local vocabulary for composing the exact MoGe graph.""" - def __init__(self, trt: Any, network: Any, state: dict[str, Any]) -> None: + def __init__( + self, trt: Any, network: Any, state: dict[str, Any], *, fast_path: bool = False + ) -> None: self.trt = trt self.network = network self.state = state + self.fast_path = fast_path # TensorRT may retain host weight views until serialization finishes. self._host_weights: list[np.ndarray] = [] @@ -180,6 +240,41 @@ def dynamic_slice( layer.mode = mode return layer.get_output(0) + def gather(self, tensor: Any, indices: Any, axis: int, name: str) -> Any: + return self._layer( + self.network.add_gather(tensor, indices, axis), "gather", name + ).get_output(0) + + def nearest_sample_indices(self, size: Any, name: str) -> Any: + positions = self.constant( + np.arange(_FOCAL_RECOVERY_SIZE, dtype=np.int64), + dtype=np.int64, + name=f"{name}.positions", + ) + scaled = self.binary(positions, size, self.trt.ElementWiseOperation.PROD, f"{name}.scaled") + divisor = self.shape_value(_FOCAL_RECOVERY_SIZE, f"{name}.divisor") + return self.binary( + scaled, + divisor, + self.trt.ElementWiseOperation.FLOOR_DIV, + f"{name}.indices", + ) + + def is_finite(self, tensor: Any, name: str) -> Any: + absolute = self.unary(tensor, self.trt.UnaryOperation.ABS, f"{name}.abs") + if tensor.dtype == self.trt.float16: + constant_dtype = np.float16 + elif tensor.dtype == self.trt.float32: + constant_dtype = np.float32 + else: + raise ValueError(f"MoGe finite check {name!r} requires FP16 or FP32 input") + infinity = self.constant( + np.full((1,) * len(tuple(tensor.shape)), np.inf, dtype=constant_dtype), + dtype=constant_dtype, + name=f"{name}.infinity", + ) + return self.binary(absolute, infinity, self.trt.ElementWiseOperation.LESS, name) + def resize( self, tensor: Any, @@ -236,12 +331,33 @@ def convolution( *, stride: int = 1, replicate_padding: int = 0, + compute_dtype: Any | None = None, ) -> Any: weight = self._array(f"{module}.weight") if weight.ndim != 4: raise ValueError(f"MoGe convolution {module!r} does not have a 4D kernel") bias = self._array(f"{module}.bias", (weight.shape[0],)) - padded = self.replicate_pad(tensor, replicate_padding, f"{name}.pad") + compute_dtype = compute_dtype or self.trt.float32 + if compute_dtype == self.trt.float16: + weight = np.ascontiguousarray(weight, dtype=np.float16) + bias = np.ascontiguousarray(bias, dtype=np.float16) + self._host_weights.extend((weight, bias)) + elif compute_dtype != self.trt.float32: + raise ValueError(f"Unsupported MoGe convolution compute dtype: {compute_dtype}") + tensor = self.cast(tensor, compute_dtype, f"{name}.input_cast") + zero_pad = module in _ZERO_PAD_SELECTION + if zero_pad and ( + replicate_padding != 1 + or stride != 1 + or tuple(int(value) for value in weight.shape[2:]) != (3, 3) + ): + raise ValueError( + f"MoGe selected zero-pad convolution {module!r} must be stride-1 3x3 " + "with one pixel of source replicate padding" + ) + padded = ( + tensor if zero_pad else self.replicate_pad(tensor, replicate_padding, f"{name}.pad") + ) layer = self._layer( self.network.add_convolution_nd( padded, @@ -254,15 +370,30 @@ def convolution( name, ) layer.stride_nd = (stride, stride) - layer.padding_nd = (0, 0) + layer.padding_nd = (1, 1) if zero_pad else (0, 0) return layer.get_output(0) - def deconvolution(self, tensor: Any, module: str, name: str) -> Any: + def deconvolution( + self, + tensor: Any, + module: str, + name: str, + *, + compute_dtype: Any | None = None, + ) -> Any: weight = self._array(f"{module}.weight") if weight.ndim != 4: raise ValueError(f"MoGe deconvolution {module!r} does not have a 4D kernel") output_channels = int(weight.shape[1]) bias = self._array(f"{module}.bias", (output_channels,)) + compute_dtype = compute_dtype or self.trt.float32 + if compute_dtype == self.trt.float16: + weight = np.ascontiguousarray(weight, dtype=np.float16) + bias = np.ascontiguousarray(bias, dtype=np.float16) + self._host_weights.extend((weight, bias)) + elif compute_dtype != self.trt.float32: + raise ValueError(f"Unsupported MoGe deconvolution compute dtype: {compute_dtype}") + tensor = self.cast(tensor, compute_dtype, f"{name}.input_cast") layer = self._layer( self.network.add_deconvolution_nd( tensor, @@ -278,15 +409,92 @@ def deconvolution(self, tensor: Any, module: str, name: str) -> Any: layer.padding_nd = (0, 0) return layer.get_output(0) - def linear(self, tensor: Any, module: str, name: str) -> Any: + def fused_half_pixel_resample( + self, + tensor: Any, + module: str, + name: str, + *, + compute_dtype: Any, + ) -> Any: + """Fuse bilinear x2, replicate padding and a 3x3 convolution exactly.""" + + weight = self._array(f"{module}.weight") + bias = self._array(f"{module}.bias", (int(weight.shape[0]),)) + if compute_dtype == self.trt.float16: + # Match the source convolution's effective checkpoint precision + # before composing its weights with the exact bilinear kernel. + weight = np.ascontiguousarray(weight, dtype=np.float16) + bias = np.ascontiguousarray(bias, dtype=np.float16) + fused_weight = np.ascontiguousarray( + _fuse_half_pixel_x2_conv_weight(weight.astype(np.float32)), + dtype=np.float16, + ) + elif compute_dtype == self.trt.float32: + fused_weight = _fuse_half_pixel_x2_conv_weight(weight) + else: + raise ValueError(f"Unsupported MoGe fused resample dtype: {compute_dtype}") + self._host_weights.extend((weight, bias, fused_weight)) + + tensor = self.cast(tensor, compute_dtype, f"{name}.input_cast") + # Replicating one low-resolution pixel supplies the two high-resolution + # border samples consumed by the original post-resize replicate pad. + tensor = self.replicate_pad(tensor, 1, f"{name}.input_pad") + layer = self._layer( + self.network.add_deconvolution_nd( + tensor, + int(weight.shape[0]), + (6, 6), + self.trt.Weights(fused_weight), + self.trt.Weights(bias), + ), + "fused resample deconvolution", + name, + ) + layer.stride_nd = (2, 2) + # For padded input H+2, (H+2-1)*2 + 6 - 2*4 == 2H. + layer.padding_nd = (4, 4) + return layer.get_output(0) + + def linear( + self, + tensor: Any, + module: str, + name: str, + *, + compute_dtype: Any | None = None, + output_dtype: Any | None = None, + ) -> Any: weight = self._array(f"{module}.weight") if weight.ndim != 2: raise ValueError(f"MoGe linear {module!r} does not have a 2D weight") output_width, input_width = (int(value) for value in weight.shape) bias = self._array(f"{module}.bias", (output_width,)) + compute_dtype = compute_dtype or self.trt.float32 + if compute_dtype == self.trt.float16: + constant_dtype = np.float16 + weight = np.ascontiguousarray(weight, dtype=np.float16) + bias = np.ascontiguousarray(bias, dtype=np.float16) + self._host_weights.extend((weight, bias)) + elif compute_dtype != self.trt.float32: + raise ValueError(f"Unsupported MoGe linear compute dtype: {compute_dtype}") + else: + constant_dtype = np.float32 + tensor = self.cast(tensor, compute_dtype, f"{name}.input_cast") rank = len(tuple(tensor.shape)) - matrix_shape = (1,) * max(0, rank - 2) + (output_width, input_width) - rhs = self.constant(weight.reshape(matrix_shape), name=f"{name}.weight") + restore_shape = None + if rank > 2: + input_shape = self.shape(tensor, f"{name}.input_shape") + leading = [ + self.shape_index(input_shape, index, f"{name}.output_dim_{index}") + for index in range(rank - 1) + ] + restore_shape = self.shape_concat( + [*leading, self.shape_value(output_width, f"{name}.output_width")], + f"{name}.output_shape", + ) + tensor = self.reshape(tensor, (-1, input_width), f"{name}.input_rows") + rhs = self.constant(weight, dtype=constant_dtype, name=f"{name}.weight") product = self._layer( self.network.add_matrix_multiply( tensor, @@ -297,13 +505,21 @@ def linear(self, tensor: Any, module: str, name: str) -> Any: "matrix multiply", f"{name}.matmul", ).get_output(0) - bias_shape = (1,) * (rank - 1) + (output_width,) - bias_tensor = self.constant(bias.reshape(bias_shape), name=f"{name}.bias") - return self.binary( + bias_tensor = self.constant( + bias.reshape(1, output_width), dtype=constant_dtype, name=f"{name}.bias" + ) + result = self.binary( product, bias_tensor, self.trt.ElementWiseOperation.SUM, f"{name}.bias_add" ) + if restore_shape is not None: + result = self.reshape(result, restore_shape, f"{name}.restore") + if output_dtype is not None: + result = self.cast(result, output_dtype, f"{name}.output_cast") + return result def layer_norm(self, tensor: Any, module: str, name: str) -> Any: + if self.fast_path: + tensor = self.cast(tensor, self.trt.float32, f"{name}.input_fp32") rank = len(tuple(tensor.shape)) width = int(self.state[f"{module}.weight"].numel()) parameter_shape = (1,) * (rank - 1) + (width,) @@ -438,7 +654,16 @@ def uv(self, height: Any, width: Any, aspect: Any, name: str) -> Any: def attention(self, hidden: Any, layer_index: int, total_tokens: Any) -> Any: prefix = f"encoder.backbone.blocks.{layer_index}.attn" name = f"vit.block.{layer_index}.attention" - qkv = self.linear(hidden, f"{prefix}.qkv", f"{name}.qkv") + if self.fast_path: + qkv = self.linear( + hidden, + f"{prefix}.qkv", + f"{name}.qkv", + compute_dtype=self.trt.float16, + output_dtype=self.trt.float32, + ) + else: + qkv = self.linear(hidden, f"{prefix}.qkv", f"{name}.qkv") component_shape = self.shape_concat( [ self.shape_value(1, f"{name}.batch"), @@ -476,6 +701,10 @@ def attention(self, hidden: Any, layer_index: int, total_tokens: Any) -> Any: ] scale = self.constant([[[[0.125]]]], name=f"{name}.scale") q = self.binary(q, scale, self.trt.ElementWiseOperation.PROD, f"{name}.q_scaled") + if self.fast_path: + q = self.cast(q, self.trt.float16, f"{name}.q_fp16") + k = self.cast(k, self.trt.float16, f"{name}.k_fp16") + v = self.cast(v, self.trt.float16, f"{name}.v_fp16") add_attention_v2 = getattr(self.network, "add_attention_v2", None) if callable(add_attention_v2): layer = add_attention_v2( @@ -490,7 +719,7 @@ def attention(self, hidden: Any, layer_index: int, total_tokens: Any) -> Any: q, k, v, self.trt.AttentionNormalizationOp.SOFTMAX, False ) attention = self._layer(layer, "IAttention", name) - attention.decomposable = True + attention.decomposable = not self.fast_path if hasattr(attention, "query_form"): attention.query_form = self.trt.AttentionIOForm.PADDED_BHND attention.key_value_form = self.trt.AttentionIOForm.PADDED_BHND @@ -508,6 +737,14 @@ def attention(self, hidden: Any, layer_index: int, total_tokens: Any) -> Any: f"{name}.context", first_transpose=(0, 2, 1, 3), ) + if self.fast_path: + return self.linear( + context, + f"{prefix}.proj", + f"{name}.projection", + compute_dtype=self.trt.float16, + output_dtype=self.trt.float16, + ) return self.linear(context, f"{prefix}.proj", f"{name}.projection") def transformer_block(self, hidden: Any, index: int, total_tokens: Any) -> Any: @@ -521,6 +758,8 @@ def transformer_block(self, hidden: Any, index: int, total_tokens: Any) -> Any: shape=(1, 1, _HIDDEN), name=f"{name}.ls1", ) + if self.fast_path: + gamma1 = self.cast(gamma1, self.trt.float16, f"{name}.ls1_fp16") attention = self.binary( attention, gamma1, self.trt.ElementWiseOperation.PROD, f"{name}.scaled_attention" ) @@ -528,15 +767,34 @@ def transformer_block(self, hidden: Any, index: int, total_tokens: Any) -> Any: hidden, attention, self.trt.ElementWiseOperation.SUM, f"{name}.attention_residual" ) normalized = self.layer_norm(hidden, f"{prefix}.norm2", f"{name}.norm2") - mlp = self.linear(normalized, f"{prefix}.mlp.fc1", f"{name}.mlp.fc1") + if self.fast_path: + mlp = self.linear( + normalized, + f"{prefix}.mlp.fc1", + f"{name}.mlp.fc1", + compute_dtype=self.trt.float16, + ) + else: + mlp = self.linear(normalized, f"{prefix}.mlp.fc1", f"{name}.mlp.fc1") mlp = self.gelu(mlp, f"{name}.mlp.gelu") - mlp = self.linear(mlp, f"{prefix}.mlp.fc2", f"{name}.mlp.fc2") + if self.fast_path: + mlp = self.linear( + mlp, + f"{prefix}.mlp.fc2", + f"{name}.mlp.fc2", + compute_dtype=self.trt.float16, + output_dtype=self.trt.float16, + ) + else: + mlp = self.linear(mlp, f"{prefix}.mlp.fc2", f"{name}.mlp.fc2") gamma2 = self.weight_constant( f"{prefix}.ls2.gamma", expected=(_HIDDEN,), shape=(1, 1, _HIDDEN), name=f"{name}.ls2", ) + if self.fast_path: + gamma2 = self.cast(gamma2, self.trt.float16, f"{name}.ls2_fp16") mlp = self.binary(mlp, gamma2, self.trt.ElementWiseOperation.PROD, f"{name}.scaled_mlp") return self.binary(hidden, mlp, self.trt.ElementWiseOperation.SUM, f"{name}.mlp_residual") @@ -582,6 +840,7 @@ def projected_intermediate( image, f"encoder.output_projections.{projection_index}", f"{name}.projection", + compute_dtype=self.trt.float16 if self.fast_path else self.trt.float32, ) return projected, class_token @@ -655,7 +914,11 @@ def encoder(self, image: Any) -> tuple[Any, Any, Any, Any, Any]: pixels = self.binary(pixels, mean, self.trt.ElementWiseOperation.SUB, "input.center") pixels = self.binary(pixels, std, self.trt.ElementWiseOperation.DIV, "input.normalize") patches = self.convolution( - pixels, "encoder.backbone.patch_embed.proj", "vit.patch_embed", stride=_PATCH + pixels, + "encoder.backbone.patch_embed.proj", + "vit.patch_embed", + stride=_PATCH, + compute_dtype=self.trt.float16 if self.fast_path else self.trt.float32, ) patch_tokens = self.binary( base_h, base_w, self.trt.ElementWiseOperation.PROD, "vit.patch_tokens" @@ -685,6 +948,8 @@ def encoder(self, image: Any) -> tuple[Any, Any, Any, Any, Any]: expected=(1, 1, _HIDDEN), name="vit.class_token", ) + if self.fast_path: + class_token = self.cast(class_token, self.trt.float16, "vit.class_token_fp16") token_concat = self._layer( self.network.add_concatenation([class_token, hidden]), "token concat", "vit.tokens" ) @@ -726,12 +991,17 @@ def encoder(self, image: Any) -> tuple[Any, Any, Any, Any, Any]: "vit.position.tokens", ) position_concat.axis = 1 + position_tokens = position_concat.get_output(0) + if self.fast_path: + position_tokens = self.cast(position_tokens, self.trt.float16, "vit.position_fp16") hidden = self.binary( hidden, - position_concat.get_output(0), + position_tokens, self.trt.ElementWiseOperation.SUM, "vit.tokens_plus_position", ) + if self.fast_path: + hidden = self.cast(hidden, self.trt.float16, "vit.residual_fp16") captured: list[Any] = [] last_class = None @@ -755,20 +1025,44 @@ def encoder(self, image: Any) -> tuple[Any, Any, Any, Any, Any]: class_vector = self.reshape(last_class, (1, _HIDDEN), "vit.class_vector") return encoded, class_vector, base_h, base_w, aspect - def residual_conv_block(self, tensor: Any, module: str, name: str) -> Any: + def residual_conv_block( + self, tensor: Any, module: str, name: str, *, compute_dtype: Any + ) -> Any: hidden = self.relu(tensor, f"{name}.relu1") hidden = self.convolution( - hidden, f"{module}.layers.2", f"{name}.conv1", replicate_padding=1 + hidden, + f"{module}.layers.2", + f"{name}.conv1", + replicate_padding=1, + compute_dtype=compute_dtype, ) hidden = self.relu(hidden, f"{name}.relu2") hidden = self.convolution( - hidden, f"{module}.layers.5", f"{name}.conv2", replicate_padding=1 + hidden, + f"{module}.layers.5", + f"{name}.conv2", + replicate_padding=1, + compute_dtype=compute_dtype, ) return self.binary(hidden, tensor, self.trt.ElementWiseOperation.SUM, f"{name}.residual") - def resample(self, tensor: Any, module: str, level: int, name: str) -> Any: + def resample( + self, tensor: Any, module: str, level: int, name: str, *, compute_dtype: Any + ) -> Any: if level < 3: - tensor = self.deconvolution(tensor, f"{module}.0", f"{name}.deconvolution") + tensor = self.deconvolution( + tensor, + f"{module}.0", + f"{name}.deconvolution", + compute_dtype=compute_dtype, + ) + elif level == 3 and self.fast_path: + return self.fused_half_pixel_resample( + tensor, + f"{module}.1", + f"{name}.fused_deconvolution", + compute_dtype=compute_dtype, + ) else: shape = self.shape(tensor, f"{name}.input_shape") height = self.shape_index(shape, 2, f"{name}.height") @@ -779,7 +1073,13 @@ def resample(self, tensor: Any, module: str, level: int, name: str) -> Any: tensor = self.resize_nchw_to_hw( tensor, output_h, output_w, self.trt.InterpolationMode.LINEAR, f"{name}.resize" ) - return self.convolution(tensor, f"{module}.1", f"{name}.convolution", replicate_padding=1) + return self.convolution( + tensor, + f"{module}.1", + f"{name}.convolution", + replicate_padding=1, + compute_dtype=compute_dtype, + ) def conv_stack( self, @@ -788,12 +1088,16 @@ def conv_stack( num_res_blocks: tuple[int, ...], *, final_projection: bool, + compute_dtype: Any, ) -> list[Any]: outputs: list[Any] = [] current = None for level, feature in enumerate(inputs): projected = self.convolution( - feature, f"{prefix}.input_blocks.{level}", f"{prefix}.level.{level}.input" + feature, + f"{prefix}.input_blocks.{level}", + f"{prefix}.level.{level}.input", + compute_dtype=compute_dtype, ) current = ( projected @@ -810,6 +1114,7 @@ def conv_stack( current, f"{prefix}.res_blocks.{level}.{block}", f"{prefix}.level.{level}.block.{block}", + compute_dtype=compute_dtype, ) output = current if final_projection and level == len(inputs) - 1: @@ -817,6 +1122,7 @@ def conv_stack( current, f"{prefix}.output_blocks.{level}", f"{prefix}.level.{level}.output", + compute_dtype=compute_dtype, ) outputs.append(output) if level < len(inputs) - 1: @@ -825,6 +1131,7 @@ def conv_stack( f"{prefix}.resamplers.{level}", level, f"{prefix}.level.{level}.resample", + compute_dtype=compute_dtype, ) return outputs @@ -840,6 +1147,8 @@ def outputs(self, image: Any) -> tuple[Any, Any, Any]: base_w, multiplier, self.trt.ElementWiseOperation.PROD, f"uv.level.{level}.width" ) uv = self.uv(height, width, aspect, f"uv.level.{level}") + if self.fast_path: + uv = self.cast(uv, self.trt.float16, f"uv.level.{level}.fp16") if level == 0: concat = self._layer( self.network.add_concatenation([features[0], uv]), @@ -851,14 +1160,27 @@ def outputs(self, image: Any) -> tuple[Any, Any, Any]: else: features.append(uv) + decoder_dtype = self.trt.float16 if self.fast_path else self.trt.float32 neck = self.conv_stack( - features, "neck", (0, 2, 2, 2, 0), final_projection=False + features, + "neck", + (0, 2, 2, 2, 0), + final_projection=False, + compute_dtype=decoder_dtype, ) points = self.conv_stack( - neck, "points_head", (0, 1, 1, 1, 0), final_projection=True + neck, + "points_head", + (0, 1, 1, 1, 0), + final_projection=True, + compute_dtype=decoder_dtype, )[-1] mask = self.conv_stack( - neck, "mask_head", (0, 1, 1, 1, 0), final_projection=True + neck, + "mask_head", + (0, 1, 1, 1, 0), + final_projection=True, + compute_dtype=decoder_dtype, )[-1] scale = self.linear(class_vector, "scale_head.0", "scale_head.0") scale = self.relu(scale, "scale_head.1") @@ -869,71 +1191,120 @@ def outputs(self, image: Any) -> tuple[Any, Any, Any]: input_shape = self.shape(image, "output.input_shape") input_h = self.shape_index(input_shape, 2, "output.height") input_w = self.shape_index(input_shape, 3, "output.width") - points = self.resize_nchw_to_hw( + raw_points = self.resize_nchw_to_hw( points, input_h, input_w, self.trt.InterpolationMode.LINEAR, "output.points_resize" ) mask = self.resize_nchw_to_hw( mask, input_h, input_w, self.trt.InterpolationMode.LINEAR, "output.mask_resize" ) - points_shape = self.shape_concat( - [ - self.shape_value(1, "output.points_batch"), - input_h, - input_w, - self.shape_value(3, "output.points_channels"), - ], - "output.points_shape", - ) - points = self.reshape( - points, - points_shape, - "output.points_nhwc", - first_transpose=(0, 2, 3, 1), - ) xy_shape = self.shape_concat( [ self.shape_value(1, "output.xy_batch"), + self.shape_value(2, "output.xy_channels"), input_h, input_w, - self.shape_value(2, "output.xy_channels"), ], "output.xy_shape", ) z_shape = self.shape_concat( [ self.shape_value(1, "output.z_batch"), + self.shape_value(1, "output.z_channels"), input_h, input_w, - self.shape_value(1, "output.z_channels"), ], "output.z_shape", ) - xy = self.dynamic_slice(points, (0, 0, 0, 0), xy_shape, "output.xy") - z = self.dynamic_slice(points, (0, 0, 0, 2), z_shape, "output.z") - z = self.unary(z, self.trt.UnaryOperation.EXP, "output.z_exp") - xy = self.binary(xy, z, self.trt.ElementWiseOperation.PROD, "output.xy_scaled") - point_concat = self._layer( - self.network.add_concatenation([xy, z]), "point concat", "output.points_remap" - ) - point_concat.axis = 3 - points = point_concat.get_output(0) + raw_xy = self.dynamic_slice(raw_points, (0, 0, 0, 0), xy_shape, "output.raw_xy") + raw_z = self.dynamic_slice(raw_points, (0, 2, 0, 0), z_shape, "output.raw_z") + z = self.unary(raw_z, self.trt.UnaryOperation.EXP, "output.z_exp") + xy = self.binary(raw_xy, z, self.trt.ElementWiseOperation.PROD, "output.xy_scaled") mask_shape = self.shape_concat( [self.shape_value(1, "output.mask_batch"), input_h, input_w], "output.mask_shape", ) + affine_depth = self.reshape(z, mask_shape, "output.affine_depth_squeeze") + affine_depth = self.cast(affine_depth, self.trt.float32, "output.affine_depth_fp32") + + row_indices = self.nearest_sample_indices(input_h, "output.focal_rows") + column_indices = self.nearest_sample_indices(input_w, "output.focal_columns") + sampled_xy = self.gather(xy, row_indices, 2, "output.focal_xy_rows") + sampled_xy = self.gather(sampled_xy, column_indices, 3, "output.focal_xy_columns") + sampled_z = self.gather(z, row_indices, 2, "output.focal_z_rows") + sampled_z = self.gather(sampled_z, column_indices, 3, "output.focal_z_columns") + sampled_concat = self._layer( + self.network.add_concatenation([sampled_xy, sampled_z]), + "sampled point concat", + "output.focal_samples_nchw", + ) + sampled_concat.axis = 1 + focal_samples = self.reshape( + sampled_concat.get_output(0), + (1, _FOCAL_RECOVERY_SIZE, _FOCAL_RECOVERY_SIZE, 3), + "output.focal_samples_nhwc", + first_transpose=(0, 2, 3, 1), + ) + focal_samples = self.cast(focal_samples, self.trt.float32, "output.focal_samples_fp32") + + x = self.dynamic_slice(xy, (0, 0, 0, 0), z_shape, "output.valid.x") + y = self.dynamic_slice(xy, (0, 1, 0, 0), z_shape, "output.valid.y") + x_finite = self.is_finite(x, "output.valid.x_finite") + y_finite = self.is_finite(y, "output.valid.y_finite") + z_finite = self.is_finite(z, "output.valid.z_finite") + points_finite = self.binary( + x_finite, + y_finite, + self.trt.ElementWiseOperation.AND, + "output.valid.xy_finite", + ) + points_finite = self.binary( + points_finite, + z_finite, + self.trt.ElementWiseOperation.AND, + "output.valid.xyz_finite", + ) + points_finite = self.reshape(points_finite, mask_shape, "output.valid.points_squeeze") + + # Keep the legacy sigmoid and FP16->FP32 boundary. For a tiny positive + # FP16 logit, sigmoid can round to exactly 0.5, so logit > 0 is not an + # exact replacement for the public mask predicate. mask = self.reshape(mask, mask_shape, "output.mask_squeeze") mask = self.sigmoid(mask, "output.mask_sigmoid") + mask = self.cast(mask, self.trt.float32, "output.mask_fp32") + mask_threshold = self.constant( + [[[0.5]]], dtype=np.float32, name="output.valid.mask_threshold" + ) + mask_selected = self.binary( + mask, + mask_threshold, + self.trt.ElementWiseOperation.GREATER, + "output.valid.mask_selected", + ) + # GREATER is ordered: NaN compares false. Sigmoid maps finite values + # and +/-infinity to finite probabilities, so the legacy isfinite(mask) + # term cannot reject anything that this comparison would select. + valid = self.binary( + points_finite, + mask_selected, + self.trt.ElementWiseOperation.AND, + "output.valid.selected", + ) + valid = self.cast(valid, self.trt.float16, "output.valid_fp16") + scale = self.unary(scale, self.trt.UnaryOperation.EXP, "output.metric_scale_exp") scale = self.reshape(scale, (1,), "output.metric_scale_squeeze") - return points, mask, scale + scale = self.cast(scale, self.trt.float32, "output.metric_scale_fp32") + return affine_depth, valid, focal_samples, scale def _build_native_engine( state: dict[str, Any], *, + precision: str, verbose: bool, ) -> bytes: + fast_path = precision == "fp16" trt = trt_compat.get_trt() logger = trt.Logger(trt.Logger.INFO if verbose else trt.Logger.WARNING) builder = trt.Builder(logger) @@ -942,10 +1313,11 @@ def _build_native_engine( ) if network is None: raise RuntimeError("TensorRT failed to create the MoGe network") - image = network.add_input("image", trt.float32, (1, -1, -1, 3)) + input_dims = (1, -1, -1, 3) + image = network.add_input("image", trt.float32, input_dims) if image is None: raise RuntimeError("TensorRT rejected the MoGe image input") - graph = _NativeMogeGraph(trt, network, state) + graph = _NativeMogeGraph(trt, network, state, fast_path=fast_path) input_shape = graph.shape(image, "input_hwc.shape") input_h = graph.shape_index(input_shape, 1, "input_hwc.height") input_w = graph.shape_index(input_shape, 2, "input_hwc.width") @@ -964,47 +1336,62 @@ def _build_native_engine( "input_hwc.to_nchw", first_transpose=(0, 3, 1, 2), ) - points, mask, metric_scale = graph.outputs(image) + affine_depth, valid, focal_samples, metric_scale = graph.outputs(image) for name, tensor in ( - ("points", points), - ("mask", mask), + ("affine_depth", affine_depth), + ("valid", valid), + ("focal_samples", focal_samples), ("metric_scale", metric_scale), ): - if tensor.dtype != trt.float32: - tensor = graph.cast(tensor, trt.float32, f"output.{name}_fp32") tensor.name = name network.mark_output(tensor) config = builder.create_builder_config() tf32 = getattr(trt.BuilderFlag, "TF32", None) if tf32 is not None: - config.clear_flag(tf32) + if fast_path: + config.set_flag(tf32) + else: + config.clear_flag(tf32) config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 16 << 30) profile = builder.create_optimization_profile() - profile.set_shape( - "image", - (1, _MIN_IMAGE_SIZE, _MIN_IMAGE_SIZE, 3), - (1, _OPT_IMAGE_SIZE, _OPT_IMAGE_SIZE, 3), - (1, _MAX_IMAGE_SIZE, _MAX_IMAGE_SIZE, 3), - ) + if fast_path: + profile.set_shape( + "image", + (1, _FAST_MIN_IMAGE_HEIGHT, _FAST_MIN_IMAGE_WIDTH, 3), + (1, _FAST_OPT_IMAGE_HEIGHT, _FAST_OPT_IMAGE_WIDTH, 3), + (1, _FAST_MAX_IMAGE_HEIGHT, _FAST_MAX_IMAGE_WIDTH, 3), + ) + else: + profile.set_shape( + "image", + (1, _MIN_IMAGE_SIZE, _MIN_IMAGE_SIZE, 3), + (1, _OPT_IMAGE_SIZE, _OPT_IMAGE_SIZE, 3), + (1, _MAX_IMAGE_SIZE, _MAX_IMAGE_SIZE, 3), + ) if not profile: raise RuntimeError("Failed to configure the MoGe dynamic image profile") config.add_optimization_profile(profile) if hasattr(config, "builder_optimization_level"): - # TRT 11.2 optimization level 3 tries to absorb the dynamic FP32 - # decomposable-attention chain into one Myelin ForeignNode, whose - # dynamic BMM fallback has no implementation. Level 0 preserves the - # native IAttention decomposition and builds the full 64..2048 profile. - config.builder_optimization_level = 0 + # Level 3 enables the FP16 fused-attention fast path. Level 0 keeps the + # broad dynamic FP32 attention graph decomposed for reliable builds. + config.builder_optimization_level = 3 if fast_path else 0 if hasattr(config, "avg_timing_iterations"): - config.avg_timing_iterations = 1 + config.avg_timing_iterations = 3 if hasattr(config, "max_aux_streams"): config.max_aux_streams = 0 if verbose: + profile_label = ( + f"{_FAST_MIN_IMAGE_WIDTH}x{_FAST_MIN_IMAGE_HEIGHT}.." + f"{_FAST_MAX_IMAGE_WIDTH}x{_FAST_MAX_IMAGE_HEIGHT}" + f"@{_FAST_OPT_IMAGE_WIDTH}x{_FAST_OPT_IMAGE_HEIGHT}" + if fast_path + else f"{_MIN_IMAGE_SIZE}..{_MAX_IMAGE_SIZE}" + ) print( "[trtmc build] Building native MoGe TensorRT graph " f"({network.num_layers} layers, num_tokens={_NUM_TOKENS}, " - f"profile={_MIN_IMAGE_SIZE}..{_MAX_IMAGE_SIZE}) ...", + f"precision={precision}, profile={profile_label}) ...", file=sys.stderr, ) plan = builder.build_serialized_network(network, config) @@ -1025,10 +1412,8 @@ def build_moge_engine( checkpoint_path = model_root / _CHECKPOINT if not checkpoint_path.is_file(): raise FileNotFoundError(f"MoGe checkpoint not found: {checkpoint_path}") - if precision != "fp32": - raise ValueError( - "The native MoGe-2 ViT-L accuracy contract supports precision='fp32' only" - ) + if precision not in {"fp32", "fp16"}: + raise ValueError("The native MoGe-2 ViT-L builder supports precision='fp32' or 'fp16' only") checkpoint_sha256 = _checkpoint_digest(checkpoint_path) if checkpoint_sha256 != _CHECKPOINT_SHA256: raise ValueError( @@ -1045,5 +1430,6 @@ def build_moge_engine( ) return _build_native_engine( checkpoint["model"], + precision=precision, verbose=verbose, ) diff --git a/python/tensorrt_model_connect/families/moge/tests/test_family.py b/python/tensorrt_model_connect/families/moge/tests/test_family.py index ae3207f05b..f8450be7a7 100644 --- a/python/tensorrt_model_connect/families/moge/tests/test_family.py +++ b/python/tensorrt_model_connect/families/moge/tests/test_family.py @@ -11,6 +11,7 @@ from types import SimpleNamespace import tomllib +import numpy as np import pytest from tensorrt_model_connect import engine_builder @@ -49,12 +50,13 @@ def test_config_adapter_claims_one_flat_checkpoint(tmp_path: Path) -> None: assert not family_plugin.plugin.matches("MoGeModel") -def test_plugin_keeps_model_state_local_and_rejects_unimplemented_modes( +def test_plugin_keeps_model_state_local_and_rejects_unsupported_quantization( tmp_path: Path, ) -> None: (tmp_path / "model.pt").write_bytes(b"checkpoint") config = SimpleNamespace(raw={}) + assert family_plugin.plugin.default_build_precision == "fp32" assert family_plugin.plugin.load_weights(str(tmp_path), config) == { "model_dir": str(tmp_path.resolve()) } @@ -107,6 +109,10 @@ def test_production_builder_is_fixed_and_tensor_rt_native_only() -> None: "add_plugin", "get_plugin_registry", "trtmc_moge_", + "add_quantize", + "add_dequantize", + "fp8_scale_map", + "_fp8_dense_selection", ): assert forbidden not in lowered for required in ( @@ -122,16 +128,139 @@ def test_production_builder_is_fixed_and_tensor_rt_native_only() -> None: "GELU_ERF", "first_transpose=(0, 3, 1, 2)", "_NUM_TOKENS = 1800", + "_FOCAL_RECOVERY_SIZE = 64", + "_FAST_MIN_IMAGE_HEIGHT = 540", + "_FAST_MIN_IMAGE_WIDTH = 608", + "_FAST_OPT_IMAGE_HEIGHT = 1080", + "_FAST_OPT_IMAGE_WIDTH = 1920", + "_FAST_MAX_IMAGE_HEIGHT = 2160", + "_FAST_MAX_IMAGE_WIDTH = 3840", + "attention.decomposable = not self.fast_path", + "compute_dtype=self.trt.float16", + "output_dtype=self.trt.float16", + 'tensor = self.cast(tensor, self.trt.float32, f"{name}.input_fp32")', + 'hidden = self.cast(hidden, self.trt.float16, "vit.residual_fp16")', + "compute_dtype=self.trt.float16 if self.fast_path else self.trt.float32", + "config.builder_optimization_level = 3 if fast_path else 0", + "config.avg_timing_iterations = 3", + "ElementWiseOperation.FLOOR_DIV", + "add_gather", + '"output.valid_fp16"', + "np.full((1,) * len(tuple(tensor.shape))", + "[[[0.5]]]", + '"output.raw_xy"', + '"output.raw_z"', + '"output.focal_samples_nchw"', + '"output.affine_depth_fp32"', + '"output.focal_samples_fp32"', + '"output.mask_sigmoid"', + '"output.mask_fp32"', ): assert required in source - for output in ("points", "mask", "metric_scale"): + assert '"output.points_nhwc"' not in source + assert '"output.points_remap"' not in source + assert '"output.valid.mask_finite"' not in source + assert '"output.valid_int8"' not in source + for output in ("affine_depth", "valid", "focal_samples", "metric_scale"): assert f'("{output}",' in source -def test_build_rejects_unqualified_precision_and_wrong_checkpoint(tmp_path: Path) -> None: +def test_focal_sample_index_contract_covers_observed_shapes() -> None: + observed_sizes = ( + (608, 1080), + (612, 1080), + (1066, 1920), + (1076, 1920), + (1078, 1920), + (1080, 1840), + (1080, 1904), + (1080, 1906), + (1080, 1912), + (1080, 1918), + (1080, 1920), + (1264, 1080), + (1428, 1080), + (1440, 1080), + (1674, 1080), + (1904, 1080), + (1906, 1080), + (1912, 1080), + (1918, 1074), + (1918, 1080), + (1920, 1076), + (1920, 1078), + (1920, 1080), + (2688, 1508), + (3840, 2156), + (3840, 2160), + ) + + assert len(observed_sizes) == 26 + sample_size = model_module._FOCAL_RECOVERY_SIZE + assert sample_size == 64 + for width, height in observed_sizes: + for size in (width, height): + indices = tuple(index * size // sample_size for index in range(sample_size)) + assert indices[0] == 0 + assert indices[-1] == (sample_size - 1) * size // sample_size + assert all(0 <= index < size for index in indices) + assert all(left <= right for left, right in zip(indices, indices[1:])) + + +def test_slim_graph_retains_legacy_mask_rounding_and_ieee_finite_edges() -> None: + tiny_positive = np.nextafter(np.float16(0.0), np.float16(1.0)) + logits = np.asarray([-np.inf, -0.0, tiny_positive, np.inf, np.nan], dtype=np.float16) + with np.errstate(over="ignore", invalid="ignore"): + probabilities = np.asarray( + 1.0 / (1.0 + np.exp(-logits.astype(np.float32))), dtype=np.float16 + ) + legacy_selected = np.isfinite(probabilities) & (probabilities > np.float16(0.5)) + ordered_probability_selected = probabilities > np.float16(0.5) + logit_selected = logits > np.float16(0.0) + + np.testing.assert_array_equal(legacy_selected, np.asarray([False, False, False, True, False])) + np.testing.assert_array_equal(ordered_probability_selected, legacy_selected) + np.testing.assert_array_equal(logit_selected, np.asarray([False, False, True, True, False])) + assert probabilities[2] == np.float16(0.5) + + values = np.asarray( + [-np.finfo(np.float16).max, np.finfo(np.float16).max, -np.inf, np.inf, np.nan], + dtype=np.float16, + ) + with np.errstate(invalid="ignore"): + ordered_finite = np.abs(values) < np.float16(np.inf) + np.testing.assert_array_equal(ordered_finite, np.isfinite(values)) + + +def test_fp16_valid_output_uses_exact_zero_and_one_bit_patterns() -> None: + valid = np.asarray([False, True], dtype=np.bool_).astype(np.float16) + np.testing.assert_array_equal(valid.view(np.uint16), np.asarray([0x0000, 0x3C00], np.uint16)) + + +def test_slim_sample_gather_preserves_the_legacy_fp16_cast_boundary() -> None: + height, width = 67, 83 + affine_nchw = np.arange(3 * height * width, dtype=np.float32).reshape(1, 3, height, width) + affine_nchw = np.asarray(affine_nchw / 257.0, dtype=np.float16) + rows = np.asarray([index * height // 64 for index in range(64)]) + columns = np.asarray([index * width // 64 for index in range(64)]) + + legacy_nhwc = np.transpose(affine_nchw, (0, 2, 3, 1)).astype(np.float32) + legacy_samples = legacy_nhwc[:, rows, :, :][:, :, columns, :] + sampled_nchw = affine_nchw[:, :, rows, :][:, :, :, columns] + slim_samples = np.transpose(sampled_nchw, (0, 2, 3, 1)).astype(np.float32) + np.testing.assert_array_equal(slim_samples, legacy_samples) + + legacy_depth = legacy_nhwc[..., 2] + slim_depth = affine_nchw[:, 2, :, :].astype(np.float32) + np.testing.assert_array_equal(slim_depth, legacy_depth) + + +def test_build_rejects_unknown_precision_and_wrong_checkpoint(tmp_path: Path) -> None: (tmp_path / "model.pt").write_bytes(b"wrong checkpoint") - with pytest.raises(ValueError, match="supports precision='fp32' only"): - model_module.build_moge_engine(str(tmp_path), precision="fp16") + with pytest.raises(ValueError, match="supports precision='fp32' or 'fp16' only"): + model_module.build_moge_engine(str(tmp_path), precision="bf16") with pytest.raises(ValueError, match="checkpoint SHA-256 mismatch"): model_module.build_moge_engine(str(tmp_path), precision="fp32") + with pytest.raises(ValueError, match="checkpoint SHA-256 mismatch"): + model_module.build_moge_engine(str(tmp_path), precision="fp16") diff --git a/python/tensorrt_model_connect/families/moge/tests/test_fused_resample.py b/python/tensorrt_model_connect/families/moge/tests/test_fused_resample.py new file mode 100644 index 0000000000..4a8dae3993 --- /dev/null +++ b/python/tensorrt_model_connect/families/moge/tests/test_fused_resample.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only algebra and graph contracts for the exact level-3 resample fusion.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +from tensorrt_model_connect.families.moge import model as model_module + + +def _half_pixel_resize_x2(tensor: np.ndarray) -> np.ndarray: + batch, channels, height, width = tensor.shape + output = np.empty((batch, channels, 2 * height, 2 * width), dtype=tensor.dtype) + for output_y in range(2 * height): + source_y = (output_y + 0.5) * 0.5 - 0.5 + lower_y_raw = int(np.floor(source_y)) + fraction_y = source_y - lower_y_raw + lower_y = min(height - 1, max(0, lower_y_raw)) + upper_y = min(height - 1, max(0, lower_y_raw + 1)) + for output_x in range(2 * width): + source_x = (output_x + 0.5) * 0.5 - 0.5 + lower_x_raw = int(np.floor(source_x)) + fraction_x = source_x - lower_x_raw + lower_x = min(width - 1, max(0, lower_x_raw)) + upper_x = min(width - 1, max(0, lower_x_raw + 1)) + output[:, :, output_y, output_x] = ( + tensor[:, :, lower_y, lower_x] * (1.0 - fraction_y) * (1.0 - fraction_x) + + tensor[:, :, lower_y, upper_x] * (1.0 - fraction_y) * fraction_x + + tensor[:, :, upper_y, lower_x] * fraction_y * (1.0 - fraction_x) + + tensor[:, :, upper_y, upper_x] * fraction_y * fraction_x + ) + return output + + +def _replicate_conv3x3(tensor: np.ndarray, weight: np.ndarray, bias: np.ndarray) -> np.ndarray: + padded = np.pad(tensor, ((0, 0), (0, 0), (1, 1), (1, 1)), mode="edge") + output = np.empty( + (tensor.shape[0], weight.shape[0], tensor.shape[2], tensor.shape[3]), + dtype=tensor.dtype, + ) + for output_y in range(tensor.shape[2]): + for output_x in range(tensor.shape[3]): + patch = padded[:, :, output_y : output_y + 3, output_x : output_x + 3] + output[:, :, output_y, output_x] = np.einsum("nchw,ochw->no", patch, weight) + bias + return output + + +def _deconvolution_stride2_padding4( + tensor: np.ndarray, weight: np.ndarray, bias: np.ndarray +) -> np.ndarray: + full_height = (tensor.shape[2] - 1) * 2 + 6 + full_width = (tensor.shape[3] - 1) * 2 + 6 + full = np.zeros( + (tensor.shape[0], weight.shape[1], full_height, full_width), + dtype=tensor.dtype, + ) + for input_y in range(tensor.shape[2]): + for input_x in range(tensor.shape[3]): + contribution = np.einsum("ni,iohw->nohw", tensor[:, :, input_y, input_x], weight) + full[ + :, + :, + 2 * input_y : 2 * input_y + 6, + 2 * input_x : 2 * input_x + 6, + ] += contribution + output_height = (tensor.shape[2] - 1) * 2 - 8 + 6 + output_width = (tensor.shape[3] - 1) * 2 - 8 + 6 + return full[:, :, 4 : 4 + output_height, 4 : 4 + output_width] + bias[None, :, None, None] + + +def _reference(tensor: np.ndarray, weight: np.ndarray, bias: np.ndarray) -> np.ndarray: + return _replicate_conv3x3(_half_pixel_resize_x2(tensor), weight, bias) + + +def _fused(tensor: np.ndarray, weight: np.ndarray, bias: np.ndarray) -> np.ndarray: + padded = np.pad(tensor, ((0, 0), (0, 0), (1, 1), (1, 1)), mode="edge") + fused_weight = model_module._fuse_half_pixel_x2_conv_weight(weight) + return _deconvolution_stride2_padding4(padded, fused_weight, bias) + + +@pytest.mark.parametrize( + ("dtype", "tolerance"), + ((np.float64, 1.0e-11), (np.float32, 5.0e-5)), +) +@pytest.mark.parametrize("height,width", ((1, 1), (1, 3), (2, 1), (2, 2), (3, 5), (7, 4))) +def test_fused_resample_matches_half_pixel_replicate_reference( + dtype: Any, tolerance: float, height: int, width: int +) -> None: + random = np.random.default_rng(1000 + 10 * height + width) + tensor = random.standard_normal((1, 3, height, width)).astype(dtype) + weight = random.standard_normal((2, 3, 3, 3)).astype(dtype) + bias = random.standard_normal((2,)).astype(dtype) + + reference = _reference(tensor, weight, bias) + fused = _fused(tensor, weight, bias) + + assert reference.shape == fused.shape == (1, 2, 2 * height, 2 * width) + np.testing.assert_allclose(fused, reference, rtol=0.0, atol=tolerance) + + +@pytest.mark.parametrize("height,width", ((1, 1), (2, 3), (5, 4))) +def test_fused_resample_preserves_corner_edge_and_center_impulses(height: int, width: int) -> None: + random = np.random.default_rng(2000 + 10 * height + width) + weight = random.standard_normal((2, 2, 3, 3)).astype(np.float64) + bias = np.zeros((2,), dtype=np.float64) + positions = { + (0, 0), + (0, width - 1), + (height - 1, 0), + (height - 1, width - 1), + (height // 2, width // 2), + } + for input_y, input_x in positions: + tensor = np.zeros((1, 2, height, width), dtype=np.float64) + tensor[0, 0, input_y, input_x] = 1.0 + tensor[0, 1, input_y, input_x] = -0.5 + np.testing.assert_allclose( + _fused(tensor, weight, bias), + _reference(tensor, weight, bias), + rtol=0.0, + atol=1.0e-11, + ) + + +@dataclass +class _FakeTensor: + dtype: Any + shape: tuple[int, ...] + + +class _FakeLayer: + def __init__(self, output: _FakeTensor) -> None: + self.output = output + self.name = "" + self.stride_nd: tuple[int, ...] | None = None + self.padding_nd: tuple[int, ...] | None = None + + def get_output(self, index: int) -> _FakeTensor: + assert index == 0 + return self.output + + +class _FakeNetwork: + def __init__(self, dtype: Any) -> None: + self.dtype = dtype + self.calls: list[dict[str, Any]] = [] + + def add_deconvolution_nd( + self, + tensor: _FakeTensor, + output_channels: int, + kernel_size: tuple[int, int], + weight: np.ndarray, + bias: np.ndarray, + ) -> _FakeLayer: + height = (tensor.shape[2] - 1) * 2 + kernel_size[0] - 8 + width = (tensor.shape[3] - 1) * 2 + kernel_size[1] - 8 + layer = _FakeLayer(_FakeTensor(self.dtype, (1, output_channels, height, width))) + self.calls.append( + { + "tensor": tensor, + "output_channels": output_channels, + "kernel_size": kernel_size, + "weight": weight, + "bias": bias, + "layer": layer, + } + ) + return layer + + +class _FakeTrt: + float16 = "float16" + float32 = "float32" + + @staticmethod + def Weights(array: np.ndarray) -> np.ndarray: + return array + + +@pytest.mark.parametrize("height,width", ((1, 1), (3, 5), (32, 57), (57, 32))) +def test_fused_resample_graph_contract_is_dynamic_and_native( + monkeypatch: pytest.MonkeyPatch, height: int, width: int +) -> None: + network = _FakeNetwork(_FakeTrt.float32) + graph = model_module._NativeMogeGraph(_FakeTrt, network, {}, fast_path=True) + source_weight = np.arange(2 * 3 * 3 * 3, dtype=np.float32).reshape(2, 3, 3, 3) + source_bias = np.asarray((0.25, -0.5), dtype=np.float32) + arrays = {"resampler.weight": source_weight, "resampler.bias": source_bias} + monkeypatch.setattr(graph, "_array", lambda name, expected=None: arrays[name]) + pad_calls: list[tuple[int, str]] = [] + + def fake_pad(tensor: _FakeTensor, padding: int, name: str) -> _FakeTensor: + pad_calls.append((padding, name)) + return _FakeTensor( + tensor.dtype, + (tensor.shape[0], tensor.shape[1], tensor.shape[2] + 2, tensor.shape[3] + 2), + ) + + monkeypatch.setattr(graph, "replicate_pad", fake_pad) + tensor = _FakeTensor(_FakeTrt.float32, (1, 3, height, width)) + + output = graph.fused_half_pixel_resample( + tensor, + "resampler", + "level3", + compute_dtype=_FakeTrt.float32, + ) + + assert output.shape == (1, 2, 2 * height, 2 * width) + assert pad_calls == [(1, "level3.input_pad")] + assert len(network.calls) == 1 + call = network.calls[0] + assert call["output_channels"] == 2 + assert call["kernel_size"] == (6, 6) + assert call["layer"].stride_nd == (2, 2) + assert call["layer"].padding_nd == (4, 4) + np.testing.assert_array_equal( + call["weight"], model_module._fuse_half_pixel_x2_conv_weight(source_weight) + ) + np.testing.assert_array_equal(call["bias"], source_bias) + + +def test_level3_fast_path_uses_fused_native_deconvolution_only() -> None: + source = Path(model_module.__file__).read_text(encoding="utf-8") + assert "elif level == 3 and self.fast_path:" in source + assert 'f"{name}.fused_deconvolution"' in source + assert "self.network.add_deconvolution_nd" in source + assert "layer.stride_nd = (2, 2)" in source + assert "layer.padding_nd = (4, 4)" in source + for stack in ('"neck"', '"points_head"', '"mask_head"'): + assert stack in source + assert "add_plugin" not in source.lower() diff --git a/python/tensorrt_model_connect/families/moge/tests/test_selective_zero_padding.py b/python/tensorrt_model_connect/families/moge/tests/test_selective_zero_padding.py new file mode 100644 index 0000000000..3ccf8b0a7f --- /dev/null +++ b/python/tensorrt_model_connect/families/moge/tests/test_selective_zero_padding.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only contracts for the family-owned selective zero-padding lane.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +from tensorrt_model_connect.families.moge import model as model_module + + +EXPECTED_ZERO_PAD_SELECTION = frozenset( + { + "mask_head.res_blocks.1.0.layers.5", + "mask_head.res_blocks.2.0.layers.2", + "mask_head.res_blocks.2.0.layers.5", + "mask_head.res_blocks.3.0.layers.2", + "mask_head.res_blocks.3.0.layers.5", + "mask_head.resamplers.1.1", + "mask_head.resamplers.2.1", + "neck.res_blocks.2.0.layers.2", + "neck.res_blocks.2.0.layers.5", + "neck.res_blocks.2.1.layers.2", + "neck.res_blocks.2.1.layers.5", + "neck.res_blocks.3.0.layers.2", + "neck.res_blocks.3.0.layers.5", + "neck.res_blocks.3.1.layers.2", + "neck.res_blocks.3.1.layers.5", + "points_head.res_blocks.1.0.layers.2", + "points_head.res_blocks.1.0.layers.5", + "points_head.res_blocks.2.0.layers.2", + "points_head.res_blocks.2.0.layers.5", + "points_head.res_blocks.3.0.layers.2", + "points_head.res_blocks.3.0.layers.5", + "points_head.resamplers.0.1", + "points_head.resamplers.1.1", + "points_head.resamplers.2.1", + } +) +FUSED_LEVEL3_RESAMPLERS = frozenset( + { + "neck.resamplers.3.1", + "points_head.resamplers.3.1", + "mask_head.resamplers.3.1", + } +) + + +def _decoder_replicate_modules() -> frozenset[str]: + modules = set() + for prefix, block_counts in ( + ("neck", (0, 2, 2, 2, 0)), + ("points_head", (0, 1, 1, 1, 0)), + ("mask_head", (0, 1, 1, 1, 0)), + ): + modules.update(f"{prefix}.resamplers.{level}.1" for level in range(4)) + for level, count in enumerate(block_counts): + for block in range(count): + modules.add(f"{prefix}.res_blocks.{level}.{block}.layers.2") + modules.add(f"{prefix}.res_blocks.{level}.{block}.layers.5") + return frozenset(modules) + + +@dataclass +class _FakeTensor: + dtype: Any + shape: tuple[int, ...] + label: str + + +class _FakeLayer: + def __init__(self, output: _FakeTensor) -> None: + self.output = output + self.name = "" + self.stride_nd = None + self.padding_nd = None + + def get_output(self, index: int) -> _FakeTensor: + assert index == 0 + return self.output + + +class _FakeNetwork: + def __init__(self) -> None: + self.calls = [] + + def add_convolution_nd( + self, + tensor: _FakeTensor, + output_channels: int, + kernel_shape: tuple[int, int], + weight: np.ndarray, + bias: np.ndarray, + ) -> _FakeLayer: + layer = _FakeLayer(_FakeTensor(tensor.dtype, (1, output_channels, 8, 8), "output")) + self.calls.append( + { + "tensor": tensor, + "kernel_shape": kernel_shape, + "weight": weight, + "bias": bias, + "layer": layer, + } + ) + return layer + + +class _FakeTrt: + float16 = "float16" + float32 = "float32" + + @staticmethod + def Weights(array: np.ndarray) -> np.ndarray: + return array + + +def test_zero_pad_selection_is_exact_unique_and_topology_owned() -> None: + replicate_modules = _decoder_replicate_modules() + + assert len(replicate_modules) == 36 + assert model_module._ZERO_PAD_SELECTION == EXPECTED_ZERO_PAD_SELECTION + assert len(model_module._ZERO_PAD_SELECTION) == 24 + assert model_module._ZERO_PAD_SELECTION < replicate_modules + assert model_module._ZERO_PAD_SELECTION.isdisjoint(FUSED_LEVEL3_RESAMPLERS) + assert FUSED_LEVEL3_RESAMPLERS < replicate_modules + + +def test_only_selected_modules_use_native_zero_padding( + monkeypatch: pytest.MonkeyPatch, +) -> None: + network = _FakeNetwork() + graph = model_module._NativeMogeGraph(_FakeTrt, network, {}, fast_path=True) + weight = np.zeros((4, 4, 3, 3), dtype=np.float32) + bias = np.zeros((4,), dtype=np.float32) + monkeypatch.setattr( + graph, + "_array", + lambda name, expected=None: weight if name.endswith(".weight") else bias, + ) + monkeypatch.setattr(graph, "cast", lambda tensor, dtype, name: tensor) + pad_calls = [] + + def fake_replicate_pad(tensor: _FakeTensor, padding: int, name: str) -> _FakeTensor: + pad_calls.append((padding, name)) + return _FakeTensor(tensor.dtype, (1, 4, 10, 10), "replicate-padded") + + monkeypatch.setattr(graph, "replicate_pad", fake_replicate_pad) + + for module in sorted(_decoder_replicate_modules()): + pad_calls.clear() + network.calls.clear() + tensor = _FakeTensor(_FakeTrt.float32, (1, 4, 8, 8), "source") + graph.convolution( + tensor, + module, + f"graph.{module}", + replicate_padding=1, + compute_dtype=_FakeTrt.float32, + ) + assert len(network.calls) == 1 + call = network.calls[0] + if module in EXPECTED_ZERO_PAD_SELECTION: + assert pad_calls == [] + assert call["tensor"] is tensor + assert call["layer"].padding_nd == (1, 1) + else: + assert pad_calls == [(1, f"graph.{module}.pad")] + assert call["tensor"].label == "replicate-padded" + assert call["layer"].padding_nd == (0, 0) + + +def test_selected_module_rejects_an_unexpected_convolution_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + network = _FakeNetwork() + graph = model_module._NativeMogeGraph(_FakeTrt, network, {}, fast_path=True) + weight = np.zeros((4, 4, 3, 3), dtype=np.float32) + bias = np.zeros((4,), dtype=np.float32) + monkeypatch.setattr( + graph, + "_array", + lambda name, expected=None: weight if name.endswith(".weight") else bias, + ) + monkeypatch.setattr(graph, "cast", lambda tensor, dtype, name: tensor) + tensor = _FakeTensor(_FakeTrt.float32, (1, 4, 8, 8), "source") + + with pytest.raises(ValueError, match="must be stride-1 3x3"): + graph.convolution( + tensor, + next(iter(EXPECTED_ZERO_PAD_SELECTION)), + "invalid", + replicate_padding=0, + compute_dtype=_FakeTrt.float32, + ) + + +def test_selective_zero_patch_contains_no_linear_or_level4_fusion() -> None: + source = Path(model_module.__file__).read_text(encoding="utf-8") + for forbidden in ( + "_ENABLE_LEVEL12_LINEAR_FUSION", + "_ENABLE_LEVEL4_HEAD_FUSION", + "_compose_deconv2_replicate_conv3", + ): + assert forbidden not in source diff --git a/src/runtime/models/moge/pipeline.cpp b/src/runtime/models/moge/pipeline.cpp index 69f54b2e20..e2f2ff61f7 100644 --- a/src/runtime/models/moge/pipeline.cpp +++ b/src/runtime/models/moge/pipeline.cpp @@ -21,11 +21,33 @@ namespace { constexpr int32_t kFocalRecoverySize = 64; constexpr double kDenominatorEpsilon = 1.0e-9; -constexpr int32_t kMinImageSize = 64; -constexpr int32_t kMaxImageSize = 2048; constexpr float kMinAspectRatio = 0.5F; constexpr float kMaxAspectRatio = 2.0F; -constexpr float kMaskThreshold = 0.5F; + +struct ImageProfileBounds { + int32_t min_height{0}; + int32_t min_width{0}; + int32_t max_height{0}; + int32_t max_width{0}; +}; + +bool valid_image_profile_shape(const std::vector& shape) { + constexpr auto kInt32Max = std::numeric_limits::max(); + return shape.size() == 4U && shape[0] == 1 && shape[1] > 0 && shape[1] <= kInt32Max && + shape[2] > 0 && shape[2] <= kInt32Max && shape[3] == 3; +} + +ImageProfileBounds image_profile_bounds(const ITrtModule& model) { + const auto profile = model.profile_idx(); + const auto minimum = model.input_profile_shape("image", profile, ProfileShapeSelector::kMin); + const auto maximum = model.input_profile_shape("image", profile, ProfileShapeSelector::kMax); + if (!valid_image_profile_shape(minimum) || !valid_image_profile_shape(maximum) || + minimum[1] > maximum[1] || minimum[2] > maximum[2]) { + throw std::runtime_error("MogePipeline: invalid TensorRT image profile"); + } + return {static_cast(minimum[1]), static_cast(minimum[2]), + static_cast(maximum[1]), static_cast(maximum[2])}; +} struct FocalSample { double u{0.0}; @@ -44,7 +66,26 @@ double normalized_coordinate(int32_t index, int32_t size, double span) { return span * (2.0 * index + 1.0 - size) / size; } -std::vector make_focal_samples(const float* points, const std::vector& mask, +int32_t nearest_sample_index(int32_t output_index, int32_t input_size) { + return std::min(input_size - 1, static_cast(static_cast(output_index) * + input_size / kFocalRecoverySize)); +} + +bool valid_focal_neighborhood(const uint16_t* valid, int32_t height, int32_t width, int32_t y, + int32_t x) { + if (y <= 0 || x <= 0 || y >= height - 1 || x >= width - 1) + return false; + for (int32_t neighbor_y = y - 1; neighbor_y <= y + 1; ++neighbor_y) { + for (int32_t neighbor_x = x - 1; neighbor_x <= x + 1; ++neighbor_x) { + const auto neighbor = static_cast(neighbor_y) * width + neighbor_x; + if (valid[neighbor] == 0) + return false; + } + } + return true; +} + +std::vector make_focal_samples(const float* sampled_points, const uint16_t* valid, int32_t height, int32_t width) { const double aspect = static_cast(width) / height; const double diagonal_factor = std::sqrt(1.0 + aspect * aspect); @@ -57,18 +98,18 @@ std::vector make_focal_samples(const float* points, const std::vect // floor(output_index * input_size / output_size). The official MoGe // recovery downsamples points, UVs, and the mask this way to 64x64. for (int32_t out_y = 0; out_y < kFocalRecoverySize; ++out_y) { - const int32_t y = std::min(height - 1, static_cast(static_cast(out_y) * - height / kFocalRecoverySize)); + const int32_t y = nearest_sample_index(out_y, height); for (int32_t out_x = 0; out_x < kFocalRecoverySize; ++out_x) { - const int32_t x = std::min(width - 1, static_cast(static_cast(out_x) * - width / kFocalRecoverySize)); + const int32_t x = nearest_sample_index(out_x, width); const auto pixel = static_cast(y) * width + x; - if (mask[pixel] == 0) + if (valid[pixel] == 0) continue; - const auto point = pixel * 3U; - const double px = points[point]; - const double py = points[point + 1U]; - const double pz = points[point + 2U]; + if (!valid_focal_neighborhood(valid, height, width, y, x)) + continue; + const auto sample = (static_cast(out_y) * kFocalRecoverySize + out_x) * 3U; + const double px = sampled_points[sample]; + const double py = sampled_points[sample + 1U]; + const double pz = sampled_points[sample + 2U]; if (!std::isfinite(px) || !std::isfinite(py) || !std::isfinite(pz)) continue; samples.push_back({normalized_coordinate(x, width, span_x), @@ -171,29 +212,24 @@ const float* require_float_output(const TensorMap& outputs, const char* name, return static_cast(tensor.data); } +const uint16_t* require_float16_output(const TensorMap& outputs, const char* name, + const std::vector& shape) { + const auto iterator = outputs.find(name); + if (iterator == outputs.end()) + throw std::runtime_error(std::string("MoGe engine did not return required output '") + + name + "'"); + const auto& tensor = iterator->second; + if (tensor.data == nullptr || tensor.dtype != DType::kFloat16 || tensor.shape != shape) { + throw std::runtime_error(std::string("MoGe output contract mismatch for '") + name + "'"); + } + return static_cast(tensor.data); +} + void validate_metric_scale(float metric_scale) { if (!std::isfinite(metric_scale) || metric_scale <= 0.0F) throw std::invalid_argument("MoGe metric scale must be finite and positive"); } -bool valid_raw_geometry_pixel(const float* affine_points, const float* mask_probabilities, - std::size_t pixel) { - const auto point = pixel * 3U; - return std::isfinite(mask_probabilities[pixel]) && mask_probabilities[pixel] > kMaskThreshold && - std::isfinite(affine_points[point]) && std::isfinite(affine_points[point + 1U]) && - std::isfinite(affine_points[point + 2U]); -} - -std::vector make_raw_mask(const float* affine_points, const float* mask_probabilities, - std::size_t area) { - std::vector raw_mask(area, 0); - for (std::size_t pixel = 0; pixel < area; ++pixel) { - raw_mask[pixel] = static_cast( - valid_raw_geometry_pixel(affine_points, mask_probabilities, pixel)); - } - return raw_mask; -} - struct MogeCalibration { double span_x{1.0}; double span_y{1.0}; @@ -243,8 +279,8 @@ void write_invalid_geometry(moge::GeometryResult& result, std::size_t pixel, flo result.points[point + 2U] = infinity; } -void populate_geometry_result(moge::GeometryResult& result, const float* affine_points, - const std::vector& raw_mask, const FocalShift& recovered, +void populate_geometry_result(moge::GeometryResult& result, const float* affine_depth, + const uint16_t* valid_pixels, const FocalShift& recovered, const MogeCalibration& calibration, float metric_scale) { const float infinity = std::numeric_limits::infinity(); for (int32_t y = 0; y < result.height; ++y) { @@ -253,8 +289,9 @@ void populate_geometry_result(moge::GeometryResult& result, const float* affine_ const auto pixel = static_cast(y) * result.width + x; const auto point = pixel * 3U; const double depth_unscaled = - static_cast(affine_points[point + 2U]) + recovered.shift; - const bool valid = valid_metric_depth(raw_mask[pixel], depth_unscaled); + static_cast(affine_depth[pixel]) + recovered.shift; + const bool valid = + valid_metric_depth(static_cast(valid_pixels[pixel] != 0), depth_unscaled); result.mask[pixel] = static_cast(valid); if (!valid) { write_invalid_geometry(result, pixel, infinity); @@ -270,9 +307,9 @@ void populate_geometry_result(moge::GeometryResult& result, const float* affine_ } } -bool supported_image_size(int32_t height, int32_t width) { - return height >= kMinImageSize && width >= kMinImageSize && height <= kMaxImageSize && - width <= kMaxImageSize; +bool supported_image_size(int32_t height, int32_t width, const ImageProfileBounds& profile) { + return height >= profile.min_height && width >= profile.min_width && + height <= profile.max_height && width <= profile.max_width; } bool supported_aspect_ratio(int32_t height, int32_t width) { @@ -281,13 +318,14 @@ bool supported_aspect_ratio(int32_t height, int32_t width) { } bool valid_rgb_value(float value) { - return std::isfinite(value) && value >= 0.0F && value <= 1.0F; + return value >= 0.0F && value <= 1.0F; } -void validate_image_input(const float* pixels, int32_t height, int32_t width) { +void validate_image_input(const float* pixels, int32_t height, int32_t width, + const ImageProfileBounds& profile) { if (pixels == nullptr) throw std::invalid_argument("MoGe image pointer is null"); - if (!supported_image_size(height, width)) + if (!supported_image_size(height, width, profile)) throw std::invalid_argument("MoGe image dimensions are outside the bundle profile"); if (!supported_aspect_ratio(height, width)) throw std::invalid_argument("MoGe image aspect ratio is outside the supported range"); @@ -298,10 +336,9 @@ void validate_image_input(const float* pixels, int32_t height, int32_t width) { } } -std::optional recover_focal_shift(const float* affine_points, - const std::vector& mask, int32_t height, - int32_t width) { - const auto samples = make_focal_samples(affine_points, mask, height, width); +std::optional recover_focal_shift(const float* focal_samples, const uint16_t* valid, + int32_t height, int32_t width) { + const auto samples = make_focal_samples(focal_samples, valid, height, width); if (samples.size() < 2U) return std::nullopt; @@ -336,20 +373,17 @@ std::optional recover_focal_shift(const float* affine_points, return FocalShift{static_cast(current.focal), static_cast(shift)}; } -moge::GeometryResult postprocess_geometry(const float* affine_points, - const float* mask_probabilities, float metric_scale, +moge::GeometryResult postprocess_geometry(const float* affine_depth, const uint16_t* valid, + const float* focal_samples, float metric_scale, int32_t height, int32_t width) { const auto area = static_cast(height) * width; validate_metric_scale(metric_scale); - const auto raw_mask = make_raw_mask(affine_points, mask_probabilities, area); - - const auto recovered = recover_focal_shift(affine_points, raw_mask, height, width); + const auto recovered = recover_focal_shift(focal_samples, valid, height, width); if (!recovered) throw std::runtime_error("MoGe could not recover camera focal length and shift"); const auto calibration = make_calibration(*recovered, height, width); auto result = make_geometry_result(height, width, area, calibration); - populate_geometry_result(result, affine_points, raw_mask, *recovered, calibration, - metric_scale); + populate_geometry_result(result, affine_depth, valid, *recovered, calibration, metric_scale); return result; } @@ -359,18 +393,26 @@ MogePipeline::MogePipeline(std::unique_ptr model, std::string model_ : model_(std::move(model)), model_id_(std::move(model_id)) { if (!model_ || !model_->ok()) throw std::runtime_error("MogePipeline: invalid model"); + const auto profile = image_profile_bounds(*model_); + min_image_height_ = profile.min_height; + min_image_width_ = profile.min_width; + max_image_height_ = profile.max_height; + max_image_width_ = profile.max_width; } moge::GeometryResult MogePipeline::estimate_geometry(const float* pixels, int32_t height, int32_t width) { - validate_image_input(pixels, height, width); - + validate_image_input( + pixels, height, width, + {min_image_height_, min_image_width_, max_image_height_, max_image_width_}); Tensor image{const_cast(pixels), {1, height, width, 3}, DType::kFloat32}; const auto outputs = model_->forward({{"image", image}}); - const auto* points = require_float_output(outputs, "points", {1, height, width, 3}); - const auto* mask = require_float_output(outputs, "mask", {1, height, width}); + const auto* affine_depth = require_float_output(outputs, "affine_depth", {1, height, width}); + const auto* valid = require_float16_output(outputs, "valid", {1, height, width}); + const auto* focal_samples = require_float_output( + outputs, "focal_samples", {1, kFocalRecoverySize, kFocalRecoverySize, 3}); const auto* scale = require_float_output(outputs, "metric_scale", {1}); - return postprocess_geometry(points, mask, scale[0], height, width); + return postprocess_geometry(affine_depth, valid, focal_samples, scale[0], height, width); } } // namespace trtmc diff --git a/src/runtime/models/moge/pipeline.h b/src/runtime/models/moge/pipeline.h index 3046719657..dc4d97c3e9 100644 --- a/src/runtime/models/moge/pipeline.h +++ b/src/runtime/models/moge/pipeline.h @@ -28,6 +28,10 @@ class MogePipeline final : public IPipeline, public moge::IGeometryEstimator { private: std::unique_ptr model_; std::string model_id_; + int32_t min_image_height_{0}; + int32_t min_image_width_{0}; + int32_t max_image_height_{0}; + int32_t max_image_width_{0}; }; } // namespace trtmc diff --git a/tests/cpp/models/moge/test_moge_pipeline.cpp b/tests/cpp/models/moge/test_moge_pipeline.cpp index 6eb5ffc88e..a1ae6220df 100644 --- a/tests/cpp/models/moge/test_moge_pipeline.cpp +++ b/tests/cpp/models/moge/test_moge_pipeline.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -50,16 +51,45 @@ std::vector affine_points(int32_t height, int32_t width, float focal, flo return points; } +std::vector affine_depth(const std::vector& points) { + std::vector depth(points.size() / 3U); + for (std::size_t pixel = 0; pixel < depth.size(); ++pixel) + depth[pixel] = points[pixel * 3U + 2U]; + return depth; +} + +std::vector focal_samples(const std::vector& points, int32_t height, int32_t width) { + constexpr int32_t sample_size = 64; + std::vector samples(static_cast(sample_size) * sample_size * 3U); + for (int32_t out_y = 0; out_y < sample_size; ++out_y) { + const int32_t y = static_cast(static_cast(out_y) * height / sample_size); + for (int32_t out_x = 0; out_x < sample_size; ++out_x) { + const int32_t x = + static_cast(static_cast(out_x) * width / sample_size); + const auto source = (static_cast(y) * width + x) * 3U; + const auto target = (static_cast(out_y) * sample_size + out_x) * 3U; + std::copy_n(points.data() + source, 3, samples.data() + target); + } + } + return samples; +} + class FakeMogeModule final : public trtmc::ITrtModule { public: - FakeMogeModule(int32_t height, int32_t width, bool invalidate_pixel = false) - : height_(height), width_(width), points_(affine_points(height, width, 0.8F, 1.25F)), - mask_(static_cast(height) * width, 0.9F) { + FakeMogeModule(int32_t height, int32_t width, bool invalidate_pixel = false, + trtmc::DType valid_dtype = trtmc::DType::kFloat16) + : height_(height), width_(width), min_height_(height), min_width_(width), + max_height_(height), max_width_(width), + points_(affine_points(height, width, 0.8F, 1.25F)), depth_(affine_depth(points_)), + samples_(focal_samples(points_, height, width)), + valid_(static_cast(height) * width, uint16_t{0x3C00}), + valid_dtype_(valid_dtype) { if (invalidate_pixel) - mask_[5] = 0.1F; + valid_[5] = 0; } trtmc::TensorMap forward(const trtmc::TensorMap& inputs) override { + ++forward_count_; const auto image = inputs.find("image"); if (image != inputs.end()) { input_shape = image->second.shape; @@ -67,8 +97,9 @@ class FakeMogeModule final : public trtmc::ITrtModule { input_values.assign(values, values + image->second.numel()); } return { - {"points", {points_.data(), {1, height_, width_, 3}, trtmc::DType::kFloat32}}, - {"mask", {mask_.data(), {1, height_, width_}, trtmc::DType::kFloat32}}, + {"affine_depth", {depth_.data(), {1, height_, width_}, trtmc::DType::kFloat32}}, + {"valid", {valid_.data(), {1, height_, width_}, valid_dtype_}}, + {"focal_samples", {samples_.data(), {1, 64, 64, 3}, trtmc::DType::kFloat32}}, {"metric_scale", {scale_.data(), {1}, trtmc::DType::kFloat32}}, }; } @@ -79,36 +110,72 @@ class FakeMogeModule final : public trtmc::ITrtModule { cudaStream_t stream() const override { return nullptr; } void enable_cuda_graph() override {} bool cuda_graph_active() const override { return false; } - int32_t profile_idx() const override { return 0; } + int32_t profile_idx() const override { return profile_index_; } std::vector input_info() const override { return {}; } std::vector output_info() const override { return {}; } bool has_input(const std::string& name) const override { return name == "image"; } bool has_output(const std::string& name) const override { - return name == "points" || name == "mask" || name == "metric_scale"; + return name == "affine_depth" || name == "valid" || name == "focal_samples" || + name == "metric_scale"; + } + trtmc::DType tensor_dtype(const std::string& name) const override { + return name == "valid" ? valid_dtype_ : trtmc::DType::kFloat32; } - trtmc::DType tensor_dtype(const std::string&) const override { return trtmc::DType::kFloat32; } std::vector tensor_shape(const std::string& name) const override { if (name == "image") return {1, height_, width_, 3}; - if (name == "points") - return {1, height_, width_, 3}; - if (name == "mask") + if (name == "affine_depth" || name == "valid") return {1, height_, width_}; + if (name == "focal_samples") + return {1, 64, 64, 3}; if (name == "metric_scale") return {1}; throw std::runtime_error("unknown fake tensor"); } - std::vector input_profile_shape(const std::string&, int32_t, - trtmc::ProfileShapeSelector) const override { + std::vector input_profile_shape(const std::string&, int32_t profile_index, + trtmc::ProfileShapeSelector selector) const override { + last_queried_profile_index_ = profile_index; + ++profile_query_count_; + if (selector == trtmc::ProfileShapeSelector::kMin) + return {1, min_height_, min_width_, 3}; + if (selector == trtmc::ProfileShapeSelector::kMax) + return {1, max_height_, max_width_, 3}; return {1, height_, width_, 3}; } - int32_t optimization_profile_count() const override { return 1; } + int32_t optimization_profile_count() const override { return profile_index_ + 1; } void* device_ptr(const std::string&) const override { return nullptr; } void bind_external(const std::string&, void*) override {} bool ok() const override { return true; } void keep_alive(std::shared_ptr) override {} - void invalidate_all() { std::fill(mask_.begin(), mask_.end(), 0.1F); } + void invalidate_all() { std::fill(valid_.begin(), valid_.end(), uint16_t{0}); } + void set_valid(int32_t y, int32_t x, bool value) { + valid_.at(static_cast(y) * width_ + x) = value ? uint16_t{0x3C00} : 0; + } + void set_valid_neighborhood(int32_t y, int32_t x) { + for (int32_t neighbor_y = y - 1; neighbor_y <= y + 1; ++neighbor_y) { + for (int32_t neighbor_x = x - 1; neighbor_x <= x + 1; ++neighbor_x) + set_valid(neighbor_y, neighbor_x, true); + } + } + void set_focal_sample(int32_t y, int32_t x, float px, float py, float pz) { + const auto sample = (static_cast(y) * 64U + x) * 3U; + samples_.at(sample) = px; + samples_.at(sample + 1U) = py; + samples_.at(sample + 2U) = pz; + } + uint16_t valid_bits(std::size_t pixel) const { return valid_.at(pixel); } + int32_t forward_count() const { return forward_count_; } + int32_t last_queried_profile_index() const { return last_queried_profile_index_; } + int32_t profile_query_count() const { return profile_query_count_; } + void set_profile_index(int32_t profile_index) { profile_index_ = profile_index; } + void set_profile_bounds(int32_t min_height, int32_t min_width, int32_t max_height, + int32_t max_width) { + min_height_ = min_height; + min_width_ = min_width; + max_height_ = max_height; + max_width_ = max_width; + } std::vector input_shape; std::vector input_values; @@ -116,8 +183,19 @@ class FakeMogeModule final : public trtmc::ITrtModule { private: int32_t height_; int32_t width_; + int32_t min_height_; + int32_t min_width_; + int32_t max_height_; + int32_t max_width_; + int32_t forward_count_{0}; + int32_t profile_index_{0}; + mutable int32_t last_queried_profile_index_{-1}; + mutable int32_t profile_query_count_{0}; std::vector points_; - std::vector mask_; + std::vector depth_; + std::vector samples_; + std::vector valid_; + trtmc::DType valid_dtype_; std::vector scale_{2.0F}; }; @@ -163,11 +241,13 @@ void test_pipeline_recovers_metric_geometry_from_hwc_input() { check(dynamic_cast(&pipeline) != nullptr, "MoGe exposes its family-owned geometry contract"); check(std::string(pipeline.model_id()) == "moge-2-vitl", "MoGe model id"); + check(module_ptr->valid_bits(0) == 0x3C00, "MoGe valid=true uses FP16 one bits"); } void test_invalid_mask_materializes_infinity() { constexpr int32_t size = 64; auto module = std::make_unique(size, size, true); + auto* module_ptr = module.get(); trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); auto image = rgb_image(size, size); @@ -176,6 +256,24 @@ void test_invalid_mask_materializes_infinity() { check(result.mask[5] == 0, "MoGe invalid pixel mask cleared"); check(std::isinf(result.depth[5]), "MoGe invalid depth is infinity"); check(std::isinf(result.points[15]), "MoGe invalid point is infinity"); + check(module_ptr->valid_bits(5) == 0x0000, "MoGe valid=false uses FP16 zero bits"); + check(module_ptr->valid_bits(0) == 0x3C00, "MoGe retained valid pixel uses FP16 one bits"); +} + +void test_legacy_int8_valid_contract_is_rejected() { + constexpr int32_t size = 64; + auto module = std::make_unique(size, size, false, trtmc::DType::kInt8); + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + auto image = rgb_image(size, size); + + try { + (void)pipeline.estimate_geometry(image.data(), size, size); + check(false, "MoGe rejects legacy INT8 valid output"); + } catch (const std::runtime_error& error) { + check(std::string(error.what()).find("output contract mismatch for 'valid'") != + std::string::npos, + "MoGe legacy INT8 valid rejection is explicit"); + } } void test_focal_recovery_failure_is_reported() { @@ -194,12 +292,140 @@ void test_focal_recovery_failure_is_reported() { } } +void test_invalid_rgb_values_are_rejected() { + constexpr int32_t size = 64; + const std::vector invalid_values = { + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + -0.01F, + 1.01F, + }; + for (const float value : invalid_values) { + auto module = std::make_unique(size, size); + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + auto image = rgb_image(size, size); + image[0] = value; + try { + (void)pipeline.estimate_geometry(image.data(), size, size); + check(false, "MoGe rejects non-finite or out-of-range RGB input"); + } catch (const std::invalid_argument& error) { + check(std::string(error.what()).find("RGB input values") != std::string::npos, + "MoGe RGB rejection is explicit"); + } + } +} + +void test_loaded_engine_profile_bounds_are_enforced_before_forward() { + constexpr int32_t height = 64; + constexpr int32_t width = 80; + auto module = std::make_unique(height, width); + module->set_profile_index(2); + module->set_profile_bounds(height, width, 128, 160); + auto* module_ptr = module.get(); + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + check(module_ptr->last_queried_profile_index() == 2 && module_ptr->profile_query_count() == 2, + "MoGe reads min and max bounds from the active engine profile"); + + for (const auto& [input_height, input_width] : + {std::pair{height - 1, width}, std::pair{height, width - 1}, std::pair{129, width}, + std::pair{height, 161}}) { + auto image = rgb_image(input_height, input_width); + try { + (void)pipeline.estimate_geometry(image.data(), input_height, input_width); + check(false, "MoGe rejects dimensions outside the loaded engine profile"); + } catch (const std::invalid_argument& error) { + check(std::string(error.what()).find("outside the bundle profile") != std::string::npos, + "MoGe profile-bound rejection is explicit"); + } + } + check(module_ptr->forward_count() == 0, + "MoGe rejects profile-incompatible dimensions before TensorRT forward"); +} + +void test_invalid_loaded_engine_profile_is_rejected() { + auto module = std::make_unique(64, 80); + module->set_profile_bounds(128, 160, 64, 80); + try { + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + check(false, "MoGe rejects an invalid loaded engine profile"); + } catch (const std::runtime_error& error) { + check(std::string(error.what()).find("invalid TensorRT image profile") != std::string::npos, + "MoGe invalid engine profile rejection is explicit"); + } +} + +void test_focal_sampling_excludes_mapped_image_edges() { + constexpr int32_t size = 64; + auto module = std::make_unique(size, size); + for (int32_t index = 0; index < size; ++index) { + module->set_focal_sample(0, index, 1000.0F, -1000.0F, 0.1F); + module->set_focal_sample(size - 1, index, 1000.0F, -1000.0F, 0.1F); + module->set_focal_sample(index, 0, 1000.0F, -1000.0F, 0.1F); + module->set_focal_sample(index, size - 1, 1000.0F, -1000.0F, 0.1F); + } + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + auto image = rgb_image(size, size); + + const auto result = pipeline.estimate_geometry(image.data(), size, size); + + check(close(result.intrinsics[0], 0.5656854F), + "MoGe focal sampling excludes mapped image-edge points"); + check(close(result.intrinsics[4], 0.5656854F), + "MoGe image-edge exclusion preserves normalized fy"); +} + +void test_focal_sampling_excludes_invalid_three_by_three_neighborhood() { + constexpr int32_t size = 64; + constexpr int32_t sample_y = 20; + constexpr int32_t sample_x = 20; + auto module = std::make_unique(size, size); + module->set_focal_sample(sample_y, sample_x, 1000.0F, -1000.0F, 0.1F); + module->set_valid(sample_y, sample_x + 1, false); + auto* module_ptr = module.get(); + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + auto image = rgb_image(size, size); + + const auto result = pipeline.estimate_geometry(image.data(), size, size); + + check(module_ptr->valid_bits(static_cast(sample_y) * size + sample_x) == 0x3C00, + "MoGe focal center remains valid when its neighbor is invalid"); + check(close(result.intrinsics[0], 0.5656854F), + "MoGe focal sampling excludes a center with an invalid neighbor"); +} + +void test_focal_sampling_retains_complete_interior_neighborhoods() { + constexpr int32_t size = 64; + auto module = std::make_unique(size, size); + module->invalidate_all(); + for (int32_t y : {16, 32, 48}) { + for (int32_t x : {16, 32, 48}) + module->set_valid_neighborhood(y, x); + } + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + auto image = rgb_image(size, size); + + const auto result = pipeline.estimate_geometry(image.data(), size, size); + + check(close(result.intrinsics[0], 0.5656854F), + "MoGe focal sampling retains complete interior neighborhoods"); + check(result.mask[32U * size + 32U] == 1, + "MoGe retained interior focal center remains valid geometry"); +} + } // namespace int main() { test_pipeline_recovers_metric_geometry_from_hwc_input(); test_invalid_mask_materializes_infinity(); + test_legacy_int8_valid_contract_is_rejected(); test_focal_recovery_failure_is_reported(); + test_invalid_rgb_values_are_rejected(); + test_loaded_engine_profile_bounds_are_enforced_before_forward(); + test_invalid_loaded_engine_profile_is_rejected(); + test_focal_sampling_excludes_mapped_image_edges(); + test_focal_sampling_excludes_invalid_three_by_three_neighborhood(); + test_focal_sampling_retains_complete_interior_neighborhoods(); if (g_failures != 0) { std::cerr << g_failures << " MoGe pipeline test(s) failed\n"; return 1;