AIEVectorToPointerLoops.cpp

Go to the documentation of this file.

//===- AIEVectorToPointerLoops.cpp -----------------------------*- C++ -*-===//
//
// Copyright (C) 2025 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This pass transforms vector.load/store operations with loop-carried indices
// to use ptr dialect operations (ptr.to_ptr, ptr.ptr_add, ptr.from_ptr).
//
// Goal: Make pointer increment patterns explicit to help LLVM backend generate
// efficient post-increment addressing modes (GEP fusion).
//
// Transformation:
//   Before:
//     scf.for iter_args(%idx = %0) -> (index) {
//       %vec = vector.load %memref[%idx]
//       %next_idx = arith.addi %idx, %stride
//       scf.yield %next_idx
//     }
//
//   After:
//     %base_ptr = ptr.to_ptr %memref
//     %init_ptr = ptr.ptr_add %base_ptr, %0
//     scf.for iter_args(%ptr = %init_ptr) -> (!ptr.ptr<...>) {
//       %memref_tmp = ptr.from_ptr %ptr
//       %vec = vector.load %memref_tmp[%c0]
//       %next_ptr = ptr.ptr_add %ptr, %stride
//       scf.yield %next_ptr
//     }
//
//===----------------------------------------------------------------------===//
 
#include "aie/Dialect/AIE/IR/AIEDialect.h"
#include "aie/Dialect/AIE/Transforms/AIEPasses.h"
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Ptr/IR/PtrOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
 
#define DEBUG_TYPE "aie-vector-to-pointer-loops"
 
using namespace mlir;
using namespace xilinx;
using namespace xilinx::AIE;
 
