[MLIR][OpenMP] Conversion pattern to convert OpenMP ops to OpenMP ops with LLVM dialect

[kiranchandramohan]

Summary

Adding a conversion pattern for the parallel Operation. This will help the conversion of parallel operation with standard dialect to parallel operation with llvm dialect. The type conversion of the block arguments in a parallel region is controlled by the pattern for the parallel Operation. Without this pattern, a parallel Operation with block arguments cannot be converted from standard to LLVM dialect. Other OpenMP operations without regions are marked as legal. When the translation of new OpenMP operations with regions are added then patterns for these operations can also be added.
This conversion uses all the standard to llvm patterns. Patterns of other dialects can be added later if needed.

See discussion in https://llvm.discourse.group/t/about-using-omp-dialect-in-mlir/1518/9

[rriddle]

mlir/include/mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h

11

Are these includes necessary? I would expect that forward declarations would be enough.

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

16

Can you trim these includes?

25

Classes should be placed in an anonymous namespace.

36

Did you intend to convert the result types here? Did you intend to pass in the converted operands, in operands as well?

38

The last .getOperation() shouldn't be necessary.

41

nit: Drop trivial braces.

mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir

3

This looks like it is testing something already tested by the StandardToLLVM conversion pass. Can you trim this test down to checking what the patterns that you have added?

[kiranchandramohan]

Addressed comments from @rriddle.

mlir/include/mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h

11

Removed these but had to add <memory> for the unique_ptr

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

36

The result types are actually empty for the Parallel operation. So I have changed it here.

I saw that the operands of the parallel region are already converted so i did not want to do that in the conversion pattern also. This conversion pattern is only for converting the block arguments inside the parallel region to LLVM.
My understanding here might not be fully correct. If so please let me know if the operands of the parallel region also should be converted.

mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir

3

I only have one function in the test now.
The conversion pattern is necessary for changing block arguments to llvm. So the test checks the block arguments and its uses inside the openmp region.

[kiranchandramohan]

Resubmitting with full context.

[rriddle]

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

26

nit: I'd just remove this, there isn't much of a need to protect it here.

36

In general, you should always use the updated operands unless you explicitly don't want to. When type conversions are involved the conversion infra won't implicitly replace operands with new values of a different type, so you should make sure to always use the updated operands in your patterns.

mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir

4

nit: Can you add a CHECK-LABEL here?

[kiranchandramohan]

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

36

Thanks @rriddle. Is there a single function which can convert all the operands (from standard to llvm) or do i have to individually for each of the operands. Here some of the operands are attributes. What would be a good reference code?

[rriddle]

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

36

Is there a single function which can convert all the operands (from standard to llvm) or do i have to individually for each of the operands.

The converted operands are already being passed to the pattern via the operands parameters. This parameter contains the values that replaced the operands of the original operation, with the order corresponding 1-1 with the order of the original operand list. You can use the omp::ParallelOp::Adaptor class to provide named accessors(just like the proper omp::ParallelOp class) into this remapped operand list.

Here some of the operands are attributes.

nit on terminology, operands are never attributes. An argument, i.e. the thing you specify in ODS/Tablegen, may be either an Attribute or an Operand but these two are completely different concepts. An Operand is an SSA value, where as an Attribute is constant metadata.

What would be a good reference code?

Any of the other Foo(e.g. Standard)->LLVM conversion patterns are a good reference for LLVM conversions. For patterns in general, a lot of things in Conversion/ can be useful.

[kiranchandramohan]

Addressed comments from @rriddle.

[kiranchandramohan]

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

36

Thanks for the explanation. I am using the converted operands passed to the pattern to create the new Operation as follows.

auto newOp = rewriter.create<omp::ParallelOp>(
    curOp.getLoc(), ArrayRef<Type>(), operands, curOp.getAttrs());

Thanks for the correction about the attributes and operands!

mlir/include/mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h

+//===- OpenMPToLLVM.h - Utils to convert from the OpenMP dialect ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_CONVERSION_OPENMPTOLLVM_OPENMPTOLLVM_H_
+#define MLIR_CONVERSION_OPENMPTOLLVM_OPENMPTOLLVM_H_
+
+#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+class MLIRContext;
+class ModuleOp;
+template <typename T>
+class OperationPass;
+
+/// Populate the given list with patterns that convert from OpenMP to LLVM.
+void populateOpenMPToLLVMConversionPatterns(MLIRContext *context,
+                                            LLVMTypeConverter &converter,
+                                            OwningRewritePatternList &patterns);
+
+/// Create a pass to convert OpenMP operations to the LLVMIR dialect.
+std::unique_ptr<OperationPass<ModuleOp>> createConvertOpenMPToLLVMPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_OPENMPTOLLVM_OPENMPTOLLVM_H_

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 #include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h"
 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
