How can CompiledCodeFunctions be optimized and parallelized?

Question

Fast generation of random variates as test case

The earlier question (195435) about performance tuning with regard to fast generation of random variates for a gamma distribution has stirred my curiosity and below I am using slightly simpler code from the 3rd ed. of Numerical Recipes (Normaldev) to test different versions of a numerical routine to generate random variates for a normal distribution.

Different Implementations

The following implementations use the same core body, but differ with regard to the way a function is declared in Mathematica:

• normaldev is an uncompiled implementation
• normaldevC uses Compile with CompilationTarget → "C"
• normaldevCLP uses Compile as above, but adds listablility and parallelization
• normaldevLLVM finally uses the new compiler and FunctionCompile

Code (`body` has been escaped as\`body\` ):

{ normaldev, normaldevC, normaldevCLP, normaldevLLVM } = With[
    {
      funcBody = "
                While[ cond,
                    u = RandomReal[];
                    v = 1.7156 * (RandomReal[] - 0.5);
                    x = u -0.449871;
                    y = Abs[v] + 0.386595;
                    q = x*x + y *(0.19600*y - 0.25482 * x);
                    If[ q > 0.27597 && ( q > 0.27846 || v*v > -4. * Log[u] * u * u ), cond = True, cond = False ]
                ];
                mu + sig * v/u"
    },
    Module[
        { normaldev, normaldevC, normaldevCLP, normaldevLLVM }
        ,
        normaldev = StringTemplate["
        Function[ {mu, sig },
            Module[
                {
                    u = 1.0,
                    v = 1.0,
                    x = 1.0,
                    y = 1.0,
                    q = 1.0,
                    cond = True
                }
                ,
                \`body\`
            ]
        ]"
        ];
        normaldevC = StringTemplate["
        Compile[ { { mu, _Real }, { sig, _Real } },
            Module[
                {
                    u = 1.0,
                    v = 1.0,
                    x = 1.0,
                    y = 1.0,
                    q = 1.0,
                    cond = True
                }
                ,
                \`body\`
            ],
            CompilationTarget -> \"C\"
        ]"
        ];
        normaldevCLP = StringTemplate["
        Compile[ { { mu, _Real }, { sig, _Real } },
            Module[
                {
                    u = 1.0,
                    v = 1.0,
                    x = 1.0,
                    y = 1.0,
                    q = 1.0,
                    cond = True
                }
                ,
                \`body\`
            ],
            CompilationTarget -> \"C\",
            RuntimeAttributes -> {Listable},
            Parallelization -> True
        ]"
        ];
        normaldevLLVM = StringTemplate["
        FunctionCompile[
            Function[
                {
                    Typed[ mu,  \"Real64\" ],
                    Typed[ sig, \"Real64\" ]
                }
                ,
                Module[
                    {
                        Typed[ u,    \"Real64\"  ] = 1.0,
                        Typed[ v,    \"Real64\"  ] = 1.0,
                        Typed[ x,    \"Real64\"  ] = 1.0,
                        Typed[ y,    \"Real64\"  ] = 1.0,
                        Typed[ q,    \"Real64\"  ] = 1.0,
                        Typed[ cond, \"Boolean\" ] = True
                    }
                    ,
                    \`body\`
                ]
            ],
            CompilerOptions -> {
                \"AbortHandling\" -> False,
                \"LLVMOptimization\" -> \"ClangOptimization\"[3]
            }
        ]"
        ];
        Map[ ToExpression @ TemplateApply[ #, <| "body" -> funcBody |>]&, { normaldev, normaldevC, normaldevCLP, normaldevLLVM } ]
    ]
];

Performance evaluation

Using RepeatedTiming we can compare run times for these implementations, which I have sorted below from slowest to fastest.

RepeatedTiming @ Through[ {Mean, StandardDeviation}@ Table[ normaldev[0.,1.], {1000000}] ]
(* {0.535271,{-0.00048065,1.00115}} *)
RepeatedTiming @ Through[ {Mean, StandardDeviation}@ Table[ normaldevLLVM[0., 1.], {1000000}] ]
(* {0.199849,{0.00011759,1.00026}} *)
RepeatedTiming @ Through[ {Mean, StandardDeviation}@ Table[ normaldevC[0.,1.], {1000000}] ]
(* {0.18677,{-0.00244043,0.998699}} *)
With[ { args =  Transpose @ ConstantArray[{0.0,1.0},1000000]},
    RepeatedTiming @ Through[ {Mean, StandardDeviation}@ ( normaldevCLP @@ args )]
]
(* {0.119437,{0.000838859,0.999507}} *)

Questions

1. What would be the recommended implementation for the LLVM version if speed is the main goal, e.g., setting options, using Native`Random instead of RandomReal?
2. How can the benefits of listability and parallelization be obtained for Compiled CodeFunction?

Answer 1

In general, the weak spots of interpreted languages are loop constructs and function calls. So I'd suggest to push also the Table into the compiled code to generate many random numbers at once.

But the weakness of normaldevCLP lies at another point: You first generate a constant array of means and standard deviations. Then you send all this data to the library function. This is two real numbers for each real number that you want to generate. You can avoid this by using

normaldevCLP2 = StringTemplate["
        Compile[ { { mu, _Real }, { sig, _Real }, { n, _Integer } },
            Table[
                Module[
                    {
                        u = 1.0,
                        v = 1.0,
                        x = 1.0,
                        y = 1.0,
                        q = 1.0,
                        cond = True
                    },
                    `body`
                ]
            ,{i,1,n}],
            CompilationTarget -> \"C\",
            RuntimeAttributes -> {Listable},
            Parallelization -> True,
            RuntimeOptions -> \"Speed\"
        ]"];

On a 8-core machine for example, you would invoke the compiled function later like this:

n = 1000000;
\[Mu] = 0.;
\[Sigma] = 1.;
AdevCLP2 = normaldevCLP2[\[Mu], \[Sigma], ConstantArray[Quotient[n, 8], 8]];

On my machine, this yields a 2.5-fold speedup compared to

AdevCLP = normaldevCLP[ConstantArray[\[Mu], n], ConstantArray[\[Sigma], n]]

"What would be the recommended implementation for the LLVM version if speed is the main goal [...] ?"

I'd recommend not to use FunctionCompile at all. IMHO, it is completely ill-designed, in particular when you compare its clunky syntax with the ease of use of numba + numpy in Python. And I have not seen any example, yet, where FunctionCompile's performance was comparable to Compile LibraryLink. The only advantage of FunctionCompile seems to be that it allows for way more data types. And, btw., what the heck is the point of wrapping options of LLVM with an additional abstract layer like

 \"LLVMOptimization\" -> \"ClangOptimization\"[3]

?

This is just inflexible and unmaintainable. Why not simply using the approach of CreateLibrary by just sending the string(s) of compiler options directly to the compiler like, e.g.,

"CompileOptions" -> {
  , "-Xpreprocessor -fopenmp -fopenmp-simd"
  , "-Ofast"
  , "-flto"
  }

?

If you really look for performance, then better use the good ol' LibraryLink. It's syntax is awful, but one gets used to it. Here some example code to generate normally distributed numbers with the C++ standard library and OpenMP for parallelization:

Needs["CCompilerDriver`"];
Module[{lib, file, name},
name = "normalLibraryLink";

Print["Compiling " <> name <> "..."];

file = Export[FileNameJoin[{$TemporaryDirectory, name <> ".cpp"}],

"
#include"WolframLibrary.h"
#include <omp.h>
#include <random>
EXTERN_C DLLEXPORT int " <> name <> 
    "(WolframLibraryData libData, mint Argc, MArgument *Args, MArgument Res)
{
    // Grab the input arguments.
    mreal mu           = MArgument_getReal(Args[0]);
    mreal sigma        = MArgument_getReal(Args[1]);
    MTensor A          = MArgument_getMTensor(Args[2]);
    mint  thread_count = MArgument_getInteger(Args[3]);
// Using OpenMP for parallelization.
#pragma omp parallel num_threads( thread_count )
{
    // Slow, but truely random number generator for seeding.
    std::random_device r;

    // std::random_device generates 32-bit numbers, so we should use two of them to seed a 64-bit random number engine.
    std::seed_seq seed { r(), r() };

    // A fast pseudorandom number generator, seeded by a random number.
    // 64-bit Mersenne Twister by Matsumoto and Nishimura, 2000
    std::mt19937_64 random_engine ( seed );

    std::normal_distribution&lt;mreal&gt; dist (mu, sigma);

    // Gamma distribution is available by
    // std::gamma_distribution&lt;mreal&gt; dist (alpha, beta);
    // See https://cplusplus.com/reference/random/ for further random distributions.

    // Compute the data range [i_begin,i_end[ for this thread.
    const mint n       = libData-&gt;MTensor_getDimensions(A)[0];  
    const mint thread  = omp_get_thread_num();
    const mint i_begin = (n / thread_count) * (thread  ) + (n % \ thread_count * (thread  )) / thread_count;
    const mint i_end   = (n / thread_count) * (thread+1) + (n % \ thread_count * (thread+1)) / thread_count;

    // Get the pointer to the start of shared buffer supplied by the tensor argument.
    mreal * a = libData-&gt;MTensor_getRealData(A);

    // Fill the range [i_begin,i_end[ of a with random numbers drawn from the distribution dist.
    for( mint i = i_begin; i &lt; i_end; ++i )
    {
        a[i] = dist( random_engine );
    }
}

// Tell the library that it has to let go A.
libData-&gt;MTensor_disown(A);

return LIBRARY_NO_ERROR;

}",
   "Text"
    ];
lib = CreateLibrary[{file}, name
   , "TargetDirectory" -> $TemporaryDirectory
   , "ShellOutputFunction" -> Print
   , "CompileOptions" -> {
     "-Wall"
     , "-Wextra"
     , "-Wno-unused-parameter"
     , "-mmacosx-version-min=12.0"
     , "-std=c++11"
     , "-Xpreprocessor -fopenmp -fopenmp-simd"
     , "-Ofast"
     , "-flto"
     , "-mcpu=apple-m1"
     , "-mtune=native"
     }
   , "LinkerOptions" -> {"-lm", "-ldl", "-lomp"}
   , "IncludeDirectories" -> {"/opt/local/include/libomp"}
   , "LibraryDirectories" -> {"/opt/local/lib/libomp"}
   (*,"ShellCommandFunction"[Rule]Print*)
   , "ShellOutputFunction" -> Print
    ];
    Print["Compilation done."];
normalLibraryLink = LibraryFunctionLoad[lib, name, 
  {
    Real,                 (*mu*)
    Real,                 (*sigma*)
    {Real, 1, &quot;Shared&quot;},  (*vector A; passed by reference. Caution: A gets modified!*)
    Integer               (*thread_count*)
  },
  &quot;Void&quot;                  (*no return value*)
]

]

You can run it by first letting Mathematica allocate some array

a = ConstantArray[0., n];

Then you invoke normalLibraryLink to fill it with random numbers, using 8 threads:

normalLibraryLink[0., 1., a, 8]; // RepeatedTiming

0.00177217

The advantage is that the array a is passed as "Shared", which means "by reference". So the library can modify the array directly without any allocations or copy operations. You can fill the array a; do something with the random numbers; and then refill the same array again. On my machine this is six times faster than AdevCLP2 and 130 times faster than AdevLLVM.

The compile options are selected for macos on an Apple M1 processor and with OpenMP installed via Macports. Of course, you have to adapt them appropriately for other platforms.