src/topology/torus_model.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/topology/torus_model.h"

#include <algorithm>
#include <cmath>
#include <iomanip>
#include <sstream>

namespace tenzr {

  std::string TorusDimensions::ToString() const {
    return std::to_string(kx) + "x" + std::to_string(ky) + "x" + std::to_string(kz) + " (" + std::to_string(TotalNodes()) + " nodes)";
  }

  TorusCostModel::TorusCostModel(TorusDimensions dims, NetworkParameters params) : dims_(dims), params_(params) {}

  double TorusCostModel::PredictTimeSeconds(uint64_t payload_bytes) const {
    double kx = static_cast<double>(dims_.kx);
    double ky = static_cast<double>(dims_.ky);
    double kz = static_cast<double>(dims_.kz);
    double n = static_cast<double>(payload_bytes);

    // Mathematical formula from Section 5:
    // T_torus ≈ 2 * ((Kx - 1) + (Ky - 1) + (Kz - 1)) * alpha +
    //           2 * ((Kx - 1)/Kx + (Ky - 1)/Ky + (Kz - 1)/Kz) * (N / beta)
    double latency_term = 2.0 * ((kx - 1.0) + (ky - 1.0) + (kz - 1.0)) * params_.alpha_latency;
    double bw_multiplier = (kx > 1.0 ? (kx - 1.0) / kx : 0.0) + (ky > 1.0 ? (ky - 1.0) / ky : 0.0) + (kz > 1.0 ? (kz - 1.0) / kz : 0.0);
    double bw_term = 2.0 * bw_multiplier * (n / params_.beta_bandwidth);

    return latency_term + bw_term;
  }

  double TorusCostModel::PredictRingTimeSeconds(uint64_t payload_bytes, double ring_bandwidth) const {
    double p = static_cast<double>(dims_.TotalNodes());
    if (p <= 1.0) {
      return 0.0;
    }
    double n = static_cast<double>(payload_bytes);

    if (ring_bandwidth <= 0.0) {
      // A single ring only has 1 physical link in each direction, whereas 3D torus uses 3 concurrent dimensions
      ring_bandwidth = params_.beta_bandwidth / 3.0;
    }

    // Standard ring: 2 * (P - 1) * alpha + 2 * (P - 1)/P * (N / beta_ring)
    double latency_term = 2.0 * (p - 1.0) * params_.alpha_latency;
    double bw_term = 2.0 * ((p - 1.0) / p) * (n / ring_bandwidth);
    return latency_term + bw_term;
  }

  double TorusCostModel::SpeedupOverRing(uint64_t payload_bytes, double ring_bandwidth) const {
    double t_torus = PredictTimeSeconds(payload_bytes);
    if (t_torus <= 0.0) {
      return 1.0;
    }
    double t_ring = PredictRingTimeSeconds(payload_bytes, ring_bandwidth);
    return t_ring / t_torus;
  }

  TorusDimensions TorusCostModel::OptimizeDimensions(int64_t total_nodes) {
    if (total_nodes <= 1) {
      return TorusDimensions{1, 1, 1};
    }

    int64_t best_kx = 1, best_ky = 1, best_kz = total_nodes;
    int64_t best_hop_sum = total_nodes * 3;

    // Find 3D factorization minimizing (kx + ky + kz)
    for (int64_t x = 1; x * x * x <= total_nodes * 4; ++x) {
      if (total_nodes % x != 0) {
        continue;
      }
      int64_t rem = total_nodes / x;
      for (int64_t y = 1; y * y <= rem * 4; ++y) {
        if (rem % y != 0) {
          continue;
        }
        int64_t z = rem / y;
        int64_t hop_sum = (x - 1) + (y - 1) + (z - 1);
        if (hop_sum < best_hop_sum) {
          best_hop_sum = hop_sum;
          best_kx = x;
          best_ky = y;
          best_kz = z;
        }
      }
    }

    return TorusDimensions{best_kx, best_ky, best_kz};
  }

  std::string TorusCostModel::GenerateReport(uint64_t payload_bytes) const {
    double t_torus = PredictTimeSeconds(payload_bytes);
    double t_ring = PredictRingTimeSeconds(payload_bytes);
    double speedup = (t_torus > 0.0) ? (t_ring / t_torus) : 1.0;

    std::ostringstream oss;
    oss << std::fixed << std::setprecision(4);
    oss << "=== Tenzr Torus Topology Cost Analysis ===\n";
    oss << "Dimensions: " << dims_.ToString() << "\n";
    oss << "Network alpha (latency): " << (params_.alpha_latency * 1.0e6) << " us\n";
    oss << "Network beta (bandwidth): " << (params_.beta_bandwidth / 1.0e9) << " GB/s\n";
    oss << "Payload size: " << (static_cast<double>(payload_bytes) / (1024.0 * 1024.0)) << " MB (" << payload_bytes << " bytes)\n";
    oss << "Predicted Torus Allreduce time: " << (t_torus * 1000.0) << " ms\n";
    oss << "Predicted Ring Allreduce time:  " << (t_ring * 1000.0) << " ms\n";
    oss << "Analytical Speedup: " << speedup << "x\n";
    return oss.str();
  }

}  // namespace tenzr