flang/unittests/Runtime/CUDA/AllocatorCUF.cpp

   1 //===-- flang/unittests/Runtime/AllocatableCUF.cpp ---------------*- C++-*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #include "gtest/gtest.h"
  10 #include "../../../runtime/terminator.h"
  11 #include "flang/Common/Fortran.h"
  12 #include "flang/Runtime/CUDA/allocator.h"
  13 #include "flang/Runtime/CUDA/descriptor.h"
  14 #include "flang/Runtime/allocatable.h"
  15 #include "flang/Runtime/allocator-registry.h"
  16
  17 #include "cuda_runtime.h"
  18
  19 using namespace Fortran::runtime;
  20 using namespace Fortran::runtime::cuda;
  21
  22 static OwningPtr<Descriptor> createAllocatable(
  23     Fortran::common::TypeCategory tc, int kind, int rank = 1) {
  24   return Descriptor::Create(TypeCode{tc, kind}, kind, nullptr, rank, nullptr,
  25       CFI_attribute_allocatable);
  26 }
  27
  28 TEST(AllocatableCUFTest, SimpleDeviceAllocate) {
  29   using Fortran::common::TypeCategory;
  30   RTNAME(CUFRegisterAllocator)();
  31   // REAL(4), DEVICE, ALLOCATABLE :: a(:)
  32   auto a{createAllocatable(TypeCategory::Real, 4)};
  33   a->SetAllocIdx(kDeviceAllocatorPos);
  34   EXPECT_EQ((int)kDeviceAllocatorPos, a->GetAllocIdx());
  35   EXPECT_FALSE(a->HasAddendum());
  36   RTNAME(AllocatableSetBounds)(*a, 0, 1, 10);
  37   RTNAME(AllocatableAllocate)
  38   (*a, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
  39   EXPECT_TRUE(a->IsAllocated());
  40   RTNAME(AllocatableDeallocate)
  41   (*a, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
  42   EXPECT_FALSE(a->IsAllocated());
  43 }
  44
  45 TEST(AllocatableCUFTest, SimplePinnedAllocate) {
  46   using Fortran::common::TypeCategory;
  47   RTNAME(CUFRegisterAllocator)();
  48   // INTEGER(4), PINNED, ALLOCATABLE :: a(:)
  49   auto a{createAllocatable(TypeCategory::Integer, 4)};
  50   EXPECT_FALSE(a->HasAddendum());
  51   a->SetAllocIdx(kPinnedAllocatorPos);
  52   EXPECT_EQ((int)kPinnedAllocatorPos, a->GetAllocIdx());
  53   EXPECT_FALSE(a->HasAddendum());
  54   RTNAME(AllocatableSetBounds)(*a, 0, 1, 10);
  55   RTNAME(AllocatableAllocate)
  56   (*a, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
  57   EXPECT_TRUE(a->IsAllocated());
  58   RTNAME(AllocatableDeallocate)
  59   (*a, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
  60   EXPECT_FALSE(a->IsAllocated());
  61 }
  62
  63 TEST(AllocatableCUFTest, DescriptorAllocationTest) {
  64   using Fortran::common::TypeCategory;
  65   RTNAME(CUFRegisterAllocator)();
  66   // REAL(4), DEVICE, ALLOCATABLE :: a(:)
  67   auto a{createAllocatable(TypeCategory::Real, 4)};
  68   Descriptor *desc = nullptr;
  69   desc = RTNAME(CUFAllocDesciptor)(a->SizeInBytes());
  70   EXPECT_TRUE(desc != nullptr);
  71   RTNAME(CUFFreeDesciptor)(desc);
  72 }