Valentin Clement (バレンタイン クレメン) f4d87c42a6
[flang][cuda] Add asyncId to allocate entry point (#134947)
2025-04-09 10:52:02 -07:00

61 lines
2.2 KiB
C++

//===-- unittests/Runtime/CUDA/Allocatable.cpp ------------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#include "flang/Runtime/allocatable.h"
#include "cuda_runtime.h"
#include "gtest/gtest.h"
#include "flang-rt/runtime/allocator-registry.h"
#include "flang-rt/runtime/descriptor.h"
#include "flang-rt/runtime/terminator.h"
#include "flang/Runtime/CUDA/allocator.h"
#include "flang/Runtime/CUDA/common.h"
#include "flang/Runtime/CUDA/descriptor.h"
#include "flang/Support/Fortran.h"
using namespace Fortran::runtime;
using namespace Fortran::runtime::cuda;
static OwningPtr<Descriptor> createAllocatable(
Fortran::common::TypeCategory tc, int kind, int rank = 1) {
return Descriptor::Create(TypeCode{tc, kind}, kind, nullptr, rank, nullptr,
CFI_attribute_allocatable);
}
TEST(AllocatableCUFTest, SimpleDeviceAllocatable) {
using Fortran::common::TypeCategory;
RTNAME(CUFRegisterAllocator)();
// REAL(4), DEVICE, ALLOCATABLE :: a(:)
auto a{createAllocatable(TypeCategory::Real, 4)};
a->SetAllocIdx(kDeviceAllocatorPos);
EXPECT_EQ((int)kDeviceAllocatorPos, a->GetAllocIdx());
EXPECT_FALSE(a->HasAddendum());
RTNAME(AllocatableSetBounds)(*a, 0, 1, 10);
// Emulate a device descriptor for the purpose of unit testing part of the
// code.
Descriptor *device_desc;
CUDA_REPORT_IF_ERROR(cudaMalloc(&device_desc, a->SizeInBytes()));
RTNAME(AllocatableAllocate)
(*a, kNoAsyncId, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
EXPECT_TRUE(a->IsAllocated());
RTNAME(CUFDescriptorSync)(device_desc, a.get(), __FILE__, __LINE__);
cudaDeviceSynchronize();
EXPECT_EQ(cudaSuccess, cudaGetLastError());
RTNAME(AllocatableDeallocate)
(*a, /*hasStat=*/false, /*errMsg=*/nullptr, __FILE__, __LINE__);
EXPECT_FALSE(a->IsAllocated());
RTNAME(CUFDescriptorSync)(device_desc, a.get(), __FILE__, __LINE__);
cudaDeviceSynchronize();
EXPECT_EQ(cudaSuccess, cudaGetLastError());
}