LowerAffineAtomicRmwPass.cpp

Go to the documentation of this file.

//===------------------------------------------------------------------------ //
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements a pass to lower custom ops generated by the Enzyme AD
// procedure to the MemRef dialect.
//===----------------------------------------------------------------------===//
 
#include "Dialect/Ops.h"
#include "Passes/Passes.h"
#include "Passes/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
 
namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_LOWERAFFINEATOMICRMWPASS
#include "Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir
 
using namespace mlir;
 
namespace {
struct LowerAffineAtomicRmwPass
    : public enzyme::impl::LowerAffineAtomicRmwPassBase<
          LowerAffineAtomicRmwPass> {
  void runOnOperation() override {
    getOperation()->walk([&](enzyme::AffineAtomicRMWOp rmw) {
      OpBuilder builder(rmw);
      SmallVector<Value> indices;
      enzyme::computeAffineIndices(builder, rmw.getLoc(), rmw.getMap(),
                                   rmw.getIndices(), indices);
      // The enzyme atomic rather than the memref one: it is what carries the
      // ordering and the alignment this op holds. The kind is the one the op
      // names, which is not always an accumulation.
      auto lowered = enzyme::AtomicRMWOp::create(
          builder, rmw.getLoc(), rmw.getResult().getType(), rmw.getKind(),
          rmw.getOrdering(), rmw.getValue(), rmw.getMemref(), indices,
          rmw.getAlignmentAttr(), rmw.getFastmath());
      rmw.getResult().replaceAllUsesWith(lowered.getResult());
      rmw->erase();
    });
  };
};
} // end anonymous namespace