OpPythonBindingGen.cpp | Explore in Territory

//===- OpPythonBindingGen.cpp - Generator of Python API for MLIR Ops ------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// OpPythonBindingGen uses ODS specification of MLIR ops to generate Python
// binding classes wrapping a generic operation API.
//
//===----------------------------------------------------------------------===//

#include "OpGenHelpers.h"

#include "mlir/TableGen/GenInfo.h"
#include "mlir/TableGen/Operator.h"
#include "llvm/ADT/StringSet.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/TableGen/Error.h"
#include "llvm/TableGen/Record.h"

usingnamespacemlir;
usingnamespacemlir::tblgen;
formatv;
Record;
RecordKeeper;

/// File header and includes.
///   {0} is the dialect namespace.
constexpr const char *fileHeader = …;

/// Template for dialect class:
///   {0} is the dialect namespace.
constexpr const char *dialectClassTemplate = …;

constexpr const char *dialectExtensionTemplate = …;

/// Template for operation class:
///   {0} is the Python class name;
///   {1} is the operation name.
constexpr const char *opClassTemplate = …;

/// Template for class level declarations of operand and result
/// segment specs.
///   {0} is either "OPERAND" or "RESULT"
///   {1} is the segment spec
/// Each segment spec is either None (default) or an array of integers
/// where:
///   1 = single element (expect non sequence operand/result)
///   0 = optional element (expect a value or std::nullopt)
///   -1 = operand/result is a sequence corresponding to a variadic
constexpr const char *opClassSizedSegmentsTemplate = …;

/// Template for class level declarations of the _ODS_REGIONS spec:
///   {0} is the minimum number of regions
///   {1} is the Python bool literal for hasNoVariadicRegions
constexpr const char *opClassRegionSpecTemplate = …;

/// Template for single-element accessor:
///   {0} is the name of the accessor;
///   {1} is either 'operand' or 'result';
///   {2} is the position in the element list.
constexpr const char *opSingleTemplate = …;

/// Template for single-element accessor after a variable-length group:
///   {0} is the name of the accessor;
///   {1} is either 'operand' or 'result';
///   {2} is the total number of element groups;
///   {3} is the position of the current group in the group list.
/// This works for both a single variadic group (non-negative length) and an
/// single optional element (zero length if the element is absent).
constexpr const char *opSingleAfterVariableTemplate = …;

/// Template for an optional element accessor:
///   {0} is the name of the accessor;
///   {1} is either 'operand' or 'result';
///   {2} is the total number of element groups;
///   {3} is the position of the current group in the group list.
/// This works if we have only one variable-length group (and it's the optional
/// operand/result): we can deduce it's absent if the `len(operation.{1}s)` is
/// smaller than the total number of groups.
constexpr const char *opOneOptionalTemplate = …;

/// Template for the variadic group accessor in the single variadic group case:
///   {0} is the name of the accessor;
///   {1} is either 'operand' or 'result';
///   {2} is the total number of element groups;
///   {3} is the position of the current group in the group list.
constexpr const char *opOneVariadicTemplate = …;

/// First part of the template for equally-sized variadic group accessor:
///   {0} is the name of the accessor;
///   {1} is either 'operand' or 'result';
///   {2} is the total number of non-variadic groups;
///   {3} is the total number of variadic groups;
///   {4} is the number of non-variadic groups preceding the current group;
///   {5} is the number of variadic groups preceding the current group.
constexpr const char *opVariadicEqualPrefixTemplate = …;

/// Second part of the template for equally-sized case, accessing a single
/// element:
///   {0} is either 'operand' or 'result'.
constexpr const char *opVariadicEqualSimpleTemplate = …;

/// Second part of the template for equally-sized case, accessing a variadic
/// group:
///   {0} is either 'operand' or 'result'.
constexpr const char *opVariadicEqualVariadicTemplate = …;

