torch-mlir/include/npcomp/Dialect/Basicpy/IR/BasicpyOps.td

//===- BasicPyOps.td - Basic Python ops --------------------*- tablegen -*-===//
//
// This file is licensed 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
//
//===----------------------------------------------------------------------===//

#ifndef NPCOMP_DIALECT_BASICPY_IR_BASICPY_OPS
#define NPCOMP_DIALECT_BASICPY_IR_BASICPY_OPS

include "BasicpyDialect.td"
include "mlir/Interfaces/CallInterfaces.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/IR/SymbolInterfaces.td"

//===----------------------------------------------------------------------===//
// Predicates
//===----------------------------------------------------------------------===//

def BoolOrI1Type : AnyTypeOf<[Basicpy_BoolType, I1], "Python bool or i1">;

//===----------------------------------------------------------------------===//
// Binary operation enum
// The name matches the operation name in the python AST ("Add", "Mult", etc).
//===----------------------------------------------------------------------===//

def BINOP_ADD      : StrEnumAttrCase<"Add">;
def BINOP_BITAND   : StrEnumAttrCase<"BitAnd">;
def BINOP_BITOR    : StrEnumAttrCase<"BitOr">;
def BINOP_BITXOR   : StrEnumAttrCase<"BitXor">;
def BINOP_DIV      : StrEnumAttrCase<"Div">;
def BINOP_FLOORDIV : StrEnumAttrCase<"FloorDiv">;
def BINOP_LSHIFT   : StrEnumAttrCase<"LShift">;
def BINOP_MATMULT  : StrEnumAttrCase<"MatMult">;
def BINOP_MOD      : StrEnumAttrCase<"Mod">;
def BINOP_MULT     : StrEnumAttrCase<"Mult">;
def BINOP_RSHIFT   : StrEnumAttrCase<"RShift">;
def BINOP_SUB      : StrEnumAttrCase<"Sub">;

def BinaryOperationAttr : StrEnumAttr<
    "BinaryOperation", "Operation for a binary expression", [
      BINOP_ADD,
      BINOP_BITAND,
      BINOP_BITOR,
      BINOP_BITXOR,
      BINOP_DIV,
      BINOP_FLOORDIV,
      BINOP_LSHIFT,
      BINOP_MATMULT,
      BINOP_MOD,
      BINOP_MULT,
      BINOP_RSHIFT,
      BINOP_SUB,
    ]> {
  let cppNamespace = "::mlir::NPCOMP::Basicpy";
}

//===----------------------------------------------------------------------===//
// Comparison operation enum
// The name matches the operation name in the python AST ("Lt", "Gt", etc).
//===----------------------------------------------------------------------===//

def CMPOP_EQ      : StrEnumAttrCase<"Eq">;
def CMPOP_GT      : StrEnumAttrCase<"Gt">;
def CMPOP_GTE     : StrEnumAttrCase<"GtE">;
def CMPOP_IN      : StrEnumAttrCase<"In">;
def CMPOP_IS      : StrEnumAttrCase<"Is">;
def CMPOP_ISNOT   : StrEnumAttrCase<"IsNot">;
def CMPOP_LT      : StrEnumAttrCase<"Lt">;
def CMPOP_LTE     : StrEnumAttrCase<"LtE">;
def CMPOP_NEQ     : StrEnumAttrCase<"NotEq">;
def CMPOP_NOTIN   : StrEnumAttrCase<"NotIn">;

def CompareOperationAttr : StrEnumAttr<
    "CompareOperation", "Comparison operator", [
      CMPOP_EQ,
      CMPOP_GT,
      CMPOP_GTE,
      CMPOP_IN,
      CMPOP_IS,
      CMPOP_ISNOT,
      CMPOP_LT,
      CMPOP_LTE,
      CMPOP_NEQ,
      CMPOP_NOTIN,
    ]> {
  let cppNamespace = "::mlir::NPCOMP::Basicpy";
}

