LLZKConversionUtils.h

Go to the documentation of this file.

//===-- LLZKConversionUtils.h -----------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
//
// Shared utilities for dialect converting transformations.
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "llzk/Dialect/Array/IR/Types.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
#include "llzk/Util/Concepts.h"
 
#include <mlir/IR/PatternMatch.h>
#include <mlir/IR/SymbolTable.h>
#include <mlir/Transforms/DialectConversion.h>
 
#include <llvm/ADT/DenseMap.h>
#include <llvm/ADT/StringSet.h>
#include <llvm/ADT/Twine.h>
 
#include <optional>
#include <string>
 
namespace llzk {

inline mlir::DictionaryAttr
withFunctionNameAttr(mlir::DictionaryAttr attrs, llvm::StringRef attrName, llvm::StringRef name) {
  mlir::NamedAttrList newAttrs(attrs);
  newAttrs.set(attrName, mlir::StringAttr::get(attrs.getContext(), name));
  return newAttrs.getDictionary(attrs.getContext());
}

inline mlir::DictionaryAttr
withFunctionArgNameAttr(mlir::DictionaryAttr attrs, llvm::StringRef name) {
  return withFunctionNameAttr(attrs, function::ARG_NAME_ATTR_NAME, name);
}

inline mlir::DictionaryAttr
withFunctionResNameAttr(mlir::DictionaryAttr attrs, llvm::StringRef name) {
  return withFunctionNameAttr(attrs, function::RES_NAME_ATTR_NAME, name);
}

inline std::string
reserveUniqueAttrName(llvm::StringSet<> &usedNames, llvm::StringRef desiredName) {
  if (!usedNames.contains(desiredName)) {
    usedNames.insert(desiredName);
    return desiredName.str();
  }
 
  for (unsigned suffix = 1;; ++suffix) {
    std::string candidate = (desiredName + "#" + llvm::Twine(suffix)).str();
    if (!usedNames.contains(candidate)) {
      usedNames.insert(candidate);
      return candidate;
    }
  }
}

inline std::optional<mlir::StringAttr>
getAttrAtIndexWithName(mlir::ArrayAttr attrs, unsigned index, llvm::StringRef attrName) {
  if (!attrs || index >= attrs.size()) {
    return std::nullopt;
  }
  if (auto dictAttr = llvm::dyn_cast<mlir::DictionaryAttr>(attrs[index])) {
    if (auto nameAttr = llvm::dyn_cast_if_present<mlir::StringAttr>(dictAttr.get(attrName))) {
      return nameAttr;
    }
  }
  return std::nullopt;
}

struct SplitFunctionNameInfo {
  llvm::SmallVector<std::optional<llvm::StringRef>> originalNames;
  llvm::SmallVector<llvm::StringRef> existingNames;
  llvm::SmallVector<llvm::SmallVector<std::string>> splitNameSuffixes;
};

template <typename GetNameAttrFn, typename GetSplitSuffixesFn>
inline SplitFunctionNameInfo collectSplitFunctionNameInfo(
    mlir::ArrayRef<mlir::Type> origTypes, GetNameAttrFn &&getNameAttr,
    GetSplitSuffixesFn &&getSplitSuffixes
) {
  SplitFunctionNameInfo info;
  info.originalNames.reserve(origTypes.size());
  info.splitNameSuffixes.reserve(origTypes.size());
  for (auto [i, type] : llvm::enumerate(origTypes)) {
    if (std::optional<mlir::StringAttr> nameAttr = getNameAttr(i)) {
      info.originalNames.push_back(nameAttr->getValue());
      info.existingNames.push_back(nameAttr->getValue());
    } else {
      info.originalNames.push_back(std::nullopt);
    }
    info.splitNameSuffixes.push_back(getSplitSuffixes(type));
  }
  return info;
}

inline mlir::ArrayAttr replicateFunctionNameAttrsAsNeeded(
    mlir::ArrayAttr origAttrs, const llvm::SmallVector<size_t> &originalIdxToSize,
    const llvm::SmallVector<mlir::Type> &newTypes, llvm::StringRef functionNameAttrName,
    llvm::ArrayRef<std::optional<llvm::StringRef>> origNames = {},
    llvm::ArrayRef<llvm::StringRef> existingNames = {},
    llvm::ArrayRef<llvm::SmallVector<std::string>> splitNameSuffixes = {}
) {
  if (!origAttrs) {
    return nullptr;
  }
  assert(originalIdxToSize.size() == origAttrs.size());
  if (originalIdxToSize.size() == newTypes.size()) {
    return nullptr;
  }
 
  llvm::SmallVector<mlir::Attribute> newAttrs;
  llvm::StringSet<> usedNames;
  if (!origNames.empty()) {
    for (llvm::StringRef name : existingNames) {
      usedNames.insert(name);
    }
  }
 
  for (auto [i, s] : llvm::enumerate(originalIdxToSize)) {
    mlir::Attribute attr = origAttrs[i];
    if (!origNames.empty() && !splitNameSuffixes.empty() && s != 1 && origNames[i]) {
      assert(i < splitNameSuffixes.size());
      assert(splitNameSuffixes[i].size() == s);
      auto dictAttr = llvm::cast<mlir::DictionaryAttr>(attr);
      llvm::StringRef name = *origNames[i];
      for (llvm::StringRef suffix : splitNameSuffixes[i]) {
        std::string desiredName = (llvm::Twine(name) + suffix).str();
        newAttrs.push_back(withFunctionNameAttr(
            dictAttr, functionNameAttrName, reserveUniqueAttrName(usedNames, desiredName)
        ));
      }
      continue;
    }
    newAttrs.append(s, attr);
  }
  return mlir::ArrayAttr::get(origAttrs.getContext(), newAttrs);
}

inline function::CallOp createCallPreservingInstantiationOperands(
    mlir::Location loc, mlir::TypeRange newResultTypes, function::CallOp oldCall,
    llvm::ArrayRef<mlir::ValueRange> mapOperands, mlir::ValueRange argOperands,
    mlir::ConversionPatternRewriter &rewriter
) {
  llvm::SmallVector<mlir::Attribute> templateParams;
  if (mlir::ArrayAttr templateParamsAttr = oldCall.getTemplateParamsAttr()) {
    templateParams.append(templateParamsAttr.begin(), templateParamsAttr.end());
  }
 
  function::CallOp newCall;
  if (oldCall.getMapOperands().empty()) {
    newCall = rewriter.create<function::CallOp>(
        loc, newResultTypes, oldCall.getCalleeAttr(), argOperands, templateParams
    );
  } else {
    newCall = rewriter.create<function::CallOp>(
        loc, newResultTypes, oldCall.getCalleeAttr(), mapOperands, oldCall.getNumDimsPerMapAttr(),
        argOperands, templateParams
    );
  }
 
  newCall->setDiscardableAttrs(oldCall->getDiscardableAttrDictionary());
  return newCall;
}

inline static mlir::Type replaceAffineMapArrayDimsWithWildcards(mlir::Type type) {
  auto arrTy = llvm::dyn_cast<array::ArrayType>(type);
  if (!arrTy) {
    return type;
  }
 
  mlir::Builder builder(arrTy.getContext());
  llvm::SmallVector<mlir::Attribute> dims;
  dims.reserve(arrTy.getDimensionSizes().size());
  for (mlir::Attribute dimSize : arrTy.getDimensionSizes()) {
    if (llvm::isa<mlir::AffineMapAttr>(dimSize)) {
      dims.push_back(builder.getIndexAttr(mlir::ShapedType::kDynamic));
    } else {
      dims.push_back(dimSize);
    }
  }
 
  return arrTy.cloneWith(replaceAffineMapArrayDimsWithWildcards(arrTy.getElementType()), dims);
}

class FunctionTypeConverter {
 
protected:
  virtual llvm::SmallVector<mlir::Type> convertInputs(mlir::ArrayRef<mlir::Type> origTypes) = 0;
  virtual llvm::SmallVector<mlir::Type> convertResults(mlir::ArrayRef<mlir::Type> origTypes) = 0;
 
