[OpenMP][MLIR] Add lowering of omp::FlagsAttr to LLVM-IR

[agozillon]

Summary

The omp::FlagsAttr contains OpenMP RTL flags
given by a user to the compiler and a frontend
(flang currently) then populates the omp::FlagsAttr
which will then lower these to LLVM globals
which are utilised by the OpenMP runtime.

[agozillon]

Patch 2/2 after splitting https://reviews.llvm.org/D145264 this patch provides the lowering of the previously added omp::FlagsAttr

[agozillon]

Hi would it please be possible to have a reviewer take a look at this patch whenever they can spare a moment. This is the lowering segment for the FlagsAttr attribute which is currently upstream, I believe this is the segment of the broken up patch https://reviews.llvm.org/D145264 that needs most review as it's not had anyone look at it yet. However, it's also the shortest!

Primarily just tests with the addition of the amendOperation infrastructure to the OpenMPToLLVMIRTranslation.cpp infrastructure which is used to lower FlagsAttr to globals via the OpenMP IR builder.

Thank you very much in advance for any reviewers time.

[agozillon]

• [Flang][Test] Remove LLVM tests from RTL-Flags doesn't seem to be the correct place

[agozillon]

Removed the llvm lowering specific tests from the flags attribute test, as tests similar to these are done in in openmp-llvm.mlir and it's a more appropriate location.

The tests should now pass as everything required is upstream or in this patch! There is possibly going to be a single mlir test failure in the phabricator buildbot check, however, it's an unrelated failing test I picked up from the current commit I've rebased the patch on.

[agozillon]

Just a small ping for reviewer attention please, if at all possible, to push this patch a little further along the pipeline!

[agozillon]

