[flang] Introduce omp.target_allocmem and omp.target_freemem omp dialect ops. #145464
+484
−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. |
You can test this locally with the following command:git-clang-format --diff HEAD~1 HEAD --extensions cpp -- flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp flang/lib/Optimizer/Dialect/FIROps.cpp mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cppView the diff from clang-format here.diff --git a/flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp b/flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp
index a04c5d7eb..14cc7bb51 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp
@@ -142,8 +142,8 @@ static mlir::LLVM::LLVMFuncOp getOmpTargetAlloc(mlir::Operation *op) {
/*isVarArg=*/false));
}
-static mlir::Type
-convertObjectType(const fir::LLVMTypeConverter &converter, mlir::Type firType) {
+static mlir::Type convertObjectType(const fir::LLVMTypeConverter &converter,
+ mlir::Type firType) {
if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(firType))
return converter.convertBoxTypeAsStruct(boxTy);
return converter.convertType(firType);
@@ -189,8 +189,9 @@ computeElementDistance(mlir::Location loc, mlir::Type llvmObjectType,
}
static mlir::Value genTypeSizeInBytes(mlir::Location loc, mlir::Type idxTy,
- mlir::ConversionPatternRewriter &rewriter,
- mlir::Type llTy, const mlir::DataLayout &dataLayout) {
+ mlir::ConversionPatternRewriter &rewriter,
+ mlir::Type llTy,
+ const mlir::DataLayout &dataLayout) {
return computeElementDistance(loc, llTy, idxTy, rewriter, dataLayout);
}
@@ -224,8 +225,10 @@ genAllocationScaleSize(OP op, mlir::Type ity,
}
static mlir::Value integerCast(const fir::LLVMTypeConverter &converter,
- mlir::Location loc, mlir::ConversionPatternRewriter &rewriter,
- mlir::Type ty, mlir::Value val, bool fold = false) {
+ mlir::Location loc,
+ mlir::ConversionPatternRewriter &rewriter,
+ mlir::Type ty, mlir::Value val,
+ bool fold = false) {
auto valTy = val.getType();
// If the value was not yet lowered, lower its type so that it can
// be used in getPrimitiveTypeSizeInBits.
@@ -261,11 +264,13 @@ struct TargetAllocMemOpConversion
auto ity = lowerTy().indexType();
mlir::Type dataTy = fir::unwrapRefType(heapTy);
mlir::Type llvmObjectTy = convertObjectType(lowerTy(), dataTy);
- mlir::Type llvmPtrTy = mlir::LLVM::LLVMPointerType::get(allocmemOp.getContext(), 0);
+ mlir::Type llvmPtrTy =
+ mlir::LLVM::LLVMPointerType::get(allocmemOp.getContext(), 0);
if (fir::isRecordWithTypeParameters(fir::unwrapSequenceType(dataTy)))
TODO(loc, "omp.target_allocmem codegen of derived type with length "
"parameters");
- mlir::Value size = genTypeSizeInBytes(loc, ity, rewriter, llvmObjectTy, lowerTy().getDataLayout());
+ mlir::Value size = genTypeSizeInBytes(loc, ity, rewriter, llvmObjectTy,
+ lowerTy().getDataLayout());
if (auto scaleSize = genAllocationScaleSize(allocmemOp, ity, rewriter))
size = rewriter.create<mlir::LLVM::MulOp>(loc, ity, size, scaleSize);
for (mlir::Value opnd : adaptor.getOperands().drop_front())
@@ -281,7 +286,8 @@ struct TargetAllocMemOpConversion
loc, llvmPtrTy,
mlir::SmallVector<mlir::Value, 2>({size, allocmemOp.getDevice()}),
addLLVMOpBundleAttrs(rewriter, allocmemOp->getAttrs(), 2));
- rewriter.replaceOpWithNewOp<mlir::LLVM::PtrToIntOp>(allocmemOp, rewriter.getIntegerType(64), callOp.getResult());
+ rewriter.replaceOpWithNewOp<mlir::LLVM::PtrToIntOp>(
+ allocmemOp, rewriter.getIntegerType(64), callOp.getResult());
return mlir::success();
}
};
|
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
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