test using compiled function in math nodetemp-llvm-testing

4 files changed, 163 insertions, 87 deletions

diff --git a/source/blender/functions/FN_llvm.hh b/source/blender/functions/FN_llvm.hh
index 2d02af906b0..417719a21b5 100644
--- a/source/blender/functions/FN_llvm.hh
+++ b/source/blender/functions/FN_llvm.hh
@@ -18,6 +18,8 @@
 
 namespace blender::fn {
 
-void playground();
+using AddFuncType = void (*)(const float *, float, float *, int);
 
-}
+AddFuncType get_compiled_add_function();
+
+}  // namespace blender::fn
diff --git a/source/blender/functions/intern/llvm.cc b/source/blender/functions/intern/llvm.cc
index c52abf56228..324f9f3bad6 100644
--- a/source/blender/functions/intern/llvm.cc
+++ b/source/blender/functions/intern/llvm.cc
@@ -99,98 +99,113 @@ static llvm::Function &create_add_loop_function(llvm::Module &module)
   return function;
 }
 
-void playground()
-{
-  std::cout << "Start\n";
+class FunctionOptimizer {
+ private:
+  llvm::LLVMContext context_;
+  std::unique_ptr<llvm::ExecutionEngine> execution_engine_;
+
+ public:
+  AddFuncType generated_function;
 
-  static bool initialized = []() {
-    /* Set assembly syntax flavour. */
+ public:
+  void initialize()
+  {
     char const *args[] = {"some-random-name-for-the-parser", "--x86-asm-syntax=intel"};
     llvm::cl::ParseCommandLineOptions(std::size(args), args);
 
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmPrinter();
     llvm::InitializeNativeTargetAsmParser();
-    return true;
-  }();
-  UNUSED_VARS(initialized);
-
-  llvm::LLVMContext context;
-
-  std::unique_ptr<llvm::Module> module = std::make_unique<llvm::Module>("My Module", context);
-  llvm::Module *module_ptr = &*module;
-  std::string error;
-  const llvm::Target *target = llvm::TargetRegistry::lookupTarget(
-      llvm::sys::getDefaultTargetTriple(), error);
-  llvm::TargetOptions target_options;
-  llvm::StringMap<bool> cpu_feature_map;
-  llvm::sys::getHostCPUFeatures(cpu_feature_map);
-  llvm::SubtargetFeatures subtarget_features;
-  for (const auto &item : cpu_feature_map) {
-    subtarget_features.AddFeature(item.getKey(), item.getValue());
   }
-  llvm::TargetMachine *target_machine = target->createTargetMachine(
-      llvm::sys::getDefaultTargetTriple(),
-      llvm::sys::getHostCPUName().str(),
-      subtarget_features.getString(),
-      target_options,
-      {},
-      {},
-      llvm::CodeGenOpt::Aggressive);
-
-  module->setDataLayout(target_machine->createDataLayout());
-  module->setTargetTriple(target_machine->getTargetTriple().normalize());
-  std::cout << "Target Feature: " << target_machine->getTargetFeatureString().str() << "\n";
-  std::cout << "CPU: " << target_machine->getTargetCPU().str() << "\n";
-
-  std::cout << module->getTargetTriple() << "\n";
-  std::cout << module->getDataLayout().getStringRepresentation() << "\n";
-
-  llvm::Function &add_function = create_add_loop_function(*module);
-
-  BLI_assert(!llvm::verifyModule(*module, &llvm::outs()));
-  add_function.print(llvm::outs());
-
-  llvm::LoopAnalysisManager loop_analysis_manager;
-  llvm::FunctionAnalysisManager function_analysis_manager;
-  llvm::CGSCCAnalysisManager cgscc_anaylysis_manager;
-  llvm::ModuleAnalysisManager module_analysis_manager;
-
-  llvm::PassBuilder pass_builder{false, target_machine};
-
-  function_analysis_manager.registerPass([&] { return pass_builder.buildDefaultAAPipeline(); });
-
-  pass_builder.registerModuleAnalyses(module_analysis_manager);
-  pass_builder.registerCGSCCAnalyses(cgscc_anaylysis_manager);
-  pass_builder.registerFunctionAnalyses(function_analysis_manager);
-  pass_builder.registerLoopAnalyses(loop_analysis_manager);
-  pass_builder.crossRegisterProxies(loop_analysis_manager,
-                                    function_analysis_manager,
-                                    cgscc_anaylysis_manager,
-                                    module_analysis_manager);
-
-  llvm::ModulePassManager module_pass_manager = pass_builder.buildPerModuleDefaultPipeline(
-      llvm::PassBuilder::OptimizationLevel::O3);
-
-  module_pass_manager.run(*module, module_analysis_manager);
-
-  add_function.print(llvm::outs());
-
-  llvm::EngineBuilder engine_builder{std::move(module)};
-  engine_builder.setOptLevel(llvm::CodeGenOpt::Aggressive);
-  std::unique_ptr<llvm::ExecutionEngine> ee{engine_builder.create(target_machine)};
-  ee->finalizeObject();
-
-  const uint64_t function_ptr = ee->getFunctionAddress(add_function.getName().str());
-  using FuncType = void (*)(const float *, float, float *, int);
-  const FuncType generated_function = (FuncType)function_ptr;
-  UNUSED_VARS(generated_function, module_ptr);
-
-  LLVMTargetMachineEmitToFile((LLVMTargetMachineRef)ee->getTargetMachine(),
-                              llvm::wrap(module_ptr),
-                              (char *)"C:\\Users\\jacques\\Documents\\machine_code.txt",
-                              LLVMAssemblyFile,
-                              nullptr);
+
+  void generate_function()
+  {
+    std::unique_ptr<llvm::Module> module_owned = std::make_unique<llvm::Module>("My Module",
+                                                                                context_);
+    llvm::Module &ir_module = *module_owned;
+    std::string error;
+    const llvm::Target *target = llvm::TargetRegistry::lookupTarget(
+        llvm::sys::getDefaultTargetTriple(), error);
+    llvm::TargetOptions target_options;
+    llvm::StringMap<bool> cpu_feature_map;
+    llvm::sys::getHostCPUFeatures(cpu_feature_map);
+    llvm::SubtargetFeatures subtarget_features;
+    for (const auto &item : cpu_feature_map) {
+      subtarget_features.AddFeature(item.getKey(), item.getValue());
+    }
+    llvm::TargetMachine *target_machine = target->createTargetMachine(
+        llvm::sys::getDefaultTargetTriple(),
+        llvm::sys::getHostCPUName().str(),
+        subtarget_features.getString(),
+        target_options,
+        {},
+        {},
+        llvm::CodeGenOpt::Aggressive);
+
+    ir_module.setDataLayout(target_machine->createDataLayout());
+    ir_module.setTargetTriple(target_machine->getTargetTriple().normalize());
+
+    std::cout << ir_module.getTargetTriple() << "\n";
+    std::cout << ir_module.getDataLayout().getStringRepresentation() << "\n";
+
+    llvm::Function &add_function = create_add_loop_function(ir_module);
+
+    BLI_assert(!llvm::verifyModule(ir_module, &llvm::outs()));
+    add_function.print(llvm::outs());
+
+    llvm::LoopAnalysisManager loop_analysis_manager;
+    llvm::FunctionAnalysisManager function_analysis_manager;
+    llvm::CGSCCAnalysisManager cgscc_anaylysis_manager;
+    llvm::ModuleAnalysisManager module_analysis_manager;
+
+    llvm::PassBuilder pass_builder{false, target_machine};
+
+    function_analysis_manager.registerPass([&] { return pass_builder.buildDefaultAAPipeline(); });
+
+    pass_builder.registerModuleAnalyses(module_analysis_manager);
+    pass_builder.registerCGSCCAnalyses(cgscc_anaylysis_manager);
+    pass_builder.registerFunctionAnalyses(function_analysis_manager);
+    pass_builder.registerLoopAnalyses(loop_analysis_manager);
+    pass_builder.crossRegisterProxies(loop_analysis_manager,
+                                      function_analysis_manager,
+                                      cgscc_anaylysis_manager,
+                                      module_analysis_manager);
+
+    llvm::ModulePassManager module_pass_manager = pass_builder.buildPerModuleDefaultPipeline(
+        llvm::PassBuilder::OptimizationLevel::O3);
+
+    module_pass_manager.run(ir_module, module_analysis_manager);
+
+    add_function.print(llvm::outs());
+
+    llvm::EngineBuilder engine_builder{std::move(module_owned)};
+    engine_builder.setOptLevel(llvm::CodeGenOpt::Aggressive);
+    execution_engine_ = std::unique_ptr<llvm::ExecutionEngine>(
+        engine_builder.create(target_machine));
+    execution_engine_->finalizeObject();
+
+    const uint64_t function_ptr = execution_engine_->getFunctionAddress(
+        add_function.getName().str());
+
+    this->generated_function = (AddFuncType)function_ptr;
+
+    LLVMTargetMachineEmitToFile((LLVMTargetMachineRef)target_machine,
+                                llvm::wrap(&ir_module),
+                                (char *)"C:\\Users\\jacques\\Documents\\machine_code.txt",
+                                LLVMAssemblyFile,
+                                nullptr);
+  }
+};
+
+AddFuncType get_compiled_add_function()
+{
+  static FunctionOptimizer *optimizer = []() {
+    static FunctionOptimizer optimizer;
+    optimizer.initialize();
+    optimizer.generate_function();
+    return &optimizer;
+  }();
+  return optimizer->generated_function;
 }
 
 }  // namespace blender::fn
diff --git a/source/blender/modifiers/intern/MOD_nodes.cc b/source/blender/modifiers/intern/MOD_nodes.cc
index a52fc2703db..592bdc5e2ae 100644
--- a/source/blender/modifiers/intern/MOD_nodes.cc
+++ b/source/blender/modifiers/intern/MOD_nodes.cc
@@ -1109,7 +1109,6 @@ static void modifyGeometry(ModifierData *md,
                            const ModifierEvalContext *ctx,
                            GeometrySet &geometry_set)
 {
-  blender::fn::playground();
   NodesModifierData *nmd = reinterpret_cast<NodesModifierData *>(md);
   if (nmd->node_group == nullptr) {
     return;
diff --git a/source/blender/nodes/shader/nodes/node_shader_math.cc b/source/blender/nodes/shader/nodes/node_shader_math.cc
index e683d55edc8..7d7408d90f8 100644
--- a/source/blender/nodes/shader/nodes/node_shader_math.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_math.cc
@@ -26,6 +26,8 @@
 #include "NOD_math_functions.hh"
 #include "NOD_socket_search_link.hh"
 
+#include "FN_llvm.hh"
+
 /* **************** SCALAR MATH ******************** */
 
 namespace blender::nodes::node_shader_math_cc {
@@ -92,9 +94,67 @@ static int gpu_shader_math(GPUMaterial *mat,
   return 0;
 }
 
+class AddMultiFunction : public MultiFunction {
+  bool use_compiled_;
+
+ public:
+  AddMultiFunction(bool use_compiled) : use_compiled_(use_compiled)
+  {
+    static fn::MFSignature signature = create_signature();
+    this->set_signature(&signature);
+  }
+
+  static fn::MFSignature create_signature()
+  {
+    fn::MFSignatureBuilder signature{"Add"};
+    signature.single_input<float>("A");
+    signature.single_input<float>("B");
+    signature.single_output<float>("Result");
+    return signature.build();
+  }
+
+  void call(IndexMask mask, fn::MFParams params, fn::MFContext UNUSED(context)) const override
+  {
+    const VArray<float> a = params.readonly_single_input<float>(0, "A");
+    const VArray<float> b = params.readonly_single_input<float>(1, "B");
+    MutableSpan<float> result = params.uninitialized_single_output<float>(2, "Result");
+
+    if (a.is_span() && b.is_single() && mask.is_range() && mask.as_range().start() == 0) {
+      const float *a_data = a.get_internal_span().data();
+      float b_data = b.get_internal_single();
+      float *result_data = result.data();
+      int64_t size = mask.size();
+      if (use_compiled_) {
+        fn::AddFuncType func = fn::get_compiled_add_function();
+        func(a_data, b_data, result_data, size);
+      }
+      else {
+        for (const int64_t i : IndexRange(size)) {
+          result_data[i] = a_data[i] + b_data;
+        }
+      }
+    }
+    else {
+      for (const int i : mask) {
+        result[i] = a[i] + b[i];
+      }
+    }
+  }
+};
+
 static const blender::fn::MultiFunction *get_base_multi_function(bNode &node)
 {
   const int mode = node.custom1;
+
+  if (mode == NODE_MATH_ADD) {
+    static AddMultiFunction fn{false};
+    return &fn;
+  }
+  if (mode == NODE_MATH_MULTIPLY) {
+    static AddMultiFunction fn{true};
+    return &fn;
+  }
+
   const blender::fn::MultiFunction *base_fn = nullptr;
 
   blender::nodes::try_dispatch_float_math_fl_to_fl(