[MLIR][OpenMP] Add support for depend clause

[psoni2628]

Summary

This patch adds support for the OpenMP 4.0 depend clause (in, out,
inout) of the task construct to the definition of the OpenMP MLIR
dialect and translation from MLIR to LLVM IR using OMPIRBuilder.

[kiranchandramohan]

Looks mostly good. A few questions/comments.

Would we be able to construct task graphs based on task dependencies? Or would this be possible only at runtime?
https://discourse.llvm.org/t/modeling-task-graphs-in-mlir/4080

Will this cover all kinds of dependencies?

Also call out in the summary that you are adding support for the depend clause in this patch.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

701

Nit: You can move this close to its use.

708

Nit: Can you add a comment why both map to DepInout?

722

Nit: What are these nullptrs?

mlir/test/Target/LLVMIR/openmp-llvm.mlir

2265–2269

Can we reduce the code here to the minimum?

[psoni2628]

• Address review comments

[psoni2628]

In D142730#4091360, @kiranchandramohan wrote:

Looks mostly good. A few questions/comments.

Would we be able to construct task graphs based on task dependencies? Or would this be possible only at runtime?
https://discourse.llvm.org/t/modeling-task-graphs-in-mlir/4080

Will this cover all kinds of dependencies?

Also call out in the summary that you are adding support for the depend clause in this patch.

I think you would be able to construct a task graph based on dependencies. In some situations where values are only known at run time, it may be challenging to figure out dependencies statically.

This covers all kinds of dependencies in OpenMP 4.0. There are newer ones like mutexinoutset introduced in OpenMP 5.0 that we can add in a later patch.

Could you please clarify your last point here? I believe the summary already mentions that the patch is adding support for the depend clause.

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

708

Sure, I have added a comment. I wanted to add this commit as a reference for the comment: https://reviews.llvm.org/rG92e82f9cce62ae50e72530563811f0d3bd98e892.

mlir/test/Target/LLVMIR/openmp-llvm.mlir

2265–2269

I have removed the 2 unused arguments. Besides that, we could remove some of the non-depend-specific task code (e.g. outlined_fn) copied from the test above this one, but I think it is worth keeping to ensure the there are no regressions in that when using the depend clause. Let me know if you agree.

[kiranchandramohan]

LG.

This covers all kinds of dependencies in OpenMP 4.0. There are newer ones like mutexinoutset introduced in OpenMP 5.0 that we can add in a later patch.

Could you clarify what is covered and what is not covered in the summary?

Could you please clarify your last point here? I believe the summary already mentions that the patch is adding support for the depend clause.

The summary only talks about translation.

[kiranchandramohan]

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

132

Please add tests for this in the following file.
mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir

[psoni2628]

• Fix a comment Print Reduction Clause -> Print Depend Clause
• Add a test to mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir as suggested by reviewer comments

[psoni2628]

In D142730#4102431, @kiranchandramohan wrote:

LG.

This covers all kinds of dependencies in OpenMP 4.0. There are newer ones like mutexinoutset introduced in OpenMP 5.0 that we can add in a later patch.

Could you clarify what is covered and what is not covered in the summary?

Could you please clarify your last point here? I believe the summary already mentions that the patch is adding support for the depend clause.

The summary only talks about translation.

Please let me know if the summary is good now.

flang/lib/Lower/OpenMP.cpp

