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

#include <sstream>

namespace tenzr::sharding {

  std::string ShardingAnalysis::FormatReport() const {
    std::ostringstream oss;
    oss << "=== Sharding Analysis (Shardy vs GSPMD) ===\n";
    oss << "Mesh Declaration:   " << mesh_decl << "\n";
    oss << "Shardy Dialect:     " << shardy_attr << "\n";
    oss << "GSPMD Attribute:    " << gspmd_attr << "\n";
    oss << "Constraint Op:      " << shardy_constraint_op << "\n";
    oss << "Replicated:         " << (is_replicated ? "yes" : "no") << "\n";
    oss << "Unconstrained:      " << (is_unconstrained ? "yes" : "no") << "\n";
    return oss.str();
  }

  ShardingAnalysis AnalyzeSharding(const NamedSharding& sharding, const Shape& tensor_shape, std::string_view var_name, std::string_view tensor_type) {
    ShardingAnalysis sa;
    sa.mesh_decl = sharding.mesh().ToShardyMeshDecl();
    sa.shardy_attr = sharding.ToShardyString();
    sa.gspmd_attr = sharding.ToGspmdString(tensor_shape);
    sa.shardy_constraint_op = sharding.ToShardyConstraint("%res", var_name, tensor_type);

    sa.is_replicated = true;
    for (size_t i = 0; i < sharding.Spec().Rank(); ++i) {
      if (!sharding.Spec().IsReplicated(i)) {
        sa.is_replicated = false;
      }
      if (sharding.Spec().IsUnconstrained(i)) {
        sa.is_unconstrained = true;
      }
    }
    return sa;
  }

  std::string FormatShardyModuleHeader(const Mesh& mesh, std::string_view func_name, const std::vector<NamedSharding>& arg_shardings,
                                       const std::vector<Shape>& arg_shapes, std::string_view mesh_sym) {
    std::ostringstream oss;
    oss << "module @" << func_name << " attributes {mhlo.num_partitions = " << mesh.NumDevices() << " : i32, mhlo.num_replicas = 1 : i32} {\n";
    oss << "  " << mesh.ToShardyMeshDecl(mesh_sym) << "\n";
    oss << "  func.func public @main(";

    for (size_t i = 0; i < arg_shapes.size(); ++i) {
      if (i > 0) {
        oss << ", ";
      }
      oss << "%arg" << i << ": tensor<";
      for (size_t d = 0; d < arg_shapes[i].Rank(); ++d) {
        if (d > 0) {
          oss << "x";
        }
        oss << arg_shapes[i][d];
      }
      oss << "xf32>";
      if (i < arg_shardings.size()) {
        oss << " {sdy.sharding = " << arg_shardings[i].ToShardyString(mesh_sym) << "}";
      }
    }
    oss << ") {\n";
    return oss.str();
  }

}  // namespace tenzr::sharding