[MLIR][OpenMP] Set AllocaIP correctly for constructs nested in task

[kiranchandramohan]

Summary

Fixes #60874

[NimishMishra]

LGTM

[psoni2628]

This does not appear to solve the issue in D146766. I tried the following test case below and it still segmentation faults at run time. Could you please check if the test case works for you?

subroutine foo()
  implicit none
  integer::x,y
  x=0
  y=2
   !$omp task shared(x,y)
   y=x
   !$omp end task
   !$omp taskwait
   print*,"y=",y
end subroutine foo

program p
implicit none
call foo()
end program p

[kiranchandramohan]

In D151027#4359205, @psoni2628 wrote:

This does not appear to solve the issue in D146766. I tried the following test case below and it still segmentation faults at run time. Could you please check if the test case works for you?

subroutine foo()
  implicit none
  integer::x,y
  x=0
  y=2
   !$omp task shared(x,y)
   y=x
   !$omp end task
   !$omp taskwait
   print*,"y=",y
end subroutine foo

program p
implicit none
call foo()
end program p

Yes, you are right. It does not fix the issue. I kind of forgot that the issue in D146766 is a runtime crash. Anyway this helps fix a crash during compilation for openmp constructs nested in task. So is it OK to proceed with this? I will amend the summary and remove the mention of D146766.

[psoni2628]

In D151027#4359216, @kiranchandramohan wrote:

In D151027#4359205, @psoni2628 wrote:

This does not appear to solve the issue in D146766. I tried the following test case below and it still segmentation faults at run time. Could you please check if the test case works for you?

subroutine foo()
  implicit none
  integer::x,y
  x=0
  y=2
   !$omp task shared(x,y)
   y=x
   !$omp end task
   !$omp taskwait
   print*,"y=",y
end subroutine foo

program p
implicit none
call foo()
end program p

Yes, you are right. It does not fix the issue. I kind of forgot that the issue in D146766 is a runtime crash. Anyway this helps fix a crash during compilation for openmp constructs nested in task. So is it OK to proceed with this? I will amend the summary and remove the mention of D146766.

Yes, other than that, it works fine.

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

convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
   LogicalResult bodyGenStatus = success();
   if (taskOp.getIfExpr() || taskOp.getFinalExpr() || taskOp.getUntiedAttr() ||
       taskOp.getMergeableAttr() || taskOp.getInReductions() ||
       taskOp.getPriority() || !taskOp.getAllocateVars().empty()) {
     return taskOp.emitError("unhandled clauses for translation to LLVM IR");
   }
   auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codegenIP) {
+    // Save the alloca insertion point on ModuleTranslation stack for use in
+    // nested regions.
+    LLVM::ModuleTranslation::SaveStack<OpenMPAllocaStackFrame> frame(
+        moduleTranslation, allocaIP);
+
     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 :

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

 
 // CHECK: define i32 @[[wrapper_fn]](i32 %{{.+}}, ptr %[[task_data:.+]]) {
 // CHECK:   call void @[[outlined_fn]](ptr %[[task_data]])
 // CHECK:   ret i32 0
 // CHECK: }
 
 // -----
 
+llvm.func @par_task_(%arg0: !llvm.ptr<i32> {fir.bindc_name = "a"}) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  omp.task   {
+    omp.parallel   {
+      llvm.store %0, %arg0 : !llvm.ptr<i32>
+      omp.terminator
+    }
+    omp.terminator
+  }
+  llvm.return
+}
+
+// CHECK-LABEL: @par_task_
+// CHECK: %[[TASK_ALLOC:.*]] = call ptr @__kmpc_omp_task_alloc({{.*}}ptr @par_task_..omp_par.wrapper)
+// CHECK: call i32 @__kmpc_omp_task({{.*}}, ptr %[[TASK_ALLOC]])
+// CHECK-LABEL: define internal void @par_task_..omp_par
+// CHECK: %[[ARG_ALLOC:.*]] = alloca { ptr }, align 8
+// CHECK: call void ({{.*}}) @__kmpc_fork_call({{.*}}, ptr @par_task_..omp_par..omp_par, ptr %[[ARG_ALLOC]])
+// CHECK: define internal void @par_task_..omp_par..omp_par
+// CHECK: define i32 @par_task_..omp_par.wrapper
+// CHECK: call void @par_task_..omp_par
+// -----
+
 llvm.func @foo() -> ()
 
 llvm.func @omp_taskgroup(%x: i32, %y: i32, %zaddr: !llvm.ptr<i32>) {
   omp.taskgroup {
     llvm.call @foo() : () -> ()
     omp.terminator
   }
   llvm.return