ConversionUtils.cpp

#include "lib/Utils/ConversionUtils.h"

#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>

#include "lib/Dialect/LWE/IR/LWEAttributes.h"
#include "lib/Dialect/Secret/IR/SecretOps.h"
#include "llvm/include/llvm/ADT/TypeSwitch.h"         // from @llvm-project
#include "llvm/include/llvm/Support/ErrorHandling.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"   // from @llvm-project
#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Func/Transforms/FuncConversions.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/SCF/Transforms/Patterns.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
#include "mlir/include/mlir/IR/IRMapping.h"              // from @llvm-project
#include "mlir/include/mlir/IR/OpDefinition.h"           // from @llvm-project
#include "mlir/include/mlir/IR/OperationSupport.h"       // from @llvm-project
#include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
#include "mlir/include/mlir/IR/Region.h"                 // from @llvm-project
#include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
#include "mlir/include/mlir/IR/Verifier.h"               // from @llvm-project
#include "mlir/include/mlir/IR/Visitors.h"               // from @llvm-project
#include "mlir/include/mlir/Interfaces/FunctionInterfaces.h"  // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"           // from @llvm-project
#include "mlir/include/mlir/Support/LogicalResult.h"  // from @llvm-project
#include "mlir/include/mlir/Transforms/DialectConversion.h"  // from @llvm-project

namespace mlir {
namespace heir {

using ::mlir::func::CallOp;
using ::mlir::func::FuncOp;
using ::mlir::func::ReturnOp;

LogicalResult convertAnyOperand(const TypeConverter *typeConverter,
                                Operation *op, ArrayRef<Value> operands,
                                ConversionPatternRewriter &rewriter) {
  if (typeConverter->isLegal(op)) {
    return failure();
  }

  SmallVector<Type> newOperandTypes;
  if (failed(
          typeConverter->convertTypes(op->getOperandTypes(), newOperandTypes)))
    return failure();

  SmallVector<Type> newResultTypes;
  if (failed(typeConverter->convertTypes(op->getResultTypes(), newResultTypes)))
    return failure();

  SmallVector<std::unique_ptr<Region>, 1> regions;
  IRMapping mapping;
  for (auto &r : op->getRegions()) {
    Region *newRegion = new Region(op);
    rewriter.cloneRegionBefore(r, *newRegion, newRegion->end(), mapping);
    if (failed(rewriter.convertRegionTypes(newRegion, *typeConverter)))
      return failure();
    regions.emplace_back(newRegion);
  }

  Operation *newOp = rewriter.create(OperationState(
      op->getLoc(), op->getName().getStringRef(), operands, newResultTypes,
      op->getAttrs(), op->getSuccessors(), regions));

  rewriter.replaceOp(op, newOp);
  return success();
}

struct ConvertExtract : public OpConversionPattern<tensor::ExtractOp> {
  ConvertExtract(mlir::MLIRContext *context)
      : OpConversionPattern<tensor::ExtractOp>(context) {}

  using OpConversionPattern::OpConversionPattern;

