Integrating the Wolfram Compiler with the Wolfram Engine

Daniel Sanchez, WRSA

(RowDefault[, (Spacer1), , (Spacer1), ])  #WolframTechConf

Abstract

This talk covers integrating the Wolfram Compiler with the Wolfram Engine, focusing on simplifying the workflow of turning top-level code into compiled code for better performance.

Brief overview of the Wolfram Compiler

Call the Wolfram Engine from compiled code

DownValues compilation

The Wolfram Compiler

a long term project for compilation of the Wolfram Language

converts Wolfram Language programs into low level machine code

integrates generated code efficiently with the rest of the Wolfram Language.
Useful for extending the Wolfram Language itself.

The Wolfram Compiler

Compile a function:

func=FunctionCompile@
Function[{Typed[arg,"Integer64"]},
Module[{factFun},
factFun=Function[{arg1},If[arg1===1,1,arg1*factFun[arg1-1]]];
factFun[arg]
]
]

Can be executed just as any function in the Wolfram Language:

func[10]

Arguments with the wrong types result in failure:

func["10"]

The Wolfram Compiler

Non-trivial code requires declaring functions:

decl=
FunctionDeclaration[
restrictFlow,
Typed[{"Real64","Real64"}->"PackedArray"::["Real64",2]]@
Function[{x,y},
Module[{z},
  z=1-1/(x+I y)^2;
  {{x,y},{Re[z],-Im[z]}}
]
]
];

Use the declaration as a local auxiliary definition:

func=FunctionCompile[decl,
Function[{Typed[x0,"Real64"],Typed[x1,"Real64"],Typed[y0,"Real64"], Typed[y1,"Real64"]},
Table[
Flatten[
Table[
{restrictFlow[x+d/4,y], restrictFlow[x+d/4,-y]},
{x,x0,x1},
{y,y0,y1}
],
2],
{d, 0,3, 0.3}
]
]
]

A vector plot based on the data that was created:

ListVectorPlot[First[EchoTiming[func[-7,5,0.5, 3]]]]

The Wolfram Compiler

What if one already has code and doesn’t want to rewrite it for compilation?

restrictFlow[x_, y_]:=
Module[{z},
  z=1-1/(x+I y)^2;
  {{x,y},{Re[z],-Im[z]}}
];

fun[x0_,x1_,y0_, y1_]:=
Table[
Flatten[
Table[
{restrictFlow[x+d/4,y], restrictFlow[x+d/4,-y]},
{x,x0,x1},
{y,y0,y1}
],
2],
{d, 0,3, 0.3}
];

ListVectorPlot[First[EchoTiming[fun[-7,5,0.5, 3]]]]

Calling the Wolfram Engine from compiled code

interaction between evaluator code and compiled code:

an interface sits at the boundary, managing communication between the evaluator and compiled code:

Inert Expressions

Expression is the one uniform data type in the Wolfram Engine

The “InertExpression” type exists to work with Wolfram Language expressions

