mlir/lib/Conversion/MemRefToEmitC/MemRefToEmitC.cpp

//===- MemRefToEmitC.cpp - MemRef to EmitC conversion ---------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements patterns to convert memref ops into emitc ops.
//
//===----------------------------------------------------------------------===//

#include "mlir/Conversion/MemRefToEmitC/MemRefToEmitC.h"

#include "mlir/Conversion/ConvertToEmitC/ToEmitCInterface.h"
#include "mlir/Dialect/EmitC/IR/EmitC.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeRange.h"
#include "mlir/IR/Value.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/ADT/STLExtras.h"
#include <cstdint>
#include <numeric>

using namespace mlir;

static bool isMemRefTypeLegalForEmitC(MemRefType memRefType) {
  return memRefType.hasStaticShape() && memRefType.getLayout().isIdentity() &&
         !llvm::is_contained(memRefType.getShape(), 0);
}

namespace {
/// Implement the interface to convert MemRef to EmitC.
struct MemRefToEmitCDialectInterface : public ConvertToEmitCPatternInterface {
  MemRefToEmitCDialectInterface(Dialect *dialect)
      : ConvertToEmitCPatternInterface(dialect) {}

  /// Hook for derived dialect interface to provide conversion patterns
  /// and mark dialect legal for the conversion target.
  void populateConvertToEmitCConversionPatterns(
      ConversionTarget &target, TypeConverter &typeConverter,
      RewritePatternSet &patterns, std::optional<bool> lowerToCpp) const final {
    populateMemRefToEmitCConversionPatterns(patterns, typeConverter);
  }
};
} // namespace

void mlir::registerConvertMemRefToEmitCInterface(DialectRegistry &registry) {
  registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
    dialect->addInterfaces<MemRefToEmitCDialectInterface>();
  });
}

//===----------------------------------------------------------------------===//
// Conversion Patterns
//===----------------------------------------------------------------------===//

namespace {
struct ConvertAlloca final : public OpConversionPattern<memref::AllocaOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(memref::AllocaOp op, OpAdaptor operands,
                  ConversionPatternRewriter &rewriter) const override {

    if (!op.getType().hasStaticShape()) {
      return rewriter.notifyMatchFailure(
          op.getLoc(), "cannot transform alloca with dynamic shape");
    }

    if (op.getAlignment().value_or(1) > 1) {
      // TODO: Allow alignment if it is not more than the natural alignment
      // of the C array.
      return rewriter.notifyMatchFailure(
          op.getLoc(), "cannot transform alloca with alignment requirement");
    }

    Type resultTy = getTypeConverter()->convertType(op.getType());
    if (!resultTy)
      return rewriter.notifyMatchFailure(op.getLoc(), "cannot convert type");

    auto noInit = emitc::OpaqueAttr::get(getContext(), "");
    // Rank-0 path
    if (op.getType().getRank() == 0) {
      auto pointerTy = dyn_cast<emitc::PointerType>(resultTy);
      assert(pointerTy && "expected rank-0 MemRef to convert to pointer");
      Type elemTy = pointerTy.getPointee();
      auto var = emitc::VariableOp::create(
          rewriter, op.getLoc(), emitc::LValueType::get(elemTy), noInit);

      auto ptr = emitc::AddressOfOp::create(rewriter, op.getLoc(), resultTy,
                                            var.getResult());

      rewriter.replaceOp(op, ptr.getResult());
      return success();
    }
    // Rank > 0 path
    rewriter.replaceOpWithNewOp<emitc::VariableOp>(op, resultTy, noInit);
    return success();
  }
};

