fix failures seen in target-private-multiple-variables.f90

Author: bhandarkar-pranav

flang/lib/Optimizer/OpenMP/MapsForPrivatizedSymbols.cpp

@@ -7,6 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 #include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Optimizer/OpenMP/Passes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
@@ -41,10 +42,27 @@ class MapsForPrivatizedSymbolsPass
     uint64_t mapTypeTo = static_cast<
         std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
         llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO);
-
+    Operation *definingOp = var.getDefiningOp();
+    auto declOp = llvm::dyn_cast_or_null<hlfir::DeclareOp>(definingOp);
+    assert(declOp &&
+           "Expected defining Op of privatized var to be hlfir.declare");
+    Value varPtr = declOp.getOriginalBase();
+    Value varBase = declOp.getBase();
+    llvm::errs() << "varPtr = ";
+    varPtr.dump();
+    llvm::errs() << " type -> ";
+    varPtr.getType().dump();
+    llvm::errs() << "\n";
+    llvm::errs() << "varBase = ";
+    varBase.dump();
+    llvm::errs() << " type -> ";
+    varBase.getType().dump();
+    llvm::errs() << "\n";
     return builder.create<omp::MapInfoOp>(
-        loc, var.getType(), var,
-        mlir::TypeAttr::get(fir::unwrapRefType(var.getType())),
+        loc, varPtr.getType(), varPtr,
+        mlir::TypeAttr::get(
+            llvm::cast<mlir::omp::PointerLikeType>(varPtr.getType())
+                .getElementType()),
         /*varPtrPtr=*/mlir::Value{},
         /*members=*/mlir::SmallVector<mlir::Value>{},
         /*member_index=*/mlir::DenseIntElementsAttr{},
@@ -57,41 +75,67 @@ class MapsForPrivatizedSymbolsPass
   }
   void addMapInfoOp(omp::TargetOp targetOp, omp::MapInfoOp mapInfoOp) {
     mlir::Location loc = targetOp.getLoc();
+    llvm::errs() << "Adding mapInfoOp -> ";
+    mapInfoOp.dump();
+    llvm::errs() << "\n";
+
     targetOp.getMapVarsMutable().append(mlir::ValueRange{mapInfoOp});
     size_t numMapVars = targetOp.getMapVars().size();
     targetOp.getRegion().insertArgument(numMapVars - 1, mapInfoOp.getType(),
                                         loc);
   }
+  void addMapInfoOps(omp::TargetOp targetOp,
+                     llvm::SmallVectorImpl<omp::MapInfoOp> &mapInfoOps) {
+    for (auto mapInfoOp : mapInfoOps)
+      addMapInfoOp(targetOp, mapInfoOp);
+  }
   void runOnOperation() override {
     MLIRContext *context = &getContext();
     OpBuilder builder(context);
+    llvm::DenseMap<Operation *, llvm::SmallVector<omp::MapInfoOp, 4>>
+        mapInfoOpsForTarget;
     getOperation()->walk([&](omp::TargetOp targetOp) {
       if (targetOp.getPrivateVars().empty())
         return;
-
+      llvm::errs() << "Func is \n";
+      targetOp.getOperation()->getParentOp()->getParentOp()->dump();
+      llvm::errs() << "\n";
       OperandRange privVars = targetOp.getPrivateVars();
       std::optional<ArrayAttr> privSyms = targetOp.getPrivateSyms();
-
+      SmallVector<omp::MapInfoOp, 4> mapInfoOps;
       for (auto [privVar, privSym] : llvm::zip_equal(privVars, *privSyms)) {
 
         SymbolRefAttr privatizerName = llvm::cast<SymbolRefAttr>(privSym);
         omp::PrivateClauseOp privatizer =
             SymbolTable::lookupNearestSymbolFrom<omp::PrivateClauseOp>(
                 targetOp, privatizerName);
-
-        assert(mlir::isa<fir::ReferenceType>(privVar.getType()) &&
-               "Privatized variable should be a reference.");
+        llvm::errs() << "privVar = ";
+        privVar.dump();
+        llvm::errs() << "\n";
+        llvm::errs() << "privVar.getType() = ";
+        privVar.getType().dump();
+        llvm::errs() << "\n";
+        // assert(mlir::isa<fir::ReferenceType>(privVar.getType()) &&
+        //        "Privatized variable should be a reference.");
         if (!privatizerNeedsMap(privatizer)) {
-          return;
+          continue;
         }
         builder.setInsertionPoint(targetOp);
         mlir::Location loc = targetOp.getLoc();
         omp::MapInfoOp mapInfoOp = createMapInfo(loc, privVar, builder);
-        addMapInfoOp(targetOp, mapInfoOp);
+        mapInfoOps.push_back(mapInfoOp);
         LLVM_DEBUG(llvm::dbgs() << "MapsForPrivatizedSymbolsPass created ->\n");
         LLVM_DEBUG(mapInfoOp.dump());
       }
+      if (!mapInfoOps.empty()) {
+        mapInfoOpsForTarget.insert({targetOp.getOperation(), mapInfoOps});
+      }
     });
+    if (!mapInfoOpsForTarget.empty()) {
+      for (auto &[targetOp, mapInfoOps] : mapInfoOpsForTarget) {
+        addMapInfoOps(static_cast<omp::TargetOp>(targetOp), mapInfoOps);
+      }
+    }
   }
 };
 } // namespace