From 187d9f8cd9d32ca0021619153dcb3ffb0b78b987 Mon Sep 17 00:00:00 2001
From: Kiran Chandramohan <kiran.chandramohan@arm.com>
Date: Sun, 12 Sep 2021 10:13:14 +0000
Subject: [PATCH] [OpenMP][MLIR] Add a conversion pattern for the master op

The conversion pattern is particularly useful for conversion of
block arguments in the master op.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D109610
---
 mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp        |  5 +++--
 mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir | 13 +++++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
index bf725c2..e0b3ed8 100644
--- a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
+++ b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
@@ -47,7 +47,8 @@ struct RegionOpConversion : public ConvertOpToLLVMPattern<OpType> {
 
 void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                   RewritePatternSet &patterns) {
-  patterns.add<RegionOpConversion<omp::ParallelOp>,
+  patterns.add<RegionOpConversion<omp::MasterOp>,
+               RegionOpConversion<omp::ParallelOp>,
                RegionOpConversion<omp::WsLoopOp>>(converter);
 }
 
@@ -69,7 +70,7 @@ void ConvertOpenMPToLLVMPass::runOnOperation() {
   populateOpenMPToLLVMConversionPatterns(converter, patterns);
 
   LLVMConversionTarget target(getContext());
-  target.addDynamicallyLegalOp<omp::ParallelOp, omp::WsLoopOp>(
+  target.addDynamicallyLegalOp<omp::MasterOp, omp::ParallelOp, omp::WsLoopOp>(
       [&](Operation *op) { return converter.isLegal(&op->getRegion(0)); });
   target.addLegalOp<omp::TerminatorOp, omp::TaskyieldOp, omp::FlushOp,
                     omp::BarrierOp, omp::TaskwaitOp>();
diff --git a/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir b/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir
index e9a46cb..f572935 100644
--- a/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir
+++ b/mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir
@@ -1,5 +1,18 @@
 // RUN: mlir-opt -convert-openmp-to-llvm %s  -split-input-file | FileCheck %s
 
+// CHECK-LABEL: llvm.func @master_block_arg
+func @master_block_arg() {
+  // CHECK: omp.master
+  omp.master {
+  // CHECK-NEXT: ^[[BB0:.*]](%[[ARG1:.*]]: i64, %[[ARG2:.*]]: i64):
+  ^bb0(%arg1: index, %arg2: index):
+    // CHECK-NEXT: "test.payload"(%[[ARG1]], %[[ARG2]]) : (i64, i64) -> ()
+    "test.payload"(%arg1, %arg2) : (index, index) -> ()
+    omp.terminator
+  }
+  return
+}
+
 // CHECK-LABEL: llvm.func @branch_loop
 func @branch_loop() {
   %start = constant 0 : index
-- 
2.7.4