using type = T; // The data type of the argument

std::string name;

std::optional<T> default_value;

// short_circuit: Symbolic values are not Python scalars.

if (v->isSymbolic()) {

return false;

}

// Check if the dtype is compatible with a Python scalar.

// e.g., Python scalar cannot distiguish between ComplexDouble and

// ComplexFloat. In this case, define_scalar must be used to specify custom

// dtype.

PrimDataType value_dtype(std::get<PrimDataType>(v->dtype().type));

switch (value_dtype) {

case PrimDataType::Bool:

case PrimDataType::Int:

case PrimDataType::Index:

case PrimDataType::ComplexDouble:

case PrimDataType::Double:

return true;

default:

return false;

}

public:

PythonPrinter(std::ostream& os) : os_(&os) {}

// Generate a python string for a string value.

std::string toString(const std::string& s) {

return s;

}

// Generate a python string for a boolean value.

std::string toString(bool b) {

return b ? "True" : "False";

}

// Generate a python string for an int64_t value.

std::string toString(int64_t i) {

return std::to_string(i);

}

// Generate a python string for an size_t value.

std::string toString(size_t i) {

return std::to_string(i);

}

// Generate a python string for a complex double value.

std::string toString(std::complex<double> c) {

std::stringstream ss;

ss << std::showpoint << std::real(c) << "+" << std::showpoint

<< std::imag(c) << "j";

return ss.str();

}

// Generate a python string for a double value.

std::string toString(double d) {

if (std::isinf(d)) {

if (std::signbit(d)) {

return "float(\"-inf\")";

} else {

return "float(\"inf\")";

}

} else if (std::isnan(d)) {

return "float(\"nan\")";

} else {

std::stringstream ss;

ss << std::showpoint << d;

return ss.str();

}

}

// Generate a python string for a Datatype.

std::string toString(DataType dtype) {

return dtypeToPyString(std::get<PrimDataType>(dtype.type));

}

// Generate a python string for a PolymorphicValue with simple types.

std::string toString(const PolymorphicValue& pv) {

if (pv.is<bool>()) {

return toString(pv.as<bool>());

} else if (pv.is<int64_t>()) {

return toString(pv.as<int64_t>());

} else if (pv.is<std::complex<double>>()) {

return toString(pv.as<std::complex<double>>());

} else if (pv.is<double>()) {

return toString(pv.as<double>());

} else if (pv.is<std::monostate>()) {

return "None";

} else {

NVF_THROW("Unsupported PolymorphicValue type");

}

}

// Generate a unique name for a Val. Map val to name to track Val's lifetime.

std::string toString(const nvfuser::Val* v, bool is_lvalue = false) {

std::stringstream ss;

if (v == nullptr) {

return "None";

} else if (v->isA<TensorView>()) {

ss << "tv" << v->name();

} else if (!is_lvalue && isPythonScalar(v)) {

ss << toString(v->value());

} else {

ss << "c" << v->name();

}

return ss.str();

}

// Generate a python string for an optional value.

template <typename T>

std::string toString(std::optional<T> optional, bool skip_none = true) {

if (optional.has_value()) {

return toString(optional.value());

} else if (!skip_none) {

return "None";

} else {

return "";

}

}

// Generate a python string for a keyword argument.

template <typename T>

std::string toString(

const std::string& name,

T value,

const std::string& separator) {

std::string result = toString(value);

if (result.empty()) {

return "";

}

return separator + name + "=" + result;

}

// Generate a python list of values.

template <typename T>

std::string toString(const std::vector<T>& vec, bool is_list = true) {

std::stringstream ss;

if (is_list) {

ss << "[";

}

for (auto&& [i, val] : enumerate(vec)) {

if constexpr (is_optional_v<T>) {

ss << toString(val, /*skip_none=*/false);

} else {

ss << toString(val);

}

if (i < std::ssize(vec) - 1) {

ss << ", ";

}

}

if (is_list) {

ss << "]";

}

return ss.str();

}

// Generate a python list of values.

