chore: import upstream snapshot with attribution

2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
@@ -0,0 +1,188 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * 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.
+ */
+
+#include "sampleDevice.h"
+
+#include <iomanip>
+
+#if !defined(_WIN32)
+#include <dlfcn.h>
+#endif
+
+namespace sample
+{
+
+namespace
+{
+// Subset of NVML types/constants needed to query Confidential Compute state.
+// Declared locally so we do not introduce a build-time dependency on nvml.h
+// or libnvidia-ml; functions are resolved via dlopen at runtime.
+using NvmlReturnT = int32_t;
+constexpr NvmlReturnT kNVML_SUCCESS = 0;
+constexpr uint32_t kNVML_CC_FEATURE_ENABLED = 1;
+
+struct NvmlConfComputeSystemState
+{
+    uint32_t environment;
+    uint32_t ccFeature;
+    uint32_t devToolsMode;
+};
+
+using NvmlInitFn = NvmlReturnT (*)();
+using NvmlShutdownFn = NvmlReturnT (*)();
+using NvmlGetCcStateFn = NvmlReturnT (*)(NvmlConfComputeSystemState*);
+
+bool queryConfidentialCompute()
+{
+#if defined(_WIN32)
+    return false;
+#else
+    void* handle = dlopen("libnvidia-ml.so.1", RTLD_LAZY | RTLD_LOCAL);
+    if (handle == nullptr)
+    {
+        return false;
+    }
+
+    auto init = reinterpret_cast<NvmlInitFn>(dlsym(handle, "nvmlInit_v2"));
+    auto shutdown = reinterpret_cast<NvmlShutdownFn>(dlsym(handle, "nvmlShutdown"));
+    auto getState = reinterpret_cast<NvmlGetCcStateFn>(dlsym(handle, "nvmlSystemGetConfComputeState"));
+
+    bool enabled = false;
+    if (init != nullptr && shutdown != nullptr && getState != nullptr && init() == kNVML_SUCCESS)
+    {
+        NvmlConfComputeSystemState state{};
+        if (getState(&state) == kNVML_SUCCESS && state.ccFeature == kNVML_CC_FEATURE_ENABLED)
+        {
+            enabled = true;
+        }
+        shutdown();
+    }
+
+    dlclose(handle);
+    return enabled;
+#endif
+}
+} // namespace
+
+bool isConfidentialComputeEnabled()
+{
+    static bool const kCC_ENABLED = queryConfidentialCompute();
+    return kCC_ENABLED;
+}
+
+// Construct GPU UUID string in the same format as nvidia-smi does.
+std::string getUuidString(cudaUUID_t uuid)
+{
+    constexpr int32_t kUUID_SIZE = sizeof(cudaUUID_t);
+    static_assert(kUUID_SIZE == 16, "Unexpected size for cudaUUID_t!");
+
+    std::ostringstream ss;
+    std::vector<int32_t> const splits = {0, 4, 6, 8, 10, kUUID_SIZE};
+
+    ss << "GPU" << std::hex << std::setfill('0');
+    for (int32_t splitIdx = 0; splitIdx < static_cast<int32_t>(splits.size()) - 1; ++splitIdx)
+    {
+        ss << "-";
+        for (int32_t byteIdx = splits[splitIdx]; byteIdx < splits[splitIdx + 1]; ++byteIdx)
+        {
+            ss << std::setw(2) << +static_cast<uint8_t>(uuid.bytes[byteIdx]);
+        }
+    }
+    return ss.str();
+}
+
+void setCudaDevice(int32_t device, std::ostream& os)
+{
+    os << "=== Device Information ===" << std::endl;
+
+    // Get the number of visible GPUs.
+    int32_t nbDevices{-1};
+    CHECK(cudaGetDeviceCount(&nbDevices));
+
+    if (nbDevices <= 0)
+    {
+        os << "Cannot find any available devices (GPUs)!" << std::endl;
+        exit(EXIT_FAILURE);
+    }
+
+    // Print out the GPU name and PCIe bus ID of each GPU.
+    os << "Available Devices: " << std::endl;
+    cudaDeviceProp properties{};
+    for (int32_t deviceIdx = 0; deviceIdx < nbDevices; ++deviceIdx)
+    {
+        cudaDeviceProp tempProperties;
+        CHECK(cudaGetDeviceProperties(&tempProperties, deviceIdx));
+
+        // clang-format off
+        os << "  Device " << deviceIdx << ": \"" << tempProperties.name << "\" UUID: "
+           << getUuidString(tempProperties.uuid) << std::endl;
+        // clang-format on
+
+        // Record the properties of the desired GPU.
+        if (deviceIdx == device)
+        {
+            properties = tempProperties;
+        }
+    }
+
+    // Exit with error if the requested device ID does not exist.
+    if (device < 0 || device >= nbDevices)
+    {
+        os << "Cannot find device ID " << device << "!" << std::endl;
+        exit(EXIT_FAILURE);
+    }
+
+    // Set to the corresponding GPU.
+    CHECK(cudaSetDevice(device));
+
+    // clang-format off
+    os << "Selected Device: "      << properties.name                                               << std::endl;
+    os << "Selected Device ID: "   << device                                                        << std::endl;
+    os << "Selected Device UUID: " << getUuidString(properties.uuid)                                << std::endl;
+    os << "Compute Capability: "   << properties.major << "." << properties.minor                   << std::endl;
+    os << "SMs: "                  << properties.multiProcessorCount                                << std::endl;
+    os << "Device Global Memory: " << (properties.totalGlobalMem >> 20) << " MiB"                   << std::endl;
+    os << "Shared Memory per SM: " << (properties.sharedMemPerMultiprocessor >> 10) << " KiB"       << std::endl;
+    os << "Memory Bus Width: "     << properties.memoryBusWidth << " bits"
+                        << " (ECC " << (properties.ECCEnabled != 0 ? "enabled" : "disabled") << ")" << std::endl;
+    int32_t clockRate = 0;
+    int32_t memoryClockRate = 0;
+    CHECK(cudaDeviceGetAttribute(&clockRate, cudaDevAttrClockRate, device));
+    CHECK(cudaDeviceGetAttribute(&memoryClockRate, cudaDevAttrMemoryClockRate, device));
+    os << "Application Compute Clock Rate: "   << clockRate / 1000000.0F << " GHz"       << std::endl;
+    os << "Application Memory Clock Rate: "    << memoryClockRate / 1000000.0F << " GHz" << std::endl;
+    os << std::endl;
+    os << "Note: The application clock rates do not reflect the actual clock rates that the GPU is "
+                                                                         << "currently running at." << std::endl;
+    // clang-format on
+}
+
+int32_t getCudaDriverVersion()
+{
+    int32_t version{-1};
+    CHECK(cudaDriverGetVersion(&version));
+    return version;
+}
+
+int32_t getCudaRuntimeVersion()
+{
+    int32_t version{-1};
+    CHECK(cudaRuntimeGetVersion(&version));
+    return version;
+}
+
+} // namespace sample