PyHIP/sample/axpy.py at main · MiloLurati/PyHIP · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
from pyhip import hip, hiprtc
import ctypes
from itertools import repeat
import time

size = 50000000
x = 2
y = 3


def create_array():
    cpu = list()
    gpu = (ctypes.c_int * size)()
    j = 1
    for i in repeat(0, size):
        cpu.append(j)
        gpu[j-1] = j
        j = j + 1
    return cpu, gpu


def cpu_axpy(a):
    result = list()
    j = 0
    for i in repeat(0, size):
        result.append(a[j] * x + y)
        j = j + 1
    return result


def gpu_axpy(res):
    source = """
    #if (HIP_VERSION_MAJOR == 4 && HIP_VERSION_MINOR <= 2)
    #include <hip/hip_runtime.h> // Needed for older ROCm versions (at least 4.2)
    #endif
    extern "C" __global__ void axpy(int *a, int x, int y, size_t size) {
      int i = blockDim.x * blockIdx.x + threadIdx.x;
      if(i < size) {
        a[i] = a[i] * x + y;
      }
    }
    """
    prog = hiprtc.hiprtcCreateProgram(source, 'axpy', [], [])
    device_properties = hip.hipGetDeviceProperties(0)
    print(f"Compiling kernel for {device_properties.gcnArchName}")
    hiprtc.hiprtcCompileProgram(
        prog, [f'--offload-arch={device_properties.gcnArchName}'])
    code = hiprtc.hiprtcGetCode(prog)
    module = hip.hipModuleLoadData(code)
    kernel = hip.hipModuleGetFunction(module, 'axpy')
    ptr = hip.hipMalloc(4 * size)

    class PackageStruct(ctypes.Structure):
        _fields_ = [("a", ctypes.c_void_p),
                    ("x", ctypes.c_int),
                    ("y", ctypes.c_int),
                    ("size", ctypes.c_size_t)]

    hip.hipMemcpy_htod(ptr, ctypes.byref(
        res), ctypes.sizeof(ctypes.c_int) * size)
    struct = PackageStruct(ptr, 2, 3, size)
    block = int(size/1024) + 1
    hip.hipModuleLaunchKernel(kernel, block, 1, 1, 1024, 1, 1, 0, 0, struct)
    hip.hipMemcpy_dtoh(ctypes.byref(res), ptr,
                       ctypes.sizeof(ctypes.c_int) * size)
    hip.hipFree(ptr)
    return res


if __name__ == "__main__":
    cpu_array, gpu_array = create_array()

    start = time.time()
    cpu_result = cpu_axpy(cpu_array)
    cend = time.time()
    gres = gpu_axpy(gpu_array)
    gend = time.time()
    gpu_result = list()

    # Copy back for comparision
    for i in gres:
        gpu_result.append(i)

    if cpu_result != gpu_result:
        print("Error in vaidation: ", cpu_result, gpu_result)
    else:
        print("Passed")
        print("CPU Time: ", (cend - start), " Gpu time: ", (gend - cend))