[flang] Introduce omp.target_allocmem and omp.target_freemem omp dialect ops. #145464
+490
−1
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
llvmbot
added
flang
|
@llvm/pr-subscribers-mlir-llvm @llvm/pr-subscribers-flang-codegen Author: Chaitanya (skc7) ChangesThis PR introduces two new ops in fir dialect, fir.omp_target_allocmem and fir.omp_target_freemem. Example: The work in this PR is C-P/inspired from @ivanradanov commit from coexecute implementation: Full diff: https://github.com/llvm/llvm-project/pull/145464.diff 5 Files Affected:
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index 8ac847dd7dd0a..93d617027e30b 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -517,6 +517,69 @@ def fir_ZeroOp : fir_OneResultOp<"zero_bits", [NoMemoryEffect]> {
let assemblyFormat = "type($intype) attr-dict";
}
+def fir_OmpTargetAllocMemOp : fir_Op<"omp_target_allocmem",
+ [MemoryEffects<[MemAlloc<DefaultResource>]>, AttrSizedOperandSegments]> {
+ let summary = "allocate storage on an openmp device for an object of a given type";
+
+ let description = [{
+ Creates a heap memory reference suitable for storing a value of the
+ given type, T. The heap refernce returned has type `!fir.heap<T>`.
+ The memory object is in an undefined state. `omp_target_allocmem` operations must
+ be paired with `omp_target_freemem` operations to avoid memory leaks.
+
+ ```
+ %device = arith.constant 0 : i32
+ %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3xi32>
+ ```
+ }];
+
+ let arguments = (ins
+ Arg<AnyIntegerType>:$device,
+ TypeAttr:$in_type,
+ OptionalAttr<StrAttr>:$uniq_name,
+ OptionalAttr<StrAttr>:$bindc_name,
+ Variadic<AnyIntegerType>:$typeparams,
+ Variadic<AnyIntegerType>:$shape
+ );
+ let results = (outs fir_HeapType);
+
+ let hasCustomAssemblyFormat = 1;
+ let hasVerifier = 1;
+
+ let extraClassDeclaration = [{
+ mlir::Type getAllocatedType();
+ bool hasLenParams() { return !getTypeparams().empty(); }
+ bool hasShapeOperands() { return !getShape().empty(); }
+ unsigned numLenParams() { return getTypeparams().size(); }
+ operand_range getLenParams() { return getTypeparams(); }
+ unsigned numShapeOperands() { return getShape().size(); }
+ operand_range getShapeOperands() { return getShape(); }
+ static mlir::Type getRefTy(mlir::Type ty);
+ }];
+}
+
+def fir_OmpTargetFreeMemOp : fir_Op<"omp_target_freemem",
+ [MemoryEffects<[MemFree]>]> {
+ let summary = "free a heap object on an openmp device";
+
+ let description = [{
+ Deallocates a heap memory reference that was allocated by an `omp_target_allocmem`.
+ The memory object that is deallocated is placed in an undefined state
+ after `fir.omp_target_freemem`.
+ ```
+ %device = arith.constant 0 : i32
+ %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3xi32>
+ fir.omp_target_freemem %device, %1 : i32, !fir.heap<!fir.array<3x3xi32>>
+ ```
+ }];
+
+ let arguments = (ins
+ Arg<AnyIntegerType, "", [MemFree]>:$device,
+ Arg<fir_HeapType, "", [MemFree]>:$heapref
+ );
+ let assemblyFormat = "$device `,` $heapref attr-dict `:` type($device) `,` qualified(type($heapref))";
+}
+
//===----------------------------------------------------------------------===//
// Terminator operations
//===----------------------------------------------------------------------===//
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index a3de3ae9d116a..042ade6b1e0a1 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -1168,6 +1168,105 @@ struct FreeMemOpConversion : public fir::FIROpConversion<fir::FreeMemOp> {
};
} // namespace
+static mlir::LLVM::LLVMFuncOp getOmpTargetAlloc(mlir::Operation *op) {
+ auto module = op->getParentOfType<mlir::ModuleOp>();
+ if (mlir::LLVM::LLVMFuncOp mallocFunc =
+ module.lookupSymbol<mlir::LLVM::LLVMFuncOp>("omp_target_alloc"))
+ return mallocFunc;
+ mlir::OpBuilder moduleBuilder(module.getBodyRegion());
+ auto i64Ty = mlir::IntegerType::get(module->getContext(), 64);
+ auto i32Ty = mlir::IntegerType::get(module->getContext(), 32);
+ return moduleBuilder.create<mlir::LLVM::LLVMFuncOp>(
+ moduleBuilder.getUnknownLoc(), "omp_target_alloc",
+ mlir::LLVM::LLVMFunctionType::get(
+ mlir::LLVM::LLVMPointerType::get(module->getContext()),
+ {i64Ty, i32Ty},
+ /*isVarArg=*/false));
+}
+
+namespace {
+struct OmpTargetAllocMemOpConversion
+ : public fir::FIROpConversion<fir::OmpTargetAllocMemOp> {
+ using FIROpConversion::FIROpConversion;
+
+ mlir::LogicalResult
+ matchAndRewrite(fir::OmpTargetAllocMemOp heap, OpAdaptor adaptor,
+ mlir::ConversionPatternRewriter &rewriter) const override {
+ mlir::Type heapTy = heap.getType();
+ mlir::LLVM::LLVMFuncOp mallocFunc = getOmpTargetAlloc(heap);
+ mlir::Location loc = heap.getLoc();
+ auto ity = lowerTy().indexType();
+ mlir::Type dataTy = fir::unwrapRefType(heapTy);
+ mlir::Type llvmObjectTy = convertObjectType(dataTy);
+ if (fir::isRecordWithTypeParameters(fir::unwrapSequenceType(dataTy)))
+ TODO(loc, "fir.omp_target_allocmem codegen of derived type with length "
+ "parameters");
+ mlir::Value size = genTypeSizeInBytes(loc, ity, rewriter, llvmObjectTy);
+ if (auto scaleSize = genAllocationScaleSize(heap, ity, rewriter))
+ size = rewriter.create<mlir::LLVM::MulOp>(loc, ity, size, scaleSize);
+ for (mlir::Value opnd : adaptor.getOperands().drop_front())
+ size = rewriter.create<mlir::LLVM::MulOp>(
+ loc, ity, size, integerCast(loc, rewriter, ity, opnd));
+ auto mallocTyWidth = lowerTy().getIndexTypeBitwidth();
+ auto mallocTy =
+ mlir::IntegerType::get(rewriter.getContext(), mallocTyWidth);
+ if (mallocTyWidth != ity.getIntOrFloatBitWidth())
+ size = integerCast(loc, rewriter, mallocTy, size);
+ heap->setAttr("callee", mlir::SymbolRefAttr::get(mallocFunc));
+ rewriter.replaceOpWithNewOp<mlir::LLVM::CallOp>(
+ heap, ::getLlvmPtrType(heap.getContext()),
+ mlir::SmallVector<mlir::Value, 2>({size, heap.getDevice()}),
+ addLLVMOpBundleAttrs(rewriter, heap->getAttrs(), 2));
+ return mlir::success();
+ }
+
+ /// Compute the allocation size in bytes of the element type of
+ /// \p llTy pointer type. The result is returned as a value of \p idxTy
+ /// integer type.
+ mlir::Value genTypeSizeInBytes(mlir::Location loc, mlir::Type idxTy,
+ mlir::ConversionPatternRewriter &rewriter,
+ mlir::Type llTy) const {
+ return computeElementDistance(loc, llTy, idxTy, rewriter, getDataLayout());
+ }
+};
+} // namespace
+
+static mlir::LLVM::LLVMFuncOp getOmpTargetFree(mlir::Operation *op) {
+ auto module = op->getParentOfType<mlir::ModuleOp>();
+ if (mlir::LLVM::LLVMFuncOp freeFunc =
+ module.lookupSymbol<mlir::LLVM::LLVMFuncOp>("omp_target_free"))
+ return freeFunc;
+ mlir::OpBuilder moduleBuilder(module.getBodyRegion());
+ auto i32Ty = mlir::IntegerType::get(module->getContext(), 32);
+ return moduleBuilder.create<mlir::LLVM::LLVMFuncOp>(
+ moduleBuilder.getUnknownLoc(), "omp_target_free",
+ mlir::LLVM::LLVMFunctionType::get(
+ mlir::LLVM::LLVMVoidType::get(module->getContext()),
+ {getLlvmPtrType(module->getContext()), i32Ty},
+ /*isVarArg=*/false));
+}
+
+namespace {
+struct OmpTargetFreeMemOpConversion
+ : public fir::FIROpConversion<fir::OmpTargetFreeMemOp> {
+ using FIROpConversion::FIROpConversion;
+
+ mlir::LogicalResult
+ matchAndRewrite(fir::OmpTargetFreeMemOp freemem, OpAdaptor adaptor,
+ mlir::ConversionPatternRewriter &rewriter) const override {
+ mlir::LLVM::LLVMFuncOp freeFunc = getOmpTargetFree(freemem);
+ mlir::Location loc = freemem.getLoc();
+ freemem->setAttr("callee", mlir::SymbolRefAttr::get(freeFunc));
+ rewriter.create<mlir::LLVM::CallOp>(
+ loc, mlir::TypeRange{},
+ mlir::ValueRange{adaptor.getHeapref(), freemem.getDevice()},
+ addLLVMOpBundleAttrs(rewriter, freemem->getAttrs(), 2));
+ rewriter.eraseOp(freemem);
+ return mlir::success();
+ }
+};
+} // namespace
+
// Convert subcomponent array indices from column-major to row-major ordering.
static llvm::SmallVector<mlir::Value>
convertSubcomponentIndices(mlir::Location loc, mlir::Type eleTy,
@@ -4274,7 +4373,8 @@ void fir::populateFIRToLLVMConversionPatterns(
GlobalLenOpConversion, GlobalOpConversion, InsertOnRangeOpConversion,
IsPresentOpConversion, LenParamIndexOpConversion, LoadOpConversion,
LocalitySpecifierOpConversion, MulcOpConversion, NegcOpConversion,
- NoReassocOpConversion, SelectCaseOpConversion, SelectOpConversion,
+ NoReassocOpConversion, OmpTargetAllocMemOpConversion,
+ OmpTargetFreeMemOpConversion, SelectCaseOpConversion, SelectOpConversion,
SelectRankOpConversion, SelectTypeOpConversion, ShapeOpConversion,
ShapeShiftOpConversion, ShiftOpConversion, SliceOpConversion,
StoreOpConversion, StringLitOpConversion, SubcOpConversion,
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index ecfa2939e96a6..9335a4b041ac8 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -106,24 +106,38 @@ static bool verifyTypeParamCount(mlir::Type inType, unsigned numParams) {
return false;
}
-/// Parser shared by Alloca and Allocmem
-///
+/// Parser shared by Alloca, Allocmem and OmpTargetAllocmem
+/// boolean flag isTargetOp is used to identify omp_target_allocmem
/// operation ::= %res = (`fir.alloca` | `fir.allocmem`) $in_type
/// ( `(` $typeparams `)` )? ( `,` $shape )?
/// attr-dict-without-keyword
+/// operation ::= %res = (`fir.omp_target_alloca`) $device : devicetype,
+/// $in_type ( `(` $typeparams `)` )? ( `,` $shape )?
+/// attr-dict-without-keyword
template <typename FN>
-static mlir::ParseResult parseAllocatableOp(FN wrapResultType,
- mlir::OpAsmParser &parser,
- mlir::OperationState &result) {
+static mlir::ParseResult
+parseAllocatableOp(FN wrapResultType, mlir::OpAsmParser &parser,
+ mlir::OperationState &result, bool isTargetOp = false) {
+ auto &builder = parser.getBuilder();
+ bool hasOperands = false;
+ std::int32_t typeparamsSize = 0;
+ // Parse device number as a new operand
+ if (isTargetOp) {
+ mlir::OpAsmParser::UnresolvedOperand deviceOperand;
+ mlir::Type deviceType;
+ if (parser.parseOperand(deviceOperand) || parser.parseColonType(deviceType))
+ return mlir::failure();
+ if (parser.resolveOperand(deviceOperand, deviceType, result.operands))
+ return mlir::failure();
+ if (parser.parseComma())
+ return mlir::failure();
+ }
mlir::Type intype;
if (parser.parseType(intype))
return mlir::failure();
- auto &builder = parser.getBuilder();
result.addAttribute("in_type", mlir::TypeAttr::get(intype));
llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand> operands;
llvm::SmallVector<mlir::Type> typeVec;
- bool hasOperands = false;
- std::int32_t typeparamsSize = 0;
if (!parser.parseOptionalLParen()) {
// parse the LEN params of the derived type. (<params> : <types>)
if (parser.parseOperandList(operands, mlir::OpAsmParser::Delimiter::None) ||
@@ -147,13 +161,19 @@ static mlir::ParseResult parseAllocatableOp(FN wrapResultType,
parser.resolveOperands(operands, typeVec, parser.getNameLoc(),
result.operands))
return mlir::failure();
+
mlir::Type restype = wrapResultType(intype);
if (!restype) {
parser.emitError(parser.getNameLoc(), "invalid allocate type: ") << intype;
return mlir::failure();
}
- result.addAttribute("operandSegmentSizes", builder.getDenseI32ArrayAttr(
- {typeparamsSize, shapeSize}));
+ llvm::SmallVector<std::int32_t> segmentSizes;
+ if (isTargetOp)
+ segmentSizes.push_back(1);
+ segmentSizes.push_back(typeparamsSize);
+ segmentSizes.push_back(shapeSize);
+ result.addAttribute("operandSegmentSizes",
+ builder.getDenseI32ArrayAttr(segmentSizes));
if (parser.parseOptionalAttrDict(result.attributes) ||
parser.addTypeToList(restype, result.types))
return mlir::failure();
@@ -385,6 +405,56 @@ llvm::LogicalResult fir::AllocMemOp::verify() {
return mlir::success();
}
+//===----------------------------------------------------------------------===//
+// OmpTargetAllocMemOp
+//===----------------------------------------------------------------------===//
+
+mlir::Type fir::OmpTargetAllocMemOp::getAllocatedType() {
+ return mlir::cast<fir::HeapType>(getType()).getEleTy();
+}
+
+mlir::Type fir::OmpTargetAllocMemOp::getRefTy(mlir::Type ty) {
+ return fir::HeapType::get(ty);
+}
+
+mlir::ParseResult
+fir::OmpTargetAllocMemOp::parse(mlir::OpAsmParser &parser,
+ mlir::OperationState &result) {
+ return parseAllocatableOp(wrapAllocMemResultType, parser, result, true);
+}
+
+void fir::OmpTargetAllocMemOp::print(mlir::OpAsmPrinter &p) {
+ p << " ";
+ p.printOperand(getDevice());
+ p << " : ";
+ p << getDevice().getType();
+ p << ", ";
+ p << getInType();
+ if (!getTypeparams().empty()) {
+ p << '(' << getTypeparams() << " : " << getTypeparams().getTypes() << ')';
+ }
+ for (auto sh : getShape()) {
+ p << ", ";
+ p.printOperand(sh);
+ }
+ p.printOptionalAttrDict((*this)->getAttrs(),
+ {"in_type", "operandSegmentSizes"});
+}
+
+llvm::LogicalResult fir::OmpTargetAllocMemOp::verify() {
+ llvm::SmallVector<llvm::StringRef> visited;
+ if (verifyInType(getInType(), visited, numShapeOperands()))
+ return emitOpError("invalid type for allocation");
+ if (verifyTypeParamCount(getInType(), numLenParams()))
+ return emitOpError("LEN params do not correspond to type");
+ mlir::Type outType = getType();
+ if (!mlir::dyn_cast<fir::HeapType>(outType))
+ return emitOpError("must be a !fir.heap type");
+ if (fir::isa_unknown_size_box(fir::dyn_cast_ptrEleTy(outType)))
+ return emitOpError("cannot allocate !fir.box of unknown rank or type");
+ return mlir::success();
+}
+
//===----------------------------------------------------------------------===//
// ArrayCoorOp
//===----------------------------------------------------------------------===//
diff --git a/flang/test/Fir/omp_target_allocmem.fir b/flang/test/Fir/omp_target_allocmem.fir
new file mode 100644
index 0000000000000..5140c91c9510c
--- /dev/null
+++ b/flang/test/Fir/omp_target_allocmem.fir
@@ -0,0 +1,28 @@
+// RUN: %flang_fc1 -emit-llvm %s -o - | FileCheck %s
+
+// CHECK-LABEL: define ptr @omp_target_allocmem_array_of_nonchar(
+// CHECK: call ptr @omp_target_alloc(i64 36, i32 0)
+func.func @omp_target_allocmem_array_of_nonchar() -> !fir.heap<!fir.array<3x3xi32>> {
+ %device = arith.constant 0 : i32
+ %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3xi32>
+ return %1 : !fir.heap<!fir.array<3x3xi32>>
+}
+
+// CHECK-LABEL: define ptr @omp_target_allocmem_array_of_char(
+// CHECK: call ptr @omp_target_alloc(i64 90, i32 0)
+func.func @omp_target_allocmem_array_of_char() -> !fir.heap<!fir.array<3x3x!fir.char<1,10>>> {
+ %device = arith.constant 0 : i32
+ %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3x!fir.char<1,10>>
+ return %1 : !fir.heap<!fir.array<3x3x!fir.char<1,10>>>
+}
+
+// CHECK-LABEL: define ptr @omp_target_allocmem_array_of_dynchar(
+// CHECK-SAME: i32 %[[len:.*]])
+// CHECK: %[[mul1:.*]] = sext i32 %[[len]] to i64
+// CHECK: %[[mul2:.*]] = mul i64 9, %[[mul1]]
+// CHECK: call ptr @omp_target_alloc(i64 %[[mul2]], i32 0)
+func.func @omp_target_allocmem_array_of_dynchar(%l: i32) -> !fir.heap<!fir.array<3x3x!fir.char<1,?>>> {
+ %device = arith.constant 0 : i32
+ %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3x!fir.char<1,?>>(%l : i32)
+ return %1 : !fir.heap<!fir.array<3x3x!fir.char<1,?>>>
+}
diff --git a/flang/test/Fir/omp_target_freemem.fir b/flang/test/Fir/omp_target_freemem.fir
new file mode 100644
index 0000000000000..02e136076a9cf
--- /dev/null
+++ b/flang/test/Fir/omp_target_freemem.fir
@@ -0,0 +1,28 @@
+// RUN: %flang_fc1 -emit-llvm %s -o - | FileCheck %s
+
+// CHECK-LABEL: define void @omp_target_allocmem_array_of_nonchar(
+// CHECK: call void @omp_target_free(ptr {{.*}}, i32 0)
+func.func @omp_target_allocmem_array_of_nonchar() -> () {
+ %device = arith.constant 0 : i32
+ %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3xi32>
+ fir.omp_target_freemem %device, %1 : i32, !fir.heap<!fir.array<3x3xi32>>
+ return
+}
+
+// CHECK-LABEL: define void @omp_target_allocmem_array_of_char(
+// CHECK: call void @omp_target_free(ptr {{.*}}, i32 0)
+func.func @omp_target_allocmem_array_of_char() -> () {
+ %device = arith.constant 0 : i32
+ %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3x!fir.char<1,10>>
+ fir.omp_target_freemem %device, %1 : i32, !fir.heap<!fir.array<3x3x!fir.char<1,10>>>
+ return
+}
+
+// CHECK-LABEL: define void @omp_target_allocmem_array_of_dynchar(
+// CHECK: call void @omp_target_free(ptr {{.*}}, i32 0)
+func.func @omp_target_allocmem_array_of_dynchar(%l: i32) -> () {
+ %device = arith.constant 0 : i32
+ %1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3x!fir.char<1,?>>(%l : i32)
+ fir.omp_target_freemem %device, %1 : i32, !fir.heap<!fir.array<3x3x!fir.char<1,?>>>
+ return
+}
|
tblah
reviewed
Jun 26, 2025
|
Thanks @skc7 for the PR! I have a few concerns about this approach, please let me know if I am missing anything:
I am probably missing some context here though so my question might not be totally informed. |
|
@tblah @ergawy This PR is a pre-requisite for workdistribute construct implementation and lowering in flang. I would try to experiment adding this op to openmp mlir dialect as suggested and let you know if it fits this workdistribute implementation. |
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
Hi @ergawy For implementing workdistribute lowering, in a new pass "lower-workdistribute" : LINK, Need to move fir.allocmem nested in omp.target outside it (or hoist it outside omp.target) Have updated the PR to implement the approach 1, to introduce omp.target_alloc_mem and omp.target_free_mem in openMP dialect. Also, their conversion to llvm ir has been implemented. Approach 2, is to directly lower to a LLVM dialect, call to "omp_target_alloc". So, went ahead with approach 1. Please let me know your comments on this. |
skc7
changed the title
ergawy
reviewed
Jul 9, 2025
| const mlir::DataLayout &dataLayout) { | ||
| llvm::TypeSize size = dataLayout.getTypeSize(llvmObjectType); | ||
| unsigned short alignment = dataLayout.getTypeABIAlignment(llvmObjectType); | ||
| std::int64_t distance = llvm::alignTo(size, alignment); |
There was a problem hiding this comment.
I think we are using a deprected conversion here (see: https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/Support/TypeSize.h#L360).
| } | ||
|
|
||
| static llvm::SmallVector<mlir::NamedAttribute> | ||
| addLLVMOpBundleAttrs(mlir::ConversionPatternRewriter &rewriter, |
There was a problem hiding this comment.
It seems this function does more than what the name suggests (i.e. copies old attributes, drops the old operandSegmentSizes and recreates it, and adds and empty op_bundle_sizes).
Should we do the attribute changes in-place (i.e. at the call site) instead? Since this function does not seem reusable in other scenarios.
There was a problem hiding this comment.
I see this is adapted from CodeGen.cpp though, so see my comment below.
| /*isVarArg=*/false)); | ||
| } | ||
|
|
||
| static mlir::Type convertObjectType(const fir::LLVMTypeConverter &converter, |
There was a problem hiding this comment.
It seems this function and some of the ones below are shared between CodeGen.cpp and CodeGenOpenMP.cpp. Can we move them to a shared location, e.g. flang/Optimizer/Support/Utils.h/.cpp?
| mlir::Type dataTy = fir::unwrapRefType(heapTy); | ||
| mlir::Type llvmObjectTy = convertObjectType(lowerTy(), dataTy); | ||
| mlir::Type llvmPtrTy = | ||
| mlir::LLVM::LLVMPointerType::get(allocmemOp.getContext(), 0); |
There was a problem hiding this comment.
Suggested change
| mlir::LLVM::LLVMPointerType::get(allocmemOp.getContext(), 0); | |
| mlir::LLVM::LLVMPointerType::get(allocmemOp.getContext()); |
0 is the default value for the address space.
Comment on lines
+282
to
+283
| if (mallocTyWidth != ity.getIntOrFloatBitWidth()) | ||
| size = integerCast(lowerTy(), loc, rewriter, mallocTy, size); |
There was a problem hiding this comment.
I think this if condition is not covered by the introduced tests.
| matchAndRewrite(mlir::omp::TargetAllocMemOp allocmemOp, OpAdaptor adaptor, | ||
| mlir::ConversionPatternRewriter &rewriter) const override { | ||
| mlir::Type heapTy = allocmemOp.getAllocatedType(); | ||
| mlir::LLVM::LLVMFuncOp mallocFunc = getOmpTargetAlloc(allocmemOp); |
There was a problem hiding this comment.
nit: move before use location below.
| } | ||
|
|
||
| // FIR Op specific conversion for TargetAllocMemOp | ||
| struct TargetAllocMemOpConversion |
There was a problem hiding this comment.
I think this conversion pattern should not replace omp.target_allocmem. Instead it should only convert the fir-specific allocation type to i64 and leave the lowering of omp.target_allocmem to call ptr @omp_target_alloc(...) to convertTargetAllocMemOp in OpenMPToLLVMIRTranslation.cpp. What we have now means that we do the same conversion in 2 different places. The preferred place would be OpenMPToLLVMIRTranslation.cpp since this is where all OpenMP to LLVM conversions are done.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR introduces two new ops in fir dialect, fir.omp_target_allocmem and fir.omp_target_freemem.
fir.omp_target_allocmem: Allocates heap memory on device. Will be lowered to omp_target_alloc call in llvm.
fir.omp_target_freemem: Deallocates heap memory on device. Will be lowered to omp+target_free call in llvm.
Example:
%device = arith.constant 0 : i32
%1 = fir.omp_target_allocmem %device : i32, !fir.array<3x3xi32>
fir.omp_target_freemem %device, %1 : i32, !fir.heap<!fir.array<3x3xi32>>
The work in this PR is C-P/inspired from @ivanradanov commit from coexecute implementation:
Add fir omp target alloc and free ops
Lower omp_target_{alloc,free} to llvm