ReluIndexExtraction.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 "llvm/Support/raw_ostream.h"

#include <iostream>

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

namespace {

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

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

  bool isReluInAffineFor(AffineForOp loop, Value &memrefIn, Value &memrefOut) {

    Value iv = loop.getInductionVar();
    bool found = false;

    loop.getBody()->walk([&](scf::IfOp scfIf) {
      if (found)
        return;

      // Verify that the scf if condition is given by an arith.cmpi sgt
      auto cmp = scfIf.getCondition().getDefiningOp<arith::CmpIOp>();
      if (!cmp || cmp.getPredicate() != arith::CmpIPredicate::sgt)
        return;

      // Verify that the comparison is done on a value loaded with affine.load
      auto load = cmp.getLhs().getDefiningOp<AffineLoadOp>();
      if (!load)
        return;

      // Verify that the load is done on a 1D vector at the induction variable position
      if (load.getIndices().size() != 1 || load.getIndices()[0] != iv)
        return;

      Value inputMemref = load.getMemRef();

      // Verify that the right-hand side of the comparison is a constant equal to 0
      auto rhs = cmp.getRhs().getDefiningOp<arith::ConstantOp>();
      if (!rhs)
        return;
      auto rhsAttr = rhs.getValue().dyn_cast<IntegerAttr>();
      if (!rhsAttr || !rhsAttr.getValue().isZero())
        return;

      // Verify that the scf if has an else block
      if (!scfIf.elseBlock())
        return;

      // Verify that the then block stores the loaded value into the output
      AffineStoreOp thenStore = nullptr;
      scfIf.thenBlock()->walk([&](AffineStoreOp store) {
        thenStore = store;
      });
      if (!thenStore)
        return;

      if (thenStore.getValue() != load.getResult())
        return;

      // Verify that the store is done on a 1D vector at the induction variable position
      if (thenStore.getIndices().size() != 1 || thenStore.getIndices()[0] != iv)
        return;

      Value outputMemref = thenStore.getMemRef();

      // Verify that the else block stores a constant equal to 0 into the output
      AffineStoreOp elseStore = nullptr;
      scfIf.elseBlock()->walk([&](AffineStoreOp store) {
        elseStore = store;
      });
      if (!elseStore)
        return;

      auto elseConstant = elseStore.getValue().getDefiningOp<arith::ConstantOp>();
      if (!elseConstant)
        return;
      auto elseAttr = elseConstant.getValue().dyn_cast<IntegerAttr>();
      if (!elseAttr || !elseAttr.getValue().isZero())
        return;

      if (elseStore.getIndices().size() != 1 || elseStore.getIndices()[0] != iv)
        return;

      if (elseStore.getMemRef() != outputMemref)
        return;

      memrefIn = inputMemref;
      memrefOut = outputMemref;
      found = true;
    });

    return found;
  }

  void runOnOperation() override {
    ModuleOp module = getOperation();

    for (FuncOp func : module.getOps<FuncOp>()) {
      SmallVector<AffineForOp, 1> loopsToErase;

      for (AffineForOp loop : func.getOps<AffineForOp>()) {
        Value memrefIn, memrefOut;

        if (!isReluInAffineFor(loop, memrefIn, memrefOut))
          continue;

        auto memrefType = memrefIn.getType().dyn_cast<MemRefType>();
        if (!memrefType || memrefType.getRank() != 1)
          continue;

        auto numberElements = memrefType.getShape()[0];
        constexpr int numberEntries = 4;

        int fullBlocks = numberElements / numberEntries;
        int remainder = numberElements % numberEntries;

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

        AffineExpr d0 = getAffineDimExpr(0, ctx);
        Value zero = builder.create<arith::ConstantIntOp>(loc, 0, memrefType.getElementType());

        auto blockLoop = builder.create<AffineForOp>(loc, 0, fullBlocks);
        builder.setInsertionPointToStart(blockLoop.getBody());

        Value block = blockLoop.getInductionVar();

        SmallVector<Value, 4> entries;
        for (int i = 0; i < numberEntries; ++i) {
          AffineMap map = AffineMap::get(1, 0, d0 * numberEntries + i, ctx);
          entries.push_back(builder.create<AffineApplyOp>(loc, map, block));
        }

        for (Value entry : entries) {
          Value val = builder.create<AffineLoadOp>(loc, memrefIn, entry);
          Value cmp = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, val, zero);

          auto scfIf = builder.create<scf::IfOp>(loc, cmp, true);
          builder.setInsertionPointToStart(scfIf.thenBlock());
          builder.create<AffineStoreOp>(loc, val, memrefOut, entry);
          builder.setInsertionPointToStart(scfIf.elseBlock());
          builder.create<AffineStoreOp>(loc, zero, memrefOut, entry);
          builder.setInsertionPointAfter(scfIf);
        }

        if (remainder == 0) {
          loopsToErase.push_back(loop);
          continue;
        }

        builder.setInsertionPointAfter(blockLoop);

        for (int i = 0; i < remainder; ++i) {
          Value entry = builder.create<arith::ConstantIndexOp>(loc, fullBlocks * numberEntries + i);

          Value val = builder.create<AffineLoadOp>(loc, memrefIn, entry);
          Value cmp = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sgt, val, zero);

          auto scfIf = builder.create<scf::IfOp>(loc, cmp, true);
          builder.setInsertionPointToStart(scfIf.thenBlock());
          builder.create<AffineStoreOp>(loc, val, memrefOut, entry);
          builder.setInsertionPointToStart(scfIf.elseBlock());
          builder.create<AffineStoreOp>(loc, zero, memrefOut, entry);
          builder.setInsertionPointAfter(scfIf);
        }

        loopsToErase.push_back(loop);
      }

      for (AffineForOp l : loopsToErase)
        l.erase();
    }
  }
};

} // namespace

std::unique_ptr<mlir::Pass> cei::reluIndexExtractionPass() {
  return std::make_unique<ReluIndexExtractionPass>();
}