Just a small ping to keep this patch on peoples radar and to try and get some reviewer attention (I am aware you're all quite busy, so thank you very much for any time you can spare) for more review feedback to push this patch further along the pipeline!

[jdoerfert]

LG

[agozillon]

@jdoerfert Thank you very much!

[kiranchandramohan]

See comment inline.

Also, I am a bit confused about when we are using the amendOperation translation vs setting the OpenMPIRBuilder config (with data from the OffloadModuleInterface). @agozillon @skatrak @jsjodin

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

1548

Is the operation unused here?

I think it will be good to add a return failure() if the operation is not a module.

[agozillon]

In D147344#4299066, @kiranchandramohan wrote:

See comment inline.

Also, I am a bit confused about when we are using the amendOperation translation vs setting the OpenMPIRBuilder config (with data from the OffloadModuleInterface). @agozillon @skatrak @jsjodin

I'm not sure what you mean in this case (I am probably being dumb and missing your point, so I apologies), this runtime flag generation is seperate from the config file, the only thing this does is utilise the OMPIRBuilder to create these globals in the exact same way Clang does at the moment. The attribute itself contains the values that in Clang are frontend language options, so can be generated directly from them. In Flangs and the OpenMP dialects case, we need an attribute to contain and carry them down. This particular amendOperation, will be called when the Module operation is processed and only if the mlir::omp::FlagsAttr is present on the module, it is only present on the module if Flang is called with -fopenmp-is-device.

Perhaps the lowering itself could be pushed into the OMPIRBuilder's finalize method, by filling the OMPIRBuilder's config file with these flags during config setup, however, I am unsure how that would impact Clang at the moment.

As an probably unneeded aside the IR Config comes into play in other functions, such as getAddrOfDeclareTargetVar & registerTargetGlobalVariable: https://reviews.llvm.org/D149162 which need some toggles in both Clang and Flang to specify if it's SIMD/generating for target etc. and flags carried down from Flang such as isDevice help to populate this config to generate the IR! It also helps to have the config for the builder to have some state, so as not to not have an excessive amount of arguments passed in to every call (and have to query the module everywhere for isDevice or other flags).

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

1548

Happy to add this return failure in a follow up commit if it's not a module, although I believe the attribute is never added to anything except modules, but it does prevent accidents in the future.

[kiranchandramohan]

In D147344#4299147, @agozillon wrote:

In D147344#4299066, @kiranchandramohan wrote:

See comment inline.

Also, I am a bit confused about when we are using the amendOperation translation vs setting the OpenMPIRBuilder config (with data from the OffloadModuleInterface). @agozillon @skatrak @jsjodin

I'm not sure what you mean in this case (I am probably being dumb and missing your point, so I apologies), this runtime flag generation is seperate from the config file, the only thing this does is utilise the OMPIRBuilder to create these globals in the exact same way Clang does at the moment. The attribute itself contains the values that in Clang are frontend language options, so can be generated directly from them. In Flangs and the OpenMP dialects case, we need an attribute to contain and carry them down. This particular amendOperation, will be called when the Module operation is processed and only if the mlir::omp::FlagsAttr is present on the module, it is only present on the module if Flang is called with -fopenmp-is-device.

There are two concerns:

1. There are two paths for translating OpenMP information in the module. (the config + OffloadModuleInterface and the amendOperation translation)
2. Directly accessing the attribute on the Module without using the interface.

This is more evident in requires translation (D147219) where both paths are used.

The good part about amendOperation is that it is available in OpenMPToLLVMIRTranslation.
Ideally, if convertOperation of the Module was the first thing being done, and is available in OpenMPToLLVMIRTranslation then everything could be done there. But I guess this is not possible since BuiltinToLLVMIRTranslation.cpp already handles it. I wonder if BuiltinToLLVMIRTranslation can forward to OpenMPToLLVMIRTranslation.

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

1548

Yes, please add this.

[agozillon]

In D147344#4301606, @kiranchandramohan wrote:

In D147344#4299147, @agozillon wrote:

In D147344#4299066, @kiranchandramohan wrote:

See comment inline.

Also, I am a bit confused about when we are using the amendOperation translation vs setting the OpenMPIRBuilder config (with data from the OffloadModuleInterface). @agozillon @skatrak @jsjodin

I'm not sure what you mean in this case (I am probably being dumb and missing your point, so I apologies), this runtime flag generation is seperate from the config file, the only thing this does is utilise the OMPIRBuilder to create these globals in the exact same way Clang does at the moment. The attribute itself contains the values that in Clang are frontend language options, so can be generated directly from them. In Flangs and the OpenMP dialects case, we need an attribute to contain and carry them down. This particular amendOperation, will be called when the Module operation is processed and only if the mlir::omp::FlagsAttr is present on the module, it is only present on the module if Flang is called with -fopenmp-is-device.

There are two concerns:

1. There are two paths for translating OpenMP information in the module. (the config + OffloadModuleInterface and the amendOperation translation)

I believe I see your concerns for 1, in this case, however I believe amendOperation is being used as a trigger to translate the information in both cases (D147344 & D147219), so as to keep the triggering logic seperate from the OMPIRBuilder as Clang does just now.
It might be possible to add a few config options and move the lowering of these Flags into the IRBuilders finalize method, I am not sure if it would affect Clang, so it'd be something I have to look into first. And I'm not sure that'd be possible for @skatrak's case as requires is more complex than the flags at the least. And I am not sure all cases will be similarly handle-able, a lot of these attributes are information from the frontend with their own constraints.

I think with Jan's current change of initializing the config file with data retrieved from the OffloadModule it's possible to just directly use, isEmbedded in place of a lot of cases where isDevice is utilised via the OffloadModuleInterface, perhaps we wish to just set the
config on initialisation from the stored OffloadModuleInterface and then only use the config internally inside of the OpenMPToLLVMIRTranslation file. The issue with this is, that it may lead to a lot of extra information being stored in the Config that Clang does not need and the intent is to try and keep it similarly of use for both compile flows to LLVM-IR. Although, I imagine a few extra config values aren't the end of the world.

2. Directly accessing the attribute on the Module without using the interface.

As for 2) if it is directly related to the amendOperation I personally am not sure it's a huge deal to be accessing this information without the module interface inside of the amendOperation for it's own specialisation (omp::FlagsAttr in this case) inside of the OpenMP lowering phase, this seems like one of the few cases it'd be fine to do so. However, it does leave the case where someone applies the attribute directly to an operation which is not a module, resulting in duplicate lowering (soon to be protected by the return failure you recommended), however, this requires a lot of effort to do already now that the registration of the attribute resides inside of an interface that is only applied to the module currently.