//===----------------------------------------------------------------------===//
// Ops
//===----------------------------------------------------------------------===//

def Basicpy_BinaryCompareOp : Basicpy_Op<"binary_compare", []> {
  let summary = "Performs a comparison between two operands";
  let description = [{
    This op performs only one step of a potentially multi-step short
    circuit comparison.
    See: https://docs.python.org/3/reference/expressions.html#comparisons
  }];
  let arguments = (ins
    AnyType:$left,
    AnyType:$right,
    CompareOperationAttr:$operation
  );
  let results = (outs
    Basicpy_BoolType:$result
  );
  let assemblyFormat = "$left $operation $right attr-dict `:` type(operands)";
}

def Basicpy_BinaryExprOp : Basicpy_Op<"binary_expr", []> {
  let summary = "Binary expression";
  let description = [{
    An expression between two operands as generated by the AST BinOp node.
  }];
  let arguments = (ins
    AnyType:$left,
    AnyType:$right,
    BinaryOperationAttr:$operation
  );
  let results = (outs
    AnyType:$result
  );
  let assemblyFormat = "$left $operation $right attr-dict `:` functional-type(operands, results)";
}

def Basicpy_BoolCastOp : Basicpy_Op<"bool_cast", [NoSideEffect]> {
  let summary = "Casts between BoolType and i1";
  let description = [{
    When interfacing with lower level dialect or progressively lowering
    the Python BoolType away, it is often necessary to cast between it and
    i1, which is used to represent bool-ness at lower levels.
  }];
  let arguments = (ins BoolOrI1Type:$operand);
  let results = (outs BoolOrI1Type:$result);
  let assemblyFormat = "$operand attr-dict `:` type(operands) `->` type(results)";
}

def Basicpy_BoolConstantOp : Basicpy_Op<"bool_constant", [
    ConstantLike, NoSideEffect]> {
  let summary = "A boolean constant";
  let description = [{
    A constant of type !basicpy.BoolType that can take either an i1 value
    of 0 (False) or 1 (True).
  }];
  let arguments = (ins I1Attr:$value);
  let results = (outs
    Basicpy_BoolType:$result
  );
  let assemblyFormat = "$value attr-dict";
  let hasFolder = 1;
}

def Basicpy_BytesConstantOp : Basicpy_Op<"bytes_constant", [
    ConstantLike, NoSideEffect]> {
  let summary = "Constant bytes value";
  let description = [{
    A bytes value of BytesType. The value is represented by a StringAttr.
  }];
  let arguments = (ins
    StrAttr:$value
  );
  let results = (outs
    Basicpy_BytesType:$result
  );
  let assemblyFormat = "$value attr-dict";
  let hasFolder = 1;
}

def Basicpy_ExecOp : Basicpy_Op<"exec", [
    SingleBlockImplicitTerminator<"ExecDiscardOp">]> {
  let summary = "Evaluates an expression being executed as a statement";
  let description = [{
    The result is discarded. Typically expressions are no-side-effect and can
    be re-ordered as needed. Embedding one in an exec op ensures that its
    placement in program order is preserved.
  }];
  let regions = (region SizedRegion<1>:$body);

  let skipDefaultBuilders = 1;
  let builders = [
    OpBuilder<"OpBuilder &builder, OperationState &result">,
  ];
  let extraClassDeclaration = [{
    OpBuilder getBodyBuilder() {
      Block* body = getBody(0);
      return OpBuilder::atBlockEnd(body);
    }
  }];
}

def Basicpy_ExecDiscardOp : Basicpy_Op<"exec_discard", [
    NoSideEffect, ReturnLike, Terminator]> {
  let summary = "Terminator for an exec block";
  let description = [{
    Discards results and terminates an exec block.
  }];
  let arguments = (ins Variadic<AnyType>:$operands);
  let assemblyFormat = "operands attr-dict `:` type(operands)";
}

