[llvm-branch-commits] [clang] [llvm] [mlir] [MLIR][OpenMP] Add LLVM translation support for OpenMP UserDefinedMappers (PR #124746)

Akash Banerjee via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Fri Jan 31 03:37:33 PST 2025


================
@@ -3421,6 +3441,85 @@ static void genMapInfos(llvm::IRBuilderBase &builder,
   }
 }
 
+static llvm::Expected<llvm::Function *>
+emitUserDefinedMapper(Operation *declMapperOp, llvm::IRBuilderBase &builder,
+                      LLVM::ModuleTranslation &moduleTranslation);
+
+static llvm::Expected<llvm::Function *>
+getOrCreateUserDefinedMapperFunc(Operation *declMapperOp,
+                                 llvm::IRBuilderBase &builder,
+                                 LLVM::ModuleTranslation &moduleTranslation) {
+  llvm::DenseMap<const Operation *, llvm::Function *> userDefMapperMap;
+  auto iter = userDefMapperMap.find(declMapperOp);
+  if (iter != userDefMapperMap.end())
+    return iter->second;
+  llvm::Expected<llvm::Function *> mapperFunc =
+      emitUserDefinedMapper(declMapperOp, builder, moduleTranslation);
+  if (!mapperFunc)
+    return mapperFunc.takeError();
+  userDefMapperMap.try_emplace(declMapperOp, *mapperFunc);
+  return userDefMapperMap.lookup(declMapperOp);
+}
+
+static llvm::Expected<llvm::Function *>
+emitUserDefinedMapper(Operation *op, llvm::IRBuilderBase &builder,
+                      LLVM::ModuleTranslation &moduleTranslation) {
+  auto declMapperOp = cast<omp::DeclareMapperOp>(op);
+  auto declMapperInfoOp =
+      *declMapperOp.getOps<omp::DeclareMapperInfoOp>().begin();
+  DataLayout dl = DataLayout(declMapperOp->getParentOfType<ModuleOp>());
+  llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+  llvm::Type *varType =
+      moduleTranslation.convertType(declMapperOp.getVarType());
+  std::string mapperName = ompBuilder->createPlatformSpecificName(
+      {"omp_mapper", declMapperOp.getSymName()});
+  SmallVector<Value> mapVars = declMapperInfoOp.getMapVars();
+
+  using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
+
+  // Fill up the arrays with all the mapped variables.
+  MapInfosTy combinedInfo;
+  auto genMapInfoCB =
+      [&](InsertPointTy codeGenIP, llvm::Value *ptrPHI,
+          llvm::Value *unused2) -> llvm::OpenMPIRBuilder::MapInfosOrErrorTy {
+    builder.restoreIP(codeGenIP);
+    moduleTranslation.mapValue(declMapperOp.getRegion().getArgument(0), ptrPHI);
+    moduleTranslation.mapBlock(&declMapperOp.getRegion().front(),
+                               builder.GetInsertBlock());
+    if (failed(moduleTranslation.convertBlock(declMapperOp.getRegion().front(),
+                                              /*ignoreArguments=*/true,
+                                              builder)))
+      return llvm::make_error<PreviouslyReportedError>();
+    MapInfoData mapData;
+    collectMapDataFromMapOperands(mapData, mapVars, moduleTranslation, dl,
+                                  builder);
+    genMapInfos(builder, moduleTranslation, dl, combinedInfo, mapData);
+
+    // Drop the mapping that is no longer necessary so that the same region can
+    // be processed multiple times.
+    moduleTranslation.forgetMapping(declMapperOp.getRegion());
+    return combinedInfo;
+  };
+
+  auto customMapperCB = [&](unsigned i, llvm::Function **mapperFunc) {
+    if (combinedInfo.Mappers[i]) {
+      // Call the corresponding mapper function.
+      llvm::Expected<llvm::Function *> newFn = getOrCreateUserDefinedMapperFunc(
----------------
TIFitis wrote:

I've updated the offloading test to trigger this recursion.

Worklist for something like this might be an overkill. Clang implements declare mapper in a similar recursive manner. Also, cyclic recursions are not possible in legal code.

Here's what the generated llvm IR looks like for the new test:

```
define void @_QQmain() {
  %.offload_baseptrs = alloca [2 x ptr], align 8
  %.offload_ptrs = alloca [2 x ptr], align 8
  %.offload_mappers = alloca [2 x ptr], align 8
  %1 = alloca { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, align 8
  %2 = alloca i32, i64 1, align 4
  %3 = alloca i32, i64 1, align 4
  %4 = alloca i32, i64 1, align 4
  br label %5

5:                                                ; preds = %9, %0
  %6 = phi i32 [ %18, %9 ], [ 1, %0 ]
  %7 = phi i64 [ %19, %9 ], [ 1024, %0 ]
  %8 = icmp sgt i64 %7, 0
  br i1 %8, label %9, label %20

9:                                                ; preds = %5
  store i32 %6, ptr %4, align 4
  %10 = load i32, ptr %4, align 4
  %11 = sext i32 %10 to i64
  %12 = sub nsw i64 %11, 1
  %13 = mul nsw i64 %12, 1
  %14 = mul nsw i64 %13, 1
  %15 = add nsw i64 %14, 0
  %16 = getelementptr i32, ptr @_QFEobj, i64 %15
  store i32 1, ptr %16, align 4
  %17 = load i32, ptr %4, align 4
  %18 = add nsw i32 %17, 1
  %19 = sub i64 %7, 1
  br label %5

20:                                               ; preds = %5
  store i32 %6, ptr %4, align 4
  store { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] } { ptr @_QFEobj, i64 ptrtoint (ptr getelementptr (i32, ptr null, i32 1) to i64), i32 20240719, i8 1, i8 9, i8 0, i8 0, [1 x [3 x i64]] [[3 x i64] [i64 1, i64 1024, i64 ptrtoint (ptr getelementptr (i32, ptr null, i32 1) to i64)]] }, ptr %1, align 8
  %21 = call i32 @_FortranASumInteger4(ptr %1, ptr @_QQclXf46b0d060c890540d012b521bc3468aa, i32 21, i32 0, ptr null)
  store i32 %21, ptr %2, align 4
  store i32 0, ptr %3, align 4
  %22 = getelementptr inbounds [2 x ptr], ptr %.offload_baseptrs, i32 0, i32 0
  store ptr %3, ptr %22, align 8
  %23 = getelementptr inbounds [2 x ptr], ptr %.offload_ptrs, i32 0, i32 0
  store ptr %3, ptr %23, align 8
  %24 = getelementptr inbounds [2 x ptr], ptr %.offload_mappers, i64 0, i64 0
  store ptr null, ptr %24, align 8
  %25 = getelementptr inbounds [2 x ptr], ptr %.offload_baseptrs, i32 0, i32 1
  store ptr @_QFEobj, ptr %25, align 8
  %26 = getelementptr inbounds [2 x ptr], ptr %.offload_ptrs, i32 0, i32 1
  store ptr @_QFEobj, ptr %26, align 8
  %27 = getelementptr inbounds [2 x ptr], ptr %.offload_mappers, i64 0, i64 1
  store ptr @.omp_mapper._QQFmy_mapper2, ptr %27, align 8
  %28 = getelementptr inbounds [2 x ptr], ptr %.offload_baseptrs, i32 0, i32 0
  %29 = getelementptr inbounds [2 x ptr], ptr %.offload_ptrs, i32 0, i32 0
  call void @__tgt_target_data_begin_mapper(ptr @7, i64 -1, i32 2, ptr %28, ptr %29, ptr @.offload_sizes, ptr @.offload_maptypes, ptr @.offload_mapnames, ptr %.offload_mappers)
  br label %omp.data.region

omp.data.region3:                                 ; preds = %omp.data.region1
  store i32 %43, ptr %4, align 4
  br label %omp.region.cont

omp.data.region2:                                 ; preds = %omp.data.region1
  store i32 %43, ptr %4, align 4
  %30 = load i32, ptr %3, align 4
  %31 = load i32, ptr %4, align 4
  %32 = sext i32 %31 to i64
  %33 = sub nsw i64 %32, 1
  %34 = mul nsw i64 %33, 1
  %35 = mul nsw i64 %34, 1
  %36 = add nsw i64 %35, 0
  %37 = getelementptr i32, ptr @_QFEobj, i64 %36
  %38 = load i32, ptr %37, align 4
  %39 = add i32 %30, %38
  store i32 %39, ptr %3, align 4
  %40 = load i32, ptr %4, align 4
  %41 = add nsw i32 %40, 1
  %42 = sub i64 %44, 1
  br label %omp.data.region1

omp.data.region1:                                 ; preds = %omp.data.region2, %omp.data.region
  %43 = phi i32 [ %41, %omp.data.region2 ], [ 1, %omp.data.region ]
  %44 = phi i64 [ %42, %omp.data.region2 ], [ 1024, %omp.data.region ]
  %45 = icmp sgt i64 %44, 0
  br i1 %45, label %omp.data.region2, label %omp.data.region3

omp.data.region:                                  ; preds = %20
  br label %omp.data.region1

omp.region.cont:                                  ; preds = %omp.data.region3
  %46 = getelementptr inbounds [2 x ptr], ptr %.offload_baseptrs, i32 0, i32 0
  %47 = getelementptr inbounds [2 x ptr], ptr %.offload_ptrs, i32 0, i32 0
  call void @__tgt_target_data_end_mapper(ptr @7, i64 -1, i32 2, ptr %46, ptr %47, ptr @.offload_sizes, ptr @.offload_maptypes, ptr @.offload_mapnames, ptr %.offload_mappers)
  %48 = call ptr @_FortranAioBeginExternalListOutput(i32 6, ptr @_QQclXf46b0d060c890540d012b521bc3468aa, i32 32)
  %49 = call i1 @_FortranAioOutputAscii(ptr %48, ptr @_QQclX53756D206F6E20686F73743A20202020, i64 16)
  %50 = load i32, ptr %2, align 4
  %51 = call i1 @_FortranAioOutputInteger32(ptr %48, i32 %50)
  %52 = call i32 @_FortranAioEndIoStatement(ptr %48)
  %53 = call ptr @_FortranAioBeginExternalListOutput(i32 6, ptr @_QQclXf46b0d060c890540d012b521bc3468aa, i32 33)
  %54 = call i1 @_FortranAioOutputAscii(ptr %53, ptr @_QQclX53756D206F6E206465766963653A2020, i64 16)
  %55 = load i32, ptr %3, align 4
  %56 = call i1 @_FortranAioOutputInteger32(ptr %53, i32 %55)
  %57 = call i32 @_FortranAioEndIoStatement(ptr %53)
  %58 = load i32, ptr %3, align 4
  %59 = load i32, ptr %2, align 4
  %60 = icmp eq i32 %58, %59
  br i1 %60, label %61, label %65

61:                                               ; preds = %omp.region.cont
  %62 = call ptr @_FortranAioBeginExternalListOutput(i32 6, ptr @_QQclXf46b0d060c890540d012b521bc3468aa, i32 36)
  %63 = call i1 @_FortranAioOutputAscii(ptr %62, ptr @_QQclX546573742070617373656421, i64 12)
  %64 = call i32 @_FortranAioEndIoStatement(ptr %62)
  br label %69

65:                                               ; preds = %omp.region.cont
  %66 = call ptr @_FortranAioBeginExternalListOutput(i32 6, ptr @_QQclXf46b0d060c890540d012b521bc3468aa, i32 38)
  %67 = call i1 @_FortranAioOutputAscii(ptr %66, ptr @_QQclX54657374206661696C656421, i64 12)
  %68 = call i32 @_FortranAioEndIoStatement(ptr %66)
  br label %69

69:                                               ; preds = %61, %65
  ret void
}

; Function Attrs: noinline nounwind
define internal void @.omp_mapper._QQFmy_mapper2(ptr noundef %0, ptr noundef %1, ptr noundef %2, i64 noundef %3, i64 noundef %4, ptr noundef %5) #0 {
entry:
  %6 = udiv exact i64 %3, 4096
  %7 = getelementptr %_QFTmytype, ptr %2, i64 %6
  %omp.arrayinit.isarray = icmp sgt i64 %6, 1
  %8 = and i64 %4, 8
  %9 = icmp ne ptr %1, %2
  %10 = and i64 %4, 16
  %11 = icmp ne i64 %10, 0
  %12 = and i1 %9, %11
  %13 = or i1 %omp.arrayinit.isarray, %12
  %.omp.array..init..delete = icmp eq i64 %8, 0
  %14 = and i1 %13, %.omp.array..init..delete
  br i1 %14, label %.omp.array..init, label %omp.arraymap.head

.omp.array..init:                                 ; preds = %entry
  %15 = mul nuw i64 %6, 4096
  %16 = and i64 %4, -4
  %17 = or i64 %16, 512
  call void @__tgt_push_mapper_component(ptr %0, ptr %1, ptr %2, i64 %15, i64 %17, ptr %5)
  br label %omp.arraymap.head

omp.arraymap.head:                                ; preds = %.omp.array..init, %entry
  %omp.arraymap.isempty = icmp eq ptr %2, %7
  br i1 %omp.arraymap.isempty, label %omp.done, label %omp.arraymap.body

omp.arraymap.body:                                ; preds = %omp.type.end, %omp.arraymap.head
  %omp.arraymap.ptrcurrent = phi ptr [ %2, %omp.arraymap.head ], [ %omp.arraymap.next, %omp.type.end ]
  %18 = getelementptr %_QFTmytype, ptr %omp.arraymap.ptrcurrent, i32 0, i32 0
  %array_offset = getelementptr inbounds [1024 x i32], ptr %18, i64 0, i64 0
  %19 = call i64 @__tgt_mapper_num_components(ptr %0)
  %20 = shl i64 %19, 48
  %21 = add nuw i64 2, %20
  %22 = and i64 %4, 3
  %23 = icmp eq i64 %22, 0
  br i1 %23, label %omp.type.alloc, label %omp.type.alloc.else

omp.type.alloc:                                   ; preds = %omp.arraymap.body
  %24 = and i64 %21, -4
  br label %omp.type.end

omp.type.alloc.else:                              ; preds = %omp.arraymap.body
  %25 = icmp eq i64 %22, 1
  br i1 %25, label %omp.type.to, label %omp.type.to.else

omp.type.to:                                      ; preds = %omp.type.alloc.else
  %26 = and i64 %21, -3
  br label %omp.type.end

omp.type.to.else:                                 ; preds = %omp.type.alloc.else
  %27 = icmp eq i64 %22, 2
  br i1 %27, label %omp.type.from, label %omp.type.end

omp.type.from:                                    ; preds = %omp.type.to.else
  %28 = and i64 %21, -2
  br label %omp.type.end

omp.type.end:                                     ; preds = %omp.type.from, %omp.type.to.else, %omp.type.to, %omp.type.alloc
  %omp.maptype = phi i64 [ %24, %omp.type.alloc ], [ %26, %omp.type.to ], [ %28, %omp.type.from ], [ %21, %omp.type.to.else ]
  call void @.omp_mapper._QQFmy_mapper1(ptr %0, ptr %omp.arraymap.ptrcurrent, ptr %array_offset, i64 4096, i64 %omp.maptype, ptr @3) #1
  %omp.arraymap.next = getelementptr %_QFTmytype, ptr %omp.arraymap.ptrcurrent, i32 1
  %omp.arraymap.isdone = icmp eq ptr %omp.arraymap.next, %7
  br i1 %omp.arraymap.isdone, label %omp.arraymap.exit, label %omp.arraymap.body

omp.arraymap.exit:                                ; preds = %omp.type.end
  %omp.arrayinit.isarray1 = icmp sgt i64 %6, 1
  %29 = and i64 %4, 8
  %.omp.array..del..delete = icmp ne i64 %29, 0
  %30 = and i1 %omp.arrayinit.isarray1, %.omp.array..del..delete
  br i1 %30, label %.omp.array..del, label %omp.done

.omp.array..del:                                  ; preds = %omp.arraymap.exit
  %31 = mul nuw i64 %6, 4096
  %32 = and i64 %4, -4
  %33 = or i64 %32, 512
  call void @__tgt_push_mapper_component(ptr %0, ptr %1, ptr %2, i64 %31, i64 %33, ptr %5)
  br label %omp.done

omp.done:                                         ; preds = %.omp.array..del, %omp.arraymap.exit, %omp.arraymap.head
  ret void
}

; Function Attrs: noinline nounwind
define internal void @.omp_mapper._QQFmy_mapper1(ptr noundef %0, ptr noundef %1, ptr noundef %2, i64 noundef %3, i64 noundef %4, ptr noundef %5) #0 {
entry:
  %6 = udiv exact i64 %3, 4096
  %7 = getelementptr %_QFTmytype, ptr %2, i64 %6
  %omp.arrayinit.isarray = icmp sgt i64 %6, 1
  %8 = and i64 %4, 8
  %9 = icmp ne ptr %1, %2
  %10 = and i64 %4, 16
  %11 = icmp ne i64 %10, 0
  %12 = and i1 %9, %11
  %13 = or i1 %omp.arrayinit.isarray, %12
  %.omp.array..init..delete = icmp eq i64 %8, 0
  %14 = and i1 %13, %.omp.array..init..delete
  br i1 %14, label %.omp.array..init, label %omp.arraymap.head

.omp.array..init:                                 ; preds = %entry
  %15 = mul nuw i64 %6, 4096
  %16 = and i64 %4, -4
  %17 = or i64 %16, 512
  call void @__tgt_push_mapper_component(ptr %0, ptr %1, ptr %2, i64 %15, i64 %17, ptr %5)
  br label %omp.arraymap.head

omp.arraymap.head:                                ; preds = %.omp.array..init, %entry
  %omp.arraymap.isempty = icmp eq ptr %2, %7
  br i1 %omp.arraymap.isempty, label %omp.done, label %omp.arraymap.body

omp.arraymap.body:                                ; preds = %omp.type.end, %omp.arraymap.head
  %omp.arraymap.ptrcurrent = phi ptr [ %2, %omp.arraymap.head ], [ %omp.arraymap.next, %omp.type.end ]
  %18 = getelementptr %_QFTmytype, ptr %omp.arraymap.ptrcurrent, i32 0, i32 0
  %array_offset = getelementptr inbounds [1024 x i32], ptr %18, i64 0, i64 0
  %19 = call i64 @__tgt_mapper_num_components(ptr %0)
  %20 = shl i64 %19, 48
  %21 = add nuw i64 1, %20
  %22 = and i64 %4, 3
  %23 = icmp eq i64 %22, 0
  br i1 %23, label %omp.type.alloc, label %omp.type.alloc.else

omp.type.alloc:                                   ; preds = %omp.arraymap.body
  %24 = and i64 %21, -4
  br label %omp.type.end

omp.type.alloc.else:                              ; preds = %omp.arraymap.body
  %25 = icmp eq i64 %22, 1
  br i1 %25, label %omp.type.to, label %omp.type.to.else

omp.type.to:                                      ; preds = %omp.type.alloc.else
  %26 = and i64 %21, -3
  br label %omp.type.end

omp.type.to.else:                                 ; preds = %omp.type.alloc.else
  %27 = icmp eq i64 %22, 2
  br i1 %27, label %omp.type.from, label %omp.type.end

omp.type.from:                                    ; preds = %omp.type.to.else
  %28 = and i64 %21, -2
  br label %omp.type.end

omp.type.end:                                     ; preds = %omp.type.from, %omp.type.to.else, %omp.type.to, %omp.type.alloc
  %omp.maptype = phi i64 [ %24, %omp.type.alloc ], [ %26, %omp.type.to ], [ %28, %omp.type.from ], [ %21, %omp.type.to.else ]
  call void @__tgt_push_mapper_component(ptr %0, ptr %omp.arraymap.ptrcurrent, ptr %array_offset, i64 4096, i64 %omp.maptype, ptr @5)
  %omp.arraymap.next = getelementptr %_QFTmytype, ptr %omp.arraymap.ptrcurrent, i32 1
  %omp.arraymap.isdone = icmp eq ptr %omp.arraymap.next, %7
  br i1 %omp.arraymap.isdone, label %omp.arraymap.exit, label %omp.arraymap.body

omp.arraymap.exit:                                ; preds = %omp.type.end
  %omp.arrayinit.isarray1 = icmp sgt i64 %6, 1
  %29 = and i64 %4, 8
  %.omp.array..del..delete = icmp ne i64 %29, 0
  %30 = and i1 %omp.arrayinit.isarray1, %.omp.array..del..delete
  br i1 %30, label %.omp.array..del, label %omp.done

.omp.array..del:                                  ; preds = %omp.arraymap.exit
  %31 = mul nuw i64 %6, 4096
  %32 = and i64 %4, -4
  %33 = or i64 %32, 512
  call void @__tgt_push_mapper_component(ptr %0, ptr %1, ptr %2, i64 %31, i64 %33, ptr %5)
  br label %omp.done

omp.done:                                         ; preds = %.omp.array..del, %omp.arraymap.exit, %omp.arraymap.head
  ret void
}
```

https://github.com/llvm/llvm-project/pull/124746


More information about the llvm-branch-commits mailing list