[MLIR][OpenMP] Fix compiler crash when translating omp.target to LLVM IR

[skatrak]

Summary

This patch fixes a compiler crash that would happen during translation to LLVM IR if the optional map argument of the omp.target operation was not present. A unit test is added to ensure this has been fixed.

[agozillon]

LGTM, nice catch! My patch series changes the way map is accessed soon but it'll take a while to land I imagine. However, in the meantime perhaps there are other operations that suffer similarly at the moment, like the enter/exit data operations, I believe they use getMapType as well currently, but I don't know if they access the getMapType information in the same way (or if they can be tested in the same way, I don't know how/if they lower fully to LLVM-IR at the moment, I'm unfamiliar with them outside of knowing they have map information that needs processed).

[skatrak]

In D158722#4614261, @agozillon wrote:

LGTM, nice catch! My patch series changes the way map is accessed soon but it'll take a while to land I imagine. However, in the meantime perhaps there are other operations that suffer similarly at the moment, like the enter/exit data operations, I believe they use getMapType as well currently, but I don't know if they access the getMapType information in the same way (or if they can be tested in the same way, I don't know how/if they lower fully to LLVM-IR at the moment, I'm unfamiliar with them outside of knowing they have map information that needs processed).

Thank you Andrew for looking at this. As far as I can see, the only ops for which the map_types argument is optional are omp.target_data and omp.target, and for omp.target_data the result of getMapTypes() is always checked to exist before using it. For omp.target_{enter|exit}_data the argument is required, so using the result of getMapTypes() is not going to crash the compiler there. I believe this small patch is the only place that we missed that check.

I'll wait until Monday to land this patch, in case someone else has any comments before accepting it.

[agozillon]

In D158722#4616571, @skatrak wrote:

In D158722#4614261, @agozillon wrote:

LGTM, nice catch! My patch series changes the way map is accessed soon but it'll take a while to land I imagine. However, in the meantime perhaps there are other operations that suffer similarly at the moment, like the enter/exit data operations, I believe they use getMapType as well currently, but I don't know if they access the getMapType information in the same way (or if they can be tested in the same way, I don't know how/if they lower fully to LLVM-IR at the moment, I'm unfamiliar with them outside of knowing they have map information that needs processed).

Thank you Andrew for looking at this. As far as I can see, the only ops for which the map_types argument is optional are omp.target_data and omp.target, and for omp.target_data the result of getMapTypes() is always checked to exist before using it. For omp.target_{enter|exit}_data the argument is required, so using the result of getMapTypes() is not going to crash the compiler there. I believe this small patch is the only place that we missed that check.

I'll wait until Monday to land this patch, in case someone else has any comments before accepting it.

Sounds good! Thank you for looking into it further.

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

convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   int32_t defaultValTeams = -1;
   int32_t defaultValThreads = 0;
 
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
       findAllocaInsertPoint(builder, moduleTranslation);
 
   DataLayout DL = DataLayout(opInst.getParentOfType<ModuleOp>());
   SmallVector<Value> mapOperands = targetOp.getMapOperands();
-  ArrayAttr mapTypes = targetOp.getMapTypes().value();
+  ArrayAttr mapTypes = targetOp.getMapTypes().value_or(nullptr);
 
   llvm::OpenMPIRBuilder::MapInfosTy combinedInfos;
   auto genMapInfoCB = [&](llvm::OpenMPIRBuilder::InsertPointTy codeGenIP)
       -> llvm::OpenMPIRBuilder::MapInfosTy & {
     builder.restoreIP(codeGenIP);
     genMapInfos(builder, moduleTranslation, DL, combinedInfos, mapOperands,
                 mapTypes, {}, {}, true);
     return combinedInfos;

mlir/test/Target/LLVMIR/omptarget-region-llvm.mlir

omp.target map((to -> %3 : !llvm.ptr<i32>), (to -> %5 : !llvm.ptr<i32>), (from -> %7 : !llvm.ptr<i32>)) {
       %8 = llvm.load %3 : !llvm.ptr<i32>
       %9 = llvm.load %5 : !llvm.ptr<i32>
       %10 = llvm.add %8, %9  : i32
       llvm.store %10, %7 : !llvm.ptr<i32>
       omp.terminator
     }
     llvm.return
   }
+
+  llvm.func @omp_target_no_map() {
+    omp.target {
+      omp.terminator
+    }
+    llvm.return
+  }
 }
 
-// CHECK: call i32 @__tgt_target_kernel(ptr @4, i64 -1, i32 -1, i32 0, ptr @.__omp_offloading_[[DEV:.*]]_[[FIL:.*]]_omp_target_region__l[[LINE:.*]].region_id, ptr %kernel_args)
+// CHECK: define void @omp_target_region_()
+// CHECK: call i32 @__tgt_target_kernel(ptr @4, i64 -1, i32 -1, i32 0, ptr @.__omp_offloading_[[DEV:.*]]_[[FIL:.*]]_omp_target_region__l[[LINE1:.*]].region_id, ptr %kernel_args)
+
+// CHECK: br i1 %{{.*}}, label %omp_offload.failed, label %omp_offload.cont
+// CHECK: omp_offload.failed:
+// CHECK: call void @__omp_offloading_[[DEV]]_[[FIL]]_omp_target_region__l[[LINE1]](ptr %{{.*}}, ptr %{{.*}}, ptr %{{.*}})
+// CHECK: omp_offload.cont:
+
+// CHECK: define void @omp_target_no_map()
+// CHECK: call i32 @__tgt_target_kernel(ptr @4, i64 -1, i32 -1, i32 0, ptr @.__omp_offloading_[[DEV:.*]]_[[FIL:.*]]_omp_target_no_map_l[[LINE2:.*]].region_id, ptr %kernel_args)
 
 // CHECK: br i1 %{{.*}}, label %omp_offload.failed, label %omp_offload.cont
 // CHECK: omp_offload.failed:
-// CHECK: call void @__omp_offloading_[[DEV]]_[[FIL]]_omp_target_region__l[[LINE]](ptr %{{.*}}, ptr %{{.*}}, ptr %{{.*}})
+// CHECK: call void @__omp_offloading_[[DEV]]_[[FIL]]_omp_target_no_map_l[[LINE2]]()
 // CHECK: omp_offload.cont:
 
-// CHECK: define internal void @__omp_offloading_[[DEV]]_[[FIL]]_omp_target_region__l[[LINE]](ptr %[[ADDR_A:.*]], ptr %[[ADDR_B:.*]], ptr %[[ADDR_C:.*]])
+// CHECK: define internal void @__omp_offloading_[[DEV]]_[[FIL]]_omp_target_region__l[[LINE1]](ptr %[[ADDR_A:.*]], ptr %[[ADDR_B:.*]], ptr %[[ADDR_C:.*]])
 // CHECK: %[[VAL_A:.*]] = load i32, ptr %[[ADDR_A]], align 4
 // CHECK: %[[VAL_B:.*]] = load i32, ptr %[[ADDR_B]], align 4
 // CHECK: %[[SUM:.*]] = add i32 %[[VAL_A]], %[[VAL_B]]
 // CHECK: store i32 %[[SUM]], ptr %[[ADDR_C]], align 4
 
+// CHECK: define internal void @__omp_offloading_[[DEV]]_[[FIL]]_omp_target_no_map_l[[LINE2]]()
+// CHECK: ret void