Handling for transfer_n_pin and trivially_replicable

Signed-off-by: Alex Vasile <[email protected]>

Author: Alex-Vasile

sharktank/sharktank/ops/_registry.py

@@ -270,7 +270,18 @@ def __call__(self, *args, **kwargs):
         selected_override, *results = trampoline(self, *args, **kwargs)
         if _ENABLE_TEST_LAST_OP_DISPATCH:
             global _TEST_LAST_OP_DISPATCH
-            _TEST_LAST_OP_DISPATCH = selected_override
+
+            if hasattr(selected_override, "_trivially_replicable_wrapper"):
+                # For trivially replicable wrappers, don't set _TEST_LAST_OP_DISPATCH
+                # the inner calls (which occured already)will set it to the actual op.
+                # NOTE: This assumes that all shards called the same op.
+                pass
+            else:
+                # For wrappers such as `transfer_n_pin`, we set _TEST_LAST_OP_DISPATCH to the original op (not the wrapper).
+                _TEST_LAST_OP_DISPATCH = getattr(
+                    selected_override, "_unwrapped", selected_override
+                )
+
         arity = len(results)
         if arity == 1:
             return results[0]

sharktank/sharktank/ops/sharded_impls.py

@@ -109,6 +109,18 @@ def func_wrapper(*args: Tuple, **kwargs: Dict[str, Any]):
             return res
 
         func_wrapper._impl_name = getattr(f, "_impl_name", None)  # For impl selection
+
+        if hasattr(f, "_trivially_replicable_wrapper"):
+            # If wrapping a trivially replicable function, we do not know what underlying op will be called on each shard,
+            # since we don't dispatch based on shards.
+            # Instead label this wrapper as a trivially replicable wrapper so that
+            # _TEST_LAST_OP_DISPATCH tracking can handle it correctly.
+            # _TEST_LAST_OP_DISPATCH will not update for this wrapper, but instead allow the last inner call to set it.
+            func_wrapper._trivially_replicable_wrapper = f._trivially_replicable_wrapper
+        else:
+            # We know the underlying op will be called, set __unwrapped__ to the original op
+            # so that _TEST_LAST_OP_DISPATCH tracking can handle it correctly.
+            func_wrapper._unwrapped = f
         return func_wrapper
 
     def wrap_override(signature_dispatcher_override):

sharktank/sharktank/ops/utils.py

@@ -152,6 +152,10 @@ def fn(a: torch.Tensor) -> torch.Tensor:
     def wrapper(*args, **kwargs):
         return call_trivially_replicable(fn, args, kwargs)
 
+    # Mark the trivially replicable wrapper as such so that _TEST_LAST_OP_DISPATCH tracking handles it correctly.
+    # There is no way for us to know which op will be called on each shard, so we cannot set _unwrapped the way we can for `transfer_n_pin`.
+    wrapper._trivially_replicable_wrapper = True
+
     return wrapper
 
 

sharktank/tests/ops/dispatch_test.py

@@ -0,0 +1,70 @@
+# Copyright 2025 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from typing import Callable
+import unittest
+import torch
+
+from parameterized import parameterized
+
+from sharktank import ops
+from sharktank.ops.sharded_impls import zeros_replicated
+from sharktank.types import DefaultPrimitiveTensor
+from sharktank.ops.default_impls import abs_default, cos_default, zeros_default
+from sharktank.types.tensors import InferenceTensor, ReplicatedTensor
+from sharktank.utils.testing import assert_tensor_close
+
+
+class DispatchTest(unittest.TestCase):
+    def setUp(self):
+        ops._registry._test_enable_last_op_dispatch(True)
+
+    def tearDown(self):
+        ops._registry._test_enable_last_op_dispatch(False)
+
+    def make_tensor(self, shards: int) -> InferenceTensor:
+        tensor = torch.tensor([1.0, 2.0, 3.0])
+        if shards == 1:
+            return DefaultPrimitiveTensor(data=tensor)
+        else:
+            return ReplicatedTensor(ts=tensor, shard_count=shards)
+
+    @parameterized.expand(
+        [
+            (None, zeros_default),
+            ((2, 3), zeros_replicated._unwrapped),
+        ]
+    )
+    def test_non_trivially_replicable_op(
+        self, devices: tuple[int, ...] | None, expected_dispatch: Callable
+    ):
+        ops.zeros(1, devices=devices)
+
+        last_dispatch_after_zeros = ops._registry._test_get_last_op_dispatch()
+        self.assertIs(last_dispatch_after_zeros, expected_dispatch)
+
+    @parameterized.expand([(1,), (2,)])
+    def test_trivially_replicable_op(self, shard_count: int):
+        self.make_tensor(shard_count).abs()
+
+        last_dispatch = ops._registry._test_get_last_op_dispatch()
+        self.assertIs(last_dispatch, abs_default)
+
+    @parameterized.expand([(1,), (2,)])
+    def test_multiple_dispatch(self, shard_count: int):
+        tensor = self.make_tensor(shard_count)
+
+        tensor.abs()
+        last_dispatch_after_abs = ops._registry._test_get_last_op_dispatch()
+        self.assertIs(last_dispatch_after_abs, abs_default)
+
+        tensor.cos()
+        last_dispatch_after_cos = ops._registry._test_get_last_op_dispatch()
+        self.assertIs(last_dispatch_after_cos, cos_default)
+
+
+if __name__ == "__main__":
+    unittest.main()