[omp][mlir] Introduce TargetAllocMem and TargetFreeMem ops in openMP mlir dialect

Author: skc7

flang/include/flang/Optimizer/Dialect/FIROps.td

@@ -517,69 +517,6 @@ 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
 //===----------------------------------------------------------------------===//

flang/lib/Optimizer/CodeGen/CodeGen.cpp

@@ -1224,105 +1224,6 @@ 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,
@@ -4469,8 +4370,7 @@ void fir::populateFIRToLLVMConversionPatterns(
       GlobalLenOpConversion, GlobalOpConversion, InsertOnRangeOpConversion,
       IsPresentOpConversion, LenParamIndexOpConversion, LoadOpConversion,
       LocalitySpecifierOpConversion, MulcOpConversion, NegcOpConversion,
-      NoReassocOpConversion, OmpTargetAllocMemOpConversion,
-      OmpTargetFreeMemOpConversion, SelectCaseOpConversion, SelectOpConversion,
+      NoReassocOpConversion, SelectCaseOpConversion, SelectOpConversion,
       SelectRankOpConversion, SelectTypeOpConversion, ShapeOpConversion,
       ShapeShiftOpConversion, ShiftOpConversion, SliceOpConversion,
       StoreOpConversion, StringLitOpConversion, SubcOpConversion,

flang/lib/Optimizer/CodeGen/CodeGenOpenMP.cpp

@@ -125,10 +125,171 @@ struct PrivateClauseOpConversion
     return mlir::success();
   }
 };
