[Fix][Relax] Dispatch median sorting on GPU

python/tvm/relax/backend/dispatch_sort_scan.py

@@ -22,7 +22,7 @@
 
 import tvm_ffi
 
-from tvm import relax, topi
+from tvm import relax, te, topi
 from tvm.contrib.thrust import can_use_thrust
 from tvm.relax import expr_functor
 from tvm_tirx.ir import GlobalVar, Op
@@ -33,6 +33,35 @@
 from .utils import BackendDispatcher
 
 
+def _te_median_from_sorted(values, indices, axis, keepdims, flattened_rank):
+    """Select the lower middle value and index from sorted tensors."""
+    if flattened_rank is not None:
+        out_shape = [1] * flattened_rank if keepdims else []
+    elif keepdims:
+        out_shape = [1 if i == axis else dim for i, dim in enumerate(values.shape)]
+    else:
+        out_shape = [dim for i, dim in enumerate(values.shape) if i != axis]
+    middle = (values.shape[axis] - 1) // 2
+
+    def coordinates(out):
+        if flattened_rank is not None:
+            return (middle,)
+        coord = list(out)
+        if keepdims:
+            coord[axis] = middle
+        else:
+            coord.insert(axis, middle)
+        return tuple(coord)
+
+    result = te.compute(out_shape, lambda *out: values[coordinates(out)], name="median_value")
+    if flattened_rank is not None:
+        return result
+    result_index = te.compute(
+        out_shape, lambda *out: indices[coordinates(out)], name="median_index"
+    )
+    return result, result_index
+
+
 @expr_functor.mutator
 class SortScanDispatcher(BackendDispatcher):
     """Dispatcher to dispatch sort and scan."""
@@ -121,6 +150,58 @@ def visit_call_(self, call: relax.Call) -> relax.Expr:
                 dtype=call.attrs.dtype,
                 **kwargs,
             )
+        if call.op.name == "relax.median":
+            try:
+                tgt = self._get_target(call.ty)
+            except ValueError:
+                # A target-less dispatch previously left median unchanged.
+                return super().visit_call_(call)
+            axes = call.attrs.axis
+            shape = call.args[0].ty.shape
+            if not self.is_gpu_target(tgt) or not isinstance(shape, relax.ShapeExpr):
+                return super().visit_call_(call)
+            if axes is not None and len(axes) != 1:
+                return super().visit_call_(call)
+
+            flattened_rank = None
+            data = call.args[0]
+            if axes is None:
+                flattened_rank = len(shape.values)
+                size = reduce(mul, shape.values, 1)
+                data = self.builder_.emit(relax.op.reshape(data, relax.ShapeExpr([size])))
+                axis = 0
+            else:
+                axis = int(axes[0]) % len(shape.values)
+
+            kwargs = {}
+            te_sort = topi.gpu.argsort
+            if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
+                te_sort = topi.gpu.argsort_thrust
+                kwargs["workspace"] = self.allocate_workspace(call)
+            with tgt:
+                sorted_pair = self.builder_.emit(
+                    self.builder_.call_te(
+                        te_sort,
+                        data,
+                        axis=axis,
+                        is_ascend=True,
+                        dtype="int64",
+                        ret_type="both",
+                        **kwargs,
+                    )
+                )
+                sorted_values = self.builder_.emit(relax.TupleGetItem(sorted_pair, 0))
+                sorted_indices = self.builder_.emit(relax.TupleGetItem(sorted_pair, 1))
+                result = self.builder_.call_te(
+                    _te_median_from_sorted,
+                    sorted_values,
+                    sorted_indices,
+                    axis,
+                    bool(call.attrs.keepdims),
+                    flattened_rank,
+                )
+            self._append_calls_to_update(result, tgt)
+            return result
         if call.op.name == "relax.topk":
             tgt = self._get_target(call.ty)
             te_func = topi.topk

tests/python/relax/test_backend_dispatch_median.py

