Unverified Commit 568e1ad6 authored by Sergio Afonso's avatar Sergio Afonso Committed by GitHub
Browse files

[Flang][MLIR][OpenMP] Add explicit shared memory (de-)allocation ops (#161862) 

This patch introduces the `omp.alloc_shared_mem` and
`omp.free_shared_mem` operations to represent explicit allocations and
deallocations of shared memory across threads in a team, mirroring the
existing `omp.target_allocmem` and `omp.target_freemem`.

The `omp.alloc_shared_mem` op goes through the same Flang-specific
transformations as `omp.target_allocmem`, so that the size of the buffer
can be properly calculated when translating to LLVM IR.

The corresponding runtime functions produced for these new operations
are `__kmpc_alloc_shared` and `__kmpc_free_shared`, which previously
could only be created for implicit allocations (e.g. privatized and
reduction variables).
parent f6012dd7

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

+23 −0
Original line number Diff line number Diff line
@@ -3266,6 +3266,17 @@ public: 
  LLVM_ABI CallInst *createOMPFree(const LocationDescription &Loc, Value *Addr,

                                   Value *Allocator, std::string Name = "");



  /// Create a runtime call for kmpc_alloc_shared.

  ///

  /// \param Loc The insert and source location description.

  /// \param Size Size of allocated memory space.

  /// \param Name Name of call Instruction.

  ///

  /// \returns CallInst to the kmpc_alloc_shared call.

  LLVM_ABI CallInst *createOMPAllocShared(const LocationDescription &Loc,

                                          Value *Size,

                                          const Twine &Name = Twine(""));



  /// Create a runtime call for kmpc_alloc_shared.

  ///

  /// \param Loc The insert and source location description.
@@ -3277,6 +3288,18 @@ public: 
                                          Type *VarType,

                                          const Twine &Name = Twine(""));



  /// Create a runtime call for kmpc_free_shared.

  ///

  /// \param Loc The insert and source location description.

  /// \param Addr Value obtained from the corresponding kmpc_alloc_shared call.

  /// \param Size Size of allocated memory space.

  /// \param Name Name of call Instruction.

  ///

  /// \returns CallInst to the kmpc_free_shared call.

  LLVM_ABI CallInst *createOMPFreeShared(const LocationDescription &Loc,

                                         Value *Addr, Value *Size,

                                         const Twine &Name = Twine(""));



  /// Create a runtime call for kmpc_free_shared.

  ///

  /// \param Loc The insert and source location description.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

+21 −8
Original line number Diff line number Diff line
@@ -7930,32 +7930,45 @@ CallInst *OpenMPIRBuilder::createOMPFree(const LocationDescription &Loc, 
}



CallInst *OpenMPIRBuilder::createOMPAllocShared(const LocationDescription &Loc,

                                                Type *VarType,

                                                Value *Size,

                                                const Twine &Name) {

  IRBuilder<>::InsertPointGuard IPG(Builder);

  updateToLocation(Loc);



  const DataLayout &DL = M.getDataLayout();

  Value *Args[] = {Builder.getInt64(DL.getTypeAllocSize(VarType))};

  Value *Args[] = {Size};

  Function *Fn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_alloc_shared);

  CallInst *Call = Builder.CreateCall(Fn, Args, Name);

  Call->addRetAttr(

      Attribute::getWithAlignment(M.getContext(), DL.getPrefTypeAlign(Int64)));

  Call->addRetAttr(Attribute::getWithAlignment(

      M.getContext(), M.getDataLayout().getPrefTypeAlign(Int64)));

  return Call;

}



CallInst *OpenMPIRBuilder::createOMPAllocShared(const LocationDescription &Loc,

                                                Type *VarType,

                                                const Twine &Name) {

  return createOMPAllocShared(

      Loc, Builder.getInt64(M.getDataLayout().getTypeAllocSize(VarType)), Name);

}



CallInst *OpenMPIRBuilder::createOMPFreeShared(const LocationDescription &Loc,

                                               Value *Addr, Type *VarType,

                                               Value *Addr, Value *Size,

                                               const Twine &Name) {

  IRBuilder<>::InsertPointGuard IPG(Builder);

  updateToLocation(Loc);



  Value *Args[] = {

      Addr, Builder.getInt64(M.getDataLayout().getTypeAllocSize(VarType))};

  Value *Args[] = {Addr, Size};

  Function *Fn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_free_shared);

  return Builder.CreateCall(Fn, Args, Name);

}



CallInst *OpenMPIRBuilder::createOMPFreeShared(const LocationDescription &Loc,

                                               Value *Addr, Type *VarType,

                                               const Twine &Name) {

  return createOMPFreeShared(

      Loc, Addr, Builder.getInt64(M.getDataLayout().getTypeAllocSize(VarType)),

      Name);

}



