//===- PolygeistMem2Reg.cpp - MemRef DataFlow Optimization pass ------ -*-===//
//
// 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 forward memref stores to loads, thereby
// potentially getting rid of intermediate memref's entirely.
// TODO: In the future, similar techniques could be used to eliminate
// dead memref store's and perform more complex forwarding when support for
// SSA scalars live out of 'affine.for'/'affine.if' statements is available.
//===----------------------------------------------------------------------===//
#include "Enzyme/MLIR/Interfaces/AutoDiffOpInterface.h"
#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Dominance.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/Passes.h"
#include "src/enzyme_ad/jax/Dialect/Ops.h"
#include "src/enzyme_ad/jax/Passes/Passes.h"
#include "src/enzyme_ad/jax/Utils.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallPtrSet.h"
#include <algorithm>
#include <deque>
#include <iostream>
#include <set>

#define DEBUG_TYPE "mem2reg"

using namespace mlir;
using namespace mlir::arith;
using namespace mlir::enzymexla;
using namespace mlir::enzyme;

namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_POLYGEISTMEM2REG
#include "src/enzyme_ad/jax/Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir

// Whether two accesses name the same place, one lies wholly within the other,
// they may share bytes, or they share none.
enum class Match { Exact, Contains, Maybe, None };

bool operator<(Value lhs, Value rhs) {
  if (auto lhsBA = dyn_cast<BlockArgument>(lhs)) {
    if (auto rhsBA = dyn_cast<BlockArgument>(rhs)) {
      if (lhsBA.getOwner() != rhsBA.getOwner())
        return lhsBA.getOwner() < rhsBA.getOwner();
      else
        return lhsBA.getArgNumber() < rhsBA.getArgNumber();
    } else {
      return true;
    }
  }
  auto lhsOR = cast<OpResult>(lhs);
  if (auto rhsBA = dyn_cast<BlockArgument>(rhs)) {
    return false;
  } else {
    auto rhsOR = cast<OpResult>(rhs);
    if (lhsOR.getOwner() != rhsOR.getOwner())
      return lhsOR.getOwner() < rhsOR.getOwner();
    else
      return lhsOR.getResultNumber() < rhsOR.getResultNumber();
  }
}
class Offset {
public:
  enum class Type { Value, Index, Affine } type;
  mlir::Value val;
  size_t idx;
  AffineExpr aff;
  SmallVector<Value> dim;
  SmallVector<Value> sym;
  Offset(size_t constant) {
    idx = constant;
    type = Type::Index;
  }
  // An index counts from the start of what it is an offset into, so anything
  // below that is the expression it is instead, which takes a context to name.
  Offset(int64_t constant, MLIRContext *ctx) {
    if (constant >= 0) {
      idx = constant;
      type = Type::Index;
      return;
    }
    aff = getAffineConstantExpr(constant, ctx);
    type = Type::Affine;
  }
  Offset(AffineExpr expr, SmallVector<Value> dims, SmallVector<Value> syms) {
    if (auto cst = dyn_cast<AffineConstantExpr>(expr)) {
      if (cst.getValue() >= 0) {
        idx = cst.getValue();
        type = Type::Index;
      } else {
        aff = expr;
        type = Type::Affine;
      }
      return;
    }
    if (auto d = dyn_cast<AffineDimExpr>(expr)) {
      val = dims[d.getPosition()];
      type = Type::Value;
      return;
    }
    if (auto sy = dyn_cast<AffineSymbolExpr>(expr)) {
      val = syms[sy.getPosition()];
      type = Type::Value;
      return;
    }
    type = Type::Affine;
    aff = expr;
    dim = std::move(dims);
    sym = std::move(syms);
  }
  Offset(mlir::Value v) {
    std::optional<int64_t> constant;
    if (auto op = v.getDefiningOp<ConstantIntOp>())
      constant = op.value();
    else if (auto op = v.getDefiningOp<ConstantIndexOp>())
      constant = op.value();

    if (constant) {
      // An index counts from the start of what it is an offset into, so
      // anything below that is the expression it is instead.
      if (*constant >= 0) {
        idx = *constant;
        type = Type::Index;
      } else {
        aff = getAffineConstantExpr(*constant, v.getContext());
        type = Type::Affine;
      }
      return;
    }
    val = v;
    type = Type::Value;
  }
  Offset(AffineExpr op, unsigned numDims, unsigned numSymbols,
         mlir::OperandRange vals) {
    // An index counts from the start of what it is an offset into, so anything
    // below that is the expression it is instead, of nothing.
    if (auto opc = dyn_cast<AffineConstantExpr>(op)) {
      if (opc.getValue() >= 0) {
        idx = opc.getValue();
        type = Type::Index;
      } else {
        aff = op;
        type = Type::Affine;
      }
      return;
    }
    if (auto opd = dyn_cast<AffineDimExpr>(op)) {
      val = vals[opd.getPosition()];
      type = Type::Value;
      return;
    }
    if (auto ops = dyn_cast<AffineSymbolExpr>(op)) {
      val = vals[numDims + ops.getPosition()];
      type = Type::Value;
      return;
    }

    aff = op;
    for (unsigned i = 0; i < numDims; i++)
      dim.push_back(vals[i]);

    for (unsigned i = numDims; i < numSymbols; i++)
      sym.push_back(vals[i]);

    type = Type::Affine;
  }
  bool isZero() const { return type == Type::Index && idx == 0; }

  bool operator<(const Offset o) const {
    if (type != o.type) {
      return type < o.type;
    } else {
      switch (type) {
      case Offset::Type::Affine:
        if (aff == o.aff) {
          for (auto pair : llvm::zip(dim, o.dim)) {
            if (std::get<0>(pair) != std::get<1>(pair))
              return std::get<0>(pair).getAsOpaquePointer() <
                     std::get<1>(pair).getAsOpaquePointer();
          }
          for (auto pair : llvm::zip(sym, o.sym)) {
            if (std::get<0>(pair) != std::get<1>(pair))
              return std::get<0>(pair).getAsOpaquePointer() <
                     std::get<1>(pair).getAsOpaquePointer();
          }
          return false;
        } else
          return hash_value(aff) < hash_value(o.aff);
      case Offset::Type::Value:
        return val.getAsOpaquePointer() < o.val.getAsOpaquePointer();
      case Offset::Type::Index:
        return idx < o.idx;
      default:
        llvm_unreachable("unknown offset type");
      }
    }
  }
};

// The extent of a value, when the layout can say what it is.
static std::optional<uint64_t> typeSize(mlir::Type ty, const DataLayout &dl) {
  // Asking for the size of anything the layout cannot measure is fatal rather
  // than an error to handle, so this asks only about what it can: what
  // getDefaultTypeSizeInBits answers for, plus whatever brings its own answer.
  if (!ty ||
      !(ty.isIntOrFloat() ||
        isa<IndexType, VectorType, ComplexType, DataLayoutTypeInterface>(ty)))
    return std::nullopt;
  return dl.getTypeSize(ty);
}

// What a single index of an access counts in: bytes of the type a
// getelementptr steps through, or a dimension of a memref.
class OffsetType {
public:
  enum class Kind { Bytes, Dim } kind = Kind::Bytes;
  // Bytes: the size an index of this component multiplies.
  uint64_t stride = 0;
  // Dim: how many elements the dimension holds.
  int64_t dimSize = ShapedType::kDynamic;

  static OffsetType bytes(uint64_t stride) {
    OffsetType res;
    res.kind = Kind::Bytes;
    res.stride = stride;
    return res;
  }
  static OffsetType dim(int64_t dimSize) {
    OffsetType res;
    res.kind = Kind::Dim;
    res.dimSize = dimSize;
    return res;
  }

  bool operator<(const OffsetType &o) const {
    if (kind != o.kind)
      return kind < o.kind;
    switch (kind) {
    case Kind::Bytes:
      return stride < o.stride;
    case Kind::Dim:
      return dimSize < o.dimSize;
    }
    llvm_unreachable("Unknown offset unit");
  }
  bool operator==(const OffsetType &o) const {
    return !(*this < o) && !(o < *this);
  }
};

// Where an access lands within an allocation: one Offset per index, each with
// the unit it counts in, and the type of the value the access reads or writes.
// `unknown` is an offset which could be anywhere, the same as an Offset holding
// a value except that no value has to exist to name it -- which is what adding
// two offsets counted in units that cannot be mixed produces.
class OffsetTree {
  bool unknownOffset = false;
  std::vector<Offset> offsets;
  std::vector<OffsetType> units;
  mlir::Type base;

  void normalize() {
    if (unknownOffset) {
      offsets.clear();
      units.clear();
      return;
    }

    // Counting in dimensions of a memref and in bytes of a pointer are two
    // ways of saying where something is, and one path takes one of them.
    assert((!hasDim() || !hasByte()) &&
           "an offset counts in dimensions or in bytes, not both");

    // Naming the start of an allocation is naming no offset into it at all,
    // whatever the rank of the access that got there. Holding one form of it is
    // what lets an access at [0, 0] of a view meet the index-free access an
    // llvm.load performs.
    if (isZero()) {
      offsets.clear();
      units.clear();
      return;
    }
    // A constant lands at one byte however many indices reached it, which is
    // what lets paths of different shapes that reach it meet.
    if (auto bytes = constantBytes()) {
      offsets.assign(1, Offset((size_t)*bytes));
      units.assign(1, OffsetType::bytes(1));
      return;
    }

    // The indices run most significant first, which is what lets the product of
    // what the ones after an index cover be what that index steps over. A
    // constant displacement -- a field of a struct -- is no stride at all and
    // sits wherever the walk into the type put it.
#ifndef NDEBUG
    std::optional<uint64_t> last;
    for (auto &&[off, unit] : llvm::zip(offsets, units)) {
      if (unit.kind != OffsetType::Kind::Bytes ||
          off.type == Offset::Type::Index)
        continue;
      assert((!last || *last >= unit.stride) &&
             "byte offsets are held most significant first");
      last = unit.stride;
    }
#endif
  }

public:
  OffsetTree(mlir::Type base = {}, std::vector<Offset> offsets = {},
             std::vector<OffsetType> units = {})
      : offsets(std::move(offsets)), units(std::move(units)), base(base) {
    assert(this->offsets.size() == this->units.size());
    normalize();
  }

  // An offset which could be anywhere: it may be any other, and is none.
  static OffsetTree unknown() {
    OffsetTree res;
    res.unknownOffset = true;
    return res;
  }

  bool isUnknown() const { return unknownOffset; }
  mlir::Type getBase() const { return base; }

  bool hasDim() const {
    return llvm::any_of(units, [](const OffsetType &unit) {
      return unit.kind == OffsetType::Kind::Dim;
    });
  }
  // An offset that goes nowhere counts in no units, so it sits happily beside
  // one counted in dimensions: both name a place, and neither disagrees with
  // the other about how to get there.
  bool hasDimOrZero() const { return hasDim() || isZero(); }

  // A byte offset of no bytes is the placeholder for not moving at all, which
  // is not a way of counting anything.
  bool hasByte() const {
    return llvm::any_of(units, [](const OffsetType &unit) {
      return unit.kind == OffsetType::Kind::Bytes && unit.stride != 0;
    });
  }

  void print(llvm::raw_ostream &o) const;

  // Whether this names the start of what it is an offset into.
  bool isZero() const {
    if (unknownOffset)
      return false;
    return llvm::all_of(offsets,
                        [](const Offset &off) { return off.isZero(); });
  }

  // A single constant offset in bytes, when that is all this is.
  std::optional<uint64_t> constantBytes() const {
    if (unknownOffset)
      return std::nullopt;
    uint64_t total = 0;
    for (auto [off, unit] : llvm::zip(offsets, units)) {
      if (off.type != Offset::Type::Index ||
          unit.kind != OffsetType::Kind::Bytes)
        return std::nullopt;
      total += off.idx * unit.stride;
    }
    return total;
  }

  // The expression an offset is, over shared lists of the values everything is
  // built from: whatever it needs that is already named is reused, and what is
  // not is added, so several offsets can be put into one numbering.
  static AffineExpr asAffine(const Offset &off, SmallVector<Value> &dims,
                             SmallVector<Value> &syms, MLIRContext *ctx) {
    auto exprFor = [&](Value val, bool asSymbol) -> AffineExpr {
      auto dim = llvm::find(dims, val);
      if (dim != dims.end())
        return getAffineDimExpr(std::distance(dims.begin(), dim), ctx);
      auto sym = llvm::find(syms, val);
      if (sym != syms.end())
        return getAffineSymbolExpr(std::distance(syms.begin(), sym), ctx);
      if (asSymbol) {
        syms.push_back(val);
        return getAffineSymbolExpr(syms.size() - 1, ctx);
      }
      dims.push_back(val);
      return getAffineDimExpr(dims.size() - 1, ctx);
    };

    switch (off.type) {
    case Offset::Type::Index:
      return getAffineConstantExpr(off.idx, ctx);
    case Offset::Type::Value:
      // Nothing is known about a value an index just is, which is what a
      // symbol is for; a dimension is something an affine map may index by.
      return exprFor(off.val, /*asSymbol=*/true);
    case Offset::Type::Affine: {
      SmallVector<AffineExpr> dimMap, symMap;
      for (Value val : off.dim)
        dimMap.push_back(exprFor(val, /*asSymbol=*/false));
      for (Value val : off.sym)
        symMap.push_back(exprFor(val, /*asSymbol=*/true));
      return off.aff.replaceDimsAndSymbols(dimMap, symMap);
    }
    }
    llvm_unreachable("Unknown offset type");
  }

  // Two offsets of the same thing, counted the same way, add up: as
  // expressions of everything either of them is built from.
  static Offset addOffsets(const Offset &lhs, const Offset &rhs) {
    if (lhs.isZero())
      return rhs;
    if (rhs.isZero())
      return lhs;
    if (lhs.type == Offset::Type::Index && rhs.type == Offset::Type::Index)
      return Offset(lhs.idx + rhs.idx);

    MLIRContext *ctx = nullptr;
    for (const Offset *off : {&lhs, &rhs}) {
      if (off->type == Offset::Type::Value)
        ctx = off->val.getContext();
      else if (off->type == Offset::Type::Affine)
        ctx = off->aff.getContext();
    }

    SmallVector<Value> dims, syms;
    AffineExpr expr = asAffine(lhs, dims, syms, ctx);
    expr = expr + asAffine(rhs, dims, syms, ctx);
    return Offset(simplifyAffineExpr(expr, dims.size(), syms.size()), dims,
                  syms);
  }

  // How far from the start of what it is an offset into this lands, in bytes,
  // when every index says where it goes.
  std::optional<uint64_t> constantOffset(const DataLayout &dl) const {
    if (unknownOffset)
      return std::nullopt;
    auto element = typeSize(base, dl);
    uint64_t total = 0;
    for (auto &&[i, off] : llvm::enumerate(offsets)) {
      if (off.type != Offset::Type::Index)
        return std::nullopt;
      int64_t step = stepOf(i);
      if (step == ShapedType::kDynamic)
        return std::nullopt;
      // A dimension steps by that many of what the access reads.
      uint64_t bytes = step;
      if (units[i].kind == OffsetType::Kind::Dim) {
        if (!element)
          return std::nullopt;
        bytes = step * *element;
      }
      total += off.idx * bytes;
    }
    return total;
  }

