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);
      rmw.getResult().replaceAllUsesWith(
          memref::AtomicRMWOp::create(builder, rmw.getLoc(),
                                      arith::AtomicRMWKind::addf,
                                      rmw.getValue(), rmw.getMemref(), indices)
              .getResult());
      rmw->erase();
    });
  };
};
} // end anonymous namespace