Line data Source code
1 : /* -------------------------------------------------------------------------- *
2 : * lepton *
3 : * -------------------------------------------------------------------------- *
4 : * This is part of the lepton expression parser. *
5 : * See https://openmm.org/development. *
6 : * *
7 : * Portions copyright (c) 2013-2022 Stanford University and the Authors. *
8 : * Authors: Peter Eastman *
9 : * Contributors: *
10 : * *
11 : * Permission is hereby granted, free of charge, to any person obtaining a *
12 : * copy of this software and associated documentation files (the "Software"), *
13 : * to deal in the Software without restriction, including without limitation *
14 : * the rights to use, copy, modify, merge, publish, distribute, sublicense, *
15 : * and/or sell copies of the Software, and to permit persons to whom the *
16 : * Software is furnished to do so, subject to the following conditions: *
17 : * *
18 : * The above copyright notice and this permission notice shall be included in *
19 : * all copies or substantial portions of the Software. *
20 : * *
21 : * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
22 : * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, *
23 : * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL *
24 : * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, *
25 : * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR *
26 : * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE *
27 : * USE OR OTHER DEALINGS IN THE SOFTWARE. *
28 : * -------------------------------------------------------------------------- */
29 :
30 : #include "CompiledExpression.h"
31 : #include "Operation.h"
32 : #include "ParsedExpression.h"
33 : #ifdef __PLUMED_HAS_ASMJIT
34 : #include "asmjit/a64.h"
35 : #include "asmjit/x86.h"
36 : #endif
37 : #include <utility>
38 : #include <cstdlib>
39 : #include <limits>
40 :
41 : namespace PLMD {
42 : using namespace lepton;
43 : using namespace std;
44 : #ifdef __PLUMED_HAS_ASMJIT
45 : using namespace asmjit;
46 : #endif
47 :
48 65477 : bool lepton::useAsmJit() {
49 : #ifdef __PLUMED_HAS_ASMJIT
50 356 : static const bool use=[](){
51 356 : if(auto s=std::getenv("PLUMED_USE_ASMJIT")) {
52 3 : auto ss=std::string(s);
53 3 : if(ss=="yes") return true;
54 1 : if(ss=="no") return false;
55 0 : throw Exception("PLUMED_USE_ASMJIT variable is set to " + ss + "; should be yes or no");
56 : }
57 : return true; // by default use asmjit
58 65477 : }();
59 65477 : return use;
60 : #else
61 : return false;
62 : #endif
63 : }
64 :
65 27908 : AsmJitRuntimePtr::AsmJitRuntimePtr()
66 : #ifdef __PLUMED_HAS_ASMJIT
67 27908 : : ptr(useAsmJit()?new asmjit::JitRuntime:nullptr)
68 : #endif
69 27908 : {}
70 :
71 27908 : AsmJitRuntimePtr::~AsmJitRuntimePtr()
72 : {
73 : #ifdef __PLUMED_HAS_ASMJIT
74 27908 : if(useAsmJit()) delete static_cast<asmjit::JitRuntime*>(ptr);
75 : #endif
76 27908 : }
77 :
78 284 : CompiledExpression::CompiledExpression() : jitCode(NULL) {
79 284 : }
80 :
81 9304 : CompiledExpression::CompiledExpression(const ParsedExpression& expression) : jitCode(NULL) {
82 9304 : ParsedExpression expr = expression.optimize(); // Just in case it wasn't already optimized.
83 9304 : vector<pair<ExpressionTreeNode, int> > temps;
84 9304 : compileExpression(expr.getRootNode(), temps);
85 : int maxArguments = 1;
86 46229 : for (int i = 0; i < (int) operation.size(); i++)
87 36925 : if (operation[i]->getNumArguments() > maxArguments)
88 5639 : maxArguments = operation[i]->getNumArguments();
89 9304 : argValues.resize(maxArguments);
90 : #ifdef __PLUMED_HAS_ASMJIT
91 9304 : if(useAsmJit()) generateJitCode();
92 : #endif
93 18608 : }
94 :
95 27908 : CompiledExpression::~CompiledExpression() {
96 139109 : for (int i = 0; i < (int) operation.size(); i++)
97 111201 : if (operation[i] != NULL)
98 111201 : delete operation[i];
99 111632 : }
100 :
101 18320 : CompiledExpression::CompiledExpression(const CompiledExpression& expression) : jitCode(NULL) {
102 18320 : *this = expression;
103 18320 : }
104 :
105 18574 : CompiledExpression& CompiledExpression::operator=(const CompiledExpression& expression) {
106 18574 : arguments = expression.arguments;
107 18574 : target = expression.target;
108 : variableIndices = expression.variableIndices;
109 : variableNames = expression.variableNames;
110 18574 : workspace.resize(expression.workspace.size());
111 18574 : argValues.resize(expression.argValues.size());
112 18574 : operation.resize(expression.operation.size());
113 92850 : for (int i = 0; i < (int) operation.size(); i++)
114 74276 : operation[i] = expression.operation[i]->clone();
115 18574 : setVariableLocations(variablePointers);
116 18574 : return *this;
117 : }
118 :
119 59454 : void CompiledExpression::compileExpression(const ExpressionTreeNode& node, vector<pair<ExpressionTreeNode, int> >& temps) {
120 59454 : if (findTempIndex(node, temps) != -1)
121 10129 : return; // We have already processed a node identical to this one.
122 :
123 : // Process the child nodes.
124 :
125 49325 : vector<int> args;
126 99475 : for (int i = 0; i < node.getChildren().size(); i++) {
127 50150 : compileExpression(node.getChildren()[i], temps);
128 50150 : args.push_back(findTempIndex(node.getChildren()[i], temps));
129 : }
130 :
131 : // Process this node.
132 :
133 49325 : if (node.getOperation().getId() == Operation::VARIABLE) {
134 12400 : variableIndices[node.getOperation().getName()] = (int) workspace.size();
135 24800 : variableNames.insert(node.getOperation().getName());
136 : }
137 : else {
138 36925 : int stepIndex = (int) arguments.size();
139 36925 : arguments.push_back(vector<int>());
140 36925 : target.push_back((int) workspace.size());
141 36925 : operation.push_back(node.getOperation().clone());
142 36925 : if (args.size() == 0)
143 1348 : arguments[stepIndex].push_back(0); // The value won't actually be used. We just need something there.
144 : else {
145 : // If the arguments are sequential, we can just pass a pointer to the first one.
146 :
147 : bool sequential = true;
148 50150 : for (int i = 1; i < args.size(); i++)
149 14573 : if (args[i] != args[i-1]+1)
150 : sequential = false;
151 35577 : if (sequential)
152 25927 : arguments[stepIndex].push_back(args[0]);
153 : else
154 9650 : arguments[stepIndex] = args;
155 : }
156 : }
157 49325 : temps.push_back(make_pair(node, (int) workspace.size()));
158 49325 : workspace.push_back(0.0);
159 : }
160 :
161 109604 : int CompiledExpression::findTempIndex(const ExpressionTreeNode& node, vector<pair<ExpressionTreeNode, int> >& temps) {
162 1872402 : for (int i = 0; i < (int) temps.size(); i++)
163 1823077 : if (temps[i].first == node)
164 : return i;
165 : return -1;
166 : }
167 :
168 4214 : const set<string>& CompiledExpression::getVariables() const {
169 4214 : return variableNames;
170 : }
171 :
172 78167 : double& CompiledExpression::getVariableReference(const string& name) {
173 : map<string, double*>::iterator pointer = variablePointers.find(name);
174 78167 : if (pointer != variablePointers.end())
175 100 : return *pointer->second;
176 : map<string, int>::iterator index = variableIndices.find(name);
177 78067 : if (index == variableIndices.end())
178 39207 : throw Exception("getVariableReference: Unknown variable '"+name+"'");
179 64998 : return workspace[index->second];
180 : }
181 :
182 18674 : void CompiledExpression::setVariableLocations(map<string, double*>& variableLocations) {
183 : variablePointers = variableLocations;
184 18674 : static const bool asmjit=useAsmJit();
185 18674 : if(asmjit) {
186 : #ifdef __PLUMED_HAS_ASMJIT
187 : // Rebuild the JIT code.
188 :
189 18350 : if (workspace.size() > 0)
190 18350 : generateJitCode();
191 : #endif
192 : }
193 : // Make a list of all variables we will need to copy before evaluating the expression.
194 :
195 18674 : variablesToCopy.clear();
196 43547 : for (map<string, int>::const_iterator iter = variableIndices.begin(); iter != variableIndices.end(); ++iter) {
197 24873 : map<string, double*>::iterator pointer = variablePointers.find(iter->first);
198 24873 : if (pointer != variablePointers.end())
199 100 : variablesToCopy.push_back(make_pair(&workspace[iter->second], pointer->second));
200 : }
201 18674 : }
202 :
203 85788486 : double CompiledExpression::evaluate() const {
204 : #ifdef __PLUMED_HAS_ASMJIT
205 85788486 : if(jitCode) return jitCode();
206 : #endif
207 16443 : for (int i = 0; i < variablesToCopy.size(); i++)
208 0 : *variablesToCopy[i].first = *variablesToCopy[i].second;
209 :
210 : // Loop over the operations and evaluate each one.
211 :
212 49272 : for (int step = 0; step < operation.size(); step++) {
213 : const vector<int>& args = arguments[step];
214 32829 : if (args.size() == 1)
215 30198 : workspace[target[step]] = operation[step]->evaluate(&workspace[args[0]], dummyVariables);
216 : else {
217 8097 : for (int i = 0; i < args.size(); i++)
218 5466 : argValues[i] = workspace[args[i]];
219 2631 : workspace[target[step]] = operation[step]->evaluate(&argValues[0], dummyVariables);
220 : }
221 : }
222 16443 : return workspace[workspace.size()-1];
223 : }
224 :
225 : #ifdef __PLUMED_HAS_ASMJIT
226 : #if defined(__aarch64__) || defined(_M_ARM64)
227 : static void generateSingleArgCall(a64::Compiler&, arm::Vec&, arm::Vec&, double (*)(double));
228 : static void generateTwoArgCall(a64::Compiler&, arm::Vec&, arm::Vec&, arm::Vec&, double (*)(double, double));
229 : #else
230 : static void generateSingleArgCall(x86::Compiler&, x86::Xmm&, x86::Xmm&, double (*)(double));
231 : static void generateTwoArgCall(x86::Compiler&, x86::Xmm&, x86::Xmm&, x86::Xmm&, double (*)(double, double));
232 : #endif
233 3010 : static double evaluateOperation(Operation* op, double* args) {
234 3010 : static map<string, double> dummyVariables;
235 3010 : return op->evaluate(args, dummyVariables);
236 : }
237 :
238 27492 : void CompiledExpression::findPowerGroups(vector<vector<int> >& groups, vector<vector<int> >& groupPowers, vector<int>& stepGroup) {
239 : // Identify every step that raises an argument to an integer power.
240 :
241 27492 : vector<int> stepPower(operation.size(), 0);
242 27492 : vector<int> stepArg(operation.size(), -1);
243 138021 : for (int step = 0; step < operation.size(); step++) {
244 110529 : Operation& op = *operation[step];
245 : int power = 0;
246 110529 : if (op.getId() == Operation::SQUARE)
247 : power = 2;
248 97745 : else if (op.getId() == Operation::CUBE)
249 : power = 3;
250 97502 : else if (op.getId() == Operation::POWER_CONSTANT) {
251 631 : double realPower = dynamic_cast<const Operation::PowerConstant*>(&op)->getValue();
252 631 : if (realPower == (int) realPower)
253 : power = (int) realPower;
254 : }
255 576 : if (power != 0) {
256 13603 : stepPower[step] = power;
257 13603 : stepArg[step] = arguments[step][0];
258 : }
259 : }
260 :
261 : // Find groups that operate on the same argument and whose powers have the same sign.
262 :
263 27492 : stepGroup.resize(operation.size(), -1);
264 138021 : for (int i = 0; i < operation.size(); i++) {
265 110529 : if (stepGroup[i] != -1)
266 231 : continue;
267 110298 : vector<int> group, power;
268 2582910 : for (int j = i; j < operation.size(); j++) {
269 2472612 : if (stepArg[i] == stepArg[j] && stepPower[i]*stepPower[j] > 0) {
270 13603 : stepGroup[j] = groups.size();
271 13603 : group.push_back(j);
272 13603 : power.push_back(stepPower[j]);
273 : }
274 : }
275 110298 : groups.push_back(group);
276 110298 : groupPowers.push_back(power);
277 : }
278 27492 : }
279 :
280 : #if defined(__aarch64__) || defined(_M_ARM64)
281 : void CompiledExpression::generateJitCode() {
282 : auto& runtime = *static_cast<asmjit::JitRuntime*>(runtimeptr.get());
283 : if (jitCode) {
284 : runtime.release(jitCode);
285 : jitCode = nullptr;
286 : }
287 : CodeHolder code;
288 : code.init(runtime.environment());
289 : a64::Compiler c(&code);
290 : c.addFunc(FuncSignatureT<double>());
291 : vector<arm::Vec> workspaceVar(workspace.size());
292 : for (int i = 0; i < (int) workspaceVar.size(); i++)
293 : workspaceVar[i] = c.newVecD();
294 : arm::Gp argsPointer = c.newIntPtr();
295 : c.mov(argsPointer, imm(&argValues[0]));
296 : vector<vector<int> > groups, groupPowers;
297 : vector<int> stepGroup;
298 : findPowerGroups(groups, groupPowers, stepGroup);
299 :
300 : // Load the arguments into variables.
301 :
302 : for (set<string>::const_iterator iter = variableNames.begin(); iter != variableNames.end(); ++iter) {
303 : map<string, int>::iterator index = variableIndices.find(*iter);
304 : arm::Gp variablePointer = c.newIntPtr();
305 : c.mov(variablePointer, imm(&getVariableReference(index->first)));
306 : c.ldr(workspaceVar[index->second], arm::ptr(variablePointer, 0));
307 : }
308 :
309 : // Make a list of all constants that will be needed for evaluation.
310 :
311 : vector<int> operationConstantIndex(operation.size(), -1);
312 : for (int step = 0; step < (int) operation.size(); step++) {
313 : // Find the constant value (if any) used by this operation.
314 :
315 : Operation& op = *operation[step];
316 : double value;
317 : if (op.getId() == Operation::CONSTANT)
318 : value = dynamic_cast<Operation::Constant&>(op).getValue();
319 : else if (op.getId() == Operation::ADD_CONSTANT)
320 : value = dynamic_cast<Operation::AddConstant&>(op).getValue();
321 : else if (op.getId() == Operation::MULTIPLY_CONSTANT)
322 : value = dynamic_cast<Operation::MultiplyConstant&>(op).getValue();
323 : else if (op.getId() == Operation::RECIPROCAL)
324 : value = 1.0;
325 : else if (op.getId() == Operation::STEP)
326 : value = 1.0;
327 : else if (op.getId() == Operation::DELTA)
328 : value = std::numeric_limits<double>::infinity();
329 : else if (op.getId() == Operation::POWER_CONSTANT) {
330 : if (stepGroup[step] == -1)
331 : value = dynamic_cast<Operation::PowerConstant&>(op).getValue();
332 : else
333 : value = 1.0;
334 : }
335 : else
336 : continue;
337 :
338 : // See if we already have a variable for this constant.
339 :
340 : for (int i = 0; i < (int) constants.size(); i++)
341 : if (value == constants[i]) {
342 : operationConstantIndex[step] = i;
343 : break;
344 : }
345 : if (operationConstantIndex[step] == -1) {
346 : operationConstantIndex[step] = constants.size();
347 : constants.push_back(value);
348 : }
349 : }
350 :
351 : // Load constants into variables.
352 :
353 : vector<arm::Vec> constantVar(constants.size());
354 : if (constants.size() > 0) {
355 : arm::Gp constantsPointer = c.newIntPtr();
356 : c.mov(constantsPointer, imm(&constants[0]));
357 : for (int i = 0; i < (int) constants.size(); i++) {
358 : constantVar[i] = c.newVecD();
359 : c.ldr(constantVar[i], arm::ptr(constantsPointer, 8*i));
360 : }
361 : }
362 :
363 : // Evaluate the operations.
364 :
365 : vector<bool> hasComputedPower(operation.size(), false);
366 : for (int step = 0; step < (int) operation.size(); step++) {
367 : if (hasComputedPower[step])
368 : continue;
369 :
370 : // When one or more steps involve raising the same argument to multiple integer
371 : // powers, we can compute them all together for efficiency.
372 :
373 : if (stepGroup[step] != -1) {
374 : vector<int>& group = groups[stepGroup[step]];
375 : vector<int>& powers = groupPowers[stepGroup[step]];
376 : arm::Vec multiplier = c.newVecD();
377 : if (powers[0] > 0)
378 : c.fmov(multiplier, workspaceVar[arguments[step][0]]);
379 : else {
380 : c.fdiv(multiplier, constantVar[operationConstantIndex[step]], workspaceVar[arguments[step][0]]);
381 : for (int i = 0; i < powers.size(); i++)
382 : powers[i] = -powers[i];
383 : }
384 : vector<bool> hasAssigned(group.size(), false);
385 : bool done = false;
386 : while (!done) {
387 : done = true;
388 : for (int i = 0; i < group.size(); i++) {
389 : if (powers[i]%2 == 1) {
390 : if (!hasAssigned[i])
391 : c.fmov(workspaceVar[target[group[i]]], multiplier);
392 : else
393 : c.fmul(workspaceVar[target[group[i]]], workspaceVar[target[group[i]]], multiplier);
394 : hasAssigned[i] = true;
395 : }
396 : powers[i] >>= 1;
397 : if (powers[i] != 0)
398 : done = false;
399 : }
400 : if (!done)
401 : c.fmul(multiplier, multiplier, multiplier);
402 : }
403 : for (int step : group)
404 : hasComputedPower[step] = true;
405 : continue;
406 : }
407 :
408 : // Evaluate the step.
409 :
410 : Operation& op = *operation[step];
411 : vector<int> args = arguments[step];
412 : if (args.size() == 1) {
413 : // One or more sequential arguments. Fill out the list.
414 :
415 : for (int i = 1; i < op.getNumArguments(); i++)
416 : args.push_back(args[0]+i);
417 : }
418 :
419 : // Generate instructions to execute this operation.
420 :
421 : switch (op.getId()) {
422 : case Operation::CONSTANT:
423 : c.fmov(workspaceVar[target[step]], constantVar[operationConstantIndex[step]]);
424 : break;
425 : case Operation::ADD:
426 : c.fadd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
427 : break;
428 : case Operation::SUBTRACT:
429 : c.fsub(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
430 : break;
431 : case Operation::MULTIPLY:
432 : c.fmul(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
433 : break;
434 : case Operation::DIVIDE:
435 : c.fdiv(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
436 : break;
437 : case Operation::POWER:
438 : generateTwoArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]], pow);
439 : break;
440 : case Operation::NEGATE:
441 : c.fneg(workspaceVar[target[step]], workspaceVar[args[0]]);
442 : break;
443 : case Operation::SQRT:
444 : c.fsqrt(workspaceVar[target[step]], workspaceVar[args[0]]);
445 : break;
446 : case Operation::EXP:
447 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], exp);
448 : break;
449 : case Operation::LOG:
450 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], log);
451 : break;
452 : case Operation::SIN:
453 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], sin);
454 : break;
455 : case Operation::COS:
456 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], cos);
457 : break;
458 : case Operation::TAN:
459 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], tan);
460 : break;
461 : case Operation::ASIN:
462 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], asin);
463 : break;
464 : case Operation::ACOS:
465 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], acos);
466 : break;
467 : case Operation::ATAN:
468 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], atan);
469 : break;
470 : case Operation::ATAN2:
471 : generateTwoArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]], atan2);
472 : break;
473 : case Operation::SINH:
474 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], sinh);
475 : break;
476 : case Operation::COSH:
477 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], cosh);
478 : break;
479 : case Operation::TANH:
480 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], tanh);
481 : break;
482 : case Operation::STEP:
483 : c.fcmge(workspaceVar[target[step]], workspaceVar[args[0]], imm(0));
484 : c.and_(workspaceVar[target[step]], workspaceVar[target[step]], constantVar[operationConstantIndex[step]]);
485 : break;
486 : case Operation::DELTA:
487 : c.fcmeq(workspaceVar[target[step]], workspaceVar[args[0]], imm(0));
488 : c.and_(workspaceVar[target[step]], workspaceVar[target[step]], constantVar[operationConstantIndex[step]]);
489 : break;
490 : case Operation::SQUARE:
491 : c.fmul(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[0]]);
492 : break;
493 : case Operation::CUBE:
494 : c.fmul(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[0]]);
495 : c.fmul(workspaceVar[target[step]], workspaceVar[target[step]], workspaceVar[args[0]]);
496 : break;
497 : case Operation::RECIPROCAL:
498 : c.fdiv(workspaceVar[target[step]], constantVar[operationConstantIndex[step]], workspaceVar[args[0]]);
499 : break;
500 : case Operation::ADD_CONSTANT:
501 : c.fadd(workspaceVar[target[step]], workspaceVar[args[0]], constantVar[operationConstantIndex[step]]);
502 : break;
503 : case Operation::MULTIPLY_CONSTANT:
504 : c.fmul(workspaceVar[target[step]], workspaceVar[args[0]], constantVar[operationConstantIndex[step]]);
505 : break;
506 : case Operation::POWER_CONSTANT:
507 : generateTwoArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], constantVar[operationConstantIndex[step]], pow);
508 : break;
509 : case Operation::MIN:
510 : c.fmin(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
511 : break;
512 : case Operation::MAX:
513 : c.fmax(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
514 : break;
515 : case Operation::ABS:
516 : c.fabs(workspaceVar[target[step]], workspaceVar[args[0]]);
517 : break;
518 : case Operation::FLOOR:
519 : c.frintm(workspaceVar[target[step]], workspaceVar[args[0]]);
520 : break;
521 : case Operation::CEIL:
522 : c.frintp(workspaceVar[target[step]], workspaceVar[args[0]]);
523 : break;
524 : case Operation::SELECT:
525 : c.fcmeq(workspaceVar[target[step]], workspaceVar[args[0]], imm(0));
526 : c.bsl(workspaceVar[target[step]], workspaceVar[args[2]], workspaceVar[args[1]]);
527 : break;
528 : default:
529 : // Just invoke evaluateOperation().
530 :
531 : for (int i = 0; i < (int) args.size(); i++)
532 : c.str(workspaceVar[args[i]], arm::ptr(argsPointer, 8*i));
533 : arm::Gp fn = c.newIntPtr();
534 : c.mov(fn, imm((void*) evaluateOperation));
535 : InvokeNode* invoke;
536 : c.invoke(&invoke, fn, FuncSignatureT<double, Operation*, double*>());
537 : invoke->setArg(0, imm(&op));
538 : invoke->setArg(1, imm(&argValues[0]));
539 : invoke->setRet(0, workspaceVar[target[step]]);
540 : }
541 : }
542 : c.ret(workspaceVar[workspace.size()-1]);
543 : c.endFunc();
544 : Error error = c.finalize();
545 : if (!error)
546 : error = runtime.add(&jitCode, &code);
547 : if (error)
548 : throw Exception(string("AsmJit compilation failed: ") + DebugUtils::errorAsString(error));
549 : }
550 :
551 : static void generateSingleArgCall(a64::Compiler& c, arm::Vec& dest, arm::Vec& arg, double (*function)(double)) {
552 : arm::Gp fn = c.newIntPtr();
553 : c.mov(fn, imm((void*) function));
554 : InvokeNode* invoke;
555 : c.invoke(&invoke, fn, FuncSignatureT<double, double>());
556 : invoke->setArg(0, arg);
557 : invoke->setRet(0, dest);
558 : }
559 :
560 : static void generateTwoArgCall(a64::Compiler& c, arm::Vec& dest, arm::Vec& arg1, arm::Vec& arg2, double (*function)(double, double)) {
561 : arm::Gp fn = c.newIntPtr();
562 : c.mov(fn, imm((void*) function));
563 : InvokeNode* invoke;
564 : c.invoke(&invoke, fn, FuncSignatureT<double, double, double>());
565 : invoke->setArg(0, arg1);
566 : invoke->setArg(1, arg2);
567 : invoke->setRet(0, dest);
568 : }
569 : #else
570 27492 : void CompiledExpression::generateJitCode() {
571 27492 : const CpuInfo& cpu = CpuInfo::host();
572 27492 : if (!cpu.hasFeature(CpuFeatures::X86::kAVX))
573 0 : return;
574 : auto& runtime = *static_cast<asmjit::JitRuntime*>(runtimeptr.get());
575 27492 : if (jitCode) {
576 : runtime.release(jitCode);
577 100 : jitCode = nullptr;
578 : }
579 27492 : CodeHolder code;
580 27492 : code.init(runtime.environment());
581 27492 : x86::Compiler c(&code);
582 27492 : FuncNode* funcNode = c.addFunc(FuncSignatureT<double>());
583 : funcNode->frame().setAvxEnabled();
584 27492 : vector<x86::Xmm> workspaceVar(workspace.size());
585 174910 : for (int i = 0; i < (int) workspaceVar.size(); i++)
586 147418 : workspaceVar[i] = c.newXmmSd();
587 : x86::Gp argsPointer = c.newIntPtr();
588 0 : c.mov(argsPointer, imm(&argValues[0]));
589 27492 : vector<vector<int> > groups, groupPowers;
590 27492 : vector<int> stepGroup;
591 27492 : findPowerGroups(groups, groupPowers, stepGroup);
592 :
593 : // Load the arguments into variables.
594 :
595 : x86::Gp variablePointer = c.newIntPtr();
596 64381 : for (set<string>::const_iterator iter = variableNames.begin(); iter != variableNames.end(); ++iter) {
597 : map<string, int>::iterator index = variableIndices.find(*iter);
598 73778 : c.mov(variablePointer, imm(&getVariableReference(index->first)));
599 36889 : c.vmovsd(workspaceVar[index->second], x86::ptr(variablePointer, 0, 0));
600 : }
601 :
602 : // Make a list of all constants that will be needed for evaluation.
603 :
604 27492 : vector<int> operationConstantIndex(operation.size(), -1);
605 138021 : for (int step = 0; step < (int) operation.size(); step++) {
606 : // Find the constant value (if any) used by this operation.
607 :
608 110529 : Operation& op = *operation[step];
609 : double value;
610 110529 : if (op.getId() == Operation::CONSTANT)
611 3838 : value = dynamic_cast<Operation::Constant&>(op).getValue();
612 106691 : else if (op.getId() == Operation::ADD_CONSTANT)
613 7148 : value = dynamic_cast<Operation::AddConstant&>(op).getValue();
614 99543 : else if (op.getId() == Operation::MULTIPLY_CONSTANT)
615 25026 : value = dynamic_cast<Operation::MultiplyConstant&>(op).getValue();
616 74517 : else if (op.getId() == Operation::RECIPROCAL)
617 1652 : value = 1.0;
618 72865 : else if (op.getId() == Operation::STEP)
619 1075 : value = 1.0;
620 71790 : else if (op.getId() == Operation::DELTA)
621 738 : value = std::numeric_limits<double>::infinity();
622 71052 : else if (op.getId() == Operation::ABS) {
623 : long long mask = 0x7FFFFFFFFFFFFFFF;
624 46 : value = *reinterpret_cast<double*>(&mask);
625 : }
626 71006 : else if (op.getId() == Operation::POWER_CONSTANT) {
627 631 : if (stepGroup[step] == -1)
628 55 : value = dynamic_cast<Operation::PowerConstant&>(op).getValue();
629 : else
630 576 : value = 1.0;
631 : }
632 : else
633 70375 : continue;
634 :
635 : // See if we already have a variable for this constant.
636 :
637 388801 : for (int i = 0; i < (int) constants.size(); i++)
638 354910 : if (value == constants[i]) {
639 6263 : operationConstantIndex[step] = i;
640 6263 : break;
641 : }
642 40154 : if (operationConstantIndex[step] == -1) {
643 33891 : operationConstantIndex[step] = constants.size();
644 33891 : constants.push_back(value);
645 : }
646 : }
647 :
648 : // Load constants into variables.
649 :
650 27492 : vector<x86::Xmm> constantVar(constants.size());
651 27492 : if (constants.size() > 0) {
652 : x86::Gp constantsPointer = c.newIntPtr();
653 14028 : c.mov(constantsPointer, imm(&constants[0]));
654 48019 : for (int i = 0; i < (int) constants.size(); i++) {
655 33991 : constantVar[i] = c.newXmmSd();
656 33991 : c.vmovsd(constantVar[i], x86::ptr(constantsPointer, 8*i, 0));
657 : }
658 : }
659 :
660 : // Evaluate the operations.
661 :
662 27492 : vector<bool> hasComputedPower(operation.size(), false);
663 138021 : for (int step = 0; step < (int) operation.size(); step++) {
664 110529 : if (hasComputedPower[step])
665 13603 : continue;
666 :
667 : // When one or more steps involve raising the same argument to multiple integer
668 : // powers, we can compute them all together for efficiency.
669 :
670 110298 : if (stepGroup[step] != -1) {
671 13372 : vector<int>& group = groups[stepGroup[step]];
672 : vector<int>& powers = groupPowers[stepGroup[step]];
673 : x86::Xmm multiplier = c.newXmmSd();
674 13372 : if (powers[0] > 0)
675 13366 : c.vmovsd(multiplier, workspaceVar[arguments[step][0]], workspaceVar[arguments[step][0]]);
676 : else {
677 6 : c.vdivsd(multiplier, constantVar[operationConstantIndex[step]], workspaceVar[arguments[step][0]]);
678 18 : for (int i = 0; i < powers.size(); i++)
679 12 : powers[i] = -powers[i];
680 : }
681 13372 : vector<bool> hasAssigned(group.size(), false);
682 : bool done = false;
683 27222 : while (!done) {
684 : done = true;
685 55232 : for (int i = 0; i < group.size(); i++) {
686 28010 : if (powers[i]%2 == 1) {
687 14308 : if (!hasAssigned[i])
688 13603 : c.vmovsd(workspaceVar[target[group[i]]], multiplier, multiplier);
689 : else
690 705 : c.vmulsd(workspaceVar[target[group[i]]], workspaceVar[target[group[i]]], multiplier);
691 : hasAssigned[i] = true;
692 : }
693 28010 : powers[i] >>= 1;
694 28010 : if (powers[i] != 0)
695 : done = false;
696 : }
697 27222 : if (!done)
698 : c.vmulsd(multiplier, multiplier, multiplier);
699 : }
700 26975 : for (int step : group)
701 : hasComputedPower[step] = true;
702 : continue;
703 : }
704 :
705 : // Evaluate the step.
706 :
707 96926 : Operation& op = *operation[step];
708 96926 : vector<int> args = arguments[step];
709 96926 : if (args.size() == 1) {
710 : // One or more sequential arguments. Fill out the list.
711 :
712 82501 : for (int i = 1; i < op.getNumArguments(); i++)
713 14559 : args.push_back(args[0]+i);
714 : }
715 :
716 : // Generate instructions to execute this operation.
717 :
718 96926 : switch (op.getId()) {
719 3838 : case Operation::CONSTANT:
720 3838 : c.vmovsd(workspaceVar[target[step]], constantVar[operationConstantIndex[step]], constantVar[operationConstantIndex[step]]);
721 96832 : break;
722 : case Operation::ADD:
723 16863 : c.vaddsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
724 : break;
725 : case Operation::SUBTRACT:
726 4031 : c.vsubsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
727 : break;
728 : case Operation::MULTIPLY:
729 14127 : c.vmulsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
730 : break;
731 : case Operation::DIVIDE:
732 8213 : c.vdivsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
733 : break;
734 : case Operation::POWER:
735 6 : generateTwoArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]], pow);
736 : break;
737 2710 : case Operation::NEGATE:
738 2710 : c.vxorps(workspaceVar[target[step]], workspaceVar[target[step]], workspaceVar[target[step]]);
739 2710 : c.vsubsd(workspaceVar[target[step]], workspaceVar[target[step]], workspaceVar[args[0]]);
740 : break;
741 : case Operation::SQRT:
742 1645 : c.vsqrtsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[0]]);
743 : break;
744 : case Operation::EXP:
745 5399 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], exp);
746 : break;
747 : case Operation::LOG:
748 495 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], log);
749 : break;
750 : case Operation::SIN:
751 1291 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], sin);
752 : break;
753 : case Operation::COS:
754 1879 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], cos);
755 : break;
756 : case Operation::TAN:
757 6 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], tan);
758 : break;
759 : case Operation::ASIN:
760 3 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], asin);
761 : break;
762 : case Operation::ACOS:
763 57 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], acos);
764 : break;
765 : case Operation::ATAN:
766 3 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], atan);
767 : break;
768 : case Operation::ATAN2:
769 79 : generateTwoArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]], atan2);
770 : break;
771 : case Operation::SINH:
772 123 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], sinh);
773 : break;
774 : case Operation::COSH:
775 45 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], cosh);
776 : break;
777 : case Operation::TANH:
778 104 : generateSingleArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], tanh);
779 : break;
780 1075 : case Operation::STEP:
781 1075 : c.vxorps(workspaceVar[target[step]], workspaceVar[target[step]], workspaceVar[target[step]]);
782 1075 : c.vcmpsd(workspaceVar[target[step]], workspaceVar[target[step]], workspaceVar[args[0]], imm(18)); // Comparison mode is _CMP_LE_OQ = 18
783 1075 : c.vandps(workspaceVar[target[step]], workspaceVar[target[step]], constantVar[operationConstantIndex[step]]);
784 : break;
785 738 : case Operation::DELTA:
786 738 : c.vxorps(workspaceVar[target[step]], workspaceVar[target[step]], workspaceVar[target[step]]);
787 738 : c.vcmpsd(workspaceVar[target[step]], workspaceVar[target[step]], workspaceVar[args[0]], imm(16)); // Comparison mode is _CMP_EQ_OS = 16
788 738 : c.vandps(workspaceVar[target[step]], workspaceVar[target[step]], constantVar[operationConstantIndex[step]]);
789 : break;
790 : case Operation::SQUARE:
791 0 : c.vmulsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[0]]);
792 : break;
793 : case Operation::CUBE:
794 0 : c.vmulsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[0]]);
795 0 : c.vmulsd(workspaceVar[target[step]], workspaceVar[target[step]], workspaceVar[args[0]]);
796 : break;
797 : case Operation::RECIPROCAL:
798 1652 : c.vdivsd(workspaceVar[target[step]], constantVar[operationConstantIndex[step]], workspaceVar[args[0]]);
799 : break;
800 7148 : case Operation::ADD_CONSTANT:
801 7148 : c.vaddsd(workspaceVar[target[step]], workspaceVar[args[0]], constantVar[operationConstantIndex[step]]);
802 : break;
803 25026 : case Operation::MULTIPLY_CONSTANT:
804 25026 : c.vmulsd(workspaceVar[target[step]], workspaceVar[args[0]], constantVar[operationConstantIndex[step]]);
805 : break;
806 55 : case Operation::POWER_CONSTANT:
807 55 : generateTwoArgCall(c, workspaceVar[target[step]], workspaceVar[args[0]], constantVar[operationConstantIndex[step]], pow);
808 : break;
809 : case Operation::MIN:
810 6 : c.vminsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
811 : break;
812 : case Operation::MAX:
813 6 : c.vmaxsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[1]]);
814 : break;
815 46 : case Operation::ABS:
816 46 : c.vandpd(workspaceVar[target[step]], workspaceVar[args[0]], constantVar[operationConstantIndex[step]]);
817 : break;
818 3 : case Operation::FLOOR:
819 3 : c.vroundsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[0]], imm(1));
820 3 : break;
821 3 : case Operation::CEIL:
822 3 : c.vroundsd(workspaceVar[target[step]], workspaceVar[args[0]], workspaceVar[args[0]], imm(2));
823 3 : break;
824 157 : case Operation::SELECT:
825 : {
826 : x86::Xmm mask = c.newXmmSd();
827 : c.vxorps(mask, mask, mask);
828 157 : c.vcmpsd(mask, mask, workspaceVar[args[0]], imm(0)); // Comparison mode is _CMP_EQ_OQ = 0
829 157 : c.vblendvps(workspaceVar[target[step]], workspaceVar[args[1]], workspaceVar[args[2]], mask);
830 : break;
831 : }
832 : default:
833 : // Just invoke evaluateOperation().
834 :
835 188 : for (int i = 0; i < (int) args.size(); i++)
836 94 : c.vmovsd(x86::ptr(argsPointer, 8*i, 0), workspaceVar[args[i]]);
837 : x86::Gp fn = c.newIntPtr();
838 94 : c.mov(fn, imm((void*) evaluateOperation));
839 : InvokeNode* invoke;
840 0 : c.invoke(&invoke, fn, FuncSignatureT<double, Operation*, double*>());
841 94 : invoke->setArg(0, imm(&op));
842 94 : invoke->setArg(1, imm(&argValues[0]));
843 94 : invoke->setRet(0, workspaceVar[target[step]]);
844 : }
845 : }
846 27492 : c.ret(workspaceVar[workspace.size()-1]);
847 27492 : c.endFunc();
848 27492 : Error error = c.finalize();
849 27492 : if (!error)
850 27492 : error = runtime.add(&jitCode, &code);
851 27492 : if (error)
852 0 : throw Exception(string("AsmJit compilation failed: ") + DebugUtils::errorAsString(error));
853 54984 : }
854 :
855 9405 : static void generateSingleArgCall(x86::Compiler& c, x86::Xmm& dest, x86::Xmm& arg, double (*function)(double)) {
856 : x86::Gp fn = c.newIntPtr();
857 9405 : c.mov(fn, imm((void*) function));
858 : InvokeNode* invoke;
859 9405 : c.invoke(&invoke, fn, FuncSignatureT<double, double>());
860 9405 : invoke->setArg(0, arg);
861 : invoke->setRet(0, dest);
862 9405 : }
863 :
864 140 : static void generateTwoArgCall(x86::Compiler& c, x86::Xmm& dest, x86::Xmm& arg1, x86::Xmm& arg2, double (*function)(double, double)) {
865 : x86::Gp fn = c.newIntPtr();
866 140 : c.mov(fn, imm((void*) function));
867 : InvokeNode* invoke;
868 140 : c.invoke(&invoke, fn, FuncSignatureT<double, double, double>());
869 140 : invoke->setArg(0, arg1);
870 : invoke->setArg(1, arg2);
871 : invoke->setRet(0, dest);
872 140 : }
873 : #endif
874 : #endif
875 :
876 : } // namespace PLMD
|