  // Where `o` starts within what this names, when everything it names lies
  // inside and both say where they land.
  std::optional<uint64_t> containsAt(const OffsetTree &o,
                                     const DataLayout &dl) const {
    auto mine = constantOffset(dl), theirs = o.constantOffset(dl);
    auto mySize = typeSize(base, dl), oSize = typeSize(o.base, dl);
    if (!mine || !theirs || !mySize || !oSize)
      return std::nullopt;
    if (*theirs < *mine || *theirs + *oSize > *mine + *mySize)
      return std::nullopt;
    return *theirs - *mine;
  }

  // Composing this path with a further one. The further one names the value
  // that is read or written, so the result is of its type.
  OffsetTree add(const OffsetTree &o, const DataLayout &dl) const {
    if (unknownOffset || o.unknownOffset)
      return unknown();

    // This path does not move, so it lands wherever the access does.
    if (isZero())
      return o;

    // Dimensions count what the access reads, so two paths through them only
    // compose when what they count is the same size.
    if ((hasDim() || o.hasDim()) && base != o.base) {
      auto lhsSize = typeSize(base, dl), rhsSize = typeSize(o.base, dl);
      if (!lhsSize || !rhsSize || *lhsSize != *rhsSize)
        return unknown();
    }

    // The access is at the start of what this path reached, so it lands here,
    // reading what the access reads.
    if (o.isZero())
      return OffsetTree(o.base, offsets, units);

    // Moving the same way twice is moving once by the two together.
    if (units == o.units) {
      std::vector<Offset> sum;
      for (auto &&[lhs, rhs] : llvm::zip(offsets, o.offsets))
        sum.push_back(addOffsets(lhs, rhs));
      return OffsetTree(o.base, std::move(sum), units);
    }

    // Two constant amounts, however many indices reached either, are one.
    auto lhsBytes = constantBytes(), rhsBytes = o.constantBytes();
    if (lhsBytes && rhsBytes)
      return OffsetTree(o.base, {Offset((size_t)(*lhsBytes + *rhsBytes))},
                        {OffsetType::bytes(1)});

    // Otherwise a constant amount of bytes joins a path that steps by a size it
    // is a multiple of: it is that many of those steps, taken innermost, which
    // is where the smallest of them is.
    auto join = [](const OffsetTree &path, uint64_t bytes,
                   mlir::Type base) -> OffsetTree {
      if (path.units.empty() ||
          path.units.back().kind != OffsetType::Kind::Bytes)
        return unknown();
      uint64_t stride = path.units.back().stride;
      if (!stride || bytes % stride)
        return unknown();
      std::vector<Offset> offsets = path.offsets;
      offsets.back() =
          addOffsets(path.offsets.back(), Offset((size_t)(bytes / stride)));
      return OffsetTree(base, std::move(offsets), path.units);
    };
    if (rhsBytes)
      return join(*this, *rhsBytes, o.base);
    if (lhsBytes)
      return join(o, *lhsBytes, o.base);
    return unknown();
  }

  // The number of elements (if dim) or bytes moved
  // by changing the i-th index by 1
  int64_t stepOf(int64_t i) const {
    assert(i < units.size());
    if (units[i].kind == OffsetType::Kind::Bytes) {
      return units[i].stride;
    } else {
      int64_t cur = 1;
      for (int64_t j = i + 1; j < (int64_t)units.size(); j++) {
        if (units[j].dimSize == ShapedType::kDynamic) {
          return ShapedType::kDynamic;
        }
        cur *= units[j].dimSize;
      }
      return cur;
    }
  }

  // The maximum value allowed as an arg for the i-th index
  int64_t maxValue(int64_t i) const {
    if (units[i].kind == OffsetType::Kind::Bytes) {
      if (i == 0)
        return ShapedType::kDynamic;
      auto prevStep = stepOf(i - 1);
      auto curStep = stepOf(i);
      if (prevStep == ShapedType::kDynamic || curStep == ShapedType::kDynamic ||
          curStep == 0)
        return ShapedType::kDynamic;
      assert(prevStep % curStep == 0);
      assert(prevStep != curStep);
      return prevStep / curStep;
    } else {
      return units[i].dimSize;
    }
  }

  // The first index guaranteed not to be touched by this access, if any
  int64_t reachBound() const {
    if (units.empty())
      return 0;

    int64_t step = stepOf(0), extent = maxValue(0);
    assert(step != 0);
    if (step == ShapedType::kDynamic || extent == ShapedType::kDynamic)
      return ShapedType::kDynamic;
    return step * extent;
  }

  // Whether a constant index at `at` puts this past everything `o` can name.
  bool beyond(size_t at, std::optional<uint64_t> mySize, const OffsetTree &o,
              std::optional<uint64_t> oSize, bool sameExtent) const {
    int64_t step = stepOf(at), bound = o.reachBound();
    assert(step != 0);
    if (step == ShapedType::kDynamic || bound == ShapedType::kDynamic)
      return false;

    // Reading different amounts, the two only line up once both are counted in
    // bytes rather than in the elements each reads.
    if (!sameExtent) {
      if (!mySize || !oSize)
        return false;
      if (hasDim())
        step *= (int64_t)*mySize;
      if (o.hasDim())
        bound *= (int64_t)*oSize;
    }

    // A dimension counts the elements the other reads, so its bound already
    // spans what is read at the last of them. A byte offset is only where the
    // other starts, and what it reads there comes after that.
    if (!o.hasDim()) {
      if (!oSize)
        return false;
      // Unless it goes nowhere, in which case it counts in no units at all and
      // what it reaches has to be said in whatever this one counts in: a single
      // element when that is dimensions of what it reads, its own extent
      // otherwise -- which is what the step has been put into above when the
      // two read different amounts.
      bound += o.units.empty() && hasDim() && sameExtent ? 1 : (int64_t)*oSize;
    }

    uint64_t idx = offsets[at].idx;
    if (idx > (uint64_t)(INT64_MAX / step))
      return false;
    return (int64_t)idx * step >= bound;
  }

  // Whether this index or one outside it is known to move, which is what puts
  // an access past everything the indices within it can reach. The outer ones
  // are the ones before it.
  int64_t smallestStrideUpTo(size_t at) const {
    assert(at < offsets.size());

    if (!hasDim()) {
      int64_t best = INT64_MAX;
      for (size_t i = 0; i < at; i++) {
        if (offsets[i].isZero())
          continue;
        assert(units[i].kind == OffsetType::Kind::Bytes);
        if (units[i].stride < best) {
          best = units[i].stride;
        }
      }
      return best;
    } else {
      if (at == 0)
        return INT64_MAX;
      size_t cur = 1;
      for (size_t i = at; i < offsets.size(); i++) {
        if (units[i].dimSize == ShapedType::kDynamic) {
          continue;
        }
        cur *= units[i].dimSize;
      }
      return (int64_t)cur;
    }
  }

  // One index against another, given what a step of each covers, how far all of
  // each index reaches, how much each access reads, and whether either is known
  // to have moved at this index or an outer one. What the shapes leave open is
  // kDynamic, which compares to nothing.
  static Match compareIndex(const Offset &lhs, int64_t lhsStep,
                            std::optional<uint64_t> lhsSize,
                            int64_t lhsSmallestStride, const Offset &rhs,
                            int64_t rhsStep, std::optional<uint64_t> rhsSize,
                            int64_t rhsSmallestStride, bool isDim,
                            bool sameExtent) {

    // The lhs step is smaller than the rhs step, and it can be proven that the
    // lhs stride cannot possibly equal a single rhs stride.
    if (lhsStep != ShapedType::kDynamic && rhsStep != ShapedType::kDynamic &&
        lhsStep < rhsStep && rhsStep % lhsStep == 0) {
      if (lhs.type == Offset::Type::Index && lhs.idx != 0 &&
          (int64_t)lhs.idx < rhsStep / lhsStep) {
        return Match::None;
      }
    }

    // Same the other way
    if (rhsStep != ShapedType::kDynamic && lhsStep != ShapedType::kDynamic &&
        rhsStep < lhsStep && lhsStep % rhsStep == 0) {
      if (rhs.type == Offset::Type::Index && rhs.idx != 0 &&
          (int64_t)rhs.idx < lhsStep / rhsStep) {
        return Match::None;
      }
    }

    if (lhsStep != ShapedType::kDynamic && lhsStep == rhsStep &&
        (sameExtent || !isDim)) {
      int64_t difference = INT64_MAX;
      if (lhs.type != rhs.type)
        return Match::Maybe;
      switch (lhs.type) {
      case Offset::Type::Affine:
        if (lhs.dim == rhs.dim && lhs.sym == rhs.sym) {
          if (lhs.aff == rhs.aff) {
            return Match::Exact;
          }
          if (auto cst = dyn_cast<AffineConstantExpr>(rhs.aff - lhs.aff)) {
            difference = std::abs(cst.getValue());
          }
        }
        break;
      case Offset::Type::Value:
        if (lhs.val == rhs.val)
          return Match::Exact;
        break;
      case Offset::Type::Index:
        if (lhs.idx == rhs.idx) {
          return Match::Exact;
        }
        difference =
            (lhs.idx < rhs.idx ? rhs.idx - lhs.idx : lhs.idx - rhs.idx);
        break;
      }

      if (difference == INT64_MAX)
        return Match::Maybe;

      if (isDim) {
        if (difference < lhsSmallestStride && difference < rhsSmallestStride) {
          return Match::None;
        }
      } else {
        if (sameExtent || (lhsSize && rhsSize)) {
          if (!sameExtent) {
            difference +=
                *lhsSize > *rhsSize ? *lhsSize - *rhsSize : *rhsSize - *lhsSize;
          }
          if (difference < lhsSmallestStride &&
              difference < rhsSmallestStride) {
            return Match::None;
          }
        }
      }
    }

    return Match::Maybe;
  }

  Match matches(const OffsetTree &o, const DataLayout &dl,
                uint64_t *containedAt = nullptr) const {
    if (unknownOffset || o.unknownOffset)
      return Match::Maybe;

    // A different type at the same place may still be the same value, since
    // castToType converts between spellings of one extent. One type covers the
    // same bytes as itself whether or not the layout can say how many; two of
    // them do only if it can say so, and two extents of different length
    // starting together share no more than a part.
    std::optional<uint64_t> size = typeSize(base, dl);
    std::optional<uint64_t> osize =
        base == o.base ? size : typeSize(o.base, dl);
    bool sameExtent = base == o.base || (size && osize && *size == *osize);
    if (!sameExtent && (!size || !osize))
      return Match::Maybe;

    // Lying wholly within is a question of where each lands and how far it
    // reaches, which needs nothing of how either counts its way there.
    if (!sameExtent && containedAt)
      if (auto at = containsAt(o, dl)) {
        *containedAt = *at;
        return Match::Contains;
      }

    // Two offsets that both go nowhere name the same place, so what is left is
    // how far each reaches from it: the same extent is the same access, and a
    // smaller one lies inside the larger.
    if (isZero() && o.isZero()) {
      if (sameExtent)
        return Match::Exact;
      if (containedAt && size && osize && *osize <= *size) {
        *containedAt = 0;
        return Match::Contains;
      }
      return Match::Maybe;
    }

    bool isDim = true;

    if (hasDimOrZero() && o.hasDimOrZero()) {
      isDim = true;
    } else if (!hasDim() && !o.hasDim()) {
      isDim = false;
    } else {
      return Match::Maybe;
    }

    // The indices run most significant first, so the walk runs the other way:
    // what an index steps over is everything the ones after it cover, which is
    // what the running product holds by the time it is reached.
    ssize_t i = offsets.size() - 1, j = o.offsets.size() - 1;

    bool exact = sameExtent;

    while (i >= 0 && j >= 0) {
      // An index that does not move lines up with anything, though whatever it
      // would have stepped over still lies inside the ones outside it.
      if (offsets[i].isZero()) {
        i--;
        continue;
      }
      if (o.offsets[j].isZero()) {
        j--;
        continue;
      }

      auto lhsStep = stepOf(i);
      auto rhsStep = o.stepOf(j);

      auto lhsStride = smallestStrideUpTo(i);
      auto rhsStride = o.smallestStrideUpTo(j);
      auto res =
          compareIndex(offsets[i], lhsStep, size, lhsStride, o.offsets[j],
                       rhsStep, osize, rhsStride, isDim, sameExtent);
      if (res == Match::None)
        return Match::None;

      if (res == Match::Exact) {
        i--;
        j--;
        continue;
      }

      exact = false;

      if (sameExtent && lhsStep != rhsStep && lhsStep != ShapedType::kDynamic &&
          rhsStep != ShapedType::kDynamic) {
        if (lhsStep < rhsStep) {
          i--;
        } else {
          j--;
        }
        continue;
      }

      i--;
      j--;
    }

    // Whatever is left over of one path moves on from where the other stopped,
    // which is wherever the indices it did have put it.
    for (; i >= 0; i--) {
      if (offsets[i].isZero()) {
        continue;
      }
      exact = false;
      if (offsets[i].type == Offset::Type::Index &&
          beyond(i, size, o, osize, sameExtent)) {
        return Match::None;
      }
    }
    for (; j >= 0; j--) {
      if (o.offsets[j].isZero()) {
        continue;
      }
      exact = false;
      if (o.offsets[j].type == Offset::Type::Index &&
          o.beyond(j, osize, *this, size, sameExtent)) {
        return Match::None;
      }
    }

    if (exact) {
      return Match::Exact;
    }

    return Match::Maybe;
  }

  bool operator<(const OffsetTree &o) const {
    if (unknownOffset != o.unknownOffset)
      return unknownOffset < o.unknownOffset;
    if (offsets.size() != o.offsets.size())
      return offsets.size() < o.offsets.size();
    for (auto &&[i, off] : llvm::enumerate(offsets)) {
      if (off < o.offsets[i])
        return true;
      if (o.offsets[i] < off)
        return false;
      if (units[i] < o.units[i])
        return true;
      if (o.units[i] < units[i])
        return false;
    }
    return false;
  }
};

static OffsetTree accessOffsets(Type accessed) { return OffsetTree(accessed); }

// The pointer a memcpy/memmove reads from; a memset reads from nothing.
static Value transferSource(Operation *op) {
  if (auto cpy = dyn_cast<LLVM::MemcpyOp>(op))
    return cpy.getSrc();
  if (auto mv = dyn_cast<LLVM::MemmoveOp>(op))
    return mv.getSrc();
  return nullptr;
}

static Value transferDest(Operation *op) {
  if (auto cpy = dyn_cast<LLVM::MemcpyOp>(op))
    return cpy.getDst();
  if (auto mv = dyn_cast<LLVM::MemmoveOp>(op))
    return mv.getDst();
  if (auto ms = dyn_cast<LLVM::MemsetOp>(op))
    return ms.getDst();
  return nullptr;
}

// How many bytes a transfer moves, when that is a known constant.
static std::optional<uint64_t> transferLength(Operation *op) {
  Value len;
  if (auto cpy = dyn_cast<LLVM::MemcpyOp>(op))
    len = cpy.getLen();
  else if (auto mv = dyn_cast<LLVM::MemmoveOp>(op))
    len = mv.getLen();
  else if (auto ms = dyn_cast<LLVM::MemsetOp>(op))
    len = ms.getLen();
  else
    return std::nullopt;

  APInt val;
  if (!matchPattern(len, m_ConstantInt(&val)))
    return std::nullopt;
  return val.getZExtValue();
}