  // Convert a tensor.extract that would type-convert to extracting a tensor to
  // a tensor.extract_slice operation instead. Specifically, this targets
  // extracting SourceType from tensor<...xSourceType>  when SourceType would be
  // type converted to tensor<...>.
  LogicalResult matchAndRewrite(
      tensor::ExtractOp op, OpAdaptor adaptor,
      ConversionPatternRewriter &rewriter) const override {
    // replace tensor.extract %t[%i] from tensor<shape x SourceType>
    // with an equivalent tensor.slice from tensor<shape x resultshape>
    auto shape = op.getTensor().getType().getShape();
    auto resultType = mlir::cast<RankedTensorType>(
        getTypeConverter()->convertType(op.getResult().getType()));
    auto resultShape = resultType.getShape();

    // expand op's list of indices by appending as many zeros as there are
    // dimension in resultShape
    SmallVector<OpFoldResult> offsets;
    offsets.append(op.getIndices().begin(), op.getIndices().end());
    for (size_t i = 0; i < resultShape.size(); ++i) {
      offsets.push_back(rewriter.getIndexAttr(0));
    }

    // expand resultShape by prepending as many ones as there are dimensions in
    // shape
    SmallVector<OpFoldResult> sizes;
    for (size_t i = 0; i < shape.size(); ++i) {
      sizes.push_back(rewriter.getIndexAttr(1));
    }
    for (int64_t i : resultShape) {
      sizes.push_back(rewriter.getIndexAttr(i));
    }

    // strides are all 1, and we need as many as there are dimensions in
    // both shape and resultShape together
    SmallVector<OpFoldResult> strides;
    for (size_t i = 0; i < shape.size() + resultShape.size(); ++i) {
      strides.push_back(rewriter.getIndexAttr(1));
    }

    rewriter.replaceOpWithNewOp<tensor::ExtractSliceOp>(
        op, resultType, adaptor.getTensor(), offsets, sizes, strides);

    return success();
  }
};

struct ConvertInsert : public OpConversionPattern<tensor::InsertOp> {
  ConvertInsert(mlir::MLIRContext *context)
      : OpConversionPattern<tensor::InsertOp>(context) {}

  using OpConversionPattern::OpConversionPattern;

  // Convert a tensor.insert that would type-convert to inserting a tensor to
  // a tensor.insert_slice operation instead. Specifically, this targets
  // inserting SourceType into tensor<...xSourceType>  when SourceType would be
  // type converted to tensor<...>.
  LogicalResult matchAndRewrite(
      tensor::InsertOp op, OpAdaptor adaptor,
      ConversionPatternRewriter &rewriter) const override {
    // replace tensor.insert %s into %t[%i] with tensor<shape x SourceType>
    // with an equivalent tensor.insert_slice with tensor<shape x resultshape>
    auto shape = op.getDest().getType().getShape();
    auto resultType = mlir::cast<RankedTensorType>(
        getTypeConverter()->convertType(op.getScalar().getType()));
    auto resultShape = resultType.getShape();

    // expand op's list of indices by appending as many zeros as there are
    // dimension in resultShape
    SmallVector<OpFoldResult> offsets;
    offsets.append(op.getIndices().begin(), op.getIndices().end());
    for (size_t i = 0; i < resultShape.size(); ++i) {
      offsets.push_back(rewriter.getIndexAttr(0));
    }

    // expand resultShape by prepending as many ones as there are dimensions in
    // shape
    SmallVector<OpFoldResult> sizes;
    for (size_t i = 0; i < shape.size(); ++i) {
      sizes.push_back(rewriter.getIndexAttr(1));
    }
    for (int64_t i : resultShape) {
      sizes.push_back(rewriter.getIndexAttr(i));
    }

    // strides are all 1, and we need as many as there are dimensions in
    // both shape and resultShape together
    SmallVector<OpFoldResult> strides;
    for (size_t i = 0; i < shape.size() + resultShape.size(); ++i) {
      strides.push_back(rewriter.getIndexAttr(1));
    }

    rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
        op, adaptor.getScalar(), adaptor.getDest(), offsets, sizes, strides);

