TensorRT-LLMs/tests/functional/test_gather_nd.py

# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed 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.
import os
import sys
import unittest

# isort: off
import torch
# isort: on
from parameterized import parameterized

import tensorrt_llm
from tensorrt_llm import Tensor

sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from utils.util import create_session, run_session


class TestGatherND(unittest.TestCase):

    def setUp(self):
        tensorrt_llm.logger.set_level('error')

    @parameterized.expand([
        (
            [
                [91, 92, 93, 95, 94, 96, 97, 00,
                 00],  # 7 effective tokens and 2 ignored.
                [93, 94, 95, 92, 95, 96, 93, 97, 96],  # 9 effective tokens
            ],
            [
                [
                    [0, 1, 2, 3],
                    [0, 1, 4, 5],
                    [0, 1, 2, 6],
                ],
                [
                    [0, 1, 2, 3],
                    [0, 4, 5, 6],
                    [0, 1, 7, 8],
                ],
            ],
            [  # Assuming a batch of two sequences, each has 3 beams of 4 tokens.
                [
                    [91, 92, 93, 95],
                    [91, 92, 94, 96],
                    [91, 92, 93, 97],
                ],
                [
                    [93, 94, 95, 92],
                    [93, 95, 96, 93],
                    [93, 94, 97, 96],
                ],
            ],
        ),
        ([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], [[1, 0, 1], [0,
                                                            1, 0]], [[[2, 3],
                                                                      [0, 1],
                                                                      [2, 3]],
                                                                     [[4, 5],
                                                                      [6, 7],
                                                                      [4, 5]]]),
        (
            torch.rand((2, 9, 4), dtype=torch.float32, device="cuda"),
            torch.tensor([[[0, 1, 2, 3, 4, 5], [0, 1, 3, 4, 5, 6],
                           [0, 1, 4, 5, 6, 7], [0, 2, 3, 4, 6, 8]],
                          [[0, 1, 2, 3, 4, 5], [0, 1, 3, 4, 5, 7],
                           [0, 2, 3, 5, 6, 7], [0, 3, 4, 5, 6, 7]]],
                         device="cuda"),
            [],
        ),
    ])
    def test_gatherND(self, data, indices, ref):
        dtype = "float32"
        data = data if isinstance(data,
                                  torch.Tensor) else torch.tensor(data).cuda()
        indices = indices if isinstance(
            indices, torch.Tensor) else torch.tensor(indices).cuda()
        ref = ref if isinstance(ref, torch.Tensor) else torch.tensor(ref).cuda()
        indices = indices.unsqueeze(-1)  # needed for TRT gatherND

        # construct trt network
        builder = tensorrt_llm.Builder()
        network = builder.create_network()
        with tensorrt_llm.net_guard(network):

            d = Tensor(name='d',
                       shape=data.shape,
                       dtype=tensorrt_llm.torch_dtype_to_trt(data.dtype))
            idx = Tensor(name='idx',
                         shape=indices.shape,
                         dtype=tensorrt_llm.torch_dtype_to_trt(indices.dtype))

            output = tensorrt_llm.functional.gather_nd(d, idx, 1)
            output.mark_output('output')

        # trt run
        session = create_session(builder, network, precision=dtype)
        inputs = {'d': data, 'idx': indices}
        outputs = run_session(session, inputs)
        # compare diff
        indices = indices.squeeze(-1)
        tref = torch.stack([data[i, indices[i]] for i in range(data.shape[0])])
        if ref.numel() == 0:
            torch.testing.assert_close(tref, outputs['output'])
        else:
            torch.testing.assert_close(ref, outputs['output'])
            torch.testing.assert_close(ref, tref)

    @parameterized.expand([
        ([[91, 92, 93, 95, -1, -1, 94, 96, -1, -1, -1, 97],
          [93, 94, 95, 92, -1, 95, 96, 93, -1, -1, 97, 96]], [[0, 0], [0, 1],
                                                              [0, 2], [0, 3],
                                                              [0, 6], [0, 7],
                                                              [0, 11], [1, 0],
                                                              [1, 1], [1, 2],
                                                              [1, 3], [1, 5],
                                                              [1, 6], [1, 7],
                                                              [1, 10], [1, 11]],
         [91, 92, 93, 95, 94, 96, 97, 93, 94, 95, 92, 95, 96, 93, 97, 96])
    ])
    def test_gatherND_b0(self, data, indices, ref):
        dtype = "float32"
        data = data if isinstance(data,
                                  torch.Tensor) else torch.tensor(data).cuda()
        indices = indices if isinstance(
            indices, torch.Tensor) else torch.tensor(indices).cuda()
        ref = ref if isinstance(ref, torch.Tensor) else torch.tensor(ref).cuda()

        # construct trt network
        builder = tensorrt_llm.Builder()
        network = builder.create_network()
        with tensorrt_llm.net_guard(network):
            d = Tensor(name='d',
                       shape=data.shape,
                       dtype=tensorrt_llm.torch_dtype_to_trt(data.dtype))
            idx = Tensor(name='idx',
                         shape=indices.shape,
                         dtype=tensorrt_llm.torch_dtype_to_trt(indices.dtype))

            output = tensorrt_llm.functional.gather_nd(d, idx, 0)
            output.mark_output('output')

        # trt run
        session = create_session(builder, network, precision=dtype)
        inputs = {'d': data, 'idx': indices}
        outputs = run_session(session, inputs)

        # compare diff
        tref = data[indices[:, 0], indices[:, 1]]
        if ref.numel() == 0:
            torch.testing.assert_close(tref, outputs['output'])
        else:
            torch.testing.assert_close(ref, outputs['output'])
            torch.testing.assert_close(ref, tref)

    def test_gatherND_selectH(self):
        dtype = "float32"
        # This usecase is used to gather for validated end-tokens (diff stopping point for diff seqs)
        data = torch.rand((2, 9, 4), dtype=torch.float32, device="cuda")
        indices = torch.randint(9, size=(2, ), dtype=torch.int32, device="cuda")
        indices = torch.stack(
            [torch.arange(2, dtype=torch.int32).cuda(), indices], dim=1)

        # construct trt network
        builder = tensorrt_llm.Builder()
        network = builder.create_network()
        with tensorrt_llm.net_guard(network):
            d = Tensor(name='d',
                       shape=data.shape,
                       dtype=tensorrt_llm.torch_dtype_to_trt(data.dtype))
            idx = Tensor(name='idx',
                         shape=indices.shape,
                         dtype=tensorrt_llm.torch_dtype_to_trt(indices.dtype))

            output = tensorrt_llm.functional.gather_nd(d, idx, 0)
            output.mark_output('output')

        # trt run
        session = create_session(builder, network, precision=dtype)
        inputs = {'d': data, 'idx': indices}
        outputs = run_session(session, inputs)

        # pytorch run
        ref = data[indices[:, 0], indices[:, 1]]

        # compare diff
        torch.testing.assert_close(ref, outputs['output'])