Conv2dIndexExtraction.cpp

#include "Cei/CeiPasses.h"

#include "mlir/Dialect/Affine/Passes.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/IR/IntegerSet.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"

#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Pass/Pass.h"
#include "llvm/ADT/Sequence.h"

#include "mlir/IR/IRMapping.h"

#include "llvm/Support/raw_ostream.h"

#include <iostream>

using namespace mlir;
using namespace mlir::affine;
using namespace mlir::func;

namespace {

class Conv2dIndexExtractionPass
    : public mlir::PassWrapper<Conv2dIndexExtractionPass,
                               mlir::OperationPass<mlir::ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(Conv2dIndexExtractionPass)

  void getDependentDialects(mlir::DialectRegistry &registry) const override {
    registry.insert<mlir::affine::AffineDialect,
                    mlir::scf::SCFDialect,
                    mlir::arith::ArithDialect>();
  }

  /// Represents a term in a convolution expression: a load operation with its coefficient
  struct ConvTerm {
    AffineLoadOp load;
    int64_t coeff;
  };

  /// Represents a convolution access pattern: offset from center and its coefficient (kernel weight)
  struct ConvAccess {
    int64_t offset;
    int64_t coeff;
  };

  /// Recursively collects all load operations and their coefficients from an expression tree.
  ///
  /// This function parses arithmetic expressions (add, sub, mul) to extract load operations
  /// and accumulate their coefficients. For convolution, this identifies input accesses with
  /// their kernel weights. For example: 2*a + 3*b - c becomes three terms with coefficients
  /// 2, 3, and -1.
  static bool collectTerms(Value v,
                         int64_t coeff,
                         SmallVectorImpl<ConvTerm> &terms) {

    if (auto load = v.getDefiningOp<AffineLoadOp>()) {
      terms.push_back({load, coeff});
      return true;
    }

    if (auto addi = v.getDefiningOp<arith::AddIOp>()) {
      return collectTerms(addi.getLhs(), coeff, terms) &&
             collectTerms(addi.getRhs(), coeff, terms);
    }

    if (auto subi = v.getDefiningOp<arith::SubIOp>()) {
      return collectTerms(subi.getLhs(), coeff, terms) &&
             collectTerms(subi.getRhs(), -coeff, terms);
    }

    if (auto muli = v.getDefiningOp<arith::MulIOp>()) {

      APInt cst;

      // matchPattern + m_ConstantInt check if the Value is an arith.constant and extract its value
      if (matchPattern(muli.getLhs(), m_ConstantInt(&cst))) {

        return collectTerms(
            muli.getRhs(),
            coeff * cst.getSExtValue(), // multiply the accumulated coefficient
            terms);
      }

      if (matchPattern(muli.getRhs(), m_ConstantInt(&cst))) {

        return collectTerms(
            muli.getLhs(),
            coeff * cst.getSExtValue(), // absorb the constant into the coefficient
            terms);
      }

      return false;
    }

    return false;
  }

  /// Computes the constant offset between a load and store operation.
  ///
  /// In 2D convolution, each output element is computed from multiple surrounding input elements.
  /// The STORE operation writes to the target output position, while each LOAD operation reads
  /// from one of the surrounding input positions that contribute to this output. This function
  /// identifies which neighbor (relative position in the kernel) each load accesses.
  ///
  /// Workflow:
  ///   - Store position: the target element being computed (center of kernel application)
  ///   - Load position: one of the many surrounding input elements used in the computation
  ///   - Offset: relative position of the load from the store (identifies kernel position)
  ///
  /// Example for a 3x3 convolution with 80 columns:
  ///   Store: output[i,j] = ... → affine.store at [%j + %i * 80]
  ///   Load:  input[i,j-1]      → affine.load at [%j + %i * 80 - 1]
  ///
  ///   Affine maps:
  ///     Load:  (d0, d1) -> (d1 + d0 * 80 - 1)
  ///     Store: (d0, d1) -> (d1 + d0 * 80)
  ///     Difference: (d1 + d0*80 - 1) - (d1 + d0*80) = -1
  ///
  ///   The offset -1 means this load accesses the element to the left of the output position,
  ///   corresponding to a specific position in the convolution kernel.
  static bool getConstantOffset(AffineLoadOp load,
                              AffineStoreOp store,
                              int64_t &offset) {

    AffineMap loadMap = load.getAffineMap();
    AffineMap storeMap = store.getAffineMap();

    if (loadMap.getNumResults() != 1 || storeMap.getNumResults() != 1)
      return false;

    AffineExpr loadExpr = loadMap.getResult(0);
    AffineExpr storeExpr = storeMap.getResult(0);

    auto diff = simplifyAffineExpr(loadExpr - storeExpr,
                                   loadMap.getNumDims(),
                                   loadMap.getNumSymbols());

    auto constExpr = diff.dyn_cast<AffineConstantExpr>();

    if (!constExpr)
      return false;

    offset = constExpr.getValue();

    return true;
  }