static Value calculateMemrefTotalSizeBytes(Location loc, MemRefType memrefType,
                                           OpBuilder &builder,
                                           Type convertedElementType) {
  assert(isMemRefTypeLegalForEmitC(memrefType) &&
         "incompatible memref type for EmitC conversion");

  emitc::CallOpaqueOp elementSize = emitc::CallOpaqueOp::create(
      builder, loc, emitc::SizeTType::get(builder.getContext()),
      builder.getStringAttr("sizeof"), ValueRange{},
      ArrayAttr::get(builder.getContext(),
                     {TypeAttr::get(convertedElementType)}));

  IndexType indexType = builder.getIndexType();
  int64_t numElements = llvm::product_of(memrefType.getShape());
  emitc::ConstantOp numElementsValue = emitc::ConstantOp::create(
      builder, loc, indexType, builder.getIndexAttr(numElements));

  Type sizeTType = emitc::SizeTType::get(builder.getContext());
  emitc::MulOp totalSizeBytes = emitc::MulOp::create(
      builder, loc, sizeTType, elementSize.getResult(0), numElementsValue);

  return totalSizeBytes.getResult();
}

static emitc::AddressOfOp
createPointerFromEmitcArray(Location loc, OpBuilder &builder,
                            TypedValue<emitc::ArrayType> arrayValue) {

  emitc::ConstantOp zeroIndex = emitc::ConstantOp::create(
      builder, loc, builder.getIndexType(), builder.getIndexAttr(0));

  emitc::ArrayType arrayType = arrayValue.getType();
  llvm::SmallVector<mlir::Value> indices(arrayType.getRank(), zeroIndex);
  emitc::SubscriptOp subPtr =
      emitc::SubscriptOp::create(builder, loc, arrayValue, ValueRange(indices));
  emitc::AddressOfOp ptr = emitc::AddressOfOp::create(
      builder, loc, emitc::PointerType::get(arrayType.getElementType()),
      subPtr);

  return ptr;
}

static Value getMemRefPointer(Value v) {
  if (isa<emitc::PointerType>(v.getType()))
    return v;

  // If `v` is defined through an unrealized cast and the source of that cast
  // is `emitc.ptr`, return the pointer.
  if (auto cast = v.getDefiningOp<UnrealizedConversionCastOp>())
    if (cast.getNumOperands() == 1 &&
        isa<emitc::PointerType>(cast.getOperand(0).getType()))
      return cast.getOperand(0);
  return Value();
}

static Value computeRowMajorLinearIndex(ImplicitLocOpBuilder &builder,
                                        MemRefType memrefType,
                                        ValueRange indices) {
  ArrayRef<int64_t> shape = memrefType.getShape();

  Type idxType =
      indices.empty() ? builder.getIndexType() : indices[0].getType();

  Value linearIndex =
      indices.empty()
          ? emitc::ConstantOp::create(builder, idxType, builder.getIndexAttr(0))
          : indices[0];

  if (indices.empty())
    return linearIndex;

  for (auto [dim, idx] : llvm::zip(shape.drop_front(), indices.drop_front())) {
    Value dimSize =
        emitc::ConstantOp::create(builder, idxType, builder.getIndexAttr(dim));
    linearIndex = emitc::MulOp::create(builder, idxType, linearIndex, dimSize);
    linearIndex = emitc::AddOp::create(builder, idxType, linearIndex, idx);
  }
  return linearIndex;
}

struct ConvertAlloc final : public OpConversionPattern<memref::AllocOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(memref::AllocOp allocOp, OpAdaptor operands,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = allocOp.getLoc();
    MemRefType memrefType = allocOp.getType();
    if (!isMemRefTypeLegalForEmitC(memrefType)) {
      return rewriter.notifyMatchFailure(
          loc, "incompatible memref type for EmitC conversion");
    }

    Type sizeTType = emitc::SizeTType::get(rewriter.getContext());
    Type elementType =
        getTypeConverter()->convertType(memrefType.getElementType());
    if (!elementType) {
      return rewriter.notifyMatchFailure(
          loc, "failed to convert memref element type");
    }
    IndexType indexType = rewriter.getIndexType();
    Value totalSizeBytes =
        calculateMemrefTotalSizeBytes(loc, memrefType, rewriter, elementType);

    emitc::CallOpaqueOp allocCall;
    StringAttr allocFunctionName;
    Value alignmentValue;
    SmallVector<Value, 2> argsVec;
    if (allocOp.getAlignment()) {
      allocFunctionName = rewriter.getStringAttr(alignedAllocFunctionName);
      alignmentValue = emitc::ConstantOp::create(
          rewriter, loc, sizeTType,
          rewriter.getIntegerAttr(indexType,
                                  allocOp.getAlignment().value_or(0)));
      argsVec.push_back(alignmentValue);
    } else {
      allocFunctionName = rewriter.getStringAttr(mallocFunctionName);
    }

    argsVec.push_back(totalSizeBytes);
    ValueRange args(argsVec);

    allocCall = emitc::CallOpaqueOp::create(
        rewriter, loc,
        emitc::PointerType::get(
            emitc::OpaqueType::get(rewriter.getContext(), "void")),
        allocFunctionName, args);

    emitc::PointerType targetPointerType = emitc::PointerType::get(elementType);
    emitc::CastOp castOp = emitc::CastOp::create(
        rewriter, loc, targetPointerType, allocCall.getResult(0));

    rewriter.replaceOp(allocOp, castOp);
    return success();
  }
};

