[Mlir-commits] [mlir] [mlir][affine] Add pass --affine-raise-from-memref (PR #114032)

[Clément Fournier]

================
@@ -0,0 +1,168 @@
+
+
+#include "mlir/Dialect/Affine/Analysis/Utils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Passes.h"
+#include "mlir/Dialect/Affine/Transforms/Transforms.h"
+#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/LogicalResult.h"
+#include <algorithm>
+#include <cstddef>
+#include <functional>
+#include <iterator>
+#include <memory>
+#include <optional>
+
+namespace mlir {
+namespace affine {
+#define GEN_PASS_DEF_RAISEMEMREFDIALECT
+#include "mlir/Dialect/Affine/Passes.h.inc"
+} // namespace affine
+} // namespace mlir
+
+#define DEBUG_TYPE "raise-memref-to-affine"
+
+using namespace mlir;
+using namespace mlir::affine;
+
+namespace {
+
+static std::optional<size_t>
+findInListOrAdd(Value value, llvm::SmallVectorImpl<Value> &dims,
+                const std::function<bool(Value)> &isValidElement) {
+
+  Value *loopIV = std::find(dims.begin(), dims.end(), value);
+  if (loopIV != dims.end()) {
+    // found an IV that already has an index
+    return {std::distance(dims.begin(), loopIV)};
+  }
+  if (isValidElement(value)) {
+    // push this IV in the parameters
+    size_t idx = dims.size();
+    dims.push_back(value);
+    return idx;
+  }
+  return std::nullopt;
+}
+
+static LogicalResult toAffineExpr(Value value, AffineExpr &result,
+                                  llvm::SmallVectorImpl<Value> &affineDims,
+                                  llvm::SmallVectorImpl<Value> &affineSymbols) {
+  using namespace matchers;
+  IntegerAttr::ValueType cst;
+  if (matchPattern(value, m_ConstantInt(&cst))) {
+    result = getAffineConstantExpr(cst.getSExtValue(), value.getContext());
+    return success();
+  }
+  Value lhs;
+  Value rhs;
+  if (matchPattern(value, m_Op<arith::AddIOp>(m_Any(&lhs), m_Any(&rhs))) ||
+      matchPattern(value, m_Op<arith::MulIOp>(m_Any(&lhs), m_Any(&rhs)))) {
+    AffineExpr lhsE;
+    AffineExpr rhsE;
+    if (succeeded(toAffineExpr(lhs, lhsE, affineDims, affineSymbols)) &&
+        succeeded(toAffineExpr(rhs, rhsE, affineDims, affineSymbols))) {
----------------
oowekyala wrote:

I understand your thought, although here I think recursion is the most natural way to write this so I would  keep it that way unless you really insist. There are other examples of recursion a bit everywhere, eg [here](https://github.com/oowekyala/llvm-project/blob/a722306e9a01f0e20655941c79f3d253e57aa022/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp#L180), and I would argue that the depth of the index expressions for memref access should be small enough (it likely has a simple form like `a + b * c`, which has depth 3). An iterative implementation would require an explicit stack and would be much more complex. 

https://github.com/llvm/llvm-project/pull/114032