//===- Shape.td - Shape operations definition --------------*- 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 is the operation definition file for Shape dialect operations.
//
//===----------------------------------------------------------------------===//

#ifndef SHAPE_OPS
#define SHAPE_OPS

include "mlir/Dialect/Shape/IR/ShapeBase.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/IR/OpAsmInterface.td"
include "mlir/IR/SymbolInterfaces.td"

//===----------------------------------------------------------------------===//
// Shape op definitions
//===----------------------------------------------------------------------===//

// Base class for the operation in this dialect
class Shape_Op<string mnemonic, list<OpTrait> traits = []> :
    Op<ShapeDialect, mnemonic, traits>;

def Shape_AddOp : Shape_Op<"add",
    [Commutative, NoSideEffect,
     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Addition of sizes and indices";
  let description = [{
    Adds two sizes or indices. If either operand is an error it will be
    propagated to the result. The operands can be of type `size` or `index`. If
    at least one of the operands can hold an error, i.e. if it is of type `size`,
    the result must be of type `size`. If error propagation is not possible
    because both operands are of type `index` then the result may be of type
    `size` or `index`.
  }];

  let arguments = (ins Shape_SizeOrIndexType:$lhs, Shape_SizeOrIndexType:$rhs);
  let results = (outs Shape_SizeOrIndexType:$result);

  let assemblyFormat = [{
    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
  }];

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

  let extraClassDeclaration = [{
    // Returns when two result types are compatible for this op; method used by
    // InferTypeOpInterface
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];

  let hasFolder = 1;
}

def Shape_BroadcastOp : Shape_Op<"broadcast", [Commutative, NoSideEffect]> {
  let summary = "Returns the broadcasted output shape of two or more inputs";
  let description = [{
    Returns the broadcasted shape for input shapes or extent tensors. The rest
    of this description is simplified for the 2 input case but can be extended
    to more inputs. Both operands can be of type `shape.shape` or
    `tensor<?xindex>`. The result is of type `shape.shape` and, if both
    operands are tensors, may be of type `tensor<?xindex>`.

    If the two operand shapes are of different rank the smaller one is padded
    with 1's from the left. The resulting broadcasted shape is then defined as

        result[i] = lhs[i] if lhs[i] == rhs[i]
                  = lhs[i] if rhs[i] == 1
                  = rhs[i] if lhs[i] == 1.

    In case the resulting shape is undefined, i.e. if corresponding extents are
    different from each other but none is 1, the result is an error shape.
    Likewise error values are propagated if any of the operands holds an error
    value. If the result type is an extent tensor (and can therefore not hold
    the error value) the behavior may be undefined. The optional string
    attribute can be used to describe the error case.
  }];

  let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$shapes,
                       OptionalAttr<StrAttr>:$error);
  let results = (outs Shape_ShapeOrExtentTensorType:$result);

  let builders = [OpBuilder<(ins "Value":$shape)>];

  let assemblyFormat = [{
    $shapes attr-dict `:` type($shapes) `->` type($result)
  }];

  let builders = [OpBuilder<(ins "::mlir::Type":$result,
                                "::mlir::Value":$lhs, "::mlir::Value":$rhs,
                                "/*optional*/ ::mlir::StringAttr":$error), [{
      build($_builder, $_state, result, ::llvm::makeArrayRef({lhs, rhs}), error);
    }]>
  ];

  let hasFolder = 1;
  let hasCanonicalizer = 1;
  let verifier = [{ return ::verify(*this); }];
}

