llvm/mlir/lib/Dialect/Index/IR/IndexOps.cpp

//===- IndexOps.cpp - Index operation definitions --------------------------==//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Index/IR/IndexOps.h"
#include "mlir/Dialect/Index/IR/IndexAttrs.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/Utils/InferIntRangeCommon.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/ADT/TypeSwitch.h"

usingnamespacemlir;
usingnamespacemlir::index;

//===----------------------------------------------------------------------===//
// IndexDialect
//===----------------------------------------------------------------------===//

void IndexDialect::registerOperations() { … }

Operation *IndexDialect::materializeConstant(OpBuilder &b, Attribute value,
                                             Type type, Location loc) { … }

//===----------------------------------------------------------------------===//
// Fold Utilities
//===----------------------------------------------------------------------===//

/// Fold an index operation irrespective of the target bitwidth. The
/// operation must satisfy the property:
///
/// ```
/// trunc(f(a, b)) = f(trunc(a), trunc(b))
/// ```
///
/// For all values of `a` and `b`. The function accepts a lambda that computes
/// the integer result, which in turn must satisfy the above property.
static OpFoldResult foldBinaryOpUnchecked(
    ArrayRef<Attribute> operands,
    function_ref<std::optional<APInt>(const APInt &, const APInt &)>
        calculate) { … }

/// Fold an index operation only if the truncated 64-bit result matches the
/// 32-bit result for operations that don't satisfy the above property. These
/// are operations where the upper bits of the operands can affect the lower
/// bits of the results.
///
/// The function accepts a lambda that computes the integer result in both
/// 64-bit and 32-bit. If either call returns `std::nullopt`, the operation is
/// not folded.
static OpFoldResult foldBinaryOpChecked(
    ArrayRef<Attribute> operands,
    function_ref<std::optional<APInt>(const APInt &, const APInt &lhs)>
        calculate) { … }

//===----------------------------------------------------------------------===//
// AddOp
//===----------------------------------------------------------------------===//

OpFoldResult AddOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// SubOp
//===----------------------------------------------------------------------===//

OpFoldResult SubOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// MulOp
//===----------------------------------------------------------------------===//

OpFoldResult MulOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// DivSOp
//===----------------------------------------------------------------------===//

OpFoldResult DivSOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// DivUOp
//===----------------------------------------------------------------------===//

OpFoldResult DivUOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// CeilDivSOp
//===----------------------------------------------------------------------===//

/// Compute `ceildivs(n, m)` as `x = m > 0 ? -1 : 1` and then
/// `n*m > 0 ? (n+x)/m + 1 : -(-n/m)`.
static std::optional<APInt> calculateCeilDivS(const APInt &n, const APInt &m) { … }

OpFoldResult CeilDivSOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// CeilDivUOp
//===----------------------------------------------------------------------===//

OpFoldResult CeilDivUOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// FloorDivSOp
//===----------------------------------------------------------------------===//

/// Compute `floordivs(n, m)` as `x = m < 0 ? 1 : -1` and then
/// `n*m < 0 ? -1 - (x-n)/m : n/m`.
static std::optional<APInt> calculateFloorDivS(const APInt &n, const APInt &m) { … }

OpFoldResult FloorDivSOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// RemSOp
//===----------------------------------------------------------------------===//

OpFoldResult RemSOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// RemUOp
//===----------------------------------------------------------------------===//

OpFoldResult RemUOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// MaxSOp
//===----------------------------------------------------------------------===//

OpFoldResult MaxSOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// MaxUOp
//===----------------------------------------------------------------------===//

OpFoldResult MaxUOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// MinSOp
//===----------------------------------------------------------------------===//

OpFoldResult MinSOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// MinUOp
//===----------------------------------------------------------------------===//

OpFoldResult MinUOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// ShlOp
//===----------------------------------------------------------------------===//

OpFoldResult ShlOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// ShrSOp
//===----------------------------------------------------------------------===//

OpFoldResult ShrSOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// ShrUOp
//===----------------------------------------------------------------------===//

OpFoldResult ShrUOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// AndOp
//===----------------------------------------------------------------------===//

OpFoldResult AndOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// OrOp
//===----------------------------------------------------------------------===//

OpFoldResult OrOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// XOrOp
//===----------------------------------------------------------------------===//

OpFoldResult XOrOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// CastSOp
//===----------------------------------------------------------------------===//

static OpFoldResult
foldCastOp(Attribute input, Type type,
           function_ref<APInt(const APInt &, unsigned)> extFn,
           function_ref<APInt(const APInt &, unsigned)> extOrTruncFn) { … }

bool CastSOp::areCastCompatible(TypeRange lhsTypes, TypeRange rhsTypes) { … }

OpFoldResult CastSOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// CastUOp
//===----------------------------------------------------------------------===//

bool CastUOp::areCastCompatible(TypeRange lhsTypes, TypeRange rhsTypes) { … }

OpFoldResult CastUOp::fold(FoldAdaptor adaptor) { … }

//===----------------------------------------------------------------------===//
// CmpOp
//===----------------------------------------------------------------------===//

/// Compare two integers according to the comparison predicate.
bool compareIndices(const APInt &lhs, const APInt &rhs,
                    IndexCmpPredicate pred) { … }

/// `cmp(max/min(x, cstA), cstB)` can be folded to a constant depending on the
/// values of `cstA` and `cstB`, the max or min operation, and the comparison
/// predicate. Check whether the value folds in both 32-bit and 64-bit
/// arithmetic and to the same value.
static std::optional<bool> foldCmpOfMaxOrMin(Operation *lhsOp,
                                             const APInt &cstA,
                                             const APInt &cstB, unsigned width,
                                             IndexCmpPredicate pred) { … }

/// Return the result of `cmp(pred, x, x)`
static bool compareSameArgs(IndexCmpPredicate pred) { … }

OpFoldResult CmpOp::fold(FoldAdaptor adaptor) { … }

/// Canonicalize
/// `x - y cmp 0` to `x cmp y`. or `x - y cmp 0` to `x cmp y`.
/// `0 cmp x - y` to `y cmp x`. or `0 cmp x - y` to `y cmp x`.
LogicalResult CmpOp::canonicalize(CmpOp op, PatternRewriter &rewriter) { … }

//===----------------------------------------------------------------------===//
// ConstantOp
//===----------------------------------------------------------------------===//

void ConstantOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) { … }

OpFoldResult ConstantOp::fold(FoldAdaptor adaptor) { … }

void ConstantOp::build(OpBuilder &b, OperationState &state, int64_t value) { … }

//===----------------------------------------------------------------------===//
// BoolConstantOp
//===----------------------------------------------------------------------===//

OpFoldResult BoolConstantOp::fold(FoldAdaptor adaptor) { … }

void BoolConstantOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) { … }

//===----------------------------------------------------------------------===//
// ODS-Generated Definitions
//===----------------------------------------------------------------------===//

#define GET_OP_CLASSES
#include "mlir/Dialect/Index/IR/IndexOps.cpp.inc"