  virtual mlir::ArrayAttr
  convertInputAttrs(mlir::ArrayAttr origAttrs, llvm::SmallVector<mlir::Type> newTypes) = 0;
  virtual mlir::ArrayAttr
  convertResultAttrs(mlir::ArrayAttr origAttrs, llvm::SmallVector<mlir::Type> newTypes) = 0;
 
  virtual void processBlockArgs(mlir::Block &entryBlock, mlir::RewriterBase &rewriter) = 0;
 
public:
  virtual ~FunctionTypeConverter() = default;
 
  void convert(function::FuncDefOp op, mlir::RewriterBase &rewriter) {
    // Update in/out types of the function
    mlir::FunctionType oldTy = op.getFunctionType();
    llvm::SmallVector<mlir::Type> newInputs = convertInputs(oldTy.getInputs());
    llvm::SmallVector<mlir::Type> newResults = convertResults(oldTy.getResults());
    mlir::FunctionType newTy = mlir::FunctionType::get(
        oldTy.getContext(), mlir::TypeRange(newInputs), mlir::TypeRange(newResults)
    );
    if (newTy == oldTy) {
      return; // nothing to change
    }
 
    // Pre-condition: arg/result count equals corresponding attribute count
    assert(!op.getResAttrsAttr() || op.getResAttrsAttr().size() == op.getNumResults());
    assert(!op.getArgAttrsAttr() || op.getArgAttrsAttr().size() == op.getNumArguments());
    rewriter.modifyOpInPlace(op, [&]() {
      op.setFunctionType(newTy);
 
      // If any input or result types were added, ensure the attributes are updated too.
      if (mlir::ArrayAttr newArgAttrs = convertInputAttrs(op.getArgAttrsAttr(), newInputs)) {
        op.setArgAttrsAttr(newArgAttrs);
      }
      if (mlir::ArrayAttr newResAttrs = convertResultAttrs(op.getResAttrsAttr(), newResults)) {
        op.setResAttrsAttr(newResAttrs);
      }
    });
    // Post-condition: arg/result count equals corresponding attribute count
    assert(!op.getResAttrsAttr() || op.getResAttrsAttr().size() == op.getNumResults());
    assert(!op.getArgAttrsAttr() || op.getArgAttrsAttr().size() == op.getNumArguments());
 
    // If the function has a body, ensure the entry block arguments match the function inputs.
    if (mlir::Region *body = op.getCallableRegion()) {
      mlir::Block &entryBlock = body->front();
      bool blockArgsNeedUpdate =
          !std::cmp_equal(entryBlock.getNumArguments(), newInputs.size()) ||
          llvm::any_of(llvm::zip_equal(entryBlock.getArgumentTypes(), newInputs), [](auto pair) {
        return std::get<0>(pair) != std::get<1>(pair);
      });
      if (blockArgsNeedUpdate) {
        processBlockArgs(entryBlock, rewriter);
        // Post-condition: block args must match function inputs in both arity and type.
        assert(std::cmp_equal(entryBlock.getNumArguments(), newInputs.size()));
        for (unsigned i = 0, e = entryBlock.getNumArguments(); i < e; ++i) {
          assert(entryBlock.getArgument(i).getType() == newInputs[i]);
        }
      }
    }
  }
};

template <
    typename ImplClass, HasInterface<component::MemberRefOpInterface> MemberRefOpClass,
    typename GenHeaderType, typename IdType>
class SplitAggregateInMemberRefOp : public mlir::OpConversionPattern<MemberRefOpClass> {
public:
  using MemberInfo = std::pair<mlir::StringAttr, mlir::Type>;
  using LocalMemberReplacementMap = llvm::DenseMap<IdType, MemberInfo>;
  using MemberReplacementMap = llvm::DenseMap<
      component::StructDefOp, llvm::DenseMap<mlir::StringAttr, LocalMemberReplacementMap>>;
 
private:
  mlir::SymbolTableCollection &tables;
  const MemberReplacementMap &repMapRef;
 
