// 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/execution/collective.h"
#include <algorithm>
#include <cmath>
#include <cstring>
#include <numeric>
#include "tenzr/common/error.h"
#include "tenzr/frontend/memory_pool.h"
namespace tenzr::execution {
// Communicator
Communicator::Communicator(int rank, int size) : rank_(rank), size_(size) {
TENZR_CHECK(size >= 1, "Communicator size must be at least 1");
TENZR_CHECK(rank >= 0 && rank < size, "Communicator rank out of bounds");
}
void Communicator::Barrier() const {
// In-process barrier point
}
Communicator& Communicator::World() {
static Communicator world_comm(0, 1);
return world_comm;
}
// DeviceMesh
DeviceMesh::DeviceMesh(std::vector<int64_t> shape, std::vector<std::string> dim_names) : shape_(std::move(shape)), dim_names_(std::move(dim_names)) {
TENZR_CHECK(!shape_.empty(), "DeviceMesh shape cannot be empty");
num_devices_ = 1;
for (int64_t d : shape_) {
TENZR_CHECK(d > 0, "DeviceMesh dimension sizes must be positive");
num_devices_ *= d;
}
if (dim_names_.size() < shape_.size()) {
for (size_t i = dim_names_.size(); i < shape_.size(); ++i) {
dim_names_.push_back("dim_" + std::to_string(i));
}
}
}
int64_t DeviceMesh::DimSize(std::string_view name) const {
int64_t idx = DimIndex(name);
return shape_[static_cast<size_t>(idx)];
}
int64_t DeviceMesh::DimIndex(std::string_view name) const {
for (size_t i = 0; i < dim_names_.size(); ++i) {
if (dim_names_[i] == name) {
return static_cast<int64_t>(i);
}
}
TENZR_THROW("DeviceMesh dimension name '" + std::string(name) + "' not found");
}
std::vector<int64_t> DeviceMesh::DeviceCoords(int64_t device_id) const {
TENZR_CHECK(device_id >= 0 && device_id < num_devices_, "Device ID out of bounds in DeviceMesh");
std::vector<int64_t> coords(shape_.size());
int64_t rem = device_id;
for (int i = static_cast<int>(shape_.size()) - 1; i >= 0; --i) {
coords[static_cast<size_t>(i)] = rem % shape_[static_cast<size_t>(i)];
rem /= shape_[static_cast<size_t>(i)];
}
return coords;
}
int64_t DeviceMesh::DeviceId(const std::vector<int64_t>& coords) const {
TENZR_CHECK(coords.size() == shape_.size(), "Coords rank mismatch in DeviceMesh");
int64_t id = 0;
for (size_t i = 0; i < coords.size(); ++i) {
TENZR_CHECK(coords[i] >= 0 && coords[i] < shape_[i], "Coordinate out of bounds in DeviceMesh");
id = id * shape_[i] + coords[i];
}
return id;
}
// Collective Operations
Tensor AllReduce(const Tensor& input, ReductionOp /*op*/, const Communicator* comm) {
if (!comm || comm->Size() <= 1) {
return input.Clone();
}
return input.Clone();
}
std::vector<Tensor> AllReduce(const std::vector<Tensor>& replica_tensors, ReductionOp op) {
TENZR_CHECK(!replica_tensors.empty(), "all_reduce requires at least one replica tensor");
const size_t n_replicas = replica_tensors.size();
if (n_replicas == 1) {
return {replica_tensors[0].Clone()};
}
const auto& base_shape = replica_tensors[0].shape();
const auto base_dtype = replica_tensors[0].Dtype();
const size_t num_elems = static_cast<size_t>(base_shape.NumElements());
for (size_t k = 1; k < n_replicas; ++k) {
TENZR_CHECK(replica_tensors[k].shape() == base_shape, "Shape mismatch across replicas in all_reduce");
TENZR_CHECK(replica_tensors[k].Dtype() == base_dtype, "DType mismatch across replicas in all_reduce");
}
TENZR_CHECK(base_dtype == DType::kFloat32, "all_reduce currently requires kFloat32");
std::vector<std::vector<float>> replica_data;
replica_data.reserve(n_replicas);
for (size_t k = 0; k < n_replicas; ++k) {
replica_data.push_back(replica_tensors[k].ToVector<float>());
}
std::vector<float> reduced(num_elems);
for (size_t i = 0; i < num_elems; ++i) {
float acc = replica_data[0][i];
for (size_t k = 1; k < n_replicas; ++k) {
float val = replica_data[k][i];
switch (op) {
case ReductionOp::kSum:
acc += val;
break;
case ReductionOp::kProduct:
acc *= val;
break;
case ReductionOp::kMin:
acc = std::min(acc, val);
break;
case ReductionOp::kMax:
acc = std::max(acc, val);
break;
}
}
reduced[i] = acc;
}
std::vector<Tensor> results;
results.reserve(n_replicas);
for (size_t k = 0; k < n_replicas; ++k) {
auto res = Tensor::Parameter(base_shape, base_dtype, "reduced_replica_" + std::to_string(k));
res.SetData<float>(reduced);
results.push_back(std::move(res));
}
return results;
}
Tensor AllGather(const Tensor& input, int64_t /*axis*/, const Communicator* comm) {
if (!comm || comm->Size() <= 1) {
return input.Clone();
}
return input.Clone();
}
Tensor AllGather(const std::vector<Tensor>& replica_tensors, int64_t axis) {
TENZR_CHECK(!replica_tensors.empty(), "all_gather requires at least one replica tensor");
const size_t n_replicas = replica_tensors.size();
if (n_replicas == 1) {
return replica_tensors[0].Clone();
}
const auto& base_shape = replica_tensors[0].shape();
const auto base_dtype = replica_tensors[0].Dtype();
int64_t rank = static_cast<int64_t>(base_shape.Rank());
if (axis < 0) {
axis += rank;
}
TENZR_CHECK(axis >= 0 && axis < rank, "all_gather axis out of range");
// Verify shape matching except for gathered axis
int64_t total_gathered_dim = 0;
for (size_t k = 0; k < n_replicas; ++k) {
TENZR_CHECK(replica_tensors[k].Dtype() == base_dtype, "DType mismatch in all_gather");
const auto& s = replica_tensors[k].shape();
TENZR_CHECK(static_cast<int64_t>(s.Rank()) == rank, "Rank mismatch in all_gather");
for (int64_t d = 0; d < rank; ++d) {
if (d != axis) {
TENZR_CHECK(s.Dim(static_cast<size_t>(d)) == base_shape.Dim(static_cast<size_t>(d)), "Shape mismatch on non-gathered axis in all_gather");
}
}
total_gathered_dim += s.Dim(static_cast<size_t>(axis));
}
std::vector<int64_t> out_dims = base_shape.Dims();
out_dims[static_cast<size_t>(axis)] = total_gathered_dim;
Shape out_shape(out_dims);
TENZR_CHECK(base_dtype == DType::kFloat32, "all_gather currently requires kFloat32");
if (axis == 0) {
// Simple sequential concatenation along outer dimension
std::vector<float> gathered_data;
gathered_data.reserve(static_cast<size_t>(out_shape.NumElements()));
for (const auto& rep : replica_tensors) {
auto vec = rep.ToVector<float>();
gathered_data.insert(gathered_data.end(), vec.begin(), vec.end());
}
auto out = Tensor::Parameter(out_shape, base_dtype, "all_gather_result");
out.SetData<float>(gathered_data);
return out;
}
// General axis gathering
int64_t outer_size = 1;
for (int64_t d = 0; d < axis; ++d) {
outer_size *= base_shape.Dim(static_cast<size_t>(d));
}
int64_t inner_size = 1;
for (int64_t d = axis + 1; d < rank; ++d) {
inner_size *= base_shape.Dim(static_cast<size_t>(d));
}
std::vector<float> gathered_data(static_cast<size_t>(out_shape.NumElements()));
for (int64_t o = 0; o < outer_size; ++o) {
int64_t out_axis_offset = 0;
for (const auto& rep : replica_tensors) {
auto rep_vec = rep.ToVector<float>();
int64_t rep_axis_dim = rep.shape().Dim(static_cast<size_t>(axis));
for (int64_t a = 0; a < rep_axis_dim; ++a) {
int64_t rep_idx = (o * rep_axis_dim + a) * inner_size;
int64_t out_idx = (o * total_gathered_dim + (out_axis_offset + a)) * inner_size;
std::memcpy(&gathered_data[static_cast<size_t>(out_idx)], &rep_vec[static_cast<size_t>(rep_idx)], static_cast<size_t>(inner_size) * sizeof(float));
}
out_axis_offset += rep_axis_dim;
}
}
auto out = Tensor::Parameter(out_shape, base_dtype, "all_gather_result");
out.SetData<float>(gathered_data);
return out;
}
std::vector<Tensor> ReduceScatter(const std::vector<Tensor>& replica_tensors, ReductionOp op, int64_t axis) {
TENZR_CHECK(!replica_tensors.empty(), "reduce_scatter requires at least one replica tensor");
size_t n_replicas = replica_tensors.size();
auto reduced_replicas = AllReduce(replica_tensors, op);
const auto& reduced = reduced_replicas[0];
const auto& shape = reduced.shape();
int64_t rank = static_cast<int64_t>(shape.Rank());
if (axis < 0) {
axis += rank;
}
TENZR_CHECK(axis >= 0 && axis < rank, "reduce_scatter axis out of range");
int64_t axis_dim = shape.Dim(static_cast<size_t>(axis));
TENZR_CHECK(axis_dim % static_cast<int64_t>(n_replicas) == 0, "Dimension along axis must be divisible by replica count in reduce_scatter");
int64_t chunk_dim = axis_dim / static_cast<int64_t>(n_replicas);
std::vector<Tensor> scattered;
scattered.reserve(n_replicas);
for (size_t k = 0; k < n_replicas; ++k) {
std::vector<Slice> slices;
slices.reserve(static_cast<size_t>(rank));
for (int64_t d = 0; d < rank; ++d) {
if (d == axis) {
int64_t start = static_cast<int64_t>(k) * chunk_dim;
int64_t stop = start + chunk_dim;
slices.push_back(Slice{start, stop, 1});
} else {
slices.push_back(Slice{0, shape.Dim(static_cast<size_t>(d)), 1});
}
}
scattered.push_back(reduced.slice(slices));
}
return scattered;
}
std::vector<Tensor> Broadcast(const Tensor& root_tensor, int num_replicas, int /*root_rank*/) {
TENZR_CHECK(num_replicas >= 1, "broadcast num_replicas must be at least 1");
std::vector<Tensor> results;
results.reserve(static_cast<size_t>(num_replicas));
for (int i = 0; i < num_replicas; ++i) {
results.push_back(root_tensor.Clone());
}
return results;
}
} // namespace tenzr::execution