@*
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

@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
