xref: /llvm-project/flang/unittests/Runtime/CUDA/Memory.cpp (revision 4cb2a519db10f54815c8a4ccd5accbedc1cdfd07)
1 //===-- flang/unittests/Runtime/Memory.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/CUDA/memory.h"
10 #include "gtest/gtest.h"
11 #include "../../../runtime/terminator.h"
12 #include "../tools.h"
13 #include "flang/Common/Fortran.h"
14 #include "flang/Runtime/CUDA/allocator.h"
15 #include "flang/Runtime/CUDA/common.h"
16 #include "flang/Runtime/CUDA/descriptor.h"
17 #include "flang/Runtime/allocatable.h"
18 #include "flang/Runtime/allocator-registry.h"
19 
20 #include "cuda_runtime.h"
21 
22 using namespace Fortran::runtime;
23 using namespace Fortran::runtime::cuda;
24 
25 TEST(MemoryCUFTest, SimpleAllocTramsferFree) {
26   int *dev = (int *)RTNAME(CUFMemAlloc)(
27       sizeof(int), kMemTypeDevice, __FILE__, __LINE__);
28   EXPECT_TRUE(dev != 0);
29   int host = 42;
30   RTNAME(CUFDataTransferPtrPtr)
31   ((void *)dev, (void *)&host, sizeof(int), kHostToDevice, __FILE__, __LINE__);
32   host = 0;
33   RTNAME(CUFDataTransferPtrPtr)
34   ((void *)&host, (void *)dev, sizeof(int), kDeviceToHost, __FILE__, __LINE__);
35   EXPECT_EQ(42, host);
36   RTNAME(CUFMemFree)((void *)dev, kMemTypeDevice, __FILE__, __LINE__);
37 }
38 
39 static OwningPtr<Descriptor> createAllocatable(
40     Fortran::common::TypeCategory tc, int kind, int rank = 1) {
41   return Descriptor::Create(TypeCode{tc, kind}, kind, nullptr, rank, nullptr,
42       CFI_attribute_allocatable);
43 }
44 
45 TEST(MemoryCUFTest, CUFDataTransferDescDesc) {
46   using Fortran::common::TypeCategory;
47   RTNAME(CUFRegisterAllocator)();
48   // INTEGER(4), DEVICE, ALLOCATABLE :: a(:)
49   auto dev{createAllocatable(TypeCategory::Integer, 4)};
50   dev->SetAllocIdx(kDeviceAllocatorPos);
51   EXPECT_EQ((int)kDeviceAllocatorPos, dev->GetAllocIdx());
52   RTNAME(AllocatableSetBounds)(*dev, 0, 1, 10);
53   RTNAME(AllocatableAllocate)
54   (*dev, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
55   EXPECT_TRUE(dev->IsAllocated());
56 
57   // Create temp array to transfer to device.
58   auto x{MakeArray<TypeCategory::Integer, 4>(std::vector<int>{10},
59       std::vector<int32_t>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})};
60   RTNAME(CUFDataTransferDescDesc)
61   (dev.get(), x.get(), kHostToDevice, __FILE__, __LINE__);
62 
63   // Retrieve data from device.
64   auto host{MakeArray<TypeCategory::Integer, 4>(std::vector<int>{10},
65       std::vector<int32_t>{0, 0, 0, 0, 0, 0, 0, 0, 0, 0})};
66   RTNAME(CUFDataTransferDescDesc)
67   (host.get(), dev.get(), kDeviceToHost, __FILE__, __LINE__);
68 
69   for (unsigned i = 0; i < 10; ++i) {
70     EXPECT_EQ(*host->ZeroBasedIndexedElement<std::int32_t>(i), (std::int32_t)i);
71   }
72 }
73