// 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/frontend/op_registry.h"
#include <algorithm>
#include <utility>
#include "tenzr/common/error.h"
namespace tenzr {
OpRegistry& OpRegistry::Instance() {
static OpRegistry reg;
return reg;
}
OpRegistry::OpRegistry() { RegisterBuiltins(); }
void OpRegistry::RegisterBuiltins() {
// 1. Built-in standard unary elementwise operations lowering to StableHLO
const std::vector<std::pair<std::string, std::string>> unary_builtins = {
{"tanh", "stablehlo.tanh"}, {"exp", "stablehlo.exponential"}, {"log", "stablehlo.log"}, {"sin", "stablehlo.sine"},
{"cos", "stablehlo.cosine"}, {"abs", "stablehlo.abs"}, {"neg", "stablehlo.negate"}, {"sqrt", "stablehlo.sqrt"},
{"rsqrt", "stablehlo.rsqrt"}, {"ceil", "stablehlo.ceil"}, {"floor", "stablehlo.floor"}, {"round", "stablehlo.round_nearest_afz"},
{"sign", "stablehlo.sign"}, {"cbrt", "stablehlo.cbrt"}, {"tan", "stablehlo.tan"}, {"expm1", "stablehlo.exponential_minus_one"},
{"log1p", "stablehlo.log_plus_one"}, {"logical_not", "stablehlo.not"},
};
for (const auto& [name, hlo_op] : unary_builtins) {
OpDef def;
def.name = name;
def.Category = OpCategory::kUnaryElementwise;
def.StablehloOp = hlo_op;
RegisterOp(std::move(def));
}
// Unary predicate op: is_finite
{
OpDef def;
def.name = "is_finite";
def.Category = OpCategory::kUnaryElementwise;
def.StablehloOp = "stablehlo.is_finite";
def.dtype_fn = [](const std::vector<DType>&) -> DType { return DType::kBool; };
RegisterOp(std::move(def));
}
// 2. Built-in binary elementwise operations lowering to StableHLO
const std::vector<std::pair<std::string, std::string>> binary_builtins = {
{"pow", "stablehlo.power"},
{"remainder", "stablehlo.remainder"},
{"atan2", "stablehlo.atan2"},
{"minimum", "stablehlo.minimum"},
{"maximum", "stablehlo.maximum"},
{"bitwise_and", "stablehlo.and"},
{"bitwise_or", "stablehlo.or"},
{"bitwise_xor", "stablehlo.xor"},
{"shift_left", "stablehlo.shift_left"},
{"shift_right_arithmetic", "stablehlo.shift_right_arithmetic"},
{"shift_right_logical", "stablehlo.shift_right_logical"},
};
for (const auto& [name, hlo_op] : binary_builtins) {
OpDef def;
def.name = name;
def.Category = OpCategory::kBinaryElementwise;
def.StablehloOp = hlo_op;
RegisterOp(std::move(def));
}
// 3. Tier 3: Clamp, Select, and Comparisons
{
OpDef def;
def.name = "clamp";
def.Category = OpCategory::kCustom;
def.StablehloOp = "stablehlo.clamp";
def.shape_fn = [](const std::vector<Shape>& in_shapes) -> Shape {
TENZR_CHECK(in_shapes.size() == 3, "clamp requires 3 operands (min, operand, max)");
return Shape::Broadcast(Shape::Broadcast(in_shapes[0], in_shapes[1]), in_shapes[2]);
};
def.dtype_fn = [](const std::vector<DType>& in_dtypes) -> DType {
TENZR_CHECK(in_dtypes.size() == 3, "clamp requires 3 operands");
return in_dtypes[1];
};
def.emit_fn = [](const std::string& res_name, const std::string& res_type, const std::vector<std::string>& in_names,
const std::vector<std::string>& /*in_types*/) -> std::string {
TENZR_CHECK(in_names.size() == 3, "clamp requires 3 operands");
return " " + res_name + " = stablehlo.clamp " + in_names[0] + ", " + in_names[1] + ", " + in_names[2] + " : " + res_type + "\n";
};
RegisterOp(std::move(def));
}
{
OpDef def;
def.name = "select";
def.Category = OpCategory::kCustom;
def.StablehloOp = "stablehlo.select";
def.shape_fn = [](const std::vector<Shape>& in_shapes) -> Shape {
TENZR_CHECK(in_shapes.size() == 3, "select requires 3 operands (pred, on_true, on_false)");
return Shape::Broadcast(in_shapes[1], in_shapes[2]);
};
def.dtype_fn = [](const std::vector<DType>& in_dtypes) -> DType {
TENZR_CHECK(in_dtypes.size() == 3, "select requires 3 operands");
return PromoteTypes(in_dtypes[1], in_dtypes[2]);
};
def.emit_fn = [](const std::string& res_name, const std::string& res_type, const std::vector<std::string>& in_names,
const std::vector<std::string>& in_types) -> std::string {
TENZR_CHECK(in_names.size() == 3, "select requires 3 operands");
return " " + res_name + " = stablehlo.select " + in_names[0] + ", " + in_names[1] + ", " + in_names[2] + " : " + in_types[0] + ", " + res_type +
"\n";
};
RegisterOp(std::move(def));
}
const std::vector<std::pair<std::string, std::string>> comparisons = {
{"compare_eq", "EQ"}, {"compare_ne", "NE"}, {"compare_ge", "GE"}, {"compare_gt", "GT"}, {"compare_le", "LE"}, {"compare_lt", "LT"},
};
for (const auto& [name, dir] : comparisons) {
OpDef def;
def.name = name;
def.Category = OpCategory::kCustom;
def.StablehloOp = "stablehlo.compare";
def.shape_fn = [](const std::vector<Shape>& in_shapes) -> Shape {
TENZR_CHECK(in_shapes.size() >= 2, "comparison requires 2 inputs");
return Shape::Broadcast(in_shapes[0], in_shapes[1]);
};
def.dtype_fn = [](const std::vector<DType>&) -> DType { return DType::kBool; };
def.emit_fn = [direction = dir](const std::string& res_name, const std::string& res_type, const std::vector<std::string>& in_names,
const std::vector<std::string>& in_types) -> std::string {
TENZR_CHECK(in_names.size() >= 2, "comparison requires 2 inputs");
std::string type_kind = "FLOAT";
if (in_types[0].find("i32") != std::string::npos || in_types[0].find("i64") != std::string::npos || in_types[0].find("i1") != std::string::npos) {
type_kind = "SIGNED";
}
return " " + res_name + " = stablehlo.compare " + direction + ", " + in_names[0] + ", " + in_names[1] + ", " + type_kind + " : (" + in_types[0] +
", " + in_types[1] + ") -> " + res_type + "\n";
};
RegisterOp(std::move(def));
}
}
void OpRegistry::RegisterOp(OpDef def) {
TENZR_CHECK(!def.name.empty(), "Operator name cannot be empty");
// Supply canonical defaults for UnaryElementwise if not explicitly provided
if (def.Category == OpCategory::kUnaryElementwise) {
if (!def.shape_fn) {
def.shape_fn = [](const std::vector<Shape>& in_shapes) -> Shape {
TENZR_CHECK(!in_shapes.empty(), "Unary elementwise operator requires at least 1 input");
return in_shapes[0];
};
}
if (!def.dtype_fn) {
def.dtype_fn = [](const std::vector<DType>& in_dtypes) -> DType {
TENZR_CHECK(!in_dtypes.empty(), "Unary elementwise operator requires at least 1 input");
return in_dtypes[0];
};
}
if (!def.emit_fn && !def.StablehloOp.empty()) {
def.emit_fn = [hlo = def.StablehloOp](const std::string& res_name, const std::string& res_type, const std::vector<std::string>& in_names,
const std::vector<std::string>& in_types) -> std::string {
TENZR_CHECK(!in_names.empty(), "Unary elementwise operator requires at least 1 input variable");
if (!in_types.empty() && in_types[0] != res_type) {
return " " + res_name + " = " + hlo + " " + in_names[0] + " : (" + in_types[0] + ") -> " + res_type + "\n";
}
return " " + res_name + " = " + hlo + " " + in_names[0] + " : " + res_type + "\n";
};
}
} else if (def.Category == OpCategory::kBinaryElementwise) {
if (!def.shape_fn) {
def.shape_fn = [](const std::vector<Shape>& in_shapes) -> Shape {
TENZR_CHECK(in_shapes.size() >= 2, "Binary elementwise operator requires at least 2 inputs");
return Shape::Broadcast(in_shapes[0], in_shapes[1]);
};
}
if (!def.dtype_fn) {
def.dtype_fn = [](const std::vector<DType>& in_dtypes) -> DType {
TENZR_CHECK(in_dtypes.size() >= 2, "Binary elementwise operator requires at least 2 inputs");
return PromoteTypes(in_dtypes[0], in_dtypes[1]);
};
}
if (!def.emit_fn && !def.StablehloOp.empty()) {
def.emit_fn = [hlo = def.StablehloOp](const std::string& res_name, const std::string& res_type, const std::vector<std::string>& in_names,
const std::vector<std::string>& /*in_types*/) -> std::string {
TENZR_CHECK(in_names.size() >= 2, "Binary elementwise operator requires at least 2 inputs");
return " " + res_name + " = " + hlo + " " + in_names[0] + ", " + in_names[1] + " : " + res_type + "\n";
};
}
}
registry_[def.name] = std::move(def);
}
const OpDef* OpRegistry::Lookup(std::string_view name) const {
auto it = registry_.find(std::string(name));
return (it != registry_.end()) ? &it->second : nullptr;
}
bool OpRegistry::HasOp(std::string_view name) const { return registry_.find(std::string(name)) != registry_.end(); }
std::vector<std::string> OpRegistry::RegisteredOps() const {
std::vector<std::string> names;
names.reserve(registry_.size());
for (const auto& [name, _] : registry_) {
names.push_back(name);
}
std::sort(names.begin(), names.end());
return names;
}
OpRegistrationBuilder::OpRegistrationBuilder(std::string name) { def_.name = std::move(name); }
OpRegistrationBuilder& OpRegistrationBuilder::Category(OpCategory cat) {
def_.Category = cat;
return *this;
}
OpRegistrationBuilder& OpRegistrationBuilder::StablehloOp(std::string op) {
def_.StablehloOp = std::move(op);
OpRegistry::Instance().RegisterOp(def_);
return *this;
}
OpRegistrationBuilder& OpRegistrationBuilder::shape(ShapeInferenceFn fn) {
def_.shape_fn = std::move(fn);
OpRegistry::Instance().RegisterOp(def_);
return *this;
}
OpRegistrationBuilder& OpRegistrationBuilder::Dtype(DTypeInferenceFn fn) {
def_.dtype_fn = std::move(fn);
OpRegistry::Instance().RegisterOp(def_);
return *this;
}
OpRegistrationBuilder& OpRegistrationBuilder::Emit(EmitFn fn) {
def_.emit_fn = std::move(fn);
OpRegistry::Instance().RegisterOp(def_);
return *this;
}
} // namespace tenzr