llvm/mlir/lib/Dialect/Polynomial/IR/PolynomialOps.cpp

//===- PolynomialOps.cpp - Polynomial dialect ops ---------------*- C++ -*-===//
//
// 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/Polynomial/IR/PolynomialOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Polynomial/IR/Polynomial.h"
#include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"
#include "mlir/Dialect/Polynomial/IR/PolynomialTypes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/PatternMatch.h"
#include "llvm/ADT/APInt.h"

usingnamespacemlir;
usingnamespacemlir::polynomial;

void FromTensorOp::build(OpBuilder &builder, OperationState &result,
                         Value input, RingAttr ring) { … }

LogicalResult FromTensorOp::verify() { … }

LogicalResult ToTensorOp::verify() { … }

LogicalResult MulScalarOp::verify() { … }

/// Test if a value is a primitive nth root of unity modulo cmod.
bool isPrimitiveNthRootOfUnity(const APInt &root, const APInt &n,
                               const APInt &cmod) { … }

/// Verify that the types involved in an NTT or INTT operation are
/// compatible.
static LogicalResult verifyNTTOp(Operation *op, RingAttr ring,
                                 RankedTensorType tensorType,
                                 std::optional<PrimitiveRootAttr> root) { … }

LogicalResult NTTOp::verify() { … }

LogicalResult INTTOp::verify() { … }

ParseResult ConstantOp::parse(OpAsmParser &parser, OperationState &result) { … }

void ConstantOp::print(OpAsmPrinter &p) { … }

LogicalResult ConstantOp::inferReturnTypes(
    MLIRContext *context, std::optional<mlir::Location> location,
    ConstantOp::Adaptor adaptor,
    llvm::SmallVectorImpl<mlir::Type> &inferredReturnTypes) { … }

//===----------------------------------------------------------------------===//
// TableGen'd canonicalization patterns
//===----------------------------------------------------------------------===//

namespace {
#include "PolynomialCanonicalization.inc"
} // namespace

void SubOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                        MLIRContext *context) { … }

void NTTOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                        MLIRContext *context) { … }

void INTTOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                         MLIRContext *context) { … }