std::string generateOutputs(const std::vector<const nvfuser::Val*>& vec) {

std::stringstream ss;

for (auto&& [i, val] : enumerate(vec)) {

if (val == nullptr) {

ss << "_";

} else {

NVF_ERROR(

!isPythonScalar(val),

"A constant scalar Val* cannot be an output lvalue. Got\t",

val->toString());

ss << toString(val, /*is_lvalue=*/true);

}

if (i < std::ssize(vec) - 1) {

ss << ", ";

}

}

return ss.str();

}

// Generate a python list of values.

template <typename... Ts>

std::string generateList(std::tuple<Ts...> const& args) {

if (sizeof...(Ts) == 0) {

return "";

}

std::stringstream ss;

std::apply(

[&](Ts const&... tuple_args) {

size_t i = 0;

(((ss << (i > 0 ? ", " : "") << toString(tuple_args)), ++i), ...);

},

args);

return ss.str();

}

// Generate a python list of values with string keyword arguments.

template <typename... Ts>

std::string generateNamedList(

const std::vector<std::string>& argument_names,

std::tuple<Ts...> const& args) {

NVF_ERROR(

argument_names.size() == sizeof...(Ts),

"Input argument names and args must have the same size.");

// Use std::apply to unpack tuple of arguments into a lambda. The lambda

// contains a C++17 fold expression on a comma operator that writes each

// argument to stringstream and increments argument position.

std::stringstream ss;

std::apply(

[this, &ss, &argument_names](Ts const&... tuple_args) {

size_t i = 0;

(((ss << toString(

argument_names[i], tuple_args, (i > 0 ? ", " : ""))),

++i),

...);

},

args);

return ss.str();

}

// Generate a python list of values with string keyword arguments.

// * A tuple of default argument values is provided to the function.

// * If the default argument is optional and equal to the provided argument,

// then skip printing the keyword-argument pair.

template <typename... Ds, typename... Ts>

std::string generateNamedList(

std::tuple<Ds...> const& default_args,

std::tuple<Ts...> const& args) {

NVF_ERROR(

sizeof...(Ds) == sizeof...(Ts),

"The default and given arguments must have the same size.");

// This immediately-invoked generic lambda uses a C++17 fold expression

// to emulate a loop over the tuple elements.

//

// 1. `std::make_index_sequence` generates a compile-time sequence of

// indices (0, 1, 2...).

// 2. The lambda accepts this sequence, deducing the indices into the

// template pack `Is...`.

// 3. A fold expression over the comma operator expands the code for each

// index.

// 4. A ternary operator `(condition ? ... : ...)` performs the conditional

// logic.

// 5. If condition is true, another comma operator

// `(write_to_stream, increment_counters)` chains the side effects of

// writing to the stringstream and advancing the counters.

// 6. If condition is false, increment `printed_arg_pos` counter by zero.

std::stringstream ss;

[&]<std::size_t... Is>(std::index_sequence<Is...>) {

size_t printed_arg_pos = 0;

(((!std::get<Is>(default_args).default_value.has_value() ||

std::get<Is>(default_args).default_value.value() != std::get<Is>(args))

? ((ss << toString(

std::get<Is>(default_args).name,

std::get<Is>(args),

(printed_arg_pos > 0 ? ", " : ""))),

++printed_arg_pos)

: (printed_arg_pos += 0)),

...);

}(std::make_index_sequence<sizeof...(Ts)>{}); // Generate indices 0..N-1

return ss.str();

}

// Generate a python operation with a list of inputs and outputs.

void generateOperation(

const std::string& op_name,

const std::vector<const nvfuser::Val*>& inputs,

const std::vector<const nvfuser::Val*>& outputs) {

(*os_) << kTab;

if (!outputs.empty()) {

(*os_) << generateOutputs(outputs) << " = ";

}

(*os_) << op_name << "(" << toString(inputs, /*is_list=*/false) << ")\n";

}

// Generate a python operation with a list of inputs and outputs.

// A string keyword argument is added for each input. The default_kwargs

// argument allows skipping arguments if it isn't strictly necessary.

template <

typename... arg_types,

typename... default_kwarg_types,

typename... kwargs_types>