@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for GPU median dispatch."""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm_tirx
+from tvm import relax
+from tvm.testing import device_enabled
+
+
+def make_median(shape, dtype, axis, keepdims):
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", relax.TensorType(shape, dtype))
+    with bb.function("main", [x]):
+        with bb.dataflow():
+            out = bb.emit(relax.op.median(x, axis=axis, keepdims=keepdims))
+            out = bb.emit_output(out)
+        bb.emit_func_output(out)
+    return bb.get()
+
+
+@pytest.mark.parametrize("axis", [None, -1, 0])
+@pytest.mark.parametrize("keepdims", [False, True])
+def test_dispatch_median_cuda_type(axis, keepdims):
+    mod = make_median((2, 3), "float32", axis, keepdims)
+    with tvm_tirx.target.Target({"kind": "cuda", "arch": "sm_80"}):
+        dispatched = relax.backend.DispatchSortScan()(mod)
+    tvm_tirx.ir.assert_structural_equal(mod["main"].ret_ty, dispatched["main"].ret_ty)
+    assert "R.median(" not in dispatched.script()
+
+
+def test_dispatch_median_cpu_unchanged():
+    mod = make_median((2, 3), "float32", 1, False)
+    with tvm_tirx.target.Target("llvm"):
+        dispatched = relax.backend.DispatchSortScan()(mod)
+    tvm_tirx.ir.assert_structural_equal(mod, dispatched)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm_tirx.runtime.enabled("llvm") or not device_enabled("cuda"),
+    reason="CUDA and LLVM are required",
+)
+@pytest.mark.parametrize(
+    "shape,axis,keepdims,dtype",
+    [
+        ((5,), None, False, "float32"),
+        ((2, 3), None, True, "float32"),
+        ((3, 4), -1, False, "float32"),
+        ((3, 4), 0, True, "float32"),
+        ((2, 3, 4), 1, False, "float16"),
+        ((2, 65), -1, True, "int32"),
+        ((), None, True, "float32"),
+    ],
+)
+def test_median_cuda_values_and_indices(shape, axis, keepdims, dtype):
+    rng = np.random.default_rng(20261104)
+    size = int(np.prod(shape))
+    data = rng.permutation(size).astype(dtype).reshape(shape)
+    mod = make_median(shape, dtype, axis, keepdims)
+    dev = tvm_tirx.cuda(0)
+    arch = "sm_" + dev.compute_version.replace(".", "")
+    target = tvm_tirx.target.Target(
+        {"kind": "cuda", "arch": arch, "host": {"kind": "llvm", "mcpu": "generic"}}
+    )
+    with target:
+        exe = tvm.compile(mod, target=target)
+    vm = relax.VirtualMachine(exe, dev)
+    result = vm["main"](tvm_tirx.runtime.tensor(data, dev))
+    if axis is None:
+        expected = np.sort(data.ravel())[(data.size - 1) // 2]
+        expected = np.asarray(expected, dtype=dtype).reshape([1] * len(shape) if keepdims else ())
+        actual = result.numpy()
+        assert actual.shape == expected.shape
+        np.testing.assert_array_equal(actual, expected)
+    else:
+        sorted_indices = np.argsort(data, axis=axis, kind="stable")
+        expected_index = np.take(sorted_indices, (shape[axis] - 1) // 2, axis=axis)
+        if keepdims:
+            expected_index = np.expand_dims(expected_index, axis=axis)
+        gather_index = expected_index if keepdims else np.expand_dims(expected_index, axis=axis)
+        expected_value = np.take_along_axis(data, gather_index, axis=axis)
+        if not keepdims:
+            expected_value = np.squeeze(expected_value, axis=axis)
+        actual_value, actual_index = result[0].numpy(), result[1].numpy()
+        assert actual_value.shape == expected_value.shape
+        assert actual_index.shape == expected_index.shape
+        np.testing.assert_array_equal(actual_value, expected_value)
+        np.testing.assert_array_equal(actual_index, expected_index)
+    dev.sync()
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(
+    not tvm_tirx.runtime.enabled("llvm") or not device_enabled("cuda"),
+    reason="CUDA and LLVM are required",
+)
+@pytest.mark.parametrize("axis", [None, 0])
+def test_median_cuda_symbolic_rows(axis):
+    n = tvm_tirx.tirx.Var("n", "int64")
+    mod = make_median((n, 4), "float32", axis, True)
+    dev = tvm_tirx.cuda(0)
+    target = tvm_tirx.target.Target(
+        {
+            "kind": "cuda",
+            "arch": "sm_" + dev.compute_version.replace(".", ""),
+            "host": {"kind": "llvm", "mcpu": "generic"},
+        }
+    )
+    with target:
+        exe = tvm.compile(mod, target=target)
+    vm = relax.VirtualMachine(exe, dev)
+    for rows in (3, 5):
+        data = np.arange(rows * 4, dtype="float32").reshape(rows, 4)[::-1].copy()
+        result = vm["main"](tvm_tirx.runtime.tensor(data, dev))
+        if axis is None:
+            expected = np.sort(data.ravel())[(data.size - 1) // 2].reshape((1, 1))
+            actual = result.numpy()
+            assert actual.shape == expected.shape
+            np.testing.assert_array_equal(actual, expected)
+        else:
+            order = np.argsort(data, axis=0, kind="stable")
+            expected_index = order[(rows - 1) // 2 : (rows + 1) // 2]
+            expected_value = np.take_along_axis(data, expected_index, axis=0)
+            actual_value, actual_index = result[0].numpy(), result[1].numpy()
+            assert actual_value.shape == expected_value.shape
+            assert actual_index.shape == expected_index.shape
+            np.testing.assert_array_equal(actual_value, expected_value)
+            np.testing.assert_array_equal(actual_index, expected_index)
+    dev.sync()
+
+
+def test_dispatch_median_without_target_unchanged():
+    mod = make_median((2, 3), "float32", None, True)
+    dispatched = relax.backend.DispatchSortScan()(mod)
+    tvm_tirx.ir.assert_structural_equal(mod, dispatched)