This is more evident in requires translation (D147219) where both paths are used.

The good part about amendOperation is that it is available in OpenMPToLLVMIRTranslation.
Ideally, if convertOperation of the Module was the first thing being done, and is available in OpenMPToLLVMIRTranslation then everything could be done there. But I guess this is not possible since BuiltinToLLVMIRTranslation.cpp already handles it. I wonder if BuiltinToLLVMIRTranslation can forward to OpenMPToLLVMIRTranslation.

I can't really speak for this option, it'd be interesting, but I don't know how easy it would be to move the convertOperation of the module in the existing order of things, it may have side affects on other dialects. Or how how easy it would be to forward it, at the moment I believe @skatrak had to implement a default builtin.module lowering for convertOperation. Perhaps it's possible to have an OpenMP OffloadModuleInterface variant, I am unsure.

I can try to re-work this lowering into the OMPIRBuilders finalize method and create a new patch to replace this if I succeed in moving it into the finalize method without breaking something in Clang or Flang, if that's something we wish to do, I am more than happy to do it?

Although, for me personally the more MLIR way of doing things would be to lower the attribute itself utilising the IRBuilder and not push it into the OMPIRBuilder where things become rather opaque for anyone not in the know. If there was a single collective way to lower everything on the module as you said via convertOperation, then that would be a very ideal way to do things I believe.

[skatrak]

I'll try to give some of my thoughts on the different issues being discussed here.

Also, I am a bit confused about when we are using the amendOperation translation vs setting the OpenMPIRBuilder config (with data from the OffloadModuleInterface).

To me, it makes sense to use both approaches simultaneously. Using amendOperation to do code generation for the actual attributes, if necessary, and store some of the information from these global attributes in the OpenMPIRBuilder::Config structure when it can influence codegen for other operations.

Perhaps the lowering itself could be pushed into the OMPIRBuilder's finalize method, by filling the OMPIRBuilder's config file with these flags during config setup, however, I am unsure how that would impact Clang at the moment.

This is one thing we could certainly do, but I feel like it does obscure the fact that there is IR being generated specifically for certain dialect attributes (as @agozillon pointed out in another comment), and it will require making changes in clang as well. By relying on amendOperation for this, we're more clearly documenting the circumstances under which that IR is produced, and we're not forcing ourselves to store every global attribute in the configuration structure. With our current approach, we only need to store there the information that can impact the generation of code for other operations and not information that will just be used to produce some LLVM IR metadata, globals and such.

There are two concerns:

1. There are two paths for translating OpenMP information in the module. (the config + OffloadModuleInterface and the amendOperation translation)
2. Directly accessing the attribute on the Module without using the interface.

This is more evident in requires translation (D147219) where both paths are used.

I think I can explain a bit what the thinking was behind the approach followed by that patch. The requires directive has a couple aspects to it:

1. If present, it has an effect over the whole compilation unit.
2. The information about the clauses that are set by the user can be used to influence how other constructs/operations are lowered.
3. Some of the clauses supported by the directive (currently the ones that relate to target offloading) must be registered with the OpenMP runtime library through a function call where these clauses are represented as flags. This runtime call is made within a new registration function that is marked as constructor for the LLVM module, making sure that it runs first.