def Shape_ConstShapeOp : Shape_Op<"const_shape",
    [ConstantLike, NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Creates a constant shape or extent tensor";
  let description = [{
    Creates a constant shape or extent tensor. The individual extents are given
    as the `shape` attribute. The number of these values equals the shape's
    rank.

    ```mlir
    %0 = shape.const_shape [] : !shape.shape
    %1 = shape.const_shape [1, 2, 3] : !shape.shape
    %2 = shape.const_shape [4, 5, 6] : tensor<3xindex>
    ```
  }];
  let arguments = (ins IndexElementsAttr:$shape);
  let results = (outs Shape_ShapeOrExtentTensorType:$result);

  // TODO: Move this to main so that all shape ops implement these.
  let printer = [{ return ::print(p, *this); }];
  let parser = [{ return ::parse$cppClass(parser, result); }];
  let hasFolder = 1;
  let hasCanonicalizer = 1;

  let extraClassDeclaration = [{
    // InferTypeOpInterface:
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];
}

def Shape_ConstSizeOp : Shape_Op<"const_size", [
    ConstantLike,
    NoSideEffect,
    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>
  ]> {
  let summary = "Creates a constant of type `shape.size`";
  let description = [{
    Creates a `shape.size` type representing the constant size given by `value`.

    ```mlir
    %x = shape.const_size 10
    ```
  }];

  let arguments = (ins IndexAttr:$value);
  let results = (outs Shape_SizeType:$result);

  let builders = [OpBuilder<(ins "int64_t":$value)>];

  let assemblyFormat = "$value attr-dict";
  let hasFolder = 1;
}

def Shape_DivOp : Shape_Op<"div", [NoSideEffect,
                           DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Division of sizes and indices";
  let description = [{
    Divides two sizes or indices. If either operand is an error it will be
    propagated to the result. The operands can be of type `size` or `index`.
    If at least one of the operands can hold an error, i.e. if it is of type
    `size`, the result must be of type `size`. If error propagation is not
    possible because both operands are of type `index` then the result may be
    of type  `size` or `index`. If both operands and result are of type `index`,
    their runtime values could be negative. The result is rounded toward
    negative infinity, i.e. floor(lhs / rhs), such that

        div(lhs, rhs) * rhs + mod(lhs, rhs) = lhs

    always holds. If any of the values is of type `size`, the behavior for
    negative value is undefined.
  }];

  let arguments = (ins Shape_SizeOrIndexType:$lhs,
                       Shape_SizeOrIndexType:$rhs);
  let results = (outs Shape_SizeOrIndexType:$result);

  let assemblyFormat = [{
    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
  }];

  let verifier = [{ return ::verifySizeOrIndexOp(*this); }];
  let hasFolder = 1;

  let extraClassDeclaration = [{
    // Returns when two result types are compatible for this op; method used by
    // InferTypeOpInterface
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];
}

def Shape_ShapeEqOp : Shape_Op<"shape_eq",
    [NoSideEffect, Commutative, InferTypeOpInterface]> {
  let summary = "Returns whether the input shapes or extent tensors are equal";
  let description = [{
    Takes one or more shape or extent tensor operands and determines whether
    they are equal. When extent tensors are compared to shapes they are regarded
    as their equivalent non-error shapes. Error shapes can be tested for
    equality like any other shape value, meaning that the error value is equal
    to itself.
  }];

  let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$shapes);
  let results = (outs I1:$result);

  // Convenience builder alias for the binary version.
  let builders = [
  OpBuilder<(ins "::mlir::Value":$lhs, "::mlir::Value":$rhs),
    [{ build($_builder, $_state, ::llvm::makeArrayRef({lhs, rhs})); }]>,
  ];
  let extraClassDeclaration = [{
    // TODO: This should really be automatic. Figure out how to not need this defined.
    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
    ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
    ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
    ::llvm::SmallVectorImpl<::mlir::Type>&inferredReturnTypes) {
      inferredReturnTypes.push_back(::mlir::IntegerType::get(context,
                                                             /*width=*/1));
      return success();
    };
  }];

  let assemblyFormat = "$shapes attr-dict `:` type($shapes)";
  let hasFolder = 1;
}

def Shape_FromExtentsOp : Shape_Op<"from_extents", [NoSideEffect]> {
  let summary = "Creates a shape from extents";
  let description = [{
    Creates a shape from multiple SSA values representing the extents of
    the shape.

    ```mlir
    // Rank 2 shape.
    %s0 = shape.from_extents %a, %b
    // Rank 0 shape.
    %s1 = shape.from_extents
    ```
  }];
  let arguments = (ins Variadic<Shape_SizeOrIndexType>:$extents);
  let results = (outs Shape_ShapeType:$shape);

  let assemblyFormat = "$extents attr-dict `:` type($extents)";

  let hasFolder = 1;
}

def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", [NoSideEffect]> {
  let summary = "Creates a shape from a tensor of extents";
  let description = [{
    Creates a shape from a 1D integral tensor of extents. The rank of the
    resulting shape equals the number of elements in the tensor, and the
    extents match the values of the elements.
  }];

  let arguments = (ins 1DTensorOf<[Index]>:$input);
  let results = (outs Shape_ShapeType:$result);

  let assemblyFormat = "$input attr-dict `:` type($input)";
}