+#include "mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h"
 #include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
 #include "mlir/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVMPass.h"
 #include "mlir/Conversion/ShapeToSCF/ShapeToSCF.h"
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"

mlir/include/mlir/Conversion/Passes.td

 
 def ConvertLinalgToSPIRV : Pass<"convert-linalg-to-spirv", "ModuleOp"> {
   let summary = "Convert Linalg ops to SPIR-V ops";
   let constructor = "mlir::createLinalgToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];
 }
 
 //===----------------------------------------------------------------------===//
+// OpenMPToLLVM
+//===----------------------------------------------------------------------===//
+
+def ConvertOpenMPToLLVM : Pass<"convert-openmp-to-llvm", "ModuleOp"> {
+  let summary = "Convert the OpenMP ops to OpenMP ops with LLVM dialect";
+  let constructor = "mlir::createConvertOpenMPToLLVMPass()";
+}
+
+//===----------------------------------------------------------------------===//
 // SCFToStandard
 //===----------------------------------------------------------------------===//
 
 def SCFToStandard : Pass<"convert-scf-to-std"> {
   let summary = "Convert SCF dialect to Standard dialect, replacing structured"
                 " control flow with a CFG";
   let constructor = "mlir::createLowerToCFGPass()";
   let dependentDialects = ["StandardOpsDialect"];

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(AVX512ToLLVM)
 add_subdirectory(GPUCommon)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)
 add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LinalgToStandard)
+add_subdirectory(OpenMPToLLVM)
 add_subdirectory(SCFToGPU)
 add_subdirectory(SCFToSPIRV)
 add_subdirectory(SCFToStandard)
 add_subdirectory(ShapeToSCF)
 add_subdirectory(ShapeToStandard)
 add_subdirectory(SPIRVToLLVM)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(VectorToROCDL)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)

mlir/lib/Conversion/OpenMPToLLVM/CMakeLists.txt

+add_mlir_conversion_library(MLIROpenMPToLLVM
+  OpenMPToLLVM.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/OpenMPToLLVM
+
+  DEPENDS
+  MLIRConversionPassIncGen
+  intrinsics_gen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRLLVMIR
+  MLIROpenMP
+  MLIRStandardToLLVM
+  MLIRTransforms
+  )

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

