mirror of
https://github.com/NVIDIA/nvbench.git
synced 2026-03-14 20:27:24 +00:00
b5e4b4ba31
Per PR review suggestion: - `cuda.parallel` - device-wide algorithms/Thrust - `cuda.cooperative` - Cooperative algorithsm/CUB - `cuda.bench` - Benchmarking/NVBench
114 lines
3.3 KiB
Python
114 lines
3.3 KiB
Python
# Copyright 2025 NVIDIA Corporation
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 with the LLVM exception
|
|
# (the "License"); you may not use this file except in compliance with
|
|
# the License.
|
|
#
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://llvm.org/foundation/relicensing/LICENSE.txt
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
|
|
import sys
|
|
|
|
import cuda.bench as bench
|
|
import cuda.bindings.driver as driver
|
|
import cuda.core.experimental as core
|
|
import cupy as cp
|
|
import cutlass
|
|
import numpy as np
|
|
|
|
|
|
def as_bindings_Stream(cs: bench.CudaStream) -> driver.CUstream:
|
|
return driver.CUstream(cs.addressof())
|
|
|
|
|
|
def as_core_Stream(cs: bench.CudaStream) -> core.Stream:
|
|
return core.Stream.from_handle(cs.addressof())
|
|
|
|
|
|
def make_cp_array(
|
|
arr_h: np.ndarray, dev_buf: core.Buffer, dev_id: int | None
|
|
) -> cp.ndarray:
|
|
cp_memview = cp.cuda.UnownedMemory(
|
|
int(dev_buf.handle), dev_buf.size, dev_buf, -1 if dev_id is None else dev_id
|
|
)
|
|
zero_offset = 0
|
|
return cp.ndarray(
|
|
arr_h.shape,
|
|
dtype=arr_h.dtype,
|
|
memptr=cp.cuda.MemoryPointer(cp_memview, zero_offset),
|
|
)
|
|
|
|
|
|
def cutlass_gemm(state: bench.State) -> None:
|
|
n = state.get_int64("N")
|
|
r = state.get_int64("R")
|
|
|
|
alpha = state.get_float64("alpha")
|
|
|
|
dt = np.float64
|
|
A_h = np.random.randn(n, r).astype(dt)
|
|
B_h = np.copy(A_h.mT)
|
|
C_h = np.eye(n, dtype=dt)
|
|
D_h = np.zeros_like(C_h)
|
|
|
|
if n >= 1024:
|
|
# allow more time for large inputs
|
|
state.set_timeout(360)
|
|
|
|
dev_id = state.get_device()
|
|
cs = state.get_stream()
|
|
s = as_bindings_Stream(cs)
|
|
core_s = as_core_Stream(cs)
|
|
|
|
A_d = core.DeviceMemoryResource(dev_id).allocate(A_h.nbytes, core_s)
|
|
B_d = core.DeviceMemoryResource(dev_id).allocate(B_h.nbytes, core_s)
|
|
C_d = core.DeviceMemoryResource(dev_id).allocate(C_h.nbytes, core_s)
|
|
D_d = core.DeviceMemoryResource(dev_id).allocate(D_h.nbytes, core_s)
|
|
|
|
driver.cuMemcpyAsync(A_d.handle, A_h.ctypes.data, A_h.nbytes, s)
|
|
driver.cuMemcpyAsync(B_d.handle, B_h.ctypes.data, B_h.nbytes, s)
|
|
driver.cuMemcpyAsync(C_d.handle, C_h.ctypes.data, C_h.nbytes, s)
|
|
driver.cuMemcpyAsync(D_d.handle, D_h.ctypes.data, D_h.nbytes, s)
|
|
|
|
A_cp = make_cp_array(A_h, A_d, dev_id)
|
|
B_cp = make_cp_array(B_h, B_d, dev_id)
|
|
C_cp = make_cp_array(C_h, C_d, dev_id)
|
|
D_cp = make_cp_array(D_h, D_d, dev_id)
|
|
|
|
plan = cutlass.op.Gemm(
|
|
A=A_cp,
|
|
B=B_cp,
|
|
C=C_cp,
|
|
D=D_cp,
|
|
element=dt,
|
|
alpha=alpha,
|
|
beta=1,
|
|
layout=cutlass.LayoutType.RowMajor,
|
|
)
|
|
# warm-up to ensure compilation is not timed
|
|
plan.run(stream=s)
|
|
|
|
def launcher(launch: bench.Launch) -> None:
|
|
s = as_bindings_Stream(launch.get_stream())
|
|
plan.run(stream=s, sync=False)
|
|
|
|
state.exec(launcher)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
gemm_b = bench.register(cutlass_gemm)
|
|
gemm_b.add_int64_axis("R", [16, 64, 256])
|
|
gemm_b.add_int64_axis("N", [256, 512, 1024, 2048])
|
|
|
|
gemm_b.add_float64_axis("alpha", [1e-2])
|
|
|
|
bench.run_all_benchmarks(sys.argv)
|