__init__.py 1.68 KB
Newer Older
dugupeiwen's avatar
dugupeiwen committed
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
"""
Module that deals with HSA in a high level way
"""
import os
import numba.testing
from .api import *
from .stubs import atomic

def is_available():
    """Returns a boolean to indicate the availability of a HSA runtime.

    This will force initialization of the driver if it hasn't been
    initialized. It also checks that a toolchain is present.
    """
    from .hsadrv.driver import hsa

    from .hlc import hlc, libhlc
    has_a_toolchain = False

    try:
       libhlc.HLC()
       has_a_toolchain = True
    except:
        try:
            cmd = hlc.CmdLine().check_tooling()
            has_a_toolchain = True
        except:
            pass

    return hsa.is_available and has_a_toolchain

if is_available():
    from .hsadrv.driver import hsa
    agents = list(hsa.agents)
else:
    agents = []

38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
# sugon: adapt for numba-0.58, refer to numba/cuda/initialize.py, shoule move to numba/roc/initialize.py.
# TODO: suppot ROCmDispatcher completely.
def initialize_all():

    from numba.roc.decorators import jit
    from numba.core import dispatcher
    from numba.roc.descriptor import HSATargetDesc
    from numba.core.target_extension import (target_registry,
                                             dispatcher_registry,
                                             jit_registry)

    class ROCmDispatcher(dispatcher.Dispatcher):
        targetdescr = HSATargetDesc('ROCm')
    roc_target = target_registry["ROCm"]
    jit_registry[roc_target] = jit
    dispatcher_registry[roc_target] = ROCmDispatcher
initialize_all()

dugupeiwen's avatar
dugupeiwen committed
56
57
58
59
60
def test(*args, **kwargs):
    if not is_available():
        raise RuntimeError("HSA is not detected")

    return numba.testing.test("numba.hsa.tests", *args, **kwargs)