sglang_v0.5.2/pytorch_2.8.0/third_party/NVTX/tests/RunTest.cpp

139 lines
4.3 KiB
C++

/*
* SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
* Licensed under the Apache License v2.0 with LLVM Exceptions.
* See https://nvidia.github.io/NVTX/LICENSE.txt for license information.
*/
#include "PathHelper.h"
#include "DllHelper.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <string>
static bool SetEnvVar(const char* name, const char* value)
{
#if defined(_WIN32)
auto result = _putenv_s(name, value);
#else
auto result = setenv(name, value, 1);
#endif
return result == 0;
}
static int MainInternal(int argc, const char** argv)
{
const std::string testArg("-t");
const std::string injectionArg("-i");
std::string testName;
std::string injectionName;
auto oldArgv = argv;
++argv;
while (*argv)
{
if (*argv == testArg ) { ++argv; if (*argv) testName = *argv; else return 100; }
else if (*argv == injectionArg) { ++argv; if (*argv) injectionName = *argv; else return 101; }
else break;
++argv;
}
argc -= static_cast<int>(argv - oldArgv);
if (testName.empty())
{
return 103;
}
printf("RunTest:\n");
std::string test = AbsolutePathToLibraryInCurrentProcessPath(testName);
printf(" - Using test: %s\n", test.c_str());
std::string injection;
if (!injectionName.empty())
{
const char* injectionVar = (sizeof(void*) == 8)
? "NVTX_INJECTION64_PATH"
: "NVTX_INJECTION32_PATH";
// Passing - for the injection means to use the test library as its own injection
injection = (injectionName == "-")
? test
: AbsolutePathToLibraryInCurrentProcessPath(injectionName);
bool success = SetEnvVar(injectionVar, injection.c_str());
if (!success) return 102;
}
printf(" - Using injection: %s\n", injection.empty() ? "<none>" : injection.c_str());
DLL_HANDLE hDll = DLL_OPEN(test.c_str());
if (!hDll) return 104;
using pfnRunTest_t = int(*)(int, const char**);
auto pfnRunTest = reinterpret_cast<pfnRunTest_t>(GET_DLL_FUNC(hDll, "RunTest"));
if (!pfnRunTest) return 105;
int result = pfnRunTest(argc, argv); // Forward remaining args
if (result) return result;
return 0;
}
int main(int argc, const char** argv)
{
int result = MainInternal(argc, argv);
if (result == 0)
{
printf("RunTest PASSED\n");
}
else
{
// For error codes known to this test driver, print useful error descriptions.
// Otherwise, rely on test to print information about errors.
switch (result)
{
case 100:
puts("RunTest: -t requires an argument, the base name of the library to use as a test");
break;
case 101:
puts("RunTest: -i requires an argument, the base name of the library to use as an injection");
break;
case 102:
puts("RunTest: Failed to set NVTX injection environment variable");
break;
case 103:
puts("RunTest: Missing required argument: -t <base name of library to use as a test>");
break;
case 104:
puts("RunTest: Test library failed to load");
#ifndef _WIN32
printf(" dlerror: %s\n", dlerror());
#endif
break;
case 105:
puts("RunTest: Test library loaded, but does not export required entry point RunTest");
break;
default:
printf("RunTest FAILED with return code: %d\n", result);
}
}
return result;
}