include/tenzr/sharding/mesh.h¶

// 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.

#ifndef TENZR_SHARDING_MESH_H_
#define TENZR_SHARDING_MESH_H_

#include <algorithm>
#include <cstdint>
#include <numeric>
#include <sstream>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/common/export.h"

namespace tenzr {

  class TENZR_API Mesh {
   public:
    Mesh() : devices_{0}, shape_{1}, axis_names_{"x"} {}

    Mesh(std::vector<int64_t> devices, std::vector<int64_t> shape, std::vector<std::string> axis_names)
        : devices_(std::move(devices)), shape_(std::move(shape)), axis_names_(std::move(axis_names)) {
      Validate();
    }

    Mesh(std::vector<int64_t> shape, std::vector<std::string> axis_names) : shape_(std::move(shape)), axis_names_(std::move(axis_names)) {
      int64_t total = 1;
      for (int64_t dim : shape_) {
        total *= dim;
      }
      devices_.resize(static_cast<size_t>(total));
      std::iota(devices_.begin(), devices_.end(), 0);
      Validate();
    }

    [[nodiscard]] size_t Ndim() const noexcept { return shape_.size(); }

    [[nodiscard]] int64_t NumDevices() const noexcept { return static_cast<int64_t>(devices_.size()); }

    [[nodiscard]] int64_t Size() const noexcept { return NumDevices(); }

    [[nodiscard]] bool Empty() const noexcept { return devices_.empty(); }

    [[nodiscard]] const std::vector<int64_t>& shape() const noexcept { return shape_; }

    [[nodiscard]] const std::vector<std::string>& AxisNames() const noexcept { return axis_names_; }

    [[nodiscard]] const std::vector<int64_t>& Devices() const noexcept { return devices_; }

    [[nodiscard]] int64_t AxisSize(std::string_view axis_name) const {
      for (size_t i = 0; i < axis_names_.size(); ++i) {
        if (axis_names_[i] == axis_name) {
          return shape_[i];
        }
      }
      TENZR_THROW("Mesh axis not found: " + std::string(axis_name));
    }

    [[nodiscard]] int64_t AxisIndex(std::string_view axis_name) const {
      for (size_t i = 0; i < axis_names_.size(); ++i) {
        if (axis_names_[i] == axis_name) {
          return static_cast<int64_t>(i);
        }
      }
      return -1;
    }

    [[nodiscard]] bool HasAxis(std::string_view axis_name) const noexcept { return AxisIndex(axis_name) >= 0; }

    [[nodiscard]] std::vector<int64_t> DeviceCoords(int64_t device_id) const {
      auto it = std::find(devices_.begin(), devices_.end(), device_id);
      if (it == devices_.end()) {
        TENZR_THROW("Device ID " + std::to_string(device_id) + " not in Mesh");
      }
      size_t flat_idx = static_cast<size_t>(std::distance(devices_.begin(), it));
      std::vector<int64_t> coords(shape_.size());
      for (int64_t d = static_cast<int64_t>(shape_.size()) - 1; d >= 0; --d) {
        coords[static_cast<size_t>(d)] = static_cast<int64_t>(flat_idx % static_cast<size_t>(shape_[static_cast<size_t>(d)]));
        flat_idx /= static_cast<size_t>(shape_[static_cast<size_t>(d)]);
      }
      return coords;
    }

    [[nodiscard]] int64_t DeviceAt(const std::vector<int64_t>& coords) const {
      if (coords.size() != shape_.size()) {
        TENZR_THROW("Coordinates rank mismatch with Mesh rank");
      }
      size_t flat_idx = 0;
      for (size_t i = 0; i < coords.size(); ++i) {
        if (coords[i] < 0 || coords[i] >= shape_[i]) {
          TENZR_THROW("Coordinate out of range for axis " + axis_names_[i]);
        }
        flat_idx = flat_idx * static_cast<size_t>(shape_[i]) + static_cast<size_t>(coords[i]);
      }
      return devices_[flat_idx];
    }

    [[nodiscard]] std::string ToString() const {
      std::ostringstream oss;
      oss << "Mesh(shape={";
      for (size_t i = 0; i < shape_.size(); ++i) {
        if (i > 0) {
          oss << ", ";
        }
        oss << "'" << axis_names_[i] << "': " << shape_[i];
      }
      oss << "}, num_devices=" << devices_.size() << ")";
      return oss.str();
    }

    [[nodiscard]] std::string ToShardyMeshDecl(std::string_view mesh_sym = "mesh") const {
      std::ostringstream oss;
      oss << "sdy.mesh @" << mesh_sym << " = <[";
      for (size_t i = 0; i < shape_.size(); ++i) {
        if (i > 0) {
          oss << ", ";
        }
        oss << "\"" << axis_names_[i] << "\"=" << shape_[i];
      }
      oss << "]>";
      return oss.str();
    }

    bool operator==(const Mesh& other) const noexcept { return shape_ == other.shape_ && axis_names_ == other.axis_names_ && devices_ == other.devices_; }

    bool operator!=(const Mesh& other) const noexcept { return !(*this == other); }

   private:
    void Validate() {
      if (shape_.size() != axis_names_.size()) {
        TENZR_THROW("Mesh shape rank must match axis_names count");
      }
      int64_t expected_total = 1;
      for (int64_t dim : shape_) {
        if (dim <= 0) {
          TENZR_THROW("Mesh dimensions must be strictly positive");
        }
        expected_total *= dim;
      }
      if (static_cast<int64_t>(devices_.size()) != expected_total) {
        TENZR_THROW("Mesh devices count (" + std::to_string(devices_.size()) + ") does not match grid product (" + std::to_string(expected_total) + ")");
      }
    }

    std::vector<int64_t> devices_;
    std::vector<int64_t> shape_;
    std::vector<std::string> axis_names_;
  };

}  // namespace tenzr

#endif  // TENZR_SHARDING_MESH_H_