def Shape_IsBroadcastableOp : Shape_Op<"is_broadcastable",
                                       [Commutative, InferTypeOpInterface]> {
  let summary = "Determines if 2+ shapes can be successfully broadcasted";
  let description = [{
    Given multiple input shapes or extent tensors, return a predicate specifying
    if they are broadcastable. This broadcastable follows the same logic as what
    shape.broadcast documents.

    Concretely, shape.is_broadcastable returning true implies that
    shape.broadcast will not give an error, and shape.cstr_broadcastable will
    not result in an assertion failure. Similarly, false implies an error or
    assertion failure.

    Example:
    ```mlir
    %true = shape.is_broadcastable [2,2], [3,1,2]
    %false = shape.is_broadcastable [2,2], [3,2]
    ```
  }];

  let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$shapes);
  let results = (outs I1:$result);

  let builders = [
  OpBuilder<(ins "::mlir::Value":$lhs, "::mlir::Value":$rhs),
    [{ build($_builder, $_state, ::llvm::makeArrayRef({lhs, rhs})); }]>,
  ];
  let extraClassDeclaration = [{
    // TODO: This should really be automatic. Figure out how to not need this defined.
    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
    ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
    ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
    ::llvm::SmallVectorImpl<::mlir::Type>&inferredReturnTypes) {
      inferredReturnTypes.push_back(::mlir::IntegerType::get(context,
                                                             /*width=*/1));
      return success();
    };
  }];

  let hasFolder = 1;
  let hasCanonicalizer = 1;

  let assemblyFormat = "$shapes attr-dict `:` type($shapes)";
}

def Shape_RankOp : Shape_Op<"rank",
    [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Gets the rank of a shape";
  let description = [{
    Returns the rank of the shape or extent tensor, i.e. the number of extents.
  }];

  let arguments = (ins Shape_ShapeOrExtentTensorType:$shape);
  let results = (outs Shape_SizeOrIndexType:$rank);

  let assemblyFormat = "$shape attr-dict `:` type($shape) `->` type($rank)";

  let hasFolder = 1;
  let hasCanonicalizer = 1;
  let verifier = [{ return ::verifySizeOrIndexOp(*this); }];

  let extraClassDeclaration = [{
    // Returns when two result types are compatible for this op; method used by
    // InferTypeOpInterface
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];
}

def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
  let summary = "Creates a dimension tensor from a shape";
  let description = [{
    Converts a shape to a 1D integral tensor of extents. The number of elements
    in the tensor equals the rank of the shape, and the elements equal the
    extents of the shape.

    If the shape represents an error, this op's behavior is undefined.
  }];

  let arguments = (ins Shape_ShapeOrExtentTensorType:$input);
  let results = (outs IndexTensor:$result);

  let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";

  let hasFolder = 1;
}