def Basicpy_FuncTemplateCallOp : Basicpy_Op<"func_template_call", []> {
  let summary = "Calls a function template";
  let description = [{
    Most function calls start with this generic calling op, which binds
    symbolically to a func_template. At this level, there are very few
    semantics associated with the call, since, often, both types and the
    specific concrete callee cannot be determined.

    Per python calling conventions, all functions return one result, even if
    None or a tuple (which may be syntactically unpacked to multiple results).

    If specified, the `argNames` operand is right aligned to the list of
    positional `args`, representing arguments that are special or have been
    passed with a keyword. The following arg names are special:
      '*': Indicates that the argument is a positional argument pack (must be
           the first arg name, if present).
      '**': Indicates that the argument is a keyword argument pack (must be
           the last arg name, if present).
  }];

  let arguments = (ins
    FlatSymbolRefAttr:$callee,
    Variadic<AnyType>:$args,
    StrArrayAttr:$arg_names);
  let results = (outs AnyType:$result);
  let assemblyFormat = "$callee `(` $args `)` `kw` $arg_names attr-dict `:` functional-type($args, results)";
  let verifier = [{ return verifyBasicpyOp(*this); }];
  let skipDefaultBuilders = 1;
  let builders = [
    OpBuilder<"OpBuilder &builder, OperationState &result">,
  ];
}

def Basicpy_FuncTemplateOp : Basicpy_Op<"func_template", [
    IsolatedFromAbove,
    SingleBlockImplicitTerminator<"FuncTemplateTerminatorOp">,
    NativeOpTrait<"SymbolTable">,
    Symbol]> {
  let summary = "Group of multiple overload-resolved concrete functions";
  let description = [{
    The outer func_template op acts as a module that can contain named concrete
    functions that are interpreted as overloads. If the function signature is
    sufficient to disambiguate (i.e. with nothing more than arity and MLIR
    argument types), then this is all that is needed. However, in many cases,
    additional attributes will need to be specified to further constrain types.
    The first matching function signature is selected to satisfy a
    `func_template_call` op.

    TODO: Define this extended constraint matching.

    Instantiation
    -------------
    Once a concrete function is selected as being applicable to a given call,
    it will typically be instantiated as a standalone, unspecialized function
    in the calling module (as a peer to the func_template). This function
    will be uniquely identified by concating the outer func_template's symbol
    name, '$', and the concrete instance's symbol name.

    Note that the function may still be unspecialized (in that it contains
    UnknownType arguments/results), and type inference is expected to further
    specialize/inline/constrain it.

    Naming
    ------
    By convention, func_templates are named to avoid collision for various
    uses:
      - Global function templates: "__global$python.qualified.name"
      - Method names: "__method$method_name"
      - Attribute getter: "__getattr$attr_name"
      - Attribute setter: "__setattr$attr_name"

    As in user-level python, for functions that bind to an instance, the first
    argument must be a concrete type for the bound instance type. In this way,
    there is one `func_template` for every unique member name and the normal
    type constraints system is used to select the overload, just as if it was
    a normal function call. It is left to utility routines to merge libraries
    in a way that preserves this invariant.

    TODO: This needs to be fleshed out more as some additional rules about
    ordering and conflict resolution are likely needed to make this correct.

    Correlation with python runtime
    -------------------------------
    When extracting a program, it is typically necessary to create weak
    references to specific python functions and correlate them back to a named
    template defined here. Often times this can just be done lexically, but
    to avoid fragility, any func_template that correlates to a python
    runtime function will have an additional attribute `py_bind` that is an
    array of StringAttr qualified names to resolve and bind to in the python
    runtime. In cases of divergence, the symbol name of the template should
    be chosen just for uniqueness (not significance).

    The qualified name format for `py_bind` attribute is:
        package.name#local.qualified.name
  }];
  let arguments = (ins);
  let regions = (region SizedRegion<1>:$body);

  let verifier = [{ return verifyBasicpyOp(*this); }];

  let skipDefaultBuilders = 1;
  let builders = [
    OpBuilder<"OpBuilder &builder, OperationState &result">,
  ];
  let extraClassDeclaration = [{
    OpBuilder getBodyBuilder() {
      Block* body = getBody(0);
      return OpBuilder::atBlockEnd(body);
    }
  }];
}

