Apple AMX instruction set (M1/M2 matrix coprocessor)

Apple AMX instruction set (M1/M2 matrix coprocessor)(github.com)

210 points by rostayob 3 years ago | 188 comments

Knowing the history of Apple an open standard and given their success with their implementation of the ARM64 ISA, it is unfortunately highly probable that they will follow the proprietary route once again.

Indeed, they are already doing it, we're lucky they weren't in a dominant position when TCP/IP or HTML were invented.

aseipp 3 years ago | |

There isn't any standard matrix multiplication instruction set so there's nothing to standardize over. Machine-learning driven instruction sets (of which MM instructions are motivated by, but not exclusively for) like this have been generally bespoke because the field moves relatively quickly compared to hardware. Every vendor generally follows some basic principles but the specifics are dependent on the workloads and models they expect e.g. quantization or how they expect to split models across accelerators. And ARM does not allow public proprietary instruction set extensions to ARM cores, one of their defining architecture features is that licensees literally are not allowed to do this.[1] The only reason Apple was allowed to do so in this case is likely 1) They negotiated it as part of their AAL (probably for a lot of money) and 2) They do not publicly document or commit to this feature in any way. It could get deleted or disabled in silicon tomorrow and Apple would be able to handle that easily, and in every other visible way they have a normal ARM64 complaint CPU core (there is the custom GIC and performance counters and some other stuff, but none of those violate the architectural license and are just IP stuff they chose to work on themselves.)

So actually the thing you're complaining about is prevented by ARM themselves; Apple cannot publicly commit to features that would fragment the architecture. They don't have to do everything identical either, though.

[1] They have publicly said they will allow some future Cortex cores to contain custom instructions, but it is quite clearly something they're very much still in control over, you won't get a blank check, especially considering almost all ARM licensees use pre-canned CPU cores and IP. You'll probably have to pay them for the extra design work. There are no known desktop/server-class CPUs that fit this profile on the current ARM roadmap, or any taped out processor, that I am aware of.

brigade 3 years ago | | |

> There isn't any standard matrix multiplication instruction set

The Scalable Matrix Extension supplement was released last year. Though obviously AMX predates it, having shipped in actual silicon 3 years ago.

AprilArcus 3 years ago | | |

> The only reason Apple was allowed to do so in this case is likely 1) They negotiated it as part of their AAL (probably for a lot of money)

Apple fronted the cash that created ARM holdings in the first place, so yes, they invested quite a lot of money (well, relative to the other senior partners Acorn and VLSI and later investors), and ARM was hardly in a position to tell them "no".

stephc_int13 3 years ago | | |

History will tell, but I have a bad feeling about "Apple Silicon".

They would not use that naming if they intended to support the official ARM ISA in the long run.

The only thing that would prevent them for going the proprietary route is if they can't.

Reason077 3 years ago | |

> "we're lucky they weren't in a dominant position when TCP/IP or HTML were invented"

TCP/IP and HTML became dominant because they were open standards. Should they have been proprietary, they would have floundered and something else would have emerged instead.

rjzzleep 3 years ago | | |

I guess people have forgotten, that Webkit came out of KHTML from the KDE team and that Apple was a nightmare when it came to contributing code back. They just released a huge dump of the whole thing.

fattire 3 years ago | | |

> TCP/IP and HTML became dominant because they were open standards. Should they have been proprietary, they would have floundered and something else would have emerged instead.

Anyone remember AppleTalk?

https://en.wikipedia.org/wiki/AppleTalk

smoldesu 3 years ago | | |

ActiveX would like to have a word with you.

ribit 3 years ago | |

Apple AMX is a private extension. It is not documented and not exposed to the developer. You have to use system-provided HPC and ML libraries to take advantage of the AMX units. This gives Apple the freedom to iterate and change implementation details at any time without breaking backwards compatibility.

I am sure they will support open standards in time as they mature, but for now there is little advantage in doing so. Not to mention that open standards are far from being a universal panacea. Remember Apple’s last serious involvement into open standards - OpenCL - which was promptly killed by Nvidia. Apple has since learned their lesson and focus on their own needs and technology stack first.

Terretta 3 years ago | |

https://en.wikipedia.org/wiki/Thunderbolt_(interface)

What I liked about this is that Apple realized it makes more sense for someone other than them to develop and supply tech for something like Thunderbolt, co-engineered with Intel for Intel to do it (and eventually open it, royalty free).

GekkePrutser 3 years ago | |

