Code didn't speed up as expected when using `mps`

TeddyHuang-00: <p>I’m really excited to try out the latest pytorch build (1.12.0.dev20220518) for the m1 gpu support, but on my device (M1 max, 64GB, 16-inch MBP), the training time per epoch on <code>cpu</code> is ~9s, but after switching to <code>mps</code>, the performance drops significantly to ~17s. Is that something we should expect, or did I just mess something up?</p>

TeddyHuang-00: <p>BTW, there were failed assertions that may relate to <code>mps</code> as well, the traceback is as follows</p>
<pre><code class="lang-auto">/AppleInternal/Library/BuildRoots/560148d7-a559-11ec-8c96-4add460b61a6/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm:343: failed assertion `unsupported datatype for constant'

Aborted
/opt/homebrew/Cellar/python@3.10/3.10.4/Frameworks/Python.framework/Versions/3.10/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
</code></pre>

philipturner: <p>What were the data types of the tensors you were using? Could you show us the script? I have a good idea of which data types MPS and MPSGraph supports. Also, I think I have seen some peculiarities with MPSGraph’s internal MLIR compiler, which might cause a crash.</p>
<p><a class="mention" href="/u/alband">@albanD</a> could you check out the memory management issue in the above comment’s log? I have previously seen apps crash when you don’t balance a <code>DispatchSemaphore</code> before releasing it. Maybe that’s an iOS-specific crash, but it should be standard practice to balance semaphore objects when using Foundation.</p>

philipturner: <blockquote>
<p>the training time per epoch on <code>cpu</code> is ~9s, but after switching to <code>mps</code> , the performance drops significantly to ~17s.</p>
</blockquote>
<p>If you happen to be using all CPU cores on the M1 Max in <code>cpu</code> mode, then you have 2.0 TFLOPS of processing power (256 GFLOPS per core) for matrix multiplication. On GPU, you have a maximum of 10.4 TFLOPS, although 80% of that is used. If a certain tensor shape works better on the CPU than the GPU, then the slowdown you experienced might be possible.</p>

albanD: <p>Hey!<br>
Thanks for the report.</p>
<ul>
<li>For the datatype crash. You might want to make sure you’re not using <code>float64</code> numbers next to the place where the crash happens as these are not supported for Tensors on mps.</li>
<li>For the slowdown, that is not expected for sure (but can happen depending on the workload as GPUs require large relatively large tasks to see speedups). If you have a simple repro for us to run locally, that would be very helpful!</li>
</ul>

philipturner: <aside class="quote no-group" data-username="albanD" data-post="5" data-topic="152016">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>float64</p>
</blockquote>
</aside>
<p>For a more extensive list of which data types do and don’t run:</p>
<ul>
<li>Avoid Float64 on all Apple devices. Even if the hardware supports <code>Double</code> physically (AMD or Intel), the Metal API doesn’t let you access it.</li>
<li>Avoid BFloat16. That is natively supported by the latest Nvidia GPUs, but not supported in Metal. Also don’t try to use TF18/TF32 or Int4.</li>
<li>All standard integer types (UInt8, UInt16, UInt32, UInt64) and their signed counterparts work natively on Apple devices. Not exactly 8-bit integers, which are cast to 16-bit integers before being stored into registers, but those aren’t going to harm performance. Yes, they run 64-bit integers on the Apple GPU and not 64-bit floats. Metal allows you to use 64-bit integers in shaders on AMD and Intel, but the arithmetic there might just happen through emulation (slow). I think that’s where I experienced the crash in MPSGraph previously - trying to run an operation on UInt64 on my Intel Mac mini.</li>
</ul>
<p>For anything else regarding compatibility of data types, the Metal feature set tables might be useful:</p>
<aside class="onebox pdf" data-onebox-src="https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf">
  <header class="source">

      <a href="https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf" target="_blank" rel="noopener nofollow ugc">developer.apple.com</a>
  </header>

  <article class="onebox-body">
    <a href="https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf" target="_blank" rel="noopener nofollow ugc"><span class="pdf-onebox-logo"></span></a>

<h3><a href="https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf" target="_blank" rel="noopener nofollow ugc">Metal-Feature-Set-Tables.pdf</a></h3>


  </article>

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

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


philipturner: <p>I remember the crash with MPSGraph more precisely. I was messing around with transcendental functions while prototyping an S4TF Metal backend, and had a configuration where I passed a UInt64 <code>MPSGraphTensor</code> into a function that normally runs with floating point in hardware. They would have to convert the integer to a float before actually running the operation. It worked for everything up to 32-bit integers, but I made a special configuration using 64-bit integers that broke the MLIR compiler. The error message was something like “pow_u64 not found”. You may need to watch out for a bug like this as well as Float64 crashes.</p>

TeddyHuang-00: <aside class="onebox allowlistedgeneric" data-onebox-src="https://pastebin.com/iaMG0aR2">
  <header class="source">
      <img src="https://pastebin.com/favicon.ico" class="site-icon" width="16" height="16">

      <a href="https://pastebin.com/iaMG0aR2" target="_blank" rel="noopener nofollow ugc">Pastebin</a>
  </header>

  <article class="onebox-body">
    <img src="https://pastebin.com/i/facebook.png" class="thumbnail" width="" height="">

<h3><a href="https://pastebin.com/iaMG0aR2" target="_blank" rel="noopener nofollow ugc">CNN.py - Pastebin.com</a></h3>

  <p>Pastebin.com is the number one paste tool since 2002. Pastebin is a website where you can store text online for a set period of time.</p>


  </article>

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

  <div style="clear: both"></div>
</aside>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://pastebin.com/vYMfsB7H">
  <header class="source">
      <img src="https://pastebin.com/favicon.ico" class="site-icon" width="16" height="16">

      <a href="https://pastebin.com/vYMfsB7H" target="_blank" rel="noopener nofollow ugc">Pastebin</a>
  </header>

  <article class="onebox-body">
    <img src="https://pastebin.com/i/facebook.png" class="thumbnail" width="" height="">

<h3><a href="https://pastebin.com/vYMfsB7H" target="_blank" rel="noopener nofollow ugc">PGD.py - Pastebin.com</a></h3>

  <p>Pastebin.com is the number one paste tool since 2002. Pastebin is a website where you can store text online for a set period of time.</p>


  </article>

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

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

<p>Here are all the code that I use to test, download and just run <code>CNN.py</code> should work. I think that I’m indeed using the Float64 datatype, which is the default I think? I’ll try using another datatype later.</p>

TeddyHuang-00: <aside class="onebox allowlistedgeneric" data-onebox-src="https://pastebin.com/WzYG0T06">
  <header class="source">
      <img src="https://pastebin.com/favicon.ico" class="site-icon" width="16" height="16">

      <a href="https://pastebin.com/WzYG0T06" target="_blank" rel="noopener nofollow ugc">Pastebin</a>
  </header>

  <article class="onebox-body">
    <img src="https://pastebin.com/i/facebook.png" class="thumbnail" width="" height="">

<h3><a href="https://pastebin.com/WzYG0T06" target="_blank" rel="noopener nofollow ugc">params.py - Pastebin.com</a></h3>

  <p>Pastebin.com is the number one paste tool since 2002. Pastebin is a website where you can store text online for a set period of time.</p>


  </article>

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

  <div style="clear: both"></div>
</aside>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://pastebin.com/venT37Ha">
  <header class="source">
      <img src="https://pastebin.com/favicon.ico" class="site-icon" width="16" height="16">

      <a href="https://pastebin.com/venT37Ha" target="_blank" rel="noopener nofollow ugc">Pastebin</a>
  </header>

  <article class="onebox-body">
    <img src="https://pastebin.com/i/facebook.png" class="thumbnail" width="" height="">

<h3><a href="https://pastebin.com/venT37Ha" target="_blank" rel="noopener nofollow ugc">utils.py - Pastebin.com</a></h3>

  <p>Pastebin.com is the number one paste tool since 2002. Pastebin is a website where you can store text online for a set period of time.</p>


  </article>

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

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

<hr>
<p>Sorry that I can only put up to 2 links in a post. so here is another half of the code</p>

Hankcs_He: <p>How about torch.bool? Bool seems to be unsupported but it’s such a common data type.</p>
<pre><code class="lang-auto">import torch

t = torch.ones(1, device=torch.device('mps'))
m = t.ne(0)
m[0] = False
print(m)

</code></pre>
<blockquote>
<p>/AppleInternal/Library/BuildRoots/560148d7-a559-11ec-8c96-4add460b61a6/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm:343: failed assertion `unsupported datatype for constant’</p>
</blockquote>