// What a transfer states about the alignment of the end it is asked about,
// which is all an access put in its place may assume; absent, it promises
// nothing beyond a byte. The destination is its first argument, the source its
// second.
static unsigned transferAlignment(Operation *op, unsigned arg) {
  auto argAttrs = op->getAttrOfType<ArrayAttr>("arg_attrs");
  if (!argAttrs || argAttrs.size() <= arg)
    return 1;
  auto dict = dyn_cast<DictionaryAttr>(argAttrs[arg]);
  if (!dict)
    return 1;
  if (auto align =
          dict.getAs<IntegerAttr>(LLVM::LLVMDialect::getAlignAttrName()))
    return align.getInt();
  return 1;
}

// What a transfer touches is a number of bytes rather than a value of some
// type, which is the same thing as an access of an integer that wide. Anything
// larger than a slot could hold is not worth naming.
static OffsetTree transferAccess(uint64_t bytes, MLIRContext *ctx) {
  if (!bytes || bytes > 4096)
    return OffsetTree::unknown();
  return OffsetTree(IntegerType::get(ctx, bytes * 8));
}

static OffsetTree accessOffsets(Type accessed, Value memory,
                                mlir::OperandRange indices,
                                const DataLayout &dl) {
  std::vector<Offset> offsets;
  std::vector<OffsetType> units;
  auto MT = cast<MemRefType>(memory.getType());
  assert(MT.getElementType() == accessed);
  for (auto &&[i, idx] : llvm::enumerate(indices)) {
    offsets.emplace_back(idx);
    units.push_back(OffsetType::dim(MT.getShape()[i]));
  }
  return OffsetTree(accessed, std::move(offsets), std::move(units));
}

static OffsetTree accessOffsets(Type accessed, Value memory, AffineMap map,
                                mlir::OperandRange operands,
                                const DataLayout &dl) {
  std::vector<Offset> offsets;
  std::vector<OffsetType> units;
  auto MT = cast<MemRefType>(memory.getType());
  assert(MT.getElementType() == accessed);
  for (auto &&[i, expr] : llvm::enumerate(map.getResults())) {
    offsets.emplace_back(expr, map.getNumDims(), map.getNumSymbols(), operands);
    units.push_back(OffsetType::dim(MT.getShape()[i]));
  }
  return OffsetTree(accessed, std::move(offsets), std::move(units));
}

// A getelementptr walks into its element type: the first index steps over the
// whole of it, and each one after that into what the previous one selected.
// What it lands on is only known when every index into an aggregate is.
static OffsetTree gepOffsets(LLVM::GEPOp gep, const DataLayout &dl) {
  std::vector<Offset> offsets;
  std::vector<OffsetType> units;
  Type cur = gep.getElemType();
  assert(cur && "a getelementptr walks into a type");

  auto step = [&](std::optional<int64_t> constant, Value dynamic,
                  uint64_t stride) {
    if (constant) {
      offsets.emplace_back(*constant, gep.getContext());
    } else {
      offsets.emplace_back(dynamic);
    }
    units.push_back(OffsetType::bytes(stride));
  };

  for (auto &&[i, arg] : llvm::enumerate(gep.getIndices())) {
    Value dynamic = dyn_cast<Value>(arg);
    std::optional<int64_t> constant;
    if (auto attr = dyn_cast<IntegerAttr>(arg))
      constant = attr.getInt();
    else if (APInt val; matchPattern(dynamic, m_ConstantInt(&val)))
      constant = val.getSExtValue();

    if (i == 0) {
      step(constant, dynamic, dl.getTypeSize(cur));
      continue;
    }

    if (auto ST = dyn_cast<LLVM::LLVMStructType>(cur)) {
      // Only a constant can select a field, and it lands at a fixed byte.
      if (!constant || *constant < 0 ||
          (size_t)*constant >= ST.getBody().size())
        return OffsetTree::unknown();
      uint64_t byte = 0;
      for (auto member : ST.getBody().take_front(*constant)) {
        if (!ST.isPacked())
          byte = llvm::alignTo(byte, dl.getTypeABIAlignment(member));
        byte += dl.getTypeSize(member);
      }
      if (!ST.isPacked())
        byte = llvm::alignTo(byte,
                             dl.getTypeABIAlignment(ST.getBody()[*constant]));
      offsets.emplace_back((size_t)byte);
      units.push_back(OffsetType::bytes(1));
      cur = ST.getBody()[*constant];
      assert(cur && "a field of a struct is of some type");
      continue;
    }

    Type element;
    if (auto AT = dyn_cast<LLVM::LLVMArrayType>(cur))
      element = AT.getElementType();
    else if (auto VT = dyn_cast<VectorType>(cur))
      element = VT.getElementType();
    else
      return OffsetTree::unknown();

    assert(element && "an aggregate is made of something");
    step(constant, dynamic, dl.getTypeSize(element));
    cur = element;
  }

  // Where the walk ended is what an access through this reads.
  assert(cur && "a getelementptr lands on a type");
  return OffsetTree(cur, std::move(offsets), std::move(units));
}

llvm::raw_ostream &operator<<(llvm::raw_ostream &o, const OffsetType unit) {
  switch (unit.kind) {
  case OffsetType::Kind::Bytes:
    return o << unit.stride << "B";
  case OffsetType::Kind::Dim:
    return o << "dim<" << unit.dimSize << ">";
  }
  llvm_unreachable("Unknown offset unit");
}

llvm::raw_ostream &operator<<(llvm::raw_ostream &o, const Offset off);

void OffsetTree::print(llvm::raw_ostream &o) const {
  if (unknownOffset) {
    o << "unknown";
    return;
  }
  o << "[";
  for (auto &&[i, off] : llvm::enumerate(offsets)) {
    if (i)
      o << ", ";
    o << off << " x " << units[i];
  }
  o << "] of " << base;
}

llvm::raw_ostream &operator<<(llvm::raw_ostream &o, const OffsetTree &tree) {
  tree.print(o);
  return o;
}

llvm::raw_ostream &operator<<(llvm::raw_ostream &o, const Offset off) {
  switch (off.type) {
  case Offset::Type::Affine:
    return o << off.aff;
  case Offset::Type::Value:
    return o << off.val;
  case Offset::Type::Index:
    return o << off.idx;
  default:
    llvm_unreachable("Unknown offset type");
  }
}

namespace {
// The store to load forwarding relies on three conditions:
//
// 1) they need to have mathematically equivalent affine access functions
// (checked after full composition of load/store operands); this implies that
// they access the same single memref element for all iterations of the common
// surrounding loop,
//
// 2) the store op should dominate the load op,
//
// 3) among all op's that satisfy both (1) and (2), the one that postdominates
// all store op's that have a dependence into the load, is provably the last
// writer to the particular memref location being loaded at the load op, and its
// store value can be forwarded to the load. Note that the only dependences
// that are to be considered are those that are satisfied at the block* of the
// innermost common surrounding loop of the <store, load> being considered.
//
// (* A dependence being satisfied at a block: a dependence that is satisfied by
// virtue of the destination operation appearing textually / lexically after
// the source operation within the body of a 'affine.for' operation; thus, a
// dependence is always either satisfied by a loop or by a block).
//
// The above conditions are simple to check, sufficient, and powerful for most
// cases in practice - they are sufficient, but not necessary --- since they
// don't reason about loops that are guaranteed to execute at least once or
// multiple sources to forward from.
//
// TODO: more forwarding can be done when support for
// loop/conditional live-out SSA values is available.
// TODO: do general dead store elimination for memref's. This pass
// currently only eliminates the stores only if no other loads/uses (other
// than dealloc) remain.
//

struct PolygeistMem2Reg
    : public enzyme::impl::PolygeistMem2RegBase<PolygeistMem2Reg> {
  using PolygeistMem2RegBase::PolygeistMem2RegBase;
  void runOnOperation() override;

  // return if changed
  bool forwardStoreToLoad(
      mlir::Value AI, OffsetTree idx,
      SmallVectorImpl<Operation *> &loadOpsToErase,
      DenseMap<Operation *, SmallVector<Operation *>> &capturedAliasing,
      SymbolTableCollection &symbolTables);
};

} // end anonymous namespace

class ValueOrPlaceholder;

static inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                                            ValueOrPlaceholder &PH);

class ValueOrPlaceholder;

using ReplaceableUse = ValueOrPlaceholder *;

using BlockMap = std::map<Block *, ReplaceableUse>;

class ReplacementHandler {
public:
  Type elType;
  std::vector<std::unique_ptr<ValueOrPlaceholder>> allocs;
  DenseMap<Value, ValueOrPlaceholder *> valueMap;
  DenseMap<Operation *, ValueOrPlaceholder *> opOperands;

  BlockMap valueAtStartOfBlock;
  BlockMap valueAtEndOfBlock;
  SmallVector<std::pair<Value, ValueOrPlaceholder *>> replaceableValues;
  ReplacementHandler(Type elType) : elType(elType) {}

  ValueOrPlaceholder *get(std::nullptr_t);
  ValueOrPlaceholder *get(Value val);
  ValueOrPlaceholder *get(Block *val);
  ValueOrPlaceholder *get(scf::IfOp val, ValueOrPlaceholder *ifVal);
  ValueOrPlaceholder *get(affine::AffineIfOp val, ValueOrPlaceholder *ifVal);
  ValueOrPlaceholder *get(scf::ExecuteRegionOp val);

  void replaceValue(Value orig, Value post);
  void replaceOpWithValue(Operation *orig, ValueOrPlaceholder *ph, Value post) {
    valueMap[post] = ph;
  }
};

class ValueOrPlaceholder {
  ReplacementHandler &metaMap;

public:
  bool overwritten;
  Value val;
  Block *valueAtStart;
  scf::ExecuteRegionOp exOp;
  Operation *ifOp;
  ValueOrPlaceholder(ValueOrPlaceholder &&) = delete;
  ValueOrPlaceholder(const ValueOrPlaceholder &) = delete;
  ValueOrPlaceholder(std::nullptr_t, ReplacementHandler &metaMap)
      : metaMap(metaMap), overwritten(true), val(nullptr),
        valueAtStart(nullptr), exOp(nullptr), ifOp(nullptr) {}
  ValueOrPlaceholder(Value val, ReplacementHandler &metaMap)
      : metaMap(metaMap), overwritten(false), val(val), valueAtStart(nullptr),
        exOp(nullptr), ifOp(nullptr) {
    assert(val);
  }
  ValueOrPlaceholder(Block *valueAtStart, ReplacementHandler &metaMap)
      : metaMap(metaMap), overwritten(false), val(nullptr),
        valueAtStart(valueAtStart), exOp(nullptr), ifOp(nullptr) {
    assert(valueAtStart);
  }
  ValueOrPlaceholder(scf::ExecuteRegionOp exOp, ReplacementHandler &metaMap)
      : metaMap(metaMap), overwritten(false), val(nullptr),
        valueAtStart(nullptr), exOp(exOp), ifOp(nullptr) {
    assert(exOp);
  }
  ValueOrPlaceholder(scf::IfOp ifOp, ReplaceableUse ifLastVal,
                     ReplacementHandler &metaMap)
      : metaMap(metaMap), overwritten(false), val(nullptr),
        valueAtStart(nullptr), exOp(nullptr), ifOp(ifOp) {
    assert(ifOp);
    if (ifLastVal)
      metaMap.opOperands[ifOp] = ifLastVal;
  }
  ValueOrPlaceholder(affine::AffineIfOp ifOp, ReplaceableUse ifLastVal,
                     ReplacementHandler &metaMap)
      : metaMap(metaMap), overwritten(false), val(nullptr),
        valueAtStart(nullptr), exOp(nullptr), ifOp(ifOp) {
    assert(ifOp);
    if (ifLastVal)
      metaMap.opOperands[ifOp] = ifLastVal;
  }
  // Return true if this represents a full expression if all block argsare
  // defined at start Append the list of blocks requiring definition to block.
  bool definedWithArg(SmallPtrSetImpl<Block *> &block) {
    if (val)
      return true;
    if (overwritten)
      return false;
    if (valueAtStart) {
      auto found = metaMap.valueAtStartOfBlock.find(valueAtStart);
      if (found != metaMap.valueAtStartOfBlock.end()) {
        if (found->second->valueAtStart != valueAtStart)
          return found->second->definedWithArg(block);
      }
      block.insert(valueAtStart);
      return true;
    }
    if (ifOp) {
      if (auto sifOp = dyn_cast<scf::IfOp>(ifOp)) {
        auto thenFind = metaMap.valueAtEndOfBlock.find(getThenBlock(sifOp));
        assert(thenFind != metaMap.valueAtEndOfBlock.end());
        assert(thenFind->second);
        if (!thenFind->second->definedWithArg(block))
          return false;

        if (hasElse(sifOp)) {
          auto elseFind = metaMap.valueAtEndOfBlock.find(getElseBlock(sifOp));
          assert(elseFind != metaMap.valueAtEndOfBlock.end());
          assert(elseFind->second);
          if (!elseFind->second->definedWithArg(block))
            return false;
        } else {
          auto opFound = metaMap.opOperands.find(sifOp);
          assert(opFound != metaMap.opOperands.end());
          auto *ifLastValue = opFound->second;
          if (!ifLastValue->definedWithArg(block))
            return false;
        }
        return true;
      } else {
        auto aifOp = cast<affine::AffineIfOp>(ifOp);
        auto thenFind = metaMap.valueAtEndOfBlock.find(getThenBlock(aifOp));
        assert(thenFind != metaMap.valueAtEndOfBlock.end());
        assert(thenFind->second);
        if (!thenFind->second->definedWithArg(block))
          return false;

        if (hasElse(aifOp)) {
          auto elseFind = metaMap.valueAtEndOfBlock.find(getElseBlock(aifOp));
          assert(elseFind != metaMap.valueAtEndOfBlock.end());
          assert(elseFind->second);
          if (!elseFind->second->definedWithArg(block))
            return false;
        } else {
          auto opFound = metaMap.opOperands.find(ifOp);
          assert(opFound != metaMap.opOperands.end());
          auto *ifLastValue = opFound->second;
          if (!ifLastValue->definedWithArg(block))
            return false;
        }
        return true;
      }
    }
    if (exOp) {
      for (auto &B : exOp.getRegion()) {
        if (auto yield = dyn_cast<scf::YieldOp>(B.getTerminator())) {
          auto found = metaMap.valueAtEndOfBlock.find(&B);
          assert(found != metaMap.valueAtEndOfBlock.end());
          assert(found->second);
          if (!found->second->definedWithArg(block))
            return false;
        }
      }
      return true;
    }
    llvm_unreachable("unhandled");
  }
  Value materialize(bool full = true) {
    if (overwritten)
      return nullptr;
    if (val)
      return val;
    if (valueAtStart) {
      auto found = metaMap.valueAtStartOfBlock.find(valueAtStart);
      if (found != metaMap.valueAtStartOfBlock.end()) {
        if (found->second->valueAtStart != valueAtStart)
          return found->second->materialize(full);
        // valueAtStart = nullptr;
        // return this->val = found->second;
      }
      if (!full)
        return nullptr;
      llvm::errs() << " could not get valueAtStart: " << valueAtStart << "; ";
      if (found != metaMap.valueAtStartOfBlock.end()) {
        llvm::errs() << " map vas: " << *found->second << "\n";
      } else {
        llvm::errs() << " no map\n";
      }
      Block *blk = valueAtStart;
      blk->dump();
      llvm_unreachable("no null");
    }
    if (exOp)
      return materializeEx(full);
    if (ifOp)
      return materializeIf(full);
    llvm_unreachable("");
  }