def Basicpy_FuncTemplateTerminatorOp : Basicpy_Op<"func_template_terminator", [
    HasParent<"Basicpy::FuncTemplateOp">,
    Terminator]> {
  let summary = "Terminator pseudo-op for the FuncTemplateOp";

  let parser = ?;
  let printer = ?;
}

def Basicpy_SlotObjectMakeOp : Basicpy_Op<"slot_object_make", [
    NoSideEffect]> {
  let summary = "Creates an instance of a SlotObject type";
  let description = [{
    SlotObjects are typically instances of built-in classes that have a fixed
    number of slots. Unlike in standard python, the types of each slot are
    tracked.

    This op has a custom assembly form which can be used when valid that
    omits the operand types (since they are equal to the types in the returned
    slot object). Example:
      %0 = basicpy.singleton : !basicpy.NoneType
      %1 = basicpy.slot_object_make(%0) ->
          !basicpy.SlotObject<slice, !basicpy.NoneType>
  }];
  let arguments = (ins
    StrAttr:$className,
    // TODO: Tighter constraints on allowable types.
    Variadic<AnyType>:$slots
  );
  let results = (outs
    Basicpy_SlotObjectType:$result
  );
}

def Basicpy_SlotObjectGetOp : Basicpy_Op<"slot_object_get", [
    NoSideEffect]> {
  let summary = "Gets a slot from a slot object";
  let description = [{
    Gets a slot from a SlotObject.

    Example:
      %0 = basicpy.slot_object_make ...
      %1 = basicpy.slot_object_get %0[1] : !basicpy.SlotObject<...>
  }];
  let arguments = (ins
    Basicpy_SlotObjectType:$object,
    IndexAttr:$index
  );
  let results = (outs
    AnyType:$result
  );
}

def Basicpy_StrConstantOp : Basicpy_Op<"str_constant", [
    ConstantLike, NoSideEffect]> {
  let summary = "Constant string value";
  let description = [{
    A string value of StrType. The value is represented by a StringAttr
    that is UTF-8 encoded.
  }];
  let arguments = (ins
    StrAttr:$value
  );
  let results = (outs
    Basicpy_StrType:$result
  );
  let assemblyFormat = "$value attr-dict";
  let hasFolder = 1;
}

def Basicpy_SingletonOp : Basicpy_Op<"singleton", [
    ConstantLike, NoSideEffect]> {
  let summary = "Constant value for a singleton type";
  let description = [{
    Some types only have a single possible value, represented by the
    SingletonAttr. This op allows creating constants of these types.
  }];
  let arguments = (ins);
  let results = (outs
    Basicpy_SingletonType:$result
  );
  let assemblyFormat = "attr-dict `:` type($result)";
  let hasFolder = 1;
}

def Basicpy_ToBooleanOp : Basicpy_Op<"to_boolean", [NoSideEffect]> {
  let summary = "Evaluates an input to an i1 boolean value";
  let description = [{
    Applies the rules for interpreting a type as a boolean, returning an i1
    indicating the truthiness of the operand. Since the output of this op
    is intended to drive lower-level control flow, the i1 type is used (not
    the user level BoolType).
  }];
  let arguments = (ins AnyType:$operand);
  let results = (outs I1:$result);
  let assemblyFormat = "$operand attr-dict `:` type($operand)";
}

def Basicpy_UnknownCastOp : Basicpy_Op<"unknown_cast", [NoSideEffect]> {
  let summary = "Casts to and from the UnknownType";
  let arguments = (ins AnyType:$operand);
  let results = (outs AnyType:$result);
  let assemblyFormat = "operands attr-dict `:` type(operands) `->` type(results)";

  let hasCanonicalizer = 1;
}

#endif // NPCOMP_DIALECT_BASICPY_IR_BASICPY_OPS