Note

Go to the end to download the full example code.

Matrix Multiplication

This tutorial demonstrates matrix multiplication 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
24 def matrix_multiplication(
25         # Pointers to input and output tensors should have `asc2.GlobalAddress` type.
26         a_ptr: asc2.GlobalAddress, b_ptr: asc2.GlobalAddress, c_ptr: asc2.GlobalAddress,
27         # Input shapes for A and B matrices
28         a_shape: asc2.ConstExpr, b_shape: asc2.ConstExpr,
29         # For optimization purposes it is recommended to pass scalar parameter as constants (e.g. `asc2.ConstExpr[int]`).
30         single_core_m: asc2.ConstExpr, single_core_n: asc2.ConstExpr, step_ka: asc2.ConstExpr, step_kb: asc2.ConstExpr,
31         base_k: asc2.ConstExpr):
32     m, k = a_shape
33     _, n = b_shape
34     # Tensor descriptor is created from `asc2.GlobalAddress` to represent entire tensor.
35     a_gm = asc2.global_tensor(a_ptr, a_shape)
36     b_gm = asc2.global_tensor(b_ptr, b_shape)
37     c_gm = asc2.global_tensor(c_ptr, [m, n])
38
39     # Create 2D tensor for cube L0C accumulator
40     acc = asc2.zeros_acc([single_core_m, single_core_n], dtype=asc2.float32)
41
42     # Python expressions are used to calculate block offset:
43     # - `asc2.block_idx()` function is used to get current AICORE index.
44     n_blocks = asc2.ceildiv(n, single_core_n)
45     m_off = single_core_m * (asc2.block_idx() / n_blocks)
46     n_off = single_core_n * (asc2.block_idx() % n_blocks)
47
48     # `unroll_factor` parameter of `asc2.range` in `for` loop can be used to manage software pipelining. Set it to `2` to enable double buffering.
49     # `parallel` parameter of `asc2.range` in `for` loop enable overlapping of store operation of `i`-th iteration and load of `i+1`-th iteration.
50     # It is user responsibility to ensure that there are no data dependencies between overlapped iterations.
51     for k_outer in range(asc2.ceildiv(k, step_kb), unroll_factor=2, parallel=True):
52         # Load B matrix from GM to a new local tensor in L1
53         b_l1 = asc2.copy_in(b_gm, [k_outer * step_kb, n_off], [step_kb, single_core_n], asc2.TensorLocation.L1)
54         for k_mid in range(asc2.ceildiv(step_kb, step_ka), unroll_factor=2, parallel=True):
55             k_off = k_outer * step_kb + k_mid * step_ka
56             # Load A matrix from GM to a new local tensor in L1
57             a_l1 = asc2.copy_in(a_gm, [m_off, k_off], [single_core_m, step_ka], asc2.TensorLocation.L1)
58             for k_l0 in range(asc2.ceildiv(step_ka, base_k), unroll_factor=2, parallel=True):
59                 # Copy A matrix from L1 to a new local tensor in L0A
60                 a_l0 = asc2.copy(a_l1, [0, k_l0 * base_k], [single_core_m, base_k], asc2.TensorLocation.L0A)
61                 # Copy B matrix from L1 to a new local tensor in L0B
62                 b_l0 = asc2.copy(b_l1, [k_mid * step_ka + k_l0 * base_k, 0], [base_k, single_core_n],
63                                  asc2.TensorLocation.L0B)
64                 # Perform matrix multiplication with updating accumulator
65                 asc2.matmul_acc(acc, a_l0, b_l0)
66
67     # `asc2.copy_out` is used to move data from L0C accumulator back to GM.
68     asc2.copy_out(acc, c_gm, [m_off, n_off])
69
70
71 if __name__ == "__main__":
72     backend = asc2.Backend.Model  # can be "Model" for simulator or "NPU" for device
73     soc_version = asc2.Platform.Ascend950PR_9599  # Device version
74     device_id = 0  # might be necessary to provide if more than one NPU device is present in the system
75     asc2.set_platform(backend, soc_version, device_id)
76
77     block_num = 16
78     dtype = torch.float32
79
80     m, k, n = 1024, 64, 16
81     single_core_m, single_core_n = 64, 16
82     step_ka, step_kb, base_k = 16, 64, 16
83
84     # Allocate tensors for A, B and result matrix C
85     a = torch.rand((m, k), dtype=dtype)
86     b = torch.rand((k, n), dtype=dtype)
87     c = torch.zeros((m, n), dtype=dtype)
88
89     matrix_multiplication[block_num](a, b, c, a.shape, b.shape, single_core_m, single_core_n, step_ka, step_kb, base_k)
90     c_ref = a @ b
91     torch.testing.assert_close(c, c_ref)

Gallery generated by Sphinx-Gallery