philipturner: <p>A two second test, and you’re correct. <code>Bool</code> is not supported as an MPSGraph constant. That’s very wierd. I can’t say the same about other tensor types though, such as ones doing logical OR and what not. It shouldn’t be hard to figure that out, though.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/e/2e0b759bdebe6983cdfc3b09a1aa1d1f8f2f8e18.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/2e0b759bdebe6983cdfc3b09a1aa1d1f8f2f8e18" title="Screen Shot 2022-07-19 at 8.01.27 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2e0b759bdebe6983cdfc3b09a1aa1d1f8f2f8e18_2_690x342.jpeg" alt="Screen Shot 2022-07-19 at 8.01.27 PM" data-base62-sha1="6zkvWRpMb4v2RN5OygTkvOii5VC" width="690" height="342" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2e0b759bdebe6983cdfc3b09a1aa1d1f8f2f8e18_2_690x342.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2e0b759bdebe6983cdfc3b09a1aa1d1f8f2f8e18_2_1035x513.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2e0b759bdebe6983cdfc3b09a1aa1d1f8f2f8e18_2_1380x684.jpeg 2x" data-dominant-color="252727"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2022-07-19 at 8.01.27 PM</span><span class="informations">1920×952 130 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

TeddyHuang-00: <p>Yeah, this test case failed too on torch 1.13.0.dev20220719, here is a quick fix to the problem I think</p>
<pre><code class="lang-python">import torch

