Skip to content

Commit eb0b046

Browse files
aterrelaterrel
authored
Merge pull request #9 from samaid/main
Adding nvmath-python notebooks and updated README
2 parents c4377b4 + 1956dba commit eb0b046

File tree

5 files changed

+745
-0
lines changed

5 files changed

+745
-0
lines changed

‎README.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,7 @@
1010
| 6 | Asynchrony: Power Iteration | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pyHPC/pyhpc-tutorial/blob/main/notebooks/6__asynchrony__power_iteration.ipynb) | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pyHPC/pyhpc-tutorial/blob/main/notebooks/6__asynchrony__power_iteration__SOLUTION.ipynb) |
1111
| 7 | Kernel Authoring: Copy | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pyHPC/pyhpc-tutorial/blob/main/notebooks/7__kernel_authoring__copy.ipynb) | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pyHPC/pyhpc-tutorial/blob/main/notebooks/7__kernel_authoring__copy__SOLUTION.ipynb) |
1212
| 8 | MPI | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pyHPC/pyhpc-tutorial/blob/main/notebooks/8__mpi.ipynb) | |
13+
| 9.1 | `nvmath-python`: Interoperability with CPU and GPU tensor libraries | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/samaid/pyhpc-tutorial/blob/main/notebooks/9_1_nvmath-python_interop.ipynb) | |
14+
| 9.2 | `nvmath-python`: Kernel fusion | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/samaid/pyhpc-tutorial/blob/main/notebooks/9_2_nvmath-python_kernel_fusion.ipynb) | |
15+
| 9.3 | `nvmath-python` stateful APIs: Amortizing task preparation costs | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/samaid/pyhpc-tutorial/blob/main/notebooks/9_3_nvmath-python_stateful_apis.ipynb) | |
16+
| 9.4 | `nvmath-python` scaling to many GPUs | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/samaid/pyhpc-tutorial/blob/main/notebooks/9_4_nvmath-python_scaling.ipynb) | |

‎notebooks/9_1_nvmath-python_interop.ipynb‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {
6+
"id": "view-in-github",
7+
"colab_type": "text"
8+
},
9+
"source": [
10+
"<a href=\"https://colab.research.google.com/github/samaid/pyhpc-tutorial/blob/main/notebooks/9_1_nvmath-python_interop.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
11+
]
12+
},
13+
{
14+
"cell_type": "markdown",
15+
"id": "7b236cf1",
16+
"metadata": {
17+
"id": "7b236cf1"
18+
},
19+
"source": [
20+
"# 9.1. `nvmath-python`: Interoperability with CPU and GPU tensor libraries\n",
21+
"The goal of this exercise is to demonstrate how easy it is to plug `nvmath-python` into existing projects that rely on popular CPU or GPU array libraries, such as NumPy, CuPy, and PyTorch, or how easy it is to start a new project where `nvmath-python` is used alongside array libraries."
22+
]
23+
},
24+
{
25+
"cell_type": "markdown",
26+
"id": "e38c312d",
27+
"metadata": {
28+
"id": "e38c312d"
29+
},
30+
"source": [
31+
"### Pure CuPy implementation\n",
32+
"\n",
33+
"This example demonstrates basic matrix multiplication of CuPy 2D arrays using `matmul`:"
34+
]
35+
},
36+
{
37+
"cell_type": "code",
38+
"execution_count": null,
39+
"id": "b796dc7e",
40+
"metadata": {
41+
"id": "b796dc7e"
42+
},
43+
"outputs": [],
44+
"source": [
45+
"import cupy as cp\n",
46+
"\n",
47+
"# Prepare sample input data for matrix matmul\n",
48+
"n, m, k = 2000, 4000, 5000\n",
49+
"a = cp.random.rand(n, k)\n",
50+
"b = cp.random.rand(k, m)\n",
51+
"\n",
52+
"# Perform matrix multiplication\n",
53+
"result = cp.matmul(a, b)\n",
54+
"\n",
55+
"# Print the result\n",
56+
"print(result)\n",
57+
"\n",
58+
"# Print CUDA device for each array\n",
59+
"print(a.device)\n",
60+
"print(b.device)\n",
61+
"print(result.device)"
62+
]
63+
},
64+
{
65+
"cell_type": "markdown",
66+
"id": "7528a6f8",
67+
"metadata": {
68+
"id": "7528a6f8"
69+
},
70+
"source": [
71+
"### Using `nvmath-python` alongside CuPy\n",
72+
"\n",
73+
"This is a slight modification of the above example, where matrix multiplications is done using corresponding `nvmath-python` implementation.\n",
74+
"\n",
75+
"Note that `nvmath-python` supports multiple frameworks, including CuPy. It uses framework's memory pool and the current stream for seamless integration. The result of each operation is a tensor of the same framework that was used to pass the inputs. It is also located on the same device as the inputs."
76+
]
77+
},
78+
{
79+
"cell_type": "code",
80+
"execution_count": null,
81+
"id": "311ee2e9",
82+
"metadata": {
83+
"id": "311ee2e9"
84+
},
85+
"outputs": [],
86+
"source": [
87+
"# The same matrix multiplication as in the previous example but using nvmath-python\n",
88+
"import nvmath\n",
89+
"\n",
90+
"# Perform matrix multiplication\n",
91+
"result = nvmath.linalg.advanced.matmul(a, b)\n",
92+
"\n",
93+
"# Print the result\n",
94+
"print(result)\n",
95+
"\n",
96+
"# Print CUDA device for each array\n",
97+
"print(a.device)\n",
98+
"print(b.device)\n",
99+
"print(result.device)\n"
100+
]
101+
},
102+
{
103+
"cell_type": "markdown",
104+
"id": "85b2ae1b",
105+
"metadata": {
106+
"id": "85b2ae1b"
107+
},
108+
"source": [
109+
"As we can see, the code looks essentially the same. If one measures the performance of above implementations, it will be nearly identical.\n",
110+
"\n",
111+
"This is because CuPy and `nvmath-python` (as well as PyTorch) all use CUDA-X Math Libraries as the engine. It is up to a user, which library to choose for solving the above matrix multiplication problem.\n",
112+
"\n",
113+
"In the next examples we will demonstrate a few examples, where `nvmath-python` may become essential in reaching peak levels of performance."
114+
]
115+
},
116+
{
117+
"cell_type": "code",
118+
"execution_count": null,
119+
"id": "bf34d34d",
120+
"metadata": {
121+
"id": "bf34d34d"
122+
},
123+
"outputs": [],
124+
"source": []
125+
}
126+
],
127+
"metadata": {
128+
"kernelspec": {
129+
"display_name": "nersc-nvmath",
130+
"language": "python",
131+
"name": "python3"
132+
},
133+
"language_info": {
134+
"codemirror_mode": {
135+
"name": "ipython",
136+
"version": 3
137+
},
138+
"file_extension": ".py",
139+
"mimetype": "text/x-python",
140+
"name": "python",
141+
"nbconvert_exporter": "python",
142+
"pygments_lexer": "ipython3",
143+
"version": "3.13.5"
144+
},
145+
"colab": {
146+
"provenance": [],
147+
"include_colab_link": true
148+
}
149+
},
150+
"nbformat": 4,
151+
"nbformat_minor": 5
152+
}

0 commit comments

Comments
 (0)