Skip to content

Commit ace4c73

Browse files
dkhaldidkhaldi
authored
[SYCL][Matrix] Add spec document for the matrix extension interface and its first implementation for AMX (#3551)
Signed-off-by: Dounia <dounia.khaldi@intel.com>
1 parent 4cd4016 commit ace4c73

File tree

3 files changed

+313
-0
lines changed

3 files changed

+313
-0
lines changed

‎sycl/doc/extensions/Matrix/README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
# Matrix Programming Extension for DPC++
2+
`matrix` is a new experimental DPC++ extension to provide unified matrix programming on different tensor hardware. The current implementation provides support of the matrix interface using Intel(R) Advanced Matrix Extensions (AMX).

‎sycl/doc/extensions/Matrix/dpcpp-joint-matrix.asciidoc

Lines changed: 310 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,310 @@
1+
# Matrix Programming Extension for DPC++: SYCL_EXT_ONEAPI_MATRIX
2+
:source-highlighter: coderay
3+
:coderay-linenums-mode: table
4+
:dpcpp: pass:[DPC++]
5+
6+
// This section needs to be after the document title.
7+
:doctype: book
8+
:toc2:
9+
:toc: left
10+
:encoding: utf-8
11+
:lang: en
12+
13+
:blank: pass:[ +]
14+
15+
// Set the default source code type in this document to C++,
16+
// for syntax highlighting purposes. This is needed because
17+
// docbook uses c++ and html5 uses cpp.
18+
:language: {basebackend@docbook:c++:cpp}
19+
20+
21+
== Notice
22+
23+
Copyright (c) 2021-2021 Intel Corporation. All rights reserved.
24+
25+
NOTE: Khronos(R) is a registered trademark and SYCL(TM) and SPIR(TM) are
26+
trademarks of The Khronos Group Inc. OpenCL(TM) is a trademark of Apple Inc.
27+
used by permission by Khronos.
28+
29+
This extension is written against the SYCL 2020 revision 3 specification. All
30+
references below to the "core SYCL specification" or to section numbers in the
31+
SYCL specification refer to that revision.
32+
33+
34+
**_NOTE:_** _This document describes the current design and API for the matrix
35+
extension to {dpcpp}. This is an initial experimental version to try out functionality
36+
and performance, and **future versions of this API may change in ways that are incompatible with this experimental version**. The current implementation provides support of the matrix interface on Intel(R) Advanced Matrix Extensions (AMX). We are going to work with the community on incrementally improving
37+
the API to bring them closer to standard C++ (aligned with the `std::mdspan` and `std::mdarray` proposals) and SYCL in the next several months._
38+
39+
## Introduction
40+
This document presents an ongoing work towards defining a unified matrix interface. This interface is intended to unify different tensor hardware: AMX in Intel CPU, Habana Gaudi and Goya tensor and gemm cores, Nvidia TPUs, IBM Power MMA. All these hardware provide low-level intrinsics or assembly to access and perform matrix operations. The goal is to provide a unified interface that is portable but also benefit from the maximum performance these different hardware can offer.
41+
42+
## Feature test macro
43+
44+
This extension provides a feature-test macro as described in the core SYCL
45+
specification section 6.3.3 "Feature test macros". Therefore, an
46+
implementation supporting this extension must predefine the macro
47+
`SYCL_EXT_ONEAPI_MATRIX` to one of the values defined in the table below.
48+
Applications can test for the existence of this macro to determine if the
49+
implementation supports this feature, or applications can test the macro's
50+
value to determine which of the extension's APIs the implementation supports.
51+
52+
[frame="none",options="header"]
53+
|======================
54+
|Value |Description
55+
|1 |Initial extension implementation on AMX. Base features are supported.
56+
|======================
57+
58+
## New `joint_matrix` class
59+
We introduce a new class called `joint_matrix`. The user needs to specify the type of the elements, shape, the memory layout, and the memory scope of the matrix. This results into the following description:
60+
61+
```c++
62+
namespace sycl::ext::intel::experimental::matrix {
63+
template <typename Group, typename T, size_t Rows=sycl::dynamic_extent, size_t Cols=sycl::dynamic_extent, matrix_layout Layout = matrix_layout::row_major>
64+
struct joint_matrix {
65+
joint_matrix(Group g) {}
66+
};
67+
}
68+
```
69+
70+
71+
#### Memory Scope
72+
In this experimental API version, we used the terminology of `joint_matrix` instead of plain `matrix` to emphasis that the matrix is shared among a group of work items and is not private to each work item. The memory scope is added as an additional template parameter and is also part of the constructor arguments.
73+
74+
IMPORTANT: In the current implementation, only the subgroup scope is supported
75+
76+
When the group is a `sycl::sub_group`, a matrix is declared as follows:
77+
78+
```c++
79+
joint_matrix<sub_group, int8_t, tM, tN> tA(sg);
80+
```
81+
82+
#### Shape
83+
The same class `joint_matrix` should handle both cases where sizes are constant (GPU case) and when sizes are variables (CPU case). Note that a AMX 2d tile register permits sizes up to 1024 (16rowsx64cols) bytes. The ability to define only one interface for both makes it possible to give the user a way to make use of the flexibility introduced by the CPU but at the same time save resources on the GPU. We use `sycl::dynamic_extent` to differentiate between static and dynamic sizes.
84+
85+
IMPORTANT: In the current implementation, only the static extent is supported
86+
87+
88+
#### Layout
89+
Besides row major and column major layouts, `matrix_layout` is flexible enough to introduce customed layouts such as symmetric or tiled layouts.
90+
91+
```c++
92+
namespace sycl::ext::intel::experimental::matrix {
93+
enum class matrix_layout {
94+
row_major,
95+
col_major,
96+
packed_a,
97+
packed_b
98+
};
99+
}
100+
```
101+
102+
AMX hardware requires B matrix to be in VNNI or 32 bits packed layout. If we multiply matrices A (M, K) and B (K, N) into a matrix C (M, N). The logical sizes are M, K, N. However, the packed shape for B tile uses the VNNI format, which is described below. The user must provide the information of packed_b layout to make the implementation allocate the right shape. The layout information for AMX should be specified in user code as follows:
103+
104+
```c++
105+
joint_matrix<sub_group, int8_t, K, N, packed_b> tB(sg);
106+
```
107+
IMPORTANT: In the current implementation, only `packed_b` layout is necessary to specify on matrix B, the layout on other matrices is ignored.
108+
109+
110+
111+
## Matrix Operations and their Execution Scope
112+
We define three new functions needed to perform the main and common operations on matrices namely, load, store, and the actual multiply and add operation. This set of functions can be easily extended if the tensor hardware implements new features.
113+
114+
The base pointer determines the starting address of the matrix to be loaded/stored. `layout` determines whether the data are being read/written in a row (`row_major`), column major (`column_major`) fashion, or if the data has already been transformed into VNNI format (`packed_a`, `packed_b`). `stride` describes the number of elements between consecutive rows for row major and packed layout, columns for column major layout.
115+
116+
Note that for getting maximum performance on AMX, prepacking data in the memory is necessary. If users did not specify the packed layouts (`packed_a` in column major case, `packed_b` in row major case), transforms done by the implementation will be slow due to extra scatter/gather operations. Hence, we expose these layouts `packed_a` and `packed_b` to the user to specify that A and/or B have already been VNNIed. The packed or VNNI layout is introduced in `VNNI layout` section below.
117+
118+
IMPORTANT: In the current implementation, the layout in the load of matrix B must be `packed_b`. Therefore, both the template parameter for the declaration of the B matrix and the call to `joint_matrix_load` for the B matrix must specify the `packed_b` layout. The layout in the load of matrices A and C must be `row_major`, and the layout in the store of matrix C must also be `row_major`.
119+
120+
Since the matrix functions are group operations (as defined in Section 4.17.3 of the SYCL specification), the matrix API has to be accessed by all the work-items in the group in a convergent control flow. The `Group` template argument can be a work-group or a subgroup. These functions will be called once by each work item in the group.
121+
122+
To be aligned with the SYCL 2020 group algorithms, an additional group argument is added to the matrix operations to designate that these functions are collective operations. The {dpcpp} syntax is the following:
123+
124+
IMPORTANT: In the current implementation, only the subgroup scope is supported. Moreover, a kernel using this extension must be decorated with the [[sycl::reqd_sub_group_size(1)]] attribute.
125+
126+
#### Load
127+
```c++
128+
namespace sycl::ext::intel::experimental::matrix {
129+
template <typename Group, typename T, size_t NumRows, size_t NumCols,
130+
matrix_layout Layout,
131+
access::address_space Space>
132+
void joint_matrix_load(Group sg, joint_matrix<Group, T, NumRows, NumCols, Layout> &res,
133+
multi_ptr<T, Space> src, size_t stride, matrix_layout layout = matrix_layout::row_major);
134+
}
135+
```
136+
This function loads data from memory to the 2d tiles of AMX that is a 2d storage.
137+
138+
139+
#### Store
140+
```c++
141+
namespace sycl::ext::intel::experimental::matrix {
142+
template <typename Group, typename T, size_t NumRows, size_t NumCols,
143+
matrix_layout Layout,
144+
access::address_space Space>
145+
void joint_matrix_store(Group sg, joint_matrix<Group, T, NumRows, NumCols, Layout> &res,
146+
multi_ptr<T, Space> src, size_t stride, matrix_layout layout = matrix_layout::row_major);
147+
}
148+
```
149+
This function stores the data from the 2d tiles back to memory.
150+
151+
#### Multiply and Add
152+
153+
```c++
154+
namespace sycl::ext::intel::experimental::matrix {
155+
template <typename Group, typename T1, typename T2, std::size_t M,
156+
std::size_t K, std::size_t N,
157+
matrix_layout LayoutA, matrix_layout LayoutB,
158+
matrix_layout LayoutC>
159+
joint_matrix<Group, T2, M, N, LayoutC> joint_matrix_mad(Group sg, joint_matrix<Group, T1, M, K, LayoutA> A,
160+
joint_matrix<Group, T1, K, N, LayoutB> B, joint_matrix<Group, T2, M, N, LayoutC> C);
161+
}
162+
```
163+
The matrix multiply and add function performs the multiply operation on the matrices `A` and `B`, accumulate the result with `C` and return the result.
164+
165+
166+
## VNNI/Packed Layout
167+
AMX compute assumes register for B tile (src1) to be in VNNI format as they need 32bit of K-data in A and B to be contiguous in memory.
168+
The VNNI blocking factor is 2 in the case of 16-bit types, and it is 4 in the case of 8-bit types. While the current implementation assumes that the matrix has been already packed by the user for performance reasons, the layout information is needed to inform the implementation about this transform. The following example illustrates how a matrix in `row_major` layout is transformed into the `packed_b` layout for a 16-bit type.
169+
170+
#### Example 1: 16-bit elements
171+
// Example of a 4 row x 4 column matrix using a 16-bit data element, in row-major layout.
172+
// Element a1 is contiguous in memory with element b1, etc.
173+
// ---------------------------------
174+
// a1, b1, c1, d1
175+
// a2, b2, c2, d2
176+
// a3, b3, c3, d3
177+
// a4, b4, c4, d4
178+
// ---------------------------------
179+
// The same matrix reformatted in packed_b layout.
180+
// Here, packing of 2 elements is needed to form 32 bits.
181+
// Element a1 is contiguous in memory with element a2, etc.
182+
// ---------------------------------
183+
// a1, a2, b1, b2, c1, c2, d1, d2
184+
// a3, a4, b3, b4, c3, c4, d3, d4
185+
186+
#### Example 2: 8-bit elements
187+
188+
// Example of a 4 row x 4 column matrix using a 8-bit data element, in row-major layout.
189+
// Element a1 is contiguous in memory with element b1, etc.
190+
// ---------------------------------
191+
// a1, b1, c1, d1
192+
// a2, b2, c2, d2
193+
// a3, b3, c3, d3
194+
// a4, b4, c4, d4
195+
// ---------------------------------
196+
// The same matrix reformatted in packed_b layout.
197+
// Here, packing of 4 elements is needed to form 32 bits.
198+
// Elements a1, a2, a3, a4 are contiguous in memory, etc.
199+
// ---------------------------------
200+
// a1, a2, a3, a4, b1, b2, b3, b4, c1, c2, c3, c4, d1, d2, d3, d4
201+
202+
203+
## Example using int8_t type
204+
```c++
205+
using namespace sycl::ext::intel::experimental::matrix;
206+
207+
queue q;
208+
range<2> G = {M, N};
209+
// For this first implementation, SG_SIZE has to be equal to one
210+
range<2> L = {1, SG_SIZE};
211+
int8_t *memA = malloc_shared<int8_t>(M*K, q);
212+
int8_t *memB = malloc_shared<int8_t>(K*N, q);
213+
Int32_t *memC = malloc_shared<int32_t>(M*N, q);
214+
// Assuming memB has already been VNNIed
215+
q.parallel_for(nd_range<2>(G, L), [=](nd_item<2> item)
216+
[[sycl::reqd_sub_group_size(SG_SIZE)]] {
217+
const auto global_idx = item.get_global_id(0);
218+
const auto global_idy = item.get_global_id(1);
219+
const auto sg_startx = global_idx - item.get_local_id(0);
220+
const auto sg_starty = global_idy - item.get_local_id(1);
221+
sub_group sg = item.get_sub_group();
222+
joint_matrix<sub_group, int8_t, tM, tK> tA(sg);
223+
// For B, since current implementation does not support non packed layout,
224+
// users need to specify the updated VNNI sizes along with the packed_b layout
225+
joint_matrix<sub_group, int8_t, tK, tN, packed_b> tB(sg);
226+
joint_matrix<sub_group, int32_t, tM, tN> tC(sg);
227+
joint_matrix_load(sg, tC, memC + sg_startx * tM * N + sg_starty, N, matrix_layout::row_major);
228+
for (int k = 0; k < K; k += tk) {
229+
joint_matrix_load(sg, tA, memA + sg_startx * tM * K + k, K, matrix_layout::row_major);
230+
joint_matrix_load(sg, tB, memB + k * N + sg_starty, N, matrix_layout::packed_b); // VNNI
231+
tC = joint_matrix_mad(sg, tA, tB, tC);
232+
}
233+
joint_matrix_store(sg, tC, memC + sg_startx * tM * N + sg_starty, N, matrix_layout::row_major);
234+
}).wait();
235+
236+
```
237+
## Implementation Status
238+
For oneAPI release 3, an AOT implementation is available on the CPU device to targets AMX hardware. we are using AMX tile intrinsics to implement the matrix load and store operations. Since we are currently emitting AMX intrinsics directly, this only enables AOT compilation.
239+
240+
Currently, this is the compilation command line needed to invoke AMX unit of Sapphire Rapids CPU:
241+
242+
```c++
243+
clang++ -fsycl -march=sapphirerapids fsycl-targets="spir64_x86_64-uknown-linux-sycldevice" -O2 matmul-int8.cpp -o matmul-int8
244+
```
245+
246+
Please refer to the section "Future Implementation Work" that talks about the future unified SPIR-V path that will enable JIT compilation.
247+
248+
### Current Implementation Restrictions
249+
This section summarizes the specific features that this implementation supports. In future versions of this API and implementation, the expectation is to provide a query interface to guide the usage of this API.
250+
251+
#### Type, Sizes, and Layouts
252+
The types supported by this AMX implementation are restricted to the types that AMX hardware support. Although the AMX hardware supports 2d tiles with a maximum size of 16x64 bytes, this current implementation can handle any size. If the matrix size is bigger than 1024 bytes, it will be stored in memory rather than mapped to a 2d tile. Performance penalty may occur in this case. In order to get the best performance with this implementation, matrix sizes should be no larger than 16x64 bytes and B matrix should be already packed (put in VNNI format).
253+
254+
More specifically, the following operation C = A*B+C can be performed on AMX with this interface where:
255+
256+
A(int8, any-size, row_major), B(int8, any-size, packed_b), C(int32, any-size, row_major)
257+
258+
or
259+
260+
A(bf16, any-size, row_major), B(bf16, any-size, packed_b), C(float, any-size, row_major).
261+
262+
No other types or layouts are supported at this time.
263+
264+
#### Memory and Execution Scope
265+
This current implementation only considers a sub-group scope. However, the sub-group size has to be equal to one in this first implementation. In this case, a kernel using this extension must be decorated with the [[sycl::reqd_sub_group_size(1)]] attribute.
266+
267+
268+
## Future Implementation Work
269+
270+
### Unified LLVM IR and SPIRV JIT Enabling
271+
To enable JIT compilation, a unified matrix IR needs to be added. Currently, there is no matrix type in LLVM IR or SPIR-V. We are working towards adding a new matrix type in both LLVM IR and SPIR-V. This JIT enabling is expected to be part of a future compiler release.
272+
273+
#### LLVM IR Extension
274+
As a short-term solution, we are extending the https://llvm.org/docs/LangRef.html#llvm-matrix-transpose-intrinsic[existing LLVM IR matrix intrinsics] to include features like VNNI layout. The current matrix intrinsics use flattened vectors to represent the matrix. Therefore, we are exploring both adding matrix type to LLVM IR and also using MLIR `vector` dialect for this work.
275+
276+
#### SPIR-V Extension
277+
The current draft proposal can be found https://gitlab.devtools.intel.com/OpenCL/opencl-extension-drafts/-/blob/master/SPV_INTEL_matrix.asciidoc[here].
278+
We are adding translation from LLVM IR matrix to SPIR-V matrix and vice versa in the LLVM to SPIR-V translator tool.
279+
280+
## Future-looking API
281+
282+
283+
### Memory scope
284+
The current experimental API uses `joint_` semantics to define the memory scope of the matrix. The long term solution is to use the proposed https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/LocalMemory/SYCL_INTEL_local_memory.asciidoc[`group_local_memory` extension] to allocate the matrix in local memory associated with a SYCL group as shown in the example below.
285+
286+
287+
```c++
288+
multi_ptr<matrix<T>, address_space::local_space> tA_ptr = group_local_memory<matrix<sub_group, int8_t, tM, tN>>(sg);
289+
```
290+
We did not utilize this extension for this matrix API version because sub-group local memory is not yet well defined in {dpcpp}. Moreover, the representation of this notion in LLVM IR and SPIR-V is not clear yet.
291+
292+
293+
## Open Questions
294+
- Besides row, col major and packed (VNNI) layout, what are the additional layouts that should absolutely be added?
295+
- Are there alternative names for the `packed_a` and `packed_b` layouts that would be clearer to distinguish between the VNNI Layout in matrix A and VNNI layout in matrix B of a matrix multiply and add operation on AMX?
296+
- Ronan Keryell: "It would be interesting to investigate whether providing also member functions would simplify the API. Provide both so it is possible to use the best one for each use case, while waiting for https://en.wikipedia.org/wiki/Uniform_Function_Call_Syntax to land into C++?"
297+
- What should the API description include: (1) only features that are implemented, (2) features that are actually part of the API: currently implemented and the ones that we expect implementing them in the future. Specifically, should the document include things like dynamic_ extent and Group? These are part of the API but are not currently implemented.
298+
299+
## TODO List
300+
- Handle sub group sizes that are bigger than one.
301+
- Add support for queries that gives information about the capabilities of the implementation on a particular device.
302+
- Once the SPIRV translator work is done, this code generation work will move to the backend along enabling JIT compilation.
303+
304+
## Revision History
305+
306+
[frame="none",options="header"]
307+
|======================
308+
|Rev |Date |Author |Changes
309+
|1 |2021-04-13 |Dounia Khaldi |Initial public working draft.
310+
|======================

‎sycl/doc/extensions/README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ DPC++ extensions status:
4242
| [Invoke SIMD](InvokeSIMD/InvokeSIMD.asciidoc) | Proposal | |
4343
| [Uniform](Uniform/Uniform.asciidoc) | Proposal | |
4444
| [Assert](Assert/SYCL_ONEAPI_ASSERT.asciidoc) | Proposal | |
45+
| [Matrix](Matrix/dpcpp-joint-matrix.asciidoc) | Partially supported(AMX AOT) | Not supported: dynamic-extent, wg and wi scopes, layouts other than packed|
4546

4647
Legend:
4748

0 commit comments

Comments
 (0)