struct ConvertDealloc final : public OpConversionPattern<memref::DeallocOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(memref::DeallocOp deallocOp, OpAdaptor operands,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = deallocOp.getLoc();
    // `free` can only be emitted when the dealloc operand is recoverable as an
    // `emitc.ptr<T>`.
    Value ptr = getMemRefPointer(operands.getMemref());
    if (!ptr) {
      return rewriter.notifyMatchFailure(
          loc, "expected pointer-backed memref for EmitC deallocation");
    }

    // The allocation APIs used by MemRefToEmitC return `void *`, and `free`
    // expects that same pointer type. Deallocation therefore only needs the
    // recovered base pointer cast back to `void *` before calling `free`.
    Type opaqueVoidPtrType = emitc::PointerType::get(
        emitc::OpaqueType::get(rewriter.getContext(), "void"));
    Value freeArg =
        emitc::CastOp::create(rewriter, loc, opaqueVoidPtrType, ptr);
    emitc::CallOpaqueOp freeCall = emitc::CallOpaqueOp::create(
        rewriter, loc, TypeRange{}, rewriter.getStringAttr(freeFunctionName),
        ValueRange{freeArg});
    rewriter.replaceOp(deallocOp, freeCall.getResults());
    return success();
  }
};

struct ConvertCopy final : public OpConversionPattern<memref::CopyOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(memref::CopyOp copyOp, OpAdaptor operands,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = copyOp.getLoc();
    MemRefType srcMemrefType = cast<MemRefType>(copyOp.getSource().getType());
    MemRefType targetMemrefType =
        cast<MemRefType>(copyOp.getTarget().getType());

    if (!isMemRefTypeLegalForEmitC(srcMemrefType))
      return rewriter.notifyMatchFailure(
          loc, "incompatible source memref type for EmitC conversion");

    if (!isMemRefTypeLegalForEmitC(targetMemrefType))
      return rewriter.notifyMatchFailure(
          loc, "incompatible target memref type for EmitC conversion");

    if (srcMemrefType.getRank() == 0) {
      assert(targetMemrefType.getRank() == 0 &&
             "target must have same rank as source");
      Type elementType =
          getTypeConverter()->convertType(srcMemrefType.getElementType());
      if (!elementType)
        return rewriter.notifyMatchFailure(loc, "cannot convert element type");

      Value srcPtr = getMemRefPointer(operands.getSource());
      Value targetPtr = getMemRefPointer(operands.getTarget());
      if (!srcPtr || !targetPtr)
        return rewriter.notifyMatchFailure(loc, "expected pointer operands");

      Value zeroIndex = emitc::ConstantOp::create(
          rewriter, loc, rewriter.getIndexType(), rewriter.getIndexAttr(0));
      Value srcLValue = emitc::SubscriptOp::create(
          rewriter, loc, cast<TypedValue<emitc::PointerType>>(srcPtr),
          zeroIndex);
      Value value =
          emitc::LoadOp::create(rewriter, loc, elementType, srcLValue);

      Value targetLValue = emitc::SubscriptOp::create(
          rewriter, loc, cast<TypedValue<emitc::PointerType>>(targetPtr),
          zeroIndex);
      rewriter.replaceOpWithNewOp<emitc::AssignOp>(copyOp, targetLValue, value);
      return success();
    }

    auto srcArrayValue =
        cast<TypedValue<emitc::ArrayType>>(operands.getSource());
    emitc::AddressOfOp srcPtr =
        createPointerFromEmitcArray(loc, rewriter, srcArrayValue);

    auto targetArrayValue =
        cast<TypedValue<emitc::ArrayType>>(operands.getTarget());
    emitc::AddressOfOp targetPtr =
        createPointerFromEmitcArray(loc, rewriter, targetArrayValue);

    Type convertedElementType =
        getTypeConverter()->convertType(srcMemrefType.getElementType());
    if (!convertedElementType) {
      return rewriter.notifyMatchFailure(
          loc, "failed to convert memref element type");
    }
    Value totalSizeInBytes = calculateMemrefTotalSizeBytes(
        loc, srcMemrefType, rewriter, convertedElementType);
    emitc::CallOpaqueOp memCpyCall =
        emitc::CallOpaqueOp::create(rewriter, loc, TypeRange{}, "memcpy",
                                    ValueRange{
                                        targetPtr.getResult(),
                                        srcPtr.getResult(),
                                        totalSizeInBytes,
                                    });

    rewriter.replaceOp(copyOp, memCpyCall.getResults());

    return success();
  }
};