  /// Checks if a specific offset exists in the convolution access pattern
  bool hasAccess (
      ArrayRef<ConvAccess> accesses,
      int64_t offset) {
    return llvm::any_of(
        accesses,
        [&](const ConvAccess &a) {
          return a.offset == offset;
        });
  };

  /// Retrieves the coefficient (kernel weight) for a specific offset in the convolution pattern.
  /// Returns 0 if the offset is not found.
  static int64_t getAccess(
    ArrayRef<ConvAccess> accesses,
    int64_t offset) {

    auto it = llvm::find_if(
        accesses,
        [&](const ConvAccess &a) {
          return a.offset == offset;
        });

    if (it == accesses.end())
      return 0;

    return it->coeff;
  }

  /// Transforms a 2D convolution operation by splitting it into 3 separate passes.
  ///
  /// For a 3x3 convolution kernel:
  ///   [f0 f1 f2]
  ///   [f3 f4 f5]
  ///   [f6 f7 f8]
  ///
  /// The transformation splits computation into 3 sequential loops, each processing 3 kernel
  /// elements (one row of the kernel). This matches the CGRA architecture where we can process
  /// 3 elements in parallel. The passes are:
  ///   1. Top row (f0, f1, f2): initializes output
  ///   2. Middle row (f3, f4, f5): accumulates to existing output
  ///   3. Bottom row (f6, f7, f8): final accumulation
  ///
  /// Each pass operates on interior points (1 to rows-1, 1 to columns-1) to avoid boundary checks.
  void transformConvolution(AffineForOp outerLoop,
                            Value input,
                            Value output,
                            ArrayRef<ConvAccess> accesses,
                            int64_t rows,
                            int64_t columns,
                            DenseMap<int64_t, Value> constantValues) {

    OpBuilder builder(outerLoop);
    builder.setInsertionPoint(outerLoop);

    Location loc = outerLoop.getLoc();
    MLIRContext *ctx = builder.getContext();

    AffineExpr d0, d1;
    bindDims(ctx, d0, d1);

    auto map_main_elem = AffineMap::get(2, 0, d1 + d0 * columns);


    int f0_position = -columns - 1;
    int f1_position = -columns;
    int f2_position = -columns + 1;

    int f3_position = -1;
    int f4_position = 0;
    int f5_position = 1;

    int f6_position = columns - 1;
    int f7_position = columns;
    int f8_position = columns + 1;

    auto accumulateAccess = [&](AffineForOp outerFor,
                                AffineForOp innerFor,
                                int64_t offset,
                                Value result) -> Value {

      if (!hasAccess(accesses, offset))
        return result;

      int64_t factor = getAccess(accesses, offset);

      auto loadMap =
          AffineMap::get(2, 0, d1 + d0 * columns + offset);

      Value loaded = builder.create<AffineLoadOp>(
          loc,
          input,
          loadMap,
          ValueRange{
              outerFor.getInductionVar(),
              innerFor.getInductionVar()});

      Value term = loaded;

      if (factor != 1) {

        auto it = constantValues.find(factor);
        Value factorVal = it->second;
        term = builder.create<arith::MulIOp>(loc, loaded, factorVal);
      }
      return builder.create<arith::AddIOp>(loc, result, term);
    };

    auto outerFor1 = builder.create<AffineForOp>(loc, 1, rows - 1);
    builder.setInsertionPointToStart(outerFor1.getBody());

    auto innerFor1 = builder.create<AffineForOp>(loc, 1, columns - 1);
    builder.setInsertionPointToStart(innerFor1.getBody());

    Value result1 = builder.create<arith::ConstantIntOp>(loc, 0, 32);

    result1 = accumulateAccess(outerFor1, innerFor1, f0_position, result1);
    result1 = accumulateAccess(outerFor1, innerFor1, f1_position, result1);
    result1 = accumulateAccess(outerFor1, innerFor1, f2_position, result1);

    builder.create<AffineStoreOp>(
        loc,
        result1,
        output,
        map_main_elem,
        ValueRange{
            outerFor1.getInductionVar(),
            innerFor1.getInductionVar()}
    );

    builder.setInsertionPointAfter(outerFor1);
    auto outerFor2 = builder.create<AffineForOp>(loc, 1, rows - 1);
    builder.setInsertionPointToStart(outerFor2.getBody());

    auto innerFor2 = builder.create<AffineForOp>(loc, 1, columns - 1);
    builder.setInsertionPointToStart(innerFor2.getBody());

    Value result2 = builder.create<AffineLoadOp>(
    loc,
    output,
    map_main_elem,
    ValueRange{
        outerFor2.getInductionVar(),
        innerFor2.getInductionVar()});

    result2 = accumulateAccess(outerFor2, innerFor2, f3_position, result2);
    result2 = accumulateAccess(outerFor2, innerFor2, f4_position, result2);
    result2 = accumulateAccess(outerFor2, innerFor2, f5_position, result2);

    builder.create<AffineStoreOp>(
        loc,
        result2,
        output,
        map_main_elem,
        ValueRange{
            outerFor2.getInductionVar(),
            innerFor2.getInductionVar()}
    );

    builder.setInsertionPointAfter(outerFor2);
    auto outerFor3 = builder.create<AffineForOp>(loc, 1, rows - 1);
    builder.setInsertionPointToStart(outerFor3.getBody());

    auto innerFor3 = builder.create<AffineForOp>(loc, 1, columns - 1);
    builder.setInsertionPointToStart(innerFor3.getBody());

    Value result3 = builder.create<AffineLoadOp>(
    loc,
    output,
    map_main_elem,
    ValueRange{
        outerFor3.getInductionVar(),
        innerFor3.getInductionVar()});

    result3 = accumulateAccess(outerFor3, innerFor3, f6_position, result3);
    result3 = accumulateAccess(outerFor3, innerFor3, f7_position, result3);
    result3 = accumulateAccess(outerFor3, innerFor3, f8_position, result3);

    builder.create<AffineStoreOp>(
        loc,
        result3,
        output,
        map_main_elem,
        ValueRange{
            outerFor3.getInductionVar(),
            innerFor3.getInductionVar()}
    );

  }

