175 lines · cpp
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-exception6//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 the40 // 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 125TEST(AllocatableAsyncTest, StreamDeviceAllocatable) {126 using Fortran::common::TypeCategory;127 RTNAME(CUFRegisterAllocator)();128 // REAL(4), DEVICE, ALLOCATABLE :: a(:)129 auto a{createAllocatable(TypeCategory::Real, 4)};130 a->SetAllocIdx(kDeviceAllocatorPos);131 EXPECT_EQ((int)kDeviceAllocatorPos, a->GetAllocIdx());132 EXPECT_FALSE(a->HasAddendum());133 RTNAME(AllocatableSetBounds)(*a, 0, 1, 10);134 135 cudaStream_t stream;136 cudaStreamCreate(&stream);137 EXPECT_EQ(cudaSuccess, cudaGetLastError());138 139 RTNAME(AllocatableAllocate)140 (*a, /*asyncObject=*/(int64_t *)&stream, /*hasStat=*/false,141 /*errMsg=*/nullptr, __FILE__, __LINE__);142 EXPECT_TRUE(a->IsAllocated());143 cudaDeviceSynchronize();144 EXPECT_EQ(cudaSuccess, cudaGetLastError());145 cudaStream_t s = RTDECL(CUFGetAssociatedStream)(a->raw().base_addr);146 EXPECT_EQ(s, stream);147 RTNAME(AllocatableDeallocate)148 (*a, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);149 EXPECT_FALSE(a->IsAllocated());150 cudaDeviceSynchronize();151 152 cudaStream_t defaultStream = 0;153 RTNAME(AllocatableAllocate)154 (*a, /*asyncObject=*/(int64_t *)&defaultStream, /*hasStat=*/false,155 /*errMsg=*/nullptr, __FILE__, __LINE__);156 EXPECT_TRUE(a->IsAllocated());157 cudaDeviceSynchronize();158 EXPECT_EQ(cudaSuccess, cudaGetLastError());159 cudaStream_t d = RTDECL(CUFGetAssociatedStream)(a->raw().base_addr);160 EXPECT_EQ(d, defaultStream);161 RTNAME(AllocatableDeallocate)162 (*a, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);163 EXPECT_FALSE(a->IsAllocated());164 cudaDeviceSynchronize();165 166 RTNAME(AllocatableAllocate)167 (*a, /*asyncObject=*/nullptr, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__,168 __LINE__);169 EXPECT_TRUE(a->IsAllocated());170 cudaDeviceSynchronize();171 EXPECT_EQ(cudaSuccess, cudaGetLastError());172 cudaStream_t empty = RTDECL(CUFGetAssociatedStream)(a->raw().base_addr);173 EXPECT_EQ(empty, nullptr);174}175