void generateKwargsOperation(

const std::string& op_name,

const std::tuple<arg_types...>& args,

const std::tuple<default_kwarg_types...>& default_kwargs,

const std::tuple<kwargs_types...>& kwargs,

const std::vector<const nvfuser::Val*>& outputs) {

std::string kwargs_str = generateNamedList(default_kwargs, kwargs);

constexpr bool any_arguments = sizeof...(arg_types) == 0;

std::string connect = (any_arguments || kwargs_str.empty()) ? "" : ", ";

(*os_) << kTab << generateOutputs(outputs) << " = " << op_name << "("

<< generateList(args) << connect << kwargs_str << ")\n";

}

// Generate a python operation with a list of inputs and outputs.

// A string is added for each keyword argument.

//

// NOTES

// ------

// - args and kwargs are a tuple, so it accepts a fixed set of arguments of

// any type at compile-time.

// - outputs is a vector of nvfuser values that is converted into a string.

template <typename... arg_types, typename... kwargs_types>

void generateKwargsOperation(

const std::string& op_name,

const std::tuple<arg_types...>& args,

const std::vector<std::string>& kwargs_names,

const std::tuple<kwargs_types...>& kwargs,

const std::vector<const nvfuser::Val*>& outputs) {

std::string connect = (sizeof...(arg_types) == 0) ? "" : ", ";

(*os_) << kTab << generateOutputs(outputs) << " = " << op_name << "("

<< generateList(args) << connect

<< generateNamedList(kwargs_names, kwargs) << ")\n";

}

// Generate a python operation with a list of inputs and a single output.

// A string is added for each keyword argument.

//

// NOTES

// ------

// - args and kwargs are a tuple, so it accepts a fixed set of arguments of

// any type at compile-time.

// - output_name is a string.

template <typename... arg_types, typename... kwargs_types>

void generateKwargsOperation(

const std::string& op_name,

const std::tuple<arg_types...>& args,

const std::vector<std::string>& kwargs_names,

const std::tuple<kwargs_types...>& kwargs,

const std::string& output_name) {

std::string connect = (sizeof...(arg_types) == 0) ? "" : ", ";

(*os_) << kTab << output_name << " = " << op_name << "("

<< generateList(args) << connect

<< generateNamedList(kwargs_names, kwargs) << ")\n";

}

// Generate a python operation with a list of inputs and outputs.

// A string is added for each keyword argument.

//

// NOTES

// ------

// - args and outputs are vectors of nvfuser values that are converted into

// strings.

// - kwargs is a tuple, so it accepts a fixed set of arguments of

// any type at compile-time.

template <typename... kwargs_types>

void generateKwargsOperation(

const std::string& op_name,

const std::vector<Val*>& args,

const std::vector<std::string>& kwargs_names,

const std::tuple<kwargs_types...>& kwargs,

const std::vector<const nvfuser::Val*>& outputs) {

std::string connect = args.empty() ? "" : ", ";

(*os_) << kTab << generateOutputs(outputs) << " = " << op_name << "("

<< toString(args) << connect

<< generateNamedList(kwargs_names, kwargs) << ")\n";

}

// Generate a python definition for a FusionDefinition.

void generateFusionDefinition() {

(*os_) << "def nvfuser_fusion(fd : FusionDefinition) -> None :\n";

}

private:

//! The stream to print the python function to.

std::ostream* os_;

//! Indentation for python code.

static constexpr const char* kTab = " ";

public:

// Returns a map from the values in the CPP fusion to its corresponding

// FusionDefinition State index.

static void print(std::ostream& os, Fusion* fusion) {

PythonTranslator translator(os, fusion);

translator.translate();

}

private:

PythonTranslator(std::ostream& os, Fusion* fusion)

: printer_(os), fusion_(fusion) {}

// The output TensorView shape can be dynamic for operations like ReshapeOp.

// Check that all dynamic scalar dependencies are handled first before

// handling the expression.

bool checkDynamicShapeDependency(const Expr* op) {

// short-circuit: Only check operations with dynamic shapes encoded in the

// output TensorView.

if (!op->isOneOf<ReshapeOp, ExpandOp, FullOp>()) {

return true;

}

NVF_ERROR_EQ(op->outputs().size(), 1);

const std::vector<IterDomain*>& logical_out_domain =

op->output(0)->as<TensorView>()->domain()->logical();

std::vector<Val*> logical_domain_extents;

std::ranges::copy(

logical_out_domain | std::views::transform([](IterDomain* id) {

return id->getMaybeExpandedExtent();

}),

std::back_inserter(logical_domain_extents));

return std::ranges::all_of(logical_domain_extents, [&](Val* v) {

return v->definition() == nullptr || visited_vals_.count(v) > 0;

});

}

