include/tenzr/txpr/txpr.h¶

// 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_TXPR_TXPR_H_
#define TENZR_TXPR_TXPR_H_

#include <cstdint>
#include <limits>
#include <string>
#include <vector>

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

namespace tenzr::txpr {

  struct VarId {
    uint32_t id{std::numeric_limits<uint32_t>::max()};

    [[nodiscard]] constexpr bool IsValid() const noexcept { return id != std::numeric_limits<uint32_t>::max(); }

    [[nodiscard]] static constexpr VarId Invalid() noexcept { return VarId{std::numeric_limits<uint32_t>::max()}; }

    constexpr bool operator==(const VarId& other) const noexcept = default;
  };

  struct VarDef {
    Shape shape;
    DType dtype{DType::kFloat32};
    std::string name;
  };

  struct TxprVar {
    VarId id{VarId::Invalid()};
    Shape shape{};
    DType dtype{DType::kFloat32};

    [[nodiscard]] constexpr bool IsValid() const noexcept { return id.IsValid(); }

    bool operator==(const TxprVar& other) const noexcept = default;
  };

  enum class OpCode : uint16_t {
    kAdd,
    kSub,
    kMul,
    kDiv,

    kMatmul,
    kConvolution,

    kReshape,
    kTranspose,
    kSlice,
    kBroadcastTo,
    kPad,
    kScatter,
    kGather,
    kDynamicSlice,
    kDynamicUpdateSlice,
    kArgSort,
    kConcatenate,

    kReduceSum,
    kReduceProd,
    kReduceMax,
    kReduceMin,
    kReduceWindow,

    kRelu,
    kSigmoid,

    kConvert,

    kConstant,

    kCustomOp,

    kAllReduce,
    kAllGather,
    kAllToAll,
    kReduceScatter
  };

  [[nodiscard]] const char* OpCodeToString(OpCode op) noexcept;

  struct AttributePayload {
    std::vector<int64_t> axes;
    std::vector<int64_t> permutation;
    std::vector<Slice> slices;
    std::vector<int64_t> edge_padding_low;
    std::vector<int64_t> edge_padding_high;
    std::vector<int64_t> interior_padding;
    int64_t scatter_axis{0};
    int64_t gather_axis{0};
    std::string scatter_update{"add"};
    int64_t sort_dimension{0};
    bool sort_descending{false};
    bool sort_return_values{false};
    int64_t concat_axis{0};
    DType target_dtype{DType::kFloat32};
    bool keepdims{false};
    std::string custom_op_name;
    double scalar_value{0.0};
    std::vector<double> tensor_data;
    std::vector<int64_t> int_tensor_data;
    int64_t iota_dimension{-1};
    std::vector<int64_t> window_dimensions;
    std::vector<int64_t> window_strides;
    std::vector<int64_t> base_dilations;
    std::vector<int64_t> window_dilations;
    std::string reduce_window_op{"add"};
    std::vector<int64_t> lhs_dilation;
    std::vector<int64_t> rhs_dilation;
    std::vector<int64_t> conv_lhs_spec;
    std::vector<int64_t> conv_rhs_spec;
    std::vector<int64_t> conv_out_spec;
    std::vector<int64_t> slice_sizes;

    // Distributed collective parameters
    std::string collective_op{"sum"};
    int64_t collective_dim{0};
    int64_t split_dim{0};
    int64_t concat_dim{0};
    int64_t split_count{1};
    std::vector<std::vector<int64_t>> replica_groups;
  };

  struct TxprEqn {
    OpCode op{OpCode::kAdd};
    std::vector<VarId> in_vars;
    std::vector<VarId> out_vars;
    AttributePayload attrs;
  };

  // Instruction is an alias for TxprEqn to maintain compatibility.
  using Instruction = TxprEqn;

  class Txpr {
   public:
    Txpr() = default;

    VarId AddVar(Shape shape, DType dtype, std::string name = "");

    VarId AppendInstruction(OpCode op, std::vector<VarId> in_vars, Shape out_shape, DType out_dtype, AttributePayload attrs = {});

    std::vector<VarId> AppendMultiOutInstruction(OpCode op, std::vector<VarId> in_vars, const std::vector<Shape>& out_shapes,
                                                 const std::vector<DType>& out_dtypes, AttributePayload attrs = {});

    void MarkInput(VarId var);

    void MarkOutput(VarId var);

    [[nodiscard]] const std::vector<VarDef>& Variables() const noexcept { return variables_; }

    [[nodiscard]] const std::vector<TxprEqn>& Equations() const noexcept { return equations_; }

    [[nodiscard]] const std::vector<TxprEqn>& Instructions() const noexcept { return equations_; }

    [[nodiscard]] const std::vector<VarId>& InBinders() const noexcept { return in_binders_; }

    [[nodiscard]] const std::vector<VarId>& OutBinders() const noexcept { return out_binders_; }

    [[nodiscard]] size_t NumVars() const noexcept { return variables_.size(); }

    [[nodiscard]] size_t NumEquations() const noexcept { return equations_.size(); }

    [[nodiscard]] size_t NumInstructions() const noexcept { return equations_.size(); }

    [[nodiscard]] const VarDef& Var(VarId id) const {
      TENZR_CHECK(id.IsValid() && id.id < variables_.size(), "VarId out of bounds in Txpr");
      return variables_[id.id];
    }

    [[nodiscard]] const VarDef& Variable(VarId id) const { return Var(id); }

    [[nodiscard]] const Shape& ShapeOf(VarId id) const { return Var(id).shape; }

    [[nodiscard]] DType DtypeOf(VarId id) const { return Var(id).dtype; }

    [[nodiscard]] bool Validate(std::string* error_msg = nullptr) const;

    [[nodiscard]] std::string ToString() const;

   private:
    std::vector<VarDef> variables_;
    std::vector<TxprEqn> equations_;
    std::vector<VarId> in_binders_;
    std::vector<VarId> out_binders_;
  };

  // Backward-compatible alias
  using LinearTape = Txpr;

}  // namespace tenzr::txpr

namespace tenzr {
  using txpr::AttributePayload;
  using txpr::Instruction;
  using txpr::LinearTape;
  using txpr::OpCode;
  using txpr::Txpr;
  using txpr::TxprEqn;
  using txpr::TxprVar;
  using txpr::VarDef;
  using txpr::VarId;
}  // namespace tenzr

#endif  // TENZR_TXPR_TXPR_H_