ThreeMMIndexExtraction.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 ThreeMMIndexExtractionPass
    : public mlir::PassWrapper<ThreeMMIndexExtractionPass,
                               mlir::OperationPass<mlir::ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ThreeMMIndexExtractionPass)

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

  // Recursively collects load operations and scalar values from a chain of multiplication operations
  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 when the innermost loop accumulates starting from 0 (useful for first operand detection)
  // Checks the loop's iter operands initialization values
  bool hasZeroInit(AffineForOp loop) {
    if (loop.getNumIterOperands() == 0)
      return false;

    for (Value init : loop.getInits()) {
      auto cst = init.getDefiningOp<arith::ConstantOp>();
      if (!cst)
        continue;

      auto attr = dyn_cast<IntegerAttr>(cst.getValue());
      if (!attr)
        continue;

      if (attr.getInt() == 0)
        return true;
    }

    return false;
  }


  // Detects matrix multiplication pattern with accumulation (C += alpha * A * B)
  // Uses affine yield operations and iter args to detect the pattern
  // Note: alpha may not exist, which means it can be assumed to be 1 (pass still applicable)
  bool detectMMWithSum(AffineForOp loop,
                       Value &A,
                       Value &B,
                       Value &C,
                       Value &alpha) {
    if (loop.getNumResults() != 1)
      return false;

    auto yield = dyn_cast<AffineYieldOp>(loop.getBody()->getTerminator());
    if (!yield || yield.getNumOperands() != 1)
      return false;

    auto add = yield.getOperand(0).getDefiningOp<arith::AddIOp>();
    if (!add)
      return false;

    Value acc = loop.getRegionIterArgs()[0];

    Value other;
    if (add.getLhs() == acc)
      other = add.getRhs();
    else if (add.getRhs() == acc)
      other = add.getLhs();
    else
      return false;

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

    collectMulOperands(other, loads, scalars);

    if (loads.size() != 2)
      return false;

    A = loads[0].getMemRef();
    B = loads[1].getMemRef();

    alpha = nullptr;
    if (!scalars.empty())
      alpha = scalars[0];

    for (auto user : loop.getResult(0).getUsers()) {
      if (auto store = dyn_cast<AffineStoreOp>(user)) {
        C = store.getMemRef();
        return true;
      }
    }

    return false;
  }

  // Detects if matrix C is scaled by a beta factor in the loop's init operands (C = beta * C)
  // Checks the initialization values of the loop's iter operands for beta scaling
  Value detectBetaScaling(AffineForOp innerLoop, Value C) {
    if (isa<BlockArgument>(C))
      return Value();

    if (innerLoop.getNumIterOperands() != 1)
      return Value();

    Value init = innerLoop.getInits()[0];

    auto muli = init.getDefiningOp<arith::MulIOp>();
    if (!muli)
      return Value();

    Value lhs = muli.getLhs();
    Value rhs = muli.getRhs();

    auto load = lhs.getDefiningOp<AffineLoadOp>();
    Value beta;

    if (load && load.getMemRef() == C) {
      beta = rhs;
    } else {
      load = rhs.getDefiningOp<AffineLoadOp>();
      if (!load || load.getMemRef() != C)
        return Value();

      beta = lhs;
    }

    return beta;
  }

  /// Transforms the first matrix multiplication (TMP = alpha * A * B) by extracting indices
  /// and blocking rows for parallelization.
  ///
  /// This function divides the rows into blocks of 3 to enable parallel processing. The block
  /// size of 3 is chosen because the CGRA is 4x4, and we multiply the same column (1 entry)
  /// by 3 rows (3 entries), using 4 entries total. For each block, it processes 3 rows
  /// simultaneously using loop-carried accumulators. After processing the full blocks, it handles
  /// the remaining rows (remainder when rows % 3 != 0) separately to ensure all rows are computed.
  /// This explicit remainder handling allows safe parallelization without boundary violations.
  void transformFirstOperand(AffineForOp innerLoop,
                             Value A,
                             Value B,
                             Value TMP,
                             Value alpha) {

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

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

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

    auto rowsFirstMatrix = outerLoop.getConstantUpperBound();
    auto colsFirstMatrix = innerLoop.getConstantUpperBound();
    auto colsSecondMatrix = middleLoop.getConstantUpperBound();

    int numberEntries = 4;
    int blockSize = numberEntries - 1;
    int numRowBlocks = (rowsFirstMatrix) / blockSize;
    int restRowBlocks = (rowsFirstMatrix) % blockSize;

    llvm::errs() << "rowsFirstMatrix = " << rowsFirstMatrix << "\n";
    llvm::errs() << "numRowBlocks = " << numRowBlocks << "\n";
    llvm::errs() << "restRowBlocks = " << restRowBlocks << "\n";

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

    auto elemType = A.getType().cast<MemRefType>().getElementType();

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

    auto colSecondLoop = builder.create<AffineForOp>(loc, 0, colsSecondMatrix);
    builder.setInsertionPointToStart(colSecondLoop.getBody());
    Value colSecond = colSecondLoop.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});

    AffineMap aIndexMap = AffineMap::get(2, 0, d0 * colsFirstMatrix + d1, ctx);
    AffineMap bIndexMap = AffineMap::get(2, 0, d0 * colsSecondMatrix + d1, ctx);

    Value effectiveAlpha = alpha;

    if (!effectiveAlpha) effectiveAlpha = builder.create<arith::ConstantOp>(loc, elemType, builder.getIntegerAttr(elemType, 1));

    Value zeroAcc = builder.create<arith::ConstantOp>(loc, elemType, builder.getIntegerAttr(elemType, 0));
    Value acc1Init = zeroAcc;
    Value acc2Init = zeroAcc;
    Value acc3Init = zeroAcc;

    ValueRange initArgs = {acc1Init, acc2Init, acc3Init};

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

          Value b_value = b.create<AffineLoadOp>(loc, B, bIndexMap, ValueRange{iv, colSecond});

          Value a1 = b.create<AffineLoadOp>(loc, A, aIndexMap, ValueRange{entry1_row_index, iv});
          Value m1 = b.create<arith::MulIOp>(loc, effectiveAlpha, b.create<arith::MulIOp>(loc, b_value, a1));
          Value newAcc1 = b.create<arith::AddIOp>(loc, acc1, m1);

          Value a2 = b.create<AffineLoadOp>(loc, A, aIndexMap, ValueRange{entry2_row_index, iv});
          Value m2 = b.create<arith::MulIOp>(loc, effectiveAlpha, b.create<arith::MulIOp>(loc, b_value, a2));
          Value newAcc2 = b.create<arith::AddIOp>(loc, acc2, m2);

          Value a3 = b.create<AffineLoadOp>(loc, A, aIndexMap, ValueRange{entry3_row_index, iv});
          Value m3 = b.create<arith::MulIOp>(loc, effectiveAlpha, b.create<arith::MulIOp>(loc, b_value, a3));
          Value newAcc3 = b.create<arith::AddIOp>(loc, acc3, m3);

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

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

    AffineMap tmpIndexMap = AffineMap::get(2, 0, d0 * colsSecondMatrix + d1, ctx);

    builder.setInsertionPointAfter(colFirstLoop);

    builder.create<AffineStoreOp>(
        loc,
        acc1Final,
        TMP,
        tmpIndexMap,
        ValueRange{entry1_row_index, colSecond});

    builder.create<AffineStoreOp>(
        loc,
        acc2Final,
        TMP,
        tmpIndexMap,
        ValueRange{entry2_row_index, colSecond});

    builder.create<AffineStoreOp>(
        loc,
        acc3Final,
        TMP,
        tmpIndexMap,
        ValueRange{entry3_row_index, colSecond});

    // REST OF ROWS:

    if (restRowBlocks == 0)
      return;

    builder.setInsertionPointAfter(rowBlockLoop);
    auto remColSecondLoop = builder.create<AffineForOp>(loc, 0, colsSecondMatrix);
    builder.setInsertionPointToStart(remColSecondLoop.getBody());

    Value remColSecond = remColSecondLoop.getInductionVar();

    AffineMap remAIndexMap = AffineMap::get(2, 0, d0 * colsFirstMatrix + d1, ctx);
    AffineMap remBIndexMap = AffineMap::get(2, 0, d0 * colsSecondMatrix + d1, ctx);
    AffineMap remTmpIndexMap = AffineMap::get(2, 0, d0 * colsSecondMatrix + d1, ctx);

    Value remEffectiveAlpha = alpha;
    if (!remEffectiveAlpha)
      remEffectiveAlpha = builder.create<arith::ConstantOp>(loc, elemType, builder.getIntegerAttr(elemType, 1));

    Value remZeroAcc = builder.create<arith::ConstantOp>(loc, elemType, builder.getIntegerAttr(elemType, 0));

    SmallVector<Value> remInitArgs;
    for (int i = 0; i < restRowBlocks; ++i)
      remInitArgs.push_back(remZeroAcc);

    auto remColFirstLoop = builder.create<AffineForOp>(
        loc,
        0,
        colsFirstMatrix,
        1,
        remInitArgs,
        [&](OpBuilder &b, Location loc, Value iv, ValueRange args) {

          SmallVector<Value> newAccs;

          Value bValue = b.create<AffineLoadOp>(loc, B, remBIndexMap, ValueRange{iv, remColSecond});

          for (int i = 0; i < restRowBlocks; ++i) {

            Value acc = args[i];
            Value row = b.create<arith::ConstantIndexOp>(loc, numRowBlocks * blockSize + i);
            Value aValue = b.create<AffineLoadOp>(loc, A, remAIndexMap, ValueRange{row, iv});
            Value mul = b.create<arith::MulIOp>(loc, remEffectiveAlpha, b.create<arith::MulIOp>(loc, bValue, aValue));

            Value newAcc = b.create<arith::AddIOp>(loc, acc, mul);
            newAccs.push_back(newAcc);
          }
          b.create<AffineYieldOp>(loc, newAccs);
    });

    SmallVector<Value> finalAccs;
    for (int i = 0; i < restRowBlocks; ++i)
      finalAccs.push_back(remColFirstLoop.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], TMP, remTmpIndexMap, ValueRange{row, remColSecond});
    }

    builder.clearInsertionPoint();

  }

  /// Transforms the second matrix multiplication (D = beta * D + TMP * C) by extracting
  /// indices and blocking rows for parallelization.
  ///
  /// Similar to transformFirstOperand, this function divides rows into blocks of 3 for parallel
  /// processing. The block size of 3 is chosen because the CGRA is 4x4, and we multiply the same
  /// column (1 entry) by 3 rows (3 entries), using 4 entries total. It first applies beta scaling
  /// to D (D = beta * D) and uses that as the initial accumulator value, then performs the matrix
  /// multiplication with TMP and C. After the main blocks, it explicitly handles the remainder rows
  /// to ensure complete coverage when the total number of rows is not divisible by 3.
  void transformSecondOperand(AffineForOp innerLoop,
                              Value TMP,
                              Value C,
                              Value D,
                              Value beta) {
    auto middleLoop = innerLoop->getParentOfType<AffineForOp>();
    auto outerLoop = middleLoop->getParentOfType<AffineForOp>();

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

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

    auto rowsFirstMatrix = outerLoop.getConstantUpperBound();
    auto colsTmp = innerLoop.getConstantUpperBound();
    auto colsC = middleLoop.getConstantUpperBound();

    int numberEntries = 4;
    int blockSize = numberEntries - 1;
    int numRowBlocks = (rowsFirstMatrix) / blockSize;
    int restRowBlocks = (rowsFirstMatrix) % blockSize;

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

    AffineMap tmpIndexMap = AffineMap::get(2, 0, d0 * colsTmp + d1, ctx);
    AffineMap cIndexMap = AffineMap::get(2, 0, d0 * colsC + d1, ctx);
    AffineMap dIndexMap = AffineMap::get(2, 0, d0 * colsC + d1, ctx);

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

    auto colSecondLoop = builder.create<AffineForOp>(loc, 0, colsC);
    builder.setInsertionPointToStart(colSecondLoop.getBody());
    Value colSecond = colSecondLoop.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});

    Value dVal1 = builder.create<AffineLoadOp>(
        loc,
        D,
        dIndexMap,
        ValueRange{entry1_row_index, colSecond});

    Value acc1Init = builder.create<arith::MulIOp>(loc, dVal1, beta);
    Value dVal2 = builder.create<AffineLoadOp>(
        loc,
        D,
        dIndexMap,
        ValueRange{entry2_row_index, colSecond});

    Value acc2Init = builder.create<arith::MulIOp>(loc, dVal2, beta);
    Value dVal3 = builder.create<AffineLoadOp>(
        loc,
        D,
        dIndexMap,
        ValueRange{entry3_row_index, colSecond});

    Value acc3Init = builder.create<arith::MulIOp>(loc, dVal3, beta);

    ValueRange initArgs = {acc1Init, acc2Init, acc3Init};

  auto colFirstLoop = builder.create<AffineForOp>(
      loc,
      0,
      colsTmp,
      1,
      initArgs,
      [&](OpBuilder &b, Location loc, Value iv, ValueRange args) {

        Value acc1 = args[0];
        Value acc2 = args[1];
        Value acc3 = args[2];

        Value cValue = b.create<AffineLoadOp>(loc, C, cIndexMap, ValueRange{iv, colSecond});

        Value tmp1 = b.create<AffineLoadOp>(loc, TMP, tmpIndexMap, ValueRange{entry1_row_index, iv});
        Value m1 = b.create<arith::MulIOp>(loc, tmp1, cValue);
        Value newAcc1 = b.create<arith::AddIOp>(loc, acc1, m1);

        Value tmp2 = b.create<AffineLoadOp>(loc, TMP, tmpIndexMap, ValueRange{entry2_row_index, iv});
        Value m2 = b.create<arith::MulIOp>(loc, tmp2, cValue);
        Value newAcc2 = b.create<arith::AddIOp>(loc, acc2, m2);

        Value tmp3 = b.create<AffineLoadOp>(loc, TMP, tmpIndexMap, ValueRange{entry3_row_index, iv});
        Value m3 = b.create<arith::MulIOp>(loc, tmp3, cValue);
        Value newAcc3 = b.create<arith::AddIOp>(loc, acc3, m3);

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

    Value d1Final = colFirstLoop.getResult(0);
    Value d2Final = colFirstLoop.getResult(1);
    Value d3Final = colFirstLoop.getResult(2);

    builder.create<AffineStoreOp>(loc, d1Final, D, dIndexMap, ValueRange{entry1_row_index, colSecond});
    builder.create<AffineStoreOp>(loc, d2Final, D, dIndexMap, ValueRange{entry2_row_index, colSecond});
    builder.create<AffineStoreOp>(loc, d3Final, D, dIndexMap, ValueRange{entry3_row_index, colSecond});

    // REST OF ROWS:

    if (restRowBlocks == 0)
      return;

    builder.setInsertionPointAfter(rowBlockLoop);
    auto remColSecondLoop = builder.create<AffineForOp>(loc, 0, colsC);
    builder.setInsertionPointToStart(remColSecondLoop.getBody());

    Value remColSecond = remColSecondLoop.getInductionVar();

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

      Value accInit = builder.create<arith::MulIOp>(loc, dValue, beta);
      remInitArgs.push_back(accInit);

  }

  auto remColFirstLoop = builder.create<AffineForOp>(
    loc,
    0,
    colsTmp,
    1,
    remInitArgs,
    [&](OpBuilder &b, Location loc, Value iv, ValueRange args) {

      SmallVector<Value> newAccs;

      Value cValue = b.create<AffineLoadOp>(loc, C, cIndexMap, ValueRange{iv, remColSecond});

      for (int i = 0; i < restRowBlocks; ++i) {

            Value acc = args[i];
            Value row = b.create<arith::ConstantIndexOp>(loc, numRowBlocks * blockSize + i);
            Value tmp = b.create<AffineLoadOp>(loc, TMP, tmpIndexMap, ValueRange{row, iv});
            Value mul = b.create<arith::MulIOp>(loc, tmp, cValue);

            Value newAcc = b.create<arith::AddIOp>(loc, acc, mul);
            newAccs.push_back(newAcc);
          }
          b.create<AffineYieldOp>(loc, newAccs);
    });

    SmallVector<Value> finalAccs;
    for (int i = 0; i < restRowBlocks; ++i)
      finalAccs.push_back(remColFirstLoop.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], D, dIndexMap, ValueRange{row, remColSecond});
    }

    builder.clearInsertionPoint();
  }


  /// Main pass entry point: walks through functions to detect and transform 2MM and 3MM patterns.
  ///
  /// The 2MM (two matrix multiplications) pattern is defined in two-mm/twomm.c:
  ///   1. First multiplication: tmp = alpha * A * B (with tmp initialized to 0)
  ///   2. Second multiplication: D = beta * D + tmp * C
  ///
  /// The 3MM (three matrix multiplications) pattern is defined in three-mm/threemm.c:
  ///   1. First multiplication: E = A * B (with E initialized to 0)
  ///   2. Second multiplication: F = C * D (with F initialized to 0)
  ///   3. Third multiplication: G = E * F (with G initialized to 0)
  ///
  /// Note: 3MM applies the first transformation of 2MM three times with alpha = 1 (implicit).
  ///
  /// Analysis starts from innermost loops and collects information from parent and grandparent loops.
  /// The pass detects these patterns and transforms them by extracting indices and applying
  /// row blocking with explicit remainder handling.
  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)
        // Note: analysis starts from innermost loop, then collects information from parent and grandparent loops
        if (llvm::any_of(loop.getBody()->getOperations(),
                         [](Operation &op) { return isa<AffineForOp>(op); }))
          return;

        Value first_operand, second_operand, result, alpha;
        if (!detectMMWithSum(loop, first_operand, second_operand, result, alpha)){
          return;
        }

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

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

        if (hasZeroInit(loop)) {
          transformFirstOperand(loop, first_operand, second_operand, result, alpha);
          loopsToErase.push_back(outerLoop);
          return;
        }

        Value beta = detectBetaScaling(loop, result);

        if (beta) {
          transformSecondOperand(loop, first_operand, second_operand, result, beta);
          loopsToErase.push_back(outerLoop);
          return;
        }
      });

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

  };
}

std::unique_ptr<mlir::Pass> cei::threeMMIndexExtractionPass() {
  return std::make_unique<ThreeMMIndexExtractionPass>();
}