[MLIR][OpenMP] Use getTargetEntryUniqueInfo from IRBuilder and fix an issue when using target op with --split-input-file option

[agozillon]

Summary

This patch seeks to swap the getTargetEntryUniqueInfo
function to the shared OMPIRBuilder function.

As well as fix an issue with getting the targetEntryUniqueInfo
when mlir-translate's --split-input-file is used and target op's
are used within the file. Essentially filename has some extra
prefix and suffix information that needs to be filtered out to
get the "cannonical" filename.

[kiranchandramohan]

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

1635–1638

We wouldn't want to do something like this. Could you provide some more info or a reproducer for this issue?

[agozillon]

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

1635–1638

It's fairly easy to recreate I believe, place a TargetOp in an mlir file and mlir-translate it to llvm-ir using -split-input-file. The end results are hitting the "Unable to get unique ID for file" error. This is due to the fact that the mlir tooling prefixes and suffixes the filename of FileLineColLoc with the extra information I am trying to remove here: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Support/ToolUtilities.cpp#L85 and the getUniqueID invocation requires a filename without the extras.

The ways I can currently see of doing it are:

1. find an alternative source for the filename, unsure where that could be from, so if you or anyone else has suggestions that'd be awesome, I think @domada is working on a patch to pass an alternative filename down (for another reason, but it will likely get used in place of FileLineColLoc) that will likely help, so perhaps it's better to just wait until that lands
2. find an alternative method to calculate the unique id portion of the entry info, I believe @jsjodin suggested a checksum of the file contents, this might still require the filename
3. Alter the mlir-translate tool or FileLineColLoc to give an alternative filename via a different method that does not include this extra information

Perhaps I am missing an alternative much easier method though, so if anyone has another idea please do suggest it! Of the ones I have listed I am leaning towards 1) if @domada believes the patches he's working on are applicable to this issue like I think they are!

The work-around is to just not use TargetOp's in -split-input-file tests until the issue is resolved.

[kiranchandramohan]

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

1635–1638

The end results are hitting the "Unable to get unique ID for file" error.

Would you know whether this is due to the mangled name created by MLIR tooling not being a real file that exists?

[kiranchandramohan]

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

1635–1638

Would getting the Filename from the module (getParentOfType<ModuleOp>, rather than the function) work?

[agozillon]

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

1635–1638

That is possible but I am unfortunately not entirely sure, I believe the tooling is just creating seperate named memory buffers to contain each segment, and this name then proliferates to the hidden FileLineColLoc information inside of the module (visible via printing the MLIR debug info), perhaps they reside somewhere as temporary files though, but I am doubtful.

Someone with more familiarity with how the tooling interconnects may be able to give a better picture on this, I'm only really sure the name is being modified by the tooling as it splits the file apart unfortunately.

1635–1638

I had the same thought, unfortunately it appears not, it contains the same modified Filename inside of its FileLineColLoc (perhaps there is an alternative filename that you can get from a Module that I am unaware of though?). Which makes sense as I believe the MLIR tooling will wrap all the subsections inside of their own implicit ModuleOp if it's not already contained in one, rather than having one singular ModuleOp.

1635–1638

I talked a little with @domada last week, it unfortunately doesn't seem like his work will be able to fix this sadly.

Which leaves us with the following options (that I can see at least with my limited knowledge, I am more than open to other suggestions):

• Passing down an alternative filename as an attribute from the frontend similar to what I believe was proposed previously and rejected. Seems the easiest and less prone to breakage due to a change in MLIR infrastructure.
• Calculate the UniqueID/EntryInfo differently, I'm not really sure how we'd go about doing this if we can't access the contents of the file, perhaps use the loaded in MLIR module data in someway to unique the information. I'm also not a huge fan of doing this differently from how Clang currently does it, at least while we're still glue-ing things together.
• Use this change for the time being on the expectation that the MLIR infrastructure doesn't change this one segment

But yes, if anyone has any other suggestions or opinions on the direction for this fix, I would greatly appreciate it.

[kiranchandramohan]