if (blockDirective.v == llvm::omp::OMPD_parallel) {
     auto orderedOp = firOpBuilder.create<mlir::omp::OrderedRegionOp>(
         currentLocation, /*simd=*/false);
     createBodyOfOp<omp::OrderedRegionOp>(orderedOp, converter, currentLocation,
                                          eval);
   } else if (blockDirective.v == llvm::omp::OMPD_task) {
     auto taskOp = firOpBuilder.create<mlir::omp::TaskOp>(
         currentLocation, ifClauseOperand, finalClauseOperand, untiedAttr,
         mergeableAttr, /*in_reduction_vars=*/ValueRange(),
-        /*in_reductions=*/nullptr, priorityClauseOperand, allocateOperands,
-        allocatorOperands);
+        /*in_reductions=*/nullptr, priorityClauseOperand, /*depends=*/nullptr,
+        /*depend_vars=*/ValueRange(), allocateOperands, allocatorOperands);
     createBodyOfOp(taskOp, converter, currentLocation, eval, &opClauseList);
   } else if (blockDirective.v == llvm::omp::OMPD_taskgroup) {
     // TODO: Add task_reduction support
     auto taskGroupOp = firOpBuilder.create<mlir::omp::TaskGroupOp>(
         currentLocation, /*task_reduction_vars=*/ValueRange(),
         /*task_reductions=*/nullptr, allocateOperands, allocatorOperands);
     createBodyOfOp(taskGroupOp, converter, currentLocation, eval,
                    &opClauseList);

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

def YieldOp : OpenMP_Op<"yield",
 
   let assemblyFormat = [{ ( `(` $results^ `:` type($results) `)` )? attr-dict}];
 }
 
 //===----------------------------------------------------------------------===//
 // 2.10.1 task Construct
 //===----------------------------------------------------------------------===//
 
