[OpenMP,MLIR] Translation of parallel operation: num_threads, if clauses 3/n

[kiranchandramohan]

Summary

This simple patch translates the num_threads and if clauses of the parallel operation. Also includes test cases.

A minor change was made to parsing of the if clause to parse AnyType and return the parsed type. Updates to test cases also.

[jdoerfert]

No complaints form my side :)

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

463–467

Unrelated but FWIW, true is the conservative correct answer here ;)

[kiranktp]

LGTM

[kiranktp]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

212

NIT: Modify default clause to proc_bind in the comment

[SouraVX]

LGTM. Thanks!

[kiranchandramohan]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

212

Thanks, Will change in the final commit.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

463–467

:) Thanks. Is it OK if this is changed when we add support for cancellation?

[anchu-rajendran]

Looks good to me!

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

parseOperandAndTypeList(OpAsmParser &parser,
 
   return success();
 }
 
 static void printParallelOp(OpAsmPrinter &p, ParallelOp op) {
   p << "omp.parallel";
 
   if (auto ifCond = op.if_expr_var())
-    p << " if(" << ifCond << ")";
+    p << " if(" << ifCond << " : " << ifCond.getType() << ")";
 
   if (auto threads = op.num_threads_var())
     p << " num_threads(" << threads << " : " << threads.getType() << ")";
 
   // Print private, firstprivate, shared and copyin parameters
   auto printDataVars = [&p](StringRef name, OperandRange vars) {
     if (vars.size()) {
       p << " " << name << "(";

 /// shared ::= `shared` operand-and-type-list
 /// copyin ::= `copyin` operand-and-type-list
 /// default ::= `default` `(` (`private` | `firstprivate` | `shared` | `none`)
 /// procBind ::= `proc_bind` `(` (`master` | `close` | `spread`) `)`
 ///
 /// Note that each clause can only appear once in the clase-list.
 static ParseResult parseParallelOp(OpAsmParser &parser,
                                    OperationState &result) {
-  OpAsmParser::OperandType ifCond;
+  std::pair<OpAsmParser::OperandType, Type> ifCond;
   std::pair<OpAsmParser::OperandType, Type> numThreads;
   llvm::SmallVector<OpAsmParser::OperandType, 4> privates;
   llvm::SmallVector<Type, 4> privateTypes;
   llvm::SmallVector<OpAsmParser::OperandType, 4> firstprivates;
   llvm::SmallVector<Type, 4> firstprivateTypes;
   llvm::SmallVector<OpAsmParser::OperandType, 4> shareds;
   llvm::SmallVector<Type, 4> sharedTypes;
   llvm::SmallVector<OpAsmParser::OperandType, 4> copyins;
static ParseResult parseParallelOp(OpAsmParser &parser,
   const int copyinClausePos = 5;
   const llvm::StringRef opName = result.name.getStringRef();
 
   while (succeeded(parser.parseOptionalKeyword(&keyword))) {
     if (keyword == "if") {
       // Fail if there was already another if condition
       if (segments[ifClausePos])
         return allowedOnce(parser, "if", opName);
-      if (parser.parseLParen() || parser.parseOperand(ifCond) ||
-          parser.parseRParen())
+      if (parser.parseLParen() || parser.parseOperand(ifCond.first) ||
+          parser.parseColonType(ifCond.second) || parser.parseRParen())
         return failure();
       segments[ifClausePos] = 1;
     } else if (keyword == "num_threads") {
       // fail if there was already another num_threads clause
       if (segments[numThreadsClausePos])
         return allowedOnce(parser, "num_threads", opName);
       if (parser.parseLParen() || parser.parseOperand(numThreads.first) ||
           parser.parseColonType(numThreads.second) || parser.parseRParen())
if (keyword == "if") {
       llvm::SmallString<16> attrval;
       // The def prefix is required for the attribute as "private" is a keyword
       // in C++
       attrval += "def";
       attrval += defval;
       auto attr = parser.getBuilder().getStringAttr(attrval);
       result.addAttribute("default_val", attr);
     } else if (keyword == "proc_bind") {
-      // fail if there was already another default clause
+      // fail if there was already another proc_bind clause
       if (procBind)
         return allowedOnce(parser, "proc_bind", opName);
       procBind = true;
       llvm::StringRef bind;
       if (parser.parseLParen() || parser.parseKeyword(&bind) ||
           parser.parseRParen())
         return failure();
       auto attr = parser.getBuilder().getStringAttr(bind);
       result.addAttribute("proc_bind_val", attr);
     } else {
       return parser.emitError(parser.getNameLoc())
              << keyword << " is not a valid clause for the " << opName
              << " operation";
     }
   }
 
   // Add if parameter
   if (segments[ifClausePos]) {
-    parser.resolveOperand(ifCond, parser.getBuilder().getI1Type(),
-                          result.operands);
+    parser.resolveOperand(ifCond.first, ifCond.second, result.operands);
   }
 
   // Add num_threads parameter
   if (segments[numThreadsClausePos]) {
     parser.resolveOperand(numThreads.first, numThreads.second, result.operands);
   }
 
   // Add private parameters

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

ModuleTranslation::convertOmpParallel(Operation &opInst,
 
   // TODO: Perform finalization actions for variables. This has to be
   // called for variables which have destructors/finalizers.
   auto finiCB = [&](InsertPointTy codeGenIP) {};
 
   // TODO: The various operands of parallel operation are not handled.
   // Parallel operation is created with some default options for now.
   llvm::Value *ifCond = nullptr;
+  if (auto ifExprVar = cast<omp::ParallelOp>(opInst).if_expr_var())
+    ifCond = valueMapping.lookup(ifExprVar);
   llvm::Value *numThreads = nullptr;
+  if (auto numThreadsVar = cast<omp::ParallelOp>(opInst).num_threads_var())
+    numThreads = valueMapping.lookup(numThreadsVar);
+  // TODO: Is the Parallel construct cancellable?
   bool isCancellable = false;
   // TODO: Determine the actual alloca insertion point, e.g., the function
   // entry or the alloca insertion point as provided by the body callback
   // above.
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP(builder.saveIP());
   builder.restoreIP(ompBuilder->CreateParallel(
       builder, allocaIP, bodyGenCB, privCB, finiCB, ifCond, numThreads,
       llvm::omp::OMP_PROC_BIND_default, isCancellable));
   return success();
 }
 
 /// Given an OpenMP MLIR operation, create the corresponding LLVM IR
 /// (including OpenMP runtime calls).

mlir/test/Dialect/OpenMP/invalid.mlir

 // RUN: mlir-opt -split-input-file -verify-diagnostics %s
 
 func @unknown_clause() {
   // expected-error@+1 {{invalid is not a valid clause for the omp.parallel operation}}
   omp.parallel invalid {
   }
 
   return
 }
 
 // -----
 
 func @if_once(%n : i1) {
   // expected-error@+1 {{at most one if clause can appear on the omp.parallel operation}}
-  omp.parallel if(%n) if(%n) {
+  omp.parallel if(%n : i1) if(%n : i1) {
   }
 
   return
 }
 
 // -----
 
 func @num_threads_once(%n : si32) {

mlir/test/Dialect/OpenMP/ops.mlir

func @omp_parallel_pretty(%data_var : memref<i32>, %if_cond : i1, %num_threads : si32) -> () {
 
   // CHECK: omp.parallel private(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>)
   omp.parallel private(%data_var : memref<i32>, %data_var : memref<i32>) firstprivate(%data_var : memref<i32>) {
     omp.terminator
   }
 
   // CHECK omp.parallel shared(%{{.*}} : memref<i32>) copyin(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>)
   omp.parallel shared(%data_var : memref<i32>) copyin(%data_var : memref<i32>, %data_var : memref<i32>) {
-    omp.parallel if(%if_cond) {
+    omp.parallel if(%if_cond: i1) {
       omp.terminator
     }
     omp.terminator
   }
 
   // CHECK omp.parallel if(%{{.*}}) num_threads(%{{.*}} : si32) private(%{{.*}} : memref<i32>) proc_bind(close)
-  omp.parallel num_threads(%num_threads : si32) if(%if_cond) 
+  omp.parallel num_threads(%num_threads : si32) if(%if_cond: i1)
                private(%data_var : memref<i32>) proc_bind(close) {
     omp.terminator
   }
 
   return
 }

mlir/test/Target/openmp-llvm.mlir

// CHECK: define internal void @[[OMP_OUTLINED_FN_2]]
   // CHECK: br label %omp.par.region1
   // CHECK-LABEL: omp.par.region1:
   // CHECK: call void @body(i64 1)
   // CHECK: call void @body(i64 42)
   // CHECK: br label %omp.par.region2
   // CHECK-LABEL: omp.par.region2:
   // CHECK: call void @body(i64 43)
   // CHECK: br label %omp.par.pre_finalize
+
+// CHECK: define void @test_omp_parallel_num_threads_1(i32 %[[NUM_THREADS_VAR_1:.*]])
+llvm.func @test_omp_parallel_num_threads_1(%arg0: !llvm.i32) -> () {
+  // CHECK: %[[GTN_NUM_THREADS_VAR_1:.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GTN_SI_VAR_1:.*]])
+  // CHECK: call void @__kmpc_push_num_threads(%struct.ident_t* @[[GTN_SI_VAR_1]], i32 %[[GTN_NUM_THREADS_VAR_1]], i32 %[[NUM_THREADS_VAR_1]])
+  // CHECK: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN_NUM_THREADS_1:.*]] to {{.*}}
+  omp.parallel num_threads(%arg0: !llvm.i32) {
+    omp.barrier
+    omp.terminator
+  }
+
+  llvm.return
+}
+
+// CHECK: define internal void @[[OMP_OUTLINED_FN_NUM_THREADS_1]]
+  // CHECK: call void @__kmpc_barrier
+
+// CHECK: define void @test_omp_parallel_num_threads_2()
+llvm.func @test_omp_parallel_num_threads_2() -> () {
+  %0 = llvm.mlir.constant(4 : index) : !llvm.i32
+  // CHECK: %[[GTN_NUM_THREADS_VAR_2:.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GTN_SI_VAR_2:.*]])
+  // CHECK: call void @__kmpc_push_num_threads(%struct.ident_t* @[[GTN_SI_VAR_2]], i32 %[[GTN_NUM_THREADS_VAR_2]], i32 4)
+  // CHECK: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN_NUM_THREADS_2:.*]] to {{.*}}
+  omp.parallel num_threads(%0: !llvm.i32) {
+    omp.barrier
+    omp.terminator
+  }
+
+  llvm.return
+}
+
+// CHECK: define internal void @[[OMP_OUTLINED_FN_NUM_THREADS_2]]
+  // CHECK: call void @__kmpc_barrier
+
+// CHECK: define void @test_omp_parallel_num_threads_3()
+llvm.func @test_omp_parallel_num_threads_3() -> () {
+  %0 = llvm.mlir.constant(4 : index) : !llvm.i32
+  // CHECK: %[[GTN_NUM_THREADS_VAR_3_1:.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GTN_SI_VAR_3_1:.*]])
+  // CHECK: call void @__kmpc_push_num_threads(%struct.ident_t* @[[GTN_SI_VAR_3_1]], i32 %[[GTN_NUM_THREADS_VAR_3_1]], i32 4)
+  // CHECK: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN_NUM_THREADS_3_1:.*]] to {{.*}}
+  omp.parallel num_threads(%0: !llvm.i32) {
+    omp.barrier
+    omp.terminator
+  }
+  %1 = llvm.mlir.constant(8 : index) : !llvm.i32
+  // CHECK: %[[GTN_NUM_THREADS_VAR_3_2:.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GTN_SI_VAR_3_2:.*]])
+  // CHECK: call void @__kmpc_push_num_threads(%struct.ident_t* @[[GTN_SI_VAR_3_2]], i32 %[[GTN_NUM_THREADS_VAR_3_2]], i32 8)
+  // CHECK: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN_NUM_THREADS_3_2:.*]] to {{.*}}
+  omp.parallel num_threads(%1: !llvm.i32) {
+    omp.barrier
+    omp.terminator
+  }
+
+  llvm.return
+}
+
+// CHECK: define internal void @[[OMP_OUTLINED_FN_NUM_THREADS_3_2]]
+  // CHECK: call void @__kmpc_barrier
+
+// CHECK: define internal void @[[OMP_OUTLINED_FN_NUM_THREADS_3_1]]
+  // CHECK: call void @__kmpc_barrier
+
+// CHECK: define void @test_omp_parallel_if_1(i32 %[[IF_VAR_1:.*]])
+llvm.func @test_omp_parallel_if_1(%arg0: !llvm.i32) -> () {
+
+// CHECK: %[[IF_COND_VAR_1:.*]] = icmp slt i32 %[[IF_VAR_1]], 0
+  %0 = llvm.mlir.constant(0 : index) : !llvm.i32
+  %1 = llvm.icmp "slt" %arg0, %0 : !llvm.i32
+
+// CHECK: %[[GTN_IF_1:.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[SI_VAR_IF_1:.*]])
+// CHECK: br i1 %[[IF_COND_VAR_1]], label %[[IF_COND_TRUE_BLOCK_1:.*]], label %[[IF_COND_FALSE_BLOCK_1:.*]]
+// CHECK: [[IF_COND_TRUE_BLOCK_1]]:
+// CHECK: br label %[[OUTLINED_CALL_IF_BLOCK_1:.*]]
+// CHECK: [[OUTLINED_CALL_IF_BLOCK_1]]:
+// CHECK: call void {{.*}} @__kmpc_fork_call(%struct.ident_t* @[[SI_VAR_IF_1]], {{.*}} @[[OMP_OUTLINED_FN_IF_1:.*]] to void
+// CHECK: br label %[[OUTLINED_EXIT_IF_1:.*]]
+// CHECK: [[OUTLINED_EXIT_IF_1]]:
+// CHECK: br label %[[OUTLINED_EXIT_IF_2:.*]]
+// CHECK: [[OUTLINED_EXIT_IF_2]]:
+// CHECK: br label %[[RETURN_BLOCK_IF_1:.*]]
+// CHECK: [[IF_COND_FALSE_BLOCK_1]]:
+// CHECK: call void @__kmpc_serialized_parallel(%struct.ident_t* @[[SI_VAR_IF_1]], i32 %[[GTN_IF_1]])
+// CHECK: call void @[[OMP_OUTLINED_FN_IF_1]]
+// CHECK: call void @__kmpc_end_serialized_parallel(%struct.ident_t* @[[SI_VAR_IF_1]], i32 %[[GTN_IF_1]])
+// CHECK: br label %[[RETURN_BLOCK_IF_1]]
+  omp.parallel if(%1 : !llvm.i1) {
+    omp.barrier
+    omp.terminator
+  }
+
+// CHECK: [[RETURN_BLOCK_IF_1]]:
+// CHECK: ret void
+  llvm.return
+}
+
+// CHECK: define internal void @[[OMP_OUTLINED_FN_IF_1]]
+  // CHECK: call void @__kmpc_barrier