+//===- OpenMPToLLVM.cpp - conversion from OpenMP to LLVM dialect ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h"
+
+#include "../PassDetail.h"
+#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
+#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Module.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir;
+
+struct ParallelOpConversion : public ConvertToLLVMPattern {
+  explicit ParallelOpConversion(MLIRContext *context,
+                                LLVMTypeConverter &typeConverter)
+      : ConvertToLLVMPattern(omp::ParallelOp::getOperationName(), context,
+                             typeConverter) {}
+
+protected:
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto curOp = cast<omp::ParallelOp>(op);
+    auto newOp = rewriter.create<omp::ParallelOp>(
+        curOp.getLoc(), curOp.getOperation()->getResultTypes(),
+        curOp.getOperation()->getOperands(), curOp.getOperation()->getAttrs());
+    rewriter.inlineRegionBefore(curOp.region(), newOp.region(),
+                                newOp.region().end());
+    if (failed(rewriter.convertRegionTypes(&newOp.region(), typeConverter))) {
+      return failure();
+    }
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+void mlir::populateOpenMPToLLVMConversionPatterns(
+    MLIRContext *context, LLVMTypeConverter &converter,
+    OwningRewritePatternList &patterns) {
+  patterns.insert<ParallelOpConversion>(context, converter);
+}
+
+namespace {
+struct ConvertOpenMPToLLVMPass
+    : public ConvertOpenMPToLLVMBase<ConvertOpenMPToLLVMPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void ConvertOpenMPToLLVMPass::runOnOperation() {
+  auto module = getOperation();
+  MLIRContext *context = &getContext();
+
+  // Convert to OpenMP operations with LLVM IR dialect
+  OwningRewritePatternList patterns;
+  LLVMTypeConverter converter(&getContext());
+  populateStdToLLVMConversionPatterns(converter, patterns);
+  populateOpenMPToLLVMConversionPatterns(context, converter, patterns);
+
+  LLVMConversionTarget target(getContext());
+  target.addDynamicallyLegalOp<omp::ParallelOp>(
+      [&](omp::ParallelOp op) { return converter.isLegal(&op.getRegion()); });
+  target.addLegalOp<omp::TerminatorOp, omp::TaskyieldOp, omp::FlushOp,
+                    omp::BarrierOp, omp::TaskwaitOp>();
+  if (failed(applyPartialConversion(module, target, patterns)))
+    signalPassFailure();
+}
+
+std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertOpenMPToLLVMPass() {
+  return std::make_unique<ConvertOpenMPToLLVMPass>();
+}

mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir

+// RUN: mlir-opt -convert-openmp-to-llvm %s -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @empty() {
+// CHECK-NEXT:  llvm.return
+// CHECK-NEXT: }
+func @empty() {
+^bb0:
+  return
+}
+
+// CHECK-LABEL:  llvm.func @branch_loop() {
+func @branch_loop() {
+// CHECK-NEXT:   %[[VAL0:[0-9]+]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  %start = constant 0 : index
+// CHECK-NEXT:   %[[VAL1:[0-9]+]] = llvm.mlir.constant(1 : index) : !llvm.i64
+  %step = constant 1 : index
+// CHECK-NEXT:   %[[VAL2:[0-9]+]] = llvm.mlir.constant(100 : index) : !llvm.i64
+  %end = constant 100 : index
+// CHECK-NEXT:   %[[VAL3:[0-9]+]] = llvm.mlir.constant(4 : index) : !llvm.i64
+  %nthreads = constant 4 : index
+// CHECK-NEXT:   omp.parallel num_threads(%[[VAL3]] : !llvm.i64) private(%[[VAL0]] : !llvm.i64, %[[VAL2]] : !llvm.i64) proc_bind(spread) {
+  omp.parallel num_threads(%nthreads : index) private(%start : index, %end : index) proc_bind(spread) {
+// CHECK-NEXT:     llvm.br ^[[BB1:.*]](%[[VAL0]] : !llvm.i64)
+    br ^bb1(%start : index)
+// CHECK-NEXT:   ^[[BB1]](%[[VAL4:[0-9]+]]: !llvm.i64):
+  ^bb1(%0: index):
+// CHECK-NEXT:     %[[VAL5:[0-9]+]] = llvm.icmp "slt" %[[VAL4]], %[[VAL2]] : !llvm.i64
+    %1 = cmpi "slt", %0, %end : index
+// CHECK-NEXT:     llvm.cond_br %[[VAL5]], ^[[BB2:.*]], ^[[BB3:.*]]
+    cond_br %1, ^bb2, ^bb3
+// CHECK-NEXT:   ^[[BB2]]:
+  ^bb2:
+// CHECK-NEXT:     %[[VAL6:[0-9]+]] = llvm.trunc %[[VAL4]] : !llvm.i64 to !llvm.i32
+    %2 = index_cast %0 : index to i32
+// CHECK-NEXT:     llvm.call @print_i32(%[[VAL6]]) : (!llvm.i32) -> ()
+    call @print_i32(%2) : (i32) -> ()
+// CHECK-NEXT:     %[[VAL7:[0-9]+]] = llvm.add %[[VAL4]], %[[VAL1]] : !llvm.i64
+    %3 = addi %0, %step : index
+// CHECK-NEXT:     llvm.br ^bb1(%[[VAL7]] : !llvm.i64)
+    br ^bb1(%3 : index)
+// CHECK-NEXT:   ^[[BB3]]:
+  ^bb3:
+// CHECK-NEXT:     omp.terminator
+    omp.terminator
+// CHECK-NEXT: }
+  }
+// CHECK-NEXT:   llvm.return
+  return
+// CHECK-NEXT: }
+}
+// CHECK-NEXT: llvm.func @print_i32(!llvm.i32)
+func @print_i32(i32)
+
+// CHECK: llvm.func @test_constructs(%[[FN_ARG0:.*]]: !llvm.i32, %[[FN_ARG1:.*]]: !llvm.i32)
+func @test_constructs(%arg0: i32, %arg1: i32) {
+// CHECK-NEXT:    omp.parallel shared(%[[FN_ARG0]] : !llvm.i32, %[[FN_ARG1]] : !llvm.i32) {
+  omp.parallel shared(%arg0: i32, %arg1: i32) {
+// CHECK-NEXT:      omp.flush
+    omp.flush
+// CHECK-NEXT:      omp.barrier
+    omp.barrier
+// CHECK-NEXT:      llvm.br ^[[BB1:.*]](%[[FN_ARG0]], %[[FN_ARG1]] : !llvm.i32, !llvm.i32)
+    br ^bb1(%arg0, %arg1 : i32, i32)
+// CHECK-NEXT:    ^[[BB1]](%[[BB_ARG0:.*]]: !llvm.i32, %[[BB_ARG1:.*]]: !llvm.i32):  // pred: ^bb0
+    ^bb1(%0: i32, %1: i32):
+// CHECK-NEXT:      omp.flush %[[BB_ARG0]], %[[BB_ARG1]] : !llvm.i32, !llvm.i32
+    omp.flush %0, %1 : i32, i32
+// CHECK-NEXT:      omp.taskyield
+    omp.taskyield
+// CHECK-NEXT:      omp.taskwait
+    omp.taskwait
+// CHECK-NEXT:      omp.terminator
+    omp.terminator
+// CHECK-NEXT:    }
+  }
+// CHECK-NEXT:    llvm.return
+  return
+// CHECK-NEXT:  }
+}