Skip to content

[MLIR][OpenMP] Add OMP Declare Mapper MLIR Op definition #117045

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Feb 18, 2025

Conversation

TIFitis
Copy link
Member

@TIFitis TIFitis commented Nov 20, 2024

This patch adds the OMP.DeclareMapperOp to MLIR.
The HLFIR/FIR lowering for Declare Mapper is available here #117046.

@llvmbot
Copy link
Member

llvmbot commented Nov 20, 2024

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-flang-openmp

Author: Akash Banerjee (TIFitis)

Changes

This patch adds the OMP.DeclareMapperOp to MLIR.


Full diff: https://github.com/llvm/llvm-project/pull/117045.diff

3 Files Affected:

  • (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td (+28)
  • (modified) mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp (+8)
  • (modified) mlir/test/Dialect/OpenMP/ops.mlir (+8)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 156e6eb371b85d..eabca0b5ac7a25 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -1652,6 +1652,34 @@ def CancellationPointOp : OpenMP_Op<"cancellation_point", clauses = [
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// 2.19.7.3 Declare Mapper Directive
+//===----------------------------------------------------------------------===//
+def DeclareMapperOp : OpenMP_Op<"declare_mapper", traits = [
+  Symbol], clauses = [
+    OpenMP_MapClause
+  ]> {
+  let summary = "declare mapper directive";
+  let description = [{
+    The declare mapper directive declares a user-defined mapper for a given
+    type, and may define a mapper-identifier that can be used in a map clause.
+  }] # clausesDescription;
+
+  let arguments = !con((ins SymbolNameAttr:$sym_name,
+                            OpenMP_PointerLikeType:$var_ptr,
+                            TypeAttr:$var_type), clausesArgs);
+
+  let builders = [
+    OpBuilder<(ins CArg<"const DeclareMapperOperands &">:$clauses)>
+  ];
+
+  // Override clause-based assemblyFormat.
+  let assemblyFormat = "$sym_name `:` $var_ptr `:` type($var_ptr) `:` $var_type" # " oilist(" #
+    clausesOptAssemblyFormat # ")  attr-dict";
+
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // 2.19.5.7 declare reduction Directive
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 94e71e089d4b18..2f40405b54b73c 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2137,6 +2137,14 @@ LogicalResult DistributeOp::verifyRegions() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// DeclareMapperOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult DeclareMapperOp::verify() {
+  return verifyMapClause(*this, getMapVars());
+}
+
 //===----------------------------------------------------------------------===//
 // DeclareReductionOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index c25a6ef4b4849b..a3f2c458b6c592 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -879,6 +879,14 @@ cleanup {
   omp.yield
 }
 
+// CHECK: %[[DECL_VAR:.*]] = llvm.alloca %{{.*}}
+// CHECK: %[[DECL_MAP_INFO:.*]] = omp.map.info var_ptr(%[[DECL_VAR]] : !llvm.ptr, !llvm.struct<"my_type", (i32)>) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
+// CHECK: omp.declare_mapper @my_mapper : %[[DECL_VAR]] : !llvm.ptr : !llvm.struct<"my_type", (i32)> map_entries(%[[DECL_MAP_INFO]] : !llvm.ptr)
+%decl_c1 = arith.constant 1 : i64
+%decl_var = llvm.alloca %decl_c1 x !llvm.struct<"my_type", (i32)> : (i64) -> !llvm.ptr
+%decl_map_info = omp.map.info var_ptr(%decl_var : !llvm.ptr, !llvm.struct<"my_type", (i32)>) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
+omp.declare_mapper @my_mapper : %decl_var : !llvm.ptr : !llvm.struct<"my_type", (i32)> map_entries(%decl_map_info : !llvm.ptr)
+
 // CHECK-LABEL: func @wsloop_reduction
 func.func @wsloop_reduction(%lb : index, %ub : index, %step : index) {
   %c1 = arith.constant 1 : i32

@llvmbot
Copy link
Member

llvmbot commented Nov 20, 2024

@llvm/pr-subscribers-mlir-openmp

Author: Akash Banerjee (TIFitis)

Changes

This patch adds the OMP.DeclareMapperOp to MLIR.


Full diff: https://github.com/llvm/llvm-project/pull/117045.diff

3 Files Affected:

  • (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td (+28)
  • (modified) mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp (+8)
  • (modified) mlir/test/Dialect/OpenMP/ops.mlir (+8)
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 156e6eb371b85d..eabca0b5ac7a25 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -1652,6 +1652,34 @@ def CancellationPointOp : OpenMP_Op<"cancellation_point", clauses = [
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// 2.19.7.3 Declare Mapper Directive
+//===----------------------------------------------------------------------===//
+def DeclareMapperOp : OpenMP_Op<"declare_mapper", traits = [
+  Symbol], clauses = [
+    OpenMP_MapClause
+  ]> {
+  let summary = "declare mapper directive";
+  let description = [{
+    The declare mapper directive declares a user-defined mapper for a given
+    type, and may define a mapper-identifier that can be used in a map clause.
+  }] # clausesDescription;
+
+  let arguments = !con((ins SymbolNameAttr:$sym_name,
+                            OpenMP_PointerLikeType:$var_ptr,
+                            TypeAttr:$var_type), clausesArgs);
+
+  let builders = [
+    OpBuilder<(ins CArg<"const DeclareMapperOperands &">:$clauses)>
+  ];
+
+  // Override clause-based assemblyFormat.
+  let assemblyFormat = "$sym_name `:` $var_ptr `:` type($var_ptr) `:` $var_type" # " oilist(" #
+    clausesOptAssemblyFormat # ")  attr-dict";
+
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // 2.19.5.7 declare reduction Directive
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 94e71e089d4b18..2f40405b54b73c 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2137,6 +2137,14 @@ LogicalResult DistributeOp::verifyRegions() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// DeclareMapperOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult DeclareMapperOp::verify() {
+  return verifyMapClause(*this, getMapVars());
+}
+
 //===----------------------------------------------------------------------===//
 // DeclareReductionOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index c25a6ef4b4849b..a3f2c458b6c592 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -879,6 +879,14 @@ cleanup {
   omp.yield
 }
 
+// CHECK: %[[DECL_VAR:.*]] = llvm.alloca %{{.*}}
+// CHECK: %[[DECL_MAP_INFO:.*]] = omp.map.info var_ptr(%[[DECL_VAR]] : !llvm.ptr, !llvm.struct<"my_type", (i32)>) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
+// CHECK: omp.declare_mapper @my_mapper : %[[DECL_VAR]] : !llvm.ptr : !llvm.struct<"my_type", (i32)> map_entries(%[[DECL_MAP_INFO]] : !llvm.ptr)
+%decl_c1 = arith.constant 1 : i64
+%decl_var = llvm.alloca %decl_c1 x !llvm.struct<"my_type", (i32)> : (i64) -> !llvm.ptr
+%decl_map_info = omp.map.info var_ptr(%decl_var : !llvm.ptr, !llvm.struct<"my_type", (i32)>) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
+omp.declare_mapper @my_mapper : %decl_var : !llvm.ptr : !llvm.struct<"my_type", (i32)> map_entries(%decl_map_info : !llvm.ptr)
+
 // CHECK-LABEL: func @wsloop_reduction
 func.func @wsloop_reduction(%lb : index, %ub : index, %step : index) {
   %c1 = arith.constant 1 : i32

Copy link
Contributor

@kiranchandramohan kiranchandramohan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it will be useful to have a bit more explanation on how the omp.declare_mapper operation is going to be used. Where will the declaration be effective? Is there a subsequent pass that will implement the mapping implied by the omp.declare_mapper?

type, and may define a mapper-identifier that can be used in a map clause.
}] # clausesDescription;

let arguments = !con((ins SymbolNameAttr:$sym_name,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am assuming sym_name is the identifier. The description sounds like it is optional but the declaration here doesn't. I think you might want to update the description above to say that you will always have an identifier in the MLIR representation.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated.

@TIFitis
Copy link
Member Author

TIFitis commented Nov 22, 2024

I think it will be useful to have a bit more explanation on how the omp.declare_mapper operation is going to be used. Where will the declaration be effective? Is there a subsequent pass that will implement the mapping implied by the omp.declare_mapper?

I am working on subsequent PRs to update the mapClauseOp and FIR lowering to support declare mapper. The omp.declare_mapper symbol will be used there to identify the mapper.

The omp.declare_mapper operation itself will be used in the OpenMPToLLVMIRTranslation along with OMPIRBuilder to enact the mapping rules defined in it. I have the OMPIRBuilder work up for review already #110001. I will add another patch for OpenMPToLLVMIRTranslation soon. With that we should have all the work in place for supporting OMP Declare Mapper.

@TIFitis
Copy link
Member Author

TIFitis commented Nov 22, 2024

Also the FIR/HLFIR lowering PR for declare mapper is available here #117046.

Copy link
Member

@skatrak skatrak left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you Akash for this work, I have some concerns at the moment with this approach. Maybe I'm not getting the full picture, though, so let me know.

//===----------------------------------------------------------------------===//
def DeclareMapperOp : OpenMP_Op<"declare_mapper", traits = [
Symbol], clauses = [
OpenMP_MapClause
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would make sense to extract the definition of the mapper identifier into a clause and simplify the operation:

class OpenMP_MapperIdentifierClauseSkip<...> : OpenMP_Clause<...> {
  let traits = [
    Symbol
  ];

  let arguments = (ins
    SymbolNameAttr:$sym_name,
    OpenMP_PointerLikeType:$mapper_id_var,
    TypeAttr:$mapper_id_type
  );

  ...
}

%decl_c1 = arith.constant 1 : i64
%decl_var = llvm.alloca %decl_c1 x !llvm.struct<"my_type", (i32)> : (i64) -> !llvm.ptr
%decl_map_info = omp.map.info var_ptr(%decl_var : !llvm.ptr, !llvm.struct<"my_type", (i32)>) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
omp.declare_mapper @my_mapper : %decl_var : !llvm.ptr : !llvm.struct<"my_type", (i32)> map_entries(%decl_map_info : !llvm.ptr)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have one major concern about this operation definition, which is perhaps a result of misunderstanding it, so let me know if that's the case.

The declare mapper construct lets a user basically define a recipe that takes one argument of a given type and use it to represent an unspecified variable of that type within one or more map clauses. That placeholder variable is then bound to specific variables when encountering another map clause for which the defined mapper is implicitly or explicitly applied. The result is as if the encountered map clause was replaced by the map clause(s) defined in the mapper, after replacing the placeholder variable with the actual variable being mapped.

As such, I'd expect its definition to be global (just like privatizers and omp.declare_reduction) and contain a region with an entry block argument for that placeholder variable (and maybe also other(s) for bounds). My understanding is that the region would construct the map arguments and omp.yield them. Users of the mapper would "call" the mapper through its global symbol and a pointer-like variable. The translation to LLVM IR of map clauses using a declared mapper would have to inline the recipe body while binding the entry block argument to the variable being mapped and treating the yielded results like any other mapped variable.

In MLIR, I'm picturing something like this:

omp.declare_mapper @my_mapper(%arg0 : !llvm.ptr) {
  %mapinfo = omp.map.info var_ptr(%arg0 : !llvm.ptr, !llvm.struct<"my_type", (i32)>) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr
  omp.yield(%mapinfo : !llvm.ptr)
}

func.func @foo(...) {
  ...
  omp.target map_entries(@my_mapper %x -> %x.mapped : !llvm.ptr) {
    ...
  }
  return
}

And for declare mappers that result in multiple mapped variables something like this:

omp.declare_mapper @my_mapper(%arg0 : !llvm.ptr) {
  %gep1 = llvm.getelementptr %arg0[0, 3] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(f32, array<10 x i32>, struct<(f32, i32)>, i32)>
  %mapinfo1 = omp.map.info var_ptr(%gep1 : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr
  %gep2 = llvm.getelementptr %arg0[0, 2, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(f32, array<10 x i32>, struct<(f32, i32)>, i32)>
  %mapinfo2 = omp.map.info var_ptr(%gep2 : !llvm.ptr, i32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr
  omp.yield(%mapinfo1, %mapinfo2 : !llvm.ptr, !llvm.ptr)
}

func.func @foo(...) {
  ...
  omp.target map_entries(@my_mapper %x -> %x.mapped.1, %x.mapped.2 : !llvm.ptr, !llvm.ptr) {
    ...
  }
  return
}

My thinking is that, without making omp.declare_mapper a recipe operation, we're not able to actually define something that can be used when mapping variables elsewhere, but rather we have to produce an ad-hoc operation at each place where the mapper is instantiated, leaving all of the work to the frontend, in which case we don't even need an explicit MLIR representation for it.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking at Clang's implementation and the OMP RTL calls, declare mappers are processed a bit differently than simply "copying" the mapClause info to wherever the declMapper is referenced.

For the sake of keeping the conversation at one place, do you mind continuing this conversation at #117046. We already have an ongoing discussion there on the same topic, where I've tried to explain how declare mapper works and how it's processed.

In short, the declMapper directive results in a new mapperFunc in the code. A ptr to this mapperFunc is then passed on to the RTL call. Generation of this mapperFunc along with passing it to the RTL is handled by the OMPIRBuilder in #110001.

So every declMapperOp eventually gets lowered to it's own function, which is referenced through the target directive map clauses.

I've also discussed the issue with decoupling the map clause from the directive in #117046. AFAIK I don't think we have currently have floating mapInfoOps that are not tied to any Op through the mapClause(mapEntries). If we introduce this "special" case of map clause where it's attached to the parentOp through a nested region, then I think we would end up having to handle this special case elsewhere such as different stages of lowering, map passes(mapInfoFinalization) etc. I don't see what we stand to gain from nesting the single var that's associated with the declMapperOp to a region instead of hoisting it. Only problem I can think of is violating the scoping rule, which I think can be enforced through the declMapperOp's verifier.

Let me know if this changes your view, or if you're still in favour of nesting mapInfoOps.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you'd really only have to teach the MapInfoFinalization pass, and only a little, as from what I recall it just works on MapInfoOp's the only case it has to know more about the operation holding it is when it's something that needs BlockArguments inserted, and there is an almost generic way of doing this I think, so it isn't a lot of work I believe. I don't think you'd have to teach the MLIR -> LLVM-IR how to handle the operation much more than you'd have to already for the current approach (if we're just talking about getting the information from the map clause at least, the function generation and embedding is it's own kettle of fish), but I could be wrong.

P.S. feel free to just consider this a comment on the other PR and reply there if you wish to reply, mostly just an informational message though, rather than trying to push things one way or another :-)

@TIFitis TIFitis force-pushed the users/akash/mapper_op branch from 00c9f3d to 4b5c92e Compare December 23, 2024 15:01
@TIFitis
Copy link
Member Author

TIFitis commented Jan 28, 2025

Polite ping for review :)

@TIFitis
Copy link
Member Author

TIFitis commented Feb 3, 2025

Polite request for review 🙂

@TIFitis TIFitis requested a review from ergawy February 3, 2025 14:20
Copy link
Member

@skatrak skatrak left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you Akash! Sorry for the delay getting back to this. I have a couple of small comments.

// CHECK: %[[DECL_MAP_INFO:.*]] = omp.map.info var_ptr(%{{.*}} : !llvm.ptr, !llvm.struct<"my_type", (i32)>) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
%decl_map_info = omp.map.info var_ptr(%arg : !llvm.ptr, !llvm.struct<"my_type", (i32)>) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = ""}
// CHECK: omp.declare_mapper_info map_entries(%[[DECL_MAP_INFO]] : !llvm.ptr)
omp.declare_mapper_info map_entries(%decl_map_info : !llvm.ptr)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is your plan for finding this operation during MLIR to LLVM IR translation?

What I'm wondering is whether the plan is to allow multiple exit points from the region (e.g. multiple blocks terminated with omp.declare_mapper.info operations), and evaluate this at runtime, or whether we want to restrict the region to be able to find this operation statically.

To that end, the operation could always contain a single block (i.e. add the SingleBlock trait) or have some logic in the verifier to make sure we have a single exit block containing this terminator reachable from all paths starting at the entry block.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should add a region verifier to DeclareMapperOp (i.e. let hasRegionVerifier = 1 and define DeclareMapperOp::verifyRegions()) to check that the last operation in its only block is a DeclareMapperInfoOp. Also, I'd suggest adding an extraClassDeclaration DeclareMapperInfoOp DeclareMapperOp::getDeclareMapperInfo() or similar, so that this op-specific logic doesn't have to be replicated by all users, but rather stay contained in the operation itself.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the suggestion, I've added both.

@TIFitis TIFitis force-pushed the users/akash/mapper_op branch from 4b5c92e to 38ce5eb Compare February 7, 2025 17:22
Copy link
Member

@skatrak skatrak left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thank you!

@TIFitis TIFitis merged commit 0e960f1 into main Feb 18, 2025
5 of 7 checks passed
@TIFitis TIFitis deleted the users/akash/mapper_op branch February 18, 2025 16:11
TIFitis added a commit that referenced this pull request Feb 18, 2025
…ve (#117046)

This patch adds HLFIR/FIR lowering support for OpenMP Declare Mapper
directive.
Depends on #117045.
wldfngrs pushed a commit to wldfngrs/llvm-project that referenced this pull request Feb 19, 2025
This patch adds the OMP.DeclareMapperOp to MLIR.
The HLFIR/FIR lowering for Declare Mapper is available here llvm#117046.
wldfngrs pushed a commit to wldfngrs/llvm-project that referenced this pull request Feb 19, 2025
…ve (llvm#117046)

This patch adds HLFIR/FIR lowering support for OpenMP Declare Mapper
directive.
Depends on llvm#117045.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants