include/tenzr/execution/collective.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_EXECUTION_COLLECTIVE_H_
#define TENZR_EXECUTION_COLLECTIVE_H_

#include <cstdint>
#include <memory>
#include <string>
#include <string_view>
#include <vector>

#include "tenzr/frontend/tensor.h"

namespace tenzr::execution {

  enum class ReductionOp { kSum, kProduct, kMin, kMax };

  class Communicator {
   public:
    Communicator(int rank, int size);
    virtual ~Communicator() = default;

    int Rank() const { return rank_; }

    int Size() const { return size_; }

    virtual void Barrier() const;

    static Communicator& World();

   private:
    int rank_{0};
    int size_{1};
  };

  class DeviceMesh {
   public:
    explicit DeviceMesh(std::vector<int64_t> shape, std::vector<std::string> dim_names = {});

    int64_t NumDevices() const { return num_devices_; }

    int64_t Ndim() const { return static_cast<int64_t>(shape_.size()); }

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

    const std::vector<std::string>& DimNames() const { return dim_names_; }

    int64_t DimSize(std::string_view name) const;

    int64_t DimIndex(std::string_view name) const;

    std::vector<int64_t> DeviceCoords(int64_t device_id) const;

    int64_t DeviceId(const std::vector<int64_t>& coords) const;

   private:
    std::vector<int64_t> shape_;
    std::vector<std::string> dim_names_;
    int64_t num_devices_{1};
  };

  // Collective Operations

  Tensor AllReduce(const Tensor& input, ReductionOp op = ReductionOp::kSum, const Communicator* comm = nullptr);

  std::vector<Tensor> AllReduce(const std::vector<Tensor>& replica_tensors, ReductionOp op = ReductionOp::kSum);

  Tensor AllGather(const Tensor& input, int64_t axis = 0, const Communicator* comm = nullptr);

  Tensor AllGather(const std::vector<Tensor>& replica_tensors, int64_t axis = 0);

  std::vector<Tensor> ReduceScatter(const std::vector<Tensor>& replica_tensors, ReductionOp op = ReductionOp::kSum, int64_t axis = 0);

  std::vector<Tensor> Broadcast(const Tensor& root_tensor, int num_replicas, int root_rank = 0);

}  // namespace tenzr::execution

#endif  // TENZR_EXECUTION_COLLECTIVE_H_