[Flang][OpenMP][MLIR] Filter emitted code depending on declare target and device

[skatrak]

Summary

This patch adds support for selecting which functions are lowered to LLVM IR from MLIR depending on declare target information and whether host or device code is being generated.

The approach proposed by this patch is to perform the filtering in two stages:

• An MLIR transformation pass, which is added to the Flang translation flow before the VerifierPass. The functions that are kept are those that match the OpenMP processor (host or device) the compiler invocation is targeting, according to the presence of the -fopenmp-is-device compiler option and declare target information. All functions contaning an omp.target are also kept, regardless of the declare target information of the function, due to the need for keeping target regions visible for both host and device compilation.
• A filtering step during translation to LLVM IR, which is peformed for those functions that were kept because of the presence of a target region inside. If the targeted OpenMP processor does not match the declare target information of the function, then it is removed from the LLVM IR after its contents have been processed and translated. Since they should only contain an omp.target operation which, in turn, should have been outlined into another LLVM function, the wrapper can be deleted at that point.

Depends on D150328 and D150329.

[skatrak]

Update and add test.

[skatrak]

Rebase and fix tests

[jdoerfert]

generating functions and then deleting them is costly and will likely not work. We need to not emit them in the first place.

[skatrak]

In D147641#4313262, @jdoerfert wrote:

generating functions and then deleting them is costly and will likely not work. We need to not emit them in the first place.

I agree that we should strive to avoid doing unnecessary work rather than discarding it after the fact. However, the current implementation of target region outlining interacts with this in a way that prevents us from doing that in all cases. Let's say we implemented this function filtering patch as a semantic analysis pass:

• We could forward the information from the frontend about whether we're compiling for the host or device (-fopenmp-is-device) to this new pass.
• We could add the pass after the ImplicitDeclareTargetCapture pass implemented in D146063 or similar, so implicit declare target information is already propagated and available.
• Using the knowledge about whether functions are only intended for the host, device or both, and what we're compiling for, we could remove from the parse tree those functions that won't need lowering.
• For the host that would mean removing declare target device_type(nohost) functions, and for the device it would mean removing declare target device_type(host) and non-declare target functions.

However, the issues come when we consider the interaction of that approach with target regions:

• Target regions (which should be lowered for the device) can appear inside of host functions.
• Reverse-offload target regions (which should eventually be lowered for the host) can appear inside of device functions and target regions.
• The first two points mean that, while generating host code, you must also look for target regions inside device functions and vice versa.
• Target region outlining currently happens very late, at the MLIR to LLVM IR translation stage, during the lowering of their parent function, when the related omp.target MLIR operation is visited. D147172 was the main patch implementing this, and D147940 is a proposal to extend it for the device.
• If we remove a function before that point due to being intended for another device, then there won't be any outlining of the target regions defined inside of it either.

I've got a couple of ideas about to handle this problem, but I'm not particularly convinced about any of them:

• We could relax the requirement of the hypothetical semantics pass to avoid discarding too early functions that contain target regions, so that they can be outlined later. But then they will have to be pruned after translating them from MLIR to LLVM IR, like this patch is doing. We would be implementing the same kind of idea in two different places so that we could reduce the number of functions that we unnecessarily lower. Perhaps it would also be possible to reduce further the work done on these functions by only lowering the target regions inside (perhaps by hijacking ModuleTranslation::convertBlock), and not the rest of operations in them.
• Another option could be to try to move part of the outlining work into a semantics pass as well, so that we wouldn't have to keep around the parent functions intended for the other device. We would have to do all of the used variable capturing there (which it's not currently implemented AFAIK), so that they can be part of the new function as well. The result of this would be that at the MLIR level there would be some compiler-generated functions with just an omp.target operation inside (not called from anywhere else), together with the rest of functions in the module that apply to the device we're compiling for (possibly containing other omp.target operations that should actually be implemented by the opposite device). We would have to mark these compiler-generated functions somehow to make sure that they are treated differently during MLIR to LLVM IR translation -- for compiler-generated functions, only the contents of the omp.target operation inside should be lowered, whereas for other functions everything should be lowered and any occurrences of that operation should become a kernel call or a reverse-offload call, ignoring the region attached to them.

