| commit | f458a88724397414192c38e64f8af76e18eee49e | [log] [tgz] |
|---|---|---|
| author | bjacob <benoitjacob@google.com> | Fri Nov 10 10:08:36 2023 -0500 |
| committer | GitHub <noreply@github.com> | Fri Nov 10 10:08:36 2023 -0500 |
| tree | 02020cab391bb4029a3f45552603e95530bc0f49 | |
| parent | f9d7599874087fa30b32446d5459b627a5024558 [diff] |
Optimized vecmat ukernel tile functions for `i16 x u4 -> i32` on AVX-512-VNNI (#15525) This kernel is parametrized in N0, allowing N0==16 and N0==32. Performance on AMD Ryzen 9 7950X3D: - With N0=16: 180 Gop/s. - With N0=32: 240 Gop/s. These numbers show that there's a nice reward for going extra large, but that's also a liability for vecmat shapes whose N dimension isn't a multiple of 32. Maybe we can keep both for now. This is currently by far our fastest vecmat tile function --- it's fast even by general-matmul standards, while usually vecmat's low arithmetic intensity relegates it to lower performance levels. It shows what's possible now that we've decoupled vecmat tile shapes from general matmul tile shapes in #15431 . That 32x8 is not a truncation of a general matmul tile shape. Other element types and CPU architectures all need to get the same treatment. The idea of this kernel is to split the LHS s16 values into high and low 8-bit components to be able to use `_mm512_dpbusd_epi32`. In itself, that doesn't reduce the number of arithmetic instructions: while each now computes a 4D dot-product instead of a 2D one as in `_mm512_dpwssd_epi32`, we now need twice more of them to do separately the high and low 8bit parts of the LHS s16 values. The real benefit is that this removes the need to extend RHS u4 values to s16. Since this is a vecmat kernel, the LHS is small and the RHS is big, so it matters to avoid RHS-processing work. It's not trivial how to use `_mm512_dpbusd_epi32`, with its quirky unsigned * signed semantics. We take advantage of the fact that our u4 RHS values, when extended to u8, do not use the top bit -- so they are also interpretable as s8 values in place. So this is specific to RHS being less-than-8-bit values (it's not specific beyond that to 4bit). Meanwhile, when we split the LHS s16 values into high and low 8bit components the high 8bits are signed s8 and the low 8bit are unsigned u8. So, for each of the combinations of operands that we have to feed `_mm512_dpbusd_epi32`, we manage to find an operand order that accomodates the instruction's requirements on signednesses.
IREE (Intermediate Representation Execution Environment, pronounced as “eerie”) is an MLIR-based end-to-end compiler and runtime that lowers Machine Learning (ML) models to a unified IR that scales up to meet the needs of the datacenter and down to satisfy the constraints and special considerations of mobile and edge deployments.
See our website for project details, user guides, and instructions on building from source.
IREE is still in its early phase. We have settled down on the overarching infrastructure and are actively improving various software components as well as project logistics. It is still quite far from ready for everyday use and is made available without any support at the moment. With that said, we welcome any kind of feedback on any communication channels!
See our website for more information.
IREE is licensed under the terms of the Apache 2.0 License with LLVM Exceptions. See LICENSE for more information.