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/Array/IR/Types.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/LLZK/IR/Ops.h"
#include "llzk/Dialect/POD/IR/Ops.h"
#include "llzk/Dialect/POD/IR/Types.h"
#include "llzk/Dialect/Polymorphic/IR/Ops.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
#include "llzk/Transforms/ConversionUtils.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/ArrayRef.h>
#include <llvm/ADT/STLExtras.h>
#include <llvm/ADT/SmallVector.h>
#include <llvm/Support/Casting.h>
 
#include <cassert>
#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;
}

struct OpClassesWithStructTypes {

  using WithGeneralBuilder = 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,
      llzk::NonDetOp,
      pod::ReadPodOp,
      pod::WritePodOp,
      polymorphic::UnifiableCastOp,
      polymorphic::ConstReadOp
      // clang-format on
      >;

  using NoGeneralBuilder = std::tuple<function::CallOp, array::CreateArrayOp, pod::NewPodOp>;
};
 
namespace {

class ExecuteRegionTypeConversionPattern
    : public mlir::OpConversionPattern<mlir::scf::ExecuteRegionOp> {
public:
  using OpConversionPattern::OpConversionPattern;
 
  mlir::LogicalResult matchAndRewrite(
      mlir::scf::ExecuteRegionOp op, OneToNOpAdaptor, mlir::ConversionPatternRewriter &rewriter
  ) const override {
    mlir::SmallVector<mlir::Type> newResultTypes;
    mlir::SmallVector<unsigned> resultOffsets {0};
    for (mlir::Type resultType : op.getResultTypes()) {
      if (mlir::failed(getTypeConverter()->convertTypes(resultType, newResultTypes))) {
        return rewriter.notifyMatchFailure(op, "could not convert result type");
      }
      resultOffsets.push_back(newResultTypes.size());
    }
 
    if (mlir::failed(rewriter.convertRegionTypes(&op.getRegion(), *getTypeConverter()))) {
      return rewriter.notifyMatchFailure(op, "could not convert region types");
    }
 
    auto newOp = mlir::scf::ExecuteRegionOp::create(rewriter, op.getLoc(), newResultTypes);
    newOp->setAttrs(op->getAttrs());
    rewriter.inlineRegionBefore(op.getRegion(), newOp.getRegion(), newOp.getRegion().end());
 
    mlir::SmallVector<mlir::ValueRange> replacements;
    replacements.reserve(op.getNumResults());
    for (unsigned i = 0; i < op.getNumResults(); ++i) {
      replacements.push_back(
          newOp.getResults().slice(resultOffsets[i], resultOffsets[i + 1] - resultOffsets[i])
      );
    }
    rewriter.replaceOpWithMultiple(op, replacements);
    return mlir::success();
  }
};

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 matchAndRewrite(
      array::CreateArrayOp op, OpAdaptor adapter, mlir::ConversionPatternRewriter &rewriter
  ) const override {
    mlir::Type newType = getTypeConverter()->convertType(op.getType());
    if (!newType) {
      return op->emitError("Could not convert Op result type.");
    }
    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
      );
    }
    return mlir::success();
  }
};

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();
  }
};

class NewPodOpClassReplacePattern : public mlir::OpConversionPattern<pod::NewPodOp> {
public:
  NewPodOpClassReplacePattern(mlir::TypeConverter &converter, mlir::MLIRContext *ctx)
      : mlir::OpConversionPattern<pod::NewPodOp>(converter, ctx, 0) {}
 
  mlir::LogicalResult matchAndRewrite(
      pod::NewPodOp op, OpAdaptor adaptor, mlir::ConversionPatternRewriter &rewriter
  ) const override {
    auto newResultType = dyn_cast_if_present<pod::PodType>(
        getTypeConverter()->convertType(op.getResult().getType())
    );
    if (!newResultType) {
      return op->emitError("Could not convert Op result types.");
    }
    replaceOpWithNewOp<pod::NewPodOp>(
        rewriter, op, newResultType, adaptor.getMapOperands(), op.getNumDimsPerMapAttr(),
        pod::getInitializedRecordValues(adaptor.getInitialValues(), op.getInitializedRecords())
    );
    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, NewPodOpClassReplacePattern>(
      tyConv, ctx
  );
  // Add builtin FunctionType and SCF op converters
  mlir::populateFunctionOpInterfaceTypeConversionPattern<function::FuncDefOp>(patterns, tyConv);
  mlir::scf::populateSCFStructuralTypeConversionsAndLegality(tyConv, patterns, target);
  patterns.add<ExecuteRegionTypeConversionPattern>(tyConv, ctx);
  target.addDynamicallyLegalOp<mlir::scf::ExecuteRegionOp>([&](mlir::scf::ExecuteRegionOp op) {
    return tyConv.isLegal(op.getResultTypes());
  });
  target.addDynamicallyLegalOp<mlir::scf::YieldOp>([&](mlir::scf::YieldOp op) {
    mlir::Operation *parent = op->getParentOp();
    if (!llvm::isa<
            mlir::scf::ExecuteRegionOp, mlir::scf::ForOp, mlir::scf::IfOp, mlir::scf::WhileOp>(
            parent
        )) {
      return true;
    }
    return tyConv.isLegal(op.getOperandTypes());
  });
  return patterns;
}
 
} // namespace llzk