rng.py 4.85 KB
Newer Older
1
import numpy as np
Martin Bauer's avatar
Martin Bauer committed
2
3
import sympy as sp

4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from pystencils import TypedSymbol
from pystencils.astnodes import LoopOverCoordinate
from pystencils.backends.cbackend import CustomCodeNode

philox_two_doubles_call = """
{result_symbols[0].dtype} {result_symbols[0].name};
{result_symbols[1].dtype} {result_symbols[1].name};
philox_double2({parameters}, {result_symbols[0].name}, {result_symbols[1].name});
"""

philox_four_floats_call = """
{result_symbols[0].dtype} {result_symbols[0].name};
{result_symbols[1].dtype} {result_symbols[1].name};
{result_symbols[2].dtype} {result_symbols[2].name};
{result_symbols[3].dtype} {result_symbols[3].name};
19
philox_float4({parameters},
20
21
22
23
24
              {result_symbols[0].name}, {result_symbols[1].name}, {result_symbols[2].name}, {result_symbols[3].name});

"""


25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def _get_philox_code(template, dialect, vector_instruction_set, time_step, offsets, keys, dim, result_symbols):
    parameters = [time_step] + [LoopOverCoordinate.get_loop_counter_symbol(i) + offsets[i]
                                for i in range(dim)] + list(keys)

    while len(parameters) < 6:
        parameters.append(0)
    parameters = parameters[:6]

    assert len(parameters) == 6

    if dialect == 'cuda' or (dialect == 'c' and vector_instruction_set is None):
        return template.format(parameters=', '.join(str(p) for p in parameters),
                               result_symbols=result_symbols)
    else:
        raise NotImplementedError("Not yet implemented for this backend")


42
43
class PhiloxTwoDoubles(CustomCodeNode):

44
    def __init__(self, dim, time_step=TypedSymbol("time_step", np.uint32), offsets=(0, 0, 0), keys=(0, 0)):
45
46
47
48
        self.result_symbols = tuple(TypedSymbol(sp.Dummy().name, np.float64) for _ in range(2))
        symbols_read = [s for s in keys if isinstance(s, sp.Symbol)]
        super().__init__("", symbols_read=symbols_read, symbols_defined=self.result_symbols)
        self._time_step = time_step
49
        self._offsets = offsets
50
        self.headers = ['"philox_rand.h"']
51
52
        self.keys = tuple(keys)
        self._args = sp.sympify((dim, time_step, keys))
53
54
55
56
57
58
59
60
        self._dim = dim

    @property
    def args(self):
        return self._args

    @property
    def undefined_symbols(self):
61
        result = {a for a in (self._time_step, *self._offsets, *self.keys) if isinstance(a, sp.Symbol)}
62
63
64
65
66
        loop_counters = [LoopOverCoordinate.get_loop_counter_symbol(i)
                         for i in range(self._dim)]
        result.update(loop_counters)
        return result

67
68
69
    def fast_subs(self, _):
        return self  # nothing to replace inside this node - would destroy intermediate "dummy" by re-creating them

70
    def get_code(self, dialect, vector_instruction_set):
71
72
        return _get_philox_code(philox_two_doubles_call, dialect, vector_instruction_set,
                                self._time_step, self._offsets, self.keys, self._dim, self.result_symbols)
73
74
75
76
77
78
79

    def __repr__(self):
        return "{}, {} <- PhiloxRNG".format(*self.result_symbols)


class PhiloxFourFloats(CustomCodeNode):

80
    def __init__(self, dim, time_step=TypedSymbol("time_step", np.uint32), offsets=(0, 0, 0), keys=(0, 0)):
81
82
83
84
        self.result_symbols = tuple(TypedSymbol(sp.Dummy().name, np.float32) for _ in range(4))
        symbols_read = [s for s in keys if isinstance(s, sp.Symbol)]
        super().__init__("", symbols_read=symbols_read, symbols_defined=self.result_symbols)
        self._time_step = time_step
85
        self._offsets = offsets
86
        self.headers = ['"philox_rand.h"']
87
        self.keys = tuple(keys)
88
        self._args = sp.sympify((dim, time_step, offsets, keys))
89
90
91
92
93
94
95
96
        self._dim = dim

    @property
    def args(self):
        return self._args

    @property
    def undefined_symbols(self):
97
        result = {a for a in (self._time_step, *self._offsets, *self.keys) if isinstance(a, sp.Symbol)}
98
99
100
101
102
        loop_counters = [LoopOverCoordinate.get_loop_counter_symbol(i)
                         for i in range(self._dim)]
        result.update(loop_counters)
        return result

103
104
105
    def fast_subs(self, _):
        return self  # nothing to replace inside this node - would destroy intermediate "dummy" by re-creating them

106
    def get_code(self, dialect, vector_instruction_set):
107
108
        return _get_philox_code(philox_four_floats_call, dialect, vector_instruction_set,
                                self._time_step, self._offsets, self.keys, self._dim, self.result_symbols)
109
110
111

    def __repr__(self):
        return "{}, {}, {}, {} <- PhiloxRNG".format(*self.result_symbols)
Martin Bauer's avatar
Martin Bauer committed
112
113


114
115
def random_symbol(assignment_list, seed=TypedSymbol("seed", np.uint32), rng_node=PhiloxTwoDoubles, *args, **kwargs):
    counter = 0
Martin Bauer's avatar
Martin Bauer committed
116
    while True:
117
        node = rng_node(*args, keys=(counter, seed), **kwargs)
Martin Bauer's avatar
Martin Bauer committed
118
119
120
121
122
123
        inserted = False
        for symbol in node.result_symbols:
            if not inserted:
                assignment_list.insert(0, node)
                inserted = True
            yield symbol