使用 NVIDIA nvmath-python 大规模运行高性能核心数学运算

import numpy as np
import nvmath

m, n, k = 10, 40, 100
a = np.random.randn(m, k)  # a is a NumPy array
b = np.random.randn(k, n)  # b is a NumPy array
c = nvmath.linalg.advanced.matmul(a, b)  # c is also a NumPy array

import cupy as cp
import numpy as np
import nvmath

N = 2048
a_gpu = cp.random.randn(N) + 1j * cp.random.randn(N)
a_cpu = np.random.randn(N) + 1j * np.random.randn(N)
c_gpu = nvmath.fft.fft(a_gpu)
c_cpu = nvmath.fft.fft(a_cpu)

import numpy as np
import nvmath
import logging

logging.basicConfig(level=logging.INFO,
    format="%(asctime)s %(levelname)-8s %(message)s", force=True)
logging.disable(logging.NOTSET)

m, n, k = 8000, 2000, 4000
a_cpu = np.random.randn(m, k).astype(np.float32)
b_cpu = np.random.randn(k, n).astype(np.float32)
d_cpu = nvmath.linalg.advanced.matmul(a_cpu, b_cpu)

2025-09-18 14:53:32,166 INFO     = SPECIFICATION PHASE =
2025-09-18 14:53:32,167 INFO     The data type of operand A is 'float32', and that of operand B is 'float32'.
2025-09-18 14:53:32,168 INFO     The input operands' memory space is cpu, and the execution space is on device 0.
...

import numpy as np
import nvmath
import logging

logging.basicConfig(level=logging.INFO,
    format="%(asctime)s %(levelname)-8s %(message)s", force=True)
logging.disable(logging.NOTSET)

N = 10000
e_cpu = (np.random.randn(N) + 1j * np.random.randn(N)).astype(np.complex64)
r_cpu = nvmath.fft.fft(e_cpu)

2025-09-18 15:46:22,295 INFO     The FFT type is C2C.
2025-09-18 15:46:22,295 INFO     The input data type is complex64, and the result data type is complex64.
2025-09-18 15:46:22,296 INFO     The specified FFT axes are (0,).
2025-09-18 15:46:22,297 INFO     The input tensor's memory space is cpu, and the execution space is cpu, with device cpu.
2025-09-18 15:46:22,298 INFO     The specified stream for the FFT ctor is None.
...

import cupy as cp
import nvmath

m, n, k = 10_000_000, 40, 10
a = cp.random.randn(m, k, dtype=cp.float32)
b = cp.random.randn(k, n, dtype=cp.float32)
c = cp.random.randn(m, n, dtype=cp.float32)
alpha, beta = 1.5, 0.5

d1 = alpha * cp.matmul(a, b) + beta * c # Multiple kernels
d2 = nvmath.linalg.advanced.matmul(a, b, c=c, alpha=alpha, beta=beta) # Single kernel

import nvmath
from nvmath.linalg.advanced import MatmulEpilog
import cupy as cp

feed_count = 10  # The operation feed count.

batch_size = 1024
m, n, k = 1024, 1024, 1024

a = cp.random.rand(batch_size, m, k, dtype=cp.float32)
b = cp.random.rand(batch_size, k, n, dtype=cp.float32)
bias = cp.random.rand(batch_size, m, 1, dtype=cp.float32)

with nvmath.linalg.advanced.Matmul(a, b) as mm:
    # 1. Planning phase
    mm.plan(epilog=MatmulEpilog(MatmulEpilog.RELU_BIAS),
            epilog_inputs={"bias": bias})

    # 2. Autotuning phase
    mm.autotune(iterations=5)

    # 3. Execution phase.
    for i in range(feed_count):
        d = mm.execute()
        # The result of the previous MM is the operand `a` of the next MM, so use
        # reset_operands_unchecked() to reset the `a` operand.
        mm.reset_operands_unchecked(a=d)

from PIL import Image
import nvmath
import cupy as cp

img = cp.asarray(Image.open("your_lovely_dog.jpg").convert("L")) / 255.0 # Gray[0,1]
wh = img.shape[0] * image.shape[1] # We must normalize by the image area
sigma_value = 20.0  # Filter size

# Implement Gaussian filter in the frequency domain
def gaussian_filter(shape, sigma):
    fy = cp.fft.fftfreq(shape[0])[:,None] # Column
    fx = cp.fft.rfftfreq(shape[1])[None,:] # Row
    return = cp.exp(-2.0 * cp.pi * cp.pi * sigma * sigma * (fx * fx + fy * fy))

# Implement FFT epilog wrapper with the pre-defined signature
def epilog_impl(data_out, offset, data, filter_data, unused): # Epilog to be compiled
    data_out[offset] = data * filter_data[offset] / wh

# Compile epilog to LTO-IR targeting the current CUDA device
epilog = nvmath.fft.compile_epilog(epilog_impl, "complex64", "complex64")

# Compute R2C FFT using nvmath-python with the compiled epilog
h_filter = gaussian_filter(img.shape, sigma)
img_fft = nvmath.fft.rfft(image, epilog={"ltoir": epilog, "data": h_filter.data.ptr})

# Compute C2R inverse FFT using nvmath-python
filtered_img = nvmath.fft.irfft(img_fft) # Visualize or save as you want

from numba import cuda
from nvmath.device import random
import cupy as cp
import math

# Pre-compile the RNGs into IR to use alongside other device code
compiled_rng = random.Compile(cc=None)

# GBM parameters
rng_seed = 7777
n_time_steps, n_paths = 252, 8192
mu, sigma, s0 = 0.003, 0.027, 100.0

# Set up CUDA kernel launch configuration
threads_per_block = 32
blocks = n_paths // threads_per_block + bool(n_paths % threads_per_block)
nthreads = threads_per_block * blocks

# RNG initialization kernel
@cuda.jit(link=compiled_rng.files, extensions=compiled_rng.extension)
def init_rng(states, seed):
    idx = cuda.grid(1)
    random.init(seed, idx, 0, states[idx])

# GBM path generation kernel
@cuda.jit(link=compiled_rng.files, extensions=compiled_rng.extension)
def generate_gbm_paths(states, paths, nsteps, mu, sigma, s0):
    idx = cuda.grid(1)
    if idx >= paths.shape[0]:
        return
    paths[idx, 0] = s0

    # Consume 4 normal variates at a time for better throughput
    for i in range(1, nsteps, 4):
        v = random.normal4(states[idx])  # Returned as float32x4 type
        vals = v.x, v.y, v.z, v.w  # Decompose into a tuple of float32
        for j in range(i, min(i + 4, nsteps)):  # Process a chunk of 4 time steps
            paths[idx, j] = paths[idx, j - 1] * math.exp(mu + sigma * vals[j - i])

# Initialize RNG
states = random.StatesPhilox4_32_10(nthreads)
init_rng[blocks, threads_per_block](states, rng_seed)

# Generate GBM paths on GPU
paths = cp.empty((n_paths, n_time_steps), dtype=cp.float32, order='F')
generate_gbm_paths[blocks, threads_per_block](states, paths, n_time_steps, mu, sigma, s0)

pip install nvmath-python[cu13]