The implementation in this patch, although not very efficient, is a very simple first approach that doesn't require significantly reworking different parts of offloading support. I'd like to know what you think about this, or if @jsjodin, @domada, @agozillon, @kiranchandramohan or others have any thoughts about it as well. It'd be good to reach some sort of consensus before committing to significant work.

[jsjodin]

I don't think moving the outlining early on is the best approach since we want this code to be shared with clang. Would it be sufficient to filter the functions that do not contain target regions after the analysis in https://reviews.llvm.org/D146063 somewhere in the front-end and let the MLIR->LLVM IR translation handle the functions that has target regions?

[skatrak]

Ping @jdoerfert. Do you have any other thoughts different to what Jan's suggesting? Just looking to make sure we're in the same page to avoid redoing the work multiple times.

[skatrak]

Rebase and update to depend on declare target work split into multiple patches.

[jsjodin]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

1035

The maskedFunction information lives only inside this function. Can we make it in to a local variable and pass it in to convertOneFunction?

[skatrak]

Ping again in case there any opinions against following Jan's suggestion before starting its implementation. In summary, the idea would be to:

• Add a semantics pass to filter out functions that are not intended for the device we're compiling, according to implicit/explicit declare target information, except for those with a target region inside.
• Keep also the approach in this first patch, so that the IR for the enclosing functions where target regions are located is deleted after the target regions themselves are outlined by the OpenMPIRBuilder.

I am thinking that there could be situations where this approach might break. Let's say there's a host-only subroutine f() called by a host-only subroutine g() that contains a target region as well (f() not called from inside the target region). Then, during compilation for the device, f() would be removed in semantics whereas g() would be kept due to it having a target region to be processed later during MLIR to LLVM IR translation. What would happen then with the call to the f() subroutine that would have become undefined?

Should we just replace the body of functions to be deleted in semantics by an empty body or a default 'return' statement instead? Or would it maybe be a preferable alternative to delete these functions altogether but also delete any calls to them? The first option could become difficult if these subroutines/functions return complex data types, and also we may end up generating many unused functions. The only issue I see with the second option is what would happen if function calls are used as part of an expression to initialize some data. I'm leaning towards just deleting all statements except variable declarations and target regions from functions intended for another device, which seems like solves all of the issues above.

Let me know if you have any more thoughts on this, and thank you for taking the time.

[jsjodin]

In D147641#4413305, @skatrak wrote:

Ping again in case there any opinions against following Jan's suggestion before starting its implementation. In summary, the idea would be to:

• Add a semantics pass to filter out functions that are not intended for the device we're compiling, according to implicit/explicit declare target information, except for those with a target region inside.
• Keep also the approach in this first patch, so that the IR for the enclosing functions where target regions are located is deleted after the target regions themselves are outlined by the OpenMPIRBuilder.

I am thinking that there could be situations where this approach might break. Let's say there's a host-only subroutine f() called by a host-only subroutine g() that contains a target region as well (f() not called from inside the target region). Then, during compilation for the device, f() would be removed in semantics whereas g() would be kept due to it having a target region to be processed later during MLIR to LLVM IR translation. What would happen then with the call to the f() subroutine that would have become undefined?

Should we just replace the body of functions to be deleted in semantics by an empty body or a default 'return' statement instead? Or would it maybe be a preferable alternative to delete these functions altogether but also delete any calls to them? The first option could become difficult if these subroutines/functions return complex data types, and also we may end up generating many unused functions. The only issue I see with the second option is what would happen if function calls are used as part of an expression to initialize some data. I'm leaning towards just deleting all statements except variable declarations and target regions from functions intended for another device, which seems like solves all of the issues above.

Let me know if you have any more thoughts on this, and thank you for taking the time.

Can we add declarations for them instead of having empty bodies?

[skatrak]

In D147641#4417720, @jsjodin wrote:

Can we add declarations for them instead of having empty bodies?

Since semantics work at the PFT level, I've looked into how to define the Fortran-equivalent to forward declarations. It seems like external procedures could be one solution to this. So it should be possible to delete all function and subroutine definitions intended for another device and add an external declaration to all their callers. These calls and external procedure declarations will remain in MLIR but they will be removed in the MLIR to LLVM IR translation step.