def Shape_GetExtentOp : Shape_Op<"get_extent",
    [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Gets the specified extent from a shape or extent tensor";
  let description = [{
    Gets the extent indexed by `dim` from the `shape` operand. If the shape is
    an error then it returns an invalid size.
  }];
  let arguments = (ins Shape_ShapeOrExtentTensorType:$shape,
                       Shape_SizeOrIndexType:$dim);
  let results = (outs Shape_SizeOrIndexType:$extent);
  let assemblyFormat = "$shape `,` $dim attr-dict `:` type($shape) `,` type($dim) `->` "
                       "type($extent)";

  let builders = [
    // Builder that allows passing a constant dimension as a simple integer.
    OpBuilder<(ins "Value":$shape, "int64_t":$dim)>
  ];

  let extraClassDeclaration = [{
    /// Get the `dim` value as integer if it is constant.
    Optional<int64_t> getConstantDim();
    /// Returns when two result types are compatible for this op; method used by
    /// InferTypeOpInterface
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];

  let hasFolder = 1;
  let verifier = [{ return ::verifySizeOrIndexOp(*this); }];
}

def Shape_IndexToSizeOp : Shape_Op<"index_to_size", [NoSideEffect]> {
  let summary = "Converts a standard index to a shape size";
  let description = [{
    Converts a standard index to a `shape.size`. This operation and its
    inverse, `size_to_index`, facilitate index conversion between the standard
    and the shape dialect.

    The behavior is undefined for negative indices.
  }];

  let arguments = (ins Index:$arg);
  let results = (outs Shape_SizeType:$result);

  let assemblyFormat = "$arg attr-dict";

  let hasFolder = 1;
  let hasCanonicalizer = 1;
}

def Shape_MaxOp : Shape_Op<"max",
    [Commutative, NoSideEffect,
     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Elementwise maximum";
  let description = [{
    Computes the elementwise maximum of two sizes or shapes with equal ranks.
    If either operand is an error, then an error will be propagated to the
    result. If the input types mismatch or the ranks do not match, then the
    result is an error.
  }];

  let arguments = (ins Shape_ShapeOrSizeType:$lhs, Shape_ShapeOrSizeType:$rhs);
  let results = (outs Shape_ShapeOrSizeType:$result);

  let assemblyFormat = [{
    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
  }];

  let hasFolder = 1;

  let extraClassDeclaration = [{
    // Returns when two result types are compatible for this op; method used by
    // InferTypeOpInterface
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];
}

def Shape_MeetOp : Shape_Op<"meet",
    [Commutative, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Returns the least general shape.shape of its operands";
  let description = [{
    An operation that computes the least general shape of input operands.
    This effectively asserts that corresponding static dimensions are equal.
    The behavior is to match each element of the `shape.shape` and propagate the
    most restrictive information, returning an invalid shape if there are
    contradictory requirements. E.g., using pseudo code

    ```
    shape.meet([*], [*]) -> [*]
    shape.meet([*], [1, ?]) -> [1, ?]
    shape.meet([1, 2], [1, ?]) -> [1, 2]
    shape.meet([*], [1, 2]) -> [1, 2]
    shape.meet([], []) -> []
    shape.meet([], [*]) -> []
    shape.meet([], [?, ?]) -> [invalid]
    shape.meet([1, ?], [2, ?, ?]) -> [invalid]
    ```

    `shape.meet` also allows specifying an optional error string, that may be
    used to return an error to the user upon mismatch of dimensions.

    ```mlir
    %c = shape.meet %a, %b, error="<reason>" : !shape.shape, !shape.shape -> !shape.shape
    ```
  }];

  let arguments = (ins Shape_ShapeOrSizeType:$arg0, Shape_ShapeOrSizeType:$arg1,
                   OptionalAttr<StrAttr>:$error);
  let results = (outs Shape_ShapeOrSizeType:$result);

  let assemblyFormat = [{
    $arg0 `,` $arg1 (`,` `error` `=` $error^)? attr-dict `:`
      type($arg0) `,` type($arg1) `->` type($result)
  }];

  let extraClassDeclaration = [{
    // Returns when two result types are compatible for this op; method used by
    // InferTypeOpInterface
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];
}

def Shape_MinOp : Shape_Op<"min",
    [Commutative, NoSideEffect,
     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Elementwise minimum";
  let description = [{
    Computes the elementwise minimum of two sizes or shapes with equal ranks.
    If either operand is an error, then an error will be propagated to the
    result. If the input types mismatch or the ranks do not match, then the
    result is an error.
  }];

  let arguments = (ins Shape_ShapeOrSizeType:$lhs, Shape_ShapeOrSizeType:$rhs);
  let results = (outs Shape_ShapeOrSizeType:$result);

  let assemblyFormat = [{
    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
  }];

  let hasFolder = 1;

  let extraClassDeclaration = [{
    // Returns when two result types are compatible for this op; method used by
    // InferTypeOpInterface
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];
}

def Shape_MulOp : Shape_Op<"mul",
    [Commutative, NoSideEffect,
     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Multiplication of sizes and indices";
  let description = [{
    Multiplies two sizes or indices. If either operand is an error it will be
    propagated to the result. The operands can be of type `size` or `index`. If
    at least one of the operands can hold an error, i.e. if it is of type `size`,
    the result must be of type `size`. If error propagation is not possible
    because both operands are of type `index` then the result may be of type
    `size` or `index`.
  }];

  let arguments = (ins Shape_SizeOrIndexType:$lhs, Shape_SizeOrIndexType:$rhs);
  let results = (outs Shape_SizeOrIndexType:$result);

  let assemblyFormat = [{
    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)
  }];

  let verifier = [{ return ::verifySizeOrIndexOp(*this); }];
  let hasFolder = 1;

  let extraClassDeclaration = [{
    // Returns when two result types are compatible for this op; method used by
    // InferTypeOpInterface
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];
}

def Shape_NumElementsOp : Shape_Op<"num_elements",
    [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Returns the number of elements for a given shape";
  let description = [{
    Returns the number of elements for a given shape which is the product of its
    extents. If the argument is of type `shape` then the result will be of type
    `size` and potential errors will be propagated. Otherwise, if the argument
    is and extent tensor `tensor<?xindex>` then the result will be of type
    `index`.
  }];

  let arguments = (ins Shape_ShapeOrExtentTensorType:$shape);
  let results = (outs Shape_SizeOrIndexType:$result);

  let assemblyFormat = "$shape attr-dict `:` type($shape) `->` type($result)";

  let hasFolder = 1;
  let verifier = [{ return ::verifySizeOrIndexOp(*this); }];
  let extraClassDeclaration = [{
    // Returns when two result types are compatible for this op; method used by
    // InferTypeOpInterface
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];
}

def Shape_ReduceOp : Shape_Op<"reduce",
    [SingleBlockImplicitTerminator<"YieldOp">]> {
  let summary = "Returns an expression reduced over a shape or extent tensor";
  let description = [{
    An operation that takes as input a shape or extent tensor, and a number of
    initial values. This operation has a region that is applied repeatedly for
    every extent of the input. Starting with the initial values, the individual
    extents are then aggregated as defined by the associated region.

    Conceptually this op performs the following reduction:

    ```
    res[] = init;
    for (int i = 0, i < shape.rank(); i++) {
      res = reduce(i, shape[i], res[0], ..., res[n]);
    }
    ```

    Where `reduce` represents the region attached and the result of the reduce
    op is the last computed output of the reduce region. As an example, the
    number of elements can be computed as follows:

    ```mlir
    func @reduce(%shape : !shape.shape, %init : !shape.size) -> !shape.size {
      %num_elements = shape.reduce(%shape, %init) -> !shape.size  {
        ^bb0(%index: index, %dim: !shape.size, %acc: !shape.size):
          %updated_acc = "shape.mul"(%acc, %dim) :
            (!shape.size, !shape.size) -> !shape.size
          shape.yield %updated_acc : !shape.size
      }
      return %num_elements : !shape.size
    }
    ```
  }];

  let arguments = (ins Shape_ShapeOrExtentTensorType:$shape,
                       Variadic<AnyType>:$initVals);
  let results = (outs Variadic<AnyType>:$result);
  let regions = (region SizedRegion<1>:$region);

  let builders = [OpBuilder<(ins "Value":$shape, "ValueRange":$initVals)>];

  let verifier = [{ return ::verify(*this); }];
  let printer = [{ return ::print(p, *this); }];
  let parser = [{ return ::parse$cppClass(parser, result); }];
}

def Shape_ShapeOfOp : Shape_Op<"shape_of",
    [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
  let summary = "Returns shape of a value or shaped type operand";

  let description = [{
    The operation takes a value or a shaped operand as an argument and it
    returns a shape or extent tensor.
  }];

  let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$arg);
  let results = (outs Shape_ShapeOrExtentTensorType:$result);

  let assemblyFormat = "$arg attr-dict `:` type($arg) `->` type($result)";

  let verifier = [{ return ::verifyShapeOrExtentTensorOp(*this); }];
  let hasCanonicalizer = 1;
  let hasFolder = 1;

  let extraClassDeclaration = [{
    // Returns when two result types are compatible for this op; method used by
    // InferTypeOpInterface
    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
  }];
}

def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [NoSideEffect]> {
  let summary = "Casts between index types of the shape and standard dialect";
  let description = [{
    Converts a `shape.size` to a standard index. This operation and its
    inverse, `index_to_size`, facilitate index conversion between the standard
    and the shape dialect. The behavior is undefined for unknown and invalid
    arguments.
  }];

  let arguments = (ins Shape_SizeOrIndexType:$arg);
  let results = (outs Index:$result);

  let assemblyFormat = "$arg attr-dict `:` type($arg)";

  let hasFolder = 1;
  let hasCanonicalizer = 1;
}

def Shape_ValueAsShapeOp : Shape_Op<"value_as_shape", [NoSideEffect]> {
  let summary = "Returns value as a shape";

  let description = [{
    The operations takes a ValueShape and returns a Shape corresponding to the value.
    If the input value cannot be shape (e.g., not a 1D tensor of integral value
    representing sizes) then this propagages the error shape. E.g.,

    ```mlir
    // The following
    %0 = arith.constant dense<[1,2]> : tensor<2xi32>
    %shape = shape.value_as_shape %0 : tensor<2xi32> -> !shape.shape
    // is equivalent to
    %shape' = shape.const_shape [1, 2] : !shape.shape
    ```

    This operation is the compliment of `shape_of` wrt ValueShape values.
  }];

  let arguments = (ins AnyTypeOf<[1DTensorOf<[AnyInteger, Index]>, Shape_ValueShapeType]>:$arg);
  let results = (outs Shape_ShapeOrExtentTensorType:$result);

  let assemblyFormat = "$arg attr-dict `:` type($arg) `->` type($result)";
}

def Shape_WithOp : Shape_Op<"with_shape", [NoSideEffect]> {
  let summary = "Returns ValueShape with given shape";
  let description = [{
    Returns ValueShape with the shape updated to match the shape operand. That
    is a new ValueShape tuple is created with value equal to `operand`'s
    value and shape equal to `shape`. If the ValueShape and given `shape` are
    non-conformant, then the returned ValueShape will represent an error of
    this mismatch. Similarly if either inputs are in an error state, then an
    error is propagated.

    Usage:
      %0 = shape.with_shape %1, %2 : tensor<...>, !shape.shape

    This is used, for example, where one combines shape function calculations
    and/or call one shape function from another. E.g.,

    ```mlir
    func @shape_foobah(%a: !shape.value_shape,
                       %b: !shape.value_shape,
                       %c: !shape.value_shape) -> !shape.shape {
      %0 = call @shape_foo(%a, %b) :
        (!shape.value_shape, !shape.value_shape) -> !shape.shape
      %1 = shape.with_shape %b, %0 : !shape.value_shape, !shape.shape
      %2 = call @shape_bah(%c, %1) :
        (!shape.value_shape, !shape.value_shape) -> !shape.shape
      return %2 : !shape.shape
    }
    ```

    This op need not be a refinement of the shape. In non-error cases the input
    ValueShape's value and shape are conformant and so too for the output, but
    the result may be less specified than `operand`'s shape as `shape` is
    merely used to construct the new ValueShape. If join behavior is desired
    then a join op should be used.
  }];

  let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$operand,
                       Shape_ShapeType:$shape);
  let results = (outs Shape_ValueShapeType:$result);

  let assemblyFormat = "operands attr-dict `:` type($operand) `,` type($shape)";
}

def Shape_YieldOp : Shape_Op<"yield",
    [HasParent<"ReduceOp, FunctionLibraryOp">,
     NoSideEffect,
     ReturnLike,
     Terminator]> {
  let summary = "Returns the value to parent op";

  let arguments = (ins Variadic<AnyType>:$operands);

  let builders = [OpBuilder<(ins),
    [{ build($_builder, $_state, llvm::None); }]>
  ];

  let verifier = [{ return ::verify(*this); }];
  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
}

// TODO: Add Ops: if_static, if_ranked

// For testing usage.
def Shape_DebugPrintOp : Shape_Op<"debug_print", []> {
  let summary = "Prints the input shape or size";
  let description = [{
    Prints the input dim or shape and passes through input.

    Note: This is intended for testing and debugging only.
  }];

  let arguments = (ins Shape_ShapeOrSizeType:$input);
  let results =  (outs Shape_ShapeOrSizeType:$output);
}

def Shape_SplitAtOp : Shape_Op<"split_at", [NoSideEffect]> {
  let summary = "Splits a shape at a given index";
  let description = [{
    Splits a shape at a given dimension `index`, returning two shapes.
    If `index` is negative, it is treated as indexing from the back of the
    shape. This negative-handling behavior is important when handling unranked
    shapes, where the positive index is not necessarily knowable due to a
    dynamic number of leading dimensions. If the result is in extent tensor form
    out of bounds indices result in undefined behavior.

    Examples:
    - split_at([4,5,6], index=0) -> [], [4,5,6]
    - split_at([4,5,6], index=1) -> [4], [5,6]
    - split_at([4,5,6], index=2) -> [4,5], [6]
    - split_at([4,5,6], index=3) -> [4,5,6], []
    - split_at([4,5,6], index=4) -> error
    - split_at([4,5,6], index=-1) -> [4,5], [6]
    - split_at([4,5,6], index=-2) -> [4], [5,6]
    - split_at([4,5,6], index=-3) -> [], [4,5,6]
    - split_at([4,5,6], index=-4) -> error

    Requires:
    - `index` is in the range [-rank(operand),rank(operand)]
  }];

  let arguments = (ins Shape_ShapeOrExtentTensorType:$operand,
                       Shape_SizeOrIndexType:$index);
  let results = (outs Shape_ShapeOrExtentTensorType:$head,
                      Shape_ShapeOrExtentTensorType:$tail);
  let hasFolder = 1;
}

def Shape_ConcatOp : Shape_Op<"concat", [NoSideEffect]> {
  let summary = "Concatenates two shapes";
  let description = [{
    Creates a shape whose dimensions consist of first the dimensions from `lhs`
    followed by the dimensions of `rhs`.

    Example:
    concat([2,3], [4,5]) -> [2,3,4,5]
    concat([], []) -> []
    concat([], [4,5,6]) -> [4,5,6]
  }];

  let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
  let results = (outs Shape_ShapeType:$result);

  let assemblyFormat = "$lhs `,` $rhs attr-dict";
  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Shape constraint related ops.
//===----------------------------------------------------------------------===//

// TODO: Move the code below and witnesses to a different file.
def Shape_AnyOp : Shape_Op<"any", [Commutative,
                                   NoSideEffect]> {
  let summary = "Return any combination of the input shapes";
  let description = [{
    This operation takes multiple input shapes or extent tensors and returns
    some combination of their dimensions. This can be best seen with examples
    below.

    The result is undefined, but still side-effect free, in cases where the
    inputs have differing ranks or differ in extents of shared dimensions.

    Example:
    ```mlir
    %s0 = shape.any [2,?], [?,3] // [2,3]
    %s1 = shape.any [?,?], [1,2] // [1,2]
    ```
  }];

  let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$inputs);
  let results = (outs Shape_ShapeOrExtentTensorType:$result);

  let assemblyFormat = "$inputs attr-dict `:` type($inputs) `->` type($result)";

  let hasFolder = 1;
}

def Shape_AssumingAllOp : Shape_Op<"assuming_all", [Commutative, NoSideEffect]> {
  let summary = "Return a logical AND of all witnesses";
  let description = [{
    Used to simplify constraints as any single failing precondition is enough
    to prevent execution.

    "assuming" operations represent an execution order restriction to the
    compiler, information for dependent code to rely on (by assuming), and
    nothing else. They should not exist after a program is fully lowered and
    ready to execute.

    Example:
    ```mlir
    %w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Passing
    %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
    %w2 = shape.cstr_eq [1,2], [1,2], [1,2] // Passing
    %wf = shape.assuming_all %w0, %w1 // Failure
    %wt = shape.assuming_all %w0, %w2 // Passing
    ```
  }];

  let arguments = (ins Variadic<Shape_WitnessType>:$inputs);
  let results = (outs Shape_WitnessType:$result);

  // Only needed while tablegen is unable to generate this for ops with variadic
  // arguments.
  let builders = [
    OpBuilder<(ins "ValueRange":$inputs)>,
  ];

  let assemblyFormat = "$inputs attr-dict";

  let hasFolder = 1;
  let hasCanonicalizer = 1;

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

def Shape_AssumingOp : Shape_Op<"assuming", [
    SingleBlockImplicitTerminator<"AssumingYieldOp">,
    DeclareOpInterfaceMethods<RegionBranchOpInterface>,
    RecursiveSideEffects]> {
  let summary = "Execute the region";
  let description = [{
    Executes the region assuming all witnesses are true.

    "assuming" operations represent an execution order restriction to the
    compiler, information for dependent code to rely on (by assuming), and
    nothing else. They should not exist after a program is fully lowered and
    ready to execute.
  }];
  let arguments = (ins Shape_WitnessType:$witness);
  let regions = (region SizedRegion<1>:$doRegion);
  let results = (outs Variadic<AnyType>:$results);

  let printer = [{ return ::print(p, *this); }];
  let parser = [{ return ::parse$cppClass(parser, result); }];
  let verifier = [{ return RegionBranchOpInterface::verifyTypes(*this); }];

  let extraClassDeclaration = [{
    // Inline the region into the region containing the AssumingOp and delete
    // the AssumingOp.
    //
    // This does no checks on the inputs to the AssumingOp.
    static void inlineRegionIntoParent(AssumingOp &op, PatternRewriter &rewriter);
  }];

  let builders = [
    OpBuilder<(ins "Value":$witness,
        CArg<"function_ref<SmallVector<Value, 2>(OpBuilder &, Location)>">)>
  ];

  let hasCanonicalizer = 1;
}

def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
       [NoSideEffect, ReturnLike, Terminator, HasParent<"AssumingOp">]> {
  let summary = "Yield operation";
  let description = [{
    This yield operation represents a return operation within the
    `shape.assuming` operation region. The operation takes variable number of
    operands and produces no results. The operand number and types must match
    the number and types of parent `shape.assuming` results.
  }];

  let arguments = (ins Variadic<AnyType>:$operands);

  let builders = [
    OpBuilder<(ins), [{ /* nothing to do */ }]>,
  ];

  let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";
}

def Shape_CstrBroadcastableOp : Shape_Op<"cstr_broadcastable",
                                         [Commutative, InferTypeOpInterface]> {
  let summary = "Determines if 2+ shapes can be successfully broadcasted";
  let description = [{
    Given input shapes or extent tensors, return a witness specifying if they
    are broadcastable. This broadcastable follows the same logic as what
    shape.broadcast documents.

    "cstr" operations represent runtime assertions.

    Example:
    ```mlir
    %w0 = shape.cstr_broadcastable [2,2], [3,1,2] // Passing
    %w1 = shape.cstr_broadcastable [2,2], [3,2] // Failure
    ```
  }];

  let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$shapes);
  let results = (outs Shape_WitnessType:$result);

  let assemblyFormat = "$shapes attr-dict `:` type($shapes)";

  let builders = [
  OpBuilder<(ins "::mlir::Value":$lhs, "::mlir::Value":$rhs),
    [{ build($_builder, $_state, ::llvm::makeArrayRef({lhs, rhs})); }]>,
  ];

  let extraClassDeclaration = [{
    // TODO: This should really be automatic. Figure out how to not need this defined.
    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
    ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
    ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
    ::llvm::SmallVectorImpl<::mlir::Type>&inferredReturnTypes) {
      inferredReturnTypes.push_back(::mlir::shape::WitnessType::get(context));
      return success();
    };
  }];

  let hasCanonicalizer = 1;
  let hasFolder = 1;
  let verifier = [{ return ::verify(*this); }];
}

def Shape_CstrEqOp : Shape_Op<"cstr_eq", [Commutative, InferTypeOpInterface]> {
  let summary = "Determines if all input shapes are equal";
  let description = [{
    Given 1 or more input shapes, determine if all shapes are the exact same.

    "cstr" operations represent runtime assertions.

    Example:
    ```mlir
    %w0 = shape.cstr_eq [1,2], [1,2], [1,2] // Passing
    %w1 = shape.cstr_eq [2,2], [1,2] // Failure
    ```
  }];
  let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$shapes);
  let results = (outs Shape_WitnessType:$result);

  let assemblyFormat = "$shapes attr-dict `:` type($shapes)";

  let extraClassDeclaration = [{
    // TODO: This should really be automatic. Figure out how to not need this defined.
    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
    ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
    ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
    ::llvm::SmallVectorImpl<::mlir::Type>&inferredReturnTypes) {
      inferredReturnTypes.push_back(::mlir::shape::WitnessType::get(context));
      return success();
    };
  }];

  let hasCanonicalizer = 1;
  let hasFolder = 1;
}

