From 3a43fd93d5db91d50e732ad5ab4981fd13a4d760 Mon Sep 17 00:00:00 2001 From: lezcano Date: Thu, 14 Sep 2023 14:58:54 +0000 Subject: [PATCH 01/11] Initial commit --- blogpost/post.md | 160 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 160 insertions(+) create mode 100644 blogpost/post.md diff --git a/blogpost/post.md b/blogpost/post.md new file mode 100644 index 00000000..5f892ab9 --- /dev/null +++ b/blogpost/post.md @@ -0,0 +1,160 @@ +# Compiling NumPy into C++ or CUDA via `torch.compile` + +Tracing through NumPy code via `torch.compile` is now possible in PyTorch 2.1. This feature leverages PyTorch's compiler to generate efficient fused vectorized code without having to modify your original code. Even more, it also allows for executing NumPy functions on CUDA just by running them through `torch.compile` under `torch.device("cuda")`! + +In this post, we go over how to use this feature and give a few tips and tricks to make the most of it. + + +## Compiling NumPy into Parallel C++ + +We will take as our running example the iteration step in a K-Means algorithm presented in this [NumPy book](https://realpython.com/numpy-array-programming/#clustering-algorithms) + +```python +import numpy as np + +def get_labels(X, means): + return np.argmin(np.linalg.norm(X - means[:, None], axis=2), axis=0) +``` + +We create a synthetic dataset with 10M random 2-D points. We can see that, given that the means are chosen appropriately, the function returns the correct cluster for all of them + +```python +npts = 10_000_000 +X = np.repeat([[5, 5], [10, 10]], [npts, npts], axis=0) +X = X + np.random.randn(*X.shape) # 2 distinct "blobs" +means = np.array([[5, 5], [10, 10]]) +pred = get_labels(X, means) +``` + +Benchmarking this function gives us a baseline of **1.26s** on an AMD 3970X. + +Compiling this function is now as easy as wrapping it with `torch.compile` and executing it with the example inputs + +```python +compiled_fn = torch.compile(get_labels) +new_pred = compiled_fn(X, means) +assert np.allclose(prediction, new_pred) +``` + +The compiled function yields a 9x speed-up when running it on 1 core. Even better, since the compiled code also runs on multiple cores, we get a **57x speed-up** when running it on 32 cores. Note that vanilla NumPy always runs on just one core. + +We may inspect the generated C++ code by running the script with `TORCH_LOGS=output_code`, and we can see that `torch.compile` was able to compile the broadcasting, together with the two reductions into just one for-loop, and it parallelizes it using OpenMP +```c++ +extern "C" void kernel(const double* in_ptr0, const long* in_ptr1, long* out_ptr0) { + #pragma omp parallel num_threads(32) + #pragma omp for + for(long i0=0L; i0<20000000L; i0+=1L) { + auto tmp0 = in_ptr0[2L*i0]; + auto tmp1 = in_ptr1[0L]; + auto tmp5 = in_ptr0[1L + (2L*i0)]; + auto tmp6 = in_ptr1[1L]; + ... +``` + +## Compiling NumPy into CUDA + +Compiling our code so that it runs on CUDA is as simple as setting locally the default dtype to be the CUDA + +```python +with torch.device("cuda"): + cuda_pred = compiled_fn(X, means) +assert np.allclose(prediction, cuda_pred) +``` + +By inspecting the generated code via `TORCH_LOGS=output_code`, we see that, rather than generating CUDA code directly, `torch.compile` generates rather readable [triton](https://triton-lang.org/main/index.html) code + +```python +def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr): + xnumel = 20000000 + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK)[:] + xmask = xindex < xnumel + x0 = xindex + tmp0 = tl.load(in_ptr0 + (2*x0), xmask) + tmp1 = tl.load(in_ptr1 + (0)) + ... +``` + +Running this small snippet on an RTX 2060 gives an **8x speed-up** over the original NumPy code. This is something, but it is not particularly impressive, given the speed-ups we have seen on CPU. Let's have a look into how to get some proper speed-ups on CUDA via a couple minor changes. + +**`float64` vs `float32`**. Many GPUs, in particular consumer-grade ones, are rather sluggish when running operations on `float64`. For this reason, changing the data generation to `float32`, the original NumPy code just gets a bit faster, about a 9%, but our CUDA code gets **40% faster**, yielding a **11x speed-up** over the plain NumPy code. + +`torch.compile`, by default, respects the NumPy semantics, and as such, it uses `np.float64` as its default dtype for all its creation ops. As discussed, this can hinder performance, so it is possible to change this default by setting + +```python +from torch._dynamo import config +config.numpy_default_float = "float32" +``` + +**CPU <> CUDA copies**. An 11x speed-up is good, but it is not even close to the CPU numbers. This is caused by a small transformation that `torch.compile` does behind the scenes. The code above takes NumPy arrays and returns NumPy arrays. All of these arrays are on CPU, but the computations are performed on the GPU. This means that every time the function is called, `torch.compie` has to copy all these arrays from CPU to the GPU, and then copy the result from CUDA back to CPU to preserve the original semantics. There is no native solution to this issue in NumPy, as NumPy does not have the notion of a `device`. That being said, we can work around it by creating a wrapper to this function so that it accepts PyTorch tensors and returns PyTorch tensors. + +```python +@torch.compile +def tensor_fn(X, means): + X, means = X.numpy(), means.numpy() + ret = get_labels(X, means) + return torch.from_numpy(ret) + +def cuda_fn(X, means): + with torch.device("cuda"): + return tensor_fn(X, means) +``` + +This function now takes tensors in CUDA memory and returns tensors in CUDA memory, but the function itself is written in NumPy! When we keep the tensors in CUDA and perform the computations in `float32`, we see a **200x speed-up** over the initial NumPy implementation on `float32` arrays. + +**Mixing NumPy and PyTorch**. In this example, we had to write a small adaptor to move the data from CPU to CUDA and back. In programs that mix PyTorch and NumPy this is already done by calling `x.detach().cpu().numpy()` (or simply `x.numpy(force=True)`). Since when running under `torch.compile` we can run NumPy code in CUDA, we can simply modify this code to call `x.numpy()` and when running it under `device("cuda")`, as we did above, it will generate efficient CUDA code from original NumPy calls without copying the data from CUDA to CPU at all. + +## Further Speed-up tricks + +**General advice**. The CUDA code we have shown is already quite efficient, but it is true that this is a rather tiny program. When dealing with larger programs, we may need to tweak parts of it to make it more efficient. A good place to start is the [`torch.compile` troubleshooting page](https://pytorch.org/docs/stable/dynamo/troubleshooting.html#performance-profiling). This showcases a number of ways to inspect the tracing process, and how to identify problematic code that may cause slow downs. + +**Advice when compiling NumPy code**. NumPy, even if it is rather similar to PyTorch, it is often used very differently. It is rather common to perform computations in NumPy and then do an if/else depending on the value of the array, or perform operations in-place, perhaps via boolean masks. These constructions, while supported by `torch.compile`, hamper its performance. Changes like moving from in-place indexing to using `np.where`, writing the code in a branchless way, or avoid using in-place ops in favor of out-of-place ops can go a long way. + +To write fast NumPy code, it is best to avoid loops, but sometimes they are unavoidable. When tracing through a loop, `torch.compile` will try to fully unroll it. This is sometimes desirable, but sometimes it may not even be possible, like when we have a dynamic stopping condition (like a while loop). In these cases, it may be best to just compile the body of the loop, perhaps compiling a few iterations at a time (loop unrolling). + +**Debugging NumPy code**. Debugging is rather tricky when a compiler is involved. To figure out whether an error you are hitting is a `torch.compile` error, or an error from the program, you can execute your NumPy program without `torch.compile` by replacing the NumPy import by `import torch._numpy as np`. This is should just be used for **debugging purposes** and is in no way a replacement for the PyTorch API, as it is **much slower**. + +## Differences between NumPy and `torch.compile`d NumPy + +**NumPy scalars**. NumPy returns NumPy scalars in almost any case where PyTorch would return a 0-D tensor (e.g. from `np.sum`). Under `torch.compile`, NumPy scalars are treated as 0-D arrays. This is just fine in most cases. The only case when their behavior diverges is when NumPy scalars are implicitly used as Python scalars. For example, +```python +>>> np.asarray(2) * [1, 2, 3] # 0-D array is an array-like +array([2, 4, 6]) +>>> u = np.int32(2) +>>> u * [1, 2, 3] # scalar decays into a Python int +[1, 2, 3, 1, 2, 3] +>>> torch.compile(lambda: u * [1, 2, 3])() +array([2, 4, 6]) # acts as a 0-D array, not as a scalar ?!?! +``` +If we compile the first two lines, we see that `torch.compile` treats `u` as a 0-D array. To recover the eager semantics, we just need to make the casting explicit +```python +>>> torch.compile(lambda: int(u) * [1, 2, 3])() +[1, 2, 3, 1, 2, 3] +``` + +**Type promotion and versioning**. NumPy's type promotion rules may be, at times, a bit surprising +```python +>>> np.asarray([1], dtype=np.int8) + 127 +array([128], dtype=int8) +>>> np.asarray([1], dtype=np.int8) + 128 +array([129], dtype=int16) +``` +These rules are changing to follow a set of rules that is closer to that of PyTorch in NumPy 2.0. The relevant technical document is [NEP 50](https://numpy.org/neps/nep-0050-scalar-promotion.html). `torch.compile` went ahead and implemented NEP 50 rather than the about-to-be-deprecated rules. + +In general, `torch.compile` will match the semantics of the last NumPy release. + +## Beyond NumPy: SciPy and scikit-learn + +In parallel to this effort, other Quansight engineers have designed, proposed and got merged a way to support PyTorch arrays within SciPy and scikit-learn. This was encountered with a big enthusiasm by the other maintainers from these libraries, as it was shown that using PyTorch as a backend would often yield considerable speed-ups. + +This can of course be combined with `torch.compile` to be able to compile programs that rely on these other libraries. + +Note that the initial support is just restricted to a few algorithms in scikit-learn and to `scipy.cluster` in SciPy. + +If you want to learn more about this effort, how to use it, or how to help moving it forward, see this post. [TODO link post] + +## Conclusion +[TODO Make sure Greg approves this wording] +PyTorch has committed since its inception to be a framework compatible with the rest of the Python ecosystem. Enabling compiling NumPy programs, and establishing the tools necessary to do the same for other prominent libraries are two more steps in this direction. Quansight and Meta continue working in this direction, improving the compatibility between PyTorch and the rest of the ecosystem. + +From Quansight, we would like to thank Meta for funding this project and all the previous work that lead to it, like improving the NumPy compatibility within PyTorch, and developing the [python Array API](https://data-apis.org/array-api/latest/). Without this consistent support, this would not have been possible. From 8a3581f5eb46411eb29abf25fc0b8c4432edf2e6 Mon Sep 17 00:00:00 2001 From: lezcano Date: Thu, 14 Sep 2023 16:23:17 +0000 Subject: [PATCH 02/11] split lines --- blogpost/post.md | 160 +++++++++++++++++++++++++++++++++++++---------- 1 file changed, 126 insertions(+), 34 deletions(-) diff --git a/blogpost/post.md b/blogpost/post.md index 5f892ab9..ba06b473 100644 --- a/blogpost/post.md +++ b/blogpost/post.md @@ -1,13 +1,19 @@ # Compiling NumPy into C++ or CUDA via `torch.compile` -Tracing through NumPy code via `torch.compile` is now possible in PyTorch 2.1. This feature leverages PyTorch's compiler to generate efficient fused vectorized code without having to modify your original code. Even more, it also allows for executing NumPy functions on CUDA just by running them through `torch.compile` under `torch.device("cuda")`! +Tracing through NumPy code via `torch.compile` is now possible in PyTorch 2.1. +This feature leverages PyTorch's compiler to generate efficient fused +vectorized code without having to modify your original code. Even more, it +also allows for executing NumPy functions on CUDA just by running them through +`torch.compile` under `torch.device("cuda")`! -In this post, we go over how to use this feature and give a few tips and tricks to make the most of it. +In this post, we go over how to use this feature and give a few tips and tricks +to make the most of it. ## Compiling NumPy into Parallel C++ -We will take as our running example the iteration step in a K-Means algorithm presented in this [NumPy book](https://realpython.com/numpy-array-programming/#clustering-algorithms) +We will take as our running example the iteration step in a K-Means algorithm +presented in this [NumPy book](https://realpython.com/numpy-array-programming/#clustering-algorithms) ```python import numpy as np @@ -16,7 +22,9 @@ def get_labels(X, means): return np.argmin(np.linalg.norm(X - means[:, None], axis=2), axis=0) ``` -We create a synthetic dataset with 10M random 2-D points. We can see that, given that the means are chosen appropriately, the function returns the correct cluster for all of them +We create a synthetic dataset with 10M random 2-D points. We can see that, +given that the means are chosen appropriately, the function returns the correct +cluster for all of them ```python npts = 10_000_000 @@ -28,7 +36,8 @@ pred = get_labels(X, means) Benchmarking this function gives us a baseline of **1.26s** on an AMD 3970X. -Compiling this function is now as easy as wrapping it with `torch.compile` and executing it with the example inputs +Compiling this function is now as easy as wrapping it with `torch.compile` and +executing it with the example inputs ```python compiled_fn = torch.compile(get_labels) @@ -36,9 +45,15 @@ new_pred = compiled_fn(X, means) assert np.allclose(prediction, new_pred) ``` -The compiled function yields a 9x speed-up when running it on 1 core. Even better, since the compiled code also runs on multiple cores, we get a **57x speed-up** when running it on 32 cores. Note that vanilla NumPy always runs on just one core. +The compiled function yields a 9x speed-up when running it on 1 core. Even +better, since the compiled code also runs on multiple cores, we get a **57x speed-up** +when running it on 32 cores. Note that vanilla NumPy always runs on +just one core. -We may inspect the generated C++ code by running the script with `TORCH_LOGS=output_code`, and we can see that `torch.compile` was able to compile the broadcasting, together with the two reductions into just one for-loop, and it parallelizes it using OpenMP +We may inspect the generated C++ code by running the script with +`TORCH_LOGS=output_code`, and we can see that `torch.compile` was able to +compile the broadcasting, together with the two reductions into just one +for-loop, and it parallelizes it using OpenMP ```c++ extern "C" void kernel(const double* in_ptr0, const long* in_ptr1, long* out_ptr0) { #pragma omp parallel num_threads(32) @@ -53,7 +68,8 @@ extern "C" void kernel(const double* in_ptr0, const long* in_ptr1, long* out_ptr ## Compiling NumPy into CUDA -Compiling our code so that it runs on CUDA is as simple as setting locally the default dtype to be the CUDA +Compiling our code so that it runs on CUDA is as simple as setting locally the +default dtype to be the CUDA ```python with torch.device("cuda"): @@ -61,7 +77,9 @@ with torch.device("cuda"): assert np.allclose(prediction, cuda_pred) ``` -By inspecting the generated code via `TORCH_LOGS=output_code`, we see that, rather than generating CUDA code directly, `torch.compile` generates rather readable [triton](https://triton-lang.org/main/index.html) code +By inspecting the generated code via `TORCH_LOGS=output_code`, we see that, +rather than generating CUDA code directly, `torch.compile` generates rather +readable [triton](https://triton-lang.org/main/index.html) code ```python def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr): @@ -75,18 +93,36 @@ def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr): ... ``` -Running this small snippet on an RTX 2060 gives an **8x speed-up** over the original NumPy code. This is something, but it is not particularly impressive, given the speed-ups we have seen on CPU. Let's have a look into how to get some proper speed-ups on CUDA via a couple minor changes. +Running this small snippet on an RTX 2060 gives an **8x speed-up** over the +original NumPy code. This is something, but it is not particularly impressive, +given the speed-ups we have seen on CPU. Let's have a look into how to get some +proper speed-ups on CUDA via a couple minor changes. -**`float64` vs `float32`**. Many GPUs, in particular consumer-grade ones, are rather sluggish when running operations on `float64`. For this reason, changing the data generation to `float32`, the original NumPy code just gets a bit faster, about a 9%, but our CUDA code gets **40% faster**, yielding a **11x speed-up** over the plain NumPy code. +**`float64` vs `float32`**. Many GPUs, in particular consumer-grade ones, are +rather sluggish when running operations on `float64`. For this reason, changing +the data generation to `float32`, the original NumPy code just gets a bit +faster, about a 9%, but our CUDA code gets **40% faster**, yielding a **11x +speed-up** over the plain NumPy code. -`torch.compile`, by default, respects the NumPy semantics, and as such, it uses `np.float64` as its default dtype for all its creation ops. As discussed, this can hinder performance, so it is possible to change this default by setting +`torch.compile`, by default, respects the NumPy semantics, and as such, it uses +`np.float64` as its default dtype for all its creation ops. As discussed, this +can hinder performance, so it is possible to change this default by setting ```python from torch._dynamo import config config.numpy_default_float = "float32" ``` -**CPU <> CUDA copies**. An 11x speed-up is good, but it is not even close to the CPU numbers. This is caused by a small transformation that `torch.compile` does behind the scenes. The code above takes NumPy arrays and returns NumPy arrays. All of these arrays are on CPU, but the computations are performed on the GPU. This means that every time the function is called, `torch.compie` has to copy all these arrays from CPU to the GPU, and then copy the result from CUDA back to CPU to preserve the original semantics. There is no native solution to this issue in NumPy, as NumPy does not have the notion of a `device`. That being said, we can work around it by creating a wrapper to this function so that it accepts PyTorch tensors and returns PyTorch tensors. +**CPU <> CUDA copies**. An 11x speed-up is good, but it is not even close to +the CPU numbers. This is caused by a small transformation that `torch.compile` +does behind the scenes. The code above takes NumPy arrays and returns NumPy +arrays. All of these arrays are on CPU, but the computations are performed on +the GPU. This means that every time the function is called, `torch.compie` has +to copy all these arrays from CPU to the GPU, and then copy the result from +CUDA back to CPU to preserve the original semantics. There is no native +solution to this issue in NumPy, as NumPy does not have the notion of a +`device`. That being said, we can work around it by creating a wrapper to this +function so that it accepts PyTorch tensors and returns PyTorch tensors. ```python @torch.compile @@ -100,23 +136,60 @@ def cuda_fn(X, means): return tensor_fn(X, means) ``` -This function now takes tensors in CUDA memory and returns tensors in CUDA memory, but the function itself is written in NumPy! When we keep the tensors in CUDA and perform the computations in `float32`, we see a **200x speed-up** over the initial NumPy implementation on `float32` arrays. +This function now takes tensors in CUDA memory and returns tensors in CUDA +memory, but the function itself is written in NumPy! When we keep the tensors +in CUDA and perform the computations in `float32`, we see a **200x speed-up** +over the initial NumPy implementation on `float32` arrays. -**Mixing NumPy and PyTorch**. In this example, we had to write a small adaptor to move the data from CPU to CUDA and back. In programs that mix PyTorch and NumPy this is already done by calling `x.detach().cpu().numpy()` (or simply `x.numpy(force=True)`). Since when running under `torch.compile` we can run NumPy code in CUDA, we can simply modify this code to call `x.numpy()` and when running it under `device("cuda")`, as we did above, it will generate efficient CUDA code from original NumPy calls without copying the data from CUDA to CPU at all. +**Mixing NumPy and PyTorch**. In this example, we had to write a small adaptor +to move the data from CPU to CUDA and back. In programs that mix PyTorch and +NumPy this is already done by calling `x.detach().cpu().numpy()` (or simply +`x.numpy(force=True)`). Since when running under `torch.compile` we can run +NumPy code in CUDA, we can simply modify this code to call `x.numpy()` and when +running it under `device("cuda")`, as we did above, it will generate efficient +CUDA code from original NumPy calls without copying the data from CUDA to CPU +at all. ## Further Speed-up tricks -**General advice**. The CUDA code we have shown is already quite efficient, but it is true that this is a rather tiny program. When dealing with larger programs, we may need to tweak parts of it to make it more efficient. A good place to start is the [`torch.compile` troubleshooting page](https://pytorch.org/docs/stable/dynamo/troubleshooting.html#performance-profiling). This showcases a number of ways to inspect the tracing process, and how to identify problematic code that may cause slow downs. - -**Advice when compiling NumPy code**. NumPy, even if it is rather similar to PyTorch, it is often used very differently. It is rather common to perform computations in NumPy and then do an if/else depending on the value of the array, or perform operations in-place, perhaps via boolean masks. These constructions, while supported by `torch.compile`, hamper its performance. Changes like moving from in-place indexing to using `np.where`, writing the code in a branchless way, or avoid using in-place ops in favor of out-of-place ops can go a long way. - -To write fast NumPy code, it is best to avoid loops, but sometimes they are unavoidable. When tracing through a loop, `torch.compile` will try to fully unroll it. This is sometimes desirable, but sometimes it may not even be possible, like when we have a dynamic stopping condition (like a while loop). In these cases, it may be best to just compile the body of the loop, perhaps compiling a few iterations at a time (loop unrolling). - -**Debugging NumPy code**. Debugging is rather tricky when a compiler is involved. To figure out whether an error you are hitting is a `torch.compile` error, or an error from the program, you can execute your NumPy program without `torch.compile` by replacing the NumPy import by `import torch._numpy as np`. This is should just be used for **debugging purposes** and is in no way a replacement for the PyTorch API, as it is **much slower**. +**General advice**. The CUDA code we have shown is already quite efficient, but +it is true that this is a rather tiny program. When dealing with larger +programs, we may need to tweak parts of it to make it more efficient. A good +place to start is the [`torch.compile` troubleshooting +page](https://pytorch.org/docs/stable/dynamo/troubleshooting.html#performance-profiling). +This showcases a number of ways to inspect the tracing process, and how to +identify problematic code that may cause slow downs. + +**Advice when compiling NumPy code**. NumPy, even if it is rather similar to +PyTorch, it is often used very differently. It is rather common to perform +computations in NumPy and then do an if/else depending on the value of the +array, or perform operations in-place, perhaps via boolean masks. These +constructions, while supported by `torch.compile`, hamper its performance. +Changes like moving from in-place indexing to using `np.where`, writing the +code in a branchless way, or avoid using in-place ops in favor of out-of-place +ops can go a long way. + +To write fast NumPy code, it is best to avoid loops, but sometimes they are +unavoidable. When tracing through a loop, `torch.compile` will try to fully +unroll it. This is sometimes desirable, but sometimes it may not even be +possible, like when we have a dynamic stopping condition (like a while loop). +In these cases, it may be best to just compile the body of the loop, perhaps +compiling a few iterations at a time (loop unrolling). + +**Debugging NumPy code**. Debugging is rather tricky when a compiler is +involved. To figure out whether an error you are hitting is a `torch.compile` +error, or an error from the program, you can execute your NumPy program without +`torch.compile` by replacing the NumPy import by `import torch._numpy as np`. +This is should just be used for **debugging purposes** and is in no way a +replacement for the PyTorch API, as it is **much slower**. ## Differences between NumPy and `torch.compile`d NumPy -**NumPy scalars**. NumPy returns NumPy scalars in almost any case where PyTorch would return a 0-D tensor (e.g. from `np.sum`). Under `torch.compile`, NumPy scalars are treated as 0-D arrays. This is just fine in most cases. The only case when their behavior diverges is when NumPy scalars are implicitly used as Python scalars. For example, +**NumPy scalars**. NumPy returns NumPy scalars in almost any case where PyTorch +would return a 0-D tensor (e.g. from `np.sum`). Under `torch.compile`, NumPy +scalars are treated as 0-D arrays. This is just fine in most cases. The only +case when their behavior diverges is when NumPy scalars are implicitly used as +Python scalars. For example, ```python >>> np.asarray(2) * [1, 2, 3] # 0-D array is an array-like array([2, 4, 6]) @@ -126,35 +199,54 @@ array([2, 4, 6]) >>> torch.compile(lambda: u * [1, 2, 3])() array([2, 4, 6]) # acts as a 0-D array, not as a scalar ?!?! ``` -If we compile the first two lines, we see that `torch.compile` treats `u` as a 0-D array. To recover the eager semantics, we just need to make the casting explicit + +If we compile the first two lines, we see that `torch.compile` treats `u` as a +0-D array. To recover the eager semantics, we just need to make the casting +explicit ```python >>> torch.compile(lambda: int(u) * [1, 2, 3])() [1, 2, 3, 1, 2, 3] ``` -**Type promotion and versioning**. NumPy's type promotion rules may be, at times, a bit surprising +**Type promotion and versioning**. NumPy's type promotion rules may be, at +times, a bit surprising ```python >>> np.asarray([1], dtype=np.int8) + 127 array([128], dtype=int8) >>> np.asarray([1], dtype=np.int8) + 128 array([129], dtype=int16) ``` -These rules are changing to follow a set of rules that is closer to that of PyTorch in NumPy 2.0. The relevant technical document is [NEP 50](https://numpy.org/neps/nep-0050-scalar-promotion.html). `torch.compile` went ahead and implemented NEP 50 rather than the about-to-be-deprecated rules. +These rules are changing to follow a set of rules that is closer to that of +PyTorch in NumPy 2.0. The relevant technical document is [NEP 50](https://numpy.org/neps/nep-0050-scalar-promotion.html). +`torch.compile` went ahead and implemented NEP 50 rather than the about-to-be-deprecated rules. In general, `torch.compile` will match the semantics of the last NumPy release. ## Beyond NumPy: SciPy and scikit-learn -In parallel to this effort, other Quansight engineers have designed, proposed and got merged a way to support PyTorch arrays within SciPy and scikit-learn. This was encountered with a big enthusiasm by the other maintainers from these libraries, as it was shown that using PyTorch as a backend would often yield considerable speed-ups. +In parallel to this effort, other Quansight engineers have designed, proposed +and got merged a way to support PyTorch arrays within SciPy and scikit-learn. +This was encountered with a big enthusiasm by the other maintainers from these +libraries, as it was shown that using PyTorch as a backend would often yield +considerable speed-ups. -This can of course be combined with `torch.compile` to be able to compile programs that rely on these other libraries. +This can of course be combined with `torch.compile` to be able to compile +programs that rely on these other libraries. -Note that the initial support is just restricted to a few algorithms in scikit-learn and to `scipy.cluster` in SciPy. +Note that the initial support is just restricted to a few algorithms in +scikit-learn and to `scipy.cluster` in SciPy. -If you want to learn more about this effort, how to use it, or how to help moving it forward, see this post. [TODO link post] +If you want to learn more about this effort, how to use it, or how to help +moving it forward, see this post. [TODO link post] -## Conclusion -[TODO Make sure Greg approves this wording] -PyTorch has committed since its inception to be a framework compatible with the rest of the Python ecosystem. Enabling compiling NumPy programs, and establishing the tools necessary to do the same for other prominent libraries are two more steps in this direction. Quansight and Meta continue working in this direction, improving the compatibility between PyTorch and the rest of the ecosystem. +## Conclusion [TODO Make sure Greg approves this wording] PyTorch has committed +since its inception to be a framework compatible with the rest of the Python +ecosystem. Enabling compiling NumPy programs, and establishing the tools +necessary to do the same for other prominent libraries are two more steps in +this direction. Quansight and Meta continue working in this direction, +improving the compatibility between PyTorch and the rest of the ecosystem. -From Quansight, we would like to thank Meta for funding this project and all the previous work that lead to it, like improving the NumPy compatibility within PyTorch, and developing the [python Array API](https://data-apis.org/array-api/latest/). Without this consistent support, this would not have been possible. +From Quansight, we would like to thank Meta for funding this project and all +the previous work that lead to it, like improving the NumPy compatibility +within PyTorch, and developing the [python Array API](https://data-apis.org/array-api/latest/). +Without this consistent support, this would not have been possible. From 33e168ab62632bc7795ffc39c15ed64b23404df6 Mon Sep 17 00:00:00 2001 From: lezcano Date: Fri, 15 Sep 2023 14:57:48 +0000 Subject: [PATCH 03/11] Address review --- blogpost/post.md | 86 +++++++++++++++++++++++++++--------------------- 1 file changed, 48 insertions(+), 38 deletions(-) diff --git a/blogpost/post.md b/blogpost/post.md index ba06b473..1eb2ebcd 100644 --- a/blogpost/post.md +++ b/blogpost/post.md @@ -1,16 +1,16 @@ -# Compiling NumPy into C++ or CUDA via `torch.compile` +# Compiling NumPy code into C++ or CUDA via `torch.compile` Tracing through NumPy code via `torch.compile` is now possible in PyTorch 2.1. This feature leverages PyTorch's compiler to generate efficient fused -vectorized code without having to modify your original code. Even more, it +vectorized code without having to modify your original code. Even more, it also allows for executing NumPy functions on CUDA just by running them through `torch.compile` under `torch.device("cuda")`! In this post, we go over how to use this feature and give a few tips and tricks -to make the most of it. +to make the most out of it. -## Compiling NumPy into Parallel C++ +## Compiling NumPy code into Parallel C++ We will take as our running example the iteration step in a K-Means algorithm presented in this [NumPy book](https://realpython.com/numpy-array-programming/#clustering-algorithms) @@ -34,7 +34,7 @@ means = np.array([[5, 5], [10, 10]]) pred = get_labels(X, means) ``` -Benchmarking this function gives us a baseline of **1.26s** on an AMD 3970X. +Benchmarking this function gives us a baseline of **1.26s** on an AMD 3970X CPU. Compiling this function is now as easy as wrapping it with `torch.compile` and executing it with the example inputs @@ -63,13 +63,13 @@ extern "C" void kernel(const double* in_ptr0, const long* in_ptr1, long* out_ptr auto tmp1 = in_ptr1[0L]; auto tmp5 = in_ptr0[1L + (2L*i0)]; auto tmp6 = in_ptr1[1L]; - ... + // Rest of the kernel omitted for brevity ``` -## Compiling NumPy into CUDA +## Compiling NumPy code into CUDA -Compiling our code so that it runs on CUDA is as simple as setting locally the -default dtype to be the CUDA +Compiling our code so that it runs on CUDA is as simple as setting the +default dtype to be CUDA ```python with torch.device("cuda"): @@ -90,7 +90,7 @@ def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr): x0 = xindex tmp0 = tl.load(in_ptr0 + (2*x0), xmask) tmp1 = tl.load(in_ptr1 + (0)) - ... + // Rest of the kernel omitted for brevity ``` Running this small snippet on an RTX 2060 gives an **8x speed-up** over the @@ -98,7 +98,7 @@ original NumPy code. This is something, but it is not particularly impressive, given the speed-ups we have seen on CPU. Let's have a look into how to get some proper speed-ups on CUDA via a couple minor changes. -**`float64` vs `float32`**. Many GPUs, in particular consumer-grade ones, are +**`float64` vs `float32`**. Many GPUs, in particular consumer-grade ones, are rather sluggish when running operations on `float64`. For this reason, changing the data generation to `float32`, the original NumPy code just gets a bit faster, about a 9%, but our CUDA code gets **40% faster**, yielding a **11x @@ -117,7 +117,7 @@ config.numpy_default_float = "float32" the CPU numbers. This is caused by a small transformation that `torch.compile` does behind the scenes. The code above takes NumPy arrays and returns NumPy arrays. All of these arrays are on CPU, but the computations are performed on -the GPU. This means that every time the function is called, `torch.compie` has +the GPU. This means that every time the function is called, `torch.compile` has to copy all these arrays from CPU to the GPU, and then copy the result from CUDA back to CPU to preserve the original semantics. There is no native solution to this issue in NumPy, as NumPy does not have the notion of a @@ -148,7 +148,8 @@ NumPy this is already done by calling `x.detach().cpu().numpy()` (or simply NumPy code in CUDA, we can simply modify this code to call `x.numpy()` and when running it under `device("cuda")`, as we did above, it will generate efficient CUDA code from original NumPy calls without copying the data from CUDA to CPU -at all. +at all. Note that the resulting code would not run without `torch.compile`. For +it to run in eager mode you would need to go rollback `x.numpy(force=True)`. ## Further Speed-up tricks @@ -158,10 +159,10 @@ programs, we may need to tweak parts of it to make it more efficient. A good place to start is the [`torch.compile` troubleshooting page](https://pytorch.org/docs/stable/dynamo/troubleshooting.html#performance-profiling). This showcases a number of ways to inspect the tracing process, and how to -identify problematic code that may cause slow downs. +identify problematic code that may cause slowdowns. -**Advice when compiling NumPy code**. NumPy, even if it is rather similar to -PyTorch, it is often used very differently. It is rather common to perform +**Advice when compiling NumPy code**. NumPy, even if rather similar to +PyTorch, is often used very differently. It is rather common to perform computations in NumPy and then do an if/else depending on the value of the array, or perform operations in-place, perhaps via boolean masks. These constructions, while supported by `torch.compile`, hamper its performance. @@ -181,7 +182,8 @@ involved. To figure out whether an error you are hitting is a `torch.compile` error, or an error from the program, you can execute your NumPy program without `torch.compile` by replacing the NumPy import by `import torch._numpy as np`. This is should just be used for **debugging purposes** and is in no way a -replacement for the PyTorch API, as it is **much slower**. +replacement for the PyTorch API, as it is **much slower** and, as a private API, +may change without notice. ## Differences between NumPy and `torch.compile`d NumPy @@ -220,33 +222,41 @@ These rules are changing to follow a set of rules that is closer to that of PyTorch in NumPy 2.0. The relevant technical document is [NEP 50](https://numpy.org/neps/nep-0050-scalar-promotion.html). `torch.compile` went ahead and implemented NEP 50 rather than the about-to-be-deprecated rules. -In general, `torch.compile` will match the semantics of the last NumPy release. +In general, `torch.compile` will match the semantics of the lastest NumPy release. ## Beyond NumPy: SciPy and scikit-learn -In parallel to this effort, other Quansight engineers have designed, proposed -and got merged a way to support PyTorch arrays within SciPy and scikit-learn. -This was encountered with a big enthusiasm by the other maintainers from these -libraries, as it was shown that using PyTorch as a backend would often yield -considerable speed-ups. +In parallel to this effort of making torch.compile understand NumPy code, other +Quansight engineers have designed and proposed a way to support PyTorch tensors +within scikit-learn and SciPy. This was received enthusiastically by other +maintainers from these libraries, as it was shown that using PyTorch as a +backend would often yield considerable speed-ups.and support in a number of +APIs. Support for PyTorch tensors across an initial set of APIs and submodules +has now been merged into both projects. -This can of course be combined with `torch.compile` to be able to compile -programs that rely on these other libraries. +This sets the stepping stone to move towards a future where PyTorch tensors +can be used within other libraries in the Python data ecosystem. Even more, +this will enable running these other libraries on GPU and even compiling code +mixing these libraries and PyTorch, similar to how we have been discussed in +this post. -Note that the initial support is just restricted to a few algorithms in -scikit-learn and to `scipy.cluster` in SciPy. +Note that this initial support is just restricted to a few algorithms in +scikit-learn and to `scipy.cluster` and `scipy.fft` in SciPy. If you want to learn more about this effort, how to use it, or how to help moving it forward, see this post. [TODO link post] -## Conclusion [TODO Make sure Greg approves this wording] PyTorch has committed -since its inception to be a framework compatible with the rest of the Python -ecosystem. Enabling compiling NumPy programs, and establishing the tools -necessary to do the same for other prominent libraries are two more steps in -this direction. Quansight and Meta continue working in this direction, -improving the compatibility between PyTorch and the rest of the ecosystem. - -From Quansight, we would like to thank Meta for funding this project and all -the previous work that lead to it, like improving the NumPy compatibility -within PyTorch, and developing the [python Array API](https://data-apis.org/array-api/latest/). -Without this consistent support, this would not have been possible. +## Conclusion + +PyTorch has committed since its inception to be a framework compatible with the +rest of the Python ecosystem. Enabling compiling NumPy programs, and +establishing the tools necessary to do the same for other prominent libraries +are two more steps in this direction. Quansight and Meta continue working in +this direction, improving the compatibility between PyTorch and the rest of the +ecosystem. + +From Quansight, we would like to thank Meta for funding this project as well as +previous work on improving NumPy compatibility within PyTorch, and the project +that led to support PyTorch within scikit-learn and SciPy. These are giant leaps +towards consolidating PyTorch as the reference framework within the open source +Python data ecosystem. From 5f6c2fc02f6d11e9dbc61a5a57460058c5a76d9b Mon Sep 17 00:00:00 2001 From: lezcano Date: Fri, 15 Sep 2023 15:01:23 +0000 Subject: [PATCH 04/11] Shoehorned QS into the first paragraph --- blogpost/post.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/blogpost/post.md b/blogpost/post.md index 1eb2ebcd..9a1a3cec 100644 --- a/blogpost/post.md +++ b/blogpost/post.md @@ -1,6 +1,7 @@ # Compiling NumPy code into C++ or CUDA via `torch.compile` -Tracing through NumPy code via `torch.compile` is now possible in PyTorch 2.1. +Quansight engineers have implemented the ability to trace through NumPy code via +`torch.compile` for PyTorch 2.1. This feature leverages PyTorch's compiler to generate efficient fused vectorized code without having to modify your original code. Even more, it also allows for executing NumPy functions on CUDA just by running them through From 279b4dd952aeb1c0bf7e2b3aed86fa3b9f7c973c Mon Sep 17 00:00:00 2001 From: lezcano Date: Fri, 15 Sep 2023 15:06:46 +0000 Subject: [PATCH 05/11] Now I'm happier with the first paragraph --- blogpost/post.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/blogpost/post.md b/blogpost/post.md index 9a1a3cec..fa72da29 100644 --- a/blogpost/post.md +++ b/blogpost/post.md @@ -1,7 +1,7 @@ # Compiling NumPy code into C++ or CUDA via `torch.compile` -Quansight engineers have implemented the ability to trace through NumPy code via -`torch.compile` for PyTorch 2.1. +Quansight engineers have implemented support for tracing through NumPy code via +`torch.compile` in PyTorch 2.1. This feature leverages PyTorch's compiler to generate efficient fused vectorized code without having to modify your original code. Even more, it also allows for executing NumPy functions on CUDA just by running them through From 1172ac0f6e5b44d9dd7e4b5ad571af19f2ff26dd Mon Sep 17 00:00:00 2001 From: lezcano Date: Fri, 15 Sep 2023 15:32:17 +0000 Subject: [PATCH 06/11] Many more tweaks --- blogpost/post.md | 70 +++++++++++++++++++++++------------------------- 1 file changed, 34 insertions(+), 36 deletions(-) diff --git a/blogpost/post.md b/blogpost/post.md index fa72da29..ab8ba31f 100644 --- a/blogpost/post.md +++ b/blogpost/post.md @@ -13,8 +13,8 @@ to make the most out of it. ## Compiling NumPy code into Parallel C++ -We will take as our running example the iteration step in a K-Means algorithm -presented in this [NumPy book](https://realpython.com/numpy-array-programming/#clustering-algorithms) +We take as our running example one step in a K-Means algorithm. +This piece of code is borrowed from this [NumPy book](https://realpython.com/numpy-array-programming/#clustering-algorithms) ```python import numpy as np @@ -47,14 +47,15 @@ assert np.allclose(prediction, new_pred) ``` The compiled function yields a 9x speed-up when running it on 1 core. Even -better, since the compiled code also runs on multiple cores, we get a **57x speed-up** -when running it on 32 cores. Note that vanilla NumPy always runs on -just one core. - -We may inspect the generated C++ code by running the script with -`TORCH_LOGS=output_code`, and we can see that `torch.compile` was able to -compile the broadcasting, together with the two reductions into just one -for-loop, and it parallelizes it using OpenMP +better, as opposed to NumPy, our generated code does take advantage of all the +cores in a processor. As such, when we run it on 32 cores, we get a **57x speed-up**. +Note that PyTorch always uses all the available cores unless explicitly restricted, +so this is the default behavior you get when using `torch.compile`. + +We may inspect the generated C++ code by running the script with the +environment variable `TORCH_LOGS=output_code`. When doing so, we can see that +`torch.compile` was able to compile the broadcasting, together with the two +reductions into just one for-loop, and parallelizes it using OpenMP ```c++ extern "C" void kernel(const double* in_ptr0, const long* in_ptr1, long* out_ptr0) { #pragma omp parallel num_threads(32) @@ -96,8 +97,8 @@ def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr): Running this small snippet on an RTX 2060 gives an **8x speed-up** over the original NumPy code. This is something, but it is not particularly impressive, -given the speed-ups we have seen on CPU. Let's have a look into how to get some -proper speed-ups on CUDA via a couple minor changes. +given the speed-ups we have seen on CPU. Let's have a look into how to squeeze +the most out of our GPU via a couple minor changes. **`float64` vs `float32`**. Many GPUs, in particular consumer-grade ones, are rather sluggish when running operations on `float64`. For this reason, changing @@ -119,8 +120,8 @@ the CPU numbers. This is caused by a small transformation that `torch.compile` does behind the scenes. The code above takes NumPy arrays and returns NumPy arrays. All of these arrays are on CPU, but the computations are performed on the GPU. This means that every time the function is called, `torch.compile` has -to copy all these arrays from CPU to the GPU, and then copy the result from -CUDA back to CPU to preserve the original semantics. There is no native +to copy all these arrays from CPU to the GPU, and then copy the result +back to CPU to preserve the original semantics. There is no native solution to this issue in NumPy, as NumPy does not have the notion of a `device`. That being said, we can work around it by creating a wrapper to this function so that it accepts PyTorch tensors and returns PyTorch tensors. @@ -143,19 +144,19 @@ in CUDA and perform the computations in `float32`, we see a **200x speed-up** over the initial NumPy implementation on `float32` arrays. **Mixing NumPy and PyTorch**. In this example, we had to write a small adaptor -to move the data from CPU to CUDA and back. In programs that mix PyTorch and +to convert tensors to ndarrays and then back to tensors. In programs that mix PyTorch and NumPy this is already done by calling `x.detach().cpu().numpy()` (or simply `x.numpy(force=True)`). Since when running under `torch.compile` we can run NumPy code in CUDA, we can simply modify this code to call `x.numpy()` and when running it under `device("cuda")`, as we did above, it will generate efficient CUDA code from original NumPy calls without copying the data from CUDA to CPU at all. Note that the resulting code would not run without `torch.compile`. For -it to run in eager mode you would need to go rollback `x.numpy(force=True)`. +it to run in eager mode one would need to rollback to `x.numpy(force=True)`. ## Further Speed-up tricks **General advice**. The CUDA code we have shown is already quite efficient, but -it is true that this is a rather tiny program. When dealing with larger +it is true that the running example is rather short. When dealing with larger programs, we may need to tweak parts of it to make it more efficient. A good place to start is the [`torch.compile` troubleshooting page](https://pytorch.org/docs/stable/dynamo/troubleshooting.html#performance-profiling). @@ -164,7 +165,7 @@ identify problematic code that may cause slowdowns. **Advice when compiling NumPy code**. NumPy, even if rather similar to PyTorch, is often used very differently. It is rather common to perform -computations in NumPy and then do an if/else depending on the value of the +computations in NumPy and then do an if/else depending on values within the array, or perform operations in-place, perhaps via boolean masks. These constructions, while supported by `torch.compile`, hamper its performance. Changes like moving from in-place indexing to using `np.where`, writing the @@ -174,9 +175,9 @@ ops can go a long way. To write fast NumPy code, it is best to avoid loops, but sometimes they are unavoidable. When tracing through a loop, `torch.compile` will try to fully unroll it. This is sometimes desirable, but sometimes it may not even be -possible, like when we have a dynamic stopping condition (like a while loop). +possible, like when we have a dynamic stopping condition, like in a while loop. In these cases, it may be best to just compile the body of the loop, perhaps -compiling a few iterations at a time (loop unrolling). +a few iterations at a time (loop unrolling). **Debugging NumPy code**. Debugging is rather tricky when a compiler is involved. To figure out whether an error you are hitting is a `torch.compile` @@ -184,7 +185,7 @@ error, or an error from the program, you can execute your NumPy program without `torch.compile` by replacing the NumPy import by `import torch._numpy as np`. This is should just be used for **debugging purposes** and is in no way a replacement for the PyTorch API, as it is **much slower** and, as a private API, -may change without notice. +**may change without notice**. ## Differences between NumPy and `torch.compile`d NumPy @@ -219,31 +220,28 @@ array([128], dtype=int8) >>> np.asarray([1], dtype=np.int8) + 128 array([129], dtype=int16) ``` -These rules are changing to follow a set of rules that is closer to that of -PyTorch in NumPy 2.0. The relevant technical document is [NEP 50](https://numpy.org/neps/nep-0050-scalar-promotion.html). +NumPy 2.0 is changing these rules to follow others that are closer to those +PyTorch. The relevant technical document is [NEP 50](https://numpy.org/neps/nep-0050-scalar-promotion.html). `torch.compile` went ahead and implemented NEP 50 rather than the about-to-be-deprecated rules. In general, `torch.compile` will match the semantics of the lastest NumPy release. ## Beyond NumPy: SciPy and scikit-learn -In parallel to this effort of making torch.compile understand NumPy code, other +In parallel to this effort of making `torch.compile` understand NumPy code, other Quansight engineers have designed and proposed a way to support PyTorch tensors within scikit-learn and SciPy. This was received enthusiastically by other maintainers from these libraries, as it was shown that using PyTorch as a -backend would often yield considerable speed-ups.and support in a number of -APIs. Support for PyTorch tensors across an initial set of APIs and submodules -has now been merged into both projects. +backend would often yield considerable speed-ups. +Both projects have now merge initial support for PyTorch tensors across a number of +APIs and submodules. This sets the stepping stone to move towards a future where PyTorch tensors can be used within other libraries in the Python data ecosystem. Even more, -this will enable running these other libraries on GPU and even compiling code -mixing these libraries and PyTorch, similar to how we have been discussed in +this will enable running these other libraries on GPUs and even compiling code +mixing these libraries and PyTorch, similar to what we have been discussed in this post. -Note that this initial support is just restricted to a few algorithms in -scikit-learn and to `scipy.cluster` and `scipy.fft` in SciPy. - If you want to learn more about this effort, how to use it, or how to help moving it forward, see this post. [TODO link post] @@ -252,12 +250,12 @@ moving it forward, see this post. [TODO link post] PyTorch has committed since its inception to be a framework compatible with the rest of the Python ecosystem. Enabling compiling NumPy programs, and establishing the tools necessary to do the same for other prominent libraries -are two more steps in this direction. Quansight and Meta continue working in -this direction, improving the compatibility between PyTorch and the rest of the +are two more steps in this direction. Quansight and Meta continue working hand on +hand, improving the compatibility between PyTorch and the rest of the ecosystem. From Quansight, we would like to thank Meta for funding this project as well as previous work on improving NumPy compatibility within PyTorch, and the project -that led to support PyTorch within scikit-learn and SciPy. These are giant leaps -towards consolidating PyTorch as the reference framework within the open source +that led to supporting PyTorch within scikit-learn and SciPy. These are giant leaps +towards consolidating PyTorch as the framework of choice within the open source Python data ecosystem. From f2b1d610d65bf56202daca16f944cfb67a6254bd Mon Sep 17 00:00:00 2001 From: lezcano Date: Mon, 18 Sep 2023 09:20:28 +0000 Subject: [PATCH 07/11] Review --- blogpost/post.md | 50 +++++++++++++++++++++++++++--------------------- 1 file changed, 28 insertions(+), 22 deletions(-) diff --git a/blogpost/post.md b/blogpost/post.md index ab8ba31f..e44b06a7 100644 --- a/blogpost/post.md +++ b/blogpost/post.md @@ -3,7 +3,7 @@ Quansight engineers have implemented support for tracing through NumPy code via `torch.compile` in PyTorch 2.1. This feature leverages PyTorch's compiler to generate efficient fused -vectorized code without having to modify your original code. Even more, it +vectorized code without having to modify your original NumPy code. Even more, it also allows for executing NumPy functions on CUDA just by running them through `torch.compile` under `torch.device("cuda")`! @@ -54,8 +54,8 @@ so this is the default behavior you get when using `torch.compile`. We may inspect the generated C++ code by running the script with the environment variable `TORCH_LOGS=output_code`. When doing so, we can see that -`torch.compile` was able to compile the broadcasting, together with the two -reductions into just one for-loop, and parallelizes it using OpenMP +`torch.compile` was able to compile the broadcasting and the two +reductions into just one for-loop, and parallelize it using OpenMP ```c++ extern "C" void kernel(const double* in_ptr0, const long* in_ptr1, long* out_ptr0) { #pragma omp parallel num_threads(32) @@ -71,7 +71,7 @@ extern "C" void kernel(const double* in_ptr0, const long* in_ptr1, long* out_ptr ## Compiling NumPy code into CUDA Compiling our code so that it runs on CUDA is as simple as setting the -default dtype to be CUDA +default device to be CUDA ```python with torch.device("cuda"): @@ -139,19 +139,23 @@ def cuda_fn(X, means): ``` This function now takes tensors in CUDA memory and returns tensors in CUDA -memory, but the function itself is written in NumPy! When we keep the tensors -in CUDA and perform the computations in `float32`, we see a **200x speed-up** -over the initial NumPy implementation on `float32` arrays. +memory, but the function itself is written in NumPy! `torch.compile` uses the +`numpy()` and the `from_numpy()` calls as hints, and optimizes them away, and +internally it simply works with PyTorch tensors without moving moving the +memory at all. When we keep the tensors in CUDA and perform the computations in +`float32`, we see a **200x speed-up** over the initial NumPy implementation on +`float32` arrays. **Mixing NumPy and PyTorch**. In this example, we had to write a small adaptor -to convert tensors to ndarrays and then back to tensors. In programs that mix PyTorch and -NumPy this is already done by calling `x.detach().cpu().numpy()` (or simply -`x.numpy(force=True)`). Since when running under `torch.compile` we can run -NumPy code in CUDA, we can simply modify this code to call `x.numpy()` and when -running it under `device("cuda")`, as we did above, it will generate efficient -CUDA code from original NumPy calls without copying the data from CUDA to CPU -at all. Note that the resulting code would not run without `torch.compile`. For -it to run in eager mode one would need to rollback to `x.numpy(force=True)`. +to convert tensors to ndarrays and then back to tensors. In programs that mix +PyTorch and NumPy converting a tensor into an ndarray is often implemented as +`x.detach().cpu().numpy()`, or simply `x.numpy(force=True)`. Since when running +under `torch.compile` we can run NumPy code in CUDA, we can implement this +conversion pattern as call to `x.numpy()`, as we did above. Doing so and +running the resulting code under `device("cuda")` will generate efficient CUDA +code from original NumPy calls without copying the data from CUDA to CPU at +all. Note that the resulting code does not run without `torch.compile`. For it +to run in eager mode one would need to rollback to `x.numpy(force=True)`. ## Further Speed-up tricks @@ -215,8 +219,8 @@ explicit **Type promotion and versioning**. NumPy's type promotion rules may be, at times, a bit surprising ```python ->>> np.asarray([1], dtype=np.int8) + 127 -array([128], dtype=int8) +>>> np.asarray([1], dtype=np.int8) + 126 +array([127], dtype=int8) >>> np.asarray([1], dtype=np.int8) + 128 array([129], dtype=int16) ``` @@ -254,8 +258,10 @@ are two more steps in this direction. Quansight and Meta continue working hand o hand, improving the compatibility between PyTorch and the rest of the ecosystem. -From Quansight, we would like to thank Meta for funding this project as well as -previous work on improving NumPy compatibility within PyTorch, and the project -that led to supporting PyTorch within scikit-learn and SciPy. These are giant leaps -towards consolidating PyTorch as the framework of choice within the open source -Python data ecosystem. +From Quansight, we would like to thank Mengwei, Voz, and Ed for their +invaluable help in integrating our work with `torch.compile`. We would also +like to thank Meta for funding this project as well as previous work on +improving NumPy compatibility within PyTorch, and the project that led to +supporting PyTorch within scikit-learn and SciPy. These are giant leaps towards +consolidating PyTorch as the framework of choice within the open source Python +data ecosystem. From 7e81913864f84c05a74cd24726d05d6af8acf28d Mon Sep 17 00:00:00 2001 From: lezcano Date: Mon, 18 Sep 2023 12:16:53 +0000 Subject: [PATCH 08/11] More errata --- blogpost/post.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/blogpost/post.md b/blogpost/post.md index e44b06a7..df2833bb 100644 --- a/blogpost/post.md +++ b/blogpost/post.md @@ -141,8 +141,8 @@ def cuda_fn(X, means): This function now takes tensors in CUDA memory and returns tensors in CUDA memory, but the function itself is written in NumPy! `torch.compile` uses the `numpy()` and the `from_numpy()` calls as hints, and optimizes them away, and -internally it simply works with PyTorch tensors without moving moving the -memory at all. When we keep the tensors in CUDA and perform the computations in +internally it simply works with PyTorch tensors without moving the memory at +all. When we keep the tensors in CUDA and perform the computations in `float32`, we see a **200x speed-up** over the initial NumPy implementation on `float32` arrays. @@ -237,7 +237,7 @@ Quansight engineers have designed and proposed a way to support PyTorch tensors within scikit-learn and SciPy. This was received enthusiastically by other maintainers from these libraries, as it was shown that using PyTorch as a backend would often yield considerable speed-ups. -Both projects have now merge initial support for PyTorch tensors across a number of +Both projects have now merged initial support for PyTorch tensors across a number of APIs and submodules. This sets the stepping stone to move towards a future where PyTorch tensors From 49ea0840945053bf339fe384434745ec204a2783 Mon Sep 17 00:00:00 2001 From: lezcano Date: Mon, 18 Sep 2023 12:19:26 +0000 Subject: [PATCH 09/11] wrap paragraphs properly --- blogpost/post.md | 78 +++++++++++++++++++++++++----------------------- 1 file changed, 40 insertions(+), 38 deletions(-) diff --git a/blogpost/post.md b/blogpost/post.md index df2833bb..e70a3a82 100644 --- a/blogpost/post.md +++ b/blogpost/post.md @@ -1,11 +1,10 @@ # Compiling NumPy code into C++ or CUDA via `torch.compile` Quansight engineers have implemented support for tracing through NumPy code via -`torch.compile` in PyTorch 2.1. -This feature leverages PyTorch's compiler to generate efficient fused -vectorized code without having to modify your original NumPy code. Even more, it -also allows for executing NumPy functions on CUDA just by running them through -`torch.compile` under `torch.device("cuda")`! +`torch.compile` in PyTorch 2.1. This feature leverages PyTorch's compiler to +generate efficient fused vectorized code without having to modify your original +NumPy code. Even more, it also allows for executing NumPy functions on CUDA +just by running them through `torch.compile` under `torch.device("cuda")`! In this post, we go over how to use this feature and give a few tips and tricks to make the most out of it. @@ -48,14 +47,16 @@ assert np.allclose(prediction, new_pred) The compiled function yields a 9x speed-up when running it on 1 core. Even better, as opposed to NumPy, our generated code does take advantage of all the -cores in a processor. As such, when we run it on 32 cores, we get a **57x speed-up**. -Note that PyTorch always uses all the available cores unless explicitly restricted, -so this is the default behavior you get when using `torch.compile`. +cores in a processor. As such, when we run it on 32 cores, we get a **57x +speed-up**. Note that PyTorch always uses all the available cores unless +explicitly restricted, so this is the default behavior you get when using +`torch.compile`. We may inspect the generated C++ code by running the script with the environment variable `TORCH_LOGS=output_code`. When doing so, we can see that -`torch.compile` was able to compile the broadcasting and the two -reductions into just one for-loop, and parallelize it using OpenMP +`torch.compile` was able to compile the broadcasting and the two reductions +into just one for-loop, and parallelize it using OpenMP + ```c++ extern "C" void kernel(const double* in_ptr0, const long* in_ptr1, long* out_ptr0) { #pragma omp parallel num_threads(32) @@ -120,11 +121,11 @@ the CPU numbers. This is caused by a small transformation that `torch.compile` does behind the scenes. The code above takes NumPy arrays and returns NumPy arrays. All of these arrays are on CPU, but the computations are performed on the GPU. This means that every time the function is called, `torch.compile` has -to copy all these arrays from CPU to the GPU, and then copy the result -back to CPU to preserve the original semantics. There is no native -solution to this issue in NumPy, as NumPy does not have the notion of a -`device`. That being said, we can work around it by creating a wrapper to this -function so that it accepts PyTorch tensors and returns PyTorch tensors. +to copy all these arrays from CPU to the GPU, and then copy the result back to +CPU to preserve the original semantics. There is no native solution to this +issue in NumPy, as NumPy does not have the notion of a `device`. That being +said, we can work around it by creating a wrapper to this function so that it +accepts PyTorch tensors and returns PyTorch tensors. ```python @torch.compile @@ -167,21 +168,20 @@ page](https://pytorch.org/docs/stable/dynamo/troubleshooting.html#performance-pr This showcases a number of ways to inspect the tracing process, and how to identify problematic code that may cause slowdowns. -**Advice when compiling NumPy code**. NumPy, even if rather similar to -PyTorch, is often used very differently. It is rather common to perform -computations in NumPy and then do an if/else depending on values within the -array, or perform operations in-place, perhaps via boolean masks. These -constructions, while supported by `torch.compile`, hamper its performance. -Changes like moving from in-place indexing to using `np.where`, writing the -code in a branchless way, or avoid using in-place ops in favor of out-of-place -ops can go a long way. +**Advice when compiling NumPy code**. NumPy, even if rather similar to PyTorch, +is often used very differently. It is rather common to perform computations in +NumPy and then do an if/else depending on values within the array, or perform +operations in-place, perhaps via boolean masks. These constructions, while +supported by `torch.compile`, hamper its performance. Changes like moving from +in-place indexing to using `np.where`, writing the code in a branchless way, or +avoid using in-place ops in favor of out-of-place ops can go a long way. To write fast NumPy code, it is best to avoid loops, but sometimes they are unavoidable. When tracing through a loop, `torch.compile` will try to fully unroll it. This is sometimes desirable, but sometimes it may not even be possible, like when we have a dynamic stopping condition, like in a while loop. -In these cases, it may be best to just compile the body of the loop, perhaps -a few iterations at a time (loop unrolling). +In these cases, it may be best to just compile the body of the loop, perhaps a +few iterations at a time (loop unrolling). **Debugging NumPy code**. Debugging is rather tricky when a compiler is involved. To figure out whether an error you are hitting is a `torch.compile` @@ -198,6 +198,7 @@ would return a 0-D tensor (e.g. from `np.sum`). Under `torch.compile`, NumPy scalars are treated as 0-D arrays. This is just fine in most cases. The only case when their behavior diverges is when NumPy scalars are implicitly used as Python scalars. For example, + ```python >>> np.asarray(2) * [1, 2, 3] # 0-D array is an array-like array([2, 4, 6]) @@ -211,6 +212,7 @@ array([2, 4, 6]) # acts as a 0-D array, not as a scalar ?!?! If we compile the first two lines, we see that `torch.compile` treats `u` as a 0-D array. To recover the eager semantics, we just need to make the casting explicit + ```python >>> torch.compile(lambda: int(u) * [1, 2, 3])() [1, 2, 3, 1, 2, 3] @@ -218,6 +220,7 @@ explicit **Type promotion and versioning**. NumPy's type promotion rules may be, at times, a bit surprising + ```python >>> np.asarray([1], dtype=np.int8) + 126 array([127], dtype=int8) @@ -232,17 +235,16 @@ In general, `torch.compile` will match the semantics of the lastest NumPy releas ## Beyond NumPy: SciPy and scikit-learn -In parallel to this effort of making `torch.compile` understand NumPy code, other -Quansight engineers have designed and proposed a way to support PyTorch tensors -within scikit-learn and SciPy. This was received enthusiastically by other -maintainers from these libraries, as it was shown that using PyTorch as a -backend would often yield considerable speed-ups. -Both projects have now merged initial support for PyTorch tensors across a number of -APIs and submodules. - -This sets the stepping stone to move towards a future where PyTorch tensors -can be used within other libraries in the Python data ecosystem. Even more, -this will enable running these other libraries on GPUs and even compiling code +In parallel to this effort of making `torch.compile` understand NumPy code, +other Quansight engineers have designed and proposed a way to support PyTorch +tensors within scikit-learn and SciPy. This was received enthusiastically by +other maintainers from these libraries, as it was shown that using PyTorch as a +backend would often yield considerable speed-ups. Both projects have now merged +initial support for PyTorch tensors across a number of APIs and submodules. + +This sets the stepping stone to move towards a future where PyTorch tensors can +be used within other libraries in the Python data ecosystem. Even more, this +will enable running these other libraries on GPUs and even compiling code mixing these libraries and PyTorch, similar to what we have been discussed in this post. @@ -254,8 +256,8 @@ moving it forward, see this post. [TODO link post] PyTorch has committed since its inception to be a framework compatible with the rest of the Python ecosystem. Enabling compiling NumPy programs, and establishing the tools necessary to do the same for other prominent libraries -are two more steps in this direction. Quansight and Meta continue working hand on -hand, improving the compatibility between PyTorch and the rest of the +are two more steps in this direction. Quansight and Meta continue working hand +on hand, improving the compatibility between PyTorch and the rest of the ecosystem. From Quansight, we would like to thank Mengwei, Voz, and Ed for their From 3b6d73440b667a490f4fa942d09c138164cac7f6 Mon Sep 17 00:00:00 2001 From: lezcano Date: Mon, 18 Sep 2023 12:47:07 +0000 Subject: [PATCH 10/11] google docs suggestion --- blogpost/post.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/blogpost/post.md b/blogpost/post.md index e70a3a82..77b3ae0f 100644 --- a/blogpost/post.md +++ b/blogpost/post.md @@ -85,7 +85,7 @@ rather than generating CUDA code directly, `torch.compile` generates rather readable [triton](https://triton-lang.org/main/index.html) code ```python -def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr): +def triton_(in_ptr0, in_ptr1, out_ptr0, XBLOCK : tl.constexpr): xnumel = 20000000 xoffset = tl.program_id(0) * XBLOCK xindex = xoffset + tl.arange(0, XBLOCK)[:] @@ -174,7 +174,7 @@ NumPy and then do an if/else depending on values within the array, or perform operations in-place, perhaps via boolean masks. These constructions, while supported by `torch.compile`, hamper its performance. Changes like moving from in-place indexing to using `np.where`, writing the code in a branchless way, or -avoid using in-place ops in favor of out-of-place ops can go a long way. +avoiding in-place ops in favor of out-of-place ops can go a long way. To write fast NumPy code, it is best to avoid loops, but sometimes they are unavoidable. When tracing through a loop, `torch.compile` will try to fully @@ -222,10 +222,10 @@ explicit times, a bit surprising ```python ->>> np.asarray([1], dtype=np.int8) + 126 +>>> np.zeros(1, dtype=np.int8) + 127 array([127], dtype=int8) ->>> np.asarray([1], dtype=np.int8) + 128 -array([129], dtype=int16) +>>> np.zeros(1, dtype=np.int8) + 128 +array([128], dtype=int16) ``` NumPy 2.0 is changing these rules to follow others that are closer to those PyTorch. The relevant technical document is [NEP 50](https://numpy.org/neps/nep-0050-scalar-promotion.html). From 8808ab9876f1b6acad92ad054b7ec3738c5c0101 Mon Sep 17 00:00:00 2001 From: lezcano Date: Mon, 18 Sep 2023 12:53:57 +0000 Subject: [PATCH 11/11] Fix var names --- blogpost/post.md | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/blogpost/post.md b/blogpost/post.md index 77b3ae0f..d76566c0 100644 --- a/blogpost/post.md +++ b/blogpost/post.md @@ -31,7 +31,7 @@ npts = 10_000_000 X = np.repeat([[5, 5], [10, 10]], [npts, npts], axis=0) X = X + np.random.randn(*X.shape) # 2 distinct "blobs" means = np.array([[5, 5], [10, 10]]) -pred = get_labels(X, means) +np_pred = get_labels(X, means) ``` Benchmarking this function gives us a baseline of **1.26s** on an AMD 3970X CPU. @@ -40,9 +40,11 @@ Compiling this function is now as easy as wrapping it with `torch.compile` and executing it with the example inputs ```python +import torch + compiled_fn = torch.compile(get_labels) -new_pred = compiled_fn(X, means) -assert np.allclose(prediction, new_pred) +torch_pred = compiled_fn(X, means) +assert np.allclose(np_pred, torch_pred) ``` The compiled function yields a 9x speed-up when running it on 1 core. Even @@ -77,7 +79,7 @@ default device to be CUDA ```python with torch.device("cuda"): cuda_pred = compiled_fn(X, means) -assert np.allclose(prediction, cuda_pred) +assert np.allclose(np_pred, cuda_pred) ``` By inspecting the generated code via `TORCH_LOGS=output_code`, we see that,