From 2e08836c43f9521cb3ffeb3f9f846175efe435c1 Mon Sep 17 00:00:00 2001
From: Jan Leyonberg <jan_sjodin@yahoo.com>
Date: Thu, 28 Mar 2024 12:15:56 -0400
Subject: [PATCH] [MLIR][OpenMP] Add case for TargetDataOp when traversing
 inner regions

---
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp       | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 456476d5532f93..9860b0f90a894d 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -3622,16 +3622,24 @@ convertInternalTargetOp(Operation *op, llvm::IRBuilderBase &builder,
 
 static LogicalResult
 convertTopLevelTargetOp(Operation *op, llvm::IRBuilderBase &builder,
-                              LLVM::ModuleTranslation &moduleTranslation) {
+                        LLVM::ModuleTranslation &moduleTranslation) {
   if (isa<omp::TargetOp>(op))
     return convertOmpTarget(*op, builder, moduleTranslation);
   if (isa<omp::TargetDataOp>(op))
     return convertOmpTargetData(op, builder, moduleTranslation);
   bool interrupted =
-      op->walk<WalkOrder::PreOrder>([&](omp::TargetOp targetOp) {
-          if (failed(convertOmpTarget(*targetOp, builder, moduleTranslation)))
-            return WalkResult::interrupt();
-          return WalkResult::skip();
+      op->walk<WalkOrder::PreOrder>([&](Operation *oper) {
+          if (isa<omp::TargetOp>(oper)) {
+            if (failed(convertOmpTarget(*oper, builder, moduleTranslation)))
+              return WalkResult::interrupt();
+            return WalkResult::skip();
+          }
+          if (isa<omp::TargetDataOp>(oper)) {
+            if (failed(convertOmpTargetData(oper, builder, moduleTranslation)))
+              return WalkResult::interrupt();
+            return WalkResult::skip();
+          }
+          return WalkResult::advance();
         }).wasInterrupted();
   return failure(interrupted);
 }