def Shape_ConstWitnessOp : Shape_Op<"const_witness", [ConstantLike, NoSideEffect]> {
  let summary = "An operation that returns a statically known witness value";
  let description = [{
  This operation represents a statically known witness result. This can be
  often used to canonicalize/fold constraint and assuming code that will always
  pass.

  ```mlir
  %0 = shape.const_shape [1,2,3]
  %1 = shape.const_shape [1,2,3]
  %w0 = shape.cstr_eq(%0, %1) // Can be folded to "const_witness true"
  %w1 = shape.const_witness true
  %w2 = shape.assuming_all(%w0, %w2) // Can be folded to "const_witness true"
  ```
  }];
  let arguments = (ins BoolAttr:$passing);
  let results = (outs Shape_WitnessType:$result);

  let assemblyFormat = "$passing attr-dict";

  let hasFolder = 1;
}

def Shape_CstrRequireOp : Shape_Op<"cstr_require", []> {
  let summary = "Represents a runtime assertion that an i1 is `true`";
  let description = [{
    Represents a runtime assertion that an i1 is true. It returns a
    !shape.witness to order this assertion.

    For simplicity, prefer using other cstr_* ops if they are available for a
    given constraint.

    Example:
    ```mlir
    %bool = ...
    %w0 = shape.cstr_require %bool, "msg" // Passing if `%bool` is true.
    ```

    Since this op can be used to express many different possible assertions
    (depending on whatever computation calculated `pred`), the `msg`
    should clarify the nature of the assertion for users.
  }];
  let arguments = (ins I1:$pred, StrAttr:$msg);
  let results = (outs Shape_WitnessType:$result);

  let assemblyFormat = "$pred `,` $msg attr-dict";

  let hasFolder = 1;
}

