Skip to content

Commit 7915feb

Browse files
samaidsamaid
authored
Add files via upload
1 parent c4377b4 commit 7915feb

4 files changed

Lines changed: 638 additions & 0 deletions

‎notebooks/9_1_nvmath-python_interop.ipynb‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "7b236cf1",
6+
"metadata": {},
7+
"source": [
8+
"# 9.1. `nvmath-python`: Interoperability with CPU and GPU tensor libraries\n",
9+
"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."
10+
]
11+
},
12+
{
13+
"cell_type": "markdown",
14+
"id": "e38c312d",
15+
"metadata": {},
16+
"source": [
17+
"### Pure CuPy implementation\n",
18+
"\n",
19+
"This example demonstrates basic matrix multiplication of CuPy 2D arrays using `matmul`:"
20+
]
21+
},
22+
{
23+
"cell_type": "code",
24+
"execution_count": null,
25+
"id": "b796dc7e",
26+
"metadata": {},
27+
"outputs": [],
28+
"source": [
29+
"import cupy as cp\n",
30+
"\n",
31+
"# Prepare sample input data for matrix matmul\n",
32+
"n, m, k = 2000, 4000, 5000\n",
33+
"a = cp.random.rand(n, k)\n",
34+
"b = cp.random.rand(k, m)\n",
35+
"\n",
36+
"# Perform matrix multiplication\n",
37+
"result = cp.matmul(a, b)\n",
38+
"\n",
39+
"# Print the result\n",
40+
"print(result)\n",
41+
"\n",
42+
"# Print CUDA device for each array\n",
43+
"print(a.device)\n",
44+
"print(b.device)\n",
45+
"print(result.device)"
46+
]
47+
},
48+
{
49+
"cell_type": "markdown",
50+
"id": "7528a6f8",
51+
"metadata": {},
52+
"source": [
53+
"### Using `nvmath-python` alongside CuPy\n",
54+
"\n",
55+
"This is a slight modification of the above example, where matrix multiplications is done using corresponding `nvmath-python` implementation.\n",
56+
"\n",
57+
"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. "
58+
]
59+
},
60+
{
61+
"cell_type": "code",
62+
"execution_count": null,
63+
"id": "311ee2e9",
64+
"metadata": {},
65+
"outputs": [],
66+
"source": [
67+
"# The same matrix multiplication as in the previous example but using nvmath-python\n",
68+
"import nvmath\n",
69+
"\n",
70+
"# Perform matrix multiplication\n",
71+
"result = nvmath.linalg.advanced.matmul(a, b)\n",
72+
"\n",
73+
"# Print the result\n",
74+
"print(result)\n",
75+
"\n",
76+
"# Print CUDA device for each array\n",
77+
"print(a.device)\n",
78+
"print(b.device)\n",
79+
"print(result.device)\n"
80+
]
81+
},
82+
{
83+
"cell_type": "markdown",
84+
"id": "85b2ae1b",
85+
"metadata": {},
86+
"source": [
87+
"As we can see, the code looks essentially the same. If one measures the performance of above implementations, it will be nearly identical. \n",
88+
"\n",
89+
"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",
90+
"\n",
91+
"In the next examples we will demonstrate a few examples, where `nvmath-python` may become essential in reaching peak levels of performance. "
92+
]
93+
},
94+
{
95+
"cell_type": "code",
96+
"execution_count": null,
97+
"id": "bf34d34d",
98+
"metadata": {},
99+
"outputs": [],
100+
"source": []
101+
}
102+
],
103+
"metadata": {
104+
"kernelspec": {
105+
"display_name": "nersc-nvmath",
106+
"language": "python",
107+
"name": "python3"
108+
},
109+
"language_info": {
110+
"codemirror_mode": {
111+
"name": "ipython",
112+
"version": 3
113+
},
114+
"file_extension": ".py",
115+
"mimetype": "text/x-python",
116+
"name": "python",
117+
"nbconvert_exporter": "python",
118+
"pygments_lexer": "ipython3",
119+
"version": "3.13.5"
120+
}
121+
},
122+
"nbformat": 4,
123+
"nbformat_minor": 5
124+
}

‎notebooks/9_2_nvmath-python_kernel_fusion.ipynb‎