  // Static check to ensure the methods are implemented in all subclasses.
  inline static void ensureImplementedAtCompile() {
    static_assert(
        sizeof(MemberRefOpClass) == 0,
        "SplitAggregateInMemberRefOp not implemented for requested type."
    );
  }
 
protected:
  using OpAdaptor = typename MemberRefOpClass::Adaptor;

  static GenHeaderType genHeader(MemberRefOpClass, mlir::ConversionPatternRewriter &) {
    ensureImplementedAtCompile();
    llvm_unreachable("must have concrete instantiation");
  }

  static void forId(
      mlir::Location, GenHeaderType &, IdType, MemberInfo, OpAdaptor,
      mlir::ConversionPatternRewriter &
  ) {
    ensureImplementedAtCompile();
    llvm_unreachable("must have concrete instantiation");
  }
 
public:
  // Suppress false positive from `clang-tidy`
  // NOLINTNEXTLINE(bugprone-crtp-constructor-accessibility)
  SplitAggregateInMemberRefOp(
      mlir::MLIRContext *ctx, mlir::SymbolTableCollection &symTables,
      const MemberReplacementMap &memberRepMap
  )
      : mlir::OpConversionPattern<MemberRefOpClass>(ctx), tables(symTables),
        repMapRef(memberRepMap) {}
 
  static bool legal(MemberRefOpClass) {
    ensureImplementedAtCompile();
    llvm_unreachable("must have concrete instantiation");
    return false;
  }
 
  mlir::LogicalResult matchAndRewrite(
      MemberRefOpClass op, OpAdaptor adaptor, mlir::ConversionPatternRewriter &rewriter
  ) const override {
    if (ImplClass::legal(op)) {
      return mlir::failure();
    }
    component::StructType tgtStructTy =
        llvm::cast<component::MemberRefOpInterface>(op.getOperation()).getStructType();
    assert(tgtStructTy);
    auto tgtStructDef = tgtStructTy.getDefinition(tables, op);
    assert(mlir::succeeded(tgtStructDef));
 
    GenHeaderType prefixResult = ImplClass::genHeader(op, rewriter);
 
    const LocalMemberReplacementMap &idToName =
        repMapRef.at(tgtStructDef->get()).at(op.getMemberNameAttr().getAttr());
    // Split the aggregate member into a series of scalar member ops.
    for (const auto &[id, newMember] : idToName) {
      ImplClass::forId(op.getLoc(), prefixResult, id, newMember, adaptor, rewriter);
    }
    if constexpr (requires { ImplClass::finalize(op, prefixResult, adaptor, rewriter); }) {
      ImplClass::finalize(op, prefixResult, adaptor, rewriter);
    }
    rewriter.eraseOp(op);
    return mlir::success();
  }
};
 
} // namespace llzk