TypeConversionPatterns.h

Go to the documentation of this file.

//===-- TypeConversionPatterns.h --------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// Copyright 2026 Project LLZK
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "llzk/Dialect/Array/IR/Ops.h"
#include "llzk/Dialect/Constrain/IR/Ops.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/Global/IR/Ops.h"
#include "llzk/Dialect/LLZK/IR/AttributeHelper.h"
#include "llzk/Dialect/Polymorphic/IR/Ops.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
 
#include <mlir/Dialect/SCF/Transforms/Patterns.h>
#include <mlir/IR/Attributes.h>
#include <mlir/IR/BuiltinAttributes.h>
#include <mlir/IR/BuiltinTypes.h>
#include <mlir/IR/MLIRContext.h>
#include <mlir/IR/Operation.h>
#include <mlir/IR/PatternMatch.h>
#include <mlir/Transforms/DialectConversion.h>
 
#include <llvm/ADT/SmallVector.h>
 
#include <tuple>
 
namespace llzk {

inline bool defaultLegalityCheck(const mlir::TypeConverter &tyConv, mlir::Operation *op) {
  // Check operand types and result types
  if (!tyConv.isLegal(op)) {
    return false;
  }
  // Check type attributes
  // Extend lifetime of temporary to suppress warnings.
  mlir::DictionaryAttr dictAttr = op->getAttrDictionary();
  for (mlir::NamedAttribute n : dictAttr.getValue()) {
    if (mlir::TypeAttr tyAttr = llvm::dyn_cast<mlir::TypeAttr>(n.getValue())) {
      mlir::Type t = tyAttr.getValue();
      if (mlir::FunctionType funcTy = llvm::dyn_cast<mlir::FunctionType>(t)) {
        if (!tyConv.isSignatureLegal(funcTy)) {
          return false;
        }
      } else {
        if (!tyConv.isLegal(t)) {
          return false;
        }
      }
    }
  }
  return true;
}

template <typename OpClass, typename Rewriter, typename... Args>
inline OpClass replaceOpWithNewOp(Rewriter &rewriter, mlir::Operation *op, Args &&...args) {
  mlir::DictionaryAttr attrs = op->getDiscardableAttrDictionary();
  OpClass newOp = rewriter.template replaceOpWithNewOp<OpClass>(op, std::forward<Args>(args)...);
  newOp->setDiscardableAttrs(attrs);
  return newOp;
}

static struct OpClassesWithStructTypes {

  const std::tuple<
      // clang-format off
      array::ArrayLengthOp,
      array::ReadArrayOp,
      array::WriteArrayOp,
      array::InsertArrayOp,
      array::ExtractArrayOp,
      constrain::EmitEqualityOp,
      constrain::EmitContainmentOp,
      component::MemberDefOp,
      component::MemberReadOp,
      component::MemberWriteOp,
      component::CreateStructOp,
      function::FuncDefOp,
      function::ReturnOp,
      global::GlobalDefOp,
      global::GlobalReadOp,
      global::GlobalWriteOp,
      polymorphic::UnifiableCastOp,
      polymorphic::ConstReadOp
      // clang-format on
      >
      WithGeneralBuilder {};

  const std::tuple<function::CallOp, array::CreateArrayOp> NoGeneralBuilder {};
 
} OpClassesWithStructTypes;
 
namespace {

template <typename OpClass>
class GeneralTypeReplacePattern : public mlir::OpConversionPattern<OpClass> {
public:
  GeneralTypeReplacePattern(mlir::TypeConverter &converter, mlir::MLIRContext *ctx)
      : mlir::OpConversionPattern<OpClass>(converter, ctx, 0) {}
 
