src/txpr/node_to_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/node_to_txpr.h"
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/frontend/node_traversal.h"
namespace tenzr::txpr {
namespace {
OpCode OpcodeFor(const std::shared_ptr<Node>& n) {
switch (n->Kind()) {
case NodeKind::kAdd:
return OpCode::kAdd;
case NodeKind::kMul:
return OpCode::kMul;
case NodeKind::kSub:
return OpCode::kSub;
case NodeKind::kDiv:
return OpCode::kDiv;
case NodeKind::kMatmul:
return OpCode::kMatmul;
case NodeKind::kConvolution:
return OpCode::kConvolution;
case NodeKind::kReshape:
return OpCode::kReshape;
case NodeKind::kTranspose:
return OpCode::kTranspose;
case NodeKind::kSlice:
return OpCode::kSlice;
case NodeKind::kPad:
return OpCode::kPad;
case NodeKind::kGather:
return OpCode::kGather;
case NodeKind::kScatter:
return OpCode::kScatter;
case NodeKind::kDynamicSlice:
return OpCode::kDynamicSlice;
case NodeKind::kDynamicUpdateSlice:
return OpCode::kDynamicUpdateSlice;
case NodeKind::kArgSort:
return OpCode::kArgSort;
case NodeKind::kConcatenate:
return OpCode::kConcatenate;
case NodeKind::kRelu:
return OpCode::kRelu;
case NodeKind::kSigmoid:
return OpCode::kSigmoid;
case NodeKind::kCustomOp:
return OpCode::kCustomOp;
case NodeKind::kConvert:
return OpCode::kConvert;
case NodeKind::kReduce: {
auto rn = std::static_pointer_cast<ReduceNode>(n);
switch (rn->Reduction()) {
case ReduceOp::kSum:
return OpCode::kReduceSum;
case ReduceOp::kProd:
return OpCode::kReduceProd;
case ReduceOp::kMax:
return OpCode::kReduceMax;
case ReduceOp::kMin:
return OpCode::kReduceMin;
}
return OpCode::kReduceSum;
}
case NodeKind::kReduceWindow:
return OpCode::kReduceWindow;
default:
TENZR_THROW("NodeToTxpr: unsupported node kind");
}
}
AttributePayload AttrsFor(const std::shared_ptr<Node>& n) {
AttributePayload a;
switch (n->Kind()) {
case NodeKind::kConvolution: {
auto c = std::static_pointer_cast<ConvolutionNode>(n);
a.window_strides = c->WindowStrides();
a.lhs_dilation = c->LhsDilation();
a.rhs_dilation = c->RhsDilation();
a.edge_padding_low = c->PaddingLow();
a.edge_padding_high = c->PaddingHigh();
break;
}
case NodeKind::kTranspose:
a.permutation = std::static_pointer_cast<TransposeNode>(n)->Permutation();
break;
case NodeKind::kSlice:
a.slices = std::static_pointer_cast<SliceNode>(n)->Slices();
break;
case NodeKind::kPad: {
auto p = std::static_pointer_cast<PadNode>(n);
a.edge_padding_low = p->EdgePaddingLow();
a.edge_padding_high = p->EdgePaddingHigh();
a.interior_padding = p->InteriorPadding();
break;
}
case NodeKind::kGather:
a.gather_axis = 0;
break;
case NodeKind::kScatter: {
auto s = std::static_pointer_cast<ScatterNode>(n);
a.scatter_axis = 0;
a.scatter_update = s->UpdateOp();
break;
}
case NodeKind::kDynamicSlice:
a.slice_sizes = std::static_pointer_cast<DynamicSliceNode>(n)->SliceSizes();
break;
case NodeKind::kArgSort: {
auto s = std::static_pointer_cast<ArgSortNode>(n);
a.sort_dimension = s->Axis();
a.sort_descending = s->Descending();
a.sort_return_values = s->ReturnValues();
break;
}
case NodeKind::kConcatenate:
a.concat_axis = std::static_pointer_cast<ConcatenateNode>(n)->Axis();
break;
case NodeKind::kCustomOp:
a.custom_op_name = n->OpName();
break;
case NodeKind::kConvert:
a.target_dtype = n->Dtype();
break;
case NodeKind::kReduce: {
auto r = std::static_pointer_cast<ReduceNode>(n);
a.axes = r->CanonicalDimensions();
a.keepdims = r->Keepdims();
break;
}
case NodeKind::kReduceWindow: {
auto r = std::static_pointer_cast<ReduceWindowNode>(n);
a.window_dimensions = r->WindowDimensions();
a.window_strides = r->WindowStrides();
a.edge_padding_low = r->PaddingLow();
a.edge_padding_high = r->PaddingHigh();
a.window_dilations = r->WindowDilations();
a.base_dilations.assign(r->WindowDimensions().size(), 1);
a.reduce_window_op = r->Reduction();
break;
}
default:
break;
}
return a;
}
std::vector<VarId> InVarsFor(const std::shared_ptr<Node>& n, const std::unordered_map<uint64_t, VarId>& node_to_var) {
if (n->Kind() == NodeKind::kArgSort) {
auto s = std::static_pointer_cast<ArgSortNode>(n);
if (s->ReturnValues()) {
// Values-only sort: a single-operand stablehlo.sort.
return {node_to_var.at(s->Values()->Id())};
}
}
std::vector<VarId> in_vars;
in_vars.reserve(n->Inputs().size());
for (const auto& in : n->Inputs()) {
in_vars.push_back(node_to_var.at(in->Id()));
}
return in_vars;
}
} // namespace
Txpr LowerNodeToTxpr(const std::shared_ptr<Node>& root) {
std::unordered_map<uint64_t, VarId> node_to_var;
return LowerNodeToTxpr(root, node_to_var);
}
Txpr LowerNodeToTxpr(const std::shared_ptr<Node>& root, std::unordered_map<uint64_t, VarId>& node_to_var) {
Txpr txpr;
if (!root) {
return txpr;
}
for (const auto& n : CollectNodesPostorder(root)) {
if (n->Kind() == NodeKind::kVariable) {
auto vn = std::static_pointer_cast<VariableNode>(n);
VarId v = txpr.AddVar(n->shape(), n->Dtype(), vn->VariableName());
txpr.MarkInput(v);
node_to_var[n->Id()] = v;
continue;
}
VarId out = txpr.AppendInstruction(OpcodeFor(n), InVarsFor(n, node_to_var), n->shape(), n->Dtype(), AttrsFor(n));
node_to_var[n->Id()] = out;
}
txpr.MarkOutput(node_to_var.at(root->Id()));
return txpr;
}
} // namespace tenzr::txpr