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