105 lines · cpp
1//===-- unittests/Runtime/CUDA/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-exception6//7//===----------------------------------------------------------------------===//8 9#include "flang/Runtime/CUDA/memory.h"10#include "cuda_runtime.h"11#include "../tools.h"12#include "gtest/gtest.h"13#include "flang-rt/runtime/allocator-registry.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/Runtime/allocatable.h"19#include "flang/Support/Fortran.h"20 21using namespace Fortran::runtime;22using namespace Fortran::runtime::cuda;23 24TEST(MemoryCUFTest, SimpleAllocTramsferFree) {25 int *dev = (int *)RTNAME(CUFMemAlloc)(26 sizeof(int), kMemTypeDevice, __FILE__, __LINE__);27 EXPECT_TRUE(dev != 0);28 int host = 42;29 RTNAME(CUFDataTransferPtrPtr)30 ((void *)dev, (void *)&host, sizeof(int), kHostToDevice, __FILE__, __LINE__);31 host = 0;32 RTNAME(CUFDataTransferPtrPtr)33 ((void *)&host, (void *)dev, sizeof(int), kDeviceToHost, __FILE__, __LINE__);34 EXPECT_EQ(42, host);35 RTNAME(CUFMemFree)((void *)dev, kMemTypeDevice, __FILE__, __LINE__);36}37 38TEST(MemoryCUFTest, AllocZero) {39 int *dev = (int *)RTNAME(CUFMemAlloc)(0, kMemTypeDevice, __FILE__, __LINE__);40 EXPECT_TRUE(dev != 0);41 RTNAME(CUFMemFree)((void *)dev, kMemTypeDevice, __FILE__, __LINE__);42}43 44static OwningPtr<Descriptor> createAllocatable(45 Fortran::common::TypeCategory tc, int kind, int rank = 1) {46 return Descriptor::Create(TypeCode{tc, kind}, kind, nullptr, rank, nullptr,47 CFI_attribute_allocatable);48}49 50TEST(MemoryCUFTest, CUFDataTransferDescDesc) {51 using Fortran::common::TypeCategory;52 RTNAME(CUFRegisterAllocator)();53 // INTEGER(4), DEVICE, ALLOCATABLE :: a(:)54 auto dev{createAllocatable(TypeCategory::Integer, 4)};55 dev->SetAllocIdx(kDeviceAllocatorPos);56 EXPECT_EQ((int)kDeviceAllocatorPos, dev->GetAllocIdx());57 RTNAME(AllocatableSetBounds)(*dev, 0, 1, 10);58 RTNAME(AllocatableAllocate)59 (*dev, /*asyncObject=*/nullptr, /*hasStat=*/false, /*errMsg=*/nullptr,60 __FILE__, __LINE__);61 EXPECT_TRUE(dev->IsAllocated());62 63 // Create temp array to transfer to device.64 auto x{MakeArray<TypeCategory::Integer, 4>(std::vector<int>{10},65 std::vector<int32_t>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})};66 RTNAME(CUFDataTransferDescDesc)67 (dev.get(), x.get(), kHostToDevice, __FILE__, __LINE__);68 69 // Retrieve data from device.70 auto host{MakeArray<TypeCategory::Integer, 4>(std::vector<int>{10},71 std::vector<int32_t>{0, 0, 0, 0, 0, 0, 0, 0, 0, 0})};72 RTNAME(CUFDataTransferDescDesc)73 (host.get(), dev.get(), kDeviceToHost, __FILE__, __LINE__);74 75 for (unsigned i = 0; i < 10; ++i) {76 EXPECT_EQ(*host->ZeroBasedIndexedElement<std::int32_t>(i), (std::int32_t)i);77 }78}79 80TEST(MemoryCUFTest, CUFDataTransferDescDescDstNotAllocated) {81 using Fortran::common::TypeCategory;82 RTNAME(CUFRegisterAllocator)();83 // INTEGER(4), DEVICE, ALLOCATABLE :: a(:)84 auto dev{createAllocatable(TypeCategory::Integer, 4)};85 dev->SetAllocIdx(kDeviceAllocatorPos);86 EXPECT_EQ((int)kDeviceAllocatorPos, dev->GetAllocIdx());87 EXPECT_FALSE(dev->IsAllocated());88 89 // Create temp array to transfer to device.90 auto x{MakeArray<TypeCategory::Integer, 4>(std::vector<int>{10},91 std::vector<int32_t>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})};92 RTNAME(CUFDataTransferDescDesc)93 (dev.get(), x.get(), kHostToDevice, __FILE__, __LINE__);94 95 // Retrieve data from device.96 auto host{MakeArray<TypeCategory::Integer, 4>(std::vector<int>{10},97 std::vector<int32_t>{0, 0, 0, 0, 0, 0, 0, 0, 0, 0})};98 RTNAME(CUFDataTransferDescDesc)99 (host.get(), dev.get(), kDeviceToHost, __FILE__, __LINE__);100 101 for (unsigned i = 0; i < 10; ++i) {102 EXPECT_EQ(*host->ZeroBasedIndexedElement<std::int32_t>(i), (std::int32_t)i);103 }104}105