/// Template for an attribute-sized group accessor:
///   {0} is the name of the accessor;
///   {1} is either 'operand' or 'result';
///   {2} is the position of the group in the group list;
///   {3} is a return suffix (expected [0] for single-element, empty for
///       variadic, and opVariadicSegmentOptionalTrailingTemplate for optional).
constexpr const char *opVariadicSegmentTemplate = …;

/// Template for a suffix when accessing an optional element in the
/// attribute-sized case:
///   {0} is either 'operand' or 'result';
constexpr const char *opVariadicSegmentOptionalTrailingTemplate = …;

/// Template for an operation attribute getter:
///   {0} is the name of the attribute sanitized for Python;
///   {1} is the original name of the attribute.
constexpr const char *attributeGetterTemplate = …;

/// Template for an optional operation attribute getter:
///   {0} is the name of the attribute sanitized for Python;
///   {1} is the original name of the attribute.
constexpr const char *optionalAttributeGetterTemplate = …;

/// Template for a getter of a unit operation attribute, returns True of the
/// unit attribute is present, False otherwise (unit attributes have meaning
/// by mere presence):
///    {0} is the name of the attribute sanitized for Python,
///    {1} is the original name of the attribute.
constexpr const char *unitAttributeGetterTemplate = …;

/// Template for an operation attribute setter:
///    {0} is the name of the attribute sanitized for Python;
///    {1} is the original name of the attribute.
constexpr const char *attributeSetterTemplate = …;

/// Template for a setter of an optional operation attribute, setting to None
/// removes the attribute:
///    {0} is the name of the attribute sanitized for Python;
///    {1} is the original name of the attribute.
constexpr const char *optionalAttributeSetterTemplate = …;

/// Template for a setter of a unit operation attribute, setting to None or
/// False removes the attribute:
///    {0} is the name of the attribute sanitized for Python;
///    {1} is the original name of the attribute.
constexpr const char *unitAttributeSetterTemplate = …;

/// Template for a deleter of an optional or a unit operation attribute, removes
/// the attribute from the operation:
///    {0} is the name of the attribute sanitized for Python;
///    {1} is the original name of the attribute.
constexpr const char *attributeDeleterTemplate = …;

constexpr const char *regionAccessorTemplate = …;

constexpr const char *valueBuilderTemplate = …;

static llvm::cl::OptionCategory
    clOpPythonBindingCat("Options for -gen-python-op-bindings");

static llvm::cl::opt<std::string>
    clDialectName("bind-dialect",
                  llvm::cl::desc("The dialect to run the generator for"),
                  llvm::cl::init(""), llvm::cl::cat(clOpPythonBindingCat));

static llvm::cl::opt<std::string> clDialectExtensionName(
    "dialect-extension", llvm::cl::desc("The prefix of the dialect extension"),
    llvm::cl::init(""), llvm::cl::cat(clOpPythonBindingCat));

AttributeClasses;

/// Checks whether `str` would shadow a generated variable or attribute
/// part of the OpView API.
static bool isODSReserved(StringRef str) { … }

/// Modifies the `name` in a way that it becomes suitable for Python bindings
/// (does not change the `name` if it already is suitable) and returns the
/// modified version.
static std::string sanitizeName(StringRef name) { … }

static std::string attrSizedTraitForKind(const char *kind) { … }

/// Emits accessors to "elements" of an Op definition. Currently, the supported
/// elements are operands and results, indicated by `kind`, which must be either
/// `operand` or `result` and is used verbatim in the emitted code.
static void emitElementAccessors(
    const Operator &op, raw_ostream &os, const char *kind,
    unsigned numVariadicGroups, unsigned numElements,
    llvm::function_ref<const NamedTypeConstraint &(const Operator &, int)>
        getElement) { … }

/// Free function helpers accessing Operator components.
static int getNumOperands(const Operator &op) { … }
static const NamedTypeConstraint &getOperand(const Operator &op, int i) { … }
static int getNumResults(const Operator &op) { … }
static const NamedTypeConstraint &getResult(const Operator &op, int i) { … }