  Value materializeEx(bool full = true) {
    assert(exOp);

    SmallVector<scf::YieldOp> yields;
    SmallVector<Value> values;
    std::set<size_t> equivalent;
    for (size_t i = 0, num = exOp.getNumResults(); i < num; i++)
      equivalent.insert(i);

    // Force early materialization in case any materializations
    // overwrite subsequent ones.
    for (auto &B : exOp.getRegion()) {
      if (auto yield = dyn_cast<scf::YieldOp>(B.getTerminator())) {
        auto found = metaMap.valueAtEndOfBlock.find(&B);
        if (found == metaMap.valueAtEndOfBlock.end()) {
          this->overwritten = true;
          this->exOp = nullptr;
          return nullptr;
        }
        assert(found->second);
        Value post = found->second->materialize(full);
        if (found->second->overwritten) {
          this->overwritten = true;
          this->exOp = nullptr;
          return nullptr;
        }
        if (!post) {
          if (full) {
            this->overwritten = true;
            this->exOp = nullptr;
          }
          return nullptr;
        }
      }
    }
    for (auto &B : exOp.getRegion()) {
      if (auto yield = dyn_cast<scf::YieldOp>(B.getTerminator())) {
        auto found = metaMap.valueAtEndOfBlock.find(&B);
        assert(found != metaMap.valueAtEndOfBlock.end());
        assert(found->second);
        Value post = found->second->materialize(full);
        assert(post);
        yields.push_back(yield);
        values.push_back(post);
        for (auto pair : llvm::enumerate(yield.getOperands()))
          if (pair.value() != post)
            equivalent.erase(pair.index());
      }
    }

    // Must contain only region invariant results.
    bool allSame = true;
    for (auto v : values)
      allSame &= v == values[0];

    // If all all paths are the same, and the value is not defined within the
    // execute region, simply return that single value, rather than creating a
    // new return.
    if (allSame) {
      if (values[0].getDefiningOp() &&
          !exOp->isAncestor(values[0].getDefiningOp())) {
        return values[0];
      }
      if (auto ba = dyn_cast<BlockArgument>(values[0]))
        if (!exOp->isAncestor(ba.getOwner()->getParentOp())) {
          return values[0];
        }
    }
    // If there's an equivalent return, don't create a new return and instead
    // use that result.
    if (equivalent.size() > 0) {
      return exOp.getResult(*equivalent.begin());
    }

    OpBuilder B(exOp.getContext());
    B.setInsertionPoint(exOp);
    SmallVector<mlir::Type, 4> tys(exOp.getResultTypes().begin(),
                                   exOp.getResultTypes().end());
    tys.push_back(metaMap.elType);
    auto nextEx = mlir::scf::ExecuteRegionOp::create(B, exOp.getLoc(), tys);

    nextEx.getRegion().takeBody(exOp.getRegion());
    for (auto pair : llvm::zip(yields, values)) {
      SmallVector<Value, 4> vals = std::get<0>(pair).getOperands();
      vals.push_back(std::get<1>(pair));
      std::get<0>(pair)->setOperands(vals);
    }

    SmallVector<mlir::Value, 3> resvals = nextEx.getResults();
    this->val = resvals.back();
    resvals.pop_back();

    for (auto pair : llvm::zip(exOp.getResults(), resvals)) {
      metaMap.replaceValue(std::get<0>(pair), std::get<1>(pair));
    }
    metaMap.replaceOpWithValue(exOp, this, this->val);
    // StoringOperations.erase(exOp);
    // StoringOperations.insert(nextEx);
    exOp.erase();
    this->exOp = nullptr;
    return this->val;
  }

  Value materializeIf(bool full = true) {
    if (auto sop = dyn_cast<scf::IfOp>(ifOp))
      return materializeIf<scf::IfOp, scf::YieldOp>(sop, full);
    return materializeIf<affine::AffineIfOp, affine::AffineYieldOp>(
        cast<affine::AffineIfOp>(ifOp), full);
  }

  template <typename IfType, typename YieldType>
  Value materializeIf(IfType ifOp, bool full = true) {
    auto thenFind = metaMap.valueAtEndOfBlock.find(getThenBlock(ifOp));
    assert(thenFind != metaMap.valueAtEndOfBlock.end());
    assert(thenFind->second);
    Value thenVal = thenFind->second->materialize(full);
    if (thenFind->second->overwritten) {
      this->overwritten = true;
      metaMap.opOperands.erase(ifOp);
      this->ifOp = nullptr;
      return nullptr;
    }
    if (!thenVal) {
      if (full) {
        this->overwritten = true;
        metaMap.opOperands.erase(ifOp);
        this->ifOp = nullptr;
      }
      return nullptr;
    }
    Value elseVal;

    if (hasElse(ifOp)) {
      auto elseFind = metaMap.valueAtEndOfBlock.find(getElseBlock(ifOp));
      assert(elseFind != metaMap.valueAtEndOfBlock.end());
      assert(elseFind->second);
      elseVal = elseFind->second->materialize(full);
      if (elseFind->second->overwritten) {
        this->overwritten = true;
        metaMap.opOperands.erase(ifOp);
        this->ifOp = nullptr;
        return nullptr;
      }
    } else {
      auto opFound = metaMap.opOperands.find(ifOp);
      assert(opFound != metaMap.opOperands.end());
      auto ifLastValue = opFound->second;
      elseVal = ifLastValue->materialize(full);
      if (ifLastValue->overwritten) {
        this->overwritten = true;
        metaMap.opOperands.erase(ifOp);
        this->ifOp = nullptr;
        return nullptr;
      }
    }

    if (!elseVal) {
      if (full) {
        this->overwritten = true;
        metaMap.opOperands.erase(ifOp);
        this->ifOp = nullptr;
      }
      return nullptr;
    }

    // Rematerialize thenVal in case it was overwritten by the getElseRegion
    //  materialization.
    thenVal = thenFind->second->materialize(full);

    if (thenVal == elseVal) {
      return thenVal;
    }

    if (hasElse(ifOp)) {
      for (auto tup : llvm::reverse(
               llvm::zip(ifOp.getResults(), getThenYield(ifOp).getOperands(),
                         getElseYield(ifOp).getOperands()))) {
        if (std::get<1>(tup) == thenVal && std::get<2>(tup) == elseVal) {
          return std::get<0>(tup);
        }
      }
    }

    OpBuilder B(ifOp.getContext());
    B.setInsertionPoint(ifOp);
    SmallVector<mlir::Type, 4> tys(ifOp.getResultTypes().begin(),
                                   ifOp.getResultTypes().end());
    tys.push_back(thenVal.getType());
    auto nextIf = cloneWithoutResults(ifOp, B, {}, tys);

    SmallVector<mlir::Value, 4> thenVals = getThenYield(ifOp).getOperands();
    thenVals.push_back(thenVal);
    getThenRegion(nextIf).takeBody(getThenRegion(ifOp));
    getThenYield(nextIf)->setOperands(thenVals);

    if (hasElse(ifOp)) {
      getElseRegion(nextIf).getBlocks().clear();
      SmallVector<mlir::Value, 4> elseVals = getElseYield(ifOp).getOperands();
      elseVals.push_back(elseVal);
      getElseRegion(nextIf).takeBody(getElseRegion(ifOp));
      getElseYield(nextIf)->setOperands(elseVals);
    } else {
      B.setInsertionPoint(&getElseRegion(nextIf).back(),
                          getElseRegion(nextIf).back().begin());
      SmallVector<mlir::Value, 4> elseVals = {elseVal};
      YieldType::create(B, ifOp.getLoc(), elseVals);
    }

    SmallVector<mlir::Value, 3> resvals = nextIf.getResults();
    this->val = resvals.back();
    resvals.pop_back();

    for (auto pair : llvm::zip(ifOp.getResults(), resvals)) {
      metaMap.replaceValue(std::get<0>(pair), std::get<1>(pair));
    }
    metaMap.replaceOpWithValue(ifOp, this, this->val);

    // StoringOperations.erase(ifOp);
    // StoringOperations.insert(nextIf);
    ifOp.erase();
    metaMap.opOperands.erase(ifOp);
    this->ifOp = nullptr;
    return this->val;
  }
};

static inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                                            ValueOrPlaceholder &PH) {
  if (PH.overwritten)
    return os << "<overwritten>";
  if (PH.val)
    return os << "val:" << PH.val;
  if (PH.valueAtStart) {
    PH.valueAtStart->print(os);
    return os;
    ;
  }
  if (PH.ifOp)
    return os << "ifOp:" << *PH.ifOp;
  if (PH.exOp)
    return os << "exOp:" << PH.exOp;
  return os;
}
ValueOrPlaceholder *ReplacementHandler::get(std::nullptr_t val) {
  ValueOrPlaceholder *PH;
  allocs.emplace_back(PH = new ValueOrPlaceholder(val, *this));
  return PH;
}
ValueOrPlaceholder *ReplacementHandler::get(Value val) {
  auto found = valueMap.find(val);
  if (found != valueMap.end())
    return found->second;
  ValueOrPlaceholder *PH;
  allocs.emplace_back(PH = new ValueOrPlaceholder(val, *this));
  valueMap.insert(std::make_pair(val, PH));
  return PH;
}
ValueOrPlaceholder *ReplacementHandler::get(Block *val) {
  ValueOrPlaceholder *PH;
  allocs.emplace_back(PH = new ValueOrPlaceholder(val, *this));
  return PH;
}
ValueOrPlaceholder *ReplacementHandler::get(scf::IfOp val,
                                            ValueOrPlaceholder *ifVal) {
  ValueOrPlaceholder *PH;
  allocs.emplace_back(PH = new ValueOrPlaceholder(val, ifVal, *this));
  return PH;
}
ValueOrPlaceholder *ReplacementHandler::get(affine::AffineIfOp val,
                                            ValueOrPlaceholder *ifVal) {
  ValueOrPlaceholder *PH;
  allocs.emplace_back(PH = new ValueOrPlaceholder(val, ifVal, *this));
  return PH;
}
ValueOrPlaceholder *ReplacementHandler::get(scf::ExecuteRegionOp val) {
  ValueOrPlaceholder *PH;
  allocs.emplace_back(PH = new ValueOrPlaceholder(val, *this));
  return PH;
}
void ReplacementHandler::replaceValue(Value orig, Value post) {
  orig.replaceAllUsesWith(post);
  auto found = valueMap.find(orig);
  if (found != valueMap.end()) {
    auto pfound = valueMap.find(post);
    if (pfound == valueMap.end()) {
      found->second->val = post;
    } else {
      for (auto &pair : valueAtStartOfBlock)
        if (pair.second == found->second)
          pair.second = pfound->second;
      for (auto &pair : valueAtEndOfBlock)
        if (pair.second == found->second)
          pair.second = pfound->second;
      for (auto &pair : replaceableValues)
        if (pair.second == found->second)
          pair.second = pfound->second;
      for (auto &pair : opOperands)
        if (pair.second == found->second)
          pair.second = pfound->second;
      valueMap.erase(found);
    }
  }
}

struct Analyzer {

  const std::set<Block *> &Good;
  const std::set<Block *> &Bad;
  const std::set<Block *> &Other;
  std::set<Block *> Legal;
  std::set<Block *> Illegal;
  size_t depth;
  Analyzer(const std::set<Block *> &Good, const std::set<Block *> &Bad,
           const std::set<Block *> &Other, std::set<Block *> Legal,
           std::set<Block *> Illegal, size_t depth = 0)
      : Good(Good), Bad(Bad), Other(Other), Legal(Legal), Illegal(Illegal),
        depth(depth) {}

  void analyze() {
    while (1) {
      std::deque<Block *> todo(Other.begin(), Other.end());
      todo.insert(todo.end(), Good.begin(), Good.end());
      while (todo.size()) {
        auto *block = todo.front();
        todo.pop_front();
        if (Legal.count(block) || Illegal.count(block))
          continue;
        bool currentlyLegal = !block->hasNoPredecessors();
        for (auto *pred : block->getPredecessors()) {
          if (Bad.count(pred)) {
            assert(!Legal.count(block));
            Illegal.insert(block);
            currentlyLegal = false;
            for (auto *succ : block->getSuccessors()) {
              todo.push_back(succ);
            }
            break;
          } else if (Good.count(pred) || Legal.count(pred)) {
            continue;
          } else if (Illegal.count(pred)) {
            Illegal.insert(block);
            currentlyLegal = false;
            for (auto *succ : block->getSuccessors()) {
              todo.push_back(succ);
            }
            break;
          } else {
            /*
            if (!Other.count(pred)) {
              pred->getParentOp()->dump();
              pred->dump();
              llvm::errs() << " - pred ptr: " << pred << "\n";
            }
            assert(Other.count(pred));
            */
            currentlyLegal = false;
            break;
          }
        }
        if (currentlyLegal) {
          Legal.insert(block);
          assert(!Illegal.count(block));
          for (auto *succ : block->getSuccessors()) {
            todo.push_back(succ);
          }
        }
      }
      bool changed = false;
      for (auto *O : Other) {
        if (Legal.count(O) || Illegal.count(O))
          continue;
        Analyzer AssumeLegal(Good, Bad, Other, Legal, Illegal, depth + 1);
        AssumeLegal.Legal.insert(O);
        AssumeLegal.analyze();
        bool currentlyLegal = true;
        for (auto *pred : O->getPredecessors()) {
          if (!AssumeLegal.Legal.count(pred) && !AssumeLegal.Good.count(pred)) {
            currentlyLegal = false;
            break;
          }
        }
        if (currentlyLegal) {
          Legal.insert(AssumeLegal.Legal.begin(), AssumeLegal.Legal.end());
          assert(!Illegal.count(O));
          changed = true;
          break;
        } else {
          Illegal.insert(O);
          assert(!Legal.count(O));
        }
      }
      if (!changed)
        break;
    }
  }
};

