torch-mlir/include/npcomp/Dialect/Basicpy/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_BASICPY_OPS
#define NPCOMP_DIALECT_BASICPY_BASICPY_OPS

include "BasicpyDialect.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_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:$lhs,
    AnyType:$rhs,
    BinaryOperationAttr:$operation
  );
  let results = (outs
    AnyType:$result
  );
  let assemblyFormat = "$operation operands 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";
}

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";
}

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_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";
}

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)";
}

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

#endif // NPCOMP_DIALECT_BASICPY_BASICPY_OPS