    return success();
  }
};

struct ConvertFromElements
    : public OpConversionPattern<tensor::FromElementsOp> {
  ConvertFromElements(mlir::MLIRContext *context)
      : OpConversionPattern<tensor::FromElementsOp>(context) {}

  using OpConversionPattern::OpConversionPattern;

  // Converts a tensor.from_elements %s0, %s1, ... : tensor<...xSourceType>
  // where SourceType would be type-converted to tensor<...> to
  // a concatenation of the converted operands (with appropriate reshape)
  LogicalResult matchAndRewrite(
      tensor::FromElementsOp op, OpAdaptor adaptor,
      ConversionPatternRewriter &rewriter) const override {
    // Expand each of the (converted) operands:
    SmallVector<Value> newOperands;
    for (auto o : adaptor.getElements()) {
      // extend tensor<...xT> to tensor<1x...xT>
      if (auto tensorType = mlir::dyn_cast<RankedTensorType>(o.getType())) {
        auto shape = tensorType.getShape();
        SmallVector<int64_t> newShape(1, 1);
        newShape.append(shape.begin(), shape.end());

        // Create a dense constant for targetShape
        auto shapeOp = rewriter.create<mlir::arith::ConstantOp>(
            op.getLoc(),
            RankedTensorType::get(newShape.size(), rewriter.getIndexType()),
            rewriter.getIndexTensorAttr(newShape));

        auto reshapeOp = rewriter.create<tensor::ReshapeOp>(
            op.getLoc(),
            RankedTensorType::get(newShape, tensorType.getElementType()), o,
            shapeOp);
        newOperands.push_back(reshapeOp);
      } else {
        newOperands.push_back(o);
      }
    }
    // Create the final tensor.concat operation
    rewriter.replaceOpWithNewOp<tensor::ConcatOp>(op, 0, newOperands);

    return success();
  }
};

void addTensorOfTensorConversionPatterns(TypeConverter &typeConverter,
                                         RewritePatternSet &patterns,
                                         ConversionTarget &target) {
  target.addDynamicallyLegalDialect<tensor::TensorDialect>(
      [&](Operation *op) { return typeConverter.isLegal(op); });

  typeConverter.addConversion([&](TensorType type) -> Type {
    if (!typeConverter.isLegal(type.getElementType())) {
      if (auto convertedType =
              typeConverter.convertType(type.getElementType())) {
        if (auto castConvertedType =
                mlir::dyn_cast<RankedTensorType>(convertedType)) {
          //  Create the combined shape
          auto polyShape = castConvertedType.getShape();
          auto tensorShape = type.getShape();
          SmallVector<int64_t, 4> combinedShape(tensorShape.begin(),
                                                tensorShape.end());
          combinedShape.append(polyShape.begin(), polyShape.end());
          auto combinedType = RankedTensorType::get(
              combinedShape, castConvertedType.getElementType());
          return combinedType;
        }
      }
    }
    return type;
  });

  target.addDynamicallyLegalDialect<affine::AffineDialect>(
      [&](Operation *op) { return typeConverter.isLegal(op); });

  patterns
      .add<ConvertAny<>, ConvertExtract, ConvertInsert, ConvertFromElements>(
          typeConverter, patterns.getContext());
}

void addStructuralConversionPatterns(TypeConverter &typeConverter,
                                     RewritePatternSet &patterns,
                                     ConversionTarget &target) {
  populateFunctionOpInterfaceTypeConversionPattern<FuncOp>(patterns,
                                                           typeConverter);
  target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
    return typeConverter.isSignatureLegal(op.getFunctionType()) &&
           typeConverter.isLegal(&op.getBody());
  });

  populateReturnOpTypeConversionPattern(patterns, typeConverter);
  target.addDynamicallyLegalOp<func::ReturnOp>(
      [&](func::ReturnOp op) { return typeConverter.isLegal(op); });

  populateCallOpTypeConversionPattern(patterns, typeConverter);
  target.addDynamicallyLegalOp<func::CallOp>(
      [&](func::CallOp op) { return typeConverter.isLegal(op); });

  populateBranchOpInterfaceTypeConversionPattern(patterns, typeConverter);
  target.markUnknownOpDynamicallyLegal([&](Operation *op) {
    return isNotBranchOpInterfaceOrReturnLikeOp(op) ||
           isLegalForBranchOpInterfaceTypeConversionPattern(op,
                                                            typeConverter) ||
           isLegalForReturnOpTypeConversionPattern(op, typeConverter);
  });

  scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter, patterns,
                                                       target);
}

int widthFromEncodingAttr(Attribute encoding) {
  return llvm::TypeSwitch<Attribute, int>(encoding)
      .Case<lwe::BitFieldEncodingAttr, lwe::UnspecifiedBitFieldEncodingAttr>(
          [](auto attr) -> int { return attr.getCleartextBitwidth(); })
      .Default([](Attribute attr) -> int {
        llvm_unreachable("Unsupported encoding attribute");
        return 0;
      });
}

Attribute TypeWithAttrTypeConverter::getValueAttr(Value value) const {
  Attribute attr;
  if (auto blockArg = dyn_cast<BlockArgument>(value)) {
    auto *parentOp = blockArg.getOwner()->getParentOp();
    auto funcOp = dyn_cast<FunctionOpInterface>(parentOp);
    if (funcOp) {
      attr = funcOp.getArgAttr(blockArg.getArgNumber(), attrName);
    }
  } else {
    auto *parentOp = value.getDefiningOp();
    attr = parentOp->getAttr(attrName);
  }
  return attr;
}

void TypeWithAttrTypeConverter::convertValueRangeTypes(
    ValueRange values, SmallVectorImpl<Type> &newTypes) const {
  newTypes.reserve(values.size());
  for (auto value : values) {
    Attribute attr = getValueAttr(value);
    auto newType = convertTypeWithAttr(value.getType(), attr);
    // this is actually unsafe...
    // all the thing should be done through the rewriter,
    // if we are using the rewriter
    value.setType(newType);
    newTypes.push_back(newType);
  }
}

void TypeWithAttrTypeConverter::convertOpResultTypes(
    Operation *op, SmallVectorImpl<Type> &newResultTypes) const {
  newResultTypes.reserve(op->getResultTypes().size());
  auto attr = op->getAttr(attrName);
  for (auto resultType : op->getResultTypes()) {
    auto newType = convertTypeWithAttr(resultType, attr);
    newResultTypes.push_back(newType);
  }
}

void TypeWithAttrTypeConverter::convertFuncArgumentAndResultTypes(
    FunctionOpInterface funcOp, SmallVectorImpl<Type> &newArgTypes,
    SmallVectorImpl<Type> &newResultTypes) const {
  for (auto argument : funcOp.getArguments()) {
    auto attr = funcOp.getArgAttr(argument.getArgNumber(), attrName);
    auto newType = convertTypeWithAttr(argument.getType(), attr);
    // this is actually unsafe...
    // we should go through rewriter.convertRegionTypes,
    // which will create unresolved_materializaton,
    // and everything is safe.
    argument.setType(newType);
    newArgTypes.push_back(newType);
  }
  // did not convert block arg/signature though..
  for (auto &block : funcOp.getBlocks()) {
    for (auto result : block.getTerminator()->getOperands()) {
      auto attr = getValueAttr(result);
      auto newType = convertTypeWithAttr(result.getType(), attr);
      result.setType(newType);
      newResultTypes.push_back(newType);
    }
  }
}

bool TypeWithAttrTypeConverter::isValueLegal(Value value) {
  auto attr = getValueAttr(value);
  return value.getType() == convertTypeWithAttr(value.getType(), attr);
}

bool TypeWithAttrTypeConverter::isOperationLegal(Operation *op) {
  for (auto operand : op->getOperands()) {
    if (!isValueLegal(operand)) {
      return false;
    }
  }
  for (auto result : op->getResults()) {
    if (!isValueLegal(result)) {
      return false;
    }
  }
  for (auto &region : op->getRegions()) {
    for (auto &block : region) {
      for (auto argument : block.getArguments()) {
        if (!isValueLegal(argument)) {
          return false;
        }
      }
      for (auto result : block.getTerminator()->getOperands()) {
        if (!isValueLegal(result)) {
          return false;
        }
      }
    }
  }
  return true;
}

bool TypeWithAttrTypeConverter::isFuncArgumentAndResultLegal(
    FunctionOpInterface funcOp) {
  for (auto argument : funcOp.getArguments()) {
    if (!isValueLegal(argument)) {
      return false;
    }
  }
  for (auto &block : funcOp.getBlocks()) {
    for (auto result : block.getTerminator()->getOperands()) {
      if (!isValueLegal(result)) {
        return false;
      }
    }
  }
  return true;
}

}  // namespace heir
}  // namespace mlir