+
+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));
+}
+
+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);
+}
+
+static llvm::SmallVector<mlir::NamedAttribute>
+addLLVMOpBundleAttrs(mlir::ConversionPatternRewriter &rewriter,
+                     llvm::ArrayRef<mlir::NamedAttribute> attrs,
+                     int32_t numCallOperands) {
+  llvm::SmallVector<mlir::NamedAttribute> newAttrs;
+  newAttrs.reserve(attrs.size() + 2);
+
+  for (mlir::NamedAttribute attr : attrs) {
+    if (attr.getName() != "operandSegmentSizes")
+      newAttrs.push_back(attr);
+  }
+
+  newAttrs.push_back(rewriter.getNamedAttr(
+      "operandSegmentSizes",
+      rewriter.getDenseI32ArrayAttr({numCallOperands, 0})));
+  newAttrs.push_back(rewriter.getNamedAttr("op_bundle_sizes",
+                                           rewriter.getDenseI32ArrayAttr({})));
+  return newAttrs;
+}
+
+static mlir::LLVM::ConstantOp
+genConstantIndex(mlir::Location loc, mlir::Type ity,
+                 mlir::ConversionPatternRewriter &rewriter,
+                 std::int64_t offset) {
+  auto cattr = rewriter.getI64IntegerAttr(offset);
+  return rewriter.create<mlir::LLVM::ConstantOp>(loc, ity, cattr);
+}
+
+static mlir::Value
+computeElementDistance(mlir::Location loc, mlir::Type llvmObjectType,
+                       mlir::Type idxTy,
+                       mlir::ConversionPatternRewriter &rewriter,
+                       const mlir::DataLayout &dataLayout) {
+  llvm::TypeSize size = dataLayout.getTypeSize(llvmObjectType);
+  unsigned short alignment = dataLayout.getTypeABIAlignment(llvmObjectType);
+  std::int64_t distance = llvm::alignTo(size, alignment);
+  return genConstantIndex(loc, idxTy, rewriter, distance);
+}
+
+static mlir::Value genTypeSizeInBytes(mlir::Location loc, mlir::Type idxTy,
+                                 mlir::ConversionPatternRewriter &rewriter,
+                                 mlir::Type llTy, const mlir::DataLayout &dataLayout) {
+  return computeElementDistance(loc, llTy, idxTy, rewriter, dataLayout);
+}
+
+template <typename OP>
+static mlir::Value
+genAllocationScaleSize(OP op, mlir::Type ity,
+                       mlir::ConversionPatternRewriter &rewriter) {
+  mlir::Location loc = op.getLoc();
+  mlir::Type dataTy = op.getInType();
+  auto seqTy = mlir::dyn_cast<fir::SequenceType>(dataTy);
+  fir::SequenceType::Extent constSize = 1;
+  if (seqTy) {
+    int constRows = seqTy.getConstantRows();
+    const fir::SequenceType::ShapeRef &shape = seqTy.getShape();
+    if (constRows != static_cast<int>(shape.size())) {
+      for (auto extent : shape) {
+        if (constRows-- > 0)
+          continue;
+        if (extent != fir::SequenceType::getUnknownExtent())
+          constSize *= extent;
+      }
+    }
+  }
+
+  if (constSize != 1) {
+    mlir::Value constVal{
+        genConstantIndex(loc, ity, rewriter, constSize).getResult()};
+    return constVal;
+  }
+  return nullptr;
+}
+
+static mlir::Value integerCast(const fir::LLVMTypeConverter &converter,
+    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.
+  if (!mlir::isa<mlir::IntegerType>(valTy))
+    valTy = converter.convertType(valTy);
+  auto toSize = mlir::LLVM::getPrimitiveTypeSizeInBits(ty);
+  auto fromSize = mlir::LLVM::getPrimitiveTypeSizeInBits(valTy);
+  if (fold) {
+    if (toSize < fromSize)
+      return rewriter.createOrFold<mlir::LLVM::TruncOp>(loc, ty, val);
+    if (toSize > fromSize)
+      return rewriter.createOrFold<mlir::LLVM::SExtOp>(loc, ty, val);
+  } else {
+    if (toSize < fromSize)
+      return rewriter.create<mlir::LLVM::TruncOp>(loc, ty, val);
+    if (toSize > fromSize)
+      return rewriter.create<mlir::LLVM::SExtOp>(loc, ty, val);
+  }
+  return val;
+}
+
+// FIR Op specific conversion for TargetAllocMemOp
+struct TargetAllocMemOpConversion
+    : public OpenMPFIROpConversion<mlir::omp::TargetAllocMemOp> {
+  using OpenMPFIROpConversion::OpenMPFIROpConversion;
+
+  llvm::LogicalResult
+  matchAndRewrite(mlir::omp::TargetAllocMemOp allocmemOp, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    mlir::Type heapTy = allocmemOp.getAllocatedType();
+    mlir::LLVM::LLVMFuncOp mallocFunc = getOmpTargetAlloc(allocmemOp);
+    mlir::Location loc = allocmemOp.getLoc();
+    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);
+    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());
+    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())
+      size = rewriter.create<mlir::LLVM::MulOp>(
+          loc, ity, size, integerCast(lowerTy(), loc, rewriter, ity, opnd));
+    auto mallocTyWidth = lowerTy().getIndexTypeBitwidth();
+    auto mallocTy =
+        mlir::IntegerType::get(rewriter.getContext(), mallocTyWidth);
+    if (mallocTyWidth != ity.getIntOrFloatBitWidth())
+      size = integerCast(lowerTy(), loc, rewriter, mallocTy, size);
+    allocmemOp->setAttr("callee", mlir::SymbolRefAttr::get(mallocFunc));
+    auto callOp = rewriter.create<mlir::LLVM::CallOp>(
+        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());
+    return mlir::success();
+  }
+};
 } // namespace
 
 void fir::populateOpenMPFIRToLLVMConversionPatterns(
     const LLVMTypeConverter &converter, mlir::RewritePatternSet &patterns) {
   patterns.add<MapInfoOpConversion>(converter);
   patterns.add<PrivateClauseOpConversion>(converter);
+  patterns.add<TargetAllocMemOpConversion>(converter);
 }