/// Emits accessors to Op operands.
static void emitOperandAccessors(const Operator &op, raw_ostream &os) { … }

/// Emits accessors Op results.
static void emitResultAccessors(const Operator &op, raw_ostream &os) { … }

/// Emits accessors to Op attributes.
static void emitAttributeAccessors(const Operator &op, raw_ostream &os) { … }

/// Template for the default auto-generated builder.
///   {0} is a comma-separated list of builder arguments, including the trailing
///       `loc` and `ip`;
///   {1} is the code populating `operands`, `results` and `attributes`,
///       `successors` fields.
constexpr const char *initTemplate = …;

/// Template for appending a single element to the operand/result list.
///   {0} is the field name.
constexpr const char *singleOperandAppendTemplate = …;
constexpr const char *singleResultAppendTemplate = …;

/// Template for appending an optional element to the operand/result list.
///   {0} is the field name.
constexpr const char *optionalAppendOperandTemplate = …;
constexpr const char *optionalAppendAttrSizedOperandsTemplate = …;
constexpr const char *optionalAppendResultTemplate = …;

/// Template for appending a list of elements to the operand/result list.
///   {0} is the field name.
constexpr const char *multiOperandAppendTemplate = …;
constexpr const char *multiOperandAppendPackTemplate = …;
constexpr const char *multiResultAppendTemplate = …;

/// Template for attribute builder from raw input in the operation builder.
///   {0} is the builder argument name;
///   {1} is the attribute builder from raw;
///   {2} is the attribute builder from raw.
/// Use the value the user passed in if either it is already an Attribute or
/// there is no method registered to make it an Attribute.
constexpr const char *initAttributeWithBuilderTemplate = …;

/// Template for attribute builder from raw input for optional attribute in the
/// operation builder.
///   {0} is the builder argument name;
///   {1} is the attribute builder from raw;
///   {2} is the attribute builder from raw.
/// Use the value the user passed in if either it is already an Attribute or
/// there is no method registered to make it an Attribute.
constexpr const char *initOptionalAttributeWithBuilderTemplate = …;

constexpr const char *initUnitAttributeTemplate = …;

/// Template to initialize the successors list in the builder if there are any
/// successors.
///   {0} is the value to initialize the successors list to.
constexpr const char *initSuccessorsTemplate = …;

/// Template to append or extend the list of successors in the builder.
///   {0} is the list method ('append' or 'extend');
///   {1} is the value to add.
constexpr const char *addSuccessorTemplate = …;

/// Returns true if the SameArgumentAndResultTypes trait can be used to infer
/// result types of the given operation.
static bool hasSameArgumentAndResultTypes(const Operator &op) { … }

/// Returns true if the FirstAttrDerivedResultType trait can be used to infer
/// result types of the given operation.
static bool hasFirstAttrDerivedResultTypes(const Operator &op) { … }

/// Returns true if the InferTypeOpInterface can be used to infer result types
/// of the given operation.
static bool hasInferTypeInterface(const Operator &op) { … }

/// Returns true if there is a trait or interface that can be used to infer
/// result types of the given operation.
static bool canInferType(const Operator &op) { … }

/// Populates `builderArgs` with result names if the builder is expected to
/// accept them as arguments.
static void
populateBuilderArgsResults(const Operator &op,
                           SmallVectorImpl<std::string> &builderArgs) { … }

/// Populates `builderArgs` with the Python-compatible names of builder function
/// arguments using intermixed attributes and operands in the same order as they
/// appear in the `arguments` field of the op definition. Additionally,
/// `operandNames` is populated with names of operands in their order of
/// appearance.
static void populateBuilderArgs(const Operator &op,
                                SmallVectorImpl<std::string> &builderArgs,
                                SmallVectorImpl<std::string> &operandNames) { … }