// Remove block arguments if possible
void removeRedundantBlockArgs(
    Value AI, Type elType,
    std::map<Block *, BlockArgument> &blocksWithAddedArgs) {
  std::deque<Block *> todo;
  for (auto &p : blocksWithAddedArgs)
    todo.push_back(p.first);

  while (todo.size()) {
    auto *block = todo.front();
    todo.pop_front();
    if (!blocksWithAddedArgs.count(block))
      continue;

    BlockArgument blockArg = blocksWithAddedArgs.find(block)->second;
    if (blockArg.getOwner() != block)
      continue;

    assert(blockArg.getOwner() == block);

    mlir::Value val = nullptr;
    bool legal = true;

    SetVector<Block *> prepred(block->getPredecessors().begin(),
                               block->getPredecessors().end());
    for (auto *pred : prepred) {
      mlir::Value pval = nullptr;

      if (auto op = dyn_cast<cf::BranchOp>(pred->getTerminator())) {
        pval = op.getOperands()[blockArg.getArgNumber()];
        if (pval.getType() != elType) {
          if (auto *def = pval.getDefiningOp())
            def->getParentRegion()->getParentOp()->dump();
          llvm::errs() << pval << " - " << AI << "\n";
        }
        assert(pval.getType() == elType);
        if (pval == blockArg)
          pval = nullptr;
      } else if (auto op = dyn_cast<cf::CondBranchOp>(pred->getTerminator())) {
        if (op.getTrueDest() == block) {
          if (blockArg.getArgNumber() >= op.getTrueOperands().size()) {
            block->dump();
            llvm::errs() << op << " ba: " << blockArg.getArgNumber() << "\n";
          }
          assert(blockArg.getArgNumber() < op.getTrueOperands().size());
          pval = op.getTrueOperands()[blockArg.getArgNumber()];
          assert(pval.getType() == elType);
          if (pval == blockArg)
            pval = nullptr;
        }
        if (op.getFalseDest() == block) {
          assert(blockArg.getArgNumber() < op.getFalseOperands().size());
          auto pval2 = op.getFalseOperands()[blockArg.getArgNumber()];
          assert(pval2.getType() == elType);
          if (pval2 != blockArg) {
            if (pval == nullptr) {
              pval = pval2;
            } else if (pval != pval2) {
              legal = false;
              break;
            }
          }
          if (pval == blockArg)
            pval = nullptr;
        }
      } else if (auto op = dyn_cast<cf::SwitchOp>(pred->getTerminator())) {
        mlir::OpBuilder subbuilder(op.getOperation());
        if (op.getDefaultDestination() == block) {
          pval = op.getDefaultOperands()[blockArg.getArgNumber()];
          if (pval == blockArg)
            pval = nullptr;
        }
        for (auto pair : llvm::enumerate(op.getCaseDestinations())) {
          if (pair.value() == block) {
            auto pval2 =
                op.getCaseOperands(pair.index())[blockArg.getArgNumber()];
            if (pval2 != blockArg) {
              if (pval == nullptr)
                pval = pval2;
              else if (pval != pval2) {
                legal = false;
                break;
              }
            }
          }
        }
        if (legal == false)
          break;
      } else {
        llvm::errs() << *pred->getParent()->getParentOp() << "\n";
        pred->dump();
        block->dump();
        llvm_unreachable("unknown branch");
      }

      assert(pval != blockArg);
      if (val == nullptr) {
        val = pval;
        if (pval)
          assert(val.getType() == elType);
      } else {
        if (pval != nullptr && val != pval) {
          legal = false;
          break;
        }
      }
    }
    if (legal)
      assert(val || block->hasNoPredecessors());

    bool used = false;
    for (auto *U : blockArg.getUsers()) {

      if (auto op = dyn_cast<cf::BranchOp>(U)) {
        size_t i = 0;
        for (auto V : op.getOperands()) {
          if (V == blockArg &&
              !(i == blockArg.getArgNumber() && op.getDest() == block)) {
            used = true;
            break;
          }
        }
        if (used)
          break;
      } else if (auto op = dyn_cast<cf::CondBranchOp>(U)) {
        size_t i = 0;
        for (auto V : op.getTrueOperands()) {
          if (V == blockArg &&
              !(i == blockArg.getArgNumber() && op.getTrueDest() == block)) {
            used = true;
            break;
          }
        }
        if (used)
          break;
        i = 0;
        for (auto V : op.getFalseOperands()) {
          if (V == blockArg &&
              !(i == blockArg.getArgNumber() && op.getFalseDest() == block)) {
            used = true;
            break;
          }
        }
      } else
        used = true;
    }
    if (!used) {
      legal = true;
    }

    if (legal) {
      for (auto *U : blockArg.getUsers()) {
        if (auto *block = U->getBlock()) {
          todo.push_back(block);
          for (auto *succ : block->getSuccessors())
            todo.push_back(succ);
        }
      }
      if (val != nullptr) {
        if (blockArg.getType() != val.getType()) {
          block->dump();
          llvm::errs() << " AI: " << AI << "\n";
          llvm::errs() << blockArg << " val " << val << "\n";
        }
        assert(blockArg.getType() == val.getType());
        blockArg.replaceAllUsesWith(val);
      } else {
      }

      SetVector<Block *> prepred(block->getPredecessors().begin(),
                                 block->getPredecessors().end());
      for (auto *pred : prepred) {
        if (auto op = dyn_cast<cf::BranchOp>(pred->getTerminator())) {
          mlir::OpBuilder subbuilder(op.getOperation());
          std::vector<Value> args(op.getOperands().begin(),
                                  op.getOperands().end());
          args.erase(args.begin() + blockArg.getArgNumber());
          assert(args.size() == op.getOperands().size() - 1);
          cf::BranchOp::create(subbuilder, op.getLoc(), op.getDest(), args);
          op.erase();
        } else if (auto op =
                       dyn_cast<cf::CondBranchOp>(pred->getTerminator())) {

          mlir::OpBuilder subbuilder(op.getOperation());
          std::vector<Value> trueargs(op.getTrueOperands().begin(),
                                      op.getTrueOperands().end());
          std::vector<Value> falseargs(op.getFalseOperands().begin(),
                                       op.getFalseOperands().end());
          if (op.getTrueDest() == block) {
            trueargs.erase(trueargs.begin() + blockArg.getArgNumber());
          }
          if (op.getFalseDest() == block) {
            falseargs.erase(falseargs.begin() + blockArg.getArgNumber());
          }
          assert(trueargs.size() < op.getTrueOperands().size() ||
                 falseargs.size() < op.getFalseOperands().size());
          cf::CondBranchOp::create(subbuilder, op.getLoc(), op.getCondition(),
                                   op.getTrueDest(), trueargs,
                                   op.getFalseDest(), falseargs);
          op.erase();
        } else if (auto op = dyn_cast<cf::SwitchOp>(pred->getTerminator())) {
          mlir::OpBuilder builder(op.getOperation());
          SmallVector<Value> defaultOps(op.getDefaultOperands().begin(),
                                        op.getDefaultOperands().end());
          if (op.getDefaultDestination() == block)
            defaultOps.erase(defaultOps.begin() + blockArg.getArgNumber());

          SmallVector<SmallVector<Value>> cases;
          SmallVector<ValueRange> vrange;
          for (auto pair : llvm::enumerate(op.getCaseDestinations())) {
            cases.emplace_back(op.getCaseOperands(pair.index()));
            if (pair.value() == block) {
              cases.back().erase(cases.back().begin() +
                                 blockArg.getArgNumber());
            }
          }
          for (auto &c : cases) {
            vrange.push_back(c);
          }
          cf::SwitchOp::create(builder, op.getLoc(), op.getFlag(),
                               op.getDefaultDestination(), defaultOps,
                               op.getCaseValuesAttr(), op.getCaseDestinations(),
                               vrange);
          op.erase();
        }
      }
      block->eraseArgument(blockArg.getArgNumber());
      blocksWithAddedArgs.erase(block);
    }
  }
}

// Recursively extract leaves
static void collectLeaves(OpBuilder &b, Value val,
                          SmallVectorImpl<Value> &leaves) {
  Type t = val.getType();

  if (auto ST = dyn_cast<LLVM::LLVMStructType>(t)) {
    for (unsigned i = 0; i < ST.getBody().size(); i++) {
      Value extracted = LLVM::ExtractValueOp::create(b, val.getLoc(), val, i);
      collectLeaves(b, extracted, leaves);
    }
  } else if (auto AT = dyn_cast<LLVM::LLVMArrayType>(t)) {
    for (unsigned i = 0; i < AT.getNumElements(); i++) {
      Value extracted = LLVM::ExtractValueOp::create(b, val.getLoc(), val, i);
      collectLeaves(b, extracted, leaves);
    }
  } else {
    leaves.push_back(val);
  }
}

Value castToType(Type elType, Value val, Operation *op) {
  if (val.getType() == elType)
    return val;

  OpBuilder b(op);
  b.setInsertionPoint(op);
  if (isa<IntegerType>(val.getType()) && isa<LLVM::LLVMPointerType>(elType)) {
    return LLVM::IntToPtrOp::create(b, val.getLoc(), elType, val);
  } else if (isa<LLVM::LLVMPointerType>(val.getType()) &&
             isa<IntegerType>(elType)) {
    return LLVM::PtrToIntOp::create(b, val.getLoc(), elType, val);
  } else if ((isa<IntegerType>(val.getType()) ||
              isa<FloatType>(val.getType())) &&
             (isa<IntegerType>(elType) || isa<FloatType>(elType))) {
    return arith::BitcastOp::create(b, val.getLoc(), elType, val);
  } else if (isa<LLVM::LLVMStructType, LLVM::LLVMArrayType>(val.getType())) {
    // If the aggregate has only one element, extract and recurse
    if (auto ST = dyn_cast<LLVM::LLVMStructType>(val.getType())) {
      if (ST.getBody().size() == 1) {
        Value extracted = LLVM::ExtractValueOp::create(b, val.getLoc(), val, 0);
        return castToType(elType, extracted, op);
      }
    } else if (auto AT = dyn_cast<LLVM::LLVMArrayType>(val.getType())) {
      if (AT.getNumElements() == 1) {
        Value extracted = LLVM::ExtractValueOp::create(b, val.getLoc(), val, 0);
        return castToType(elType, extracted, op);
      }
    }
    SmallVector<Value> leaves;
    collectLeaves(b, val, leaves);

    // Only cast if leaves are of the same type
    if (!leaves.empty()) {
      Type leafType = leaves[0].getType();
      if (llvm::all_of(leaves,
                       [&](Value v) { return v.getType() == leafType; })) {
        // Create a vector of the leaves
        auto vecType = VectorType::get(leaves.size(), leafType);
        Value vec = LLVM::PoisonOp::create(b, val.getLoc(), vecType);
        for (auto [i, leaf] : llvm::enumerate(leaves)) {
          Value idx = LLVM::ConstantOp::create(b, val.getLoc(), b.getI32Type(),
                                               b.getI32IntegerAttr(i));
          vec = LLVM::InsertElementOp::create(b, val.getLoc(), vec, leaf, idx);
        }
        return castToType(elType, vec, op);
      }
    }

  } else if (auto VT = dyn_cast<VectorType>(val.getType())) {
    DataLayout dl(op->getParentOfType<ModuleOp>());
    if (dl.getTypeSize(val.getType()) == dl.getTypeSize(elType)) {
      if (isa<IntegerType>(elType) || isa<FloatType>(elType))
        return LLVM::BitcastOp::create(b, val.getLoc(), elType, val);
    }
  } else if (auto ST = dyn_cast<LLVM::LLVMStructType>(elType)) {
    if (ST.getBody().size() == 1) {
      auto ud = LLVM::UndefOp::create(b, val.getLoc(), elType);
      auto c0 = castToType(ST.getBody()[0], val, op);
      b.setInsertionPoint(op);
      return LLVM::InsertValueOp::create(b, val.getLoc(), ud, c0,
                                         b.getDenseI64ArrayAttr({0}));
    }
  }
  llvm::errs() << " mismatched load type, needed: " << elType << " found "
               << val << "\n";
  llvm_unreachable("mismatched type");
}

// Whether every position the value is passed at carries the given attribute
// on the callee. One position without it is enough for the call to do
// whatever the attribute rules out.
static bool callArgsAllHaveAttr(CallOpInterface callOp, Value val,
                                SymbolTableCollection &symbolTables,
                                StringRef attr) {
  auto calleeAttr = dyn_cast<SymbolRefAttr>(callOp.getCallableForCallee());
  if (!calleeAttr)
    return false;
  auto callee = symbolTables.lookupSymbolIn(callOp->getParentOfType<ModuleOp>(),
                                            calleeAttr);
  auto fn = dyn_cast_or_null<FunctionOpInterface>(callee);
  if (!fn)
    return false;

  auto operands = callOp.getArgOperands();
  bool seen = false;
  for (int i = 0; i < operands.size(); i++) {
    if (operands[i] == val) {
      seen = true;
      if (i >= fn.getNumArguments() || !fn.getArgAttr(i, attr))
        return false;
    }
  }
  return seen;
}

// Check if call captures alloca instance by checking for llvm.nocapture
// attribute
bool isCallNonCapturing(CallOpInterface callOp, Value val,
                        SymbolTableCollection &symbolTables) {
  return callArgsAllHaveAttr(callOp, val, symbolTables,
                             LLVM::LLVMDialect::getNoCaptureAttrName());
}

// Whether a memory-effects attribute rules out writes that could reach a
// pointer argument's bytes. Later LLVM spells whole-function
// readonly/readnone this way. argMem covers accesses through the argument
// pointers themselves; the same bytes can also be written through an access
// classified as other -- a captured pointer, a global alias -- so both must
// be write-free. Inaccessible memory cannot alias an argument, so it alone
// may be written.
static bool noWrite(LLVM::ModRefInfo mr) {
  return mr == LLVM::ModRefInfo::NoModRef || mr == LLVM::ModRefInfo::Ref;
}
static bool argMemOnlyRead(LLVM::MemoryEffectsAttr me) {
  return me && noWrite(me.getArgMem()) && noWrite(me.getOther());
}

// nocapture only says the callee does not hold on to the pointer; an out
// parameter is nocapture and written through. Only readonly (or readnone)
// says the call leaves the slot's contents alone -- carried per argument,
// as a memory-effects attribute on the call or the callee, or as the older
// readonly/readnone function attribute spelled through passthrough.
bool isCallArgOnlyRead(CallOpInterface callOp, Value val,
                       SymbolTableCollection &symbolTables) {
  if (auto call = dyn_cast<LLVM::CallOp>(callOp.getOperation()))
    if (argMemOnlyRead(call.getMemoryEffectsAttr()))
      return true;
  if (auto calleeAttr =
          dyn_cast<SymbolRefAttr>(callOp.getCallableForCallee())) {
    if (auto fn =
            dyn_cast_or_null<LLVM::LLVMFuncOp>(symbolTables.lookupSymbolIn(
                callOp->getParentOfType<ModuleOp>(), calleeAttr))) {
      if (argMemOnlyRead(fn.getMemoryEffectsAttr()))
        return true;
      if (auto pass = fn->getAttrOfType<ArrayAttr>("passthrough"))
        for (Attribute a : pass)
          if (auto s = dyn_cast<StringAttr>(a))
            if (s.getValue() == "readonly" || s.getValue() == "readnone")
              return true;
    }
  }
  return callArgsAllHaveAttr(callOp, val, symbolTables,
                             LLVM::LLVMDialect::getReadonlyAttrName()) ||
         callArgsAllHaveAttr(callOp, val, symbolTables,
                             LLVM::LLVMDialect::getReadnoneAttrName());
}

