Note

Go to the end to download the full example code.

Fused Softmax

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

16 import asc2
17 import torch
18
19
20 # The functions which are executed on Ascend NPU must be marked with `@asc2.jit` decorator.
21 # Available parameters for @asc2.jit decorator can be seen in the documentation:
22 # https://compiler-team-ru.github.io/pyasc/python-api/rst/runtime/index.html
23 @asc2.jit(reuse_alloc=1)
24 def fused_softmax(
25         # Pointers to input and output tensors should have `asc2.GlobalAddress` type.
26         input_ptr: asc2.GlobalAddress, output_ptr: asc2.GlobalAddress,
27         # Scalar parameters are passed as Python types (e.g. `int`, `float`).
28         # For optimization purposes it is recommended to pass scalar parameter as constants (e.g. `asc2.ConstExpr[int]`).
29         # Here num_rows and num_cols are compiler-time parameter:
30         num_rows: asc2.ConstExpr, num_cols: asc2.ConstExpr,
31         # tile_shape is a list and must be asc2.ConstExpr
32         tile_shape: asc2.ConstExpr):
33
34     # Tensor descriptor is created from `asc2.GlobalAddress` to represent entire tensor.
35     in_gm = asc2.global_tensor(input_ptr, [num_rows, num_cols])
36     out_gm = asc2.global_tensor(output_ptr, [num_rows, num_cols])
37
38     # Python expressions are used to calculate offset:
39     # `asc2.block_num()` function provides number of AICOREs launched.
40     rows_per_block = asc2.ceildiv(num_rows, asc2.block_num())
41     # `asc2.block_idx()` function is used to get current AICORE index.
42     block_offset = asc2.block_idx() * rows_per_block
43     # Define the loop iterating over tiles
44     ub_loop = asc2.number(asc2.ceildiv(rows_per_block, tile_shape[0]), asc2.int_)
45
46     # `unroll_factor` parameter of `asc2.range` in `for` loop can be used to manage software pipelining. Set it to `2` to enable double buffering.
47     # `parallel` parameter of `asc2.range` in `for` loop enable overlapping of store operation of `i`-th iteration and load of `i+1`-th iteration.
48     # It is user responsibility to ensure that there are no data dependencies between overlapped iterations.
49     for i in asc2.range(ub_loop, unroll_factor=2, parallel=True):
50         row_start_offset = block_offset + i * tile_shape[0]
51         # `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
52         rows = asc2.copy_in(in_gm, [row_start_offset, 0], [tile_shape[0], tile_shape[1]], pad_value=float('-inf'))
53         # Call high-level 2D softmax
54         out = asc2.softmax(rows)
55         # `asc2.copy_out` is used to move data from UB back to GM.
56         asc2.copy_out(out, out_gm, [row_start_offset, 0])
57
58
59 if __name__ == "__main__":
60     backend = asc2.Backend.Model  # can be "Model" for simulator or "NPU" for device
61     soc_version = asc2.Platform.Ascend950PR_9599  # Device version
62     device_id = 0  # might be necessary to provide if more than one NPU device is present in the system
63     asc2.set_platform(backend, soc_version, device_id)
64
65     input_shape_2d = [256, 98]
66     rows_per_iter = 5
67     block_num = 56
68     dtype = torch.float32
69
70     # Alignment for tile_shape
71     num_rows, num_cols = input_shape_2d
72     ALIGNMENT_ELEMENTS = 32 // dtype.itemsize
73     tile_shape = [rows_per_iter, asc2.ceildiv(num_cols, ALIGNMENT_ELEMENTS) * ALIGNMENT_ELEMENTS]
74
75     # Allocate tensors
76     in_tensor = torch.randn(input_shape_2d, dtype=dtype)
77     out_tensor = torch.zeros(input_shape_2d, dtype=dtype)
78
79     # For the kernel invocation, number of AICOREs should be provided in brackets:
80     fused_softmax[block_num](in_tensor, out_tensor, input_shape_2d[0], input_shape_2d[1], tile_shape)
81
82     expected = torch.softmax(in_tensor, dim=1)
83     torch.testing.assert_close(out_tensor, expected)

Gallery generated by Sphinx-Gallery