// Gather the expressions necessary to create a scalar value.

std::vector<Expr*> gatherScalarExpressions(Val* v) {

NVF_ERROR(v != nullptr);

NVF_ERROR(v->isScalar());

// short-circuit: v does not have a definition.

if (v->definition() == nullptr) {

return {};

}

std::vector<Expr*> expression_chain;

std::unordered_set<Expr*> visited;

std::vector<Expr*> to_visit = {v->definition()};

while (!to_visit.empty()) {

Expr* e = to_visit.back();

to_visit.pop_back();

expression_chain.push_back(e);

visited.insert(e);

for (Val* input : e->inputs()) {

// short-circuit: input does not have a definition.

if (input->definition() == nullptr) {

continue;

}

// short-circuit: input definition is already visited.

if (visited.count(input->definition()) > 0) {

continue;

}

to_visit.push_back(input->definition());

}

}

return expression_chain;

}

// Gather the scalar expressions necessary to create the logical domain for a

// TensorView.

std::vector<Expr*> gatherScalarExpressions(TensorView* tv) {

NVF_ERROR(tv != nullptr);

std::vector<Expr*> logical_domain_expressions;

const std::vector<IterDomain*>& logical_out_domain =

tv->domain()->logical();

for (IterDomain* id : logical_out_domain) {

std::vector<Expr*> extent_definitions =

gatherScalarExpressions(id->getMaybeExpandedExtent());

logical_domain_expressions.insert(

logical_domain_expressions.end(),

extent_definitions.begin(),

extent_definitions.end());

}

return logical_domain_expressions;

}

// Check that all of the expression's inputs are defined in FusionDefinition.

bool checkExpressionDependencies(Expr* e) {

// short-circuit: Found an operation without all its dynamic shape

// dependencies.

if (!checkDynamicShapeDependency(e)) {

return false;

}

return std::all_of(

e->inputs().begin(), e->inputs().end(), [&](const Val* v) {

return isPythonScalar(v) || visited_vals_.count(v) > 0;

});

}

void translate() {

printer_.generateFusionDefinition();

// Add Fusion inputs to FusionDefinition

for (nvfuser::Val* v : fusion_->inputs()) {

dispatch(v);

}

// Gather all expressions in CPP Fusion.

const std::vector<nvfuser::Expr*> fusion_exprs = fusion_->exprs();

std::deque<nvfuser::Expr*> to_visit(

fusion_exprs.begin(), fusion_exprs.end());

// Scalar expressions are not handled by Fusion::exprs, so gather them

// manually.

for (Expr* e : to_visit) {

if (e->isOneOf<ReshapeOp, ExpandOp, FullOp>()) {

NVF_ERROR_EQ(e->outputs().size(), 1);

std::vector<Expr*> extent_definitions =

gatherScalarExpressions(e->output(0)->as<TensorView>());

to_visit.insert(

to_visit.end(),

extent_definitions.begin(),

extent_definitions.end());

}

}

// Topological search of Fusion expressions

size_t skip_count = 0;

std::unordered_set<nvfuser::Expr*> visited;

while (!to_visit.empty()) {

Expr* e = to_visit.front();

to_visit.pop_front();

NVF_ERROR(

skip_count <= to_visit.size(),

"Cycle detected: None of the expressions can be processed!");

// short-circuit: skip if already visited

if (visited.count(e) > 0) {

continue;

}

// TODO: short-circuit: skip Split and Merge expressions created by

// Reshape

// TODO: short-circuit: skip Resize expressions created by Slice

// TODO: direct bindings does not support scheduled expressions.

// Handle scalars and constants not generated by separate expression.

std::vector<Val*> scalars;

std::ranges::copy_if(

e->inputs(), std::back_inserter(scalars), [](Val* v) {

return v->isScalar();

});

std::ranges::for_each(scalars, [this](const Val* v) { dispatch(v); });

// short-circuit: add to back of stack if not all of the expression's

// dependencies are satisfied.

if (!checkExpressionDependencies(e)) {

++skip_count;

to_visit.push_back(e);

continue;

}

// Create string representation given inputs, outputs, and attributes.

visited.insert(e);

dispatch(e);

skip_count = 0;

}

// Add tensor outputs and handle aliased outputs

std::unordered_set<nvfuser::Val*> visited_alias_output;

for (nvfuser::Val* v : fusion_->outputs()) {

NVF_ERROR(v->isA<TensorView>());

const AliasInfo& alias_info = fusion_->getOutputAlias(v);

switch (alias_info.type) {

case AllocationType::New: {

handleOutput(v->as<TensorView>());

break;

}

case AllocationType::ReuseBuffer: {

// Only apply aliasing once

if (visited_alias_output.count(v) == 0) {

visited_alias_output.insert(v);

handleOutput(v->as<TensorView>(), alias_info);

}

// If not hide_output, then the aliased output is returned as a

// fusion output.

if (alias_info.visibility == OutputVisibility::kVisible) {

handleOutput(v->as<TensorView>());

}

break;

}

default:

NVF_THROW("Unsupported AllocationType");

}

}

}