// fopen, fclose
std::set<std::string> NoWriteFunctions = {"exit", "__errno_location"};
// This is a straightforward implementation not optimized for speed. Optimize
// if needed.
// The way into `from` that reaches what lies at `at` bytes into it and is of
// the size of `want`, as the indices an extractvalue takes. Aggregates are
// walked into; anything else is only reachable when it is the whole of it.
static bool extractPath(Type from, uint64_t at, Type want, const DataLayout &dl,
                        SmallVectorImpl<int64_t> &path) {
  auto wantSize = typeSize(want, dl), fromSize = typeSize(from, dl);
  if (!wantSize || !fromSize || at + *wantSize > *fromSize)
    return false;
  if (at == 0 && *wantSize == *fromSize)
    return true;

  if (auto ST = dyn_cast<LLVM::LLVMStructType>(from)) {
    uint64_t byte = 0;
    for (auto &&[i, member] : llvm::enumerate(ST.getBody())) {
      auto size = typeSize(member, dl);
      if (!size)
        return false;
      if (!ST.isPacked())
        byte = llvm::alignTo(byte, dl.getTypeABIAlignment(member));
      if (at >= byte && at + *wantSize <= byte + *size) {
        path.push_back(i);
        return extractPath(member, at - byte, want, dl, path);
      }
      byte += *size;
    }
    return false;
  }

  if (auto AT = dyn_cast<LLVM::LLVMArrayType>(from)) {
    auto size = typeSize(AT.getElementType(), dl);
    if (!size || !*size)
      return false;
    uint64_t index = at / *size;
    if (index >= AT.getNumElements() || at % *size + *wantSize > *size)
      return false;
    path.push_back(index);
    return extractPath(AT.getElementType(), at % *size, want, dl, path);
  }

  return false;
}

