Is the AMX accelerator used on Apple silicon?

philipturner: <p>From issue <span class="hashtag">#47702</span> on the PyTorch repository, it is not yet clear whether PyTorch already uses AMX on Apple silicon to accelerate computations. It might do this because it relies on the operating system’s BLAS library, which is Accelerate on macOS. For reasons not described here, Apple has released little documentation on the AMX ever since its debut in the A13 chip.</p>
<p>If PyTorch does already use AMX, then that is ~1.3 TFLOPS of processing power. For comparison, the M1 GPU has 2.6 TFLOPS. The issue linked above was raised partially because PyTorch lacked hardware acceleration on Apple devices for a very long time. If AMX is in fact used and has comparable performance to GPU acceleration, then many people might want to know.</p>
<p>Could anyone investigate whether the AMX is being used? You may need to learn a bit of Swift, which provides direct access to Accelerate and microsecond-level precision for profiling. Note that M1 has one AMX, while M1 Pro/Max has two. Here are some helpful links for anyone who wishes to investigate this:</p>
<ul>
<li>
<a href="https://github.com/philipturner/swift-colab" rel="noopener nofollow ugc">Swift-Colab repository</a> (has links to Swift crash course)</li>
<li><a href="https://gist.github.com/philipturner/d6fe8c9458bc22a1b7417d54e47e5f40" rel="noopener nofollow ugc">Swift for PyTorch demo</a></li>
</ul>

philipturner: <p>Some more helpful links:</p>
<ul>
<li><a href="https://medium.com/swlh/apples-m1-secret-coprocessor-6599492fc1e1" rel="noopener nofollow ugc">Explanation of AMX</a></li>
<li><a href="https://github.com/pvieito/PythonKit" rel="noopener nofollow ugc">PythonKit repository</a></li>
</ul>

psobolewskiPhD: <p>This issue suggests that AMX is used as of the MPS nightlies:</p><aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/77938">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/77938" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/77938" target="_blank" rel="noopener nofollow ugc">1.12 on M1 Pro Chip not using all the CPU cores. (device="cpu")</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2022-05-20" data-time="04:07:10" data-timezone="UTC">04:07AM - 20 May 22 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2022-05-20" data-time="13:47:20" data-timezone="UTC">01:47PM - 20 May 22 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/PkuCuipy" target="_blank" rel="noopener nofollow ugc">
          <img alt="PkuCuipy" src="https://avatars.githubusercontent.com/u/53848202?v=4" class="onebox-avatar-inline" width="20" height="20">
          PkuCuipy
        </a>
      </div>
    </div>

    <div class="labels">
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">### 🐛 Describe the bug

After updated to nightly-build PyTorch 1.12, a perform<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">ance test is made to compare `'mps'` over `'cpu'` as shown below:
```python
import torch
from tqdm import trange

DTYPE = torch.float32
MAT_SIZE = 5000
DEVICE = ["cpu", "mps"][0]      # it's CPU now

mat = torch.randn([MAT_SIZE, MAT_SIZE], dtype=DTYPE, device=DEVICE)

for i in trange(N_ITER := 100):
    mat @= mat                  # &lt;--- Main Computation HERE
    print(mat[0, 0], end="")    # avoid sync-issue when using 'mps'
```
It's true that `"mps"` is somehow faster than `"cpu"` on this M1-Pro chip.
**However**, I soon noticed that it's not utilizing all the 10 CPU cores when device="cpu"?
Specifically, `Activity Monitor.app` shows that it only use ≈200% of CPU.

After further experiments, I found some interesting facts:
1. As mentioned above, `device="cpu"` on version 1.12 will not using all CPU cores on M1 Pro Chip.
2. When switching back to version 1.11, `device="cpu"` **do** take advantage of all the CPU cores.
3. Although `2.` is true, it's actually slower than 1.12! i.e. `device="cpu"` on 1.12 uses less CPU and less power yet got better performance.
4. Althogh `1.` is true, manully running N(such as 2) instances of this script will let the performance of each script drop down to 1/N of its original. (while indeed more CPU cores are scheduled and more Watts are consumed..)

I'm wondering the reason of `1.` `2.` `3.` and `4.`, and am not sure whether it's some bug of PyTorch or some mistake in my experiments or in my mind.

### Versions

PyTorch version: 1.12.0.dev20220518
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.3.1 (arm64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2.5)
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.12 | packaged by conda-forge | (main, Mar 24 2022, 23:25:14)  [Clang 12.0.1 ] (64-bit runtime)
Python platform: macOS-12.3.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.6
[pip3] torch==1.12.0.dev20220518
[pip3] torchlibrosa==0.0.9
[pip3] torchvision==0.9.0a0
[conda] numpy                     1.21.6           py39h690d673_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge
[conda] pytorch                   1.12.0.dev20220518         py3.9_0    pytorch-nightly
[conda] torchlibrosa              0.0.9                    pypi_0    pypi
[conda] torchvision               0.9.1           py39h0a40b5a_0_cpu    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>I’ve run the same example on my vanilla M1 using a nightly and get 100% CPU use, consistent with just 1 AMX, vs 2 for the OP.</p>
<p>I did a cursory look through PRs and didn’t see Accelerate explicitly mentioned…</p>

philipturner: <p>I forgot to hyperlink the following comment on this thread. The AMX needs to use all 8 power cores for full utilization.</p>
<aside class="quote quote-modified" data-post="9" data-topic="152015">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/b5e925/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-check-mps-availability/152015/9">How to check mps availability?</a> <a class="badge-category__wrapper " href="/c/metal-performance-shader/38"><span data-category-id="38" data-drop-close="true" class="badge-category " title="This category is for any question related to MPS support on Apple hardware (both M1 and x86 with AMD machines)."><span class="badge-category__name">mps</span></span></a>
  </div>
  <blockquote>
    Is AMX already used in pytorch ? I saw you posted a <a href="https://discuss.pytorch.org/t/is-the-amx-accelerator-used-on-apple-silicon/142304">question</a> asking that and also <a href="https://github.com/pytorch/pytorch/issues/47702#issuecomment-1064257151" rel="noopener nofollow ugc">this comment</a> on GitHub. 

I would regard whoever posted that comment with skepticism. Regardless, if PyTorch uses Accelerate, then it uses the AMX. <a class="mention" href="/u/alband">@albanD</a> said that PyTorch should use Accelerate by default, although the documentation does not officially confirm this is true. For example, it seems to say that other math libraries will be chosen instead of Accelerate if they already exist on the system, because the…
  </blockquote>
</aside>