[WIP][flang][OpenMP] Translate OpenMP scopes when compiling for target device

Author: ergawy

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

@@ -542,6 +542,13 @@ static llvm::omp::ProcBindKind getProcBindKind(omp::ClauseProcBindKind kind) {
   llvm_unreachable("Unknown ClauseProcBindKind kind");
 }
 
+static void forwardArgs(LLVM::ModuleTranslation &moduleTranslation,
+                        llvm::ArrayRef<BlockArgument> blockArgs,
+                        OperandRange operands) {
+  for (auto [arg, var] : llvm::zip_equal(blockArgs, operands))
+    moduleTranslation.mapValue(arg, moduleTranslation.lookupValue(var));
+}
+
 /// Helper function to map block arguments defined by ignored loop wrappers to
 /// LLVM values and prevent any uses of those from triggering null pointer
 /// dereferences.
@@ -554,18 +561,12 @@ convertIgnoredWrapper(omp::LoopWrapperInterface opInst,
   // Map block arguments directly to the LLVM value associated to the
   // corresponding operand. This is semantically equivalent to this wrapper not
   // being present.
-  auto forwardArgs =
-      [&moduleTranslation](llvm::ArrayRef<BlockArgument> blockArgs,
-                           OperandRange operands) {
-        for (auto [arg, var] : llvm::zip_equal(blockArgs, operands))
-          moduleTranslation.mapValue(arg, moduleTranslation.lookupValue(var));
-      };
-
   return llvm::TypeSwitch<Operation *, LogicalResult>(opInst)
       .Case([&](omp::SimdOp op) {
         auto blockArgIface = cast<omp::BlockArgOpenMPOpInterface>(*op);
-        forwardArgs(blockArgIface.getPrivateBlockArgs(), op.getPrivateVars());
-        forwardArgs(blockArgIface.getReductionBlockArgs(),
+        forwardArgs(moduleTranslation, blockArgIface.getPrivateBlockArgs(),
+                    op.getPrivateVars());
+        forwardArgs(moduleTranslation, blockArgIface.getReductionBlockArgs(),
                     op.getReductionVars());
         op.emitWarning() << "simd information on composite construct discarded";
         return success();
@@ -5255,6 +5256,49 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
               return WalkResult::interrupt();
             return WalkResult::skip();
           }
+
+          // Non-target ops might nest target-related ops, therefore, we
+          // translate them as non-OpenMP scopes. Translating them is needed by
+          // nested target-related ops since they might LLVM values defined in
+          // their parent non-target ops.
+          if (isa<omp::OpenMPDialect>(oper->getDialect()) &&
+              oper->getParentOfType<LLVM::LLVMFuncOp>() &&
+              !oper->getRegions().empty()) {
+
+            // TODO Handle other ops with entry block args.
+            if (auto wsloop = dyn_cast<omp::WsloopOp>(oper)) {
+              auto blockArgIface = cast<omp::BlockArgOpenMPOpInterface>(oper);
+              if (blockArgIface) {
+                forwardArgs(moduleTranslation,
+                            blockArgIface.getPrivateBlockArgs(),
+                            wsloop.getPrivateVars());
+                forwardArgs(moduleTranslation,
+                            blockArgIface.getReductionBlockArgs(),
+                            wsloop.getReductionVars());
+              }
+            }
+
+            if (auto loopNest = dyn_cast<omp::LoopNestOp>(oper)) {
+              for (auto iv : loopNest.getIVs()) {
+                // Create fake allocas just to maintain IR validity.
+                moduleTranslation.mapValue(
+                    iv, builder.CreateAlloca(
+                            moduleTranslation.convertType(iv.getType())));
+              }
+            }
+
+            for (Region &region : oper->getRegions()) {
+              auto result = convertOmpOpRegions(
+                  region, oper->getName().getStringRef().str() + ".fake.region",
+                  builder, moduleTranslation);
+              if (result.takeError())
+                return WalkResult::interrupt();
+              builder.SetInsertPoint(result.get(), result.get()->end());
+            }
+
+            return WalkResult::skip();
+          }
+
           return WalkResult::advance();
         }).wasInterrupted();
   return failure(interrupted);