bool PolygeistMem2Reg::forwardStoreToLoad(
    mlir::Value AI, OffsetTree idx,
    SmallVectorImpl<Operation *> &loadOpsToErase,
    DenseMap<Operation *, SmallVector<Operation *>> &capturedAliasing,
    SymbolTableCollection &symbolTables) {
  bool changed = false;
  std::set<mlir::Operation *> loadOps;
  // Transfers that read exactly one slot, which act as a load of it.
  std::set<mlir::Operation *> transferLoads;
  // Loads of a piece of the slot, and how far into it that piece lies.
  DenseMap<mlir::Operation *, uint64_t> containedLoads;
  mlir::Location loc = AI.getLoc();
  std::set<mlir::Operation *> allStoreOps;
  // Reads standing for what a transfer moved. One is worth the load it costs
  // only where something went on to ask for it: left behind unasked for, it is
  // a read of the other side of the transfer, which promoting that side sees,
  // sweeps, and counts as progress -- and promoting this side writes again the
  // round after, so neither side ever settles. There are several ways out of
  // this function and the reads have to go on all of them.
  SmallVector<Operation *> transferReads;
  auto dropUnaskedReads = llvm::make_scope_exit([&transferReads]() {
    for (auto *read : llvm::reverse(transferReads))
      if (read->getResult(0).use_empty())
        read->erase();
  });

  if (idx.isUnknown())
    return changed;

  DataLayout dl = DataLayout::closest(AI.getDefiningOp());
  Type elType = idx.getBase();

  std::deque<std::pair<mlir::Value, OffsetTree>> list = {{AI, OffsetTree()}};

  SmallPtrSet<Operation *, 4> AliasingStoreOperations;

  LLVM_DEBUG(llvm::dbgs()
             << "Begin forwarding store of " << AI << " to load\n"
             << *AI.getDefiningOp()->getParentOfType<FunctionOpInterface>()
             << "\n");
  bool captured = AI.getDefiningOp<memref::GetGlobalOp>();
  bool SharedMemAddr = false;
  while (list.size()) {
    auto pair = list.front();
    auto val = pair.first;
    if (auto MT = dyn_cast<MemRefType>(val.getType())) {
      if (auto ia = dyn_cast_or_null<IntegerAttr>(MT.getMemorySpace()))
        SharedMemAddr = ia.getValue() == 5;
    } else {
      auto PT = dyn_cast<LLVM::LLVMPointerType>(val.getType());
      SharedMemAddr = PT.getAddressSpace() == 5;
    }
    auto tree = pair.second;
    list.pop_front();
    for (auto *user : val.getUsers()) {
      if (auto co = dyn_cast<mlir::memref::CastOp>(user)) {
        list.emplace_back((Value)co, tree);
        continue;
      }
      if (auto co = dyn_cast<Memref2PointerOp>(user)) {
        list.emplace_back((Value)co, tree);
        continue;
      }
      if (auto co = dyn_cast<Pointer2MemrefOp>(user)) {
        list.emplace_back((Value)co, tree);
        continue;
      }
      // If at the same index, the "hole" property applies
      // and we can go through.
      if (isa<BarrierOp>(user)) {
        continue;
      }
      if (auto co = dyn_cast<mlir::LLVM::GEPOp>(user)) {
        list.emplace_back((Value)co, tree.add(gepOffsets(co, dl), dl));
        continue;
      }
      if (auto co = dyn_cast<mlir::LLVM::BitcastOp>(user)) {
        list.emplace_back((Value)co, tree);
        continue;
      }
      if (auto co = dyn_cast<mlir::LLVM::AddrSpaceCastOp>(user)) {
        list.emplace_back((Value)co, tree);
        continue;
      }
      if (isa<LLVM::LifetimeStartOp, LLVM::LifetimeEndOp>(user)) {
        continue;
      }
      // A load naming the slot reads it; one that may only overlap it reads
      // something this knows nothing about, which is no reason to stop.
      auto matchLoad = [&](Operation *loadOp, OffsetTree accessed) {
        uint64_t at = 0;
        Type read = loadOp->getResult(0).getType();
        switch (idx.matches(tree.add(accessed, dl), dl, &at)) {
        case Match::Exact:
          // Whichever spelling of the slot is read first is the one the value
          // forwarded out of it takes.
          if (elType && read != elType)
            return;
          elType = read;
          break;
        case Match::Contains: {
          // Reading a piece of the slot is reading a piece of its value, as
          // long as there is a way into it that reaches that piece.
          Type held = elType ? elType : idx.getBase();
          SmallVector<int64_t> path;
          if (!held || !extractPath(held, at, read, dl, path))
            return;
          containedLoads[loadOp] = at;
          break;
        }
        default:
          return;
        }
        loadOps.insert(loadOp);
        LLVM_DEBUG(llvm::dbgs() << "Matching Load: " << *loadOp << "\n");
      };
      auto matchStore = [&](Operation *storeOp, OffsetTree accessed) {
        switch (idx.matches(tree.add(accessed, dl), dl)) {
        case Match::Exact:
          LLVM_DEBUG(llvm::dbgs() << "Matching Store: " << *storeOp << "\n");
          allStoreOps.insert(storeOp);
          break;
        case Match::Contains:
        case Match::Maybe:
          LLVM_DEBUG(llvm::dbgs()
                     << "Mabye Aliasing Store: " << *storeOp << "\n");
          AliasingStoreOperations.insert(storeOp);
          break;
        case Match::None:
          break;
        }
      };

      if (auto loadOp = dyn_cast<mlir::memref::LoadOp>(user)) {
        matchLoad(loadOp, accessOffsets(loadOp.getType(), loadOp.getMemRef(),
                                        loadOp.getIndices(), dl));
        continue;
      }
      if (auto loadOp = dyn_cast<mlir::LLVM::LoadOp>(user)) {
        matchLoad(loadOp, accessOffsets(loadOp.getType()));
        continue;
      }
      if (auto loadOp = dyn_cast<affine::AffineLoadOp>(user)) {
        matchLoad(loadOp, accessOffsets(loadOp.getType(), loadOp.getMemRef(),
                                        loadOp.getAffineMapAttr().getValue(),
                                        loadOp.getMapOperands(), dl));
        continue;
      }
      if (auto storeOp = dyn_cast<mlir::memref::StoreOp>(user)) {
        if (storeOp.getValue() == val)
          captured = true;
        else
          matchStore(storeOp, accessOffsets(storeOp.getValue().getType(),
                                            storeOp.getMemRef(),
                                            storeOp.getIndices(), dl));
        continue;
      }
      if (auto storeOp = dyn_cast<LLVM::StoreOp>(user)) {

        if (storeOp.getValue() == val) {
          captured = true;
        } else
          matchStore(storeOp, accessOffsets(storeOp.getValue().getType()));
        continue;
      }

      if (auto storeOp = dyn_cast<affine::AffineStoreOp>(user)) {
        if (storeOp.getValue() == val) {
          captured = true;
        } else
          matchStore(storeOp,
                     accessOffsets(storeOp.getValue().getType(),
                                   storeOp.getMemRef(),
                                   storeOp.getAffineMapAttr().getValue(),
                                   storeOp.getMapOperands(), dl));
        continue;
      }
      if (auto callOp = dyn_cast<CallOpInterface>(user)) {
        auto callee = dyn_cast<SymbolRefAttr>(callOp.getCallableForCallee());
        if (!callee || callee.getLeafReference() != "free") {
          if (!isCallNonCapturing(callOp, val, symbolTables)) {
            LLVM_DEBUG(llvm::dbgs() << "Aliasing Store: " << callOp << "\n");
            AliasingStoreOperations.insert(callOp.getOperation());
            if (!callee || !getNonCapturingFunctions().count(
                               callee.getLeafReference().str()))
              captured = true;
          } else if (!isCallArgOnlyRead(callOp, val, symbolTables)) {
            // The callee cannot hold on to the pointer, but nothing says it
            // does not write through it before returning -- an out parameter
            // is exactly this. The slot's value is unknown after the call.
            LLVM_DEBUG(llvm::dbgs() << "Aliasing Store: " << callOp << "\n");
            AliasingStoreOperations.insert(callOp.getOperation());
          }
        }
        continue;
      }
      if (isa<LLVM::MemsetOp, LLVM::MemmoveOp, LLVM::MemcpyOp>(user)) {
        auto bytes = transferLength(user);
        // Without a length there is no telling what it touches, so anything it
        // writes may be this slot.
        OffsetTree touched =
            bytes ? tree.add(transferAccess(*bytes, user->getContext()), dl)
                  : OffsetTree::unknown();

        if (transferDest(user) == val) {
          switch (idx.matches(touched, dl)) {
          case Match::Exact:
            // Writing exactly this slot writes a value that is known when what
            // it was written from is: the bytes of a fill, when they are zero,
            // or whatever the copy read.
            if (auto ms = dyn_cast<LLVM::MemsetOp>(user)) {
              APInt byte;
              if (matchPattern(ms.getVal(), m_ConstantInt(&byte)) &&
                  byte.isZero()) {
                LLVM_DEBUG(llvm::dbgs() << "Matching Store: " << *user << "\n");
                allStoreOps.insert(user);
                break;
              }
            } else {
              LLVM_DEBUG(llvm::dbgs() << "Matching Store: " << *user << "\n");
              allStoreOps.insert(user);
              break;
            }
            LLVM_FALLTHROUGH;
          case Match::Contains:
          case Match::Maybe:
            // A fill reaching over the whole slot writes all of it whatever
            // else it reaches besides, and zero bytes are a zero of whatever
            // the slot holds however far they run past it.
            if (auto ms = dyn_cast<LLVM::MemsetOp>(user)) {
              APInt byte;
              if (touched.containsAt(idx, dl) &&
                  matchPattern(ms.getVal(), m_ConstantInt(&byte)) &&
                  byte.isZero()) {
                LLVM_DEBUG(llvm::dbgs() << "Matching Store: " << *user << "\n");
                allStoreOps.insert(user);
                break;
              }
            }
            LLVM_DEBUG(llvm::dbgs() << "Aliasing Store: " << *user << "\n");
            AliasingStoreOperations.insert(user);
            break;
          case Match::None:
            break;
          }
        }

        // Reading exactly this slot is reading the value in it, whatever the
        // copy then does with it.
        if (transferSource(user) == val &&
            idx.matches(touched, dl) == Match::Exact) {
          LLVM_DEBUG(llvm::dbgs() << "Matching Load: " << *user << "\n");
          transferLoads.insert(user);
        }
        continue;
      }
      if (isa<mlir::memref::DeallocOp>(user)) {
        continue;
      }

      LLVM_DEBUG(llvm::dbgs() << "Unknown, potential store: " << *user << "\n");
      AliasingStoreOperations.insert(user);
      captured = true;
    }
  }
  if (SharedMemAddr)
    AI.getDefiningOp()->getParentOp()->walk([&](mlir::NVVM::BarrierOp op) {
      LLVM_DEBUG(llvm::dbgs() << "Unknown, potential store: " << *op << "\n");
      AliasingStoreOperations.insert(op);
    });

  if (captured) {
    if (capturedAliasing.count(AI.getDefiningOp()) == 0) {
      SmallVector<Operation *> capEffects;
      AI.getDefiningOp()->getParentOp()->walk([&](Operation *op) {
        bool opMayHaveEffect = false;
        if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
          return;
        if (auto callOp = dyn_cast<mlir::LLVM::CallOp>(op)) {
          if (callOp.getCallee() && (*callOp.getCallee() == "printf")) {
            return;
          }
        }
        SmallVector<MemoryEffects::EffectInstance, 1> effects;
        collectEffects(op, effects, /*considerBarrier*/ false);

        for (auto effect : effects) {
          // If op causes EffectType on a potentially aliasing location for
          // memOp, mark as having the effect.
          if (isa<MemoryEffects::Write>(effect.getEffect())) {
            if (!mayAlias(effect.getEffect(), AI)) {
              continue;
            }
            opMayHaveEffect = true;
            break;
          }
        }
        if (opMayHaveEffect) {
          capEffects.push_back(op);
        }
      });

      capturedAliasing[AI.getDefiningOp()] = capEffects;
    }

    for (auto op : capturedAliasing[AI.getDefiningOp()]) {
      if (allStoreOps.count(op))
        continue;
      LLVM_DEBUG(llvm::dbgs() << "Potential Op ith Effect: " << *op << "\n");
      AliasingStoreOperations.insert(op);
    }
  }

  if (loadOps.size() == 0 && transferLoads.size() == 0) {
    return changed;
  }

  assert(elType);

  assert(AI.getDefiningOp());
  Region *parentAI = AI.getDefiningOp()->getParentRegion();
  assert(parentAI);

  // A list of all regions which contain loads to be replaced.
  SmallPtrSet<Region *, 4> ContainsLoadingOperation;
  {
    SmallVector<Region *> todo;
    for (auto *load : loadOps) {
      todo.push_back(load->getParentRegion());
    }
    for (auto *load : transferLoads) {
      todo.push_back(load->getParentRegion());
    }
    while (todo.size()) {
      auto *op = todo.back();
      todo.pop_back();
      if (ContainsLoadingOperation.contains(op))
        continue;
      if (op == parentAI)
        continue;
      ContainsLoadingOperation.insert(op);
      auto *parent = op->getParentRegion();
      assert(parent);
      todo.push_back(parent);
    }
  }

  // List of operations which may store that are not storeops
  SmallPtrSet<Operation *, 4> StoringOperations;
  SmallPtrSet<Block *, 4> StoringBlocks;
  {
    std::deque<Block *> todo;
    for (const auto &pair : allStoreOps) {
      LLVM_DEBUG(llvm::dbgs() << " storing operation: " << *pair << "\n");
      todo.push_back(pair->getBlock());
    }
    for (auto *op : AliasingStoreOperations) {
      StoringOperations.insert(op);
      LLVM_DEBUG(llvm::dbgs()
                 << " aliasing storing operation: " << *op << "\n");
      todo.push_back(op->getBlock());
    }
    while (todo.size()) {
      auto *block = todo.front();
      assert(block);
      todo.pop_front();
      StoringBlocks.insert(block);
      LLVM_DEBUG(llvm::dbgs() << " initial storing block: " << block << "\n");
      if (auto *op = block->getParentOp()) {
        StoringOperations.insert(op);
        if (auto *next = op->getBlock()) {
          StoringBlocks.insert(next);
          LLVM_DEBUG(llvm::dbgs()
                     << " derived storing block: " << next << "\n");
          todo.push_back(next);
        }
      }
    }
  }

  ReplacementHandler metaMap(elType);

  // Last value stored in an individual block and the operation which stored it
  BlockMap &valueAtEndOfBlock = metaMap.valueAtEndOfBlock;

  // Last value stored in an individual block and the operation which stored it
  BlockMap &valueAtStartOfBlock = metaMap.valueAtStartOfBlock;

  SmallVector<std::pair<Value, ValueOrPlaceholder *>> &replaceableValues =
      metaMap.replaceableValues;

  std::map<Block *, BlockArgument> blocksWithAddedArgs;

  auto *emptyValue = metaMap.get(nullptr);

  // What a load takes out of the value in the slot: the whole of it, or the
  // piece of it the load names.
  auto valueFor = [&](Value orig, Value slotValue) -> Value {
    auto contained = containedLoads.find(orig.getDefiningOp());
    if (contained == containedLoads.end())
      return castToType(elType, slotValue, orig.getDefiningOp());
    SmallVector<int64_t> path;
    bool found =
        extractPath(elType, contained->second, orig.getType(), dl, path);
    (void)found;
    assert(found && "a piece that was named cannot be reached");
    OpBuilder builder(orig.getDefiningOp());
    Value piece = slotValue.getType() == elType
                      ? slotValue
                      : castToType(elType, slotValue, orig.getDefiningOp());
    if (!path.empty())
      piece =
          LLVM::ExtractValueOp::create(builder, orig.getLoc(), slotValue, path);
    return castToType(orig.getType(), piece, orig.getDefiningOp());
  };

  auto replaceValue =
      [&](Value orig, ValueOrPlaceholder *replacement) -> ValueOrPlaceholder * {
    assert(replacement);
    replacement->materialize(/*full*/ false);
    assert(orig.getType() == elType ||
           containedLoads.count(orig.getDefiningOp()));
    if (replacement->overwritten) {
      loadOps.erase(orig.getDefiningOp());
      return metaMap.get(orig);
    } else if (replacement->val) {
      changed = true;
      assert(orig != replacement->val);
      auto castVal = valueFor(orig, replacement->val);
      LLVM_DEBUG(llvm::dbgs()
                 << " replaced " << orig << " with " << castVal << "\n");
      metaMap.replaceValue(orig, castVal);
      // Record this to erase later.
      loadOpsToErase.push_back(orig.getDefiningOp());
      loadOps.erase(orig.getDefiningOp());
      return replacement;
    } else {
      assert(replacement);
      SmallPtrSet<Block *, 1> requirements;
      if (replacement->definedWithArg(requirements)) {
        replaceableValues.push_back(std::make_pair(orig, replacement));
        assert(replaceableValues.back().second);
      }
      return metaMap.get(orig);
    }
  };

  // A transfer reading the whole of a slot consumes the value in it, so once
  // that value is known the copy is a store of it into the destination. Padding
  // bytes become undef rather than being copied, the same trade LLVM's SROA
  // makes when it forwards through a copy.
  auto replaceTransfer = [&](Operation *op, ValueOrPlaceholder *incoming) {
    incoming->materialize(/*full*/ false);
    // Whatever reaches the slot covers exactly the bytes the copy moves, since
    // that is what naming the same slot means.
    if (incoming->overwritten || !incoming->val)
      return;
    OpBuilder builder(op);
    LLVM::StoreOp::create(builder, op->getLoc(), incoming->val,
                          transferDest(op), transferAlignment(op, 0));
    LLVM_DEBUG(llvm::dbgs() << " replaced " << *op << " with a store of "
                            << incoming->val << "\n");
    loadOpsToErase.push_back(op);
    changed = true;
  };

  // Start by setting valueAtEndOfBlock to the last store directly in that block
  // Note that this may miss a store within a region of an operation in that
  // block
  // endRequires denotes whether this value is needed at the end of the block
  // (yield)
  std::function<void(Block &, ValueOrPlaceholder *)> handleBlock =
      [&](Block &block, ValueOrPlaceholder *lastVal) {
        valueAtStartOfBlock.emplace(&block, lastVal);
        SmallVector<Operation *, 10> ops;
        for (auto &a : block) {
          ops.push_back(&a);
        }
        LLVM_DEBUG(llvm::dbgs()
                       << "\nstarting block: lastVal=" << *lastVal << "\n";
                   block.print(llvm::dbgs()); llvm::dbgs() << "\n";);
        for (auto *a : ops) {
          if (StoringOperations.count(a)) {
            // erase a, in case overwritten later in metamap replacement/lookup.
            StoringOperations.erase(a);
            if (auto exOp = dyn_cast<mlir::scf::ExecuteRegionOp>(a)) {
              for (auto &B : exOp.getRegion())
                handleBlock(B, (&B == &exOp.getRegion().front())
                                   ? lastVal
                                   : metaMap.get(&B));
              lastVal = metaMap.get(exOp);
              continue;
            } else if (auto ifOp = dyn_cast<mlir::scf::IfOp>(a)) {
              handleBlock(*ifOp.getThenRegion().begin(), lastVal);
              if (ifOp.getElseRegion().getBlocks().size()) {
                handleBlock(*ifOp.getElseRegion().begin(), lastVal);
                lastVal = metaMap.get(ifOp, emptyValue);
              } else {
                lastVal = metaMap.get(ifOp, lastVal);
              }
              continue;
            } else if (auto ifOp = dyn_cast<mlir::affine::AffineIfOp>(a)) {
              handleBlock(*ifOp.getThenRegion().begin(), lastVal);
              if (ifOp.getElseRegion().getBlocks().size()) {
                handleBlock(*ifOp.getElseRegion().begin(), lastVal);
                lastVal = metaMap.get(ifOp, emptyValue);
              } else {
                lastVal = metaMap.get(ifOp, lastVal);
              }
              continue;
            }
            LLVM_DEBUG(llvm::dbgs() << "erased store due to: " << *a << "\n");

            for (auto &R : a->getRegions())
              for (auto &B : R)
                handleBlock(B,
                            (&B == &R.front()) ? emptyValue : metaMap.get(&B));
            lastVal = emptyValue;
          } else if (loadOps.count(a)) {
            Value loadOp = a->getResult(0);
            auto *read = replaceValue(loadOp, lastVal);
            // What a load of the whole slot read is what is in it from here
            // on; what a load of a piece read says nothing of the rest.
            if (!containedLoads.count(a))
              lastVal = read;
          } else if (transferLoads.count(a)) {
            replaceTransfer(a, lastVal);
          } else if (isa<LLVM::MemcpyOp, LLVM::MemmoveOp, LLVM::MemsetOp>(a)) {
            // A transfer writing the slot leaves what it wrote there: a copy
            // leaves what it read, which a load of that says, and a fill of
            // zero bytes leaves a zero of whatever the slot holds. Anything
            // else it wrote is something this cannot name, and stands as an
            // overwrite.
            if (allStoreOps.count(a)) {
              OpBuilder builder(a);
              // Of what the slot holds, since that is what everything threaded
              // through it is of; of as many bytes as were moved when nothing
              // has said yet what those bytes are.
              Type moved = elType ? elType
                                  : IntegerType::get(a->getContext(),
                                                     *transferLength(a) * 8);
              Value written;
              if (!LLVM::isCompatibleType(moved)) {
                // Nothing to name it with.
              } else if (Value src = transferSource(a)) {
                written = LLVM::LoadOp::create(builder, a->getLoc(), moved, src,
                                               transferAlignment(a, 1));
                transferReads.push_back(written.getDefiningOp());
              } else {
                written = LLVM::ZeroOp::create(builder, a->getLoc(), moved);
              }
              lastVal = written ? metaMap.get(written) : emptyValue;
            }
          } else if (auto storeOp = dyn_cast<enzyme::StoreLikeInterface>(a)) {
            if (allStoreOps.count(storeOp)) {
              lastVal = metaMap.get(storeOp.getStoredValue());
            }
          } else {
            // since not storing operation the value at the start of every block
            // is lastVal. However, if lastVal contains no information, use the
            // inductive block arg in case we can do load to load forwarding.
            for (auto &R : a->getRegions())
              for (auto &B : R)
                handleBlock(
                    B, lastVal->overwritten
                           ? ((&B == &R.front()) ? lastVal : metaMap.get(&B))
                           : lastVal);
          }
        }
        LLVM_DEBUG(llvm::dbgs() << " ending block: "; block.print(llvm::dbgs());
                   llvm::dbgs() << " with val:" << *lastVal << "\n";);
        assert(lastVal);
        valueAtEndOfBlock.emplace(&block, lastVal);
      };

  {
    assert(AI.getDefiningOp());
    SmallVector<Block *> todo = {AI.getDefiningOp()->getBlock()};
    SmallPtrSet<Block *, 2> done;
    while (todo.size()) {
      Block *cur = todo.back();
      todo.pop_back();
      if (done.contains(cur))
        continue;
      done.insert(cur);
      if (cur == AI.getDefiningOp()->getBlock())
        handleBlock(*cur, emptyValue);
      else
        handleBlock(*cur, metaMap.get(cur));
      for (auto *B : cur->getSuccessors())
        todo.push_back(B);
    }
  }

  if (replaceableValues.size() == 0)
    return changed;

  // Preserve only values which could inductively be replaced
  // by injecting relevant block arguments
  SmallPtrSet<Block *, 1> PotentiallyHelpfulArgs;
  for (auto &pair : replaceableValues) {
    SmallPtrSet<Block *, 1> requirements;
    bool _tmp = pair.second->definedWithArg(requirements);
    (void)_tmp;
    assert(_tmp);
    PotentiallyHelpfulArgs.insert(requirements.begin(), requirements.end());
  }

  enum class Legality {
    Unknown = 0,
    Illegal = 2,
    Required = 1,
  };
  // Map of block with potential arg added to the legality of adding that
  // argument
  std::map<Block *, Legality> PotentialArgs;
  // Map of block with potential arg added to the users which would require it
  // to compute their end value.
  std::map<Block *, std::set<Block *>> UserMap;

  // Map of block with potential arg added to the blocks it itself would
  // require.
  std::map<Block *, std::set<Block *>> RequirementMap;

  std::deque<Block *> todo(PotentiallyHelpfulArgs.begin(),
                           PotentiallyHelpfulArgs.end());
  while (todo.size()) {
    auto *block = todo.back();
    todo.pop_back();

    if (PotentialArgs.find(block) != PotentialArgs.end())
      continue;
    PotentialArgs[block] = Legality::Unknown;

    SetVector<Block *> prepred(block->getPredecessors().begin(),
                               block->getPredecessors().end());
    assert(prepred.size());

    for (auto *Pred : prepred) {

      auto endFind = valueAtEndOfBlock.find(Pred);
      assert(endFind != valueAtEndOfBlock.end());

      // Only handle known termination blocks
      if (!isa<cf::BranchOp, cf::CondBranchOp, cf::SwitchOp>(
              Pred->getTerminator())) {
        PotentialArgs[block] = Legality::Illegal;
        break;
      }

      SmallPtrSet<Block *, 1> requirements;
      if (endFind->second->definedWithArg(requirements)) {
        for (auto *r : requirements) {
          todo.push_back(r);
          UserMap[r].insert(block);
          RequirementMap[block].insert(r);
        }
      } else {
        PotentialArgs[block] = Legality::Illegal;
        break;
      }
    }
  }

  // Mark all blocks which may have an illegal predecessor
  for (auto &pair : PotentialArgs)
    if (pair.second == Legality::Illegal)
      for (auto *next : UserMap[pair.first])
        todo.push_back(next);

  while (todo.size()) {
    auto *block = todo.back();
    todo.pop_back();
    if (PotentialArgs[block] == Legality::Illegal)
      continue;
    PotentialArgs[block] = Legality::Illegal;
    for (auto *next : UserMap[block])
      todo.push_back(next);
  }

  // Go through all loading ops and confirm that all the requirements are met
  SmallVector<std::pair<Value, ValueOrPlaceholder *>> nextReplacements;
  for (auto pair : replaceableValues) {
    SmallPtrSet<Block *, 1> requirements;
    bool _tmp = pair.second->definedWithArg(requirements);
    (void)_tmp;
    assert(_tmp);
    bool illegal = false;
    for (auto *r : requirements) {
      if (PotentialArgs[r] == Legality::Illegal) {
        illegal = true;
        break;
      }
    }
    // If illegal, remove this from the list of things to update.
    if (illegal) {
      continue;
    }
    // Otherwise mark that block arg, and all of its dependencies as required.
    for (auto *r : requirements)
      todo.push_back(r);

    while (todo.size()) {
      auto *block = todo.back();
      todo.pop_back();
      if (PotentialArgs[block] == Legality::Required)
        continue;
      PotentialArgs[block] = Legality::Required;
      for (auto *prev : RequirementMap[block])
        todo.push_back(prev);
    }
    nextReplacements.push_back(pair);
  }
  replaceableValues = nextReplacements;

  if (replaceableValues.size() == 0)
    return changed;

  SmallVector<Block *, 1> Legal;
  for (auto &pair : PotentialArgs)
    if (pair.second == Legality::Required) {
      Legal.push_back(pair.first);
    }

  for (auto *block : Legal) {
    auto startFound = valueAtStartOfBlock.find(block);
    (void)startFound;
    assert(startFound != valueAtStartOfBlock.end());
    assert(startFound->second->valueAtStart == block);
    auto arg = block->addArgument(elType, loc);
    auto *argVal = metaMap.get(arg);
    valueAtStartOfBlock[block] = argVal;
    blocksWithAddedArgs[block] = arg;
  }

  for (auto &pair : replaceableValues) {
    assert(pair.first);
    assert(pair.second);
    Value val = pair.second->materialize(true);
    assert(val);

    changed = true;
    assert(pair.first != val);
    assert(val.getType() == elType);
    val = valueFor(pair.first, val);
    assert(pair.first.getType() == val.getType() && "mismatched load type");
    LLVM_DEBUG(llvm::dbgs()
               << " replaced " << pair.first << " with " << val << "\n");
    metaMap.replaceValue(pair.first, val);

    // Record this to erase later.
    loadOpsToErase.push_back(pair.first.getDefiningOp());
    loadOps.erase(pair.first.getDefiningOp());
  }
  replaceableValues.clear();

  for (auto &pair : blocksWithAddedArgs) {
    Block *block = pair.first;
    assert(valueAtStartOfBlock.find(block) != valueAtStartOfBlock.end());

    Value maybeblockArg =
        valueAtStartOfBlock.find(block)->second->materialize(false);
    auto blockArg = dyn_cast<BlockArgument>(maybeblockArg);
    (void)blockArg;
    assert(blockArg && blockArg.getOwner() == block);

    SetVector<Block *> prepred(block->getPredecessors().begin(),
                               block->getPredecessors().end());
    for (auto *pred : prepred) {
      assert(pred && "Null predecessor");
      assert(valueAtEndOfBlock.find(pred) != valueAtEndOfBlock.end());
      assert(valueAtEndOfBlock.find(pred)->second);
      mlir::Value pval =
          valueAtEndOfBlock.find(pred)->second->materialize(true);
      // What reaches the end of a predecessor is of the slot's extent but need
      // not be spelled as the slot is, and what a block argument takes must be.
      if (pval && pval.getType() != elType)
        pval = castToType(elType, pval, pred->getTerminator());
      if (!pval || pval.getType() != elType) {
        if (auto fn =
                AI.getDefiningOp()->getParentOfType<FunctionOpInterface>())
          fn.dump();
        pred->dump();
        llvm::errs() << "pval: " << *valueAtEndOfBlock.find(pred)->second
                     << " AI: " << AI << "\n";
        if (pval)
          llvm::errs() << " mat pval: " << pval << "\n";
      }
      assert(pval && "Null last stored");
      assert(pval.getType() == elType);
      assert(pred->getTerminator());

      assert(blockArg.getOwner() == block);
      if (auto op = dyn_cast<cf::BranchOp>(pred->getTerminator())) {
        mlir::OpBuilder subbuilder(op.getOperation());
        std::vector<Value> args(op.getOperands().begin(),
                                op.getOperands().end());
        args.push_back(pval);
        cf::BranchOp::create(subbuilder, op.getLoc(), op.getDest(), args);
        op.erase();
      } else if (auto op = dyn_cast<cf::CondBranchOp>(pred->getTerminator())) {

        mlir::OpBuilder subbuilder(op.getOperation());
        std::vector<Value> trueargs(op.getTrueOperands().begin(),
                                    op.getTrueOperands().end());
        std::vector<Value> falseargs(op.getFalseOperands().begin(),
                                     op.getFalseOperands().end());
        if (op.getTrueDest() == block) {
          trueargs.push_back(pval);
        }
        if (op.getFalseDest() == block) {
          falseargs.push_back(pval);
        }
        cf::CondBranchOp::create(subbuilder, op.getLoc(), op.getCondition(),
                                 op.getTrueDest(), trueargs, op.getFalseDest(),
                                 falseargs);
        op.erase();
      } else if (auto op = dyn_cast<cf::SwitchOp>(pred->getTerminator())) {
        mlir::OpBuilder builder(op.getOperation());
        SmallVector<Value> defaultOps(op.getDefaultOperands().begin(),
                                      op.getDefaultOperands().end());

        if (op.getDefaultDestination() == block)
          defaultOps.push_back(pval);

        SmallVector<SmallVector<Value>> cases;
        for (auto pair : llvm::enumerate(op.getCaseDestinations())) {
          cases.emplace_back(op.getCaseOperands(pair.index()).begin(),
                             op.getCaseOperands(pair.index()).end());
          if (pair.value() == block) {
            cases.back().push_back(pval);
          }
        }
        SmallVector<ValueRange> vrange;
        for (auto &c : cases) {
          vrange.push_back(c);
        }
        cf::SwitchOp::create(builder, op.getLoc(), op.getFlag(),
                             op.getDefaultDestination(), defaultOps,
                             op.getCaseValuesAttr(), op.getCaseDestinations(),
                             vrange);
        op.erase();
      } else {
        llvm_unreachable("unknown pred branch");
      }
    }
  }

  removeRedundantBlockArgs(AI, elType, blocksWithAddedArgs);

  for (auto *loadOp : llvm::make_early_inc_range(loadOps)) {
    assert(loadOp);
    if (loadOp->getResult(0).use_empty()) {
      loadOpsToErase.push_back(loadOp);
      loadOps.erase(loadOp);
    }
  }
  return changed;
}

