xref: /llvm-project/flang/unittests/Runtime/CUDA/Allocatable.cpp (revision 4cb2a519db10f54815c8a4ccd5accbedc1cdfd07)
1 //===-- flang/unittests/Runtime/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 "gtest/gtest.h"
11 #include "../../../runtime/terminator.h"
12 #include "flang/Common/Fortran.h"
13 #include "flang/Runtime/CUDA/allocator.h"
14 #include "flang/Runtime/CUDA/common.h"
15 #include "flang/Runtime/CUDA/descriptor.h"
16 #include "flang/Runtime/allocator-registry.h"
17 
18 #include "cuda_runtime.h"
19 
20 using namespace Fortran::runtime;
21 using namespace Fortran::runtime::cuda;
22 
23 static 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 
29 TEST(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, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
46   EXPECT_TRUE(a->IsAllocated());
47   RTNAME(CUFDescriptorSync)(device_desc, a.get(), __FILE__, __LINE__);
48   cudaDeviceSynchronize();
49 
50   EXPECT_EQ(cudaSuccess, cudaGetLastError());
51 
52   RTNAME(AllocatableDeallocate)
53   (*a, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
54   EXPECT_FALSE(a->IsAllocated());
55 
56   RTNAME(CUFDescriptorSync)(device_desc, a.get(), __FILE__, __LINE__);
57   cudaDeviceSynchronize();
58 
59   EXPECT_EQ(cudaSuccess, cudaGetLastError());
60 }
61