  mlir::LogicalResult matchAndRewrite(
      OpClass op, typename OpClass::Adaptor adaptor, mlir::ConversionPatternRewriter &rewriter
  ) const override {
    const mlir::TypeConverter *converter = mlir::OpConversionPattern<OpClass>::getTypeConverter();
    assert(converter);
    // Convert result types
    mlir::SmallVector<mlir::Type> newResultTypes;
    if (mlir::failed(converter->convertTypes(op->getResultTypes(), newResultTypes))) {
      return op->emitError("Could not convert Op result types.");
    }
    // ASSERT: 'adaptor.getAttributes()' is empty or a subset of 'op->getAttrDictionary()' so the
    // former can be ignored without losing anything.
    assert(
        adaptor.getAttributes().empty() ||
        llvm::all_of(
            adaptor.getAttributes(), [d = op->getAttrDictionary()](mlir::NamedAttribute a) {
      return d.contains(a.getName());
    }
        )
    );
    // Convert any TypeAttr in the attribute list.
    mlir::SmallVector<mlir::NamedAttribute> newAttrs(op->getAttrDictionary().getValue());
    for (mlir::NamedAttribute &n : newAttrs) {
      if (mlir::TypeAttr t = llvm::dyn_cast<mlir::TypeAttr>(n.getValue())) {
        if (mlir::Type newType = converter->convertType(t.getValue())) {
          n.setValue(mlir::TypeAttr::get(newType));
        } else {
          return op->emitError().append("Could not convert type in attribute: ", t);
        }
      }
    }
    // Build a new Op in place of the current one
    replaceOpWithNewOp<OpClass>(
        rewriter, op, mlir::TypeRange(newResultTypes), adaptor.getOperands(),
        mlir::ArrayRef(newAttrs)
    );
    return mlir::success();
  }
};

class CreateArrayOpClassReplacePattern : public mlir::OpConversionPattern<array::CreateArrayOp> {
public:
  CreateArrayOpClassReplacePattern(mlir::TypeConverter &converter, mlir::MLIRContext *ctx)
      : mlir::OpConversionPattern<array::CreateArrayOp>(converter, ctx, 0) {}
 
  mlir::LogicalResult match(array::CreateArrayOp op) const override {
    if (getTypeConverter()->convertType(op.getType())) {
      return mlir::success();
    }
    return op->emitError("Could not convert Op result type.");
  }
 
  void rewrite(
      array::CreateArrayOp op, OpAdaptor adapter, mlir::ConversionPatternRewriter &rewriter
  ) const override {
    mlir::Type newType = getTypeConverter()->convertType(op.getType());
    assert(llvm::isa<array::ArrayType>(newType) && "CreateArrayOp must produce ArrayType result");
    mlir::DenseI32ArrayAttr numDimsPerMap = op.getNumDimsPerMapAttr();
    if (isNullOrEmpty(numDimsPerMap)) {
      replaceOpWithNewOp<array::CreateArrayOp>(
          rewriter, op, llvm::cast<array::ArrayType>(newType), adapter.getElements()
      );
    } else {
      replaceOpWithNewOp<array::CreateArrayOp>(
          rewriter, op, llvm::cast<array::ArrayType>(newType), adapter.getMapOperands(),
          numDimsPerMap
      );
    }
  }
};

class CallOpClassReplacePattern : public mlir::OpConversionPattern<function::CallOp> {
public:
  CallOpClassReplacePattern(mlir::TypeConverter &converter, mlir::MLIRContext *ctx)
      : mlir::OpConversionPattern<function::CallOp>(converter, ctx, 0) {}
 
  mlir::LogicalResult matchAndRewrite(
      function::CallOp op, OpAdaptor adapter, mlir::ConversionPatternRewriter &rewriter
  ) const override {
    mlir::SmallVector<mlir::Type> newResultTypes;
    if (mlir::failed(getTypeConverter()->convertTypes(op.getResultTypes(), newResultTypes))) {
      return op->emitError("Could not convert Op result types.");
    }
    replaceOpWithNewOp<function::CallOp>(
        rewriter, op, newResultTypes, op.getCalleeAttr(), adapter.getMapOperands(),
        op.getNumDimsPerMapAttr(), adapter.getArgOperands()
    );
    return mlir::success();
  }
};
 
template <typename I, typename NextOpClass, typename... OtherOpClasses>
inline void applyToMoreTypes(I inserter) {
  std::apply(inserter, std::tuple<NextOpClass, OtherOpClasses...> {});
}
template <typename I> inline void applyToMoreTypes(I) {}
 
} // namespace

template <typename... AdditionalOpClasses>
inline mlir::RewritePatternSet newGeneralRewritePatternSet(
    mlir::TypeConverter &tyConv, mlir::MLIRContext *ctx, mlir::ConversionTarget &target
) {
  mlir::RewritePatternSet patterns(ctx);
  auto inserter = [&](auto... opClasses) {
    patterns.add<GeneralTypeReplacePattern<decltype(opClasses)>...>(tyConv, ctx);
  };
  std::apply(inserter, OpClassesWithStructTypes.WithGeneralBuilder);
  applyToMoreTypes<decltype(inserter), AdditionalOpClasses...>(inserter);
  // Special cases for ops where GeneralTypeReplacePattern doesn't work
  patterns.add<CreateArrayOpClassReplacePattern, CallOpClassReplacePattern>(tyConv, ctx);
  // Add builtin FunctionType and SCF op converters
  mlir::populateFunctionOpInterfaceTypeConversionPattern<function::FuncDefOp>(patterns, tyConv);
  mlir::scf::populateSCFStructuralTypeConversionsAndLegality(tyConv, patterns, target);
  return patterns;
}
 
} // namespace llzk