bool isPromotable(mlir::Value AI, SymbolTableCollection &symbolTables) {
  std::deque<mlir::Value> list = {AI};

  while (list.size()) {
    auto val = list.front();
    list.pop_front();

    for (auto *U : val.getUsers()) {
      if (isa<LLVM::LifetimeStartOp, LLVM::LifetimeEndOp>(U)) {
        continue;
      }
      if (auto LO = dyn_cast<memref::LoadOp>(U)) {
        continue;
      } else if (auto LO = dyn_cast<LLVM::LoadOp>(U)) {
        continue;
      } else if (auto LO = dyn_cast<affine::AffineLoadOp>(U)) {
        continue;
      } else if (auto SO = dyn_cast<enzyme::StoreLikeInterface>(U)) {
        // Only a store *into* val keeps it promotable; storing val itself puts
        // the pointer into memory, from where anything may reach it.
        if (SO.getStoredValue() != val)
          continue;
        LLVM_DEBUG(llvm::dbgs()
                   << "non promotable " << AI << " due to " << *U << "\n");
        return false;
      } else if (isa<LLVM::MemcpyOp, LLVM::MemmoveOp, LLVM::MemsetOp>(U)) {
        // The forwarding treats these as an opaque write of the destination and
        // a read of the source, neither of which lets the pointer escape.
        continue;
      } else if (isa<memref::DeallocOp>(U)) {
        continue;
      } else if (auto callOp = dyn_cast<CallOpInterface>(U)) {
        if (auto sym = dyn_cast<SymbolRefAttr>(callOp.getCallableForCallee())) {
          if (StringAttr callee = sym.getLeafReference())
            if (isCallNonCapturing(callOp, val, symbolTables) ||
                getNonCapturingFunctions().count(callee.str()))
              continue;
        }
      } else if (auto CO = dyn_cast<memref::CastOp>(U)) {
        list.push_back(CO);
      } else if (auto CO = dyn_cast<Memref2PointerOp>(U)) {
        list.push_back(CO);
      } else if (auto CO = dyn_cast<Pointer2MemrefOp>(U)) {
        list.push_back(CO);
      } else if (auto CO = dyn_cast<LLVM::GEPOp>(U)) {
        // Accesses through an offset of the allocation are what the forwarding
        // marks as indexed: they never match a slot and only block it.
        list.push_back(CO);
      } else {
        LLVM_DEBUG(llvm::dbgs()
                   << "non promotable " << AI << " due to " << *U << "\n");
        return false;
      }
    }
  }
  return true;
}

std::vector<OffsetTree> getLastStored(mlir::Value AI, const DataLayout &dl) {
  // Where an access lands and how much of it it reads: two spellings of one
  // extent are one slot, but the whole of something and the first field of it
  // start in the same place and are not.
  auto slotOf = [&](const OffsetTree &tree) {
    return std::make_pair(typeSize(tree.getBase(), dl).value_or(0), tree);
  };
  std::map<std::pair<uint64_t, OffsetTree>, unsigned> lastStored;

  std::deque<std::pair<mlir::Value, OffsetTree>> list = {{AI, OffsetTree()}};

  while (list.size()) {
    auto [val, tree] = list.front();
    list.pop_front();
    for (auto *U : val.getUsers()) {
      if (auto SO = dyn_cast<memref::StoreOp>(U)) {
        lastStored[slotOf(
            tree.add(accessOffsets(SO.getValue().getType(), SO.getMemRef(),
                                   SO.getIndices(), dl),
                     dl))]++;
      } else if (auto LO = dyn_cast<memref::LoadOp>(U)) {
        lastStored[slotOf(tree.add(
            accessOffsets(LO.getType(), LO.getMemRef(), LO.getIndices(), dl),
            dl))]++;
      } else if (auto SO = dyn_cast<affine::AffineStoreOp>(U)) {
        lastStored[slotOf(
            tree.add(accessOffsets(SO.getValue().getType(), SO.getMemRef(),
                                   SO.getAffineMapAttr().getValue(),
                                   SO.getMapOperands(), dl),
                     dl))]++;
      } else if (auto LO = dyn_cast<affine::AffineLoadOp>(U)) {
        lastStored[slotOf(
            tree.add(accessOffsets(LO.getType(), LO.getMemRef(),
                                   LO.getAffineMapAttr().getValue(),
                                   LO.getMapOperands(), dl),
                     dl))]++;
      } else if (auto SO = dyn_cast<LLVM::StoreOp>(U)) {
        lastStored[slotOf(
            tree.add(accessOffsets(SO.getValue().getType()), dl))]++;
      } else if (auto LO = dyn_cast<LLVM::LoadOp>(U)) {
        lastStored[slotOf(tree.add(accessOffsets(LO.getType()), dl))]++;
      } else if (isa<LLVM::MemcpyOp, LLVM::MemmoveOp, LLVM::MemsetOp>(U)) {
        // What a transfer reads or fills is a slot like any other, and the
        // forwarding can only be asked about slots it is told of.
        if (auto bytes = transferLength(U))
          lastStored[slotOf(
              tree.add(transferAccess(*bytes, U->getContext()), dl))]++;
      } else if (auto GO = dyn_cast<LLVM::GEPOp>(U)) {
        list.emplace_back(GO, tree.add(gepOffsets(GO, dl), dl));
      } else if (isa<memref::CastOp, Memref2PointerOp, Pointer2MemrefOp,
                     LLVM::BitcastOp, LLVM::AddrSpaceCastOp>(U)) {
        list.emplace_back(U->getResult(0), tree);
      }
    }
  }

  // A slot is worth trying when more than one access names it -- and an access
  // of a piece of a slot is an access of that slot too, since the piece can be
  // taken out of what is in it. It counts the other way round as well, since
  // an access reaching over the whole of a slot may be all that is ever said
  // about what is in it.
  std::vector<OffsetTree> todo;
  for (auto &pair : lastStored) {
    unsigned count = pair.second;
    for (auto &other : lastStored)
      if (&other != &pair &&
          (pair.first.second.containsAt(other.first.second, dl) ||
           other.first.second.containsAt(pair.first.second, dl)))
        count += other.second;
    if (count > 1)
      todo.push_back(pair.first.second);
  }
  return todo;
}

void PolygeistMem2Reg::runOnOperation() {
  auto *f = getOperation();

  // Resolving a callee by name walks every symbol the module has, and there is
  // a callee to resolve for each call each allocation reaches. Nothing here
  // adds or removes a symbol, so one collection serves the whole run.
  SymbolTableCollection symbolTables;

  // Variable indicating that a memref has had a load removed
  // and or been deleted. Because there can be memrefs of
  // memrefs etc, we may need to do multiple passes (first
  // to eliminate the outermost one, then inner ones)
  bool changed;
  do {
    changed = false;

    // A list of memref's that are potentially dead / could be eliminated.
    SmallPtrSet<Value, 4> memrefsToErase;

    // Load op's whose results were replaced by those forwarded from stores.
    SmallVector<Operation *, 8> loadOpsToErase;

    // Walk all load's and perform store to load forwarding.
    SmallVector<mlir::Value, 4> toPromote;
    f->walk([&](mlir::memref::AllocaOp AI) {
      if (isPromotable(AI, symbolTables)) {
        toPromote.push_back(AI);
      }
    });
    f->walk([&](mlir::memref::AllocOp AI) {
      if (isPromotable(AI, symbolTables)) {
        toPromote.push_back(AI);
      }
    });
    f->walk([&](LLVM::AllocaOp AI) {
      if (isPromotable(AI, symbolTables)) {
        toPromote.push_back(AI);
      }
    });
    f->walk([&](memref::GetGlobalOp AI) {
      if (isPromotable(AI, symbolTables)) {
        toPromote.push_back(AI);
      }
    });
    DenseMap<Operation *, SmallVector<Operation *>> capturedAliasing;
    for (auto AI : toPromote) {
      LLVM_DEBUG(llvm::dbgs() << " attempting to promote " << AI << "\n");
      // A nested region may carry a layout of its own, so the sizes an offset
      // is counted in are the ones in force where the allocation is.
      auto lastStored =
          getLastStored(AI, DataLayout::closest(AI.getDefiningOp()));
      for (const auto &vec : lastStored) {
        LLVM_DEBUG(llvm::dbgs() << " + forwarding vec to promote {" << vec
                                << "} of " << AI << "\n");
        // llvm::errs() << " PRE " << AI << "\n";
        // f.dump();
        changed |= forwardStoreToLoad(AI, vec, loadOpsToErase, capturedAliasing,
                                      symbolTables);
        // llvm::errs() << " POST " << AI << "\n";
        // f.dump();
      }
      if (!AI.getDefiningOp<memref::GetGlobalOp>())
        memrefsToErase.insert(AI);
    }

    // Erase all load op's whose results were replaced with store fwd'ed ones.
    // One read may be forwarded out of more than one slot -- the whole of
    // something and a piece of it are both slots it reads -- and is only there
    // to be erased once.
    SmallPtrSet<Operation *, 8> erased;
    for (auto *loadOp : loadOpsToErase) {
      if (!erased.insert(loadOp).second)
        continue;
      changed = true;
      loadOp->erase();
    }

    // Check if the store fwd'ed memrefs are now left with only stores and can
    // thus be completely deleted. Note: the canonicalize pass should be able
    // to do this as well, but we'll do it here since we collected these anyway.
    for (auto memref : memrefsToErase) {

      // If the memref hasn't been alloc'ed in this function, skip.
      Operation *defOp = memref.getDefiningOp();
      if (!defOp ||
          !(isa<memref::AllocOp>(defOp) || isa<memref::AllocaOp>(defOp) ||
            isa<LLVM::AllocaOp>(defOp)))
        // TODO: if the memref was returned by a 'call' operation, we
        // could still erase it if the call had no side-effects.
        continue;

      std::deque<mlir::Value> list = {memref};
      std::vector<mlir::Operation *> toErase;
      bool error = false;
      while (list.size()) {
        auto val = list.front();
        list.pop_front();

        for (auto *U : val.getUsers()) {
          if (auto SO = dyn_cast<enzyme::StoreLikeInterface>(U)) {
            if (SO.getStoredValue() == val) {
              error = true;
              break;
            }
            toErase.push_back(U);
          } else if (isa<LLVM::MemsetOp, LLVM::MemcpyOp, LLVM::MemmoveOp>(U)) {
            // Writing into the allocation goes away with it; reading out of it
            // writes somewhere else and has to stay.
            if (transferSource(U) == val) {
              error = true;
              break;
            }
            toErase.push_back(U);
          } else if (isa<LLVM::LifetimeStartOp, LLVM::LifetimeEndOp>(U)) {
            toErase.push_back(U);
          } else if (isa<memref::DeallocOp>(U)) {
            toErase.push_back(U);
          } else if (isa<func::CallOp>(U) &&
                     cast<func::CallOp>(U).getCallee() == "free") {
            toErase.push_back(U);
          } else if (isa<memref::CastOp, Pointer2MemrefOp, Memref2PointerOp,
                         LLVM::GEPOp>(U)) {
            toErase.push_back(U);
            list.push_back(U->getResult(0));
          } else {
            error = true;
            break;
          }
        }
        if (error)
          break;
      }

      if (!error) {
        std::reverse(toErase.begin(), toErase.end());
        for (auto *user : toErase) {
          user->erase();
        }
        defOp->erase();
        changed = true;
      } else {
        // llvm::errs() << " failed to remove: " << memref << "\n";
      }
    }
  } while (changed);
}