Well, they did create OpenCL.

But yeah it's ages ago and they've kept everything proprietary since then :(

ribit 3 years ago | | |

I think OpenCL was a pivotal moment. Apple created the draft and donated it to Khronos, where is subsequently was stagnating before Nvidia effectively sabotaged the effort in order to push its own proprietary CUDA. Since then Apple has been focusing on their own hardware and ecosystem needs. I think the advantage if this is well illustrated by Metal which grew from being this awkward limited DX9 copy to a fully featured and very flexible yet still consistent and compact GPU API. Sometimes it pays not having to cater to minimal common denominator.

In the recent years I am more and more convinced that open standards are not a panacea. It depends on the domain. For cutting edge specialized compute open standards may even be detrimental. I’d rather have vendor-specific low-level libraries that fully expose the hardware capabilities, with open standard APIs implemented on top of them.

dweekly 3 years ago | |

Re: the "lucky" remark: they certainly tried! Remember AppleTalk and HyperCard?

KerrAvon 3 years ago | | |

This is seriously ahistorical revisionism. There was no universal networking standard in the 1980’s, and there certainly wasn’t anything universally suited for networking on microcomputers, especially not with the zero-configuration usability that Apple wanted. Remember that TCP/IP was largely a plaything for academics until the early 1990’s, and Apple had to create zero configuration and multicast DNS before they could consider deprecating AppleTalk for use on small LANs.

dan-robertson 3 years ago |

Maybe I’m just slow but it wasn’t immediately obvious to me how to use this for matrix multiplication. Let me now try to explain.

Suppose we have some matrices we would like to multiply, a_ok and b_ij (and let’s say their sizes line up with the hardware because I think those details aren’t so relevant).

Their product is

  c_ik = a_ij b_jk = sum(a_ij * b_jk for all j).

The hardware lets us cheaply compute and accumulate an outer product (see picture in OP):

  r_ij = r’_ij + p_i * q_j

Now start with r = 0 and accumulate:

  r_ik = a_i1 * b1k
       + a_i2 * b2k
       + ...
       + a_in * b_nk
       = c_ik

Each row corresponds to one AMX op on all the cells of the matrix.

Writing it out like this it seems quite straightforward. I think I was caught up on thinking about the per-cell computation too much. When computing based on cells in the output, you take a row from the left hand side and dot it with a column from the right hand side (nn dot products). Here, we take a column* from the left hand side and a row from the left hand side and outer product them (n outer products) and add up the result. Perhaps this is partly a victory for this kind of symbolic index notation. I think this would all be much less obvious if I wrote it all out as a sum of outer products with eg the tensor product symbol.

16890c 3 years ago |

So what software actually uses this / what compiler supports this if it's neither supported nor documented by Apple?

garblegarble 3 years ago | |

Applications use it via a higher level interface: Apple's Accelerate framework[1]

1: https://developer.apple.com/accelerate/

peaslock 3 years ago | | |

Speaking of which, is anyone aware of example code using the LSTM? I've been trying to get this to work, but there seems to be information missing e.g how to setup the input/output descriptors and how to manage input data: https://developer.apple.com/documentation/accelerate/bnns/us...

Someone 3 years ago | |

It’s just like with private APIs. It’s functionality that they don’t know (¿yet?) whether they want to support it in the future, but they do provide calls that use them indirectly.

There also could be bugs in the hardware that they carefully programmed around.

scarface74 3 years ago | | |

Why do people act as if “private APIs” (which is somewhat of a contradiction in terms) is nefarious? An API is something that is publicly documented that the vendor promises to support where the behavior won’t change.

Raymond Chen has been blogging for close two decades about all of the hacks that MS had to put into Windows because vendors used undocumented APIs.

michelb 3 years ago | |

Here are the docs: https://developer.apple.com/documentation/accelerate

kennywinker 3 years ago | |

Presumably first party software that qualifies. I’m thinking siri, or their AR stuff. It’s possible it’s also used by something like CoreML where there is a public facing framework that utilizes this under the hood.

I have no special info, these are just guesses.

sufiyan 3 years ago |

I don’t get why other chip manufacturers don’t go this same route. For example AVX is done on the same core that also supports integer math. Many companies have a separate GPU but AVX seems to always come prepackaged.

guipsp 3 years ago | |

There is a cost to moving data to another chip(let). I think this is the main reason

viktorcode 3 years ago | |

Qualcomm most likely will go this route (proprietary instructions) as this now officially blessed by Arm. I'd bet we'll see that in their 2023 chips.

ndesaulniers 3 years ago | | |

Isn't ARM suing Qualcomm?

ladyanita22 3 years ago |

Which compiler would be required for this (https://github.com/corsix/amx/blob/main/aarch64.h)?

I understand the limitation is not at the OS side, as nothing can be done there, but at the compiler-side (I mean that the Apple-supplied compiler doesn't compile against the AMX instruction set, so you'd need a compatible one that, I understand, doesn't exist).

Or is it just undocumented and you can actually get it to work with a Standard xcode and macOS installation given the headers provided?

dougall 3 years ago | |

This header works with standard Xcode/macOS, by taking advantage of inline assembly in a slightly-cursed way (turning register names into numbers and encoding the instruction itself).

corsix 3 years ago | | |

Slightly cursed is how I roll. I would like to know where the trick first originated from though (I found it at https://github.com/yvt/amx-rs/blob/main/src/nativeops.rs#L22 rather than inventing it de novo)

dagmx 3 years ago | |

Compilers wouldn’t target it. You’d get to it via https://developer.apple.com/documentation/accelerate

saboot 3 years ago |

Is there a comparison with other fast cpu methods of matrix multiplication?

corsix 3 years ago | |

https://www.corsix.org/content/contrasting-intel-amx-and-app... compares this with Intel’s AMX

dougall 3 years ago | |

On M1, for single-precision, one AMX P-unit is ~1.64 TFLOPs, one P-core is ~102 GFLOPS. So ~16x core-for-core. But you have four P-cores for every AMX P-unit, so more like 4x. And for double-precision that shrinks to 2x (~410 GFLOPs to ~51GFLOPs).

(This is a simplification that doesn't include the E-cores, nor the AMX E-unit, but their contribution isn't huge. I suspect AMX throughput may have doubled on M2, but I haven't verified that.)

Firadeoclus 3 years ago | | |

That's a pretty huge amount of processing power hidden away! Are these experimentally confirmed performance numbers?

So this is 16x16 single precision fused multiply-adds at 3.2GHz with a throughput of one per clock, right? (16 x 16 x 2 x 3.2e9 = 1.6384e12) And does using fp16 quadruple throughput again? That would put AMX well above the GPU for fp16 matrix multiplication! (2.6 TFLOPs fp16/fp32 for 8-core GPU, 128 multiply-add per core at 1.278GHz)

How does it compare in terms of power, can it sustain 3.2GHz indefinitely or does it hit power/thermal limits fairly quickly?

moonchild 3 years ago | |

IIRC I measured something like 10-50% performance difference (don't remember exactly, but it was somewhere in there), vs a reasonably well-regarded blas implementation. This was for dgemm specifically; I don't know if the story changes for smaller floats.

saboot 3 years ago | | |

Not too bad, I do scientific computing and choose Intel/Nvidia as their APIs for accelerated math operations are documented and supported for developers.

I've been paying attention to what Apple has been pushing with their M1/M2 chips, and I'm pretty tempted to try it out, but unless these features are documented and supported I can't feel comfortable writing programs relying on them.

bfgoodrich 3 years ago | |

SGEMM / DGEMM using AMX2 (the first M1 has AMX2. The A14 has AMX1) is approximately 100% faster than the same running with NEON, which is already a specialized vector math system.

coder543 3 years ago |

Doesn't M2 add support for SVE2 as well?

aseipp 3 years ago | |

No, there are no publicly available chips with SVE2 support at all; the closest you can get is SVE1 on Graviton 3 from AWS.

dougall 3 years ago | | |

Yeah, M2 doesn't. But I think there are some phones using the Snapdragon 8 Gen 1 that have SVE2.

zeristor 3 years ago |

Is this a repost the other day, I thought that was too good have been missed out.

Also I’m keen to see if this 60Gb.a-1 near field wireless data link for Apple Watches for diagnosis will be able to be used in some sort of MagSafe/usb for iPhones.

altairprime 3 years ago | |

(This isn’t a post about near-field, Watches, MagSafe, or USB; perhaps this comment was meant for another post?)

scarface74 3 years ago |

What is the “non proprietary” alternative and why should Apple or it’s users be forced to wait on consensus?

This is the same reason that Apple wasn’t saddled with the horrible PC “standards” before USB became ubiquitous.

Not to mention even today, Bluetooth is a shit show outside of the Apple ecosystem as far as handoff an ease of pairing.