Lines changed: 233 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,233 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "e93afc39",
6+
"metadata": {},
7+
"source": [
8+
"# 9.2. `nvmath-python`: Kernel fusion\n",
9+
"\n",
10+
"Some computational problems have lower compute to memory access instructions ratio, which can often be addressed by lowering the number of memory accesses via kernel fusion. This exercise illustrate how `nvmath-python` inherently fuses simpler operations into a single composite kernel.\n",
11+
"\n",
12+
"We illustrate this on the example of neural network propagation and the fast Fourier transfer example.\n",
13+
"\n",
14+
"## Advanced `matmul` with bias and epilog\n",
15+
"Based on **Exercise 9.1**, it is not clear why one would need to use `nvmath-python` for matrix multiplications. Indeed, for basic `matmul` operation using `nvmath-python` alongside CuPy does seem an overkill. However, in scientific computing applications and AI, `matmul`s are often used in combination with other operations. For example, in neural networks quite a common usage pattern is as follows.\n",
16+
"\n",
17+
"**Matrix-Matrix Multiplication with Bias and ReLU:**\n",
18+
"\n",
19+
"$$C = \\text{ReLU}(A \\cdot B + b^T)$$\n",
20+
"\n",
21+
"where:\n",
22+
"- $A \\in \\mathbb{R}^{m \\times k}$ is the input matrix\n",
23+
"- $B \\in \\mathbb{R}^{k \\times n}$ is the weight matrix \n",
24+
"- $b \\in \\mathbb{R}^{m}$ is the bias vector (transposed and broadcasted to $m \\times n$)\n",
25+
"- $\\text{ReLU}(x) = \\max(0, x)$ is the Rectified Linear Unit activation function\n",
26+
"- $C \\in \\mathbb{R}^{m \\times n}$ is the output matrix\n",
27+
"\n",
28+
"The bias vector $b$ is transposed to $b^T \\in \\mathbb{R}^{m \\times 1}$ and automatically broadcasted across all columns of the result matrix. The ReLU function is applied element-wise to the result of the matrix multiplication plus bias.\n",
29+
"\n"
30+
]
31+
},
32+
{
33+
"cell_type": "code",
34+
"execution_count": null,
35+
"id": "317b5ac3",
36+
"metadata": {},
37+
"outputs": [],
38+
"source": [
39+
"import cupy as cp\n",
40+
"\n",
41+
"# Define the ReLU function\n",
42+
"def relu(x):\n",
43+
" return cp.maximum(0, x)\n",
44+
"\n",
45+
"# Matrix dimensions\n",
46+
"# m, k, n = 2, 5, 4\n",
47+
"m, k, n = 2000, 1000, 4000\n",
48+
"\n",
49+
"# Create input matrix A (m x k)\n",
50+
"A = cp.random.randn(m, k)\n",
51+
"\n",
52+
"# Create weight matrix B (k x n) \n",
53+
"B = cp.random.randn(k, n)\n",
54+
"\n",
55+
"# Create bias vector b (m,) that will be transposed and broadcasted to (m x n)\n",
56+
"b = cp.random.randn(m)\n",
57+
"\n",
58+
"# Implement the formula: C = ReLU(A * B + b^T)\n",
59+
"# Kernel 1: Matrix multiplication\n",
60+
"matmul_result = cp.matmul(A, B)\n",
61+
"\n",
62+
"# Kernel 2: Add bias (broadcasting happens automatically)\n",
63+
"bias_result = matmul_result + b.reshape(-1, 1)\n",
64+
"\n",
65+
"# Kernel 3: Apply ReLU activation\n",
66+
"C = relu(bias_result)\n",
67+
"\n",
68+
"# print(f\"A: {A}\")\n",
69+
"# print(f\"B: {B}\")\n",
70+
"# print(f\"b: {b}\")\n",
71+
"# print(f\"C: {C}\")\n",
72+
"\n",
73+
"\n",
74+
"print(f\"Input matrix A shape: {A.shape}\")\n",
75+
"print(f\"Weight matrix B shape: {B.shape}\") \n",
76+
"print(f\"Bias vector b shape: {b.shape}\")\n",
77+
"print(f\"Transposed bias b^T shape: {b.reshape(-1, 1).shape}\")\n",
78+
"print(f\"Output matrix C shape: {C.shape}\")\n",
79+
"print(f\"Output matrix C device: {C.device}\")"
80+
]
81+
},
82+
{
83+
"cell_type": "markdown",
84+
"id": "39c03a45",
85+
"metadata": {},
86+
"source": [
87+
"**TODO: Validate the above implementation on small easy to comprehend inputs by manually initializing matrices and bias and by printing results step by step**\n",
88+
"\n",
89+
"`nvmath-python` leverages the power of `cuBLASLt` library that provides variety of options to implement such computational patterns. Here's an example of how one can implement the above example using `nvmath-python`:"
90+
]
91+
},
92+
{
93+
"cell_type": "code",
94+
"execution_count": null,
95+
"id": "af245b0d",
96+
"metadata": {},
97+
"outputs": [],
98+
"source": [
99+
"import nvmath\n",
100+
"\n",
101+
"# Kernel 1, 2, and 3 are fused into a single kernel\n",
102+
"C = nvmath.linalg.advanced.matmul(A, B, epilog=nvmath.linalg.advanced.MatmulEpilog.RELU_BIAS, epilog_inputs={\"bias\": b})\n",
103+
"\n",
104+
"# print(f\"A: {A}\")\n",
105+
"# print(f\"B: {B}\")\n",
106+
"# print(f\"b: {b}\")\n",
107+
"# print(f\"C: {C}\")\n",
108+
"\n",
109+
"print(C.shape)\n",
110+
"print(C.device)"
111+
]
112+
},
113+
{
114+
"cell_type": "markdown",
115+
"id": "746a4d0e",
116+
"metadata": {},
117+
"source": [
118+
"**TODO: Ensure that `nvmath-python` results are identical for the same small inputs**\n",
119+
"\n",
120+
"All three kernels are fused into a single kernel using JIT machinery behind `cuBLASLt`, which in certain problem settings may result in overall better performance due to better compute-to-memory access ratio"
121+
]
122+
},
123+
{
124+
"cell_type": "markdown",
125+
"id": "114319f3",
126+
"metadata": {},
127+
"source": [
128+
"## Using custom FFT callbacks written in Python\n",
129+
"\n",
130+
"In previous example the epilog is a predefined set of activation functions and their gradients. This example illustrates the case when the epilog is a custom Python function, which is compiled into internal intermediate representation (LTO-IR) and then fused with the FFT operation into a single kernel. \n",
131+
"\n",
132+
"Specifically, we illustrate how to perform a convolution by providing a Python callback function as an epilog to the FFT operation.\n",
133+
"\n",
134+
"To begin with, let's create some input data. We will use the batched 1D FFT and apply a sine-form filter in the frequency domain."
135+
]
136+
},
137+
{
138+
"cell_type": "code",
139+
"execution_count": null,
140+
"id": "cbea6060",
141+
"metadata": {},
142+
"outputs": [],
143+
"source": [
144+
"# Create the data for the batched 1-D FFT.\n",
145+
"B, N = 256, 1024\n",
146+
"a = cp.random.rand(B, N) + 1j * cp.random.rand(B, N)\n",
147+
"\n",
148+
"# Create the data to use as a filter.\n",
149+
"filter_data = cp.sin(a)"
150+
]
151+
},
152+
{
153+
"cell_type": "markdown",
154+
"id": "e4211da9",
155+
"metadata": {},
156+
"source": [
157+
"We also define the epilog function for forward FFT, a convolution, which corresponds to pointwise multiplication in the frequency domain. We also scale by the FFT size `N` here."
158+
]
159+
},
160+
{
161+
"cell_type": "code",
162+
"execution_count": null,
163+
"id": "e1e27f3c",
164+
"metadata": {},
165+
"outputs": [],
166+
"source": [
167+
"def convolve(data_out, offset, data, filter_data, unused):\n",
168+
" data_out[offset] = data * filter_data[offset] / N"
169+
]
170+
},
171+
{
172+
"cell_type": "markdown",
173+
"id": "c39d8361",
174+
"metadata": {},
175+
"source": [
176+
"Note we are accessing `data_out` and `filter_data` with a single `offset` integer, even though the output and `filter_data` are 2D tensors (batches of samples). Care must be taken to ensure that both arrays accessed here have the same memory layout.\n",
177+
"\n",
178+
"Next thing is to compile the epilog to intermediate representation (LTO-IR). In a system with GPUs that have different compute capability, the `compute_capability` option must be specified to the `compile_prolog` or `compile_epilog` helpers. Alternatively, the epilog can be compiled in the context of the device where the FFT to which the epilog is provided is executed. In this case we use the current device context, where the operands have been created:"
179+
]
180+
},
181+
{
182+
"cell_type": "code",
183+
"execution_count": null,
184+
"id": "1eda86b7",
185+
"metadata": {},
186+
"outputs": [],
187+
"source": [
188+
"with cp.cuda.Device():\n",
189+
" epilog = nvmath.fft.compile_epilog(convolve, \"complex128\", \"complex128\")"
190+
]
191+
},
192+
{
193+
"cell_type": "markdown",
194+
"id": "4e8418ea",
195+
"metadata": {},
196+
"source": [
197+
"Finally, we perform the convolution as the forward FFT with the compiled epilog (filter) followed by the inverse FFT transformation:"
198+
]
199+
},
200+
{
201+
"cell_type": "code",
202+
"execution_count": null,
203+
"id": "3c0e951b",
204+
"metadata": {},
205+
"outputs": [],
206+
"source": [
207+
"r = nvmath.fft.fft(a, axes=[-1], epilog={\"ltoir\": epilog, \"data\": filter_data.data.ptr})\n",
208+
"r = nvmath.fft.ifft(r, axes=[-1])"
209+
]
210+
}
211+
],
212+
"metadata": {
213+
"kernelspec": {
214+
"display_name": "nersc-nvmath",
215+
"language": "python",
216+
"name": "python3"
217+
},
218+
"language_info": {
219+
"codemirror_mode": {
220+
"name": "ipython",
221+
"version": 3
222+
},
223+
"file_extension": ".py",
224+
"mimetype": "text/x-python",
225+
"name": "python",
226+
"nbconvert_exporter": "python",
227+
"pygments_lexer": "ipython3",
228+
"version": "3.13.5"
229+
}
230+
},
231+
"nbformat": 4,
232+
"nbformat_minor": 5
233+
}

0 commit comments

Comments
 (0)