/// Populates `builderArgs` with the Python-compatible names of builder function
/// successor arguments. Additionally, `successorArgNames` is also populated.
static void
populateBuilderArgsSuccessors(const Operator &op,
                              SmallVectorImpl<std::string> &builderArgs,
                              SmallVectorImpl<std::string> &successorArgNames) { … }

/// Populates `builderLines` with additional lines that are required in the
/// builder to set up operation attributes. `argNames` is expected to contain
/// the names of builder arguments that correspond to op arguments, i.e. to the
/// operands and attributes in the same order as they appear in the `arguments`
/// field.
static void
populateBuilderLinesAttr(const Operator &op, ArrayRef<std::string> argNames,
                         SmallVectorImpl<std::string> &builderLines) { … }

/// Populates `builderLines` with additional lines that are required in the
/// builder to set up successors. successorArgNames is expected to correspond
/// to the Python argument name for each successor on the op.
static void
populateBuilderLinesSuccessors(const Operator &op,
                               ArrayRef<std::string> successorArgNames,
                               SmallVectorImpl<std::string> &builderLines) { … }

/// Populates `builderLines` with additional lines that are required in the
/// builder to set up op operands.
static void
populateBuilderLinesOperand(const Operator &op, ArrayRef<std::string> names,
                            SmallVectorImpl<std::string> &builderLines) { … }

/// Python code template for deriving the operation result types from its
/// attribute:
///   - {0} is the name of the attribute from which to derive the types.
constexpr const char *deriveTypeFromAttrTemplate = …;

/// Python code template appending {0} type {1} times to the results list.
constexpr const char *appendSameResultsTemplate = …;

/// Appends the given multiline string as individual strings into
/// `builderLines`.
static void appendLineByLine(StringRef string,
                             SmallVectorImpl<std::string> &builderLines) { … }

/// Populates `builderLines` with additional lines that are required in the
/// builder to set up op results.
static void
populateBuilderLinesResult(const Operator &op, ArrayRef<std::string> names,
                           SmallVectorImpl<std::string> &builderLines) { … }

/// If the operation has variadic regions, adds a builder argument to specify
/// the number of those regions and builder lines to forward it to the generic
/// constructor.
static void populateBuilderRegions(const Operator &op,
                                   SmallVectorImpl<std::string> &builderArgs,
                                   SmallVectorImpl<std::string> &builderLines) { … }

/// Emits a default builder constructing an operation from the list of its
/// result types, followed by a list of its operands. Returns vector
/// of fully built functionArgs for downstream users (to save having to
/// rebuild anew).
static SmallVector<std::string> emitDefaultOpBuilder(const Operator &op,
                                                     raw_ostream &os) { … }

static void emitSegmentSpec(
    const Operator &op, const char *kind,
    llvm::function_ref<int(const Operator &)> getNumElements,
    llvm::function_ref<const NamedTypeConstraint &(const Operator &, int)>
        getElement,
    raw_ostream &os) { … }

static void emitRegionAttributes(const Operator &op, raw_ostream &os) { … }

/// Emits named accessors to regions.
static void emitRegionAccessors(const Operator &op, raw_ostream &os) { … }

/// Emits builder that extracts results from op
static void emitValueBuilder(const Operator &op,
                             SmallVector<std::string> functionArgs,
                             raw_ostream &os) { … }

/// Emits bindings for a specific Op to the given output stream.
static void emitOpBindings(const Operator &op, raw_ostream &os) { … }

/// Emits bindings for the dialect specified in the command line, including file
/// headers and utilities. Returns `false` on success to comply with Tablegen
/// registration requirements.
static bool emitAllOps(const RecordKeeper &records, raw_ostream &os) { … }

static GenRegistration
    genPythonBindings("gen-python-op-bindings",
                      "Generate Python bindings for MLIR Ops", &emitAllOps);
llvm/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