Due to (1), I think it makes sense to represent the directive as an MLIR module attribute, because it makes it accessible from anywhere. The information from this attribute is stored in the OpenMPIRBuilder::Config because of (2), and it is lowered directly to LLVM IR in amendOperation because of (3). I do agree that we should try to avoid accessing the attributes that are part of the OffloadModuleInterface directly as much as possible, but in the case of the amendOperation function I agree with @agozillon that this could be the one spot where it would make sense to do.

The good part about amendOperation is that it is available in OpenMPToLLVMIRTranslation.
Ideally, if convertOperation of the Module was the first thing being done, and is available in OpenMPToLLVMIRTranslation then everything could be done there. But I guess this is not possible since BuiltinToLLVMIRTranslation.cpp already handles it. I wonder if BuiltinToLLVMIRTranslation can forward to OpenMPToLLVMIRTranslation.

The issue I see with trying to somehow make BuiltinToLLVMIRTranslation forward to OpenMPToLLVMIRTranslation is that these interfaces are tied to dialects, and currently as far as I can see at most one can be tied to each dialect. I can see the possibility of creating some kind of PrioritizedLLVMIRTranslationInterface that just acts as a container for several translations, where they are ordered by some sort of priority (lowest would be the current BuiltinToLLVMIRTranslation, and define a new OpenMPBuiltinToLLVMIRTranslation to deal exclusively with modules that have the OffloadModuleInterface attached). I would have to actually give it a shot to make sure it's possible to do, just because I'm not sure how the BuiltinToLLVMIRTranslation and OpenMPToLLVMIRTranslation would both have access to that PrioritizedLLVMIRTranslationInterface that would be added to the builtin dialect. I would probably add it inside of registerBuiltinDialectTranslation, so it would be easy to set the default translation interface, but I'm not sure how registerOpenMPDialectTranslation could get access to it as well. In any case, this seems quite an over-complicated approach to me.

On the other hand, there's the option of making the BuiltinToLLVMIRTranslation aware of the OffloadModuleInterface and make convertOperation behave differently depending on whether the interface is present or not. I also don't like this option just because it doesn't look like a good idea to make the builtin translation have a dependency on the OpenMP dialect. We can't register another translation interface for the builtin dialect defined within OpenMPToLLVMIRTranslation, because it's not possible to have more than one, and overriding the default one would result in everything else probably misbehaving.

If there was a single collective way to lower everything on the module as you said via convertOperation, then that would be a very ideal way to do things I believe.

My feeling is that if we really wanted to implement lowering of these module attributes within a convertOperation function, then we should reconsider the idea of creating an omp.module that defines all these attributes. Trying to do the same to a builtin.module doesn't quite have a clean solution, in my opinion. @agozillon already made a proof-of-concept omp.module a few months ago, if I remember correctly, and that was quite involved as well. But if that's something we end up pursuing, this may be one reason for it.

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

builder.restoreIP(ompBuilder->createTargetData(
         mapperFunc, deviceID, ifCond, processMapOpCB));
   }
 
   if (failed(processMapOpStatus))
     return processMapOpStatus;
   return bodyGenStatus;
 }
 
