[MLIR][OpenMP] Add case for TargetDataOp when traversing inner regions

ROCm · Mar 28, 2024 · 2e08836 · 2e08836
1 parent d5bf7d3
commit 2e08836
Showing 1 changed file with 13 additions and 5 deletions.
diff --git 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);
 }