(HyperlinkURL["http://reference.wolfram.com/language/ref/compiledtype/InertExpression.html", http://reference.wolfram.com/language/ref/compiledtype/InertExpression.html])

Typed[expr, "InertExpression"]

Called Inert because, unlike regular Wolfram Language expressions, objects of the type “InertExpression” don’t evaluate on its own inside compiled code

Inert Expressions

Compile a function whose argument is an expression and returns its head:

exprHeadFunc=FunctionCompile@Function[{Typed[expr,"InertExpression"]},
Head[expr]
]

The “InertExpression” type will accept any argument passed by the Wolfram Evaluator:

exprHeadFunc[#]&/@{42,z,Range[10],<||>}

Inert Expressions

Construct inert expressions inside compiled code:

constructFunc=FunctionCompile@
Function[{Typed[head,"InertExpression"],Typed[arg,"InertExpression"]},
Construct[head,arg]
]

Construct[head,arg]

The result seems to evaluate because once it gets back to top-level, it evaluates on its own:

constructFunc[Echo,"foo"]

If we don't return the inert expressions, notice how it doesn' t get evaluated :

constructFunc2=FunctionCompile@
Function[{Typed[head,"InertExpression"],Typed[arg,"InertExpression"]},
Construct[head,arg]; (* Do not return the constructed expression *)
42
]

constructFunc2[Echo,"foo"]

Inert Expressions

To evaluate inert expressions inside compiled code, use InertEvaluate:

evalConstructFunc=FunctionCompile@
Function[{Typed[head,"InertExpression"],Typed[arg,"InertExpression"]},
InertEvaluate@Construct[head,arg];
42
]

The expression gets evaluated inside compiled code:

evalConstructFunc[Echo,"foo"]

Inert Expressions

Transliterate a string inside compiled code:

transliterateC1=FunctionCompile@
Function[{Typed[expr,"InertExpression"]},
InertEvaluate@Construct[InertExpression[Transliterate],expr]
]

?InertExpression

The expression gets evaluated inside compiled code:

transliterateC1["しんばし"]

Notice the compiled function also accepts a non-string argument:

transliterateC1[x]

Inert Expressions

“String” is the type we want for the argument:

transliterateC2=FunctionCompile@
Function[{Typed[str,"String"]},
InertEvaluate@Construct[InertExpression[Transliterate],str]
]

Use Cast to convert the string to an “InertExpression” object:

transliterateC2=FunctionCompile@
Function[{Typed[str,"String"]},
InertEvaluate@Construct[InertExpression[Transliterate],Cast[str,"InertExpression"]]
]

The type of the argument is correctly validated:

Map[transliterateC2,{"しんばし",x}]

Inert Expressions

InertEvaluate returns an "InertExpression" :

transliterateLengthC=FunctionCompile@
Function[{Typed[str,"String"]},
Module[{strExpr,res},
strExpr=Cast[str,"InertExpression"];
res=InertEvaluate@Construct[InertExpression[Transliterate],strExpr];
StringLength[res]
]
]

Needs to be casted to the appropriate type:

transliterateLengthC=FunctionCompile@
Function[{Typed[str,"String"]},
Module[{strExpr,resExpr},
strExpr=Cast[str,"InertExpression"];
resExpr=InertEvaluate@Construct[InertExpression[Transliterate],strExpr];
StringLength[Cast[resExpr,"String"]]
]
]

transliterateLengthC["しんばし"]

KernelFunction

?KernelFunction

Let’s use KernelFunction to call Transliterate from compiled code:

transliterateC3=FunctionCompile@
Function[{Typed[str,"String"]},
Module[{trFun=TypeHint[KernelFunction[Transliterate],{"String"}->"String"]},
trFun[str]
]
]

Notice the compiled function returns a “String”:

transliterateC3["しんばし"]

transliterateC3[x]

KernelFunction

For repeated use of a kernel function, it’s best to declare it with a function definition:

decl=FunctionDeclaration[
Transliterate,
Typed[{"String"}->"String"]@
Function[{str},
TypeHint[KernelFunction[Transliterate],{"String"}->"String"][str]
]
];

transliterateLengthC2=FunctionCompile[decl,
Function[{Typed[str,"String"]},
StringLength[Transliterate[str]]
]
]

transliterateLengthC2["しんばし"]

KernelFunction

KernelFunction has a short-hand notation:

decl=FunctionDeclaration[
Transliterate,
Typed[{"String"}->"String"]@
KernelFunction[Transliterate]
];

transliterateLengthC2=FunctionCompile[decl,
Function[{Typed[str,"String"]},
StringLength[Transliterate[str]]
]
]

transliterateLengthC2["しんばし"]

KernelFunction

KernelFunction breaks down into InertEvaluate + Construct + InertExpression

FunctionDeclaration[
Transliterate,
Typed[{"String"}->"String"]@
KernelFunction[Transliterate]
]

FunctionDeclaration[
	Transliterate,
	Typed[
		{"String"} -> "String"
	] @
	Function[{arg1},
		Cast[
			InertEvaluate @ Construct[
				InertExpression[Transliterate],
				Cast[arg1, "InertExpression"]
			],
			"String"
		]
	]
]

KernelFunction

Now we can see how to reuse the top-level code for restrictFlow:

restrictFlow[x_, y_]:=
Module[{z},
  z=1-1/(x+I y)^2;
  {{x,y},{Re[z],-Im[z]}}
];

decl=FunctionDeclaration[
restrictFlow,
Typed[{"Real64","Real64"}->"PackedArray"::["Real64",2]]@
KernelFunction[restrictFlow]
];

func=FunctionCompile[
decl,
Function[{Typed[x0,"Real64"],Typed[x1,"Real64"],Typed[y0,"Real64"], Typed[y1,"Real64"]},
Table[
Flatten[
Table[{restrictFlow[x+d/4,y], restrictFlow[x+d/4,-y]},{x,x0,x1},{y,y0,y1}],
2],
{d, 0,3, 0.3}
]
]
]

However, there’s no real speedup here because the restricted flow code isn’t really compiled:

ListVectorPlot[First[EchoTiming[func[-7,5,0.5, 3]]]]

DownValuesFunction

DownValuesFunction allows us to compile the DownValues of a symbol:

?DownValuesFunction

Define a function in top-level and use its DownValues to compile a function for the integer-typed case:

square[x_]:=x^2;

func=FunctionCompile@
Function[Typed[arg,"MachineInteger"],
Module[{squareFun},
squareFun=TypeHint[DownValuesFunction[square],{"MachineInteger"}->"MachineInteger"];
squareFun[arg]
]
]

Use the compiled function with an integer argument:

func[100]

DownValuesFunction

Let’s see how to use DownValuesFunction for the restrict flow example:

restrictFlow[x_, y_]:=
Module[{z},
  z=1-1/(x+I y)^2;
  {{x,y},{Re[z],-Im[z]}}
];

Write a declaration for restrictFlow that uses the DownValues from top-level code:

decl=FunctionDeclaration[
restrictFlow,
Typed[{"Real64","Real64"}->"PackedArray"::["Real64",2]]@
DownValuesFunction[restrictFlow]
];

func=FunctionCompile[decl,
Function[{Typed[x0,"Real64"],Typed[x1,"Real64"],Typed[y0,"Real64"], Typed[y1,"Real64"]},
Table[
Flatten[
Table[{restrictFlow[x+d/4,y], restrictFlow[x+d/4,-y]},{x,x0,x1},{y,y0,y1}],
2],
{d, 0,3, 0.3}
]
]
]

ListVectorPlot[First[EchoTiming[func[-7,5,0.5, 3]]]]

DownValuesFunction

In most cases fun will also be defined with DownValues (not Function):

restrictFlow[x_, y_]:=
Module[{z},
  z=1-1/(x+I y)^2;
  {{x,y},{Re[z],-Im[z]}}
];

fun[x0_,x1_,y0_,y1_]:=
Table[
Flatten[
Table[{restrictFlow[x+d/4,y], restrictFlow[x+d/4,-y]},{x,x0,x1},{y,y0,y1}],
2],
{d, 0,3, 0.3}
];

Let’s time the computation in top-level code:

First@AbsoluteTiming[fun[-7,5,0.5, 3];]

Now also use DownValuesFunction for compiling  fun:

func=FunctionCompile[
FunctionDeclaration[
restrictFlow,
Typed[{"Real64","Real64"}->"PackedArray"::["Real64",2]]@
DownValuesFunction[restrictFlow]
],
Typed[{"Real64","Real64","Real64","Real64"}->"PackedArray"::["Real64",4]]@
DownValuesFunction[fun]
]

Let’s time the computation in top-level code:

First@AbsoluteTiming[func[-7,5,0.5, 3];]

/

DownValuesFunction

What if a symbol has multiple DownValues?

Clear[myFun]
myFun[x_,x_]:=1
myFun[x_,y_?EvenQ]:=2
myFun[x_,y_/;y>2]:=3
myFun[x_,2]:=4

Compile only the DownValues that accept two machine integers:

myFunC=FunctionCompile[
Typed[DownValuesFunction[myFun],{"MachineInteger","MachineInteger"}->"MachineInteger"]
]

Test the compiled function with argument combinations:

myFunC[1,1]

myFunC[1,3]

DownValuesFunction

DownValues are compiled in-order:

myFunC[1,2]

DownValues[myFun]

When arguments aren’t handled, a runtime error is thrown:

myFunC[2,1]

DownValuesFunction

What is DownValuesFunction really doing?

myFun[x_, x_] := b1
myFun[x_, y_?testQ] := b2
myFun[x_, y_ /; y > 2] := b3
myFun[x_, 2] := b4
myFun[x_, y_] := b5

Typed[DownValuesFunction[myFun],{"MachineInteger","MachineInteger"}->"MachineInteger"]

Function[{arg1, arg2},
	Module[{downValFun1, downValFun2, ..., condFun3},
		(* DownValue functions *)
		downValFun1 = Function[{x, y}, b1];
		downValFun2 = Function[{x, y}, b2];
		downValFun3 = Function[{x, y}, b3];
		downValFun4 = Function[{x, y}, b4];
		downValFun5 = Function[{x, y}, b5];

		(* Conditions *)
		condFun3 = Function[{y}, y > 2];

		Which[
			arg1 === arg2,
				downValFun1[arg1, arg2]
			,
			testQ[arg2],
				downValFun2[arg1, arg2]
			,
			condFun3[arg2],
				downValFun3[arg1, arg2]
			,
            arg2 === 2,
				downValFun4[arg1, arg2]
			,
            True,
				downValFun5[arg1, arg2] (* If there isn't a fallthrough => exception *)
        ]
    ]
  ]

Thank you! Questions?

Daniel Sánchez
daniels@wolfram.com

For general discussions about Wolfram Compiler:
	l-compiler@wolfram.com

Other example

IsPalindrome[ds_]:=
Module[{start=0,end=ds["Length"]+1},
While[++start<--end,
If[ds["Part",start]=!=ds["Part",end],
Return[False]
]
];
True
];

ds=CreateDataStructure["DynamicArray", {9,2,7,0}];

IsPalindrome[ds]

ds=CreateDataStructure["DynamicArray", {9,2,2,9}];

IsPalindrome[ds]

IsPalindrome[ds]//RepeatedTiming//First

decl=FunctionDeclaration[
IsPalindrome,
Typed[{"DynamicArray"}->"Boolean"]@
DownValuesFunction[IsPalindrome]
];

IsPalindromeC=FunctionCompile[decl,Function[{Typed[ds,"DynamicArray"]},IsPalindrome[ds]]]

IsPalindromeC[ds]//RepeatedTiming//First

/

KernelEvaluate

Yet another way of evaluating code in the Wolfram Engine is via KernelEvaluate:

?KernelEvaluate

Integrate (RefLink[NDSolve, paclet:ref/NDSolve]) into compiled code:

func=FunctionCompile@Function[{Typed[init,"Real64"],Typed[val,"Real64"]},
Module[{iPoly},
iPoly=TypeHint[
KernelEvaluate[init,
First[y/.NDSolve[{y'[x]==y[x] Cos[x+y[x]],y[0]==init},y,{x,0,30}]]
],
"InertExpression"
];
TypeHint[KernelEvaluate[{iPoly,val},iPoly[val]],"Real64"]
]]

func[0.2,0.0]

KernelEvaluate

KernelEvaluate also breaks down into InertEvaluate + Construct + InertExpression:

Function[{init},
KernelEvaluate[init,
First[y/.NDSolve[{y'[x]==y[x] Cos[x+y[x]],y[0]==init},y,{x,0,30}]]
]
]

Function[{init},
	Cast[
		InertEvaluate @ Construct[
			InertExpression[With],
			Construct[
				InertExpression[List],
				Construct[
					InertExpression[Set],
					InertExpression[init],
					Cast[init, "InertExpression"]
				]
			]
			InertExpression[First[y/.NDSolve[{y'[x]==y[x] Cos[x+y[x]],y[0]==init},y,{x,0,30}]]]
		],
		_
	]
]