Programming Techniques in the Next-Generation Wolfram Compiler

Abdul Dakkak
Kernel Developer
Wolfram Research

The Compiler In a Nutshell

Issues

Due to the Large Code Base

100k lines of Mathematica code
20k lines for Runtime
30k lines for Links

Due to the Compiler

Dealing With a Large Code Base

Separation of Concerns

Make the entire process a pipeline with components that can be run at any stage. We call these passes

Keys[$Passes]

Each pass has a set of pre and post dependencies

$Passes["CopyPropagation"]

$Passes["CopyPropagation"]["requires"]

We can look at the dependencies

Graph[Flatten[Table[pass["name"]\[DirectedEdge]#["name"]&/@pass["requires"],{pass,$Passes}]],VertexLabels->"Name",EdgeShapeFunction->"Line",VertexSize->Tiny,GraphLayout->"LayeredDigraphEmbedding",ImagePadding->15]

Can selectively run different passes

pm=CompileExpr[
 Function[{Typed[x,PackedVector[Real]],Typed[𝒴,PackedVector[Real]]},
Module[{𝒩,𝒵,λlst,testdata,subseqL,KSλ},
𝒩=Length[data];𝒵=((data-𝒴)/(Sqrt[Abs[𝒴]]));
testdata=Function[{x},0.5(Erf[(x/(Sqrt[2]))]+1)]/@𝒵;
((Sqrt[𝒩])+.12+((.11)/(Sqrt[𝒩])))Max[Abs[((Range[𝒩])/𝒩)-Sort[testdata]]]
]
],
"OptimizationLevel"->None
];

RunPass[CopyPropagationPass,pm];
pm["toString"]

Since each pass does something small, they tend to be short (usually less than 100 lines of code).
We also favor large number of files over a few long files. There are over 700 files in the project.

Develop Abstractions

A variety of visitors for expressions, IR, types, etc...

Invest in Helper Functions

Asserters, data structures, loggers, timers, etc...

AssertThat["The condition should be either None (for unconditional branch) or a value", cond
		  ]["named", "condition"
		  ]["satisfiesAnyOf", {# === None&, VariableQ, ConstantValueQ}
];

Visualizations Make Life easier

pm=CompileExpr[
 Function[{Typed[x,Integer]},
Module[{y=0,ii=1},
While[ii<x,y=If[y>x,y*2,y-3]];y
]],
"OptimizationLevel"->None
];

pm["toString"]

pm["prettyPrint"]

Main:
({{start(53):}, {({{2, %1←Argument x}, {3, %2=Copy U(0)}, {4, %3=Copy U(1)}, {6, (Underscript[whileCondition(54), ↓])}})}})
({{whileCondition(54):}, {({{13, %7←φ[(Overscript[%2, start(53)]),(Overscript[%12, endif(59)])]}, {10, %5←φ[(Overscript[%1, start(53)]),(Overscript[%14, endif(59)])]}, {9, %4←φ[(Overscript[%3, start(53)]),(Overscript[%13, endif(59)])]}, {11, %6←%4 < %5}, {12, (Underscript[whileBody(55), ↙]) %6 (Underscript[whileExit(56), ↘])}})}})
({{whileBody(55):}, {({{14, %8←%7 > %5}, {18, (Underscript[thenif(57), ↙]) %8 (Underscript[elseif(58), ↘])}})}})
({{thenif(57):}, {({{19, %10←%7 * U(2)}, {20, (Underscript[endif(59), ↓])}})}})
({{elseif(58):}, {({{21, %11←%7 - U(3)}, {22, (Underscript[endif(59), ↓])}})}})
({{endif(59):}, {({{0, %14=Copy %5}, {0, %13=Copy %4}, {23, %9←φ[(Overscript[%10, thenif(57)]),(Overscript[%11, elseif(58)])]}, {24, %12=Copy %9}, {25, (Underscript[whileCondition(54), ↓])}})}})
({{whileExit(56):}, {({{28, Return %7}})}})

pm["getElement",1,3,2]

pm["toDot"]

Caused by the Compiler

Macro System

Language level transformations and simplifications. Can be endowed with knowledge about identities and invariants

RegisterBuiltinMacro[fuse,
fuse[Map[f_,rest_]]->pure[Map][fuse[f],fuse[rest]],
fuse[Composition[f_,Function[args_,body_]]]->Function[args,f[body]],
fuse[a_]->a
];
RegisterBuiltinMacro[pure,
pure[Map][f_,pure[Map][g_,rest_]]->pure[Map][fuse@Composition[f,g],pure[rest]]
];

MacroExpand[CreateMExpr[fuse@Map[f,Map[Function[{elem},elem+1],lst]]]]

Simple transformations can reduce compute time

n=10000000;
data={N@Range[n]/n,N@Range[n]/n};
Map[ArcTan,Transpose[data]];//AbsoluteTiming

MapThread[ArcTan,data];//AbsoluteTiming

Macro System

Can programmatically generate code that you’d otherwise not write:

RegisterBuiltinMacro[dot,
dot[{{a0_,b0_},{c0_,d0_}},{{a1_,b1_},{c1_,d1_}}]->{{a0 a1+b0 c1,a0 b1+b0 d1},{a1 c0+c1 d0,b1 c0+d0 d1}},
dot[{{a0_,b0_,c0_},{d0_,e0_,f0_},{g0_,h0_,i0_}},{{a1_,b1_,c1_},{d1_,e1_,f1_},{g1_,h1_,i1_}}]->{{a0 a1+b0 d1+c0 g1,a0 b1+b0 e1+c0 h1,a0 c1+b0 f1+c0 i1},{a1 d0+d1 e0+f0 g1,b1 d0+e0 e1+f0 h1,c1 d0+e0 f1+f0 i1},{a1 g0+d1 h0+g1 i0,b1 g0+e1 h0+h1 i0,c1 g0+f1 h0+i0 i1}}
];
RegisterBuiltinMacro[inverse,
inverse[{{a_,b_},{c_,d_}}]->{{d/(-b c+a d),-b/(-b c+a d)},{-c/(-b c+a d),a/(-b c+a d)}}
];

MacroExpand[CreateMExpr[dot[inverse[{{a,b},{c,d}}],{{x,y},{z,w}}]]]

Can delete constructs and work with a much lower level Mathematica code:

MacroExpand[CreateMExpr[Map[Function[{x},Sin[x]],data]]]//MExprToFormattedString

Linter

Verify that your passes do not break some known invariants

pm=CompileExpr[
 Function[{Typed[x,Integer]}->Integer[32],
Module[{y=0,ii=1},
While[ii<x,y=If[y>x,y*2,y-3]];y
]],
"OptimizationLevel"->1
];

RunPasses[{LiveVariablesWithBugPass,LintSSAIRPass},pm];

Compiling to Accelerators

E.g. NVIDIA GPUs, Xeon Phi, HSA, FPGA

Hardware

CUDA Compile Transformations

Optimizations

CUDA Compile Flow

CUDA Extensions to Compiler

Define a new behavior for Map in the Macro system

RegisterBuiltinMacro[Map,
Map[f_, lst_] ->
		If[Compile`Internal`MacroOptions["TargetArchitecture"] === "CUDA",
			Compile`Internal`CUDA`Map[f, lst],
			...
		]
]

Add new CUDA specific passes which are run when the TargetArchitecture == “CUDA”

CUDAMemorySetupPass

CUDAInstantiateMapPass

Add a new CUDABaseArray type and define functions that operate on this type

$BuiltinTypeEnvironment["functionLookup"]["addDefinitions", {
   Primitive`CUDA`Memory`Alloc :> Primitive`ExternalFunction["BitCodeCall", "cuda_memory_malloc", {Integer[64]} -> "CUDABaseArray"[a_]],
   Primitive`CUDA`Memory`Free :> Primitive`ExternalFunction["BitCodeCall", "cuda_memory_free", {"CUDABaseArray"[a_]} -> "RuntimeError"],
   Primitive`CUDA`Memory`CopyDeviceToHost :> Primitive`ExternalFunction["BitCodeCall", "cuda_memory_copy_device_to_host", {"BaseArray"[a_], "CUDABaseArray"[a_], Integer[64]} -> "RuntimeError"],
   Primitive`CUDA`Memory`CopyHostToDevice :> Primitive`ExternalFunction["BitCodeCall", "cuda_memory_copy_host_to_device", {"CUDABaseArray"[a_], "BaseArray"[a_], Integer[64]} -> "RuntimeError"]
   }];

Implement the CUDA functions in the runtime

extern "C" error cuda_memory_free(void* mem) {
	if (!mem) return error::memory;
	return cuda::check(cudaFree(mem));
}

Propagate the driver options to all parts of the compiler

pm=CompileExpr[Function[{Typed[lst,PackedVector[Integer[64]]]},Map[Function[{Typed[g,Integer[64]]},g*g-Typed[44,Integer[64]]],lst]],"TargetArchitecture"->"CUDA","OptimizationLevel"->1];
RunPasses[{CUDAMemorySetupPass,CUDAInstantiateMapPass,OptimizeBasicBlocksPass,CopyPropagationPass, ResolveFunctionCallPass},pm,<|"TargetArchitecture"->"CUDA"|>];
gpufms=Table[If[MemberQ[{"CUDADeviceOnly","CUDAGlobal"},fm["getProperty","Target"]],
fm,Nothing],{fm,pm["getFunctionModules"]}]

StringJoin[#["toString"]&/@gpufms]

GraphicsRow[#["toDot"]&/@gpufms]

Generated LLVM IR

define i64 @"Function$122835$CUDA$Device"(i64) {
start_1:
  %var4 = alloca i64, i64 1
  %1 = call i8 @binary_times_Integer64_Integer64(i64* %var4, i64 %0, i64 %0)
  %var7 = load i64, i64* %var4
  %var41 = alloca i64, i64 1
  %2 = call i8 @binary_subtract_Integer64_Integer64(i64* %var41, i64 %var7, i64 44)
  %var72 = load i64, i64* %var41
  ret i64 %var72
}

define void @"Function$122835$CUDA"(i64*, i64*, i64) {
start_1:
  %3 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
  %4 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
  %5 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x()
  %var4 = alloca i32, i64 1
  %6 = call i8 @binary_times_Integer32_Integer32(i32* %var4, i32 %5, i32 %4)
  %var7 = load i32, i32* %var4
  %var41 = alloca i32, i64 1
  %7 = call i8 @binary_plus_Integer32_Integer32(i32* %var41, i32 %var7, i32 %3)
  %var72 = load i32, i32* %var41
  %Cast = sext i32 %var72 to i64
  %cond = icmp slt i64 %Cast, %2
  br i1 %cond, label %then_2, label %end_3

then_2:                                           ; preds = %start_1
  %8 = getelementptr i64, i64* %1, i32 %var72
  %9 = load i64, i64* %8
  %funRes = call i64 @"Function$122835$CUDA$Device"(i64 %9)
  %10 = getelementptr i64, i64* %0, i32 %var72
  store i64 %funRes, i64* %10
  br label %end_3

end_3:                                            ; preds = %then_2, %start_1
  ret void
}

Equivalent CUDA code

__global__ void function_kernel(int64_t * out, int64_t * in, int64_t len) {
    int idx = threadIdx.x + blockIdx.x * gridDim.x;
    if (idx > len) return;
    out[idx] = in[idx]*in[idx] - 44;
}

error mapfunction(mtensor output, mtensor input) {
    int64_t * device_output, * device_input;
    void * host_input = mtensor_data(input);
    void * host_output = mtensor_data(input);
    size_t len = mtensor_length(input);
    int64_t size = mtensor_bytecount(input);

    cuda_memory_malloc((void**) &device_input, size);
    cuda_memory_malloc((void**) &device_output, size);
    cuda_memory_copy_host_to_device(device_input, host_input, size);

    const int blockDimX = 256;
    dim3 blockDim(blockDimX, 1, 1);
    dim3 gridDim((blockDimX + blockDimX - 1)/blockDimX, 1, 1);

    addTwo_kernel<<<gridDim, blockDim>>>(device_output, device_input, len);

    cuda_memory_copy_device_to_host(host_output, device_output, size);

    cuda_memory_free(device_input);
    cuda_memory_free(device_output);

    return len;
}

Byproducts

ClangLink

C code to Mathematica

<<ClangLink`
InstallClang[];

tr=GetTranslationUnitForFile["/Users/abduld/code/graphviz/lib/cgraph/cgraph.h"];

{Short[Types[tr]],Short[FunctionPrototypes[tr]]}

LLVMLink

Call the LLVM API from Mathematica

AutoLibraryLink

Generate LibraryLink projects from a C header file

Extended Runtime

Supports all machine types (Integer8, ..., Integer64, UnsignedInteger8, ..., Real32, Real64, ComplexReal32, ComplexReal64)

Questions

Thank You
Questions?