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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
|
@shBang #!/usr/bin/env python3
@*
Basis.tmpl
Created by Graham Dennis on 2008-12-14.
Copyright (c) 2008-2012, Graham Dennis
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*@
@extends xpdeint.ScriptElement
@import operator
@from functools import reduce
@attr $matrixType = 'real'
@attr $supportsInPlaceOperation = False
@def costEstimate(basisReps)
@set costMultiplier = 4 if $matrixType == 'complex' else 1
@return reduce(operator.mul, [rep.latticeEstimate for rep in basisReps]) * costMultiplier
@end def
@def transformMatricesForDimReps($forwardDimRep, $backwardDimRep)
_mmt_matrix_forward = ($matrixType *)xmds_malloc(sizeof($matrixType) * ${forwardDimRep.globalLattice} * ${backwardDimRep.globalLattice});
_mmt_matrix_backward = ($matrixType *)xmds_malloc(sizeof($matrixType) * ${backwardDimRep.globalLattice} * ${forwardDimRep.globalLattice});
for (long _i0 = 0; _i0 < ${forwardDimRep.globalLattice}; _i0++) {
${transformMatricesForwardDimConstantsAtIndex(forwardDimRep, backwardDimRep, '_i0'), autoIndent=True}@slurp
for (long _i1 = 0; _i1 < ${backwardDimRep.globalLattice}; _i1++) {
${transformMatricesForDimRepsAtIndices(forwardDimRep, backwardDimRep, '_i0', '_i1'), autoIndent=True}@slurp
}
}
@end def
@def transformMatricesForwardDimConstantsAtIndex($forwardDimRep, $backwardDimRep, $index)
@end def
@def transformMatricesForDimRepsAtIndices($forwardDimRep, $backwardDimRep, $forwardIndex, $backwardIndex)
@#
_mmt_matrix_forward [${backwardIndex} * ${forwardDimRep.globalLattice} + ${forwardIndex}] = \
${forwardMatrixForDimAtIndices(forwardDimRep, backwardDimRep, forwardIndex, backwardIndex)};
_mmt_matrix_backward[${forwardIndex} * ${backwardDimRep.globalLattice} + ${backwardIndex}] = \
${backwardMatrixForDimAtIndices(forwardDimRep, backwardDimRep, forwardIndex, backwardIndex)};
@#
@end def
@def performTransform($sourceDimRep, $destDimRep, $dir = None)
@#
@set $blasTypeChar = {'real': {'single': 's', 'double': 'd'}, 'complex': {'single': 'c', 'double': 'z'}}[self.matrixType][$precision]
@set $matMultFunction = 'cblas_%sgemm' % blasTypeChar
@set $alphaBetaPrefix = {'real': '', 'complex': '&'}[self.matrixType]
const ${matrixType} alpha = 1.0;
const ${matrixType} beta = 0.0;
${matMultFunction}(
CblasRowMajor, CblasNoTrans, CblasNoTrans,
${destDimRep.globalLattice},
/* nelem */ innerLoopSize,
${sourceDimRep.globalLattice},
/* alpha */ ${alphaBetaPrefix}alpha,
/* A */ _mmt_matrix_${dir}, ${sourceDimRep.globalLattice},
/* B */ source_data + _i0 * innerLoopSize * ${sourceDimRep.globalLattice},
innerLoopSize,
/* beta */ ${alphaBetaPrefix}beta,
/* C */ dest_data + _i0 * innerLoopSize * ${destDimRep.globalLattice},
innerLoopSize
);
@#
@end def
@def transformFunctionStart
static $matrixType *_mmt_matrix_forward = NULL;
static $matrixType *_mmt_matrix_backward = NULL;
@end def
@def transformFunction(transformID, transformDict, function)
@#
@set $transformPair = transformDict['transformPair']
@set $forwardDimRep = transformPair[0][0]
@set $backwardDimRep = transformPair[1][0]
static bool _initialised = false;
${transformFunctionStart}@slurp
ptrdiff_t innerLoopSize = _postfix_lattice;
${matrixType}* const __restrict__ source_data = reinterpret_cast<${matrixType}* const>(_data_in);
${matrixType}* const __restrict__ dest_data = reinterpret_cast<${matrixType}* const>(_data_out);
if (!_initialised) {
_LOG(_SIMULATION_LOG_LEVEL, "Building matrices for ${function.description}...");
${transformMatricesForDimReps($forwardDimRep, $backwardDimRep), autoIndent=True}@slurp
_LOG(_SIMULATION_LOG_LEVEL, " done.\n");
_initialised = true;
}
if (_forward) {
for (long _i0 = 0; _i0 < _prefix_lattice; _i0++) {
${performTransform(forwardDimRep, backwardDimRep, dir='forward'), autoIndent=True}@slurp
}
} else {
for (long _i0 = 0; _i0 < _prefix_lattice; _i0++) {
${performTransform(backwardDimRep, forwardDimRep, dir='backward'), autoIndent=True}@slurp
}
}
@#
@end def
|
