.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/sum.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_sum.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_sum.py:


Sum Reduction Example
================

This example demonstrates how to implement a sum reduction operation along the last dimension using Helion.

.. GENERATED FROM PYTHON SOURCE LINES 9-11

Imports
-------

.. GENERATED FROM PYTHON SOURCE LINES 11-20

.. code-block:: Python

    from __future__ import annotations

    import torch

    import helion
    from helion._testing import run_example
    import helion.language as hl



.. GENERATED FROM PYTHON SOURCE LINES 21-23

Sum Kernel
--------

.. GENERATED FROM PYTHON SOURCE LINES 23-43

.. code-block:: Python

    @helion.kernel()
    def sum_kernel(x: torch.Tensor) -> torch.Tensor:
        """
        Sums a 2D tensor along the last dimension.

        Args:
            x: Input tensor of shape [M, N]

        Returns:
            Output tensor of shape [M] containing the sum of each row
        """
        m, n = x.shape
        out = torch.empty([m], dtype=x.dtype, device=x.device)

        for tile_m in hl.tile(m):
            out[tile_m] = x[tile_m, :].sum(-1)

        return out



.. GENERATED FROM PYTHON SOURCE LINES 44-46

Benchmark Wrapper
--------------

.. GENERATED FROM PYTHON SOURCE LINES 46-64

.. code-block:: Python

    def sum_tritonbench(x: torch.Tensor) -> torch.Tensor:
        """
        Wrapper for tritonbench that handles 1D input.

        Args:
            x: Input tensor (1D or 2D)

        Returns:
            Sum of the tensor along the last dimension
        """
        if x.ndim == 1:
            # For 1D tensors, reshape to 2D for sum_kernel
            x_2d = x.unsqueeze(0)
            result = sum_kernel(x_2d)
            return result.squeeze()
        return sum_kernel(x)



.. GENERATED FROM PYTHON SOURCE LINES 65-67

Verification Function
-------------------

.. GENERATED FROM PYTHON SOURCE LINES 67-80

.. code-block:: Python

    def check(m: int, n: int) -> None:
        """
        Verify the sum kernel implementation against PyTorch's native sum function.

        Args:
            m: First dimension of the test tensor
            n: Second dimension of the test tensor
        """
        x = torch.randn([m, n], device="cuda", dtype=torch.float32)
        kernels = {"helion": sum_kernel}
        run_example(kernels, lambda x: x.sum(-1), (x,))



.. GENERATED FROM PYTHON SOURCE LINES 81-83

Main Function
-----------

.. GENERATED FROM PYTHON SOURCE LINES 83-97

.. code-block:: Python

    def main() -> None:
        """
        Main entry point that runs the sum kernel verification with different tensor sizes.

        Tests with two configurations:
        - 512x256
        - 1024x1024
        """
        check(512, 256)
        check(1024, 1024)


    if __name__ == "__main__":
        main()


.. _sphx_glr_download_examples_sum.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: sum.ipynb <sum.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: sum.py <sum.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: sum.zip <sum.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_