triton_fk: Triton backend for joint_parameters_to_local_skeleton_state (#1536)

Summary:
Pull Request resolved: #1536

Adds a Triton local-FK to `triton_fk` and a `backend` param to `Skeleton.joint_parameters_to_local_skeleton_state` and `GpuCharacter.{joint,model}_parameters_to_local_skeleton_state` (default "torch"), matching the existing global FK (`joint_parameters_to_skeleton_state`). The local skeleton state is exactly the per-joint state the global FK computes BEFORE the prefix-product up the kinematic tree, so the new kernels reuse the existing device functions: the forward stores `_load_local_state`'s output (no parent composition), and the backward is `_store_joint_param_gradients` fed grad_output directly (no tree -- grad w.r.t. the local state IS grad_output per joint). On CUDA float32 `backend="auto"` opts into Triton; it falls back to torch under torch.jit.script/trace.

Reviewed By: cstollmeta

Differential Revision: D108091266

fbshipit-source-id: d95eb4415769d887f42683e4097351c5e04c1151

Author: cdtwigg

pymomentum/backend/triton_fk.py

@@ -715,6 +715,83 @@ def _backward_kernel(
             )
             joint -= 1
 
+    @triton.jit
+    def _local_forward_kernel(
+        joint_params,
+        offsets,
+        prerotations,
+        output,
+        batch_size,
+        n_joints,
+        BLOCK_BATCH: tl.constexpr,
+    ):
+        # Local skeleton state: the per-joint state from _load_local_state, WITHOUT
+        # the parent composition that _forward_one_joint applies (i.e. the global FK
+        # before the prefix product up the kinematic tree).
+        batch = tl.program_id(0) * BLOCK_BATCH + tl.arange(0, BLOCK_BATCH)
+        mask = batch < batch_size
+        joint = 0
+        while joint < n_joints:
+            tx, ty, tz, qx, qy, qz, qw, s = _load_local_state(
+                joint_params, offsets, prerotations, batch, joint, n_joints, mask
+            )
+            out_base = output + batch * n_joints * 8 + joint * 8
+            tl.store(out_base + 0, tx, mask=mask)
+            tl.store(out_base + 1, ty, mask=mask)
+            tl.store(out_base + 2, tz, mask=mask)
+            tl.store(out_base + 3, qx, mask=mask)
+            tl.store(out_base + 4, qy, mask=mask)
+            tl.store(out_base + 5, qz, mask=mask)
+            tl.store(out_base + 6, qw, mask=mask)
+            tl.store(out_base + 7, s, mask=mask)
+            joint += 1
+
+    @triton.jit
+    def _local_backward_kernel(
+        joint_params,
+        offsets,
+        prerotations,
+        grad_output,
+        grad_joint_params,
+        batch_size,
+        n_joints,
+        BLOCK_BATCH: tl.constexpr,
+    ):
+        # Backward of the local FK: there is no tree, so grad w.r.t. the local
+        # skeleton state IS grad_output per joint. _store_joint_param_gradients
+        # propagates it to joint_params (translation identity, rotation through the
+        # prerotation*euler quaternion multiply + euler, scale through exp2).
+        batch = tl.program_id(0) * BLOCK_BATCH + tl.arange(0, BLOCK_BATCH)
+        mask = batch < batch_size
+        joint = 0
+        while joint < n_joints:
+            _, _, _, _, _, _, _, ls = _load_local_state(
+                joint_params, offsets, prerotations, batch, joint, n_joints, mask
+            )
+            gtx, gty, gtz, gqx, gqy, gqz, gqw, gs = _load_state(
+                grad_output, batch, joint, n_joints, mask
+            )
+            _store_joint_param_gradients(
+                joint_params,
+                prerotations,
+                grad_joint_params,
+                batch,
+                joint,
+                n_joints,
+                mask,
+                True,
+                gtx,
+                gty,
+                gtz,
+                gqx,
+                gqy,
+                gqz,
+                gqw,
+                gs,
+                ls,
+            )
+            joint += 1
+
 
 def _require_triton() -> None:
     if triton is None or tl is None:
@@ -907,3 +984,134 @@ def joint_parameters_to_skeleton_state(
         joint_parents,
         active_joints,
     )
+
+
+def _validate_local_inputs(
+    joint_parameters: torch.Tensor,
+    joint_translation_offsets: torch.Tensor,
+    joint_prerotations: torch.Tensor,
+) -> None:
+    # Same as _validate_inputs but the local FK needs no kinematic tree (parents).
+    if not joint_parameters.is_cuda:
+        raise RuntimeError("Triton local FK requested, but joint_parameters is on CPU.")
+    if joint_parameters.dtype != torch.float32:
+        raise RuntimeError(
+            "Triton local FK currently only supports float32 joint parameters."
+        )
+    if joint_parameters.ndim < 2 or joint_parameters.shape[-1] != 7:
+        raise RuntimeError("joint_parameters must have shape [..., num_joints, 7].")
+    if joint_translation_offsets.device != joint_parameters.device:
+        raise RuntimeError("joint_translation_offsets must be on the same device.")
+    if joint_prerotations.device != joint_parameters.device:
+        raise RuntimeError("joint_prerotations must be on the same device.")
+    if joint_translation_offsets.dtype != torch.float32:
+        raise RuntimeError("joint_translation_offsets must be float32.")
+    if joint_prerotations.dtype != torch.float32:
+        raise RuntimeError("joint_prerotations must be float32.")
+    num_joints = joint_parameters.shape[-2]
+    if joint_translation_offsets.shape != (num_joints, 3):
+        raise RuntimeError("joint_translation_offsets must have shape [num_joints, 3].")
+    if joint_prerotations.shape != (num_joints, 4):
+        raise RuntimeError("joint_prerotations must have shape [num_joints, 4].")
+
+
+def _launch_local_forward(
+    joint_parameters: torch.Tensor,
+    joint_translation_offsets: torch.Tensor,
+    joint_prerotations: torch.Tensor,
+) -> torch.Tensor:
+    _validate_local_inputs(
+        joint_parameters, joint_translation_offsets, joint_prerotations
+    )
+    _require_triton()
+    num_joints = joint_parameters.shape[-2]
+    batch_size = joint_parameters.numel() // (num_joints * 7)
+    output = torch.empty(
+        (*joint_parameters.shape[:-1], 8),
+        device=joint_parameters.device,
+        dtype=joint_parameters.dtype,
+    )
+    grid = (triton.cdiv(batch_size, _BLOCK_BATCH),)
+    _local_forward_kernel[grid](
+        joint_parameters,
+        joint_translation_offsets,
+        joint_prerotations,
+        output,
+        batch_size,
+        num_joints,
+        BLOCK_BATCH=_BLOCK_BATCH,
+    )
+    return output
+
+
+def _launch_local_backward(
+    joint_parameters: torch.Tensor,
+    joint_translation_offsets: torch.Tensor,
+    joint_prerotations: torch.Tensor,
+    grad_output: torch.Tensor,
+) -> torch.Tensor:
+    _require_triton()
+    grad_output = grad_output.contiguous()
+    grad_joint_parameters = torch.empty_like(joint_parameters)
+    num_joints = joint_parameters.shape[-2]
+    batch_size = joint_parameters.numel() // (num_joints * 7)
+    grid = (triton.cdiv(batch_size, _BLOCK_BATCH),)
+    _local_backward_kernel[grid](
+        joint_parameters,
+        joint_translation_offsets,
+        joint_prerotations,
+        grad_output,
+        grad_joint_parameters,
+        batch_size,
+        num_joints,
+        BLOCK_BATCH=_BLOCK_BATCH,
+    )
+    return grad_joint_parameters
+
+
+class _JointParametersToLocalSkeletonState(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: torch.autograd.function.FunctionCtx,
+        joint_parameters: torch.Tensor,
+        joint_translation_offsets: torch.Tensor,
+        joint_prerotations: torch.Tensor,
+    ) -> torch.Tensor:
+        joint_parameters = joint_parameters.contiguous()
+        joint_translation_offsets = joint_translation_offsets.contiguous()
+        joint_prerotations = joint_prerotations.contiguous()
+        output = _launch_local_forward(
+            joint_parameters, joint_translation_offsets, joint_prerotations
+        )
+        ctx.save_for_backward(
+            joint_parameters, joint_translation_offsets, joint_prerotations
+        )
+        return output
+
+    @staticmethod
+    def backward(
+        ctx: torch.autograd.function.FunctionCtx,
+        grad_output: torch.Tensor,
+    ) -> tuple[torch.Tensor | None, None, None]:
+        joint_parameters, joint_translation_offsets, joint_prerotations = (
+            ctx.saved_tensors
+        )
+        grad_joint_parameters = _launch_local_backward(
+            joint_parameters,
+            joint_translation_offsets,
+            joint_prerotations,
+            grad_output,
+        )
+        return grad_joint_parameters, None, None
+
+
+def joint_parameters_to_local_skeleton_state(
+    joint_parameters: torch.Tensor,
+    joint_translation_offsets: torch.Tensor,
+    joint_prerotations: torch.Tensor,
+) -> torch.Tensor:
+    return _JointParametersToLocalSkeletonState.apply(
+        joint_parameters,
+        joint_translation_offsets,
+        joint_prerotations,
+    )

pymomentum/test/test_triton_fk_backend.py

@@ -154,6 +154,80 @@ def test_fk_with_model_parameters_matches_torch(self) -> None:
         triton_state.backward(grad_output)
         self._assert_grad_close(triton_params, torch_params)
 
+    def test_local_fk_with_model_parameters_matches_torch(self) -> None:
+        device = _cuda_device()
+        character_cpu = pym_test_utils.create_test_character()
+        character = pym_character.Character(
+            character_cpu,
+            has_parameter_transform=True,
+            has_skeleton=True,
+            has_rest_mesh=False,
+            has_skinning=False,
+        ).to(device)
+
+        torch.manual_seed(0)
+        model_params = torch.rand(
+            5,
+            character_cpu.parameter_transform.size,
+            device=device,
+            dtype=torch.float32,
+        )
+
+        torch_params = model_params.detach().clone().requires_grad_()
+        triton_params = model_params.detach().clone().requires_grad_()
+
+        torch_state = character.model_parameters_to_local_skeleton_state(
+            torch_params,
+            backend="torch",
+        )
+        triton_state = character.model_parameters_to_local_skeleton_state(
+            triton_params,
+            backend="triton",
+        )
+
+        torch.testing.assert_close(torch_state, triton_state, atol=1e-5, rtol=1e-5)
+
+        grad_output = torch.randn_like(torch_state)
+        torch_state.backward(grad_output)
+        triton_state.backward(grad_output)
+        self._assert_grad_close(triton_params, torch_params)
+
+    def test_local_fk_matches_torch(self) -> None:
+        # Direct joint-parameter local FK (no parameter transform), exercising the
+        # _local_forward_kernel / _local_backward_kernel without depending on the
+        # model-parameter transform. Covers both a single and a multi-dim batch.
+        device = _cuda_device()
+        character_cpu = pym_test_utils.create_test_character()
+        character = pym_character.Character(
+            character_cpu,
+            has_parameter_transform=False,
+            has_skeleton=True,
+            has_rest_mesh=False,
+            has_skinning=False,
+        ).to(device)
+
+        torch.manual_seed(0)
+        joint_params = torch.rand(
+            5,
+            character_cpu.skeleton.size,
+            7,
+            device=device,
+            dtype=torch.float32,
+        )
+        self._assert_local_joint_parameter_backend_close(character, joint_params)
+
+        batched_joint_params = torch.rand(
+            2,
+            3,
+            character_cpu.skeleton.size,
+            7,
+            device=device,
+            dtype=torch.float32,
+        )
+        self._assert_local_joint_parameter_backend_close(
+            character, batched_joint_params
+        )
+
     def test_inverse_fk_matches_torch(self) -> None:
         device = _cuda_device()
         character_cpu = pym_test_utils.create_test_character()
@@ -224,6 +298,30 @@ def _assert_joint_parameter_backend_close(
         triton_state.backward(grad_output)
         self._assert_grad_close(triton_params, torch_params)
 
+    def _assert_local_joint_parameter_backend_close(
+        self,
+        character: pym_character.Character,
+        joint_params: torch.Tensor,
+    ) -> None:
+        torch_params = joint_params.detach().clone().requires_grad_()
+        triton_params = joint_params.detach().clone().requires_grad_()
+
+        torch_state = character.joint_parameters_to_local_skeleton_state(
+            torch_params,
+            backend="torch",
+        )
+        triton_state = character.joint_parameters_to_local_skeleton_state(
+            triton_params,
+            backend="triton",
+        )
+
+        torch.testing.assert_close(torch_state, triton_state, atol=1e-5, rtol=1e-5)
+
+        grad_output = torch.randn_like(torch_state)
+        torch_state.backward(grad_output)
+        triton_state.backward(grad_output)
+        self._assert_grad_close(triton_params, torch_params)
+
     def _assert_grad_close(
         self,
         triton_tensor: torch.Tensor,