op_registry.h¶

Namespaces¶

Name
tenzr

Classes¶

	Name
struct	tenzr::OpDef
class	tenzr::OpRegistry Process-wide registry of custom operator definitions, keyed by name.
class	tenzr::OpRegistrationBuilder

Defines¶

	Name
	TENZR_REGISTER_OP(name) Register an operator with the global OpRegistry at static-init time.

Macros Documentation¶

define TENZR_REGISTER_OP¶

#define TENZR_REGISTER_OP(
    name
)
static const ::tenzr::OpRegistrationBuilderTENZR_CONCAT(_tenzr_reg_, __LINE__) = ::tenzr::OpRegistrationBuilder(name)

Register an operator with the global OpRegistry at static-init time.

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.

#ifndef TENZR_FRONTEND_OP_REGISTRY_H_
#define TENZR_FRONTEND_OP_REGISTRY_H_

#include <functional>
#include <memory>
#include <string>
#include <string_view>
#include <unordered_map>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/common/macros.h"
#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/shape.h"

namespace tenzr {

  enum class OpCategory { kParameter, kConstant, kUnaryElementwise, kBinaryElementwise, kContraction, kReduction, kView, kCustom };

  using ShapeInferenceFn = std::function<Shape(const std::vector<Shape>&)>;

  using DTypeInferenceFn = std::function<DType(const std::vector<DType>&)>;

  using EmitFn = std::function<std::string(const std::string& result_name, const std::string& result_type, const std::vector<std::string>& input_names,
                                           const std::vector<std::string>& input_types)>;

  struct OpDef {
    std::string name;
    OpCategory Category{OpCategory::kUnaryElementwise};
    std::string StablehloOp;
    ShapeInferenceFn shape_fn;
    DTypeInferenceFn dtype_fn;
    EmitFn emit_fn;
  };

  class OpRegistry {
   public:
    static OpRegistry& Instance();

    void RegisterOp(OpDef def);

    const OpDef* Lookup(std::string_view name) const;

    bool HasOp(std::string_view name) const;

    std::vector<std::string> RegisteredOps() const;

   private:
    OpRegistry();
    void RegisterBuiltins();

    std::unordered_map<std::string, OpDef> registry_;
  };

  class OpRegistrationBuilder {
   public:
    explicit OpRegistrationBuilder(std::string name);

    OpRegistrationBuilder& Category(OpCategory cat);

    OpRegistrationBuilder& StablehloOp(std::string op);

    OpRegistrationBuilder& shape(ShapeInferenceFn fn);

    OpRegistrationBuilder& Dtype(DTypeInferenceFn fn);

    OpRegistrationBuilder& Emit(EmitFn fn);

   private:
    OpDef def_;
  };

}  // namespace tenzr

#define TENZR_REGISTER_OP(name) static const ::tenzr::OpRegistrationBuilder TENZR_CONCAT(_tenzr_reg_, __LINE__) = ::tenzr::OpRegistrationBuilder(name)

#endif  // TENZR_FRONTEND_OP_REGISTRY_H_