T = torch.ones(1, device=torch.device('mps'))[0]
F = torch.zeros(1, device=torch.device('mps'))[0]
print(f"{T=}, {F=}")


t = torch.ones(1, device=torch.device('mps'))
m = t.ne(0)
print(m)
m[0] = F
print(m)
</code></pre>
<p>And I got the expected output like below, though a fall back warning was triggered</p>
<pre data-code-wrap="sh"><code class="lang-nohighlight">T=tensor(1., device='mps:0'), F=tensor(0., device='mps:0')
tensor([True], device='mps:0')
tensor([False], device='mps:0')
</code></pre>
<hr>
<p>Edit:</p>
<p>I think the problem is that you cannot assign any value to such a tensor on mps device, so conversion before doing so is needed to get things done.</p>

philipturner: <p>Here’s a deep investigation into the nature of the problem. I lowered the Python code into its respective calls to Metal. I’m writing the script in Swift because I’m most familiar with that language, but I hope you can follow along. I explicitly marked which code sections correspond to the Python code using <code>// PYTHON CODE:</code></p>
<p>The script is large, so you may have to scroll to see all of it.</p>
<pre><code class="lang-auto">import MetalPerformanceShadersGraph
func testTensor0D(_ torchValue: Double) -&gt; Float {
  let graph = MPSGraph()
  let tensor = graph.constant(torchValue, dataType: .float32)
  let data = graph.run(
    feeds: [:], targetTensors: [tensor], targetOperations: nil)[tensor]
  let ndarray = data!.mpsndarray()
  var output: [Float32] = .init(repeating: 9999, count: 1)
  ndarray.readBytes(&amp;output, strideBytes: nil)
  return output.first!
}