CallInst *OpenMPIRBuilder::createOMPInteropInit(

    const LocationDescription &Loc, Value *InteropVar,

    omp::OMPInteropType InteropType, Value *Device, Value *NumDependences,

mlir/include/mlir/Dialect/OpenMP/OpenMPClauses.td

+37 −0
Original line number Diff line number Diff line
@@ -922,6 +922,43 @@ class OpenMP_MapClauseSkip< 


def OpenMP_MapClause : OpenMP_MapClauseSkip<>;



//===----------------------------------------------------------------------===//

// Not in the spec: Clause-like structure to memory allocation information.

//===----------------------------------------------------------------------===//



class OpenMP_MemAllocationSizeClauseSkip<

    bit traits = false, bit arguments = false, bit assemblyFormat = false,

    bit description = false, bit extraClassDeclaration = false

  > : OpenMP_Clause<traits, arguments, assemblyFormat, description,

                    extraClassDeclaration> {



  let arguments = (ins

    TypeAttr:$mem_elem_type,

    AnySignlessInteger:$mem_array_size,

    ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$mem_alignment

  );



  let reqAssemblyFormat = [{

    $mem_array_size `x` $mem_elem_type `:` `(` type($mem_array_size) `)`

  }];



  let optAssemblyFormat = [{

    `align` `(` $mem_alignment `)`

  }];



  let description = [{

    The `mem_elem_type` is the type of the object the memory allocation refers

    to. It is used to calculate the size of the allocation.



    The `mem_array_size` is the number of objects.



    The optional `mem_alignment` is used to specify the alignment for each

    element. If not set, the `DataLayout` defaults will be used instead.

  }];

}



def OpenMP_MemAllocationSizeClause : OpenMP_MemAllocationSizeClauseSkip<>;



//===----------------------------------------------------------------------===//

// V5.2: [15.8.1] `memory-order` clause set

//===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

+70 −0
Original line number Diff line number Diff line
@@ -2323,6 +2323,76 @@ def TargetFreeMemOp : OpenMP_Op<"target_freemem", 
  let assemblyFormat = "$device `,` $heapref attr-dict `:` type($device) `,` qualified(type($heapref))";

}



//===----------------------------------------------------------------------===//

// AllocSharedMemOp

//===----------------------------------------------------------------------===//



def AllocSharedMemOp : OpenMP_Op<"alloc_shared_mem", traits = [

    MemoryEffects<[MemAlloc<DefaultResource>]>

  ], clauses = [

    OpenMP_MemAllocationSizeClause

  ]> {

  let summary = "allocate storage on shared memory for objects of a given type";



  let description = [{

    Allocates memory shared across threads of a team for an object of the given

    type. Returns a pointer representing the allocated memory. The memory is

    uninitialized after allocation. Operations must be paired with

    `omp.free_shared` to avoid memory leaks.



    ```mlir

      // Allocate an i32 vector with %size elements and aligned to 8 bytes.

      %ptr_shared = omp.alloc_shared_mem %size x i32 : (i64) align(8) -> !llvm.ptr

      // ...

      omp.free_shared_mem [%size x i32 : (i64) align(8)] %ptr_shared : !llvm.ptr

    ```

  }] # clausesDescription;



  let results = (outs OpenMP_PointerLikeType);

  let assemblyFormat = clausesReqAssemblyFormat # " oilist(" #

    clausesOptAssemblyFormat # ") `->` type(results) attr-dict";



  let extraClassDeclaration = [{

    mlir::Type getAllocatedType() { return getMemElemTypeAttr().getValue(); }

  }];

  let hasVerifier = 1;

}



//===----------------------------------------------------------------------===//

// FreeSharedMemOp

//===----------------------------------------------------------------------===//



def FreeSharedMemOp : OpenMP_Op<"free_shared_mem", traits = [

    MemoryEffects<[MemFree]>

  ], clauses = [

    OpenMP_MemAllocationSizeClause

  ]> {

  let summary = "free shared memory";



  let description = [{

    Deallocates shared memory that was previously allocated by an

    `omp.alloc_shared_mem` operation. After this operation, the deallocated

    memory is in an undefined state and should not be accessed.



    ```mlir

      // Example of allocating and freeing shared memory.

      %ptr_shared = omp.alloc_shared_mem %size x i32 : (i64) -> !llvm.ptr

      // ...

      omp.free_shared_mem [%size x i32 : (i64)] %ptr_shared : !llvm.ptr

    ```



    The `heapref` operand represents the pointer to shared memory to be

    deallocated, previously returned by `omp.alloc_shared_mem`.

  }] # clausesDescription;



  let arguments = !con(clausesArgs, (ins

    Arg<OpenMP_PointerLikeType, "", [MemFree]>:$heapref

  ));

  let assemblyFormat = "` ` `[`" # clausesReqAssemblyFormat # " oilist(" #

    clausesOptAssemblyFormat # ") `]` $heapref `:` type($heapref) attr-dict";

  let hasVerifier = 1;

}



//===----------------------------------------------------------------------===//

// workdistribute Construct

//===----------------------------------------------------------------------===//

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

+24 −7
Original line number Diff line number Diff line
@@ -4690,17 +4690,34 @@ LogicalResult ScanOp::verify() { 
}



/// Verifies align clause in allocate directive

LogicalResult verifyAlignment(Operation &op,

                              std::optional<uint64_t> alignment) {

  if (alignment.has_value()) {

    if ((alignment.value() != 0) && !llvm::has_single_bit(alignment.value()))

      return op.emitError()

             << "ALIGN value : " << alignment.value() << " must be power of 2";

  }

  return success();

}



LogicalResult AllocateDirOp::verify() {

  std::optional<uint64_t> align = this->getAlign();

  return verifyAlignment(*getOperation(), getAlign());

}



//===----------------------------------------------------------------------===//

// AllocSharedMemOp

//===----------------------------------------------------------------------===//



  if (align.has_value()) {

    if ((align.value() > 0) && !llvm::has_single_bit(align.value()))

      return emitError() << "ALIGN value : " << align.value()

                         << " must be power of 2";

LogicalResult AllocSharedMemOp::verify() {

  return verifyAlignment(*getOperation(), getMemAlignment());

}



  return success();

//===----------------------------------------------------------------------===//

// FreeSharedMemOp

//===----------------------------------------------------------------------===//



LogicalResult FreeSharedMemOp::verify() {

  return verifyAlignment(*getOperation(), getMemAlignment());

}



//===----------------------------------------------------------------------===//