gps/tools/local/circt/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td

//===-- TosaShapeOps.td - TOSA dialect utility operations --*- tablegen -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file defines shape operators for the TOSA dialect.
//
//===----------------------------------------------------------------------===//

#ifndef TOSA_SHAPE_OPS
#define TOSA_SHAPE_OPS

include "mlir/IR/OpBase.td"

include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/LoopLikeInterface.td"
include "mlir/Dialect/Tosa/IR/TosaInterfaces.td"

include "mlir/Dialect/Tosa/IR/TosaTypesBase.td"
include "mlir/Dialect/Tosa/IR/TosaOpBase.td"

// Op trait: operator has operands and results with TOSA shape type
def TosaShapeOperator : NativeOpTrait<"TosaShapeOperator"> {
  let cppNamespace = "mlir::OpTrait::tosa";
}

class Tosa_ShapeOp<string mnemonic, list<Trait> traits = []>
    : Tosa_Op<mnemonic, !listconcat(traits, [TosaShapeOperator])> {
  list<Availability> availability = [
    Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
    Extension<[Tosa_EXT_SHAPE]>,
  ];

  let assemblyFormat =
      "operands attr-dict `:` functional-type(operands, results)";
}

// op trait: shape operator has same ranks for operands and results
def TosaShapeOperatorWithSameRanks
    : NativeOpTrait<"TosaShapeOperatorWithSameRanks"> {
  let cppNamespace = "mlir::OpTrait::tosa";
}

class Tosa_ElementwiseShapeOp<string mnemonic, list<Trait> traits = []>
    : Tosa_ShapeOp<mnemonic,
                   !listconcat(traits, [TosaShapeOperatorWithSameRanks])> {
}


