feat: lower to omp dialect

Author: batzor

benchmark/benchmark.bzl

@@ -146,7 +146,10 @@ def zkir_benchmark_test(name, mlir_src, test_src, zkir_opt_flags = [], data = []
             ":" + import_name,
             "@google_benchmark//:benchmark_main",
             "@googletest//:gtest",
+            "@local_config_omp//:omp",
         ],
+        copts = ["-Xclang -fopenmp"],
+        linkopts = ["-Xclang -fopenmp"],
         tags = tags,
         data = data + [generated_obj_name],
         **kwargs

benchmark/ntt/BUILD.bazel

@@ -5,7 +5,7 @@ zkir_benchmark_test(
     mlir_src = "ntt_benchmark.mlir",
     tags = ["manual"],
     test_src = ["ntt_benchmark_test.cc"],
-    zkir_opt_flags = ["-poly-to-llvm"],
+    zkir_opt_flags = ["-poly-to-omp"],
     deps = [
         "//benchmark:BenchmarkUtils",
     ],

benchmark/ntt/ntt_benchmark_test.cc

@@ -140,12 +140,12 @@ BENCHMARK(BM_intt_mont_benchmark)->Iterations(1)->Unit(::benchmark::kSecond);
 //   L1 Data 64 KiB
 //   L1 Instruction 128 KiB
 //   L2 Unified 4096 KiB (x14)
-// Load Average: 3.97, 3.11, 2.96
+// Load Average: 1.82, 2.22, 2.39
 // ------------------------------------------------------------------------------
 // Benchmark                                    Time             CPU   Iterations
 // ------------------------------------------------------------------------------
-// BM_ntt_benchmark                          10.2 s          10.1 s             1
-// BM_intt_benchmark/iterations:1            11.1 s          11.1 s             1
-// BM_ntt_mont_benchmark                    0.190 s         0.190 s             3
-// BM_intt_mont_benchmark/iterations:1      0.316 s         0.304 s             1
+// BM_ntt_benchmark                          10.1 s          10.1 s             1
+// BM_intt_benchmark/iterations:1            10.1 s          10.0 s             1
+// BM_ntt_mont_benchmark                    0.183 s         0.183 s             4
+// BM_intt_mont_benchmark/iterations:1      0.266 s         0.214 s             1
 // NOLINTEND()

tools/zkir-opt.cpp

@@ -32,7 +32,11 @@ int main(int argc, char **argv) {
 
   mlir::PassPipelineRegistration<>(
       "poly-to-llvm", "Run passes to lower the polynomial dialect to LLVM",
-      mlir::zkir::pipelines::polyToLLVMPipelineBuilder);
+      mlir::zkir::pipelines::polyToLLVMPipelineBuilder<false>);
+  mlir::PassPipelineRegistration<>(
+      "poly-to-omp",
+      "Run passes to lower the polynomial dialect to OpenMP + LLVM",
+      mlir::zkir::pipelines::polyToLLVMPipelineBuilder<true>);
 
   return failed(mlir::MlirOptMain(argc, argv, "ZKIR optimizer\n", registry));
 }

zkir/Pipelines/BUILD.bazel

@@ -21,8 +21,10 @@ cc_library(
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:LinalgTransforms",
         "@llvm-project//mlir:MemRefTransforms",
+        "@llvm-project//mlir:OpenMPToLLVM",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:SCFToControlFlow",
+        "@llvm-project//mlir:SCFToOpenMP",
         "@llvm-project//mlir:Transforms",
     ],
 )

zkir/Pipelines/PipelineRegistration.cpp

@@ -3,7 +3,9 @@
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h"
 #include "mlir/Conversion/ConvertToLLVM/ToLLVMPass.h"
+#include "mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h"
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
+#include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
@@ -41,6 +43,7 @@ void oneShotBufferize(OpPassManager &manager) {
   manager.addPass(createCanonicalizerPass());
 }
 
+template <bool allowOpenMP>
 void polyToLLVMPipelineBuilder(OpPassManager &manager) {
   manager.addPass(zkir::poly::createPolyToField());
   manager.addPass(zkir::field::createPrimeFieldToModArith());
@@ -49,9 +52,11 @@ void polyToLLVMPipelineBuilder(OpPassManager &manager) {
 
   // Linalg
   manager.addNestedPass<FuncOp>(createConvertElementwiseToLinalgPass());
+  manager.addNestedPass<FuncOp>(createLinalgElementwiseOpFusionPass());
   // Needed to lower affine.map and affine.apply
   manager.addNestedPass<FuncOp>(affine::createAffineExpandIndexOpsPass());
   manager.addNestedPass<FuncOp>(affine::createSimplifyAffineStructuresPass());
+  manager.addPass(affine::createAffineParallelize());
   manager.addPass(createLowerAffinePass());
   manager.addNestedPass<FuncOp>(memref::createExpandOpsPass());
   manager.addNestedPass<FuncOp>(memref::createExpandStridedMetadataPass());
@@ -61,7 +66,7 @@ void polyToLLVMPipelineBuilder(OpPassManager &manager) {
 
   // Linalg must be bufferized before it can be lowered
   // But lowering to loops also re-introduces affine.apply, so re-lower that
-  manager.addNestedPass<FuncOp>(createConvertLinalgToLoopsPass());
+  manager.addNestedPass<FuncOp>(createConvertLinalgToParallelLoopsPass());
   manager.addPass(createLowerAffinePass());
   manager.addPass(createBufferizationToMemRefPass());
 
@@ -73,13 +78,17 @@ void polyToLLVMPipelineBuilder(OpPassManager &manager) {
 
   // ToLLVM
   manager.addPass(arith::createArithExpandOpsPass());
+  if constexpr (allowOpenMP) {
+    manager.addPass(createConvertSCFToOpenMPPass());
+  }
   manager.addPass(createConvertSCFToCFPass());
   manager.addNestedPass<FuncOp>(memref::createExpandStridedMetadataPass());
 
   // expand strided metadata will create affine map. Needed to lower affine.map
   // and affine.apply
   manager.addNestedPass<FuncOp>(affine::createAffineExpandIndexOpsPass());
   manager.addNestedPass<FuncOp>(affine::createSimplifyAffineStructuresPass());
+  manager.addPass(affine::createAffineParallelize());
   manager.addPass(createLowerAffinePass());
   manager.addPass(createConvertToLLVMPass());
 
@@ -90,4 +99,7 @@ void polyToLLVMPipelineBuilder(OpPassManager &manager) {
   manager.addPass(createSymbolDCEPass());
 }
 
+template void polyToLLVMPipelineBuilder<false>(mlir::OpPassManager &manager);
+template void polyToLLVMPipelineBuilder<true>(mlir::OpPassManager &manager);
+
 }  // namespace mlir::zkir::pipelines

zkir/Pipelines/PipelineRegistration.h

@@ -9,8 +9,14 @@ namespace mlir::zkir::pipelines {
 
 void oneShotBufferize(OpPassManager &manager);
 
+template <bool allowOpenMP>
 void polyToLLVMPipelineBuilder(OpPassManager &manager);
 
+extern template void polyToLLVMPipelineBuilder<false>(
+    mlir::OpPassManager &manager);
+extern template void polyToLLVMPipelineBuilder<true>(
+    mlir::OpPassManager &manager);
+
 }  // namespace mlir::zkir::pipelines
 
 #endif  // ZKIR_PIPELINES_PIPELINEREGISTRATION_H_