src/txpr/txpr.cc
Namespaces
Source code
// 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/txpr/txpr.h"
#include <sstream>
#include <unordered_set>
namespace tenzr::txpr {
const char* OpCodeToString(OpCode op) noexcept {
switch (op) {
// Binary Elementwise
case OpCode::kAdd:
return "add";
case OpCode::kSub:
return "sub";
case OpCode::kMul:
return "mul";
case OpCode::kDiv:
return "div";
// Matrix Operations
case OpCode::kMatmul:
return "matmul";
case OpCode::kConvolution:
return "convolution";
// Shape & Strides
case OpCode::kReshape:
return "reshape";
case OpCode::kTranspose:
return "transpose";
case OpCode::kSlice:
return "slice";
case OpCode::kBroadcastTo:
return "broadcast_to";
// Reductions
case OpCode::kReduceSum:
return "reduce_sum";
case OpCode::kReduceProd:
return "reduce_prod";
case OpCode::kReduceMax:
return "reduce_max";
case OpCode::kReduceMin:
return "reduce_min";
case OpCode::kReduceWindow:
return "reduce_window";
// Indexing
case OpCode::kPad:
return "pad";
case OpCode::kGather:
return "gather";
case OpCode::kDynamicSlice:
return "dynamic_slice";
case OpCode::kDynamicUpdateSlice:
return "dynamic_update_slice";
case OpCode::kScatter:
return "scatter";
case OpCode::kArgSort:
return "argsort";
case OpCode::kConcatenate:
return "concatenate";
// Activations
case OpCode::kRelu:
return "relu";
case OpCode::kSigmoid:
return "sigmoid";
// Type Conversion
case OpCode::kConvert:
return "convert";
// Constants
case OpCode::kConstant:
return "constant";
// Custom Extensibility
case OpCode::kCustomOp:
return "custom_op";
// Distributed Collectives
case OpCode::kAllReduce:
return "all_reduce";
case OpCode::kAllGather:
return "all_gather";
case OpCode::kAllToAll:
return "all_to_all";
case OpCode::kReduceScatter:
return "reduce_scatter";
}
return "unknown";
}
VarId Txpr::AddVar(Shape shape, DType dtype, std::string name) {
VarId vid{static_cast<uint32_t>(variables_.size())};
variables_.push_back(VarDef{std::move(shape), dtype, std::move(name)});
return vid;
}
VarId Txpr::AppendInstruction(OpCode op, std::vector<VarId> in_vars, Shape out_shape, DType out_dtype, AttributePayload attrs) {
VarId out_var = AddVar(std::move(out_shape), out_dtype);
equations_.push_back(TxprEqn{op, std::move(in_vars), {out_var}, std::move(attrs)});
return out_var;
}
std::vector<VarId> Txpr::AppendMultiOutInstruction(OpCode op, std::vector<VarId> in_vars, const std::vector<Shape>& out_shapes,
const std::vector<DType>& out_dtypes, AttributePayload attrs) {
TENZR_CHECK(out_shapes.size() == out_dtypes.size(), "Mismatch in out_shapes and out_dtypes count");
std::vector<VarId> out_vars;
out_vars.reserve(out_shapes.size());
for (size_t i = 0; i < out_shapes.size(); ++i) {
out_vars.push_back(AddVar(out_shapes[i], out_dtypes[i]));
}
equations_.push_back(TxprEqn{op, std::move(in_vars), out_vars, std::move(attrs)});
return out_vars;
}
void Txpr::MarkInput(VarId var) {
TENZR_CHECK(var.IsValid() && var.id < variables_.size(), "Invalid input binder VarId");
in_binders_.push_back(var);
}
void Txpr::MarkOutput(VarId var) {
TENZR_CHECK(var.IsValid() && var.id < variables_.size(), "Invalid output binder VarId");
out_binders_.push_back(var);
}
bool Txpr::Validate(std::string* error_msg) const {
std::unordered_set<uint32_t> defined_vars;
// 1. Verify input binders are unique
for (const auto& in_v : in_binders_) {
if (!in_v.IsValid() || in_v.id >= variables_.size()) {
if (error_msg) {
*error_msg = "Invalid input binder ID " + std::to_string(in_v.id);
}
return false;
}
if (!defined_vars.insert(in_v.id).second) {
if (error_msg) {
*error_msg = "Duplicate input binder ID " + std::to_string(in_v.id);
}
return false;
}
}
// 2. Verify SSA property across instructions
for (size_t i = 0; i < equations_.size(); ++i) {
const auto& eqn = equations_[i];
// Verify input variables are dominated (defined prior to use)
for (const auto& in_v : eqn.in_vars) {
if (!in_v.IsValid() || in_v.id >= variables_.size()) {
if (error_msg) {
*error_msg = "Equation " + std::to_string(i) + " (" + OpCodeToString(eqn.op) + ") references invalid input VarId " + std::to_string(in_v.id);
}
return false;
}
if (defined_vars.find(in_v.id) == defined_vars.end()) {
if (error_msg) {
*error_msg =
"Equation " + std::to_string(i) + " (" + OpCodeToString(eqn.op) + ") uses undefined VarId " + std::to_string(in_v.id) + " (SSA violation)";
}
return false;
}
}
// Verify output variables are defined uniquely
for (const auto& out_v : eqn.out_vars) {
if (!out_v.IsValid() || out_v.id >= variables_.size()) {
if (error_msg) {
*error_msg = "Equation " + std::to_string(i) + " (" + OpCodeToString(eqn.op) + ") declares invalid output VarId " + std::to_string(out_v.id);
}
return false;
}
if (!defined_vars.insert(out_v.id).second) {
if (error_msg) {
*error_msg = "VarId " + std::to_string(out_v.id) + " is redefined in equation " + std::to_string(i) + " (SSA violation)";
}
return false;
}
}
}
// 3. Verify output binders are defined
for (const auto& out_v : out_binders_) {
if (!out_v.IsValid() || out_v.id >= variables_.size()) {
if (error_msg) {
*error_msg = "Invalid output binder ID " + std::to_string(out_v.id);
}
return false;
}
if (defined_vars.find(out_v.id) == defined_vars.end()) {
if (error_msg) {
*error_msg = "Output binder ID " + std::to_string(out_v.id) + " was never defined";
}
return false;
}
}
return true;
}
std::string Txpr::ToString() const {
std::ostringstream ss;
ss << "{ lambda ; ";
// Format input binders
for (size_t i = 0; i < in_binders_.size(); ++i) {
if (i > 0) {
ss << ", ";
}
const auto& v = variables_[in_binders_[i].id];
ss << "v" << in_binders_[i].id << ":" << tenzr::ToString(v.dtype) << v.shape.ToString();
}
ss << ". let\n";
// Format equations
for (const auto& eqn : equations_) {
ss << " ";
for (size_t i = 0; i < eqn.out_vars.size(); ++i) {
if (i > 0) {
ss << ", ";
}
const auto& ov = variables_[eqn.out_vars[i].id];
ss << "v" << eqn.out_vars[i].id << ":" << tenzr::ToString(ov.dtype) << ov.shape.ToString();
}
ss << " = " << OpCodeToString(eqn.op);
if (eqn.op == OpCode::kCustomOp && !eqn.attrs.custom_op_name.empty()) {
ss << "(" << eqn.attrs.custom_op_name << ")";
}
for (size_t i = 0; i < eqn.in_vars.size(); ++i) {
ss << " v" << eqn.in_vars[i].id;
}
if (!eqn.attrs.axes.empty()) {
ss << " [axes=";
for (size_t a = 0; a < eqn.attrs.axes.size(); ++a) {
if (a > 0) {
ss << ",";
}
ss << eqn.attrs.axes[a];
}
ss << "]";
}
if (!eqn.attrs.permutation.empty()) {
ss << " [perm=";
for (size_t p = 0; p < eqn.attrs.permutation.size(); ++p) {
if (p > 0) {
ss << ",";
}
ss << eqn.attrs.permutation[p];
}
ss << "]";
}
if (eqn.op == OpCode::kConvert) {
ss << " [to=" << tenzr::ToString(eqn.attrs.target_dtype) << "]";
}
if (!eqn.attrs.custom_op_name.empty()) {
ss << " [name=" << eqn.attrs.custom_op_name << "]";
}
ss << "\n";
}
// Format output binders
ss << " in (";
for (size_t i = 0; i < out_binders_.size(); ++i) {
if (i > 0) {
ss << ", ";
}
ss << "v" << out_binders_[i].id;
}
ss << ") }";
return ss.str();
}
} // namespace tenzr::txpr