Source code for pyfr.backends.hip.provider

# -*- coding: utf-8 -*-

from pyfr.backends.base import (BaseKernelProvider,
                                BasePointwiseKernelProvider, ComputeKernel)
from pyfr.backends.hip.compiler import SourceModule
from pyfr.backends.hip.generator import HIPKernelGenerator
from pyfr.util import memoize

def get_grid_for_block(block, nrow, ncol=1):
    return (int((nrow + (-nrow % block[0])) // block[0]),
            int((ncol + (-ncol % block[1])) // block[1]), 1)

class HIPKernelProvider(BaseKernelProvider):
    def _build_kernel(self, name, src, argtypes):
        # Compile the source code and retrieve the kernel
        return SourceModule(self.backend, src).get_function(name, argtypes)

[docs]class HIPPointwiseKernelProvider(HIPKernelProvider, BasePointwiseKernelProvider): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._block1d = (64, 1, 1) self._block2d = (64, 4, 1) # Pass these block sizes to the generator class KernelGenerator(HIPKernelGenerator): block1d = self._block1d block2d = self._block2d self.kernel_generator_cls = KernelGenerator
[docs] def _instantiate_kernel(self, dims, fun, arglst): block = self._block1d if len(dims) == 1 else self._block2d grid = get_grid_for_block(block, dims[-1]) class PointwiseKernel(ComputeKernel): if any(isinstance(arg, str) for arg in arglst): def run(self, queue, **kwargs): fun.exec_async(grid, block, queue.stream_comp, *[kwargs.get(ka, ka) for ka in arglst]) else: def run(self, queue, **kwargs): fun.exec_async(grid, block, queue.stream_comp, *arglst) return PointwiseKernel()