+/// Lowers the FlagsAttr which is applied to the module on the device
+/// pass when offloading, this attribute contains OpenMP RTL globals that can
+/// be passed as flags to the frontend, otherwise they are set to default
+LogicalResult convertFlagsAttr(Operation *op, mlir::omp::FlagsAttr attribute,
+                               LLVM::ModuleTranslation &moduleTranslation) {
+  llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+
+  ompBuilder->createGlobalFlag(
+      attribute.getDebugKind() /*LangOpts().OpenMPTargetDebug*/,
+      "__omp_rtl_debug_kind");
+  ompBuilder->createGlobalFlag(
+      attribute
+          .getAssumeTeamsOversubscription() /*LangOpts().OpenMPTeamSubscription*/
+      ,
+      "__omp_rtl_assume_teams_oversubscription");
+  ompBuilder->createGlobalFlag(
+      attribute
+          .getAssumeThreadsOversubscription() /*LangOpts().OpenMPThreadSubscription*/
+      ,
+      "__omp_rtl_assume_threads_oversubscription");
+  ompBuilder->createGlobalFlag(
+      attribute.getAssumeNoThreadState() /*LangOpts().OpenMPNoThreadState*/,
+      "__omp_rtl_assume_no_thread_state");
+  ompBuilder->createGlobalFlag(
+      attribute
+          .getAssumeNoNestedParallelism() /*LangOpts().OpenMPNoNestedParallelism*/
+      ,
+      "__omp_rtl_assume_no_nested_parallelism");
+
+  return success();
+}
+
 namespace {
 
 /// Implementation of the dialect interface that converts operations belonging
 /// to the OpenMP dialect to LLVM IR.
 class OpenMPDialectLLVMIRTranslationInterface
     : public LLVMTranslationDialectInterface {
 public:
   using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface;
 
   /// Translates the given operation to LLVM IR using the provided IR builder
   /// and saving the state in `moduleTranslation`.
   LogicalResult
   convertOperation(Operation *op, llvm::IRBuilderBase &builder,
                    LLVM::ModuleTranslation &moduleTranslation) const final;
+
+  LogicalResult
+  amendOperation(Operation *op, NamedAttribute attribute,
+                 LLVM::ModuleTranslation &moduleTranslation) const final;
 };
 
 } // namespace
 
+/// Given an OpenMP MLIR attribute, create the corresponding LLVM-IR, runtime
+/// calls, or operation amendments
+LogicalResult OpenMPDialectLLVMIRTranslationInterface::amendOperation(
+    Operation *op, NamedAttribute attribute,
+    LLVM::ModuleTranslation &moduleTranslation) const {
+
+  return llvm::TypeSwitch<Attribute, LogicalResult>(attribute.getValue())
+      .Case([&](mlir::omp::FlagsAttr rtlAttr) {
+        return convertFlagsAttr(op, rtlAttr, moduleTranslation);
+      })
+      .Default([&](Attribute attr) {
+        // fall through for omp attributes that do not require lowering and/or
+        // have no concrete definition and thus no type to define a case on
+        return success();
+      });
+
+  return failure();
+}
+
 /// Given an OpenMP MLIR operation, create the corresponding LLVM IR
 /// (including OpenMP runtime calls).
 LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
     Operation *op, llvm::IRBuilderBase &builder,
     LLVM::ModuleTranslation &moduleTranslation) const {
 
   llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
 

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

omp.atomic.capture memory_order(acq_rel) {
     omp.atomic.update %arg1 : !llvm.ptr {
     ^bb0(%xval: i32):
       %newval = llvm.add %xval, %expr : i32
       omp.yield(%newval : i32)
     }
   }
   llvm.return
 }