@agozillon Could you make a quick discourse post about this in MLIR discourse?

[agozillon]

In D153781#4471209, @kiranchandramohan wrote:

@agozillon Could you make a quick discourse post about this in MLIR discourse?

No problem! I've created one here: https://discourse.llvm.org/t/request-for-input-on-a-fix-for-a-bug-utilizing-omp-targetop-in-conjunction-with-mlir-translates-split-input-file/71785 hopefully it's written up reasonably, I am terrible at creating titles unfortunately.

[agozillon]

Closing this patch as they approach isn't ideal, I'll look into a solution in the future when there's a little more free time on my TODO list schedule and open up a github PR, it's unfortunately a little far down the list as this can be worked around fairly trivially for the moment in a couple of ways.

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

ompBuilder->createGlobalFlag(
           .getAssumeNoNestedParallelism() /*LangOpts().OpenMPNoNestedParallelism*/
       ,
       "__omp_rtl_assume_no_nested_parallelism");
   ompBuilder->M.addModuleFlag(llvm::Module::Max, "openmp-device",
                               attribute.getOpenmpDeviceVersion());
   return success();
 }
 
-static bool getTargetEntryUniqueInfo(llvm::TargetRegionEntryInfo &targetInfo,
-                                     omp::TargetOp targetOp,
-                                     llvm::StringRef parentName = "") {
-  auto fileLoc = targetOp.getLoc()->findInstanceOf<FileLineColLoc>();
-
-  assert(fileLoc && "No file found from location");
-  StringRef fileName = fileLoc.getFilename().getValue();
-
-  llvm::sys::fs::UniqueID id;
-  if (auto ec = llvm::sys::fs::getUniqueID(fileName, id)) {
-    targetOp.emitError("Unable to get unique ID for file");
-    return false;
-  }
-
-  uint64_t line = fileLoc.getLine();
-  targetInfo = llvm::TargetRegionEntryInfo(parentName, id.getDevice(),
-                                           id.getFile(), line);
-  return true;
-}
-
 static bool targetOpSupported(Operation &opInst) {
   auto targetOp = cast<omp::TargetOp>(opInst);
   if (targetOp.getIfExpr()) {
     opInst.emitError("If clause not yet supported");
     return false;
   }
 
   if (targetOp.getDevice()) {
auto bodyCB = [&](InsertPointTy allocaIP,
     builder.restoreIP(codeGenIP);
     llvm::BasicBlock *exitBlock = convertOmpOpRegions(
         targetRegion, "omp.target", builder, moduleTranslation, bodyGenStatus);
     builder.SetInsertPoint(exitBlock);
     return builder.saveIP();
   };
 
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
-  StringRef parentName = opInst.getParentOfType<LLVM::LLVMFuncOp>().getName();
-  llvm::TargetRegionEntryInfo entryInfo;
 
-  if (!getTargetEntryUniqueInfo(entryInfo, targetOp, parentName))
-    return failure();
+  auto fileInfoCallBack = [&]() {
+    auto loc = opInst.getParentOfType<LLVM::LLVMFuncOp>()
+                   ->getLoc()
+                   ->findInstanceOf<FileLineColLoc>();
+    llvm::StringRef filename = loc.getFilename();
+    // handle split file from mlir-translate that for some reason
+    // manipulates the string to contain other components than
+    // filename
+    if (filename.contains("within split at")) {
+      filename.consume_front("within split at ");
+      filename = filename.rsplit(":").first;
+    }
+
+    return std::pair<std::string, uint64_t>(filename, loc.getLine());
+  };
+
+  llvm::TargetRegionEntryInfo entryInfo =
+      moduleTranslation.getOpenMPBuilder()->getTargetEntryUniqueInfo(
+          fileInfoCallBack,
+          opInst.getParentOfType<LLVM::LLVMFuncOp>().getName());
 
   int32_t defaultValTeams = -1;
   int32_t defaultValThreads = -1;
 
   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createTarget(
       ompLoc, builder.saveIP(), entryInfo, defaultValTeams, defaultValThreads,
       inputs, bodyCB));