// 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/emitter.h"
#include <algorithm>
#include <cmath>
#include <cstring>
#include <iomanip>
#include <numeric>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/common/trace.h"
#include "tenzr/frontend/op_registry.h"
#include "tenzr/txpr/txpr.h"
namespace tenzr::txpr {
namespace {
std::string GetReducerInitLiteral(OpCode op, DType dtype) {
switch (op) {
case OpCode::kReduceSum:
if (IsBoolean(dtype)) {
return "false";
}
if (IsIntegral(dtype)) {
return "0";
}
return "0.000000e+00";
case OpCode::kReduceProd:
if (IsBoolean(dtype)) {
return "true";
}
if (IsIntegral(dtype)) {
return "1";
}
return "1.000000e+00";
case OpCode::kReduceMax:
if (IsBoolean(dtype)) {
return "false";
}
if (IsSignedIntegral(dtype)) {
switch (dtype) {
case DType::kInt8:
return "-128";
case DType::kInt16:
return "-32768";
case DType::kInt32:
return "-2147483648";
case DType::kInt64:
return "-9223372036854775808";
case DType::kInt1:
return "-1";
case DType::kInt2:
return "-2";
case DType::kInt4:
return "-8";
default:
return "0";
}
}
if (IsUnsignedIntegral(dtype)) {
return "0";
}
if (dtype == DType::kFloat64) {
return "0xFFF0000000000000";
}
if (dtype == DType::kFloat32) {
return "0xFF800000";
}
if (dtype == DType::kFloat16) {
return "0xFC00";
}
if (dtype == DType::kBFloat16) {
return "0xFF80";
}
TENZR_THROW("reducer identity for " + std::string(ToString(dtype)) +
" is not supported (no infinity and no sub-byte float literal); cast to float16/float32 for reductions");
case OpCode::kReduceMin:
if (IsBoolean(dtype)) {
return "true";
}
if (IsSignedIntegral(dtype)) {
switch (dtype) {
case DType::kInt8:
return "127";
case DType::kInt16:
return "32767";
case DType::kInt32:
return "2147483647";
case DType::kInt64:
return "9223372036854775807";
case DType::kInt1:
return "0";
case DType::kInt2:
return "1";
case DType::kInt4:
return "7";
default:
return "0";
}
}
if (IsUnsignedIntegral(dtype)) {
switch (dtype) {
case DType::kUInt8:
return "255";
case DType::kUInt16:
return "65535";
case DType::kUInt32:
return "4294967295";
case DType::kUInt64:
return "18446744073709551615";
case DType::kUInt1:
return "1";
case DType::kUInt2:
return "3";
case DType::kUInt4:
return "15";
default:
return "0";
}
}
if (dtype == DType::kFloat64) {
return "0x7FF0000000000000";
}
if (dtype == DType::kFloat32) {
return "0x7F800000";
}
if (dtype == DType::kFloat16) {
return "0x7C00";
}
if (dtype == DType::kBFloat16) {
return "0x7F80";
}
TENZR_THROW("reducer identity for " + std::string(ToString(dtype)) +
" is not supported (no infinity and no sub-byte float literal); cast to float16/float32 for reductions");
default:
return "0.000000e+00";
}
}
std::string GetReducerBodyOp(OpCode op) {
switch (op) {
case OpCode::kReduceSum:
return "stablehlo.add";
case OpCode::kReduceProd:
return "stablehlo.multiply";
case OpCode::kReduceMax:
return "stablehlo.maximum";
case OpCode::kReduceMin:
return "stablehlo.minimum";
default:
return "stablehlo.add";
}
}
} // namespace
std::string TxprEmitter::NextSsa() { return "%" + std::to_string(ssa_counter_++); }
std::string TxprEmitter::GetVarSsa(VarId var) const {
auto it = var_ssas_.find(var.id);
TENZR_CHECK(it != var_ssas_.end(), "SSA name not found for VarId " + std::to_string(var.id));
return it->second;
}
std::string TxprEmitter::EnsureBroadcast(VarId var, const Shape& target_shape, const Txpr& txpr) {
std::string in_ssa = GetVarSsa(var);
const auto& in_var = txpr.Variable(var);
if (in_var.shape == target_shape) {
return in_ssa;
}
std::string bcast_ssa = NextSsa();
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
std::string out_type = target_shape.ToMlirTensorType(in_var.dtype);
size_t in_rank = in_var.shape.Rank();
size_t out_rank = target_shape.Rank();
TENZR_CHECK(in_rank <= out_rank, "Cannot broadcast tensor of higher rank to lower rank");
std::vector<size_t> dims(in_rank);
for (size_t i = 0; i < in_rank; ++i) {
dims[i] = (out_rank - in_rank) + i;
}
body_ << " " << bcast_ssa << " = stablehlo.broadcast_in_dim " << in_ssa << ", dims = [";
for (size_t i = 0; i < in_rank; ++i) {
body_ << dims[i] << (i + 1 < in_rank ? ", " : "");
}
body_ << "] : (" << in_type << ") -> " << out_type << "\n";
return bcast_ssa;
}
TxprCompiledModule TxprEmitter::Emit(const Txpr& txpr) {
TENZR_TRACE_SCOPE_CAT("compiler", "TxprEmitter::emit");
TENZR_CHECK(txpr.OutBinders().size() > 0, "Cannot emit MLIR for txpr with no outputs");
var_ssas_.clear();
body_.str("");
body_.clear();
ssa_counter_ = 0;
TxprCompiledModule result;
result.in_vars = txpr.InBinders();
result.out_vars = txpr.OutBinders();
result.in_shapes.reserve(result.in_vars.size());
result.in_dtypes.reserve(result.in_vars.size());
for (size_t i = 0; i < result.in_vars.size(); ++i) {
const auto& v = txpr.Variable(result.in_vars[i]);
result.in_shapes.push_back(v.shape);
result.in_dtypes.push_back(v.dtype);
var_ssas_[result.in_vars[i].id] = "%arg" + std::to_string(i);
}
result.out_shapes.reserve(result.out_vars.size());
result.out_dtypes.reserve(result.out_vars.size());
for (size_t i = 0; i < result.out_vars.size(); ++i) {
const auto& v = txpr.Variable(result.out_vars[i]);
result.out_shapes.push_back(v.shape);
result.out_dtypes.push_back(v.dtype);
}
// Emit instructions sequentially
for (const auto& eqn : txpr.Instructions()) {
TENZR_CHECK(!eqn.out_vars.empty(), "Instruction has no output variables");
const auto& out_var = txpr.Variable(eqn.out_vars[0]);
std::string out_type = out_var.shape.ToMlirTensorType(out_var.dtype);
std::string ssa = NextSsa();
var_ssas_[eqn.out_vars[0].id] = ssa;
switch (eqn.op) {
// Constants
case OpCode::kConstant: {
if (eqn.attrs.iota_dimension >= 0) {
// Native iota: no dense constant, computed on device (avoids
// serialising a full-shape index vector into the MLIR).
body_ << " " << ssa << " = stablehlo.iota dim = " << eqn.attrs.iota_dimension << " : " << out_type << "\n";
break;
}
if (!eqn.attrs.tensor_data.empty()) {
std::string elem = std::string(ToMlirType(out_var.dtype));
std::string flat_type = "tensor<" + std::to_string(eqn.attrs.tensor_data.size()) + "x" + elem + ">";
std::string c_ssa = (out_var.shape.Rank() == 1) ? ssa : NextSsa();
body_ << " " << c_ssa << " = stablehlo.constant dense<[";
for (size_t i = 0; i < eqn.attrs.tensor_data.size(); ++i) {
std::ostringstream ss;
ss << std::scientific << std::setprecision(out_var.dtype == DType::kFloat64 ? 17 : 6) << eqn.attrs.tensor_data[i];
body_ << ss.str() << (i + 1 < eqn.attrs.tensor_data.size() ? ", " : "");
}
body_ << "]> : " << flat_type << "\n";
if (out_var.shape.Rank() != 1) {
body_ << " " << ssa << " = \"stablehlo.reshape\"(" << c_ssa << ") : (" << flat_type << ") -> " << out_type << "\n";
}
break;
}
if (!eqn.attrs.int_tensor_data.empty()) {
std::string elem = std::string(ToMlirType(out_var.dtype));
std::string flat_type = "tensor<" + std::to_string(eqn.attrs.int_tensor_data.size()) + "x" + elem + ">";
std::string c_ssa = (out_var.shape.Rank() == 1) ? ssa : NextSsa();
body_ << " " << c_ssa << " = stablehlo.constant dense<[";
for (size_t i = 0; i < eqn.attrs.int_tensor_data.size(); ++i) {
body_ << eqn.attrs.int_tensor_data[i] << (i + 1 < eqn.attrs.int_tensor_data.size() ? ", " : "");
}
body_ << "]> : " << flat_type << "\n";
if (out_var.shape.Rank() != 1) {
body_ << " " << ssa << " = \"stablehlo.reshape\"(" << c_ssa << ") : (" << flat_type << ") -> " << out_type << "\n";
}
break;
}
std::string literal;
if (IsBoolean(out_var.dtype)) {
literal = (eqn.attrs.scalar_value != 0.0) ? "true" : "false";
} else if (IsIntegral(out_var.dtype)) {
literal = std::to_string(static_cast<int64_t>(eqn.attrs.scalar_value));
} else {
const double value = eqn.attrs.scalar_value;
if (!std::isfinite(value)) {
// MLIR float literals must be finite; emit the IEEE bit pattern.
std::ostringstream ss;
ss << std::hex << std::uppercase << std::setfill('0');
if (out_var.dtype == DType::kFloat64) {
uint64_t bits = 0;
std::memcpy(&bits, &value, sizeof(bits));
ss << "0x" << std::setw(16) << bits;
} else if (out_var.dtype == DType::kFloat16) {
tenzr::float16 h(static_cast<float>(value));
uint16_t bits = 0;
std::memcpy(&bits, &h.bits, sizeof(bits));
ss << "0x" << std::setw(4) << bits;
} else if (out_var.dtype == DType::kBFloat16) {
tenzr::bfloat16 h(static_cast<float>(value));
uint16_t bits = 0;
std::memcpy(&bits, &h.bits, sizeof(bits));
ss << "0x" << std::setw(4) << bits;
} else if (out_var.dtype == DType::kFloat32) {
const float f = static_cast<float>(value);
uint32_t bits = 0;
std::memcpy(&bits, &f, sizeof(bits));
ss << "0x" << std::setw(8) << bits;
} else {
TENZR_THROW("non-finite constant for " + std::string(ToString(out_var.dtype)) + " is not supported; cast to float16/float32");
}
literal = ss.str();
} else {
std::ostringstream ss;
ss << std::scientific << std::setprecision(6) << value;
literal = ss.str();
}
}
body_ << " " << ssa << " = stablehlo.constant dense<" << literal << "> : " << out_type << "\n";
break;
}
// Binary Elementwise
case OpCode::kAdd:
case OpCode::kSub:
case OpCode::kMul:
case OpCode::kDiv: {
const char* hlo_op = nullptr;
switch (eqn.op) {
case OpCode::kAdd:
hlo_op = "stablehlo.add";
break;
case OpCode::kSub:
hlo_op = "stablehlo.subtract";
break;
case OpCode::kMul:
hlo_op = "stablehlo.multiply";
break;
case OpCode::kDiv:
hlo_op = "stablehlo.divide";
break;
default:
break;
}
std::string lhs_ssa = EnsureBroadcast(eqn.in_vars[0], out_var.shape, txpr);
std::string rhs_ssa = EnsureBroadcast(eqn.in_vars[1], out_var.shape, txpr);
body_ << " " << ssa << " = " << hlo_op << " " << lhs_ssa << ", " << rhs_ssa << " : " << out_type << "\n";
break;
}
// Matrix Operations
case OpCode::kMatmul: {
std::string lhs_ssa = GetVarSsa(eqn.in_vars[0]);
std::string rhs_ssa = GetVarSsa(eqn.in_vars[1]);
const auto& lhs_var = txpr.Variable(eqn.in_vars[0]);
const auto& rhs_var = txpr.Variable(eqn.in_vars[1]);
std::string lhs_type = lhs_var.shape.ToMlirTensorType(lhs_var.dtype);
std::string rhs_type = rhs_var.shape.ToMlirTensorType(rhs_var.dtype);
if (lhs_var.shape.Rank() == 2 && rhs_var.shape.Rank() == 2) {
body_ << " " << ssa << " = \"stablehlo.dot_general\"(" << lhs_ssa << ", " << rhs_ssa << ") {\n"
<< " dot_dimension_numbers = #stablehlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>\n"
<< " } : (" << lhs_type << ", " << rhs_type << ") -> " << out_type << "\n";
} else {
size_t lhs_rank = lhs_var.shape.Rank();
size_t rhs_rank = rhs_var.shape.Rank();
body_ << " " << ssa << " = \"stablehlo.dot_general\"(" << lhs_ssa << ", " << rhs_ssa << ") {\n"
<< " dot_dimension_numbers = #stablehlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], "
<< "lhs_contracting_dimensions = [" << (lhs_rank - 1) << "], "
<< "rhs_contracting_dimensions = [" << (rhs_rank - 2) << "]>\n"
<< " } : (" << lhs_type << ", " << rhs_type << ") -> " << out_type << "\n";
}
break;
}
case OpCode::kConvolution: {
std::string lhs_ssa = GetVarSsa(eqn.in_vars[0]);
std::string rhs_ssa = GetVarSsa(eqn.in_vars[1]);
const auto& lhs_var = txpr.Variable(eqn.in_vars[0]);
const auto& rhs_var = txpr.Variable(eqn.in_vars[1]);
std::string lhs_type = lhs_var.shape.ToMlirTensorType(lhs_var.dtype);
std::string rhs_type = rhs_var.shape.ToMlirTensorType(rhs_var.dtype);
const size_t nspatial = eqn.attrs.window_strides.size();
auto emit_i64 = [&](const std::vector<int64_t>& v) {
body_ << "[";
for (size_t i = 0; i < v.size(); ++i) {
body_ << v[i] << (i + 1 < v.size() ? ", " : "");
}
body_ << "]";
};
body_ << " " << ssa << " = stablehlo.convolution(" << lhs_ssa << ", " << rhs_ssa << ")\n";
const std::vector<int64_t> lhs_spec = eqn.attrs.conv_lhs_spec.empty() ? std::vector<int64_t>{0, 1, 2, 3} : eqn.attrs.conv_lhs_spec;
const std::vector<int64_t> rhs_spec = eqn.attrs.conv_rhs_spec.empty() ? std::vector<int64_t>{0, 1, 2, 3} : eqn.attrs.conv_rhs_spec;
const std::vector<int64_t> out_spec = eqn.attrs.conv_out_spec.empty() ? std::vector<int64_t>{0, 1, 2, 3} : eqn.attrs.conv_out_spec;
auto spec_dims = [](const std::vector<int64_t>& spec, char batch_ch, char feat_ch) {
std::string s(4, ' ');
s[static_cast<size_t>(spec[0])] = batch_ch;
s[static_cast<size_t>(spec[1])] = feat_ch;
s[static_cast<size_t>(spec[2])] = '0';
s[static_cast<size_t>(spec[3])] = '1';
return s;
};
auto join_dims = [](const std::string& s) {
std::string r;
for (size_t i = 0; i < s.size(); ++i) {
if (i != 0) {
r += ", ";
}
r += s[i];
}
return r;
};
body_ << " dim_numbers = [" << join_dims(spec_dims(lhs_spec, 'b', 'f')) << "]x[" << join_dims(spec_dims(rhs_spec, 'o', 'i')) << "]->["
<< join_dims(spec_dims(out_spec, 'b', 'f')) << "],\n";
body_ << " window = {stride = ";
emit_i64(eqn.attrs.window_strides);
body_ << ", pad = [";
for (size_t d = 0; d < nspatial; ++d) {
const int64_t lo = d < eqn.attrs.edge_padding_low.size() ? eqn.attrs.edge_padding_low[d] : 0;
const int64_t hi = d < eqn.attrs.edge_padding_high.size() ? eqn.attrs.edge_padding_high[d] : 0;
body_ << "[" << lo << ", " << hi << "]" << (d + 1 < nspatial ? ", " : "");
}
body_ << "], lhs_dilate = ";
emit_i64(eqn.attrs.lhs_dilation);
body_ << ", rhs_dilate = ";
emit_i64(eqn.attrs.rhs_dilation);
body_ << ", reverse = [0, 0]}\n";
body_ << " {feature_group_count = 1 : i64, batch_group_count = 1 : i64, "
"precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]}\n";
body_ << " : (" << lhs_type << ", " << rhs_type << ") -> " << out_type << "\n";
break;
}
// Shape & Strides
case OpCode::kReshape: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
body_ << " " << ssa << " = \"stablehlo.reshape\"(" << in_ssa << ") : (" << in_type << ") -> " << out_type << "\n";
break;
}
case OpCode::kTranspose: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
body_ << " " << ssa << " = \"stablehlo.transpose\"(" << in_ssa << ") {permutation = array<i64: ";
const auto& perm = eqn.attrs.permutation;
for (size_t p = 0; p < perm.size(); ++p) {
body_ << perm[p] << (p + 1 < perm.size() ? ", " : "");
}
body_ << ">} : (" << in_type << ") -> " << out_type << "\n";
break;
}
case OpCode::kSlice: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
size_t rank = in_var.shape.Rank();
std::vector<int64_t> starts(rank, 0);
std::vector<int64_t> limits(rank, 0);
std::vector<int64_t> strides(rank, 1);
for (size_t i = 0; i < rank; ++i) {
int64_t d = in_var.shape.Dim(i);
if (i < eqn.attrs.slices.size()) {
const auto& s = eqn.attrs.slices[i];
int64_t st = s.start.value_or(0);
int64_t sp = s.stop.value_or(d);
if (st < 0) {
st += d;
}
if (sp < 0) {
sp += d;
}
starts[i] = std::clamp(st, int64_t(0), d);
limits[i] = std::clamp(sp, int64_t(0), d);
strides[i] = s.step;
} else {
starts[i] = 0;
limits[i] = d;
strides[i] = 1;
}
}
body_ << " " << ssa << " = \"stablehlo.slice\"(" << in_ssa << ") {";
body_ << "start_indices = array<i64: ";
for (size_t i = 0; i < rank; ++i) {
body_ << starts[i] << (i + 1 < rank ? ", " : "");
}
body_ << ">, limit_indices = array<i64: ";
for (size_t i = 0; i < rank; ++i) {
body_ << limits[i] << (i + 1 < rank ? ", " : "");
}
body_ << ">, strides = array<i64: ";
for (size_t i = 0; i < rank; ++i) {
body_ << strides[i] << (i + 1 < rank ? ", " : "");
}
body_ << ">";
body_ << "} : (" << in_type << ") -> " << out_type << "\n";
break;
}
case OpCode::kPad: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
std::string pv_ssa = GetVarSsa(eqn.in_vars[1]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
const auto& pv_var = txpr.Variable(eqn.in_vars[1]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
std::string pv_type = pv_var.shape.ToMlirTensorType(pv_var.dtype);
body_ << " " << ssa << " = \"stablehlo.pad\"(" << in_ssa << ", " << pv_ssa << ") {";
body_ << "edge_padding_low = array<i64: ";
for (size_t i = 0; i < eqn.attrs.edge_padding_low.size(); ++i) {
body_ << eqn.attrs.edge_padding_low[i] << (i + 1 < eqn.attrs.edge_padding_low.size() ? ", " : "");
}
body_ << ">, edge_padding_high = array<i64: ";
for (size_t i = 0; i < eqn.attrs.edge_padding_high.size(); ++i) {
body_ << eqn.attrs.edge_padding_high[i] << (i + 1 < eqn.attrs.edge_padding_high.size() ? ", " : "");
}
body_ << ">, interior_padding = array<i64: ";
for (size_t i = 0; i < eqn.attrs.interior_padding.size(); ++i) {
body_ << eqn.attrs.interior_padding[i] << (i + 1 < eqn.attrs.interior_padding.size() ? ", " : "");
}
body_ << ">} : (" << in_type << ", " << pv_type << ") -> " << out_type << "\n";
break;
}
case OpCode::kScatter: {
std::string operand_ssa = GetVarSsa(eqn.in_vars[0]);
std::string indices_ssa = GetVarSsa(eqn.in_vars[1]);
std::string updates_ssa = GetVarSsa(eqn.in_vars[2]);
const auto& op_var = txpr.Variable(eqn.in_vars[0]);
const auto& idx_var = txpr.Variable(eqn.in_vars[1]);
const auto& upd_var = txpr.Variable(eqn.in_vars[2]);
std::string op_type = op_var.shape.ToMlirTensorType(op_var.dtype);
std::string idx_type = idx_var.shape.ToMlirTensorType(idx_var.dtype);
std::string upd_type = upd_var.shape.ToMlirTensorType(upd_var.dtype);
std::string scalar_type = "tensor<" + std::string(ToMlirType(op_var.dtype)) + ">";
std::string base = ssa;
if (!base.empty() && base[0] == '%') {
base = base.substr(1);
}
std::string r_name = "%sc" + base + "_r";
std::string a_name = "%sc" + base + "_a";
std::string b_name = "%sc" + base + "_b";
const std::string& upd = eqn.attrs.scatter_update;
std::string region_op = upd == "multiply" ? "stablehlo.multiply"
: upd == "min" ? "stablehlo.minimum"
: upd == "max" ? "stablehlo.maximum"
: upd == "set" ? ""
: "stablehlo.add";
body_ << " " << ssa << " = \"stablehlo.scatter\"(" << operand_ssa << ", " << indices_ssa << ", " << updates_ssa << ") ({\n";
body_ << " ^bb0(" << a_name << ": " << scalar_type << ", " << b_name << ": " << scalar_type << "):\n";
if (region_op.empty()) {
body_ << " stablehlo.return " << b_name << " : " << scalar_type << "\n";
} else {
body_ << " " << r_name << " = " << region_op << " " << a_name << ", " << b_name << " : " << scalar_type << "\n";
body_ << " stablehlo.return " << r_name << " : " << scalar_type << "\n";
}
int64_t saxis = eqn.attrs.scatter_axis;
if (saxis < 0) {
saxis += static_cast<int64_t>(upd_var.shape.Rank());
}
body_ << " }) {scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [";
bool first_wd = true;
for (size_t d = 0; d < upd_var.shape.Rank(); ++d) {
if (static_cast<int64_t>(d) == saxis) {
continue;
}
body_ << (first_wd ? "" : ", ") << d;
first_wd = false;
}
body_ << "], inserted_window_dims = [" << saxis << "], scatter_dims_to_operand_dims = [" << saxis << "], index_vector_dim = 1>} : (" << op_type
<< ", " << idx_type << ", " << upd_type << ") -> " << out_type << "\n";
break;
}
case OpCode::kGather: {
std::string operand_ssa = GetVarSsa(eqn.in_vars[0]);
std::string indices_ssa = GetVarSsa(eqn.in_vars[1]);
const auto& op_var = txpr.Variable(eqn.in_vars[0]);
const auto& idx_var = txpr.Variable(eqn.in_vars[1]);
std::string op_type = op_var.shape.ToMlirTensorType(op_var.dtype);
std::string idx_type = idx_var.shape.ToMlirTensorType(idx_var.dtype);
const size_t rank = op_var.shape.Rank();
int64_t gaxis = eqn.attrs.gather_axis;
if (gaxis < 0) {
gaxis += static_cast<int64_t>(rank);
}
body_ << " " << ssa << " = \"stablehlo.gather\"(" << operand_ssa << ", " << indices_ssa << ") {\n";
body_ << " dimension_numbers = #stablehlo.gather<offset_dims = [";
for (size_t d = 1; d < rank; ++d) {
body_ << d << (d + 1 < rank ? ", " : "");
}
body_ << "], collapsed_slice_dims = [" << gaxis << "], start_index_map = [" << gaxis << "], index_vector_dim = 1>,\n";
body_ << " slice_sizes = array<i64: ";
for (size_t d = 0; d < rank; ++d) {
body_ << (static_cast<int64_t>(d) == gaxis ? 1 : op_var.shape.Dim(d)) << (d + 1 < rank ? ", " : "");
}
body_ << ">,\n indices_are_sorted = false\n } : (" << op_type << ", " << idx_type << ") -> " << out_type << "\n";
break;
}
case OpCode::kDynamicSlice: {
const auto& op_var = txpr.Variable(eqn.in_vars[0]);
std::string op_type = op_var.shape.ToMlirTensorType(op_var.dtype);
body_ << " " << ssa << " = \"stablehlo.dynamic_slice\"(" << GetVarSsa(eqn.in_vars[0]);
for (size_t i = 1; i < eqn.in_vars.size(); ++i) {
body_ << ", " << GetVarSsa(eqn.in_vars[i]);
}
body_ << ") {slice_sizes = array<i64: ";
for (size_t i = 0; i < eqn.attrs.slice_sizes.size(); ++i) {
body_ << eqn.attrs.slice_sizes[i] << (i + 1 < eqn.attrs.slice_sizes.size() ? ", " : "");
}
body_ << ">} : (" << op_type;
for (size_t i = 1; i < eqn.in_vars.size(); ++i) {
const auto& iv = txpr.Variable(eqn.in_vars[i]);
body_ << ", " << iv.shape.ToMlirTensorType(iv.dtype);
}
body_ << ") -> " << out_type << "\n";
break;
}
case OpCode::kDynamicUpdateSlice: {
const auto& op_var = txpr.Variable(eqn.in_vars[0]);
std::string op_type = op_var.shape.ToMlirTensorType(op_var.dtype);
body_ << " " << ssa << " = \"stablehlo.dynamic_update_slice\"(" << GetVarSsa(eqn.in_vars[0]) << ", " << GetVarSsa(eqn.in_vars[1]);
for (size_t i = 2; i < eqn.in_vars.size(); ++i) {
body_ << ", " << GetVarSsa(eqn.in_vars[i]);
}
body_ << ") : (" << op_type;
for (size_t i = 1; i < eqn.in_vars.size(); ++i) {
const auto& iv = txpr.Variable(eqn.in_vars[i]);
body_ << ", " << iv.shape.ToMlirTensorType(iv.dtype);
}
body_ << ") -> " << out_type << "\n";
break;
}
case OpCode::kConcatenate: {
body_ << " " << ssa << " = stablehlo.concatenate ";
for (size_t i = 0; i < eqn.in_vars.size(); ++i) {
body_ << (i ? ", " : "") << GetVarSsa(eqn.in_vars[i]);
}
body_ << ", dim = " << eqn.attrs.concat_axis << " : (";
for (size_t i = 0; i < eqn.in_vars.size(); ++i) {
const auto& iv = txpr.Variable(eqn.in_vars[i]);
body_ << (i ? ", " : "") << iv.shape.ToMlirTensorType(iv.dtype);
}
body_ << ") -> " << out_type << "\n";
break;
}
case OpCode::kArgSort: {
std::string values_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& v_var = txpr.Variable(eqn.in_vars[0]);
std::string v_type = v_var.shape.ToMlirTensorType(v_var.dtype);
std::string v_scalar = "tensor<" + std::string(ToMlirType(v_var.dtype)) + ">";
std::string base = ssa;
if (!base.empty() && base[0] == '%') {
base = base.substr(1);
}
std::string av = "%av" + base;
std::string bv = "%bv" + base;
std::string rr = "%asr" + base;
std::string cmp = eqn.attrs.sort_descending ? "GT" : "LT";
if (eqn.attrs.sort_return_values) {
// Values-only sort: single-operand stablehlo.sort (no index payload).
body_ << " " << ssa << " = \"stablehlo.sort\"(" << values_ssa << ") ({\n";
body_ << " ^bb0(" << av << ": " << v_scalar << ", " << bv << ": " << v_scalar << "):\n";
body_ << " " << rr << " = stablehlo.compare " << cmp << ", " << av << ", " << bv << " : (" << v_scalar << ", " << v_scalar
<< ") -> tensor<i1>\n";
body_ << " \"stablehlo.return\"(" << rr << ") : (tensor<i1>) -> ()\n";
body_ << " }) {dimension = " << eqn.attrs.sort_dimension << " : i64, is_stable = true} : (" << v_type << ") -> " << out_type << "\n";
var_ssas_[eqn.out_vars[0].id] = ssa;
break;
}
std::string indices_ssa = GetVarSsa(eqn.in_vars[1]);
const auto& i_var = txpr.Variable(eqn.in_vars[1]);
std::string i_type = i_var.shape.ToMlirTensorType(i_var.dtype);
std::string i_scalar = "tensor<" + std::string(ToMlirType(i_var.dtype)) + ">";
std::string ai = "%ai" + base;
std::string bi = "%bi" + base;
body_ << " " << ssa << ":2 = \"stablehlo.sort\"(" << values_ssa << ", " << indices_ssa << ") ({\n";
body_ << " ^bb0(" << av << ": " << v_scalar << ", " << bv << ": " << v_scalar << ", " << ai << ": " << i_scalar << ", " << bi << ": " << i_scalar
<< "):\n";
body_ << " " << rr << " = stablehlo.compare " << cmp << ", " << av << ", " << bv << " : (" << v_scalar << ", " << v_scalar
<< ") -> tensor<i1>\n";
body_ << " \"stablehlo.return\"(" << rr << ") : (tensor<i1>) -> ()\n";
body_ << " }) {dimension = " << eqn.attrs.sort_dimension << " : i64, is_stable = true} : (" << v_type << ", " << i_type << ") -> (" << v_type
<< ", " << i_type << ")\n";
var_ssas_[eqn.out_vars[0].id] = ssa + (eqn.attrs.sort_return_values ? "#0" : "#1");
break;
}
case OpCode::kBroadcastTo: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
body_ << " " << ssa << " = stablehlo.broadcast_in_dim " << in_ssa << ", dims = [";
for (size_t i = 0; i < eqn.attrs.axes.size(); ++i) {
body_ << eqn.attrs.axes[i] << (i + 1 < eqn.attrs.axes.size() ? ", " : "");
}
body_ << "] : (" << in_type << ") -> " << out_type << "\n";
break;
}
// Reductions
case OpCode::kReduceSum:
case OpCode::kReduceProd:
case OpCode::kReduceMax:
case OpCode::kReduceMin: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string scalar_type = "tensor<" + std::string(ToMlirType(out_var.dtype)) + ">";
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
std::string c_init = NextSsa();
std::string init_lit = GetReducerInitLiteral(eqn.op, out_var.dtype);
body_ << " " << c_init << " = stablehlo.constant dense<" << init_lit << "> : " << scalar_type << "\n";
std::string red_ssa = NextSsa();
std::string arg_lhs = "%red_lhs_" + std::to_string(eqn.out_vars[0].id);
std::string arg_rhs = "%red_rhs_" + std::to_string(eqn.out_vars[0].id);
std::string acc_ssa = "%red_acc_" + std::to_string(eqn.out_vars[0].id);
std::string body_op = GetReducerBodyOp(eqn.op);
// Compute non-keepdims shape if keepdims was requested
std::vector<int64_t> reduction_axes = eqn.attrs.axes;
if (reduction_axes.empty()) {
reduction_axes.resize(in_var.shape.Rank());
std::iota(reduction_axes.begin(), reduction_axes.end(), 0);
}
std::vector<int64_t> non_keep_dims;
std::unordered_map<int64_t, bool> is_red_axis;
for (int64_t ax : reduction_axes) {
if (ax < 0) {
ax += static_cast<int64_t>(in_var.shape.Rank());
}
is_red_axis[ax] = true;
}
for (size_t i = 0; i < in_var.shape.Rank(); ++i) {
if (!is_red_axis[static_cast<int64_t>(i)]) {
non_keep_dims.push_back(in_var.shape.Dim(i));
}
}
Shape non_keep_shape(non_keep_dims);
std::string non_keep_type = non_keep_shape.ToMlirTensorType(out_var.dtype);
body_ << " " << red_ssa << " = \"stablehlo.reduce\"(" << in_ssa << ", " << c_init << ") ({\n";
body_ << " ^bb0(" << arg_lhs << ": " << scalar_type << ", " << arg_rhs << ": " << scalar_type << "):\n";
body_ << " " << acc_ssa << " = " << body_op << " " << arg_lhs << ", " << arg_rhs << " : " << scalar_type << "\n";
body_ << " \"stablehlo.return\"(" << acc_ssa << ") : (" << scalar_type << ") -> ()\n";
body_ << " }) {\n";
if (reduction_axes.empty()) {
body_ << " dimensions = array<i64>\n";
} else {
body_ << " dimensions = array<i64: ";
for (size_t i = 0; i < reduction_axes.size(); ++i) {
int64_t ax = reduction_axes[i];
if (ax < 0) {
ax += static_cast<int64_t>(in_var.shape.Rank());
}
body_ << ax << (i + 1 < reduction_axes.size() ? ", " : "");
}
body_ << ">\n";
}
body_ << " } : (" << in_type << ", " << scalar_type << ") -> " << non_keep_type << "\n";
if (eqn.attrs.keepdims && out_var.shape != non_keep_shape) {
std::string final_ssa = NextSsa();
body_ << " " << final_ssa << " = stablehlo.reshape " << red_ssa << " : (" << non_keep_type << ") -> " << out_type << "\n";
var_ssas_[eqn.out_vars[0].id] = final_ssa;
} else {
var_ssas_[eqn.out_vars[0].id] = red_ssa;
}
break;
}
case OpCode::kReduceWindow: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
std::string init_ssa = GetVarSsa(eqn.in_vars[1]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
std::string scalar_type = "tensor<" + std::string(ToMlirType(out_var.dtype)) + ">";
const std::string& rw = eqn.attrs.reduce_window_op;
std::string body_op = rw == "max" ? "stablehlo.maximum"
: rw == "min" ? "stablehlo.minimum"
: rw == "multiply" ? "stablehlo.multiply"
: "stablehlo.add";
std::string red_ssa = NextSsa();
std::string arg_lhs = "%rw_lhs_" + std::to_string(eqn.out_vars[0].id);
std::string arg_rhs = "%rw_rhs_" + std::to_string(eqn.out_vars[0].id);
std::string acc_ssa = "%rw_acc_" + std::to_string(eqn.out_vars[0].id);
auto emit_i64_array = [&](const std::vector<int64_t>& v) {
body_ << "array<i64: ";
for (size_t i = 0; i < v.size(); ++i) {
body_ << v[i] << (i + 1 < v.size() ? ", " : "");
}
body_ << ">";
};
body_ << " " << red_ssa << " = \"stablehlo.reduce_window\"(" << in_ssa << ", " << init_ssa << ") ({\n";
body_ << " ^bb0(" << arg_lhs << ": " << scalar_type << ", " << arg_rhs << ": " << scalar_type << "):\n";
body_ << " " << acc_ssa << " = " << body_op << " " << arg_lhs << ", " << arg_rhs << " : " << scalar_type << "\n";
body_ << " \"stablehlo.return\"(" << acc_ssa << ") : (" << scalar_type << ") -> ()\n";
body_ << " }) {\n";
body_ << " window_dimensions = ";
emit_i64_array(eqn.attrs.window_dimensions);
body_ << ",\n window_strides = ";
emit_i64_array(eqn.attrs.window_strides);
body_ << ",\n base_dilations = ";
emit_i64_array(eqn.attrs.base_dilations);
body_ << ",\n window_dilations = ";
emit_i64_array(eqn.attrs.window_dilations);
body_ << ",\n padding = dense<[";
const size_t rw_rank = in_var.shape.Rank();
for (size_t d = 0; d < rw_rank; ++d) {
const int64_t lo = d < eqn.attrs.edge_padding_low.size() ? eqn.attrs.edge_padding_low[d] : 0;
const int64_t hi = d < eqn.attrs.edge_padding_high.size() ? eqn.attrs.edge_padding_high[d] : 0;
body_ << "[" << lo << ", " << hi << "]" << (d + 1 < rw_rank ? ", " : "");
}
body_ << "]> : tensor<" << rw_rank << "x2xi64>\n";
body_ << " } : (" << in_type << ", " << scalar_type << ") -> " << out_type << "\n";
var_ssas_[eqn.out_vars[0].id] = red_ssa;
break;
}
// Activations
case OpCode::kRelu: {
std::string c_zero = NextSsa();
std::string zero_literal = IsBoolean(out_var.dtype) ? "false" : (IsIntegral(out_var.dtype) ? "0" : "0.000000e+00");
body_ << " " << c_zero << " = stablehlo.constant dense<" << zero_literal << "> : " << out_type << "\n";
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
body_ << " " << ssa << " = stablehlo.maximum " << in_ssa << ", " << c_zero << " : " << out_type << "\n";
break;
}
case OpCode::kSigmoid: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
body_ << " " << ssa << " = stablehlo.logistic " << in_ssa << " : " << out_type << "\n";
break;
}
case OpCode::kConvert: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
body_ << " " << ssa << " = stablehlo.convert " << in_ssa << " : (" << in_type << ") -> " << out_type << "\n";
break;
}
case OpCode::kCustomOp: {
const auto* Definition = OpRegistry::Instance().Lookup(eqn.attrs.custom_op_name);
TENZR_CHECK(Definition != nullptr, "Missing OpDef for custom op: " + eqn.attrs.custom_op_name);
std::vector<std::string> in_ssas;
std::vector<std::string> in_types;
in_ssas.reserve(eqn.in_vars.size());
in_types.reserve(eqn.in_vars.size());
for (const auto& in_v : eqn.in_vars) {
const auto& iv = txpr.Variable(in_v);
if (Definition->Category == OpCategory::kBinaryElementwise || Definition->name == "clamp" || Definition->name == "select" ||
Definition->name.rfind("compare_", 0) == 0) {
std::string in_s = EnsureBroadcast(in_v, out_var.shape, txpr);
in_ssas.push_back(in_s);
in_types.push_back(out_var.shape.ToMlirTensorType(iv.dtype));
} else {
in_ssas.push_back(GetVarSsa(in_v));
in_types.push_back(iv.shape.ToMlirTensorType(iv.dtype));
}
}
if (Definition->emit_fn) {
body_ << Definition->emit_fn(ssa, out_type, in_ssas, in_types);
} else if (!Definition->StablehloOp.empty()) {
if (eqn.in_vars.size() == 1) {
if (in_types[0] == out_type) {
body_ << " " << ssa << " = " << Definition->StablehloOp << " " << in_ssas[0] << " : " << out_type << "\n";
} else {
body_ << " " << ssa << " = " << Definition->StablehloOp << " " << in_ssas[0] << " : (" << in_types[0] << ") -> " << out_type << "\n";
}
} else if (eqn.in_vars.size() == 2) {
body_ << " " << ssa << " = " << Definition->StablehloOp << " " << in_ssas[0] << ", " << in_ssas[1] << " : " << out_type << "\n";
} else {
TENZR_THROW("Unsupported default arity for custom op: " + Definition->name);
}
} else {
TENZR_THROW("OpDef has neither emit_fn nor stablehlo_op for: " + Definition->name);
}
break;
}
case OpCode::kAllReduce: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
std::string scalar_type = Shape({}).ToMlirTensorType(in_var.dtype);
std::ostringstream groups;
if (!eqn.attrs.replica_groups.empty()) {
size_t num_groups = eqn.attrs.replica_groups.size();
size_t group_size = eqn.attrs.replica_groups[0].size();
groups << "replica_groups = dense<[";
for (size_t g = 0; g < num_groups; ++g) {
if (g > 0) {
groups << ", ";
}
groups << "[";
for (size_t r = 0; r < eqn.attrs.replica_groups[g].size(); ++r) {
if (r > 0) {
groups << ", ";
}
groups << eqn.attrs.replica_groups[g][r];
}
groups << "]";
}
groups << "]> : tensor<" << num_groups << "x" << group_size << "xi64>";
} else {
groups << "replica_groups = dense<[[0]]> : tensor<1x1xi64>";
}
body_ << " " << ssa << " = \"stablehlo.all_reduce\"(" << in_ssa << ") <{" << groups.str() << "}> ({\n";
body_ << " ^bb0(%arg_lhs: " << scalar_type << ", %arg_rhs: " << scalar_type << "):\n";
if (eqn.attrs.collective_op == "max") {
body_ << " %red_val = stablehlo.maximum %arg_lhs, %arg_rhs : " << scalar_type << "\n";
} else if (eqn.attrs.collective_op == "min") {
body_ << " %red_val = stablehlo.minimum %arg_lhs, %arg_rhs : " << scalar_type << "\n";
} else if (eqn.attrs.collective_op == "prod" || eqn.attrs.collective_op == "multiply") {
body_ << " %red_val = stablehlo.multiply %arg_lhs, %arg_rhs : " << scalar_type << "\n";
} else {
body_ << " %red_val = stablehlo.add %arg_lhs, %arg_rhs : " << scalar_type << "\n";
}
body_ << " stablehlo.return %red_val : " << scalar_type << "\n";
body_ << " }) : (" << in_type << ") -> " << out_type << "\n";
break;
}
case OpCode::kAllGather: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
std::ostringstream groups;
if (!eqn.attrs.replica_groups.empty()) {
size_t num_groups = eqn.attrs.replica_groups.size();
size_t group_size = eqn.attrs.replica_groups[0].size();
groups << "replica_groups = dense<[";
for (size_t g = 0; g < num_groups; ++g) {
if (g > 0) {
groups << ", ";
}
groups << "[";
for (size_t r = 0; r < eqn.attrs.replica_groups[g].size(); ++r) {
if (r > 0) {
groups << ", ";
}
groups << eqn.attrs.replica_groups[g][r];
}
groups << "]";
}
groups << "]> : tensor<" << num_groups << "x" << group_size << "xi64>";
} else {
groups << "replica_groups = dense<[[0]]> : tensor<1x1xi64>";
}
body_ << " " << ssa << " = \"stablehlo.all_gather\"(" << in_ssa << ") <{all_gather_dim = " << eqn.attrs.collective_dim << " : i64, "
<< groups.str() << "}> : (" << in_type << ") -> " << out_type << "\n";
break;
}
case OpCode::kAllToAll: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
std::ostringstream groups;
if (!eqn.attrs.replica_groups.empty()) {
size_t num_groups = eqn.attrs.replica_groups.size();
size_t group_size = eqn.attrs.replica_groups[0].size();
groups << "replica_groups = dense<[";
for (size_t g = 0; g < num_groups; ++g) {
if (g > 0) {
groups << ", ";
}
groups << "[";
for (size_t r = 0; r < eqn.attrs.replica_groups[g].size(); ++r) {
if (r > 0) {
groups << ", ";
}
groups << eqn.attrs.replica_groups[g][r];
}
groups << "]";
}
groups << "]> : tensor<" << num_groups << "x" << group_size << "xi64>";
} else {
groups << "replica_groups = dense<[[0]]> : tensor<1x1xi64>";
}
body_ << " " << ssa << " = \"stablehlo.all_to_all\"(" << in_ssa << ") <{concat_dimension = " << eqn.attrs.concat_dim << " : i64, " << groups.str()
<< ", split_count = " << eqn.attrs.split_count << " : i64, split_dimension = " << eqn.attrs.split_dim << " : i64}> : (" << in_type << ") -> "
<< out_type << "\n";
break;
}
case OpCode::kReduceScatter: {
std::string in_ssa = GetVarSsa(eqn.in_vars[0]);
const auto& in_var = txpr.Variable(eqn.in_vars[0]);
std::string in_type = in_var.shape.ToMlirTensorType(in_var.dtype);
std::string scalar_type = Shape({}).ToMlirTensorType(in_var.dtype);
std::ostringstream groups;
if (!eqn.attrs.replica_groups.empty()) {
size_t num_groups = eqn.attrs.replica_groups.size();
size_t group_size = eqn.attrs.replica_groups[0].size();
groups << "replica_groups = dense<[";
for (size_t g = 0; g < num_groups; ++g) {
if (g > 0) {
groups << ", ";
}
groups << "[";
for (size_t r = 0; r < eqn.attrs.replica_groups[g].size(); ++r) {
if (r > 0) {
groups << ", ";
}
groups << eqn.attrs.replica_groups[g][r];
}
groups << "]";
}
groups << "]> : tensor<" << num_groups << "x" << group_size << "xi64>";
} else {
groups << "replica_groups = dense<[[0]]> : tensor<1x1xi64>";
}
body_ << " " << ssa << " = \"stablehlo.reduce_scatter\"(" << in_ssa << ") <{scatter_dimension = " << eqn.attrs.collective_dim << " : i64, "
<< groups.str() << "}> ({\n";
body_ << " ^bb0(%arg_lhs: " << scalar_type << ", %arg_rhs: " << scalar_type << "):\n";
if (eqn.attrs.collective_op == "max") {
body_ << " %red_val = stablehlo.maximum %arg_lhs, %arg_rhs : " << scalar_type << "\n";
} else if (eqn.attrs.collective_op == "min") {
body_ << " %red_val = stablehlo.minimum %arg_lhs, %arg_rhs : " << scalar_type << "\n";
} else if (eqn.attrs.collective_op == "prod" || eqn.attrs.collective_op == "multiply") {
body_ << " %red_val = stablehlo.multiply %arg_lhs, %arg_rhs : " << scalar_type << "\n";
} else {
body_ << " %red_val = stablehlo.add %arg_lhs, %arg_rhs : " << scalar_type << "\n";
}
body_ << " stablehlo.return %red_val : " << scalar_type << "\n";
body_ << " }) : (" << in_type << ") -> " << out_type << "\n";
break;
}
default:
TENZR_THROW("Unsupported OpCode in TapeEmitter: " + std::string(OpCodeToString(eqn.op)));
}
}
// Assemble MLIR module
std::ostringstream mlir;
mlir << "module {\n";
mlir << " func.func @main(";
for (size_t i = 0; i < result.in_vars.size(); ++i) {
if (i > 0) {
mlir << ", ";
}
mlir << "%arg" << i << ": " << result.in_shapes[i].ToMlirTensorType(result.in_dtypes[i]);
}
mlir << ") -> ";
if (result.out_vars.size() == 1) {
mlir << result.out_shapes[0].ToMlirTensorType(result.out_dtypes[0]) << " {\n";
} else {
mlir << "(";
for (size_t i = 0; i < result.out_vars.size(); ++i) {
if (i > 0) {
mlir << ", ";
}
mlir << result.out_shapes[i].ToMlirTensorType(result.out_dtypes[i]);
}
mlir << ") {\n";
}
mlir << body_.str();
mlir << " return ";
for (size_t i = 0; i < result.out_vars.size(); ++i) {
if (i > 0) {
mlir << ", ";
}
mlir << GetVarSsa(result.out_vars[i]);
}
mlir << " : ";
for (size_t i = 0; i < result.out_vars.size(); ++i) {
if (i > 0) {
mlir << ", ";
}
mlir << result.out_shapes[i].ToMlirTensorType(result.out_dtypes[i]);
}
mlir << "\n";
mlir << " }\n";
mlir << "}\n";
result.mlir_text = mlir.str();
return result;
}
} // namespace tenzr::txpr