struct ConvertGlobal final : public OpConversionPattern<memref::GlobalOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(memref::GlobalOp op, OpAdaptor operands,
                  ConversionPatternRewriter &rewriter) const override {
    MemRefType opTy = op.getType();
    if (!op.getType().hasStaticShape()) {
      return rewriter.notifyMatchFailure(
          op.getLoc(), "cannot transform global with dynamic shape");
    }

    if (op.getAlignment().value_or(1) > 1) {
      // TODO: Extend GlobalOp to specify alignment via the `alignas` specifier.
      return rewriter.notifyMatchFailure(
          op.getLoc(), "global variable with alignment requirement is "
                       "currently not supported");
    }

    Type resultTy = getTypeConverter()->convertType(opTy);

    if (!resultTy) {
      return rewriter.notifyMatchFailure(op.getLoc(),
                                         "cannot convert result type");
    }

    SymbolTable::Visibility visibility = SymbolTable::getSymbolVisibility(op);
    if (visibility != SymbolTable::Visibility::Public &&
        visibility != SymbolTable::Visibility::Private) {
      return rewriter.notifyMatchFailure(
          op.getLoc(),
          "only public and private visibility is currently supported");
    }
    // We are explicit in specifying the linkage because the default linkage
    // for constants is different in C and C++.
    bool staticSpecifier = visibility == SymbolTable::Visibility::Private;
    bool externSpecifier = !staticSpecifier;

    Attribute initialValue = operands.getInitialValueAttr();
    if (opTy.getRank() == 0) {
      auto pointerTy = dyn_cast<emitc::PointerType>(resultTy);
      assert(pointerTy && "expected rank-0 MemRef to convert to pointer");
      resultTy = pointerTy.getPointee();
      // special case for `variable : memref<i32> = dense<-1>`
      if (std::optional<Attribute> initValueAttr = op.getInitialValue()) {
        if (auto elementsAttr = llvm::dyn_cast<ElementsAttr>(*initValueAttr)) {
          initialValue = elementsAttr.getSplatValue<Attribute>();
        }
      }
    }
    if (isa_and_present<UnitAttr>(initialValue))
      initialValue = {};

    rewriter.replaceOpWithNewOp<emitc::GlobalOp>(
        op, operands.getSymName(), /*sym_visibility=*/nullptr, resultTy,
        initialValue, externSpecifier, staticSpecifier, operands.getConstant());
    return success();
  }
};