namespace {
 
/// Check if a value is an iter_arg of an scf.for loop
static bool isLoopCarriedValue(Value val, scf::ForOp forOp) {
  auto blockArgs = forOp.getRegionIterArgs();
  return llvm::is_contained(blockArgs, val);
}
 
/// Structure to track memref base and its uses in vector ops
struct MemrefVectorAccess {
  Value memref;
  SmallVector<vector::LoadOp> loads;
  SmallVector<vector::StoreOp> stores;
  SmallVector<Value> indices; // Loop-carried indices used
};
 
/// Analyze loop to find vector load/store patterns with loop-carried indices
static bool analyzeLoopForVectorAccesses(
    scf::ForOp forOp, DenseMap<Value, MemrefVectorAccess> &memrefAccesses) {
 
  bool foundPattern = false;
 
  forOp.walk([&](vector::LoadOp loadOp) {
    Value base = loadOp.getBase();
    auto indices = loadOp.getIndices();
 
    // Only handle 1D access for now. vector.load requires a unit innermost
    // stride, so one index step of a 1D access is one element.
    if (indices.size() != 1)
      return;
 
    Value idx = indices[0];
 
    // Check if index is loop-carried
    if (isLoopCarriedValue(idx, forOp)) {
      memrefAccesses[base].memref = base;
      memrefAccesses[base].loads.push_back(loadOp);
      memrefAccesses[base].indices.push_back(idx);
      foundPattern = true;
    }
  });
 
  forOp.walk([&](vector::StoreOp storeOp) {
    Value base = storeOp.getBase();
    auto indices = storeOp.getIndices();
 
    if (indices.size() != 1)
      return;
 
    Value idx = indices[0];
 
    if (isLoopCarriedValue(idx, forOp)) {
      memrefAccesses[base].memref = base;
      memrefAccesses[base].stores.push_back(storeOp);
      memrefAccesses[base].indices.push_back(idx);
      foundPattern = true;
    }
  });
 
  return foundPattern;
}
 
/// Whether the index iter_arg at `pos` can become a pointer into `memref`.
/// The body rewrite only understands the index as the sole index of a
/// top-level vector.load/store on `memref`, as an operand of a top-level
/// arith.addi (whose result is then held to the same rules), and as the value
/// the loop yields back to it. Any other use (inside a nested loop, a scalar
/// memref.load, a comparison, the loop result) would receive a pointer where
/// it expects an index.
static bool canConvertIndexIterArg(scf::ForOp forOp, unsigned pos, Value memref,
                                   DenseSet<Value> &chain) {
  if (!forOp.isDefinedOutsideOfLoop(memref))
    return false;
  if (!forOp.getResult(pos).use_empty())
    return false;
 
  Block *body = forOp.getBody();
  Operation *yield = body->getTerminator();
  SmallVector<Value> worklist = {forOp.getRegionIterArgs()[pos]};
  chain.insert(worklist.begin(), worklist.end());
  while (!worklist.empty()) {
    Value v = worklist.pop_back_val();
    for (OpOperand &use : v.getUses()) {
      Operation *user = use.getOwner();
      if (user->getBlock() != body)
        return false;
      if (user == yield) {
        if (use.getOperandNumber() != pos)
          return false;
        continue;
      }
      if (auto load = dyn_cast<vector::LoadOp>(user)) {
        if (load.getBase() != memref || load.getIndices().size() != 1 ||
            load.getIndices()[0] != v)
          return false;
        continue;
      }
      if (auto store = dyn_cast<vector::StoreOp>(user)) {
        if (store.getBase() != memref || store.getIndices().size() != 1 ||
            store.getIndices()[0] != v || store.getValueToStore() == v)
          return false;
        continue;
      }
      if (auto addi = dyn_cast<arith::AddIOp>(user)) {
        // Only pointer + offset is supported, not pointer + pointer.
        if (chain.contains(addi.getLhs()) && chain.contains(addi.getRhs()))
          return false;
        if (chain.insert(addi.getResult()).second)
          worklist.push_back(addi.getResult());
        continue;
      }
      return false;
    }
  }
  // The yielded value becomes the next pointer, so it must be in the chain.
  return chain.contains(yield->getOperand(pos));
}
 
/// Transform an scf.for loop to use pointer iter_args
struct VectorToPointerLoopsPattern : public OpRewritePattern<scf::ForOp> {
  using OpRewritePattern<scf::ForOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(scf::ForOp forOp,
                                PatternRewriter &rewriter) const override {
 
    // Skip if loop already uses pointer iter_args (already transformed)
    for (Value iterArg : forOp.getRegionIterArgs()) {
      if (llvm::isa<ptr::PtrType>(iterArg.getType()))
        return failure(); // Already transformed
    }
 
    // Analyze the loop for vector access patterns
    DenseMap<Value, MemrefVectorAccess> memrefAccesses;
    if (!analyzeLoopForVectorAccesses(forOp, memrefAccesses))
      return failure();
 
    if (memrefAccesses.empty())
      return failure();
 
    // Step 1: Identify which iter_args are indices used in vector ops. This
    // step must not modify the IR: a pattern that returns failure after
    // creating ops keeps the greedy driver from converging.
    SmallVector<unsigned> indexIterArgPositions;
    SmallVector<Value> correspondingMemrefs;
 
    // Values derived from the iter_args accepted so far. Two chains that meet
    // (e.g. %a + %b) would add a pointer to a pointer.
    DenseSet<Value> converted;
    for (auto [idx, iterArg] : llvm::enumerate(forOp.getRegionIterArgs())) {
      Value owner;
      for (auto &[memref, access] : memrefAccesses) {
        if (!llvm::is_contained(access.indices, iterArg)) {
          continue;
        }
        // A pointer is derived from a single memref, so an index shared by
        // two memrefs cannot be converted.
        if (owner) {
          return failure();
        }
        owner = memref;
      }
      if (owner) {
        DenseSet<Value> chain;
        if (!canConvertIndexIterArg(forOp, idx, owner, chain)) {
          return failure();
        }
        for (Value v : chain) {
          if (!converted.insert(v).second) {
            return failure();
          }
        }
        indexIterArgPositions.push_back(idx);
        correspondingMemrefs.push_back(owner);
      }
    }
 
    if (indexIterArgPositions.empty())
      return failure();
 
    Location loc = forOp.getLoc();
    OpBuilder::InsertionGuard guard(rewriter);
    rewriter.setInsertionPoint(forOp);
 
    // Step 2: Convert memrefs to pointers before the loop
    DenseMap<Value, Value> memrefToPtrMap;
    DenseMap<Value, Type>
        memrefToGenericTypeMap; // Track generic-space memref types
    DenseMap<Value, Value>
        memrefToConvertedMap; // Track the converted memref value (for metadata)
    auto genericSpace = ptr::GenericSpaceAttr::get(rewriter.getContext());
 
    for (auto &[memref, access] : memrefAccesses) {
      // Get the memory space from the memref type
      auto memrefType = cast<MemRefType>(memref.getType());
      Attribute memorySpace = memrefType.getMemorySpace();
 
      Value memrefToConvert = memref;
      Type genericMemrefType = memrefType;
 
      // ptr.to_ptr requires the memref and pointer memory spaces to match,
      // so cast anything not already in generic_space, including the default
      // (attribute-less) space.
      if (!llvm::isa_and_nonnull<ptr::GenericSpaceAttr>(memorySpace)) {
        // Create new memref type with generic_space
        auto newMemrefType =
            MemRefType::get(memrefType.getShape(), memrefType.getElementType(),
                            memrefType.getLayout(), genericSpace);
 
        // Insert unrealized_conversion_cast to convert to generic_space
        auto castOp = UnrealizedConversionCastOp::create(rewriter, loc,
                                                         newMemrefType, memref);
        memrefToConvert = castOp.getResult(0);
        genericMemrefType = newMemrefType;
      }
 
      // Create pointer type with generic_space
      auto ptrType = ptr::PtrType::get(rewriter.getContext(), genericSpace);
      auto ptrOp =
          ptr::ToPtrOp::create(rewriter, loc, ptrType, memrefToConvert);
      memrefToPtrMap[memref] = ptrOp.getResult();
      memrefToGenericTypeMap[memref] = genericMemrefType;
      memrefToConvertedMap[memref] =
          memrefToConvert; // Store the converted memref
    }
 
    // Step 3: Build new init args with pointers
    SmallVector<Value> newInitArgs;
    Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
 
    for (auto [idx, initArg] : llvm::enumerate(forOp.getInitArgs())) {
      auto *it = llvm::find(indexIterArgPositions, idx);
      if (it != indexIterArgPositions.end()) {
        // This is an index iter_arg - convert to pointer
        size_t pos = std::distance(indexIterArgPositions.begin(), it);
        Value memref = correspondingMemrefs[pos];
        Value basePtr = memrefToPtrMap[memref];
 
        // Get element size in bytes
        auto memrefType = cast<MemRefType>(memref.getType());
        unsigned elementSizeBits = memrefType.getElementTypeBitWidth();
        unsigned elementSizeBytes = (elementSizeBits + 7) / 8;
 
        // Scale index by element size: byteOffset = initArg * elementSizeBytes
        Value byteOffset = initArg;
        if (elementSizeBytes != 1) {
          Value elementSize =
              arith::ConstantIndexOp::create(rewriter, loc, elementSizeBytes);
          byteOffset =
              arith::MulIOp::create(rewriter, loc, initArg, elementSize);
        }
 
        // Create: ptr.ptr_add basePtr, byteOffset
        auto initPtrOp =
            ptr::PtrAddOp::create(rewriter, loc, basePtr, byteOffset);
        newInitArgs.push_back(initPtrOp.getResult());
      } else {
        // Keep as-is
        newInitArgs.push_back(initArg);
      }
    }
 
    // Step 4: Create new loop with updated signature
    auto newForOp =
        scf::ForOp::create(rewriter, loc, forOp.getLowerBound(),
                           forOp.getUpperBound(), forOp.getStep(), newInitArgs);
    // Keep the loop's attributes: unsignedCmp and discardable ones such as
    // loop_annotation.
    newForOp->setAttrs(forOp->getAttrDictionary());
 
    // Step 5: Transform loop body (simplified - doesn't handle all cases yet)
    IRMapping mapper;
    mapper.map(forOp.getInductionVar(), newForOp.getInductionVar());
 
    for (auto [oldArg, newArg] :
         llvm::zip(forOp.getRegionIterArgs(), newForOp.getRegionIterArgs())) {
      mapper.map(oldArg, newArg);
    }
 
    // The memref each pointer points into, used to scale its increments by
    // the right element size. Holds the pointer iter_args and every
    // ptr.ptr_add derived from them.
    DenseMap<Value, Value> ptrToMemref;
    for (auto [pos, memref] :
         llvm::zip(indexIterArgPositions, correspondingMemrefs)) {
      ptrToMemref[newForOp.getRegionIterArgs()[pos]] = memref;
    }
 
    rewriter.setInsertionPointToStart(newForOp.getBody());
 
    // Clone operations with transformation (c0 already created above)
    for (Operation &op : forOp.getBody()->without_terminator()) {
      // Transform vector.load operations
      if (auto loadOp = dyn_cast<vector::LoadOp>(&op)) {
        Value idx = loadOp.getIndices()[0];
        Value mappedIdx = mapper.lookupOrDefault(idx);
 
        // Check if the index is now a pointer (was transformed)
        if (llvm::isa<ptr::PtrType>(mappedIdx.getType())) {
          // Get the generic-space memref type for this base
          Type genericType = memrefToGenericTypeMap[loadOp.getBase()];
 
          // Get metadata from the converted memref (with generic_space)
          Value convertedMemref = memrefToConvertedMap[loadOp.getBase()];
          auto metadataOp = ptr::GetMetadataOp::create(
              rewriter, loadOp.getLoc(), convertedMemref);
 
          // Transform: vector.load %memref[%ptr] -> ptr.from_ptr + vector.load
          // [...[0]]
          auto fromPtrOp =
              ptr::FromPtrOp::create(rewriter, loadOp.getLoc(), genericType,
                                     mappedIdx, metadataOp.getResult());
          auto newLoad = vector::LoadOp::create(
              rewriter, loadOp.getLoc(), loadOp.getVectorType(),
              fromPtrOp.getResult(), ValueRange{c0});
          mapper.map(loadOp.getResult(), newLoad.getResult());
          continue;
        }
      }
 
      // Transform vector.store operations
      if (auto storeOp = dyn_cast<vector::StoreOp>(&op)) {
        Value idx = storeOp.getIndices()[0];
        Value mappedIdx = mapper.lookupOrDefault(idx);
 
        if (llvm::isa<ptr::PtrType>(mappedIdx.getType())) {
          // Get the generic-space memref type for this base
          Type genericType = memrefToGenericTypeMap[storeOp.getBase()];
 
          // Get metadata from the converted memref (with generic_space)
          Value convertedMemref = memrefToConvertedMap[storeOp.getBase()];
          auto metadataOp = ptr::GetMetadataOp::create(
              rewriter, storeOp.getLoc(), convertedMemref);
 
          // Transform: vector.store %val, %memref[%ptr] -> ptr.from_ptr +
          // vector.store[0]
          auto fromPtrOp =
              ptr::FromPtrOp::create(rewriter, storeOp.getLoc(), genericType,
                                     mappedIdx, metadataOp.getResult());
          Value valueToStore =
              mapper.lookupOrDefault(storeOp.getValueToStore());
          vector::StoreOp::create(rewriter, storeOp.getLoc(), valueToStore,
                                  fromPtrOp.getResult(), ValueRange{c0});
          continue;
        }
      }
 
      // Transform arith.addi to ptr.ptr_add when operating on pointers
      if (auto addiOp = dyn_cast<arith::AddIOp>(&op)) {
        Value lhs = mapper.lookupOrDefault(addiOp.getLhs());
        Value rhs = mapper.lookupOrDefault(addiOp.getRhs());
        if (llvm::isa<ptr::PtrType>(rhs.getType())) {
          std::swap(lhs, rhs);
        }
 
        // If an operand is a pointer, convert to ptr.ptr_add
        if (llvm::isa<ptr::PtrType>(lhs.getType())) {
          Value memrefForPtr = ptrToMemref.lookup(lhs);
 
          Value byteOffset = rhs;
          if (memrefForPtr) {
            auto memrefType = cast<MemRefType>(memrefForPtr.getType());
            unsigned elementSizeBits = memrefType.getElementTypeBitWidth();
            unsigned elementSizeBytes = (elementSizeBits + 7) / 8;
 
            // Scale offset by element size: byteOffset = rhs * elementSizeBytes
            if (elementSizeBytes != 1) {
              Value elementSize = arith::ConstantIndexOp::create(
                  rewriter, addiOp.getLoc(), elementSizeBytes);
              byteOffset = arith::MulIOp::create(rewriter, addiOp.getLoc(), rhs,
                                                 elementSize);
            }
          }
 
          auto ptrAddOp =
              ptr::PtrAddOp::create(rewriter, addiOp.getLoc(), lhs, byteOffset);
          ptrToMemref[ptrAddOp.getResult()] = memrefForPtr;
          mapper.map(addiOp.getResult(), ptrAddOp.getResult());
          continue;
        }
      }
 
      // Default: clone the operation
      rewriter.clone(op, mapper);
    }
 
    // Clone yield
    auto oldYield = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
    SmallVector<Value> newYieldOperands;
    for (Value operand : oldYield.getOperands()) {
      newYieldOperands.push_back(mapper.lookupOrDefault(operand));
    }
    scf::YieldOp::create(rewriter, loc, newYieldOperands);
 
    // Step 6: Replace old loop
    rewriter.replaceOp(forOp, newForOp.getResults());
 
    // NOTE: This is a simplified implementation
    // Full version needs to properly transform vector.load/store and arith.addi
 
    return success();
  }
};
 
struct AIEVectorToPointerLoopsPass
    : public PassWrapper<AIEVectorToPointerLoopsPass, OperationPass<DeviceOp>> {
 
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AIEVectorToPointerLoopsPass)
 
  StringRef getArgument() const override {
    return "aie-vector-to-pointer-loops";
  }
 
  StringRef getDescription() const override {
    return "Transform vector.load/store with loop-carried indices to use ptr "
           "dialect";
  }
 
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<ptr::PtrDialect>();
    registry.insert<scf::SCFDialect>();
    registry.insert<vector::VectorDialect>();
    registry.insert<arith::ArithDialect>();
    registry.insert<memref::MemRefDialect>();
  }
 
  void runOnOperation() override {
    DeviceOp deviceOp = getOperation();
 
    RewritePatternSet patterns(&getContext());
    patterns.add<VectorToPointerLoopsPattern>(&getContext());
 
    // Apply patterns to the entire device
    // The pattern will match scf.for loops in aie.core regions
    if (failed(applyPatternsGreedily(deviceOp, std::move(patterns)))) {
      signalPassFailure();
    }
  }
};
 
} // namespace
 
namespace xilinx {
namespace AIE {
 
std::unique_ptr<OperationPass<DeviceOp>> createAIEVectorToPointerLoopsPass() {
  return std::make_unique<AIEVectorToPointerLoopsPass>();
}
 
} // namespace AIE
} // namespace xilinx