.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "programming-guide/tutorials/02-fused-softmax.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

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

        :ref:`Go to the end <sphx_glr_download_programming-guide_tutorials_02-fused-softmax.py>`
        to download the full example code.

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

.. _sphx_glr_programming-guide_tutorials_02-fused-softmax.py:


Fused Softmax
=============

This tutorial demonstrates softmax operator and its launch on Ascend simulator.

.. GENERATED FROM PYTHON SOURCE LINES 15-84

.. code-block:: Python
   :lineno-start: 16


    import asc2
    import torch


    # The functions which are executed on Ascend NPU must be marked with `@asc2.jit` decorator.
    # Available parameters for @asc2.jit decorator can be seen in the documentation:
    # https://compiler-team-ru.github.io/pyasc/python-api/rst/runtime/index.html
    @asc2.jit(reuse_alloc=1)
    def fused_softmax(
            # Pointers to input and output tensors should have `asc2.GlobalAddress` type.
            input_ptr: asc2.GlobalAddress, output_ptr: asc2.GlobalAddress,
            # Scalar parameters are passed as Python types (e.g. `int`, `float`).
            # For optimization purposes it is recommended to pass scalar parameter as constants (e.g. `asc2.ConstExpr[int]`).
            # Here num_rows and num_cols are compiler-time parameter:
            num_rows: asc2.ConstExpr, num_cols: asc2.ConstExpr,
            # tile_shape is a list and must be asc2.ConstExpr
            tile_shape: asc2.ConstExpr):

        # Tensor descriptor is created from `asc2.GlobalAddress` to represent entire tensor.
        in_gm = asc2.global_tensor(input_ptr, [num_rows, num_cols])
        out_gm = asc2.global_tensor(output_ptr, [num_rows, num_cols])

        # Python expressions are used to calculate offset:
        # `asc2.block_num()` function provides number of AICOREs launched.
        rows_per_block = asc2.ceildiv(num_rows, asc2.block_num())
        # `asc2.block_idx()` function is used to get current AICORE index.
        block_offset = asc2.block_idx() * rows_per_block
        # Define the loop iterating over tiles
        ub_loop = asc2.number(asc2.ceildiv(rows_per_block, tile_shape[0]), asc2.int_)

        # `unroll_factor` parameter of `asc2.range` in `for` loop can be used to manage software pipelining. Set it to `2` to enable double buffering.
        # `parallel` parameter of `asc2.range` in `for` loop enable overlapping of store operation of `i`-th iteration and load of `i+1`-th iteration.
        # It is user responsibility to ensure that there are no data dependencies between overlapped iterations.
        for i in asc2.range(ub_loop, unroll_factor=2, parallel=True):
            row_start_offset = block_offset + i * tile_shape[0]
            # `asc2.copy_in` is used to create 2D tensor object to load from GM to UB and pad with '-inf' all values that are out of global tensor
            rows = asc2.copy_in(in_gm, [row_start_offset, 0], [tile_shape[0], tile_shape[1]], pad_value=float('-inf'))
            # Call high-level 2D softmax
            out = asc2.softmax(rows)
            # `asc2.copy_out` is used to move data from UB back to GM.
            asc2.copy_out(out, out_gm, [row_start_offset, 0])


    if __name__ == "__main__":
        backend = asc2.Backend.Model  # can be "Model" for simulator or "NPU" for device
        soc_version = asc2.Platform.Ascend950PR_9599  # Device version
        device_id = 0  # might be necessary to provide if more than one NPU device is present in the system
        asc2.set_platform(backend, soc_version, device_id)

        input_shape_2d = [256, 98]
        rows_per_iter = 5
        block_num = 56
        dtype = torch.float32

        # Alignment for tile_shape
        num_rows, num_cols = input_shape_2d
        ALIGNMENT_ELEMENTS = 32 // dtype.itemsize
        tile_shape = [rows_per_iter, asc2.ceildiv(num_cols, ALIGNMENT_ELEMENTS) * ALIGNMENT_ELEMENTS]

        # Allocate tensors
        in_tensor = torch.randn(input_shape_2d, dtype=dtype)
        out_tensor = torch.zeros(input_shape_2d, dtype=dtype)

        # For the kernel invocation, number of AICOREs should be provided in brackets:
        fused_softmax[block_num](in_tensor, out_tensor, input_shape_2d[0], input_shape_2d[1], tile_shape)

        expected = torch.softmax(in_tensor, dim=1)
        torch.testing.assert_close(out_tensor, expected)


.. _sphx_glr_download_programming-guide_tutorials_02-fused-softmax.py:

.. only:: html

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

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

      :download:`Download Jupyter notebook: 02-fused-softmax.ipynb <02-fused-softmax.ipynb>`

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

      :download:`Download Python source code: 02-fused-softmax.py <02-fused-softmax.py>`

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

      :download:`Download zipped: 02-fused-softmax.zip <02-fused-softmax.zip>`


.. only:: html

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

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