AIENormalizeAddressSpaces.cpp

Go to the documentation of this file.

//===- AIENormalizeAddressSpaces.cpp ----------------------------*- C++ -*-===//
//
// Copyright (C) 2021-2022 Xilinx, Inc.
// Copyright (C) 2022-2026 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#include "aie/Dialect/AIE/IR/AIEDialect.h"
#include "aie/Dialect/AIE/Transforms/AIEPasses.h"
 
#include "mlir/Dialect/Ptr/IR/PtrOps.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
 
#define DEBUG_TYPE "aie-normalize-address-spaces"
 
using namespace mlir;
using namespace xilinx;
using namespace xilinx::AIE;
 
static Type memRefToDefaultAddressSpace(Type t) {
  if (auto memRefType = llvm::dyn_cast<MemRefType>(t);
      memRefType && memRefType.getMemorySpace() != nullptr) {
    // Preserve ptr::GenericSpaceAttr - it's needed for ptr dialect
    // transformations
    if (llvm::isa<ptr::GenericSpaceAttr>(memRefType.getMemorySpace()))
      return t; // Keep as-is
 
    // Remove other memory spaces
    return MemRefType::get(memRefType.getShape(), memRefType.getElementType(),
                           memRefType.getLayout(), nullptr /* Address Space */);
  }
  return t;
}
 
#include "aie/Dialect/AIE/Transforms/AIENormalizeAddressSpaces.inc"
 
namespace xilinx::AIE {
#define GEN_PASS_DEF_AIENORMALIZEADDRESSSPACES
#include "aie/Dialect/AIE/Transforms/AIEPasses.h.inc"
} // namespace xilinx::AIE
 
struct AIENormalizeAddressSpacesPass
    : xilinx::AIE::impl::AIENormalizeAddressSpacesBase<
          AIENormalizeAddressSpacesPass> {
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<func::FuncDialect>();
  }
  void runOnOperation() override {
    DeviceOp device = getOperation();
 
    TypeConverter converter;
    converter.addConversion([&](Type type) -> std::optional<Type> {
      return memRefToDefaultAddressSpace(type);
    });
 
    ConversionTarget target(getContext());
    target.addDynamicallyLegalOp<memref::GlobalOp>([](memref::GlobalOp op) {
      return op.getType().getMemorySpace() == nullptr;
    });
 
    target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
      return converter.isSignatureLegal(op.getFunctionType());
    });
 
    RewritePatternSet patterns(&getContext());
    populateWithGenerated(patterns);
    populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                   converter);
 
    if (failed(applyPartialConversion(device, target, std::move(patterns))))
      signalPassFailure();
 
    // Convert any output types to have the default address space
    device.walk([&](Operation *op) {
      for (Value r : op->getResults())
        r.setType(memRefToDefaultAddressSpace(r.getType()));
    });
  }
};
 
std::unique_ptr<OperationPass<DeviceOp>>
AIE::createAIENormalizeAddressSpacesPass() {
  return std::make_unique<AIENormalizeAddressSpacesPass>();
}