+
+// -----
+
+// CHECK-NOT: @__omp_rtl_debug_kind = weak_odr hidden constant i32 1
+// CHECK-NOT: @__omp_rtl_assume_teams_oversubscription = weak_odr hidden constant i32 1
+// CHECK-NOT: @__omp_rtl_assume_threads_oversubscription = weak_odr hidden constant i32 1
+// CHECK-NOT: @__omp_rtl_assume_no_thread_state = weak_odr hidden constant i32 1
+// CHECK-NOT: @__omp_rtl_assume_no_nested_parallelism = weak_odr hidden constant i32 1
+module attributes {omp.flags = #omp.flags<debug_kind = 1, assume_teams_oversubscription = true, 
+                                          assume_threads_oversubscription = true, assume_no_thread_state = true, 
+                                          assume_no_nested_parallelism = true>} {}
+
+// -----
+
+// CHECK: @__omp_rtl_debug_kind = weak_odr hidden constant i32 1
+// CHECK: @__omp_rtl_assume_teams_oversubscription = weak_odr hidden constant i32 1
+// CHECK: @__omp_rtl_assume_threads_oversubscription = weak_odr hidden constant i32 1
+// 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 1
+module attributes {omp.flags = #omp.flags<debug_kind = 1, assume_teams_oversubscription = true, 
+                                          assume_threads_oversubscription = true, assume_no_thread_state = true, 
+                                          assume_no_nested_parallelism = true>, 
+                   omp.is_device = #omp.isdevice<is_device = true>} {}
+
+// -----
+
+// CHECK-NOT: @__omp_rtl_debug_kind = weak_odr hidden constant i32 0
+// CHECK-NOT: @__omp_rtl_assume_teams_oversubscription = weak_odr hidden constant i32 0
+// CHECK-NOT: @__omp_rtl_assume_threads_oversubscription = weak_odr hidden constant i32 0
+// CHECK-NOT: @__omp_rtl_assume_no_thread_state = weak_odr hidden constant i32 0
+// CHECK-NOT: @__omp_rtl_assume_no_nested_parallelism = weak_odr hidden constant i32 0
+module attributes {omp.flags = #omp.flags<>} {}
+
+// -----
+
+// CHECK: @__omp_rtl_debug_kind = weak_odr hidden constant i32 0
+// CHECK: @__omp_rtl_assume_teams_oversubscription = weak_odr hidden constant i32 0
+// CHECK: @__omp_rtl_assume_threads_oversubscription = weak_odr hidden constant i32 0
+// CHECK: @__omp_rtl_assume_no_thread_state = weak_odr hidden constant i32 0
+// CHECK: @__omp_rtl_assume_no_nested_parallelism = weak_odr hidden constant i32 0
+module attributes {omp.flags = #omp.flags<>, omp.is_device = #omp.isdevice<is_device = true>} {}
+
+// -----
+
+// CHECK-NOT: @__omp_rtl_debug_kind = weak_odr hidden constant i32 0
+// CHECK-NOT: @__omp_rtl_assume_teams_oversubscription = weak_odr hidden constant i32 0
+// CHECK-NOT: @__omp_rtl_assume_threads_oversubscription = weak_odr hidden constant i32 0
+// CHECK-NOT: @__omp_rtl_assume_no_thread_state = weak_odr hidden constant i32 0
+// CHECK-NOT: @__omp_rtl_assume_no_nested_parallelism = weak_odr hidden constant i32 0
+module attributes {omp.flags = #omp.flags<debug_kind = 0, assume_teams_oversubscription = false, 
+                                          assume_threads_oversubscription = false, assume_no_thread_state = false, 
+                                          assume_no_nested_parallelism = false>} {}
+
+// -----
+
+// CHECK: @__omp_rtl_debug_kind = weak_odr hidden constant i32 0
+// CHECK: @__omp_rtl_assume_teams_oversubscription = weak_odr hidden constant i32 0
+// CHECK: @__omp_rtl_assume_threads_oversubscription = weak_odr hidden constant i32 0
+// CHECK: @__omp_rtl_assume_no_thread_state = weak_odr hidden constant i32 0
+// CHECK: @__omp_rtl_assume_no_nested_parallelism = weak_odr hidden constant i32 0
+module attributes {omp.flags = #omp.flags<debug_kind = 0, assume_teams_oversubscription = false, 
+                                          assume_threads_oversubscription = false, assume_no_thread_state = false, 
+                                          assume_no_nested_parallelism = false>, 
+                   omp.is_device = #omp.isdevice<is_device = true>} {}
+
+// -----
+
+// CHECK-NOT: @__omp_rtl_debug_kind = weak_odr hidden constant i32 0
+// CHECK-NOT: @__omp_rtl_assume_teams_oversubscription = weak_odr hidden constant i32 1
+// CHECK-NOT: @__omp_rtl_assume_threads_oversubscription = weak_odr hidden constant i32 0
+// CHECK-NOT: @__omp_rtl_assume_no_thread_state = weak_odr hidden constant i32 1
+// CHECK-NOT: @__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: @__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>, 
+                   omp.is_device = #omp.isdevice<is_device = true>} {}