  /// Main pass entry point: detects and transforms 2D convolution patterns for CGRA optimization.
  ///
  /// 2D Convolution applies a kernel (small matrix) to an input image to produce an output:
  ///   output[i,j] = sum(kernel[m,n] * input[i+m, j+n]) for all kernel positions
  ///
  /// This pass detects convolution patterns by:
  ///   1. Identifying stores with expressions that sum multiple loads (kernel taps)
  ///   2. Computing the offset of each load relative to the store position
  ///   3. Extracting the coefficient (kernel weight) for each load
  ///
  /// The transformation splits the convolution into 3 sequential passes, each processing 3 kernel
  /// elements (one row), to match the CGRA architecture (3+1=4 elements fitting the 4x4 grid).
  /// This enables parallel processing of kernel rows while maintaining correct accumulation order.
  void runOnOperation() override {

    ModuleOp module = getOperation();

    for (func::FuncOp func : module.getOps<func::FuncOp>()) {

      // Collect all constant values for reuse (kernel coefficients)
      DenseMap<int64_t, Value> constantValues;

      module.walk([&](arith::ConstantOp cst) {
      if (auto intAttr = cst.getValue().dyn_cast<IntegerAttr>()) {
        int64_t v = intAttr.getInt();
        constantValues[v] = cst.getResult();
      }
    });

      SmallVector<AffineForOp, 4> loopsToErase;

      func.walk([&](AffineStoreOp store) {

        // Collect all terms (loads with coefficients) from the store expression
        SmallVector<ConvTerm> terms;

        if (!collectTerms(store.getValueToStore(), 1, terms))
          return;

        // Build access pattern: offset and coefficient for each kernel tap
        SmallVector<ConvAccess> accesses;

        Value input;

        for (auto &t : terms) {
          int64_t offset;

          if (!getConstantOffset(t.load, store, offset))
            return;

          accesses.push_back({offset, t.coeff});

          if (!input)
            input = t.load.getMemRef();
        }

        auto innerLoop = store->getParentOfType<AffineForOp>();
        if (!innerLoop)
          return;

        auto outerCandidate = innerLoop->getParentOfType<AffineForOp>();
        if (!outerCandidate)
          return;

        AffineForOp outerLoop = outerCandidate;

        auto rowsConst = outerLoop.getUpperBoundMap().getResult(0).dyn_cast<AffineConstantExpr>();
        auto colsConst = innerLoop.getUpperBoundMap().getResult(0).dyn_cast<AffineConstantExpr>();

        if (!rowsConst || !colsConst)
          return;

        int64_t rows = rowsConst.getValue() + 1;
        int64_t columns = colsConst.getValue() + 1;

        Value output = store.getMemRef();


        OpBuilder builder(module.getContext());
        builder.setInsertionPoint(outerLoop);

        // Ensure all kernel coefficients exist as constants (create missing ones like -1)
        for (const ConvAccess &acc : accesses) {
          int64_t coeff = acc.coeff;

          if (constantValues.count(coeff))
            continue;
          Value constVal = builder.create<arith::ConstantIntOp>(
              store.getLoc(),
              coeff,
              32
          );
          constantValues[coeff] = constVal;
        }


        transformConvolution(
            outerLoop,
            input,
            output,
            accesses,
            rows,
            columns,
            constantValues);

        loopsToErase.push_back(outerLoop);
      });

      // Erase original convolution loops after transformation
      for (auto loop : loopsToErase) {
        if (loop)
          loop.erase();
      }
    }
  }

};

} // namespace

std::unique_ptr<mlir::Pass> cei::conv2dIndexExtractionPass() {
  return std::make_unique<Conv2dIndexExtractionPass>();
}