// PYTHON CODE: T = torch.ones(1, device=torch.device('mps'))[0]
let T = testTensor0D(1.0)
// PYTHON CODE: F = torch.zeros(1, device=torch.device('mps'))[0]
let F = testTensor0D(0.0)
// PYTHON CODE: print(f"{T=}, {F=}")
print("T=\(T), F=\(F)")

extension MTLBuffer {
  func setFirstElement(to value: Float32) {
    let pointer = self.contents().assumingMemoryBound(to: Float32.self)
    pointer[0] = value
  }

  func extractFirstElement() -&gt; Float32 {
    let pointer = self.contents().assumingMemoryBound(to: Float32.self)
    return pointer[0]
  }
}

func get_m_MTLBuffer() -&gt; MTLBuffer {
  let device: MTLDevice = MTLCreateSystemDefaultDevice()!
  let commandQueue: MTLCommandQueue = device.makeCommandQueue()!
  
  let graph = MPSGraph()
  // PYTHON CODE: t = torch.ones(1, device=torch.device('mps'))
  let t = graph.constant(1, dataType: .float32)
  let zeroLiteral = graph.constant(0, dataType: .float32)
  // PYTHON CODE: m = t.ne[0]
  let m_Tensor: MPSGraphTensor = graph.notEqual(t, zeroLiteral, name: nil)

  // Execute the MPSGraph.
  let m_MTLBuffer = device.makeBuffer(
    length: 1 * MemoryLayout&lt;Float32&gt;.stride, options: .storageModeShared)!
  m_MTLBuffer.setFirstElement(to: 9999)
  let m_TensorData = MPSGraphTensorData(
    m_MTLBuffer, shape: [1], dataType: .float32)
  graph.run(
    with: commandQueue, feeds: [:], targetOperations: nil, 
    resultsDictionary: [m_Tensor: m_TensorData])
  return m_MTLBuffer
}

do {
  // Metal backing buffer that stores the tensor's data.
  let m: MTLBuffer = get_m_MTLBuffer()
  // PYTHON CODE: print(m)
  print(m.extractFirstElement())
}
</code></pre>
<p>I stopped short of translating the last two lines of <a class="mention" href="/u/teddyhuang-00">@TeddyHuang-00</a>’s code sample. If I were to perfectly translate what PyTorch is doing there, the script would become absurdly long. The gist:</p>
<ul>
<li>
<ol>
<li>Make a <code>MTLBuffer</code> whose raw data is <code>F</code>. This means it’s 4 bytes of memory, set to <code>Float32(0.0)</code>.</li>
</ol>
</li>
<li>
<ol start="2">
<li>Copy that <code>MTLBuffer</code> into a slice of <code>m</code>, emulating the subscript operator from Python.</li>
</ol>
</li>
<li>
<ol start="3">
<li>Somewhere in the process, the <code>F</code> Metal buffer becomes an <code>MPSGraphTensorData</code> and needs to set its <code>MPSDataType</code>. This may even be a call to <code>MPSGraph.constant</code>, taking an <code>NSData</code> as its first argument.</li>
</ol>
</li>
<li>
<ol start="4">
<li>In <a class="mention" href="/u/hankcs_he">@Hankcs_He</a>’s example, the <code>MPSDataType</code> happened to be <code>.bool</code>. In the <a class="mention" href="/u/teddyhuang-00">@TeddyHuang-00</a>’s example, the <code>MPSDataType</code> happened to be <code>.float32</code>. That’s why it doesn’t crash.</li>
</ol>
</li>
</ul>
<aside class="quote no-group" data-username="TeddyHuang-00" data-post="12" data-topic="152016">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/teddyhuang-00/48/49748_2.png" class="avatar"> TeddyHuang-00:</div>
<blockquote>
<pre data-code-wrap="txt"><code class="lang-txt">T=tensor(1., device='mps:0'), F=tensor(0., device='mps:0')
tensor([True], device='mps:0')
tensor([False], device='mps:0')
</code></pre>
</blockquote>
</aside>
<p>In LLDB, I get the following. I’m appending <code>0.0</code> to the output, which would appear if I finished the script. This is similar to the Terminal output quoted above.</p>
<pre data-code-wrap="txt"><code class="lang-txt">T=1.0, F=0.0
-1.0
0.0
</code></pre>
<hr>
<aside class="quote no-group" data-username="TeddyHuang-00" data-post="12" data-topic="152016">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/teddyhuang-00/48/49748_2.png" class="avatar"> TeddyHuang-00:</div>
<blockquote>
<p>And I got the expected output like below, though a fall back warning was triggered</p>
</blockquote>
</aside>
<p>Which fallback warning was triggered? I assume it’s a <code>Float64/Double</code> → <code>Float32/Float</code> warning, saying “MPS doesn’t support Double, falling back to Float”. Python floating point numbers are naturally in the FP64 format, which probably causes the warning. Is it possible to set a <code>dtype='float32'</code> parameter to suppress the warning?</p>

