GemverIndexExtraction.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/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 <iostream>

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

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

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

  // Recursively collects load operations and scalar values from a chain of multiplications
  void collectMulOperands(Value v, SmallVectorImpl<AffineLoadOp> &loads, SmallVectorImpl<Value> &scalars) {
    if (auto load = v.getDefiningOp<AffineLoadOp>()) {
      loads.push_back(load);
      return;
    }

    if (auto mul = v.getDefiningOp<arith::MulIOp>()) {
      collectMulOperands(mul.getLhs(), loads, scalars);
      collectMulOperands(mul.getRhs(), loads, scalars);
      return;
    }
    scalars.push_back(v);
  }

  /// Detects the gemver matrix-vector accumulation pattern:
  ///   result[i] = result[i] + scalar * matrix[...] * vector[j]
  ///
  /// Both phase 2 (x[i] += beta * A[j*N+i] * y[j]) and phase 4 (w[i] += alpha * A[i*N+j] * x[j])
  /// match this pattern. The function identifies the result vector, the matrix, the vector operand,
  /// and the scalar multiplier from the store's value chain.
  ///
  /// Returns true if the pattern is detected, and populates:
  ///   - result: the memref being accumulated into (x or w)
  ///   - matrix: the matrix memref (A)
  ///   - vector: the vector memref being reduced over (y or x)
  ///   - scalar: the scalar multiplier (beta or alpha)
  bool detectMatVecAccum(AffineForOp innerLoop,
                         Value &result,
                         Value &matrix,
                         Value &vector,
                         Value &scalar) {

    for (auto store : innerLoop.getBody()->getOps<AffineStoreOp>()) {

      result = store.getMemRef();

      auto add = store.getValue().getDefiningOp<arith::AddIOp>();
      if (!add)
        continue;

      // Detect: result[i] + (scalar * matrix * vector)
      auto loadResult = add.getLhs().getDefiningOp<AffineLoadOp>();
      if (!loadResult || loadResult.getMemRef() != result) {
        // Try with the right-hand operand
        loadResult = add.getRhs().getDefiningOp<AffineLoadOp>();
        if (!loadResult || loadResult.getMemRef() != result)
          continue;
      }

      Value mulSide = (add.getLhs() == loadResult.getResult()) ? add.getRhs() : add.getLhs();

      SmallVector<AffineLoadOp, 4> loads;
      SmallVector<Value, 4> scalars;

      collectMulOperands(mulSide, loads, scalars);

      // We expect exactly 2 loads (matrix and vector) and 1 scalar
      if (loads.size() != 2 || scalars.size() != 1)
        continue;

      scalar = scalars[0];

      // Distinguish matrix from vector by memref size
      // The matrix has size N*N, the vector has size N
      auto memref0Type = dyn_cast<MemRefType>(loads[0].getMemRef().getType());
      auto memref1Type = dyn_cast<MemRefType>(loads[1].getMemRef().getType());
      if (!memref0Type || !memref1Type)
        continue;

      if (memref0Type.getShape()[0] > memref1Type.getShape()[0]) {
        matrix = loads[0].getMemRef();
        vector = loads[1].getMemRef();
      } else {
        matrix = loads[1].getMemRef();
        vector = loads[0].getMemRef();
      }

      // Verify that result is a vector (not the same as matrix)
      auto resultType = dyn_cast<MemRefType>(result.getType());
      auto matrixType = dyn_cast<MemRefType>(matrix.getType());
      if (!resultType || !matrixType)
        continue;

      // Matrix should be larger (N*N) than vectors (N)
      if (matrixType.getShape()[0] <= resultType.getShape()[0])
        continue;

      return true;
    }

    return false;
  }


  /// Determines if the matrix access is transposed (A[j*N+i] vs A[i*N+j]).
  ///
  /// In gemver phase 2, the matrix is accessed as A[j*N+i] where j is the inner loop variable
  /// and i is the outer loop variable. In phase 4, it's A[i*N+j] (normal access).
  /// This affects how we construct the index maps for the blocked transformation.
  ///
  /// Returns true if the first index of the matrix load corresponds to the inner loop
  /// induction variable (which means the access is transposed).
  bool isTransposedAccess(AffineForOp innerLoop, Value matrix) {
    Value innerIV = innerLoop.getInductionVar();

    for (auto load : innerLoop.getBody()->getOps<AffineLoadOp>()) {
      if (load.getMemRef() != matrix)
        continue;

      // Verify that the load has 2 indices (column and row)
      if (load.getIndices().size() != 2)
        continue;

      // The second index is the row (multiplied by stride in the affine map)
      // If the row index is the inner loop IV, the access is transposed
      if (load.getIndices()[1] == innerIV)
        return true;

      return false;
    }
    return false;
  }

  /// Transforms a gemver matrix-vector multiplication phase by extracting indices
  /// and blocking rows for parallelization.
  ///
  /// This function handles both phase 2 and phase 4 of gemver:
  ///   Phase 2: x[i] += beta * A^T[i,j] * y[j]  (transposed access)
  ///   Phase 4: w[i] += alpha * A[i,j] * x[j]    (normal access)
  ///
  /// It divides the rows into blocks of 3 to enable parallel processing (matching the CGRA 4x4
  /// architecture: 1 shared vector entry * 3 row entries from A = 4 total). For each block,
  /// it processes 3 rows simultaneously using loop-carried accumulators. After the main blocks,
  /// it handles remainder rows when the total rows are not divisible by 3.
  void transformMatVec(AffineForOp innerLoop,
                       Value result,
                       Value matrix,
                       Value vector,
                       Value scalar,
                       bool transposed) {

    auto outerLoop = innerLoop->getParentOfType<AffineForOp>();

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

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

    auto N = outerLoop.getConstantUpperBound();
    auto reductionSize = innerLoop.getConstantUpperBound();

    int numberEntries = 4;
    int blockSize = numberEntries - 1;
    int numRowBlocks = N / blockSize;
    int restRowBlocks = N % blockSize;

    AffineExpr d0 = getAffineDimExpr(0, ctx);
    AffineExpr d1 = getAffineDimExpr(1, ctx);

    // Linear index map for the matrix: d0 * N + d1
    // Transposed vs normal is controlled by the order of operands passed to the map
    AffineMap matIndexMap = AffineMap::get(2, 0, d0 * N + d1, ctx);

    // --- MAIN BLOCKED LOOP ---

    auto rowBlockLoop = builder.create<AffineForOp>(loc, 0, numRowBlocks);
    builder.setInsertionPointToStart(rowBlockLoop.getBody());
    Value rowBlock = rowBlockLoop.getInductionVar();

    AffineMap mapEntry1Row = AffineMap::get(1, 0, d0 * blockSize + 0, ctx);
    AffineMap mapEntry2Row = AffineMap::get(1, 0, d0 * blockSize + 1, ctx);
    AffineMap mapEntry3Row = AffineMap::get(1, 0, d0 * blockSize + 2, ctx);

    Value entry1_row_index = builder.create<AffineApplyOp>(loc, mapEntry1Row, ValueRange{rowBlock});
    Value entry2_row_index = builder.create<AffineApplyOp>(loc, mapEntry2Row, ValueRange{rowBlock});
    Value entry3_row_index = builder.create<AffineApplyOp>(loc, mapEntry3Row, ValueRange{rowBlock});

    // Initialize accumulators with current result values: result[i]
    Value acc1Init = builder.create<AffineLoadOp>(loc, result, ValueRange{entry1_row_index});
    Value acc2Init = builder.create<AffineLoadOp>(loc, result, ValueRange{entry2_row_index});
    Value acc3Init = builder.create<AffineLoadOp>(loc, result, ValueRange{entry3_row_index});

    ValueRange initArgs = {acc1Init, acc2Init, acc3Init};

    auto reductionLoop = builder.create<AffineForOp>(
        loc,
        0,
        reductionSize,
        1,
        initArgs,
        [&](OpBuilder &b, Location loc, Value iv, ValueRange args) {
          Value acc1 = args[0];
          Value acc2 = args[1];
          Value acc3 = args[2];

          // Load the shared vector element: y[j] or x[j]
          Value vecVal = b.create<AffineLoadOp>(loc, vector, ValueRange{iv});

          // Load matrix elements for each row in the block
          // Transposed: A[j*N + row_i], Normal: A[row_i*N + j]
          Value m1, m2, m3;
          if (transposed) {
            m1 = b.create<AffineLoadOp>(loc, matrix, matIndexMap, ValueRange{iv, entry1_row_index});
            m2 = b.create<AffineLoadOp>(loc, matrix, matIndexMap, ValueRange{iv, entry2_row_index});
            m3 = b.create<AffineLoadOp>(loc, matrix, matIndexMap, ValueRange{iv, entry3_row_index});
          } else {
            m1 = b.create<AffineLoadOp>(loc, matrix, matIndexMap, ValueRange{entry1_row_index, iv});
            m2 = b.create<AffineLoadOp>(loc, matrix, matIndexMap, ValueRange{entry2_row_index, iv});
            m3 = b.create<AffineLoadOp>(loc, matrix, matIndexMap, ValueRange{entry3_row_index, iv});
          }

          // Compute: scalar * matrix[...] * vector[j]
          Value prod1 = b.create<arith::MulIOp>(loc, scalar, b.create<arith::MulIOp>(loc, m1, vecVal));
          Value newAcc1 = b.create<arith::AddIOp>(loc, acc1, prod1);

          Value prod2 = b.create<arith::MulIOp>(loc, scalar, b.create<arith::MulIOp>(loc, m2, vecVal));
          Value newAcc2 = b.create<arith::AddIOp>(loc, acc2, prod2);

          Value prod3 = b.create<arith::MulIOp>(loc, scalar, b.create<arith::MulIOp>(loc, m3, vecVal));
          Value newAcc3 = b.create<arith::AddIOp>(loc, acc3, prod3);

          b.create<AffineYieldOp>(loc, ValueRange{newAcc1, newAcc2, newAcc3});
        });

    Value acc1Final = reductionLoop.getResult(0);
    Value acc2Final = reductionLoop.getResult(1);
    Value acc3Final = reductionLoop.getResult(2);

    builder.setInsertionPointAfter(reductionLoop);

    // Store final accumulated values back to result vector
    builder.create<AffineStoreOp>(loc, acc1Final, result, ValueRange{entry1_row_index});
    builder.create<AffineStoreOp>(loc, acc2Final, result, ValueRange{entry2_row_index});
    builder.create<AffineStoreOp>(loc, acc3Final, result, ValueRange{entry3_row_index});

    if (restRowBlocks == 0)
      return;

    builder.setInsertionPointAfter(rowBlockLoop);

    SmallVector<Value> remInitArgs;
    for (int i = 0; i < restRowBlocks; ++i) {
      Value row = builder.create<arith::ConstantIndexOp>(loc, numRowBlocks * blockSize + i);
      Value initVal = builder.create<AffineLoadOp>(loc, result, ValueRange{row});
      remInitArgs.push_back(initVal);
    }

    auto remReductionLoop = builder.create<AffineForOp>(
        loc,
        0,
        reductionSize,
        1,
        remInitArgs,
        [&](OpBuilder &b, Location loc, Value iv, ValueRange args) {
          SmallVector<Value> newAccs;

          Value vecVal = b.create<AffineLoadOp>(loc, vector, ValueRange{iv});

          for (int i = 0; i < restRowBlocks; ++i) {
            Value acc = args[i];
            Value row = b.create<arith::ConstantIndexOp>(loc, numRowBlocks * blockSize + i);

            Value mVal;
            if (transposed) {
              mVal = b.create<AffineLoadOp>(loc, matrix, matIndexMap, ValueRange{iv, row});
            } else {
              mVal = b.create<AffineLoadOp>(loc, matrix, matIndexMap, ValueRange{row, iv});
            }

            Value prod = b.create<arith::MulIOp>(loc, scalar, b.create<arith::MulIOp>(loc, mVal, vecVal));
            Value newAcc = b.create<arith::AddIOp>(loc, acc, prod);
            newAccs.push_back(newAcc);
          }
          b.create<AffineYieldOp>(loc, newAccs);
        });

    SmallVector<Value> finalAccs;
    for (int i = 0; i < restRowBlocks; ++i)
      finalAccs.push_back(remReductionLoop.getResult(i));

    for (int i = 0; i < restRowBlocks; ++i) {
      Value row = builder.create<arith::ConstantIndexOp>(loc, numRowBlocks * blockSize + i);
      builder.create<AffineStoreOp>(loc, finalAccs[i], result, ValueRange{row});
    }

    builder.clearInsertionPoint();
  }

  /// Gemver is a linear algebra kernel that combines several operations on a matrix A
  /// and multiple vectors. The algorithm performs:
  ///   1. Updates A by adding two outer products: A = A + u1*v1^T + u2*v2^T
  ///   2. Computes x by multiplying the transpose of A by vector y: x = beta * A^T * y
  ///   3. Adds vector z to x: x = x + z
  ///   4. Computes w by multiplying A by x: w = alpha * A * x
  ///
  /// This pass only transforms phases 2 and 4, since both are matrix-vector multiplications
  /// with the pattern: result[i] += scalar * A[...] * vec[j]
  ///
  /// For each one, it groups rows into blocks of 3 for CGRA parallelization.
  void runOnOperation() override {
    ModuleOp module = getOperation();

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

      SmallVector<AffineForOp, 4> loopsToErase;

      func.walk([&](AffineForOp loop) {

        // Skip loops that contain nested loops (only process innermost loops)
        if (llvm::any_of(loop.getBody()->getOperations(),
                         [](Operation &op) { return isa<AffineForOp>(op); }))
          return;

        Value result, matrix, vector, scalar;

        if (!detectMatVecAccum(loop, result, matrix, vector, scalar))
          return;

        auto outerLoop = loop->getParentOfType<AffineForOp>();
        if (!outerLoop) return;

        // Determine if matrix access is transposed (phase 2) or normal (phase 4)
        bool transposed = isTransposedAccess(loop, matrix);

        transformMatVec(loop, result, matrix, vector, scalar, transposed);
        loopsToErase.push_back(outerLoop);
      });

      for (auto loop : loopsToErase) {
        if (loop)
          loop.erase();
      }
    }
  }

  };
}

std::unique_ptr<mlir::Pass> cei::gemverIndexExtractionPass() {
  return std::make_unique<GemverIndexExtractionPass>();
}