src/check/check.cc¶

// Copyright 2026 The Tenzr Project Authors
//
// 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
//
//      https://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 "tenzr/check/check.h"

#include <cmath>
#include <iomanip>
#include <sstream>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/runtime/hardware_manager.h"

namespace tenzr::check {

  void AssertShape(const Tensor& t, const Shape& expected, const char* file, int line) {
    if (t.shape() != expected) {
      std::ostringstream oss;
      oss << "[" << file << ":" << line << "] TENZR_ASSERT_SHAPE failed: expected " << expected.ToString() << " but got " << t.shape().ToString();
      TENZR_THROW(oss.str());
    }
  }

  void AssertDtype(const Tensor& t, DType expected, const char* file, int line) {
    if (t.Dtype() != expected) {
      std::ostringstream oss;
      oss << "[" << file << ":" << line << "] TENZR_ASSERT_DTYPE failed: expected " << ToString(expected) << " but got " << ToString(t.Dtype());
      TENZR_THROW(oss.str());
    }
  }

  void AssertDevice(const Tensor& t, const std::string& expected_device, const char* file, int line) {
    bool is_device = t.HasDeviceData();
    if (expected_device == "cpu") {
      // CPU tensor is valid either if it has host data or is on CPU device
      if (!t.HasData() && is_device) {
        auto& hw = HardwareManager::Instance();
        if (hw.IsInitialized() && hw.PlatformName() != "cpu") {
          std::ostringstream oss;
          oss << "[" << file << ":" << line << "] TENZR_ASSERT_DEVICE failed: expected cpu but tensor is on " << hw.PlatformName();
          TENZR_THROW(oss.str());
        }
      }
    } else if (expected_device == "cuda" || expected_device == "gpu") {
      if (!is_device) {
        std::ostringstream oss;
        oss << "[" << file << ":" << line << "] TENZR_ASSERT_DEVICE failed: expected device data on " << expected_device << " but tensor has no device buffer";
        TENZR_THROW(oss.str());
      }
    }
  }

  void AssertClose(const Tensor& actual, const Tensor& expected, double rtol, double atol, const char* file, int line) {
    AssertShape(actual, expected.shape(), file, line);
    AssertDtype(actual, expected.Dtype(), file, line);

    if (!actual.HasData() || !expected.HasData()) {
      std::ostringstream oss;
      oss << "[" << file << ":" << line << "] assert_close requires tensors with bound data";
      TENZR_THROW(oss.str());
    }

    auto act_vec = actual.ToVector<float>();
    auto exp_vec = expected.ToVector<float>();
    size_t n = act_vec.size();

    for (size_t i = 0; i < n; ++i) {
      float a = act_vec[i];
      float e = exp_vec[i];
      float diff = std::abs(a - e);
      float tol = static_cast<float>(atol + rtol * std::abs(static_cast<double>(e)));

      if (diff > tol || std::isnan(a) != std::isnan(e)) {
        std::ostringstream oss;
        oss << "[" << file << ":" << line << "] TENZR_ASSERT_CLOSE failed at index " << i << ": actual=" << a << ", expected=" << e << ", diff=" << diff
            << " > tol=" << tol;
        TENZR_THROW(oss.str());
      }
    }
  }

  void AssertClose(const Tensor& actual, const Tensor& expected, const char* file, int line) { AssertClose(actual, expected, 1.0e-4, 1.0e-5, file, line); }

  std::vector<DeviceTestTarget> HardwareHarness::AvailableTargets() {
    std::vector<DeviceTestTarget> targets;
    // CPU is always available
    targets.push_back({"cpu", true});

    auto& hw = HardwareManager::Instance();
#if TENZR_EXCEPTIONS_ENABLED
    try {
#endif
      if (!hw.IsInitialized()) {
        hw.Initialize();
      }
      if (hw.PlatformName() == "cuda" || hw.PlatformName() == "gpu") {
        targets.push_back({"cuda", true});
      }
#if TENZR_EXCEPTIONS_ENABLED
    } catch (...) {
      // Hardware discovery exception caught
    }
#endif

    return targets;
  }

  void HardwareHarness::RunOnAllAvailableDevices(const std::function<void(const std::string& target_name)>& test_fn) {
    auto targets = AvailableTargets();
    for (const auto& target : targets) {
      if (target.is_available) {
        test_fn(target.name);
      }
    }
  }

}  // namespace tenzr::check