src/frontend/sharding.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/frontend/sharding.h"

#include <sstream>

#include "tenzr/sharding/named_sharding.h"

namespace tenzr {

  ShardingSpec::ShardingSpec(const NamedSharding& ns) {
    mesh_shape_ = ns.mesh().shape();
    tensor_dim_to_mesh_axis_.resize(ns.Spec().Rank(), -1);
    bool any_tiled = false;
    for (size_t i = 0; i < ns.Spec().Rank(); ++i) {
      if (!ns.Spec().IsReplicated(i) && !ns.Spec().IsUnconstrained(i)) {
        const auto& dims = ns.Spec()[i];
        if (!dims.empty()) {
          int64_t ax_idx = ns.mesh().AxisIndex(dims[0]);
          tensor_dim_to_mesh_axis_[i] = static_cast<int>(ax_idx);
          any_tiled = true;
        }
      }
    }
    type_ = any_tiled ? ShardingType::kTiled : ShardingType::kReplicated;
  }

  ShardingSpec ShardingSpec::Replicated() {
    ShardingSpec spec;
    spec.type_ = ShardingType::kReplicated;
    return spec;
  }

  ShardingSpec ShardingSpec::Tiled(std::vector<int64_t> mesh_shape, std::vector<int> tensor_dim_to_mesh_axis) {
    ShardingSpec spec;
    spec.type_ = ShardingType::kTiled;
    spec.mesh_shape_ = std::move(mesh_shape);
    spec.tensor_dim_to_mesh_axis_ = std::move(tensor_dim_to_mesh_axis);
    return spec;
  }

  std::string ShardingSpec::ToString() const {
    if (type_ == ShardingType::kReplicated) {
      return "ShardingSpec(Replicated)";
    }
    std::ostringstream oss;
    oss << "ShardingSpec(Tiled, mesh=[";
    for (size_t i = 0; i < mesh_shape_.size(); ++i) {
      if (i > 0) {
        oss << ", ";
      }
      oss << mesh_shape_[i];
    }
    oss << "], axes=[";
    for (size_t i = 0; i < tensor_dim_to_mesh_axis_.size(); ++i) {
      if (i > 0) {
        oss << ", ";
      }
      oss << tensor_dim_to_mesh_axis_[i];
    }
    oss << "])";
    return oss.str();
  }

  std::string ShardingSpec::ToShardingAttribute() const {
    if (type_ == ShardingType::kReplicated) {
      return "{replicated}";
    }
    std::ostringstream oss;
    oss << "{devices=[";
    for (size_t i = 0; i < mesh_shape_.size(); ++i) {
      if (i > 0) {
        oss << ",";
      }
      oss << mesh_shape_[i];
    }
    oss << "]}";
    return oss.str();
  }

}  // namespace tenzr