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

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

  /// Extracts the 2D matrix shape (rows, cols) from a 1D memref load operation.
  ///
  /// Matrices are stored in 1D memrefs with row-major layout. This function infers the 2D shape
  /// by extracting the stride (columns per row) from the affine map and calculating rows from
  /// the total size divided by columns.
  std::optional<std::pair<int64_t, int64_t>> getMatrixShape(AffineLoadOp load) {
    auto memrefType =
        dyn_cast<MemRefType>(load.getMemRef().getType());

    if (!memrefType || memrefType.getRank() != 1)
      return std::nullopt;

    auto strideOpt = extractStride(load);
    if (!strideOpt)
      return std::nullopt;

    int64_t cols = *strideOpt;
    int64_t total = memrefType.getShape()[0];

    if (cols == 0 || total % cols != 0)
      return std::nullopt;

    int64_t rows = total / cols;

    return std::make_pair(rows, cols);
  }

  /// Extracts the stride (columns per row) from an affine load's index map.
  ///
  /// For row-major 2D matrices in 1D memrefs, the access pattern is: row * stride + col
  /// This function parses the affine map to extract the constant stride value from
  /// the multiplication term.
  std::optional<int64_t> extractStride(AffineLoadOp load) {
    AffineMap map = load.getAffineMap();

    if (map.getNumResults() != 1)
      return std::nullopt;

    AffineExpr expr = map.getResult(0);

    auto add = expr.dyn_cast<AffineBinaryOpExpr>();
    if (!add || add.getKind() != AffineExprKind::Add)
      return std::nullopt;

    auto rhs = add.getRHS();

    auto mul = rhs.dyn_cast<AffineBinaryOpExpr>();
    if (!mul || mul.getKind() != AffineExprKind::Mul)
      return std::nullopt;

    auto cst = mul.getRHS().dyn_cast<AffineConstantExpr>();
    if (!cst)
      return std::nullopt;

    return cst.getValue();
  }

  /// Detects if a store operation performs beta scaling (C = beta * C).
  ///
  /// In GEMM (C = alpha * A * B + beta * C), beta scales the existing value of C before
  /// accumulation. This function detects the pattern where a store multiplies a loaded
  /// value from C by a constant beta and stores it back to the same location.
  Value detectBetaScaling(AffineStoreOp store, Value C) {
    auto muli = store.getValueToStore().getDefiningOp<arith::MulIOp>();
    if (!muli)
      return Value();

    auto lhsLoad = muli.getLhs().getDefiningOp<AffineLoadOp>();
    auto rhsLoad = muli.getRhs().getDefiningOp<AffineLoadOp>();

    auto lhsConst = muli.getLhs().getDefiningOp<arith::ConstantOp>();
    auto rhsConst = muli.getRhs().getDefiningOp<arith::ConstantOp>();

    if (lhsLoad && lhsLoad.getMemRef() == C &&
        lhsLoad.getAffineMap() == store.getAffineMap() &&
        lhsLoad.getIndices() == store.getIndices() &&
        rhsConst) {
      return muli.getRhs();
        }

    if (rhsLoad && rhsLoad.getMemRef() == C &&
        rhsLoad.getAffineMap() == store.getAffineMap() &&
        rhsLoad.getIndices() == store.getIndices() &&
        lhsConst) {
      return muli.getLhs();
        }

    return Value();
  }

  /// Detects the GEMM accumulation pattern (accum += alpha * A * B).
  ///
  /// This function identifies the core matrix multiplication accumulation in GEMM:
  ///   accum[i] = accum[i] + alpha * A[i] * B[i]
  ///
  /// It parses the IR to extract the alpha scalar, and the load operations for matrices A and B.
  /// The pattern involves nested multiplications: outer (alpha * (A * B)) and accumulation via add.
  bool detectAccumPattern(AffineStoreOp store,
                          Value accum,
                          Value &alpha,
                          Value &A,
                          Value &B) {
    auto addi = store.getValueToStore().getDefiningOp<arith::AddIOp>();
    if (!addi)
      return false;

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

    Value mulVal;

    auto lhsLoad = lhs.getDefiningOp<AffineLoadOp>();
    auto rhsLoad = rhs.getDefiningOp<AffineLoadOp>();

    if (lhsLoad && lhsLoad.getMemRef() == accum) {
      mulVal = rhs;
    } else if (rhsLoad && rhsLoad.getMemRef() == accum) {
      mulVal = lhs;
    } else {
      return false;
    }

    auto outerMul = mulVal.getDefiningOp<arith::MulIOp>();
    if (!outerMul)
      return false;

    Value outerLhs = outerMul.getLhs();
    Value outerRhs = outerMul.getRhs();

    auto lhsLoadOuter = outerLhs.getDefiningOp<AffineLoadOp>();
    auto rhsLoadOuter = outerRhs.getDefiningOp<AffineLoadOp>();

    auto lhsMulInner = outerLhs.getDefiningOp<arith::MulIOp>();
    auto rhsMulInner = outerRhs.getDefiningOp<arith::MulIOp>();

    Value innerMulVal;
    Value standaloneLoad;

    if (lhsMulInner && rhsLoadOuter) {
      innerMulVal = outerLhs;
      standaloneLoad = outerRhs;
    } else if (rhsMulInner && lhsLoadOuter) {
      innerMulVal = outerRhs;
      standaloneLoad = outerLhs;
    } else {
      return false;
    }

    auto innerMul = innerMulVal.getDefiningOp<arith::MulIOp>();
    if (!innerMul)
      return false;

    Value innerLhs = innerMul.getLhs();
    Value innerRhs = innerMul.getRhs();

    auto lhsConst = innerLhs.getDefiningOp<arith::ConstantOp>();
    auto rhsConst = innerRhs.getDefiningOp<arith::ConstantOp>();

    lhsLoad = innerLhs.getDefiningOp<AffineLoadOp>();
    rhsLoad = innerRhs.getDefiningOp<AffineLoadOp>();

    if (lhsConst && rhsLoad) {
      alpha = innerLhs;
      A = innerRhs;
    } else if (rhsConst && lhsLoad) {
      alpha = innerRhs;
      A = innerLhs;
    } else {
      return false;
    }

    B = standaloneLoad;

    return true;
  }

  /// Transforms the GEMM operation (C = alpha * A * B + beta * C) by extracting indices
  /// and blocking rows for parallelization.
  ///
  /// This function implements the full GEMM (General Matrix Multiply) transformation:
  ///   C = alpha * A * B + beta * C
  ///
  /// It divides the rows into blocks of 3 to enable parallel processing (matching the CGRA 4x4
  /// architecture: 1 column entry * 3 row entries = 4 total). For each block, it processes 3 rows
  /// simultaneously using loop-carried accumulators initialized with beta-scaled C values. After
  /// the main blocks, it handles remainder rows when the total rows are not divisible by 3.
  void transformOperand(Value loadA,
                        Value loadB,
                        Value C,
                        Value alpha,
                        Value beta,
                        AffineForOp outerLoop) {

    auto ALoad = loadA.getDefiningOp<AffineLoadOp>();
    auto BLoad = loadB.getDefiningOp<AffineLoadOp>();

    auto shapeA = getMatrixShape(ALoad);
    auto shapeB = getMatrixShape(BLoad);

    if (!shapeA || !shapeB)
      return;

    Value A = ALoad.getMemRef();
    Value B = BLoad.getMemRef();

    auto [rowsFirstMatrix, colsFirstMatrix] = *shapeA;
    auto [_, colsSecondMatrix] = *shapeB;
    auto colsC = colsSecondMatrix;

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

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

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

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

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

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

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


    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 cVal1 = builder.create<AffineLoadOp>(loc, C, cIndexMap, ValueRange{entry1_row_index, colSecond});
    Value init1 = builder.create<arith::MulIOp>(loc, cVal1, beta);

    Value cVal2 = builder.create<AffineLoadOp>(loc, C, cIndexMap, ValueRange{entry2_row_index, colSecond});
    Value init2 = builder.create<arith::MulIOp>(loc, cVal2, beta);

    Value cVal3 = builder.create<AffineLoadOp>(loc, C, cIndexMap, ValueRange{entry3_row_index, colSecond});
    Value init3 = builder.create<arith::MulIOp>(loc, cVal3, beta);

    ValueRange initArgs = {init1, init2, init3};

    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 bValue = 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, bValue, 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, bValue, 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, bValue, 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);

    builder.create<AffineStoreOp>(loc, acc1Final, C, cIndexMap, ValueRange{entry1_row_index, colSecond});
    builder.create<AffineStoreOp>(loc, acc2Final, C, cIndexMap, ValueRange{entry2_row_index, colSecond});
    builder.create<AffineStoreOp>(loc, acc3Final, C, cIndexMap, 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 cValue = builder.create<AffineLoadOp>(
          loc,
          C,
          cIndexMap,
          ValueRange{row, remColSecond});

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

    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, bIndexMap, 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 a = b.create<AffineLoadOp>(loc, A, aIndexMap, ValueRange{row, iv});
            Value mul = b.create<arith::MulIOp>(loc, effectiveAlpha, b.create<arith::MulIOp>(loc, bValue, a));

            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], C, cIndexMap, ValueRange{row, remColSecond});
    }

    builder.clearInsertionPoint();
  }

  /// Main pass entry point: detects and transforms GEMM patterns for CGRA optimization.
  ///
  /// GEMM (General Matrix Multiply) performs: C = alpha * A * B + beta * C
  ///   - alpha: scalar multiplier for the A*B product
  ///   - beta: scalar multiplier for the existing C values
  ///   - A, B: input matrices
  ///   - C: accumulator/output matrix
  ///
  /// The pass performs two-phase detection:
  ///   1. First walk: detects beta scaling operations (C = beta * C) and maps them by memref
  ///   2. Second walk: detects accumulation patterns (C += alpha * A * B) and matches with beta
  ///
  /// When both patterns are found for the same memref, it identifies a complete GEMM operation
  /// and transforms it by extracting indices and applying row blocking (blocks of 3) to match
  /// the CGRA 4x4 architecture for parallel processing.
  void runOnOperation() override {
    ModuleOp module = getOperation();
    DenseMap<Value, Value> betaScalingMap;
    SmallVector<Operation *> loopsToErase;

    module.walk([&](func::FuncOp funcOp) {

      // First walk: detect beta scaling operations
      funcOp.walk([&](AffineStoreOp store) {
        Value memref = store.getMemRef();

        if (Value beta = detectBetaScaling(store, memref)) {
          betaScalingMap[memref] = beta;
        }
      });

      // Second walk: detect accumulation patterns and match with beta
      funcOp.walk([&](AffineStoreOp store) {
        Value accum = store.getMemRef();
        Value alpha, loadA, loadB;

        if (!detectAccumPattern(store, accum, alpha, loadA, loadB))
          return;

        auto betaIt = betaScalingMap.find(accum);

        if (betaIt != betaScalingMap.end()) {
          Value beta = betaIt->second;
          auto innerLoop = store->getParentOfType<AffineForOp>();
          auto middleLoop = innerLoop->getParentOfType<AffineForOp>();
          auto outerLoop = middleLoop->getParentOfType<AffineForOp>();

          transformOperand(loadA, loadB, accum, alpha, beta, outerLoop);

          loopsToErase.push_back(outerLoop);
        }
      });
    });

    // Erase transformed loops
    for (auto loop : loopsToErase) {
      if (loop) loop->erase();
    }
  }

  };
}

std::unique_ptr<mlir::Pass> cei::geMMIndexExtractionPass() {
  return std::make_unique<GeMMIndexExtractionPass>();
}