struct ConvertGetGlobal final
    : public OpConversionPattern<memref::GetGlobalOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(memref::GetGlobalOp op, OpAdaptor operands,
                  ConversionPatternRewriter &rewriter) const override {

    MemRefType opTy = op.getType();
    Type resultTy = getTypeConverter()->convertType(opTy);

    if (!resultTy) {
      return rewriter.notifyMatchFailure(op.getLoc(),
                                         "cannot convert result type");
    }

    if (opTy.getRank() == 0) {
      auto pointerTy = dyn_cast<emitc::PointerType>(resultTy);
      assert(pointerTy && "expected rank-0 MemRef to convert to pointer");
      Type elemTy = pointerTy.getPointee();
      emitc::LValueType lvalueType = emitc::LValueType::get(elemTy);
      emitc::GetGlobalOp globalLValue = emitc::GetGlobalOp::create(
          rewriter, op.getLoc(), lvalueType, operands.getNameAttr());
      rewriter.replaceOpWithNewOp<emitc::AddressOfOp>(op, resultTy,
                                                      globalLValue);
      return success();
    }
    rewriter.replaceOpWithNewOp<emitc::GetGlobalOp>(op, resultTy,
                                                    operands.getNameAttr());
    return success();
  }
};

struct ConvertLoad final : public OpConversionPattern<memref::LoadOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(memref::LoadOp op, OpAdaptor operands,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto resultTy = getTypeConverter()->convertType(op.getType());
    if (!resultTy) {
      return rewriter.notifyMatchFailure(loc, "cannot convert type");
    }

    auto arrayValue =
        dyn_cast<TypedValue<emitc::ArrayType>>(operands.getMemref());
    Value ptr = getMemRefPointer(operands.getMemref());
    if (!ptr && arrayValue) {
      auto subscript = emitc::SubscriptOp::create(rewriter, loc, arrayValue,
                                                  operands.getIndices());

      rewriter.replaceOpWithNewOp<emitc::LoadOp>(op, resultTy, subscript);
      return success();
    }

    if (!ptr)
      return rewriter.notifyMatchFailure(loc, "expected array or pointer type");
    MemRefType opMemrefType = cast<MemRefType>(op.getMemref().getType());
    ValueRange indices = operands.getIndices();

    ImplicitLocOpBuilder b(loc, rewriter);
    Value linearIndex = computeRowMajorLinearIndex(b, opMemrefType, indices);
    auto typedPtr = cast<TypedValue<emitc::PointerType>>(ptr);
    auto subscript =
        emitc::SubscriptOp::create(rewriter, loc, typedPtr, linearIndex);

    rewriter.replaceOpWithNewOp<emitc::LoadOp>(op, resultTy, subscript);
    return success();
  }
};

struct ConvertStore final : public OpConversionPattern<memref::StoreOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(memref::StoreOp op, OpAdaptor operands,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto arrayValue =
        dyn_cast<TypedValue<emitc::ArrayType>>(operands.getMemref());
    Value ptr = getMemRefPointer(operands.getMemref());
    if (!ptr && arrayValue) {
      auto subscript = emitc::SubscriptOp::create(rewriter, loc, arrayValue,
                                                  operands.getIndices());
      rewriter.replaceOpWithNewOp<emitc::AssignOp>(op, subscript,
                                                   operands.getValue());
      return success();
    }

    if (!ptr)
      return rewriter.notifyMatchFailure(loc, "expected array or pointer type");
    MemRefType opMemrefType = cast<MemRefType>(op.getMemref().getType());
    ValueRange indices = operands.getIndices();

    ImplicitLocOpBuilder b(loc, rewriter);
    Value linearIndex = computeRowMajorLinearIndex(b, opMemrefType, indices);
    auto typedPtr = cast<TypedValue<emitc::PointerType>>(ptr);
    auto subscript =
        emitc::SubscriptOp::create(rewriter, loc, typedPtr, linearIndex);

    rewriter.replaceOpWithNewOp<emitc::AssignOp>(op, subscript,
                                                 operands.getValue());
    return success();
  }
};

} // namespace

void mlir::populateMemRefToEmitCConversionPatterns(
    RewritePatternSet &patterns, const TypeConverter &converter) {
  patterns.add<ConvertAlloca, ConvertAlloc, ConvertCopy, ConvertDealloc,
               ConvertGlobal, ConvertGetGlobal, ConvertLoad, ConvertStore>(
      converter, patterns.getContext());
}