xref: /llvm-project/flang/unittests/Runtime/CUDA/AllocatorCUF.cpp (revision 4cb2a519db10f54815c8a4ccd5accbedc1cdfd07)
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(CUFAllocDescriptor)(a->SizeInBytes());
70   EXPECT_TRUE(desc != nullptr);
71   RTNAME(CUFFreeDescriptor)(desc);
72 }
73