update

Author: skc7

flang/include/flang/Optimizer/Support/Utils.h

@@ -200,11 +200,6 @@ std::optional<llvm::ArrayRef<int64_t>> getComponentLowerBoundsIfNonDefault(
     fir::RecordType recordType, llvm::StringRef component,
     mlir::ModuleOp module, const mlir::SymbolTable *symbolTable = nullptr);
 
-// Convert FIR type to LLVM without turning fir.box<T> into memory
-// reference.
-mlir::Type convertObjectType(const fir::LLVMTypeConverter &converter,
-                             mlir::Type firType);
-
 /// Generate a LLVM constant value of type `ity`, using the provided offset.
 mlir::LLVM::ConstantOp
 genConstantIndex(mlir::Location loc, mlir::Type ity,

flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp

@@ -127,6 +127,13 @@ struct PrivateClauseOpConversion
   }
 };
 
+static mlir::Type convertObjectType(const fir::LLVMTypeConverter &converter,
+                                  mlir::Type firType) {
+  if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(firType))
+    return converter.convertBoxTypeAsStruct(boxTy);
+  return converter.convertType(firType);
+}
+
 // FIR Op specific conversion for TargetAllocMemOp
 struct TargetAllocMemOpConversion
     : public OpenMPFIROpConversion<mlir::omp::TargetAllocMemOp> {
@@ -139,7 +146,7 @@ struct TargetAllocMemOpConversion
     mlir::Location loc = allocmemOp.getLoc();
     auto ity = lowerTy().indexType();
     mlir::Type dataTy = fir::unwrapRefType(heapTy);
-    mlir::Type llvmObjectTy = fir::convertObjectType(lowerTy(), dataTy);
+    mlir::Type llvmObjectTy = convertObjectType(lowerTy(), dataTy);
     if (fir::isRecordWithTypeParameters(fir::unwrapSequenceType(dataTy)))
       TODO(loc, "omp.target_allocmem codegen of derived type with length "
                 "parameters");

flang/lib/Optimizer/OpenMP/LowerWorkdistribute.cpp

@@ -762,37 +762,19 @@ static void moveToHost(omp::TargetOp targetOp, RewriterBase &rewriter) {
     mapping.map(arg, mapInfoOp.getVarPtr());
   }
   rewriter.setInsertionPoint(targetOp);
-  SmallVector<Operation *> opsToMove;
+  SmallVector<Operation *> opsToReplace;
+  Value device = targetOp.getDevice();
+  /*
+  if (!device) {
+    device = genI32Constant(targetOp.getLoc(), rewriter, 0);
+  }
   for (auto it = targetBlock->begin(), end = std::prev(targetBlock->end());
        it != end; ++it) {
     auto *op = &*it;
-    auto allocOp = dyn_cast<fir::AllocMemOp>(op);
-    auto freeOp = dyn_cast<fir::FreeMemOp>(op);
     fir::CallOp runtimeCall = nullptr;
     if (isRuntimeCall(op))
       runtimeCall = cast<fir::CallOp>(op);
-
-    if (allocOp || freeOp || runtimeCall) {
-      Value device = targetOp.getDevice();
-      if (!device) {
-        device = genI32Constant(it->getLoc(), rewriter, 0);
-      }
-      if (allocOp) {
-        auto tmpAllocOp = rewriter.create<fir::OmpTargetAllocMemOp>(
-            allocOp.getLoc(), allocOp.getType(), device,
-            allocOp.getInTypeAttr(), allocOp.getUniqNameAttr(),
-            allocOp.getBindcNameAttr(), allocOp.getTypeparams(),
-            allocOp.getShape());
-        auto newAllocOp = cast<fir::OmpTargetAllocMemOp>(
-            rewriter.clone(*tmpAllocOp.getOperation(), mapping));
-        mapping.map(allocOp.getResult(), newAllocOp.getResult());
-        rewriter.eraseOp(tmpAllocOp);
-      } else if (freeOp) {
-        auto tmpFreeOp = rewriter.create<fir::OmpTargetFreeMemOp>(
-            freeOp.getLoc(), device, freeOp.getHeapref());
-        rewriter.clone(*tmpFreeOp.getOperation(), mapping);
-        rewriter.eraseOp(tmpFreeOp);
-      } else if (runtimeCall) {
+      if (runtimeCall) {
         auto module = runtimeCall->getParentOfType<ModuleOp>();
         auto callee = cast<func::FuncOp>(
             module.lookupSymbol(runtimeCall.getCalleeAttr()));
@@ -824,15 +806,42 @@ static void moveToHost(omp::TargetOp targetOp, RewriterBase &rewriter) {
         Operation *newCall = rewriter.clone(*tmpCall, mapping);
         mapping.map(&*it, newCall);
         rewriter.eraseOp(tmpCall);
-      }
     } else {
       Operation *clonedOp = rewriter.clone(*op, mapping);
+      auto allocOp = dyn_cast<fir::AllocMemOp>(clonedOp);
+      auto freeOp = dyn_cast<fir::FreeMemOp>(clonedOp);
+      if (allocOp || freeOp)
+        opsToReplace.push_back(clonedOp);
       for (unsigned i = 0; i < op->getNumResults(); ++i) {
         mapping.map(op->getResult(i), clonedOp->getResult(i));
       }
     }
   }
+  for (Operation* op : opsToReplace) {
+    if (auto allocOp = dyn_cast<fir::AllocMemOp>(op)) {
+      rewriter.setInsertionPoint(allocOp);
+      auto ompAllocmemOp = rewriter.create<omp::TargetAllocMemOp>(
+          allocOp.getLoc(), rewriter.getI64Type(), device,
+          allocOp.getInTypeAttr(), allocOp.getUniqNameAttr(),
+          allocOp.getBindcNameAttr(), allocOp.getTypeparams(),
+          allocOp.getShape());
+      auto firConvertOp = rewriter.create<fir::ConvertOp>(allocOp.getLoc(), 
+                            allocOp.getResult().getType(), ompAllocmemOp.getResult());
+      rewriter.replaceOp(allocOp, firConvertOp.getResult());
+    }
+    else if (auto freeOp = dyn_cast<fir::FreeMemOp>(op)) {
+      rewriter.setInsertionPoint(freeOp);
+      auto firConvertOp = rewriter.create<fir::ConvertOp>(
+        freeOp.getLoc(),
+        rewriter.getI64Type(),
+        freeOp.getHeapref());
+      rewriter.create<omp::TargetFreeMemOp>(
+          freeOp.getLoc(), device, firConvertOp.getResult());
+      rewriter.eraseOp(freeOp);
+    }
+  }
   rewriter.eraseOp(targetOp);
+  */
 }
 
 void fissionTarget(omp::TargetOp targetOp, RewriterBase &rewriter) {

flang/lib/Optimizer/Support/Utils.cpp

@@ -51,13 +51,6 @@ std::optional<llvm::ArrayRef<int64_t>> fir::getComponentLowerBoundsIfNonDefault(
   return std::nullopt;
 }
 
-mlir::Type fir::convertObjectType(const fir::LLVMTypeConverter &converter,
-                                  mlir::Type firType) {
-  if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(firType))
-    return converter.convertBoxTypeAsStruct(boxTy);
-  return converter.convertType(firType);
-}
-
 mlir::LLVM::ConstantOp
 fir::genConstantIndex(mlir::Location loc, mlir::Type ity,
                       mlir::ConversionPatternRewriter &rewriter,

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

@@ -2205,6 +2205,7 @@ def TargetFreeMemOp : OpenMP_Op<"target_freemem",
   Arg<I64, "", [MemFree]>:$heapref
   );
   let assemblyFormat = "$device `,` $heapref attr-dict `:` type($device) `,` qualified(type($heapref))";
+}
 
 //===----------------------------------------------------------------------===//
 // workdistribute Construct

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

@@ -3974,6 +3974,7 @@ llvm::LogicalResult omp::TargetAllocMemOp::verify() {
   if (!mlir::dyn_cast<IntegerType>(outType))
     return emitOpError("must be a integer type");
   return mlir::success();
+}
 
 //===----------------------------------------------------------------------===//
 // WorkdistributeOp