// =================================================================================

// Filter Functions

// Gather all TensorViews and FusionDefinition indices

std::vector<const nvfuser::Val*> tensors() {

std::vector<const nvfuser::Val*> tensors;

std::ranges::copy_if(

visited_vals_, std::back_inserter(tensors), [](const nvfuser::Val* v) {

return v->isA<TensorView>();

});

return tensors;

}

// =================================================================================

// Create scalar for given nvfuser value. The nvfuser value must not already

// exist and have a definition. It can be a fusion input, a constant, or a

// tensor's extent.

void handle(const Val* v) final {

NVF_ERROR(v != nullptr);

// short-circuit: scalar definition has a definition

if (v->definition() != nullptr) {

return;

}

// short-circuit: value already exists in FusionDefinition

if (visited_vals_.count(v) > 0) {

return;

}

// short-circuit: print python scalar directly

if (isPythonScalar(v)) {

return;

}

visited_vals_.insert(v);

// Since scalars can come from TensorView dimension sizes, search through

// all TensorViews for an iterDomain whose extent matches the desired

// value and then use size op.

for (const nvfuser::Val* tv_val : tensors()) {

const auto* tv = tv_val->as<TensorView>();

// Get extents for each IterDomain

std::vector<IterDomain*> filtered_logical_domain =

TensorDomain::noReductions(tv->domain()->logical());

std::vector<Val*> extents;

extents.reserve(filtered_logical_domain.size());

std::ranges::copy(

filtered_logical_domain | std::views::transform([](IterDomain* id) {

return id->getMaybeExpandedExtent();

}),

std::back_inserter(extents));

// Check if value matches iterdomain extent

auto iter = std::ranges::find(extents, v);

if (iter == extents.end()) {

continue;

}

int64_t dim = std::distance(extents.begin(), iter);

static const std::vector<std::string> argument_names = {"dim"};

printer_.generateKwargsOperation(

"fd.ops.size",

std::make_tuple(tv),

argument_names,

std::make_tuple(dim),

{v});

return;

}

static const std::vector<std::string> argument_names = {"dtype"};

printer_.generateKwargsOperation(

"fd.define_scalar",

std::make_tuple(v->value()),

argument_names,

std::make_tuple(v->dtype()),

{v});

}

// Add Tensor value to Fusion Definition

void handle(const TensorView* tv) final {

NVF_ERROR(tv != nullptr);

// short-circuit: value already exists in FusionDefinition

if (visited_vals_.count(tv) > 0) {

return;

}

visited_vals_.insert(tv);

std::vector<int64_t> shape;

std::transform(

tv->domain()->logical().begin(),

tv->domain()->logical().end(),

std::back_inserter(shape),

[](IterDomain* id) {

return (id->getMaybeExpandedExtent()->isConstScalar())

? id->getMaybeExpandedExtent()->evaluate().as<int64_t>()

: -1;

});

const std::vector<int64_t>& stride_order = tv->domain()->strideOrder();

static const std::vector<std::string> argument_names = {

"shape", "contiguity", "dtype", "is_cpu", "stride_order"};

printer_.generateKwargsOperation(

"fd.define_tensor",

std::make_tuple(),

argument_names,

std::make_tuple(

shape,

tv->domain()->contiguity(),

tv->dtype(),

tv->isCpuScalar(),

(stride_order.empty()) ? std::nullopt

: std::make_optional(stride_order)),

{tv});

}