+def ClauseTaskDependIn    : I32EnumAttrCase<"taskdependin",    0>;
+def ClauseTaskDependOut   : I32EnumAttrCase<"taskdependout",   1>;
+def ClauseTaskDependInOut : I32EnumAttrCase<"taskdependinout", 2>;
+
+def ClauseTaskDepend : I32EnumAttr<
+    "ClauseTaskDepend",
+    "task depend clause",
+    [ClauseTaskDependIn, ClauseTaskDependOut, ClauseTaskDependInOut]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::omp";
+}
+def ClauseTaskDependAttr :
+  EnumAttr<OpenMP_Dialect, ClauseTaskDepend, "clause_task_depend"> {
+  let assemblyFormat = "`(` $value `)`";
+}
+def TaskDependArrayAttr :
+  TypedArrayAttrBase<ClauseTaskDependAttr, "clause_task_depend array attr"> {
+    let constBuilderCall = ?;
+  }
+
 def TaskOp : OpenMP_Op<"task", [AttrSizedOperandSegments,
                        OutlineableOpenMPOpInterface, AutomaticAllocationScope,
                        ReductionClauseInterface]> {
   let summary = "task construct";
   let description = [{
     The task construct defines an explicit task.
 
     For definitions of "undeferred task", "included task", "final task" and
let description = [{
     The `priority` clause is a hint for the priority of the generated task.
     The `priority` is a non-negative integer expression that provides a hint for
     task execution order. Among all tasks ready to be executed, higher priority
     tasks (those with a higher numerical value in the priority clause
     expression) are recommended to execute before lower priority ones. The
     default priority-value when no priority clause is specified should be
     assumed to be zero (the lowest priority).
 
+    The `depends` and `depend_vars` arguments are variadic lists of values
+    that specify the dependencies of this particular task in relation to
+    other tasks.
+
     The `allocators_vars` and `allocate_vars` arguments are a variadic list of
     values that specify the memory allocator to be used to obtain storage for
     private values.
 
   }];
 
   // TODO: depend, affinity and detach clauses
   let arguments = (ins Optional<I1>:$if_expr,
                        Optional<I1>:$final_expr,
                        UnitAttr:$untied,
                        UnitAttr:$mergeable,
                        Variadic<OpenMP_PointerLikeType>:$in_reduction_vars,
                        OptionalAttr<SymbolRefArrayAttr>:$in_reductions,
                        Optional<I32>:$priority,
+                       OptionalAttr<TaskDependArrayAttr>:$depends,
+                       Variadic<OpenMP_PointerLikeType>:$depend_vars,
                        Variadic<AnyType>:$allocate_vars,
                        Variadic<AnyType>:$allocators_vars);
   let regions = (region AnyRegion:$region);
   let assemblyFormat = [{
     oilist(`if` `(` $if_expr `)`
           |`final` `(` $final_expr `)`
           |`untied` $untied
           |`mergeable` $mergeable
           |`in_reduction` `(`
               custom<ReductionVarList>(
                 $in_reduction_vars, type($in_reduction_vars), $in_reductions
               ) `)`
           |`priority` `(` $priority `)`
           |`allocate` `(`
               custom<AllocateAndAllocator>(
                 $allocate_vars, type($allocate_vars),
                 $allocators_vars, type($allocators_vars)
               ) `)`
+          |`depend` `(`
+              custom<DependVarList>(
+                $depend_vars, type($depend_vars), $depends
+              ) `)`
     ) $region attr-dict
   }];
   let extraClassDeclaration = [{
     /// Returns the reduction variables
     SmallVector<Value> getReductionVars() {
       return SmallVector<Value>(getInReductionVars().begin(),
                                 getInReductionVars().end());
     }

mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp

rewriter.replaceOpWithNewOp<omp::ReductionOp>(
         curOp, TypeRange(), adaptor.getOperands(), curOp->getAttrs());
     return success();
   }
 };
 } // namespace
 
 void mlir::configureOpenMPToLLVMConversionLegality(
     ConversionTarget &target, LLVMTypeConverter &typeConverter) {
-  target.addDynamicallyLegalOp<mlir::omp::CriticalOp, mlir::omp::ParallelOp,
-                               mlir::omp::WsLoopOp, mlir::omp::SimdLoopOp,
-                               mlir::omp::MasterOp, mlir::omp::SectionsOp,
-                               mlir::omp::SingleOp>([&](Operation *op) {
+  target.addDynamicallyLegalOp<
+      mlir::omp::CriticalOp, mlir::omp::ParallelOp, mlir::omp::WsLoopOp,
+      mlir::omp::SimdLoopOp, mlir::omp::MasterOp, mlir::omp::SectionsOp,
+      mlir::omp::SingleOp, mlir::omp::TaskOp>([&](Operation *op) {
     return typeConverter.isLegal(&op->getRegion(0)) &&
            typeConverter.isLegal(op->getOperandTypes()) &&
            typeConverter.isLegal(op->getResultTypes());
   });
   target
       .addDynamicallyLegalOp<mlir::omp::AtomicReadOp, mlir::omp::AtomicWriteOp,
                              mlir::omp::FlushOp, mlir::omp::ThreadprivateOp>(
           [&](Operation *op) {
             return typeConverter.isLegal(op->getOperandTypes()) &&
                    typeConverter.isLegal(op->getResultTypes());
           });
   target.addDynamicallyLegalOp<mlir::omp::ReductionOp>([&](Operation *op) {
     return typeConverter.isLegal(op->getOperandTypes());
   });
 }
 
 void mlir::populateOpenMPToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                   RewritePatternSet &patterns) {
   patterns.add<
       ReductionOpConversion, RegionOpConversion<omp::CriticalOp>,
       RegionOpConversion<omp::MasterOp>, ReductionOpConversion,
       RegionOpConversion<omp::MasterOp>, RegionOpConversion<omp::ParallelOp>,
       RegionOpConversion<omp::WsLoopOp>, RegionOpConversion<omp::SectionsOp>,
       RegionOpConversion<omp::SimdLoopOp>, RegionOpConversion<omp::SingleOp>,
+      RegionOpConversion<omp::TaskOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicReadOp>,
       RegionLessOpWithVarOperandsConversion<omp::AtomicWriteOp>,
       RegionLessOpWithVarOperandsConversion<omp::FlushOp>,
       RegionLessOpWithVarOperandsConversion<omp::ThreadprivateOp>>(converter);
 }
 
 namespace {
 struct ConvertOpenMPToLLVMPass

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

if (decl.getAccumulatorType() && decl.getAccumulatorType() != varType)
              << ") to be the same type as reduction declaration ("
              << decl.getAccumulatorType() << ")";
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// Parser, printer and verifier for DependVarList
+//===----------------------------------------------------------------------===//
+
+/// depend-entry-list ::= depend-entry
+///                        | depend-entry-list `,` depend-entry
+/// depend-entry ::= depend-kind `->` ssa-id `:` type
+static ParseResult
+parseDependVarList(OpAsmParser &parser,
+                   SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
+                   SmallVectorImpl<Type> &types, ArrayAttr &dependsArray) {
+  SmallVector<ClauseTaskDependAttr> dependVec;
+  if (failed(parser.parseCommaSeparatedList([&]() {
+        StringRef keyword;
+        if (parser.parseKeyword(&keyword) || parser.parseArrow() ||
+            parser.parseOperand(operands.emplace_back()) ||
+            parser.parseColonType(types.emplace_back()))
+          return failure();
+        if (std::optional<ClauseTaskDepend> keywordDepend =
+                (symbolizeClauseTaskDepend(keyword)))
+          dependVec.emplace_back(
+              ClauseTaskDependAttr::get(parser.getContext(), *keywordDepend));
+        else
+          return failure();
+        return success();
+      })))
+    return failure();
+  SmallVector<Attribute> depends(dependVec.begin(), dependVec.end());
+  dependsArray = ArrayAttr::get(parser.getContext(), depends);
+  return success();
+}
+
+/// Print Reduction clause
+static void printDependVarList(OpAsmPrinter &p, Operation *op,
+                               OperandRange dependVars, TypeRange dependTypes,
+                               std::optional<ArrayAttr> depends) {
+
+  for (unsigned i = 0, e = depends->size(); i < e; ++i) {
+    if (i != 0)
+      p << ", ";
+    p << stringifyClauseTaskDepend(
+             (*depends)[i].cast<mlir::omp::ClauseTaskDependAttr>().getValue())
+      << " -> " << dependVars[i] << " : " << dependTypes[i];
+  }
+}
+
+static LogicalResult verifyDependVarList(Operation *op,
+                                         Optional<ArrayAttr> depends,
+                                         OperandRange dependVars) {
+  if (!dependVars.empty()) {
+    if (!depends || depends->size() != dependVars.size())
+      return op->emitOpError() << "expected as many depend values"
+                                  " as depend variables";
+  } else {
+    if (depends)
+      return op->emitOpError() << "unexpected depend values";
+    return success();
+  }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // Parser, printer and verifier for Synchronization Hint (2.17.12)
 //===----------------------------------------------------------------------===//
 
 /// Parses a Synchronization Hint clause. The value of hint is an integer
 /// which is a combination of different hints from `omp_sync_hint_t`.
 ///
 /// hint-clause = `hint` `(` hint-value `)`
 static ParseResult parseSynchronizationHint(OpAsmParser &parser,
LogicalResult ReductionOp::verify() {
   }
   return emitOpError() << "the accumulator is not used by the parent";
 }
 
 //===----------------------------------------------------------------------===//
 // TaskOp
 //===----------------------------------------------------------------------===//
 LogicalResult TaskOp::verify() {
-  return verifyReductionVarList(*this, getInReductions(), getInReductionVars());
+  LogicalResult verifyDependVars =
+      verifyDependVarList(*this, getDepends(), getDependVars());
+  return failed(verifyDependVars)
+             ? verifyDependVars
+             : verifyReductionVarList(*this, getInReductions(),
+                                      getInReductionVars());
 }
 
 //===----------------------------------------------------------------------===//
 // TaskGroupOp
 //===----------------------------------------------------------------------===//
 LogicalResult TaskGroupOp::verify() {
   return verifyReductionVarList(*this, getTaskReductions(),
                                 getTaskReductionVars());

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

if (taskOp.getIfExpr() || taskOp.getFinalExpr() || taskOp.getUntiedAttr() ||
       taskOp.getPriority() || !taskOp.getAllocateVars().empty()) {
     return taskOp.emitError("unhandled clauses for translation to LLVM IR");
   }
   auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codegenIP) {
     builder.restoreIP(codegenIP);
     convertOmpOpRegions(taskOp.getRegion(), "omp.task.region", builder,
                         moduleTranslation, bodyGenStatus);
   };
+
+  SmallVector<llvm::OpenMPIRBuilder::DependData> dds;
+  if (!taskOp.getDependVars().empty() && taskOp.getDepends()) {
+    for (auto dep :
+         llvm::zip(taskOp.getDependVars(), taskOp.getDepends()->getValue())) {
+      llvm::Value *depVal = moduleTranslation.lookupValue(std::get<0>(dep));
+      llvm::omp::RTLDependenceKindTy type;
+      switch (
+          std::get<1>(dep).cast<mlir::omp::ClauseTaskDependAttr>().getValue()) {
+      case mlir::omp::ClauseTaskDepend::taskdependin:
+        type = llvm::omp::RTLDependenceKindTy::DepIn;
+        break;
+      case mlir::omp::ClauseTaskDepend::taskdependout:
+      case mlir::omp::ClauseTaskDepend::taskdependinout:
+        type = llvm::omp::RTLDependenceKindTy::DepInOut;
+        break;
+      };
+      llvm::OpenMPIRBuilder::DependData dd(type, depVal->getType(), depVal);
+      dds.emplace_back(dd);
+    }
+  }
+
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
       findAllocaInsertPoint(builder, moduleTranslation);
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createTask(
-      ompLoc, allocaIP, bodyCB, !taskOp.getUntied()));
+      ompLoc, allocaIP, bodyCB, !taskOp.getUntied(), nullptr, nullptr, dds));
   return bodyGenStatus;
 }
 
 /// Converts an OpenMP taskgroup construct into LLVM IR using OpenMPIRBuilder.
 static LogicalResult
 convertOmpTaskgroupOp(omp::TaskGroupOp tgOp, llvm::IRBuilderBase &builder,
                       LLVM::ModuleTranslation &moduleTranslation) {
   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;

mlir/test/Dialect/OpenMP/invalid.mlir

func.func @omp_single(%data_var : memref<i32>) -> () {
   "omp.single" (%data_var) ({
     omp.barrier
   }) {operand_segment_sizes = array<i32: 1,0>} : (memref<i32>) -> ()
   return
 }
 
 // -----
 
+func.func @omp_task_depend(%data_var: memref<i32>) {
+  // expected-error @below {{op expected as many depend values as depend variables}}
+    "omp.task"(%data_var) ({
+      "omp.terminator"() : () -> ()
+    }) {depends = [], operand_segment_sizes = array<i32: 0, 0, 0, 0, 1, 0, 0>} : (memref<i32>) -> ()
+   "func.return"() : () -> ()
+}
+
+// -----
+
 func.func @omp_task(%ptr: !llvm.ptr<f32>) {
   // expected-error @below {{op expected symbol reference @add_f32 to point to a reduction declaration}}
   omp.task in_reduction(@add_f32 -> %ptr : !llvm.ptr<f32>) {
     // CHECK: "test.foo"() : () -> ()
     "test.foo"() : () -> ()
     // CHECK: omp.terminator
     omp.terminator
   }

mlir/test/Dialect/OpenMP/ops.mlir

omp.task if(%bool_var) final(%bool_var) untied
     "test.foo"() : () -> ()
     // CHECK: omp.terminator
     omp.terminator
   }
 
   return
 }
 
+// CHECK-LABEL: @omp_task_depend
+// CHECK-SAME: (%arg0: memref<i32>, %arg1: memref<i32>) {
+func.func @omp_task_depend(%arg0: memref<i32>, %arg1: memref<i32>) {
+  // CHECK:  omp.task   depend(taskdependin -> %arg0 : memref<i32>, taskdependin -> %arg1 : memref<i32>, taskdependinout -> %arg0 : memref<i32>) {
+  omp.task   depend(taskdependin -> %arg0 : memref<i32>, taskdependin -> %arg1 : memref<i32>, taskdependinout -> %arg0 : memref<i32>) {
+    // CHECK: "test.foo"() : () -> ()
+    "test.foo"() : () -> ()
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+  return
+}
+
 func.func @omp_threadprivate() {
   %0 = arith.constant 1 : i32
   %1 = arith.constant 2 : i32
   %2 = arith.constant 3 : i32
 
   // CHECK: [[ARG0:%.*]] = llvm.mlir.addressof @_QFsubEx : !llvm.ptr<i32>
   // CHECK: {{.*}} = omp.threadprivate [[ARG0]] : !llvm.ptr<i32> -> !llvm.ptr<i32>
   %3 = llvm.mlir.addressof @_QFsubEx : !llvm.ptr<i32>

mlir/test/Target/LLVMIR/openmp-llvm.mlir

 
 // CHECK: define i32 @[[wrapper_fn]](i32 %{{.+}}) {
 // CHECK:   call void @[[outlined_fn]]()
 // CHECK:   ret i32 0
 // CHECK: }
 
 // -----
 
+// CHECK-LABEL: define void @omp_task_with_deps
+// CHECK-SAME: (i32 %[[x:.+]], i32 %[[y:.+]], ptr %[[zaddr:.+]])
+// CHECK:  %[[dep_arr_addr:.+]] = alloca [1 x %struct.kmp_dep_info], align 8
+// CHECK:  %[[dep_arr_addr_0:.+]] = getelementptr inbounds [1 x %struct.kmp_dep_info], ptr %[[dep_arr_addr]], i64 0, i64 0
+// CHECK:  %[[dep_arr_addr_0_val:.+]] = getelementptr inbounds %struct.kmp_dep_info, ptr %[[dep_arr_addr_0]], i32 0, i32 0
+// CHECK:  %6 = ptrtoint ptr %2 to i64
+// CHECK:  store i64 %6, ptr %[[dep_arr_addr_0_val]], align 4
+// CHECK:  %[[dep_arr_addr_0_size:.+]] = getelementptr inbounds %struct.kmp_dep_info, ptr %[[dep_arr_addr_0]], i32 0, i32 1
+// CHECK:  store i64 8, ptr %[[dep_arr_addr_0_size]], align 4
+// CHECK:  %[[dep_arr_addr_0_kind:.+]] = getelementptr inbounds %struct.kmp_dep_info, ptr %[[dep_arr_addr_0]], i32 0, i32 2
+// CHECK: store i8 1, ptr %[[dep_arr_addr_0_kind]], align 1
+llvm.func @omp_task_with_deps(%x: i32, %y: i32, %zaddr: !llvm.ptr<i32>) {
+  // CHECK: %[[omp_global_thread_num:.+]] = call i32 @__kmpc_global_thread_num({{.+}})
+  // CHECK: %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc
+  // CHECK-SAME: (ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 0,
+  // CHECK-SAME:  i64 0, ptr @[[wrapper_fn:.+]])
+  // CHECK: call i32 @__kmpc_omp_task_with_deps(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]], {{.*}})
+  omp.task depend(taskdependin -> %zaddr : !llvm.ptr<i32>) {
+    %n = llvm.mlir.constant(1 : i64) : i64
+    %valaddr = llvm.alloca %n x i32 : (i64) -> !llvm.ptr<i32>
+    %val = llvm.load %valaddr : !llvm.ptr<i32>
+    %double = llvm.add %val, %val : i32
+    llvm.store %double, %valaddr : !llvm.ptr<i32>
+    omp.terminator
+  }
+  llvm.return
+}
+
+// CHECK: define internal void @[[outlined_fn:.+]]()
+// CHECK: task.alloca{{.*}}:
+// CHECK:   br label %[[task_body:[^, ]+]]
+// CHECK: [[task_body]]:
+// CHECK:   br label %[[task_region:[^, ]+]]
+// CHECK: [[task_region]]:
+// CHECK:   %[[alloca:.+]] = alloca i32, i64 1
+// CHECK:   %[[val:.+]] = load i32, ptr %[[alloca]]
+// CHECK:   %[[newval:.+]] = add i32 %[[val]], %[[val]]
+// CHECK:   store i32 %[[newval]], ptr %{{[^, ]+}}
+// CHECK:   br label %[[exit_stub:[^, ]+]]
+// CHECK: [[exit_stub]]:
+// CHECK:   ret void
+
+// CHECK: define i32 @[[wrapper_fn]](i32 %{{.+}}) {
+// CHECK:   call void @[[outlined_fn]]()
+// CHECK:   ret i32 0
+// CHECK: }
+
+// -----
+
 // CHECK-LABEL: define void @omp_task
 // CHECK-SAME: (i32 %[[x:.+]], i32 %[[y:.+]], ptr %[[zaddr:.+]])
 module attributes {llvm.target_triple = "x86_64-unknown-linux-gnu"} {
   llvm.func @omp_task(%x: i32, %y: i32, %zaddr: !llvm.ptr<i32>) {
     // CHECK: %[[diff:.+]] = sub i32 %[[x]], %[[y]]
     %diff = llvm.sub %x, %y : i32
     // CHECK: store i32 %[[diff]], ptr %2
     llvm.store %diff, %zaddr : !llvm.ptr<i32>