Skip to content
This repository was archived by the owner on Oct 9, 2026. It is now read-only.
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 44 additions & 9 deletions DiFfRG/Mathematica/DiFfRG/CodeTools/MakeKernel.m
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ This Function creates an integrator that evaluates (constantFlow + \[Integral]in
These are prepended to the respective methods of the integration kernel, allowing one to e.g. define specific angles one needs for the flow code.

The options \"KernelReturnTransform\" and \"ConstantReturnTransform\" (default Identity) accept a Mathematica function applied to the optimized expression before code generation, which lets you wrap the return value, e.g. \"KernelReturnTransform\" -> Re renders the kernel return as real(...).
The option \"ComputeType\" (default \"double\") is the value type the kernel is evaluated and integrated in: \"double\", \"float\", \"DiFfRG::complex<double>\" or \"DiFfRG::complex<float>\". With a float type the kernel, its literals and the integrator run in single precision, while map()/get() still write double results.
The option \"KernelTraits\" declares integrator traits on the emitted kernel class, as a list of names or name -> Boolean rules, e.g. \"KernelTraits\" -> {\"matsubara_finite_extent\"} or {\"matsubara_split\" -> True}. Each becomes a `static constexpr bool <name> = <value>;` member. \"MatsubaraEven\" stays a separate option because it carries a symbolic evenness check that a generic mechanism cannot.";

MakeKernel::Invalid = "The given arguments are invalid. See MakeKernel::usage";
Expand All @@ -30,6 +31,8 @@ This Function creates an integrator that evaluates (constantFlow + \[Integral]in

MakeKernel::notEven = "MatsubaraEven requested for kernel \"`1`\" but it is not even in \"`2`\"; emitting the standard kernel (the integrator keeps the explicit kernel(+f0)+kernel(-f0) form). This is expected when the loop contains fermionic dressings evaluated at f0-shifted arguments.";

MakeKernel::ctypedeprecated = "The option \"ctype\" is deprecated; use \"ComputeType\" -> `1` instead.";

MakeKernel::InvalidTrait = "KernelTraits entry `1` is not a C++ identifier optionally followed by -> True|False.";

Begin["`Private`"]
Expand Down Expand Up @@ -69,13 +72,44 @@ This Function creates an integrator that evaluates (constantFlow + \[Integral]in

KernelSpecQ[spec_Association] :=
Module[{validKeys, validKeyTypes},
validKeys = CheckKey[spec, "Name", StringQ[#] && StringLength[#] > 0&, "Cannot be empty"] && CheckKey[spec, "Integrator", StringQ[#] && StringLength[#] > 0&, "Cannot be empty"] && CheckKey[spec, "d", IntegerQ[#] && # >= 0&, "Must be an Integer >= 0"] && CheckKey[spec, "AD", BooleanQ, "Must be a Boolean"] && CheckKey[spec, "Device", MemberQ[{"Threads", "TBB", "GPU"}, #]&, "Must be Threads, TBB or GPU."] && CheckKey[spec, "Type", StringQ[#] && StringLength[#] > 0&, "Cannot be empty"];
validKeys = CheckKey[spec, "Name", StringQ[#] && StringLength[#] > 0&, "Cannot be empty"] && CheckKey[spec, "Integrator", StringQ[#] && StringLength[#] > 0&, "Cannot be empty"] && CheckKey[spec, "d", IntegerQ[#] && # >= 0&, "Must be an Integer >= 0"] && CheckKey[spec, "AD", BooleanQ, "Must be a Boolean"] && CheckKey[spec, "Device", MemberQ[{"Threads", "TBB", "GPU"}, #]&, "Must be Threads, TBB or GPU."] && CheckKey[spec, "Type", StringQ[#] && StringLength[#] > 0&, "Cannot be empty"] && CheckKey[spec, "ComputeType", StringQ[#] && StringLength[#] > 0&, "Cannot be empty"];
Return[validKeys];
];

GetStandardKernelDefinitions[] :=
$StandardKernelDefinitions

(* "ctype" is the old name of "ComputeType"; an explicit ctype still wins so old scripts keep working. *)
resolveComputeType[spec_Association] :=
Module[{s = spec},
If[s["ctype"] =!= Automatic,
Message[MakeKernel::ctypedeprecated, ToString[s["ctype"], InputForm]];
s["ComputeType"] = s["ctype"]
];
KeyDrop[s, "ctype"]
];

singlePrecisionQ[spec_Association] :=
StringContainsQ[spec["ComputeType"], "float"];

(* Real type of the kernel's arguments. *)
realType[spec_Association] :=
If[singlePrecisionQ[spec], "float", "double"];

(* map()/get() hand results back in double even when the integrator runs in float, so the
caller's double state vectors need no staging. *)
resultType[spec_Association] :=
StringReplace[spec["ComputeType"], "float" -> "double"];

(* Print the kernel's literals as float (0.5f, complex<float>) in single precision; a double
literal would silently promote the whole expression back to double. *)
SetAttributes[withComputePrecision, HoldRest];

withComputePrecision[spec_Association, body_] :=
Block[{FunKit`Private`$codePrecision = If[singlePrecisionQ[spec], "single", FunKit`Private`$codePrecision]},
body
];

(* Emit arbitrary integrator traits as class members. The traits are what the finite-T
integrator branches on (matsubara_even, matsubara_finite_extent, matsubara_split), and a
`requires { requires K::trait; }` detector reads FALSE for anything it cannot see -- a
Expand Down Expand Up @@ -114,7 +148,7 @@ integrator branches on (matsubara_even, matsubara_finite_extent, matsubara_split

(* Internal functions added here with Internal`*::usage *)

Options[MakeKernel] = {"Coordinates" -> {}, "CoordinateArguments" -> {}, "IntegrationVariables" -> {}, "KernelDefinitions" -> $StandardKernelDefinitions, "Regulator" -> "DiFfRG::PolynomialExpRegulator", "RegulatorOpts" -> {"", ""}, "KernelBody" -> "", "KernelReturnType" -> "auto", "KernelReturnTransform" -> Identity, "ConstantBody" -> "", "ConstantReturnType" -> "auto", "ConstantReturnTransform" -> Identity, "Parameters" -> {}, "Name" -> "", "d" -> -1, "Integrator" -> "", "AD" -> False, "ctype" -> "double", "Device" -> "TBB", "Type" -> "double", "SplitKernel" -> False, "SeparateLookups" -> False, "Decorator" -> "static KOKKOS_FUNCTION", "MatsubaraEven" -> False, "KernelTraits" -> {}};
Options[MakeKernel] = {"Coordinates" -> {}, "CoordinateArguments" -> {}, "IntegrationVariables" -> {}, "KernelDefinitions" -> $StandardKernelDefinitions, "Regulator" -> "DiFfRG::PolynomialExpRegulator", "RegulatorOpts" -> {"", ""}, "KernelBody" -> "", "KernelReturnType" -> "auto", "KernelReturnTransform" -> Identity, "ConstantBody" -> "", "ConstantReturnType" -> "auto", "ConstantReturnTransform" -> Identity, "Parameters" -> {}, "Name" -> "", "d" -> -1, "Integrator" -> "", "AD" -> False, "ComputeType" -> "double", "ctype" -> Automatic, "Device" -> "TBB", "Type" -> "double", "SplitKernel" -> False, "SeparateLookups" -> False, "Decorator" -> "static KOKKOS_FUNCTION", "MatsubaraEven" -> False, "KernelTraits" -> {}};

MakeKernel[__] :=
(
Expand All @@ -128,6 +162,7 @@ integrator branches on (matsubara_even, matsubara_finite_extent, matsubara_split
MakeKernel[kernelExpr_, constExpr_, OptionsPattern[]] :=
Module[{expr, const, exec, kernel, constant, kernelClass, kernelHeader, integratorHeader, integratorCpp, integratorTemplateParams, tparams = <|"Name" -> "...t", "Type" -> "auto&&", "Reference" -> False, "Const" -> False|>, kernelDefs = OptionValue["KernelDefinitions"], coordinates = OptionValue["Coordinates"], getArgs = OptionValue["CoordinateArguments"], intVariables = OptionValue["IntegrationVariables"], preArguments, regulator, params, adSpecs, explParamAD, arguments, outputPath, sources, returnType, returnTypePointer, spec, parameters, parametersKernel, matsubaraEvenTrait, kernelTraits},
spec = Association @@ Thread[Rule @@ {#, OptionValue[MakeKernel, #]}]& @ Keys[Options[MakeKernel]];
spec = resolveComputeType[spec];
If[Not @ KernelSpecQ[spec],
Message[MakeKernel::InvalidSpec];
Abort[]
Expand Down Expand Up @@ -166,9 +201,9 @@ GENUINELY even in that variable (kernel(+f0) == kernel(-f0)), the integrator
const = constExpr;
While[ListQ[const], const = Plus @@ const];
intVariables = FunKit`Private`prepParam /@ intVariables;
intVariables = Map[Append[#, "Type" -> "double"]&, intVariables];
intVariables = Map[Append[#, "Type" -> realType[spec]]&, intVariables];
getArgs = FunKit`Private`prepParam /@ getArgs;
getArgs = Map[Append[#, "Type" -> "double"]&, getArgs];
getArgs = Map[Append[#, "Type" -> realType[spec]]&, getArgs];
(********************************************************************)
(* First, the kernel itself *)
(********************************************************************)
Expand All @@ -188,13 +223,13 @@ GENUINELY even in that variable (kernel(+f0) == kernel(-f0)), the integrator
,
spec["Parameters"]
];
kernel =
kernel = withComputePrecision[spec,
If[TrueQ[OptionValue["SplitKernel"]] || TrueQ[OptionValue["SeparateLookups"]],
FunKit`MakeCppFunctionSplit[expr, "Name" -> "kernel", "Return" -> OptionValue["KernelReturnType"], "Suffix" -> "", "Prefix" -> "static KOKKOS_INLINE_FUNCTION", "Decorator" -> OptionValue["Decorator"], "SeparateLookups" -> OptionValue["SeparateLookups"], "Parameters" -> Join[intVariables, getArgs, parametersKernel], "Body" -> StringTemplate["using namespace DiFfRG;using namespace DiFfRG::compute;\n`1`"][OptionValue["KernelBody"]], "ReturnTransform" -> OptionValue["KernelReturnTransform"]]
,
FunKit`MakeCppFunction[expr, "Name" -> "kernel", "Return" -> OptionValue["KernelReturnType"], "Suffix" -> "", "Prefix" -> "static KOKKOS_INLINE_FUNCTION", "Parameters" -> Join[intVariables, getArgs, parametersKernel], "Body" -> StringTemplate["using namespace DiFfRG;using namespace DiFfRG::compute;\n`1`"][OptionValue["KernelBody"]], "ReturnTransform" -> OptionValue["KernelReturnTransform"]]
];
constant = FunKit`MakeCppFunction[constExpr, "Name" -> "constant", "Return" -> OptionValue["ConstantReturnType"], "Suffix" -> "", "Prefix" -> "static KOKKOS_INLINE_FUNCTION", "Parameters" -> Join[getArgs, parametersKernel], "Body" -> StringTemplate["using namespace DiFfRG;using namespace DiFfRG::compute;\n`1`"][OptionValue["ConstantBody"]], "ReturnTransform" -> OptionValue["ConstantReturnTransform"]];
]];
constant = withComputePrecision[spec, FunKit`MakeCppFunction[constExpr, "Name" -> "constant", "Return" -> OptionValue["ConstantReturnType"], "Suffix" -> "", "Prefix" -> "static KOKKOS_INLINE_FUNCTION", "Parameters" -> Join[getArgs, parametersKernel], "Body" -> StringTemplate["using namespace DiFfRG;using namespace DiFfRG::compute;\n`1`"][OptionValue["ConstantBody"]], "ReturnTransform" -> OptionValue["ConstantReturnTransform"]]];
kernelTraits = kernelTraitMembers[OptionValue["KernelTraits"]];
kernelClass = FunKit`MakeCppClass["TemplateTypes" -> {"_Regulator"}, "Name" -> OptionValue["Name"] <> "_kernel", "MembersPublic" -> Join[{"using Regulator = _Regulator;"}, matsubaraEvenTrait, kernelTraits, {kernel, constant}], "MembersPrivate" -> kernelDefs];
kernelHeader = FunKit`MakeCppHeader["Includes" -> {"DiFfRG/physics/interpolation.hh", "DiFfRG/physics/physics.hh"}, "Body" -> {"namespace DiFfRG {", kernelClass, StringTemplate["} using DiFfRG::`1`_kernel;"][spec["Name"]]}];
Expand Down Expand Up @@ -222,11 +257,11 @@ GENUINELY even in that variable (kernel(+f0) == kernel(-f0)), the integrator
];
integratorTemplateParams = TemplateParameterGeneration[spec];
integratorTemplateParams = StringRiffle[integratorTemplateParams, ", "];
returnType = spec["ctype"];
returnType = resultType[spec];
returnTypePointer = StringTemplate["`1`*"][returnType];
adSpecs =
If[spec["AD"],
Map[Merge[{#, <|"IntegratorTemplateParams" -> StringRiffle[TemplateParameterGeneration[spec, #["Replacements"]], ", "], "ReturnType" -> spec["ctype"] /. #["Replacements"], "Params" -> Last @ processParameters[params, #["Replacements"]]|>}, Last]&, $ADSpecializations]
Map[Merge[{#, <|"IntegratorTemplateParams" -> StringRiffle[TemplateParameterGeneration[spec, #["Replacements"]], ", "], "ReturnType" -> resultType[spec] /. #["Replacements"], "Params" -> Last @ processParameters[params, #["Replacements"]]|>}, Last]&, $ADSpecializations]
,
{}
];
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,16 +40,16 @@

appendKeyType[templateParameter_List, params_Association, {}] :=
Module[{},
If[KeyExistsQ[params, "ctype"],
Append[templateParameter, ToString[params["ctype"]]],
If[KeyExistsQ[params, "ComputeType"],
Append[templateParameter, ToString[params["ComputeType"]]],
Append[templateParameter, "double"]
]
]

appendKeyType[templateParameter_List, params_Association, ADReplacements_] :=
Module[{},
If[KeyExistsQ[params, "ctype"],
Append[templateParameter, ToString[params["ctype"]] /. ADReplacements],
If[KeyExistsQ[params, "ComputeType"],
Append[templateParameter, ToString[params["ComputeType"]] /. ADReplacements],
Append[templateParameter, "autodiff::real"]
]
]
Expand Down
52 changes: 51 additions & 1 deletion DiFfRG/Mathematica/DiFfRG/Tests/MakeKernelSecondOrderADTests.m
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
"Name" -> "pion",
"Integrator" -> "Integrator_p2",
"d" -> 3,
"ctype" -> "DiFfRG::complex<double>",
"ComputeType" -> "DiFfRG::complex<double>",
"Parameters" -> {
<|"Name" -> "k", "Type" -> "double", "Const" -> True, "AD" -> False|>,
<|"Name" -> "T", "Type" -> "double", "Const" -> True, "AD" -> False|>,
Expand Down Expand Up @@ -80,3 +80,53 @@
Block[{Print}, Get["FunKit`"]];
AUMPCHECK[Quiet[generatesSecondOrderComplexAD[], OptionValue::nodef]];
];

(* A float ComputeType must leave no double in the kernel (it would promote the arithmetic back),
while map()/get() keep double destinations for the caller's state vectors. *)
generatesFloatKernel[typeOption_] :=
Module[{tmp, header, kernelHeader},
tmp = FileNameJoin[{AUMPTestTempDirectory[], "generated"}];
CreateDirectory[tmp];
SetFlowDirectory[tmp <> "/"];
CreateDirectory[FileNameJoin[{tmp, "flows", "fl", "src"}], CreateIntermediateDirectories -> True];
MakeKernel[
0.5 l1 + 2 k + I l1,
"IntegrationVariables" -> {"l1"},
"Coordinates" -> {"LinearCoordinates1D<float>"},
"Name" -> "fl",
"Integrator" -> "Integrator_p2",
"d" -> 3,
typeOption -> "DiFfRG::complex<float>",
"Parameters" -> {<|"Name" -> "k", "Type" -> "float", "Const" -> True|>}
];
header = Import[FileNameJoin[{tmp, "flows", "fl", "fl.hh"}], "Text"];
kernelHeader = Import[FileNameJoin[{tmp, "flows", "fl", "kernel.hh"}], "Text"];
containsAll[
header,
{
"Integrator_p2<3, DiFfRG::complex<float>, fl_kernel<Regulator>, DiFfRG::TBB_exec> integrator;",
"map(DiFfRG::complex<double>* dest, const LinearCoordinates1D<float>& coordinates",
"void get(DiFfRG::complex<double>& dest"
}
] &&
containsAll[kernelHeader, {"const float& l1", "0.5f", "complex<float>("}] &&
StringFreeQ[kernelHeader, "double"]
];

AUMPTestCase["MakeKernel with a float ComputeType emits a pure-float kernel", {"make-kernel", "precision", "funkit", "form"},
AUMPAssume[Length @ PacletFind["FunKit"] > 0, "FunKit is not installed"];
AUMPAssume[formAvailableQ[], "FORM is not installed"];
Needs["FormTracer`"];
FormTracer`DefineFormExecutable[formExecutablePath[]];
Block[{Print}, Get["FunKit`"]];
AUMPCHECK[Quiet[generatesFloatKernel["ComputeType"], OptionValue::nodef]];
];

AUMPTestCase["MakeKernel still accepts the deprecated ctype option", {"make-kernel", "precision", "funkit", "form"},
AUMPAssume[Length @ PacletFind["FunKit"] > 0, "FunKit is not installed"];
AUMPAssume[formAvailableQ[], "FORM is not installed"];
Needs["FormTracer`"];
FormTracer`DefineFormExecutable[formExecutablePath[]];
Block[{Print}, Get["FunKit`"]];
AUMPCHECK[Quiet[generatesFloatKernel["ctype"], {OptionValue::nodef, MakeKernel::ctypedeprecated}]];
];
Original file line number Diff line number Diff line change
Expand Up @@ -3,22 +3,22 @@

AUMPTestCase["TemplateParameterGeneration supports GPU float kernels", {"template-parameters"},
AUMPCHECKEqual[
TemplateParameterGeneration[<|"d" -> 2, "Name" -> "MyKernel", "ctype" -> "float", "Device" -> "GPU"|>],
TemplateParameterGeneration[<|"d" -> 2, "Name" -> "MyKernel", "ComputeType" -> "float", "Device" -> "GPU"|>],
{"2", "float", "MyKernel_kernel<Regulator>", "DiFfRG::GPU_exec"}
];
];

AUMPTestCase["TemplateParameterGeneration supports TBB double kernels", {"template-parameters"},
AUMPCHECKEqual[
TemplateParameterGeneration[<|"d" -> 3, "Name" -> "Integrator", "ctype" -> "double", "Device" -> "TBB"|>],
TemplateParameterGeneration[<|"d" -> 3, "Name" -> "Integrator", "ComputeType" -> "double", "Device" -> "TBB"|>],
{"3", "double", "Integrator_kernel<Regulator>", "DiFfRG::TBB_exec"}
];
];

AUMPTestCase["TemplateParameterGeneration applies real AD replacements", {"template-parameters", "ad"},
AUMPCHECKEqual[
TemplateParameterGeneration[
<|"d" -> 3, "Name" -> "Integrator", "ctype" -> "double", "Device" -> "TBB"|>,
<|"d" -> 3, "Name" -> "Integrator", "ComputeType" -> "double", "Device" -> "TBB"|>,
{"double" -> "autodiff::real"}
],
{"3", "autodiff::real", "Integrator_kernel<Regulator>", "DiFfRG::TBB_exec"}
Expand All @@ -28,7 +28,7 @@
AUMPTestCase["TemplateParameterGeneration applies second-order complex AD replacements", {"template-parameters", "ad"},
AUMPCHECKEqual[
TemplateParameterGeneration[
<|"d" -> 3, "Name" -> "pion", "ctype" -> "DiFfRG::complex<double>", "Device" -> "TBB"|>,
<|"d" -> 3, "Name" -> "pion", "ComputeType" -> "DiFfRG::complex<double>", "Device" -> "TBB"|>,
{"DiFfRG::complex<double>" -> "cxReal<2, double>"}
],
{"3", "cxReal<2, double>", "pion_kernel<Regulator>", "DiFfRG::TBB_exec"}
Expand All @@ -37,12 +37,12 @@

AUMPTestCase["TemplateParameterGeneration supports Threads execution", {"template-parameters"},
AUMPCHECKEqual[
TemplateParameterGeneration[<|"d" -> 1, "Name" -> "Test", "ctype" -> "float", "Device" -> "Threads"|>],
TemplateParameterGeneration[<|"d" -> 1, "Name" -> "Test", "ComputeType" -> "float", "Device" -> "Threads"|>],
{"1", "float", "Test_kernel<Regulator>", "DiFfRG::KokkosHost_exec"}
];
];

AUMPTestCase["TemplateParameterGeneration defaults ctype to double", {"template-parameters", "defaults"},
AUMPTestCase["TemplateParameterGeneration defaults ComputeType to double", {"template-parameters", "defaults"},
AUMPCHECKEqual[
TemplateParameterGeneration[<|"d" -> 4, "Name" -> "DefaultType", "Device" -> "TBB"|>],
{"4", "double", "DefaultType_kernel<Regulator>", "DiFfRG::TBB_exec"}
Expand All @@ -51,7 +51,7 @@

AUMPTestCase["TemplateParameterGeneration defaults Device to TBB", {"template-parameters", "defaults"},
AUMPCHECKEqual[
TemplateParameterGeneration[<|"d" -> 2, "Name" -> "DefaultDevice", "ctype" -> "double"|>],
TemplateParameterGeneration[<|"d" -> 2, "Name" -> "DefaultDevice", "ComputeType" -> "double"|>],
{"2", "double", "DefaultDevice_kernel<Regulator>", "DiFfRG::TBB_exec"}
];
];
Expand Down
19 changes: 19 additions & 0 deletions DiFfRG/include/DiFfRG/common/types.hh
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,25 @@ namespace DiFfRG

template <typename CT> using ctype = typename internal::_ctype<CT>::value;

namespace internal
{
template <typename T> struct _double_precision {
using value = T;
};

template <> struct _double_precision<float> {
using value = double;
};

template <> struct _double_precision<complex<float>> {
using value = complex<double>;
};
} // namespace internal

/// The double-precision counterpart of a single-precision type (float -> double,
/// complex<float> -> complex<double>); every other type maps to itself.
template <typename T> using double_precision = typename internal::_double_precision<T>::value;

template <typename T> inline constexpr bool is_autodiff = internal::_is_autodiff<T>;
} // namespace get_type
} // namespace DiFfRG
Loading
Loading