1//===-- unittests/Runtime/CUDA/Allocatable.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 "flang/Runtime/allocatable.h"
10#include "cuda_runtime.h"
11#include "gtest/gtest.h"
12#include "flang-rt/runtime/allocator-registry.h"
13#include "flang-rt/runtime/descriptor.h"
14#include "flang-rt/runtime/terminator.h"
15#include "flang/Runtime/CUDA/allocator.h"
16#include "flang/Runtime/CUDA/common.h"
17#include "flang/Runtime/CUDA/descriptor.h"
18#include "flang/Support/Fortran.h"
19
20using namespace Fortran::runtime;
21using namespace Fortran::runtime::cuda;
22
23static OwningPtr<Descriptor> createAllocatable(
24 Fortran::common::TypeCategory tc, int kind, int rank = 1) {
25 return Descriptor::Create(TypeCode{tc, kind}, kind, nullptr, rank, nullptr,
26 CFI_attribute_allocatable);
27}
28
29TEST(AllocatableCUFTest, SimpleDeviceAllocatable) {
30 using Fortran::common::TypeCategory;
31 RTNAME(CUFRegisterAllocator)();
32 // REAL(4), DEVICE, ALLOCATABLE :: a(:)
33 auto a{createAllocatable(TypeCategory::Real, 4)};
34 a->SetAllocIdx(kDeviceAllocatorPos);
35 EXPECT_EQ((int)kDeviceAllocatorPos, a->GetAllocIdx());
36 EXPECT_FALSE(a->HasAddendum());
37 RTNAME(AllocatableSetBounds)(*a, 0, 1, 10);
38
39 // Emulate a device descriptor for the purpose of unit testing part of the
40 // code.
41 Descriptor *device_desc;
42 CUDA_REPORT_IF_ERROR(cudaMalloc(&device_desc, a->SizeInBytes()));
43
44 RTNAME(AllocatableAllocate)
45 (*a, kNoAsyncObject, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__,
46 __LINE__);
47 EXPECT_TRUE(a->IsAllocated());
48 RTNAME(CUFDescriptorSync)(device_desc, a.get(), __FILE__, __LINE__);
49 cudaDeviceSynchronize();
50
51 EXPECT_EQ(cudaSuccess, cudaGetLastError());
52
53 RTNAME(AllocatableDeallocate)
54 (*a, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
55 EXPECT_FALSE(a->IsAllocated());
56
57 RTNAME(CUFDescriptorSync)(device_desc, a.get(), __FILE__, __LINE__);
58 cudaDeviceSynchronize();
59
60 EXPECT_EQ(cudaSuccess, cudaGetLastError());
61}
62
63TEST(AllocatableCUFTest, StreamDeviceAllocatable) {
64 using Fortran::common::TypeCategory;
65 RTNAME(CUFRegisterAllocator)();
66 // REAL(4), DEVICE, ALLOCATABLE :: a(:)
67 auto a{createAllocatable(TypeCategory::Real, 4)};
68 a->SetAllocIdx(kDeviceAllocatorPos);
69 EXPECT_EQ((int)kDeviceAllocatorPos, a->GetAllocIdx());
70 EXPECT_FALSE(a->HasAddendum());
71 RTNAME(AllocatableSetBounds)(*a, 0, 1, 10);
72
73 auto b{createAllocatable(TypeCategory::Real, 4)};
74 b->SetAllocIdx(kDeviceAllocatorPos);
75 EXPECT_EQ((int)kDeviceAllocatorPos, b->GetAllocIdx());
76 EXPECT_FALSE(b->HasAddendum());
77 RTNAME(AllocatableSetBounds)(*b, 0, 1, 20);
78
79 auto c{createAllocatable(TypeCategory::Real, 4)};
80 c->SetAllocIdx(kDeviceAllocatorPos);
81 EXPECT_EQ((int)kDeviceAllocatorPos, c->GetAllocIdx());
82 EXPECT_FALSE(b->HasAddendum());
83 RTNAME(AllocatableSetBounds)(*c, 0, 1, 100);
84
85 RTNAME(AllocatableAllocate)
86 (*a, /*asyncObject=*/nullptr, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__,
87 __LINE__);
88 EXPECT_TRUE(a->IsAllocated());
89 cudaDeviceSynchronize();
90 EXPECT_EQ(cudaSuccess, cudaGetLastError());
91
92 RTNAME(AllocatableAllocate)
93 (*b, /*asyncObject=*/nullptr, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__,
94 __LINE__);
95 EXPECT_TRUE(b->IsAllocated());
96 cudaDeviceSynchronize();
97 EXPECT_EQ(cudaSuccess, cudaGetLastError());
98
99 RTNAME(AllocatableAllocate)
100 (*c, /*asyncObject=*/nullptr, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__,
101 __LINE__);
102 EXPECT_TRUE(c->IsAllocated());
103 cudaDeviceSynchronize();
104 EXPECT_EQ(cudaSuccess, cudaGetLastError());
105
106 RTNAME(AllocatableDeallocate)
107 (*b, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
108 EXPECT_FALSE(b->IsAllocated());
109 cudaDeviceSynchronize();
110 EXPECT_EQ(cudaSuccess, cudaGetLastError());
111
112 RTNAME(AllocatableDeallocate)
113 (*a, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
114 EXPECT_FALSE(a->IsAllocated());
115 cudaDeviceSynchronize();
116 EXPECT_EQ(cudaSuccess, cudaGetLastError());
117
118 RTNAME(AllocatableDeallocate)
119 (*c, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
120 EXPECT_FALSE(c->IsAllocated());
121 cudaDeviceSynchronize();
122 EXPECT_EQ(cudaSuccess, cudaGetLastError());
123}
124

source code of flang-rt/unittests/Runtime/CUDA/Allocatable.cpp