TensorExtCanonicalization.td

#ifndef INCLUDE_DIALECT_TENSOREXT_IR_TENSOREXTCANONICALIZATION_TD_
#define INCLUDE_DIALECT_TENSOREXT_IR_TENSOREXTCANONICALIZATION_TD_

include "TensorExtOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "mlir/Dialect/Tensor/IR/TensorOps.td"
include "mlir/IR/PatternBase.td"

// FIXME: refactor these helpers to a common file with TensorExtRotationAlignment.td
defvar DefOverflow = ConstantEnumCase<Arith_IntegerOverflowAttr, "none">;

def MakeSingleResultVariadic: NativeCodeCall<"{ $0 }">;

def IndexType : NativeCodeCall<"$_builder.getIndexType()">;

def MakeMatchingTypedZero : NativeCodeCall<
      "$_builder.getIntegerAttr("
        "$_builder.getIntegerType("
          "$0.getType().getIntOrFloatBitWidth()"
        "), 0)">;

def CreateIndexCastOp : NativeCodeCall<
      "$_builder.create<arith::IndexCastOp>($0.getLoc(), $1.getType(), $0)">;

def IsZero :
    Constraint<
      CPred<"llvm::cast<mlir::IntegerAttr>($0).getValue().isZero()">>;

def AreOpposites :
    Constraint<
      CPred<"llvm::cast<mlir::IntegerAttr>($0).getValue() "
            "== llvm::cast<mlir::IntegerAttr>($1).getValue()">>;

def OutOfBoundsOfTensorDim :
  Constraint<
    CPred<
      "llvm::cast<mlir::IntegerAttr>($0).getValue().getSExtValue() < 0 "
      "|| llvm::cast<mlir::IntegerAttr>($0).getValue().getSExtValue() > "
      "llvm::cast<mlir::ShapedType>($1.getType()).getShape()[0]"
    >
  >;

// rotate %t, 0 -> %t
def DropZeroRotation : Pat<
  (TensorExt_RotateOp $tensor, (ConstantLikeMatcher APIntAttr:$c0)),
  (replaceWithValue $tensor),
  [(IsZero $c0)]
>;

// rotate %t, x -> rotate %t, x mod size
def NormalizeRotationIndex : Pat<
  (TensorExt_RotateOp $tensor, (Arith_ConstantOp:$shiftOp APIntAttr:$shiftAmount)),
  (TensorExt_RotateOp $tensor,
      (Arith_RemUIOp
        $shiftOp,
        // Only works for 1D tensors: index is taken modulo the tensor length,
        // i.e., dim 0
        (CreateIndexCastOp
          (Tensor_DimOp $tensor, (Arith_ConstantOp ConstantAttr<IndexAttr, "0">)),
          $shiftOp))
      ),
  [(OutOfBoundsOfTensorDim $shiftAmount, $tensor)]
>;

// rotate %t, x + rotate %t y -> rotate %t (x+y)
def SumOfConstantRotations : Pat<
  (TensorExt_RotateOp
    (TensorExt_RotateOp $tensor, (Arith_ConstantOp:$xOp APIntAttr:$x)),
    (Arith_ConstantOp:$yOp APIntAttr:$y)),
  (TensorExt_RotateOp $tensor, (Arith_AddIOp $xOp, $yOp, DefOverflow)),
  []
>;

// A rotation followed by extraction can be extracted directly from the
// original tensor.
def RotatePlusExtractToIndexedExtract : Pat<
  (Tensor_ExtractOp
    (TensorExt_RotateOp $tensor, $shift),
    (variadic $index)),
  (Tensor_ExtractOp
    $tensor,
    (MakeSingleResultVariadic (Arith_AddIOp $shift, $index, DefOverflow)))
>;

#endif  // INCLUDE_DIALECT_TENSOREXT_IR_TENSOREXTCANONICALIZATION_TD_