kernelcreation.py 5.57 KB
Newer Older
1
from pystencils.assignment_collection import AssignmentCollection
Martin Bauer's avatar
Martin Bauer committed
2
3
4
5
from pystencils.gpucuda.indexing import indexingCreatorFromParams


def createKernel(equations, target='cpu', dataType="double", iterationSlice=None, ghostLayers=None,
6
                 cpuOpenMP=False, cpuVectorizeInfo=None,
Martin Bauer's avatar
Martin Bauer committed
7
8
9
                 gpuIndexing='block', gpuIndexingParams={}):
    """
    Creates abstract syntax tree (AST) of kernel, using a list of update equations.
10
    :param equations: either be a plain list of equations or a AssignmentCollection object
Martin Bauer's avatar
Martin Bauer committed
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
    :param target: 'cpu', 'llvm' or 'gpu'
    :param dataType: data type used for all untyped symbols (i.e. non-fields), can also be a dict from symbol name
                     to type
    :param iterationSlice: rectangular subset to iterate over, if not specified the complete non-ghost layer part of the
                           field is iterated over
    :param ghostLayers: if left to default, the number of necessary ghost layers is determined automatically
                        a single integer specifies the ghost layer count at all borders, can also be a sequence of
                        pairs [(xLowerGl, xUpperGl), .... ]

    CPU specific Parameters:
    :param cpuOpenMP: True or number of threads for OpenMP parallelization, False for no OpenMP
    :param cpuVectorizeInfo: pair of instruction set name ('sse, 'avx', 'avx512') and data type ('float', 'double')

    GPU specific Parameters
    :param gpuIndexing: either 'block' or 'line' , or custom indexing class (see gpucuda/indexing.py)
    :param gpuIndexingParams: dict with indexing parameters (constructor parameters of indexing class)
                              e.g. for 'block' one can specify {'blockSize': (20, 20, 10) }

    :return: abstract syntax tree object, that can either be printed as source code or can be compiled with
             through its compile() function
    """

    # ----  Normalizing parameters
    splitGroups = ()
35
    if isinstance(equations, AssignmentCollection):
Martin Bauer's avatar
Martin Bauer committed
36
37
38
        if 'splitGroups' in equations.simplification_hints:
            splitGroups = equations.simplification_hints['splitGroups']
        equations = equations.all_assignments
Martin Bauer's avatar
Martin Bauer committed
39
40
41
42
43
44
45
46
47
48
49
50
51

    # ----  Creating ast
    if target == 'cpu':
        from pystencils.cpu import createKernel
        from pystencils.cpu import addOpenMP
        ast = createKernel(equations, typeForSymbol=dataType, splitGroups=splitGroups,
                           iterationSlice=iterationSlice, ghostLayers=ghostLayers)
        if cpuOpenMP:
            addOpenMP(ast, numThreads=cpuOpenMP)
        if cpuVectorizeInfo:
            import pystencils.backends.simd_instruction_sets as vec
            from pystencils.vectorization import vectorize
            vecParams = cpuVectorizeInfo
Martin Bauer's avatar
Martin Bauer committed
52
            vec.selectedInstructionSet = vec.x86_vector_instruction_set(instruction_set=vecParams[0], data_type=vecParams[1])
Martin Bauer's avatar
Martin Bauer committed
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
            vectorize(ast)
        return ast
    elif target == 'llvm':
        from pystencils.llvm import createKernel
        ast = createKernel(equations, typeForSymbol=dataType, splitGroups=splitGroups,
                           iterationSlice=iterationSlice, ghostLayers=ghostLayers)
        return ast
    elif target == 'gpu':
        from pystencils.gpucuda import createCUDAKernel
        ast = createCUDAKernel(equations, typeForSymbol=dataType,
                               indexingCreator=indexingCreatorFromParams(gpuIndexing, gpuIndexingParams),
                               iterationSlice=iterationSlice, ghostLayers=ghostLayers)
        return ast
    else:
        raise ValueError("Unknown target %s. Has to be one of 'cpu', 'gpu' or 'llvm' " % (target,))


def createIndexedKernel(equations, indexFields, target='cpu', dataType="double", coordinateNames=('x', 'y', 'z'),
                        cpuOpenMP=True,
                        gpuIndexing='block', gpuIndexingParams={}):
    """
    Similar to :func:`createKernel`, but here not all cells of a field are updated but only cells with
    coordinates which are stored in an index field. This traversal method can e.g. be used for boundary handling.

    The coordinates are stored in a separated indexField, which is a one dimensional array with struct data type.
    This struct has to contain fields named 'x', 'y' and for 3D fields ('z'). These names are configurable with the
    'coordinateNames' parameter. The struct can have also other fields that can be read and written in the kernel, for
    example boundary parameters.

    indexFields: list of index fields, i.e. 1D fields with struct data type
    coordinateNames: name of the coordinate fields in the struct data type
    """

86
    if isinstance(equations, AssignmentCollection):
Martin Bauer's avatar
Martin Bauer committed
87
        equations = equations.all_assignments
Martin Bauer's avatar
Martin Bauer committed
88
89
90
    if target == 'cpu':
        from pystencils.cpu import createIndexedKernel
        from pystencils.cpu import addOpenMP
91
        ast = createIndexedKernel(equations, indexFields=indexFields, typeForSymbol=dataType,
Martin Bauer's avatar
Martin Bauer committed
92
93
94
95
96
97
98
99
100
101
102
103
104
                                  coordinateNames=coordinateNames)
        if cpuOpenMP:
            addOpenMP(ast, numThreads=cpuOpenMP)
        return ast
    elif target == 'llvm':
        raise NotImplementedError("Indexed kernels are not yet supported in LLVM backend")
    elif target == 'gpu':
        from pystencils.gpucuda import createdIndexedCUDAKernel
        ast = createdIndexedCUDAKernel(equations, indexFields, typeForSymbol=dataType, coordinateNames=coordinateNames,
                                       indexingCreator=indexingCreatorFromParams(gpuIndexing, gpuIndexingParams))
        return ast
    else:
        raise ValueError("Unknown target %s. Has to be either 'cpu' or 'gpu'" % (target,))