[skatrak]

Implement early function filtering as an MLIR pass.

[skatrak]

Update after splitting off pass infrastructure into D154194.

[skatrak]

Move MLIR transformation pass to Flang and implement 2-stage filtering to work together with early target region outlining.

[jsjodin]

flang/lib/Frontend/FrontendActions.cpp

702 ↗

(On Diff #538106)

I don't think this is needed. The FIR should already have been filtered by the front end.

flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp

48 ↗

(On Diff #538106)

This could potentially be simplified if the early outlining is run before this pass. I think this is still okay though since the first instruction in the outlined functions is a TargetOp so it should be efficient, and it still works without the early outlining. So I think it is okay to keep as-is.

[skatrak]

Address review comments and integrate with recently-landed early target region outlining. Small refactor to improve readability and minimize TargetOp search overhead.

[skatrak]

Move misplaced function.

[skatrak]

Thank you for your feedback Jan, your comments should have been addressed.

flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp

48 ↗

(On Diff #538106)

I thought about this, but I think it's still necessary to check here whether there are omp.target ops inside because the outlined functions are host functions in the device pass. So they would by default be filtered out here and the omp.target inside wouldn't reach the MLIR to LLVM IR translation stage.

But you're right. By running this after early outlining, the check should be very quick.

[jsjodin]

LGTM, Please wait for another acceptance.

flang/lib/Optimizer/Transforms/OMPFunctionFiltering.cpp

48 ↗

(On Diff #538106)

It would be possible to check the early outlining interface to see if the parent function name is set, but that solution is less desirable imo.

[kiranchandramohan]

See comment inline.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

906–916

The change in convertOneFunction is intrusive. Code here must not be aware of the OpenMP attributes.

[jsjodin]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

906–916

The change in convertOneFunction is intrusive. Code here must not be aware of the OpenMP attributes.

Would it be ensure that the omp.declare_target dialect attribute is available so that convertDialectAttributes can handle everything?

[jsjodin]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

906–916

The change in convertOneFunction is intrusive. Code here must not be aware of the OpenMP attributes.

Would it be ensure that the omp.declare_target dialect attribute is available so that convertDialectAttributes can handle everything?

906–916

The change in convertOneFunction is intrusive. Code here must not be aware of the OpenMP attributes.

Would it be ensure that the omp.declare_target dialect attribute is available so that convertDialectAttributes can handle everything?

Correction: Would it be possible to ensure that the omp.declare_target dialect attribute is available so that convertDialectAttributes can handle everything?

[skatrak]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

906–916

The change in convertOneFunction is intrusive. Code here must not be aware of the OpenMP attributes.

The problem here is that the amendOperation flow is called if there is an omp.declare_target attribute explicitly set for the function. However, functions without the attribute must also be treated as declare target with device type = host, so they must be filtered out during device compilation.

So, the options I see are either to access the attribute directly here or ensure that all functions have the attribute, either by setting it here or at an earlier stage. Though it might be enough to mark outlined functions with the attribute, which should be the only ones that will be removed at this stage.

[jsjodin]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

906–916

The change in convertOneFunction is intrusive. Code here must not be aware of the OpenMP attributes.

The problem here is that the amendOperation flow is called if there is an omp.declare_target attribute explicitly set for the function. However, functions without the attribute must also be treated as declare target with device type = host, so they must be filtered out during device compilation.

So, the options I see are either to access the attribute directly here or ensure that all functions have the attribute, either by setting it here or at an earlier stage. Though it might be enough to mark outlined functions with the attribute, which should be the only ones that will be removed at this stage.

I think the best option is to modify the early outlining to set the attribute. You can add that to this patch and a test if the current tests don't already cover this.

[agozillon]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

906–916

The type of the attribute is just a boolean though, so it has no dependencies on the OpenMPDialect, we already use something similar inside of the getOpenMPBuilder call: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp#L1311 Is it possible to do a similar segment of code inside of the destructor of ModuleTranslation, where we already have some OpenMP specific code (destruction of the OMPIRBuilder)? And is that a more acceptable approach?

[skatrak]

Remove handling of OpenMP dialect attributes from generic translation. Not needed because the early outlining pass already adds the omp.declare_target attribute to the new function, so the second-stage filter can assume it will be always explicitly set.

[kiranchandramohan]

Thanks @skatrak.

[skatrak]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

906–916

After some testing it turns out that, in this specific case, we can assume that the attribute will be explicitly set by the early outlining. So there is no need to check for implicit host functions here.

The type of the attribute is just a boolean though, so it has no dependencies on the OpenMPDialect, we already use something similar inside of the getOpenMPBuilder call: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp#L1311 Is it possible to do a similar segment of code inside of the destructor of ModuleTranslation, where we already have some OpenMP specific code (destruction of the OMPIRBuilder)? And is that a more acceptable approach?

I think that is also something we should try to remove from this file if we can. In D147219 I implemented an approach that moves the OMPIRBuilder configuration to the amendOperation flow as well, which has the advantage of allowing the use of OpenMP custom attribute types. It's not been accepted, but I think it's a better way to go about that.

[agozillon]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

906–916

While I'm all for doing that, I don't think moving it into amendOperation would work at this time, as the attributes are tied to the ModuleOp and I believe the amendOperation for the ModuleOp is ran as one of the final things in the lowering process. The information in the Config for IRBuilder is required very early on, before lowering global ops (for declare target) which happens prior to even the function lowering which is when TargetOp will currently need it. I could be incorrect though!

[skatrak]

Thank you all for the comments, I'll land this patch shortly.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

906–916

Yes, I noticed that issue as well. If you look at that patch, you can see that I moved the convertOperation() call to the module inside of translateModuleToLLVMIR() to be before translating the functions inside. Doing it in that order means that the OpenMPIRBuilder is initialized before translating the contents of the module. I believe that is a good approach to avoid having dialect-specific attributes handling in the generic ModuleTranslation, since (at least at this time) there is no particular reason why the module should be converted/amended last.

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

private:
   DenseMap<const Operation *, llvm::MDNode *> tbaaMetadataMapping;
 
   /// Stack of user-specified state elements, useful when translating operations
   /// with regions.
   SmallVector<std::unique_ptr<StackFrame>> stack;
 
   /// A cache for the symbol tables constructed during symbols lookup.
   SymbolTableCollection symbolTableCollection;
+
+  /// The set of functions that should be masked out from the output, due to not
+  /// corresponding to the target device being compiled for.
+  llvm::SmallVector<llvm::Function *> maskedFunctions;
 };
 
 namespace detail {
 /// For all blocks in the region that were converted to LLVM IR using the given
 /// ModuleTranslation, connect the PHI nodes of the corresponding LLVM IR blocks
 /// to the results of preceding blocks.
 void connectPHINodes(Region &region, const ModuleTranslation &state);
 

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

for (Block *bb : blocks) {
     if (failed(convertBlock(*bb, bb->isEntryBlock(), builder)))
       return failure();
   }
 
   // After all blocks have been traversed and values mapped, connect the PHI
   // nodes to the results of preceding blocks.
   detail::connectPHINodes(func.getBody(), *this);
 
   // Finally, convert dialect attributes attached to the function.
-  return convertDialectAttributes(func);
+  LogicalResult result = convertDialectAttributes(func);
+
+  // All functions are translated first to ensure target regions are always
+  // processed, so that they can be outlined. However, they must be deleted
+  // afterwards if the device they are intended for does not match the device we
+  // are currently generating code for.
+  if (auto offloadMod =
+          dyn_cast<mlir::omp::OffloadModuleInterface>(mlirModule)) {
+    bool isDevicePass = offloadMod.getIsDevice();
+    omp::DeclareTargetDeviceType declareType =
+        omp::DeclareTargetDeviceType::host;
+
+    if (omp::OpenMPDialect::isDeclareTarget(func.getOperation()))
+      declareType =
+          omp::OpenMPDialect::getDeclareTargetDeviceType(func.getOperation());
+
+    if ((isDevicePass && declareType == omp::DeclareTargetDeviceType::host) ||
+        (!isDevicePass && declareType == omp::DeclareTargetDeviceType::nohost))
+      maskedFunctions.push_back(llvmFunc);
+  }
+
+  return result;
 }
 
 LogicalResult ModuleTranslation::convertDialectAttributes(Operation *op) {
   for (NamedAttribute attribute : op->getDialectAttrs())
     if (failed(iface.amendOperation(op, attribute, *this)))
       return failure();
   return success();
 }
for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     // Convert visibility attribute.
     llvmFunc->setVisibility(convertVisibilityToLLVM(function.getVisibility_()));
   }
 
   return success();
 }
 
 LogicalResult ModuleTranslation::convertFunctions() {
+  maskedFunctions.clear();
+
   // Convert functions.
   for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     // Ignore external functions.
     if (function.isExternal())
       continue;
 
     if (failed(convertOneFunction(function)))
       return failure();
   }
 
+  // Delete translated functions that are intended for a device we are currently
+  // not generating code for.
+  for (auto *llvmFunction : llvm::reverse(maskedFunctions))
+    llvmFunction->eraseFromParent();
+
   return success();
 }
 
 LogicalResult ModuleTranslation::createAccessGroupMetadata() {
   return loopAnnotationTranslation->createAccessGroupMetadata();
 }
 
 void ModuleTranslation::setAccessGroupsMetadata(AccessGroupOpInterface op,

mlir/test/Target/LLVMIR/openmp-llvm.mlir

 // -----
 
 // CHECK: @__omp_rtl_debug_kind = weak_odr hidden constant i32 0
 // CHECK: @__omp_rtl_assume_teams_oversubscription = weak_odr hidden constant i32 1
 // CHECK: @__omp_rtl_assume_threads_oversubscription = weak_odr hidden constant i32 0
 // CHECK: @__omp_rtl_assume_no_thread_state = weak_odr hidden constant i32 1
 // CHECK: @__omp_rtl_assume_no_nested_parallelism = weak_odr hidden constant i32 0
 module attributes {omp.flags = #omp.flags<assume_teams_oversubscription = true, assume_no_thread_state = true>} {}
+
+// -----
+
+// CHECK:     define void @any
+// CHECK:     define void @nohost
+// CHECK-NOT: define void @host
+// CHECK-NOT: define void @no_declare_target
+module attributes {omp.is_device = #omp.isdevice<is_device = true>} {
+  llvm.func @any() -> ()
+      attributes {
+        omp.declare_target =
+          #omp.declaretarget<device_type = (any), capture_clause = (to)>
+      } {
+    llvm.return
+  }
+  llvm.func @nohost() -> ()
+      attributes {
+        omp.declare_target =
+          #omp.declaretarget<device_type = (nohost), capture_clause = (to)>
+      } {
+    llvm.call @any() : () -> ()
+    llvm.return
+  }
+  llvm.func @host() -> ()
+      attributes {
+        omp.declare_target =
+          #omp.declaretarget<device_type = (host), capture_clause = (to)>
+      } {
+    llvm.call @any() : () -> ()
+    llvm.return
+  }
+  llvm.func @no_declare_target() -> () {
+    llvm.call @host() : () -> ()
+    llvm.return
+  }
+}
+
+// -----
+
+// CHECK:     define void @any
+// CHECK-NOT: define void @nohost
+// CHECK:     define void @host
+// CHECK:     define void @no_declare_target
+module attributes {omp.is_device = #omp.isdevice<is_device = false>} {
+  llvm.func @any() -> ()
+      attributes {
+        omp.declare_target =
+          #omp.declaretarget<device_type = (any), capture_clause = (to)>
+      } {
+    llvm.return
+  }
+  llvm.func @nohost() -> ()
+      attributes {
+          omp.declare_target =
+            #omp.declaretarget<device_type = (nohost), capture_clause = (to)>
+      } {
+    llvm.call @any() : () -> ()
+    llvm.return
+  }
+  llvm.func @host() -> ()
+      attributes {
+        omp.declare_target =
+          #omp.declaretarget<device_type = (host), capture_clause = (to)>
+      } {
+    llvm.call @any() : () -> ()
+    llvm.return
+  }
+  llvm.func @no_declare_target() -> () {
+    llvm.call @host() : () -> ()
+    llvm.return
+  }
+}