//===----------------------------------------------------------------------===//
// Shape collection ops.
//===----------------------------------------------------------------------===//

def Shape_FunctionLibraryOp : Shape_Op<"function_library",
    [AffineScope, IsolatedFromAbove, NoRegionArguments, SymbolTable, Symbol,
     NoTerminator, SingleBlock]> {
  let summary = "Represents shape functions and corresponding ops";
  let description = [{
    Represents a list of shape functions and the ops whose shape transfer
    functions they represent.

    Example:

    ```mlir
    shape.function_library {
      func @same_result_shape(%arg: !shape.value_shape) -> !shape.shape {
        %0 = shape.shape_of %arg : !shape.value_shape -> !shape.shape
        return %0 : !shape.shape
      }
    } mapping {
      std.atan = @same_result_shape
    }
    ```
  }];

  let arguments = (ins SymbolNameAttr:$sym_name,
                       OptionalAttr<StrAttr>:$sym_visibility);
  let arguments = (ins DictionaryAttr:$mapping);
  let regions = (region AnyRegion:$body);

  let extraClassDeclaration = [{
    /// Returns an associated shape function for an operation if defined.
    FuncOp getShapeFunction(Operation *op);
  }];

  let builders = [OpBuilder<(ins "StringRef":$name)>];
  let skipDefaultBuilders = 1;

  let printer = [{ ::print(p, *this); }];
  let parser = [{ return ::parse$cppClass(parser, result); }];
}

#endif // SHAPE_OPS