//===----------------------------------------------------------------------===//
// Operator: AddShape
//===----------------------------------------------------------------------===//
def Tosa_AddShapeOp : Tosa_ElementwiseShapeOp<"add_shape", [Pure]> {
  let summary = "Elementwise addition of shapes.";

  let description = [{
      Elementwise addition of input1 and input2. Size of shapes must match.
  }];

  let arguments = (ins
    Tosa_Shape:$input1,
    Tosa_Shape:$input2
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: ConcatShape
//===----------------------------------------------------------------------===//
def Tosa_ConcatShapeOp : Tosa_ShapeOp<"concat_shape", [Pure]> {
  let summary = "Concatenates a list of shapes.";

  let description = [{
    Concatenates a list of shapes into a new shape with length equal to the sum
    of the lengths of the inputs.
  }];

  let arguments = (ins
    Variadic<Tosa_Shape>:$input
  );

  let results = (outs Tosa_Shape:$output);

  let hasVerifier = 1;

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: ConstShape
//===----------------------------------------------------------------------===//
def Tosa_ConstShapeOp : Tosa_ShapeOp<"const_shape", [ConstantLike, Pure]> {
  let summary = "Constant Shape operator.";

  let description = [{
    A node containing a constant shape.

    Example:

    ```mlir
    // Generic form
    %out = "tosa.const_shape"() {values = dense<0> : tensor<4xindex>} : () -> !tosa.shape<4>
    ```
  }];

  let arguments = (ins IndexElementsAttr:$values);

  let results = (outs Tosa_Shape:$output);

  list<Availability> availability = [
    Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
    Extension<[]>,
  ];

  let hasVerifier = 1;
  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: Dim
//===----------------------------------------------------------------------===//
def Tosa_DimOp : Tosa_ShapeOp<"dim", [Pure]> {
  let summary = "Extract size of dimension from input tensor.";

  let description = [{
    Returns a length 1 shape_t of the size of the input tensor for the given axis.
  }];

  let arguments = (ins
    Tosa_TensorAtLeast1D:$input1,
    I32Attr:$axis
  );

  let results = (outs Tosa_Shape:$output);

  list<Availability> availability = [
    Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
    Extension<[Tosa_EXT_SHAPE, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16, Tosa_EXT_MXFP, Tosa_EXT_INT64]>,
  ];

  let hasVerifier = 1;

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: DivCeilShape
//===----------------------------------------------------------------------===//
def Tosa_DivCeilShapeOp : Tosa_ElementwiseShapeOp<"div_ceil_shape", [Pure]> {
  let summary = "Elementwise ceiling divide of shapes.";

  let description = [{
      Elementwise divide of input1 by input2. The result of the divide is rounded up.
  }];

  let arguments = (ins
    Tosa_Shape:$input1,
    Tosa_Shape:$input2
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: DivFloorShape
//===----------------------------------------------------------------------===//
def Tosa_DivFloorShapeOp : Tosa_ElementwiseShapeOp<"div_floor_shape", [Pure]> {
  let summary = "Elementwise floor divide of shapes.";

  let description = [{
      Elementwise integer divide of input1 by input2. The result of the divide is rounded down.
  }];

  let arguments = (ins
    Tosa_Shape:$input1,
    Tosa_Shape:$input2
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: Exp2Shape
//===----------------------------------------------------------------------===//
def Tosa_Exp2ShapeOp : Tosa_ElementwiseShapeOp<"exp2_shape", [Pure]> {
  let summary = "Elementwise base-2 exponential of shapes.";

  let description = [{
      Computation of raising two to the power of each element in input.
  }];

  let arguments = (ins
    Tosa_Shape:$input
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: Log2CeilShape
//===----------------------------------------------------------------------===//
def Tosa_Log2CeilShapeOp : Tosa_ElementwiseShapeOp<"log2_ceil_shape", [Pure]> {
  let summary = "Elementwise ceil base-2 logarithm of shapes.";

  let description = [{
      Computation of the base two logarithm of each element in input. Result is rounded up.
  }];

  let arguments = (ins
    Tosa_Shape:$input
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: Log2FloorShape
//===----------------------------------------------------------------------===//
def Tosa_Log2FloorShapeOp : Tosa_ElementwiseShapeOp<"log2_floor_shape", [Pure]> {
  let summary = "Elementwise floor base-2 logarithm of shapes.";

  let description = [{
      Computation of the base two logarithm of each element in input. Result is rounded down.
  }];

  let arguments = (ins
    Tosa_Shape:$input
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: MaxShape
//===----------------------------------------------------------------------===//
def Tosa_MaxShapeOp : Tosa_ElementwiseShapeOp<"max_shape", [Pure]> {
  let summary = "Elementwise maximum of shapes.";

  let description = [{
      Elementwise maximum of input1 and input2.
  }];

  let arguments = (ins
    Tosa_Shape:$input1,
    Tosa_Shape:$input2
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: MinShape
//===----------------------------------------------------------------------===//
def Tosa_MinShapeOp : Tosa_ElementwiseShapeOp<"min_shape", [Pure]> {
  let summary = "Elementwise minimum of shapes.";

  let description = [{
      Elementwise minimum of input1 and input2.
  }];

  let arguments = (ins
    Tosa_Shape:$input1,
    Tosa_Shape:$input2
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: ModShape
//===----------------------------------------------------------------------===//
def Tosa_ModShapeOp : Tosa_ElementwiseShapeOp<"mod_shape", [Pure]> {
  let summary = "Elementwise modulo of shapes.";

  let description = [{
    Elementwise modulo of input1 divided by input2.
  }];

  let arguments = (ins
    Tosa_Shape:$input1,
    Tosa_Shape:$input2
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: MulShape
//===----------------------------------------------------------------------===//
def Tosa_MulShapeOp : Tosa_ElementwiseShapeOp<"mul_shape", [Pure]> {
  let summary = "Elementwise multiplication of shapes.";

  let description = [{
      Elementwise multiplication of input1 and input2.
  }];

  let arguments = (ins
    Tosa_Shape:$input1,
    Tosa_Shape:$input2
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: SliceShape
//===----------------------------------------------------------------------===//
def Tosa_SliceShapeOp : Tosa_ShapeOp<"slice_shape", [Pure]> {
  let summary = "Extract slice of a shape.";

  let description = [{
    Extract a shape of size from input.
  }];

  let arguments = (ins
    Tosa_Shape:$input,
    Tosa_ScalarInt32Tensor:$start,
    Tosa_ScalarInt32Tensor:$size
  );

  let results = (outs Tosa_Shape:$output);

  let hasVerifier = 1;

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: SubShape
//===----------------------------------------------------------------------===//
def Tosa_SubShapeOp : Tosa_ElementwiseShapeOp<"sub_shape", [Pure]> {
  let summary = "Elementwise subtraction of shapes.";

  let description = [{
      Elementwise subtraction of input1 and input2. Size of shapes must match.
  }];

  let arguments = (ins
    Tosa_Shape:$input1,
    Tosa_Shape:$input2
  );

  let results = (outs Tosa_Shape:$output);

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Operator: AssertEqualShape
//===----------------------------------------------------------------------===//
def Tosa_AssertEqualShapeOp
    : Tosa_ShapeOp<"assert_equal_shape", [TosaShapeOperatorWithSameRanks]> {
  let summary = "Verify two shapes are equal.";

  let description = [{
      Verify input1 and input2 are equal. If allow_broadcast is set, shapes which
      are broadcast compatible are allowed.
  }];

  let arguments = (ins Tosa_Shape:$input1, Tosa_Shape:$input2,
      BoolAttr:$allow_broadcast);

  let results = (outs);
}

#endif // TOSA_SHAPE_OPS