philipturner: <p>It’s especially wierd that <code>MPSGraph.notEqual</code> returns a Float32 tensor, not a Bool tensor. You would think it returns a <code>Bool</code> because it’s comparing two numbers. In the TensorFlow low-level C++ codebase, the operator <code>NotEqual</code> always returns a boolean. Perhaps PyTorch differs from TensorFlow in that regard?</p>

TeddyHuang-00: <aside class="quote no-group quote-modified" data-username="philipturner" data-post="13" data-topic="152016">
<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"> philipturner:</div>
<blockquote>
<aside class="quote no-group" data-username="TeddyHuang-00" data-post="12" data-topic="152016">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/teddyhuang-00/48/49748_2.png" class="avatar"> TeddyHuang-00:</div>
<blockquote>
<p>And I got the expected output like below, though a fall back warning was triggered</p>
</blockquote>
</aside>
<p>Which fallback warning was triggered? I assume it’s a <code>Float64/Double</code> → <code>Float32/Float</code> warning, saying “MPS doesn’t support Double, falling back to Float”. Python floating point numbers are naturally in the FP64 format, which probably causes the warning. Is it possible to set a <code>dtype='float32'</code> parameter to suppress the warning?</p>
</blockquote>
</aside>
<p>Actually it is not the dtype fallback warning, but a boolean operation one. The full warning message is as follow:</p>
<pre data-code-wrap="sh"><code class="lang-sh">path/to/python3.9/site-packages/torch/_tensor_str.py:103: UserWarning: The operator 'aten::bitwise_and.Tensor_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
  nonzero_finite_vals = torch.masked_select(tensor_view, torch.isfinite(tensor_view) &amp; tensor_view.ne(0))
</code></pre>

TeddyHuang-00: <aside class="quote no-group" data-username="philipturner" data-post="13" data-topic="152016">
<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"> philipturner:</div>
<blockquote>
<p>Somewhere in the process, the <code>F</code> Metal buffer becomes an <code>MPSGraphTensorData</code> and needs to set its <code>MPSDataType</code>. This may even be a call to <code>MPSGraph.constant</code>, taking an <code>NSData</code> as its first argument.</p>
</blockquote>
</aside>
<p>So… to my naive understanding, this is something that happened implicitly by the torch library during the value assignment process?</p>

philipturner: <aside class="quote no-group" data-username="TeddyHuang-00" data-post="16" data-topic="152016">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/teddyhuang-00/48/49748_2.png" class="avatar"> TeddyHuang-00:</div>
<blockquote>
<p>So… to my naive understanding, this is something that happened implicitly by the torch library during the value assignment process?</p>
</blockquote>
</aside>
<p>Correct. The library implicitly tried to create an MPSGraph tensor or tensor data that was a boolean.</p>

TeddyHuang-00: <aside class="quote no-group" data-username="philipturner" data-post="17" data-topic="152016">
<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"> philipturner:</div>
<blockquote>
<p>The library implicitly tried to create an MPSGraph tensor or tensor data that was a boolean.</p>
</blockquote>
</aside>
<p>That explains part of the problem for me! Thanks a lot for your hard work and explanation!</p>