// =================================================================================

// Utility functions

// Create a vector for the logical domain of TensorView.

// Used with ReshapeOp, ExpandOp, and FullOp handlers

std::vector<Val*> getShape(TensorView* tv) {

const std::vector<IterDomain*>& logical_out_domain =

tv->domain()->logical();

std::vector<Val*> logical_domain_extents;

// Use expanded extent if available for IterDomain.

std::ranges::copy(

logical_out_domain | std::views::transform([](IterDomain* id) {

return id->getMaybeExpandedExtent();

}),

std::back_inserter(logical_domain_extents));

return logical_domain_extents;

}

// Find integer index corresponding with reduction iterDomains

std::vector<int64_t> getReductionAxes(TensorView* tv) {

std::vector<int64_t> axes;

const std::vector<IterDomain*>& logical_domain = tv->domain()->logical();

for (int64_t dim : c10::irange((int64_t)logical_domain.size())) {

if (logical_domain.at(dim)->isReduction()) {

axes.push_back(dim);

}

}

return axes;

}

// =================================================================================

// Handle add_output variants

// Add Tensor output to FusionDefinition

void handleOutput(const TensorView* tv) {

NVF_ERROR(tv != nullptr);

printer_.generateOperation("fd.add_output", {tv}, {});

}

// Alias output Tensor with input tensor

void handleOutput(const TensorView* tv, const AliasInfo& alias_info) {

NVF_ERROR(tv != nullptr);

printer_.generateOperation(

"fd.add_output", {tv, alias_info.aliased_io}, {});

}

// =================================================================================

// Map CPP Expression classes to corresponding RecordFunctors in

// python_frontend

void handle(const UnaryOp* uop) final {

NVF_ERROR(uop != nullptr);

// short-circuit: Handle cast operation separately

if (uop->getUnaryOpType() == UnaryOpType::Cast) {

return handleCastOp(uop);

}

// Map remaining UnaryOp to python_frontend

visited_vals_.insert(uop->out());

printer_.generateOperation(

"fd.ops." + nvfuser::python::toString(uop), {uop->in()}, {uop->out()});

}

void handleCastOp(const UnaryOp* uop) {

NVF_ERROR(uop->getUnaryOpType() == UnaryOpType::Cast);

visited_vals_.insert(uop->out());

static const std::vector<std::string> argument_names = {"dtype"};

printer_.generateKwargsOperation(

"fd.ops.cast",

std::make_tuple(uop->in()),

argument_names,

std::make_tuple(uop->out()->dtype()),

{uop->out()});

}

void handle(const BinaryOp* bop) final {

NVF_ERROR(bop != nullptr);

if (visited_vals_.count(bop->out()) > 0) {

return;

}

visited_vals_.insert(bop->out());

printer_.generateOperation(

"fd.ops." + nvfuser::python::toString(bop),

{bop->lhs(), bop->rhs()},

{bop->out()});

}

void handle(const TernaryOp* top) final {

NVF_ERROR(top != nullptr);

visited_vals_.insert(top->out());

printer_.generateOperation(

"fd.ops." + nvfuser::python::toString(top),

{top->in1(), top->in2(), top->in3()},

{top->out()});

}

void handle(const ReductionOp* rop) final {

NVF_ERROR(rop != nullptr);

NVF_ERROR(rop->out()->isA<TensorView>());

visited_vals_.insert(rop->out());

// The min and max reduction operations expect the dtype argument to by

// PrimDataType::Null

DataType dtype = (rop->getReductionOpType() == BinaryOpType::Min ||

rop->getReductionOpType() == BinaryOpType::FMin ||

rop->getReductionOpType() == BinaryOpType::Max ||

rop->getReductionOpType() == BinaryOpType::FMax)

? DataType::Null

: rop->out()->dtype();

std::vector<int64_t> dims = getReductionAxes(rop->out()->as<TensorView>());

// TODO: keepdim is always False in ReductionOp because a separate

// BroadcastOp node exists if keepdim is True. Detect this pattern to

// minify the python definition.

static const auto default_args = std::make_tuple(

KeywordArgument<decltype(dims)>{

.name = "dims", .default_value = std::nullopt},

KeywordArgument<bool>{.name = "keepdim", .default_value = false},

KeywordArgument<DataType>{

.name = "dtype", .default_value = DataType::Null});

printer_.generateKwargsOperation(

"fd.ops." + nvfuser::python::toString(rop),

std::make_tuple(rop->in()),

default_args,

std::make_tuple(dims, false, dtype),

{rop->out()});

}

