source: sasmodels/sasmodels/kernelcuda.py @ b297ba9

core_shell_microgelsmagnetic_modelticket-1257-vesicle-productticket_1156ticket_1265_superballticket_822_more_unit_tests
Last change on this file since b297ba9 was b297ba9, checked in by Paul Kienzle <pkienzle@…>, 5 years ago

lint

  • Property mode set to 100644
File size: 20.5 KB
Line 
1"""
2GPU driver for C kernels (with CUDA)
3
4To select cuda, use SAS_OPENCL=cuda, or SAS_OPENCL=cuda:n for a particular
5device number.  If no device number is specified, then look for CUDA_DEVICE=n
6or a file ~/.cuda-device containing n for the device number.  Otherwise, try
7all available device numbers.
8
9TODO: docs are out of date
10
11There should be a single GPU environment running on the system.  This
12environment is constructed on the first call to :func:`env`, and the
13same environment is returned on each call.
14
15After retrieving the environment, the next step is to create the kernel.
16This is done with a call to :meth:`GpuEnvironment.make_kernel`, which
17returns the type of data used by the kernel.
18
19Next a :class:`GpuData` object should be created with the correct kind
20of data.  This data object can be used by multiple kernels, for example,
21if the target model is a weighted sum of multiple kernels.  The data
22should include any extra evaluation points required to compute the proper
23data smearing.  This need not match the square grid for 2D data if there
24is an index saying which q points are active.
25
26Together the GpuData, the program, and a device form a :class:`GpuKernel`.
27This kernel is used during fitting, receiving new sets of parameters and
28evaluating them.  The output value is stored in an output buffer on the
29devices, where it can be combined with other structure factors and form
30factors and have instrumental resolution effects applied.
31
32In order to use OpenCL for your models, you will need OpenCL drivers for
33your machine.  These should be available from your graphics card vendor.
34Intel provides OpenCL drivers for CPUs as well as their integrated HD
35graphics chipsets.  AMD also provides drivers for Intel CPUs, but as of
36this writing the performance is lacking compared to the Intel drivers.
37NVidia combines drivers for CUDA and OpenCL in one package.  The result
38is a bit messy if you have multiple drivers installed.  You can see which
39drivers are available by starting python and running:
40
41    import pyopencl as cl
42    cl.create_some_context(interactive=True)
43
44Once you have done that, it will show the available drivers which you
45can select.  It will then tell you that you can use these drivers
46automatically by setting the SAS_OPENCL environment variable, which is
47PYOPENCL_CTX equivalent but not conflicting with other pyopnecl programs.
48
49Some graphics cards have multiple devices on the same card.  You cannot
50yet use both of them concurrently to evaluate models, but you can run
51the program twice using a different device for each session.
52
53OpenCL kernels are compiled when needed by the device driver.  Some
54drivers produce compiler output even when there is no error.  You
55can see the output by setting PYOPENCL_COMPILER_OUTPUT=1.  It should be
56harmless, albeit annoying.
57"""
58from __future__ import print_function
59
60import os
61import logging
62import time
63import re
64import atexit
65
66import numpy as np  # type: ignore
67
68
69# Attempt to setup CUDA. This may fail if the pycuda package is not
70# installed or if it is installed but there are no devices available.
71try:
72    import pycuda.driver as cuda  # type: ignore
73    from pycuda.compiler import SourceModule
74    from pycuda.tools import make_default_context, clear_context_caches
75    # Ask CUDA for the default context (so that we know that one exists)
76    # then immediately throw it away in case the user doesn't want it.
77    # Note: cribbed from pycuda.autoinit
78    cuda.init()
79    context = make_default_context()
80    context.pop()
81    clear_context_caches()
82    del context
83    HAVE_CUDA = True
84    CUDA_ERROR = ""
85except Exception as exc:
86    HAVE_CUDA = False
87    CUDA_ERROR = str(exc)
88
89from . import generate
90from .kernel import KernelModel, Kernel
91
92# pylint: disable=unused-import
93try:
94    from typing import Tuple, Callable, Any
95    from .modelinfo import ModelInfo
96    from .details import CallDetails
97except ImportError:
98    pass
99# pylint: enable=unused-import
100
101# The max loops number is limited by the amount of local memory available
102# on the device.  You don't want to make this value too big because it will
103# waste resources, nor too small because it may interfere with users trying
104# to do their polydispersity calculations.  A value of 1024 should be much
105# larger than necessary given that cost grows as npts^k where k is the number
106# of polydisperse parameters.
107MAX_LOOPS = 2048
108
109
110def use_cuda():
111    # type: None -> bool
112    """Returns True if CUDA is the default compute engine."""
113    sas_opencl = os.environ.get("SAS_OPENCL", "CUDA").lower()
114    return HAVE_CUDA and sas_opencl.startswith("cuda")
115
116
117ENV = None
118def reset_environment():
119    """
120    Call to create a new OpenCL context, such as after a change to SAS_OPENCL.
121    """
122    global ENV
123    # Free any previous allocated context.
124    if ENV is not None and ENV.context is not None:
125        ENV.release()
126    ENV = GpuEnvironment() if use_cuda() else None
127
128
129def environment():
130    # type: () -> "GpuEnvironment"
131    """
132    Returns a singleton :class:`GpuEnvironment`.
133
134    This provides an OpenCL context and one queue per device.
135    """
136    if ENV is None:
137        if not HAVE_CUDA:
138            raise RuntimeError("CUDA startup failed with ***"
139                               + CUDA_ERROR + "***; using C compiler instead")
140        reset_environment()
141        if ENV is None:
142            raise RuntimeError("SAS_OPENCL=None in environment")
143    return ENV
144
145
146# PyTest is not freeing ENV, so make sure it gets freed.
147atexit.register(lambda: ENV.release() if ENV is not None else None)
148
149
150def has_type(dtype):
151    # type: (np.dtype) -> bool
152    """
153    Return true if device supports the requested precision.
154    """
155    # Assume the NVIDIA card supports 32-bit and 64-bit floats.
156    # TODO: Check if pycuda support F16.
157    return dtype in (generate.F32, generate.F64)
158
159
160FUNCTION_PATTERN = re.compile(r"""^
161  (?P<space>\s*)                       # Initial space.
162  (?P<qualifiers>^(?:\s*\b\w+\b\s*)+)  # One or more qualifiers before function.
163  (?P<function>\s*\b\w+\b\s*[(])       # Function name plus open parens.
164  """, re.VERBOSE|re.MULTILINE)
165
166MARKED_PATTERN = re.compile(r"""
167  \b(return|else|kernel|device|__device__)\b
168  """, re.VERBOSE|re.MULTILINE)
169
170
171def _add_device_tag(match):
172    # type: (None) -> str
173    # Note: Should be re.Match, but that isn't a simple type.
174    """
175    replace qualifiers with __device__ qualifiers if needed
176    """
177    qualifiers = match.group("qualifiers")
178    if MARKED_PATTERN.search(qualifiers):
179        start, end = match.span()
180        return match.string[start:end]
181    else:
182        function = match.group("function")
183        space = match.group("space")
184        return "".join((space, "__device__ ", qualifiers, function))
185
186
187def mark_device_functions(source):
188    # type: (str) -> str
189    """
190    Mark all function declarations as __device__ functions (except kernel).
191    """
192    return FUNCTION_PATTERN.sub(_add_device_tag, source)
193
194
195def show_device_functions(source):
196    # type: (str) -> str
197    """
198    Show all discovered function declarations, but don't change any.
199    """
200    for match in FUNCTION_PATTERN.finditer(source):
201        print(match.group('qualifiers').replace('\n', r'\n'),
202              match.group('function'), '(')
203    return source
204
205
206def compile_model(source, dtype, fast=False):
207    # type: (str, np.dtype, bool) -> SourceModule
208    """
209    Build a model to run on the gpu.
210
211    Returns the compiled program and its type.  The returned type will
212    be float32 even if the desired type is float64 if any of the
213    devices in the context do not support the cl_khr_fp64 extension.
214    """
215    dtype = np.dtype(dtype)
216    if not has_type(dtype):
217        raise RuntimeError("%s not supported for devices"%dtype)
218
219    source_list = [generate.convert_type(source, dtype)]
220
221    source_list.insert(0, "#define USE_SINCOS\n")
222    source = "\n".join(source_list)
223    #source = show_device_functions(source)
224    source = mark_device_functions(source)
225    #with open('/tmp/kernel.cu', 'w') as fd: fd.write(source)
226    #print(source)
227    #options = ['--verbose', '-E']
228    options = ['--use_fast_math'] if fast else None
229    program = SourceModule(source, no_extern_c=True, options=options) #, include_dirs=[...])
230
231    #print("done with "+program)
232    return program
233
234
235# For now, this returns one device in the context.
236# TODO: Create a context that contains all devices on all platforms.
237class GpuEnvironment(object):
238    """
239    GPU context for CUDA.
240    """
241    context = None # type: cuda.Context
242    def __init__(self, devnum=None):
243        # type: (int) -> None
244        env = os.environ.get("SAS_OPENCL", "").lower()
245        if devnum is None and env.startswith("cuda:"):
246            devnum = int(env[5:])
247
248        # Set the global context to the particular device number if one is
249        # given, otherwise use the default context.  Perhaps this will be set
250        # by an environment variable within autoinit.
251        if devnum is not None:
252            self.context = cuda.Device(devnum).make_context()
253        else:
254            self.context = make_default_context()
255
256        ## Byte boundary for data alignment.
257        #self.data_boundary = max(d.min_data_type_align_size
258        #                         for d in self.context.devices)
259
260        # Cache for compiled programs, and for items in context.
261        self.compiled = {}
262
263    def release(self):
264        """Free the CUDA device associated with this context."""
265        if self.context is not None:
266            self.context.pop()
267            self.context = None
268
269    def __del__(self):
270        self.release()
271
272    def has_type(self, dtype):
273        # type: (np.dtype) -> bool
274        """
275        Return True if all devices support a given type.
276        """
277        return has_type(dtype)
278
279    def compile_program(self, name, source, dtype, fast, timestamp):
280        # type: (str, str, np.dtype, bool, float) -> cl.Program
281        """
282        Compile the program for the device in the given context.
283        """
284        # Note: PyCuda (probably) caches but I'll do so as well just to
285        # save some data munging time.
286        tag = generate.tag_source(source)
287        key = "%s-%s-%s%s"%(name, dtype, tag, ("-fast" if fast else ""))
288        # Check timestamp on program.
289        program, program_timestamp = self.compiled.get(key, (None, np.inf))
290        if program_timestamp < timestamp:
291            del self.compiled[key]
292        if key not in self.compiled:
293            logging.info("building %s for CUDA", key)
294            program = compile_model(str(source), dtype, fast)
295            self.compiled[key] = (program, timestamp)
296        return program
297
298
299class GpuModel(KernelModel):
300    """
301    GPU wrapper for a single model.
302
303    *source* and *model_info* are the model source and interface as returned
304    from :func:`generate.make_source` and :func:`generate.make_model_info`.
305
306    *dtype* is the desired model precision.  Any numpy dtype for single
307    or double precision floats will do, such as 'f', 'float32' or 'single'
308    for single and 'd', 'float64' or 'double' for double.  Double precision
309    is an optional extension which may not be available on all devices.
310    Half precision ('float16','half') may be available on some devices.
311    Fast precision ('fast') is a loose version of single precision, indicating
312    that the compiler is allowed to take shortcuts.
313    """
314    info = None  # type: ModelInfo
315    source = ""  # type: str
316    dtype = None  # type: np.dtype
317    fast = False  # type: bool
318    _program = None  # type: SourceModule
319    _kernels = None  # type: Dict[str, cuda.Function]
320
321    def __init__(self, source, model_info, dtype=generate.F32, fast=False):
322        # type: (Dict[str,str], ModelInfo, np.dtype, bool) -> None
323        self.info = model_info
324        self.source = source
325        self.dtype = dtype
326        self.fast = fast
327
328    def __getstate__(self):
329        # type: () -> Tuple[ModelInfo, str, np.dtype, bool]
330        return self.info, self.source, self.dtype, self.fast
331
332    def __setstate__(self, state):
333        # type: (Tuple[ModelInfo, str, np.dtype, bool]) -> None
334        self.info, self.source, self.dtype, self.fast = state
335        self._program = self._kernels = None
336
337    def make_kernel(self, q_vectors):
338        # type: (List[np.ndarray]) -> "GpuKernel"
339        return GpuKernel(self, q_vectors)
340
341    def get_function(self, name):
342        # type: (str) -> cuda.Function
343        """
344        Fetch the kernel from the environment by name, compiling it if it
345        does not already exist.
346        """
347        if self._program is None:
348            self._prepare_program()
349        return self._kernels[name]
350
351    def _prepare_program(self):
352        # type: (str) -> None
353        env = environment()
354        timestamp = generate.ocl_timestamp(self.info)
355        program = env.compile_program(
356            self.info.name,
357            self.source['opencl'],
358            self.dtype,
359            self.fast,
360            timestamp)
361        variants = ['Iq', 'Iqxy', 'Imagnetic']
362        names = [generate.kernel_name(self.info, k) for k in variants]
363        functions = [program.get_function(k) for k in names]
364        self._kernels = {k: v for k, v in zip(variants, functions)}
365        # Keep a handle to program so GC doesn't collect.
366        self._program = program
367
368
369# TODO: Check that we don't need a destructor for buffers which go out of scope.
370class GpuInput(object):
371    """
372    Make q data available to the gpu.
373
374    *q_vectors* is a list of q vectors, which will be *[q]* for 1-D data,
375    and *[qx, qy]* for 2-D data.  Internally, the vectors will be reallocated
376    to get the best performance on OpenCL, which may involve shifting and
377    stretching the array to better match the memory architecture.  Additional
378    points will be evaluated with *q=1e-3*.
379
380    *dtype* is the data type for the q vectors. The data type should be
381    set to match that of the kernel, which is an attribute of
382    :class:`GpuProgram`.  Note that not all kernels support double
383    precision, so even if the program was created for double precision,
384    the *GpuProgram.dtype* may be single precision.
385
386    Call :meth:`release` when complete.  Even if not called directly, the
387    buffer will be released when the data object is freed.
388    """
389    def __init__(self, q_vectors, dtype=generate.F32):
390        # type: (List[np.ndarray], np.dtype) -> None
391        # TODO: Do we ever need double precision q?
392        self.nq = q_vectors[0].size
393        self.dtype = np.dtype(dtype)
394        self.is_2d = (len(q_vectors) == 2)
395        # TODO: Stretch input based on get_warp().
396        # Not doing it now since warp depends on kernel, which is not known
397        # at this point, so instead using 32, which is good on the set of
398        # architectures tested so far.
399        if self.is_2d:
400            width = ((self.nq+15)//16)*16
401            self.q = np.empty((width, 2), dtype=dtype)
402            self.q[:self.nq, 0] = q_vectors[0]
403            self.q[:self.nq, 1] = q_vectors[1]
404        else:
405            width = ((self.nq+31)//32)*32
406            self.q = np.empty(width, dtype=dtype)
407            self.q[:self.nq] = q_vectors[0]
408        self.global_size = [self.q.shape[0]]
409        #print("creating inputs of size", self.global_size)
410
411        # Transfer input value to GPU.
412        self.q_b = cuda.to_device(self.q)
413
414    def release(self):
415        # type: () -> None
416        """
417        Free the buffer associated with the q value.
418        """
419        if self.q_b is not None:
420            self.q_b.free()
421            self.q_b = None
422
423    def __del__(self):
424        # type: () -> None
425        self.release()
426
427
428class GpuKernel(Kernel):
429    """
430    Callable SAS kernel.
431
432    *model* is the GpuModel object to call
433
434    The kernel is derived from :class:`Kernel`, providing the
435    :meth:`call_kernel` method to evaluate the kernel for a given set of
436    parameters.  Because of the need to move the q values to the GPU before
437    evaluation, the kernel is instantiated for a particular set of q vectors,
438    and can be called many times without transfering q each time.
439
440    Call :meth:`release` when done with the kernel instance.
441    """
442    #: SAS model information structure.
443    info = None  # type: ModelInfo
444    #: Kernel precision.
445    dtype = None  # type: np.dtype
446    #: Kernel dimensions (1d or 2d).
447    dim = ""  # type: str
448    #: Calculation results, updated after each call to :meth:`_call_kernel`.
449    result = None  # type: np.ndarray
450
451    def __init__(self, model, q_vectors):
452        # type: (GpuModel, List[np.ndarray]) -> None
453        dtype = model.dtype
454        self.q_input = GpuInput(q_vectors, dtype)
455        self._model = model
456
457        # Attributes accessed from the outside.
458        self.dim = '2d' if self.q_input.is_2d else '1d'
459        self.info = model.info
460        self.dtype = dtype
461
462        # Converter to translate input to target type.
463        self._as_dtype = np.float64 if dtype == generate.F64 else np.float32
464
465        # Holding place for the returned value.
466        nout = 2 if self.info.have_Fq and self.dim == '1d' else 1
467        extra_q = 4  # Total weight, form volume, shell volume and R_eff.
468        self.result = np.empty(self.q_input.nq*nout + extra_q, dtype)
469
470        # Allocate result value on GPU.
471        width = ((self.result.size+31)//32)*32 * self.dtype.itemsize
472        self._result_b = cuda.mem_alloc(width)
473
474    def _call_kernel(self, call_details, values, cutoff, magnetic,
475                     effective_radius_type):
476        # type: (CallDetails, np.ndarray, float, bool, int) -> np.ndarray
477
478        # Arrange data transfer to card.
479        details_b = cuda.to_device(call_details.buffer)
480        values_b = cuda.to_device(values)
481
482        # Setup kernel function and arguments.
483        name = 'Iq' if self.dim == '1d' else 'Imagnetic' if magnetic else 'Iqxy'
484        kernel = self._model.get_function(name)
485        kernel_args = [
486            np.uint32(self.q_input.nq),  # Number of inputs.
487            None,  # Placeholder for pd_start.
488            None,  # Placeholder for pd_stop.
489            details_b,  # Problem definition.
490            values_b,  # Parameter values.
491            self.q_input.q_b,  # Q values.
492            self._result_b,   # Result storage.
493            self._as_dtype(cutoff),  # Probability cutoff.
494            np.uint32(effective_radius_type),  # R_eff mode.
495        ]
496        grid = partition(self.q_input.nq)
497
498        # Call kernel and retrieve results.
499        #print("Calling CUDA")
500        #call_details.show(values)
501        last_nap = time.clock()
502        step = 100000000//self.q_input.nq + 1
503        #step = 1000000000
504        for start in range(0, call_details.num_eval, step):
505            stop = min(start + step, call_details.num_eval)
506            #print("queuing",start,stop)
507            kernel_args[1:3] = [np.int32(start), np.int32(stop)]
508            kernel(*kernel_args, **grid)
509            if stop < call_details.num_eval:
510                sync()
511                # Allow other processes to run.
512                current_time = time.clock()
513                if current_time - last_nap > 0.5:
514                    time.sleep(0.001)
515                    last_nap = current_time
516        sync()
517        cuda.memcpy_dtoh(self.result, self._result_b)
518        #print("result", self.result)
519
520        details_b.free()
521        values_b.free()
522
523    def release(self):
524        # type: () -> None
525        """
526        Release resources associated with the kernel.
527        """
528        self.q_input.release()
529        if self._result_b is not None:
530            self._result_b.free()
531            self._result_b = None
532
533    def __del__(self):
534        # type: () -> None
535        self.release()
536
537
538def sync():
539    """
540    Overview:
541        Waits for operation in the current context to complete.
542
543    Note: Maybe context.synchronize() is sufficient.
544    """
545    # Create an event with which to synchronize.
546    done = cuda.Event()
547
548    # Schedule an event trigger on the GPU.
549    done.record()
550
551    # Make sure we don't hog resource while waiting to sync.
552    while not done.query():
553        time.sleep(0.01)
554
555    # Block until the GPU executes the kernel.
556    done.synchronize()
557
558    # Clean up the event; I don't think they can be reused.
559    del done
560
561
562def partition(n):
563    '''
564    Overview:
565        Auto grids the thread blocks to achieve some level of calculation
566    efficiency.
567    '''
568    max_gx, max_gy = 65535, 65535
569    blocksize = 32
570    #max_gx, max_gy = 5, 65536
571    #blocksize = 3
572    block = (blocksize, 1, 1)
573    num_blocks = int((n+blocksize-1)/blocksize)
574    if num_blocks < max_gx:
575        grid = (num_blocks, 1)
576    else:
577        gx = max_gx
578        gy = (num_blocks + max_gx - 1) / max_gx
579        if gy >= max_gy:
580            raise ValueError("vector is too large")
581        grid = (gx, gy)
582    #print("block", block, "grid", grid)
583    #print("waste", block[0]*block[1]*block[2]*grid[0]*grid[1] - n)
584    return dict(block=block, grid=grid)
Note: See TracBrowser for help on using the repository browser.