void handle(const ScanOp* sop) final {

NVF_ERROR(sop != nullptr);

visited_vals_.insert(sop->out());

static const auto default_args = std::make_tuple(

KeywordArgument<int64_t>{.name = "dim", .default_value = std::nullopt});

printer_.generateKwargsOperation(

"fd.ops." + toString(sop),

std::make_tuple(sop->in()),

default_args,

std::make_tuple(sop->dim()),

{sop->out()});

}

void handle(const WelfordOp* wop) final {

NVF_ERROR(wop != nullptr);

NVF_ERROR(wop->initAvg()->evaluate().as<double>() == 0.0);

NVF_ERROR(wop->initVar()->evaluate().as<double>() == 0.0);

NVF_ERROR(wop->initN()->evaluate().as<int64_t>() == 0);

visited_vals_.insert(wop->outAvg());

visited_vals_.insert(wop->outVar());

visited_vals_.insert(wop->outN());

static const std::vector<std::string> argument_names = {"dims"};

printer_.generateKwargsOperation(

"fd.ops.welford",

std::make_tuple(wop->in()),

argument_names,

std::make_tuple(getReductionAxes(wop->outAvg()->as<TensorView>())),

{wop->outAvg(), wop->outVar(), wop->outN()});

}

void handle(const BroadcastOp* bcast_op) final {

NVF_ERROR(bcast_op != nullptr);

visited_vals_.insert(bcast_op->out());

static const std::vector<std::string> broadcast_argument_names = {

"is_broadcast_dim"};

printer_.generateKwargsOperation(

"fd.ops.broadcast",

std::make_tuple(bcast_op->in()),

broadcast_argument_names,

std::make_tuple(bcast_op->getBroadcastDimFlags()),

{bcast_op->out()});

}

void handle(const MatmulOp* matmul_op) final {

NVF_ERROR(matmul_op != nullptr);

visited_vals_.insert(matmul_op->out());

printer_.generateOperation(

"fd.ops.matmul",

{matmul_op->inA(), matmul_op->inB()},

{matmul_op->out()});

}

void handle(const LinearOp* lop) final {

NVF_ERROR(lop != nullptr);

visited_vals_.insert(lop->out());

static const auto default_args = std::make_tuple(

KeywordArgument<TensorView*>{.name = "bias", .default_value = nullptr});

printer_.generateKwargsOperation(

"fd.ops.linear",

std::make_tuple(lop->inA(), lop->inB()),

default_args,

std::make_tuple(lop->bias()),

{lop->out()});

}

void handle(const GroupedMmaOp* gmm_op) final {

NVF_ERROR(gmm_op != nullptr);

TensorView* out_tv = gmm_op->out();

visited_vals_.insert(gmm_op->out());

int64_t out_block_scale_size = 0;

PrimDataType out_block_scale_dtype = DataType::BFloat16;

bool out_gamma = false;

TensorView* out_block_scale_tv = gmm_op->outScale();

if (out_block_scale_tv != nullptr) {

visited_vals_.insert(gmm_op->outScale());

const std::vector<IterDomain*>& logical =

out_block_scale_tv->getLogicalDomain();

Val* block_size_extent = logical.at(logical.size() - 1)->extent();

NVF_CHECK(

block_size_extent->isConstInt(),

"Block size extent needs to be a constant integer");

out_block_scale_size = block_size_extent->evaluate().as<int64_t>();

out_block_scale_dtype =

std::get<PrimDataType>(out_block_scale_tv->dtype().type);

}

TensorView* out_gamma_tv = gmm_op->outGamma();

if (out_gamma_tv != nullptr) {

visited_vals_.insert(gmm_op->outGamma());

out_gamma = true;

}