normalize affine arithmetic helpers

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -13,6 +13,8 @@
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/LogicalResult.h"
 
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
@@ -56,48 +58,19 @@ static LogicalResult verifyStaticWeights(ComputeOpTy computeOp, StringRef kind)
   return success();
 }
 
-static bool isConstantIndexLike(Value value) {
-  APInt constantValue;
-  return matchPattern(value, m_ConstantInt(&constantValue));
-}
-
-static bool isSupportedLaneAffineExpr(AffineExpr expr) {
-  switch (expr.getKind()) {
-  case AffineExprKind::Constant:
-  case AffineExprKind::DimId:    return true;
-  case AffineExprKind::SymbolId: return false;
-  case AffineExprKind::Add:      {
-    auto binaryExpr = cast<AffineBinaryOpExpr>(expr);
-    return isSupportedLaneAffineExpr(binaryExpr.getLHS()) && isSupportedLaneAffineExpr(binaryExpr.getRHS());
-  }
-  case AffineExprKind::Mul: {
-    auto binaryExpr = cast<AffineBinaryOpExpr>(expr);
-    return (isa<AffineConstantExpr>(binaryExpr.getLHS()) && isSupportedLaneAffineExpr(binaryExpr.getRHS()))
-        || (isa<AffineConstantExpr>(binaryExpr.getRHS()) && isSupportedLaneAffineExpr(binaryExpr.getLHS()));
-  }
-  case AffineExprKind::FloorDiv:
-  case AffineExprKind::CeilDiv:
-  case AffineExprKind::Mod:      {
-    auto binaryExpr = cast<AffineBinaryOpExpr>(expr);
-    return isa<AffineConstantExpr>(binaryExpr.getRHS()) && isSupportedLaneAffineExpr(binaryExpr.getLHS());
-  }
-  }
-  llvm_unreachable("unexpected affine expression kind");
-}
-
 static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) {
-  if (value == laneArg || isConstantIndexLike(value))
+  if (value == laneArg || matchConstantIndexValue(value))
     return true;
 
   auto affineApply = value.getDefiningOp<affine::AffineApplyOp>();
   if (affineApply) {
-    if (affineApply.getAffineMap().getNumResults() != 1 || affineApply.getAffineMap().getNumSymbols() != 0)
+    if (!isSingleResultSymbolFreeAffineMap(affineApply.getAffineMap()))
       return false;
     if (!llvm::all_of(affineApply.getMapOperands(),
                       [&](Value operand) { return isSupportedLaneOffsetExpr(operand, laneArg); })) {
       return false;
     }
-    return isSupportedLaneAffineExpr(affineApply.getAffineMap().getResult(0));
+    return isDimAndConstantAffineExpr(affineApply.getAffineMap().getResult(0));
   }
 
   auto extractOp = value.getDefiningOp<tensor::ExtractOp>();
@@ -112,8 +85,8 @@ static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) {
   auto addOp = value.getDefiningOp<arith::AddIOp>();
   if (!addOp)
     return false;
-  return (addOp.getLhs() == laneArg && isConstantIndexLike(addOp.getRhs()))
-      || (addOp.getRhs() == laneArg && isConstantIndexLike(addOp.getLhs()));
+  return (addOp.getLhs() == laneArg && matchConstantIndexValue(addOp.getRhs()))
+      || (addOp.getRhs() == laneArg && matchConstantIndexValue(addOp.getLhs()));
 }
 
 static LogicalResult

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp

@@ -21,6 +21,7 @@
 
 #include "MaterializeMergeSchedule.hpp"
 #include "Scheduling/ComputeInstanceUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -1053,7 +1054,7 @@ Value createAffineIndexValue(MaterializerState& state, const IndexedIndexPattern
   }
 
   AffineMap map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
-  return createAffineApplyOrFoldedConstant(state.rewriter, loc, map, ValueRange {index}, state.func);
+  return createOrFoldAffineApply(state.rewriter, loc, map, ValueRange {index}, state.func);
 }
 
 Value createIndexedIndexValue(
@@ -1295,47 +1296,16 @@ std::optional<unsigned> getLaneProjectedDim(ArrayRef<OpFoldResult> offsets, Valu
 }
 
 static FailureOr<int64_t> evaluateProjectedOffsetForLane(OpFoldResult value, Value laneArg, uint32_t lane) {
-  if (auto attr = dyn_cast<Attribute>(value)) {
-    auto intAttr = dyn_cast<IntegerAttr>(attr);
-    if (!intAttr)
-      return failure();
-    return intAttr.getInt();
-  }
+  if (std::optional<int64_t> constant = matchConstantIndexValue(value))
+    return *constant;
 
   Value current = cast<Value>(value);
   if (current == laneArg)
     return static_cast<int64_t>(lane);
 
-  if (auto constant = current.getDefiningOp<arith::ConstantIndexOp>())
-    return constant.value();
-
-  if (auto constant = current.getDefiningOp<arith::ConstantOp>())
-    if (auto intAttr = dyn_cast<IntegerAttr>(constant.getValue()))
-      return intAttr.getInt();
-
-  if (auto affineApply = current.getDefiningOp<affine::AffineApplyOp>()) {
-    AffineMap map = affineApply.getAffineMap();
-    if (map.getNumResults() != 1)
-      return failure();
-
-    SmallVector<Attribute, 4> operandConstants;
-    operandConstants.reserve(affineApply.getMapOperands().size());
-    for (Value operand : affineApply.getMapOperands()) {
-      FailureOr<int64_t> folded = evaluateProjectedOffsetForLane(operand, laneArg, lane);
-      if (failed(folded))
-        return failure();
-      operandConstants.push_back(IntegerAttr::get(IndexType::get(current.getContext()), *folded));
-    }
-
-    SmallVector<Attribute, 1> foldedResults;
-    if (failed(map.constantFold(operandConstants, foldedResults)) || foldedResults.size() != 1)
-      return failure();
-
-    auto constantResult = dyn_cast<IntegerAttr>(foldedResults.front());
-    if (!constantResult)
-      return failure();
-    return constantResult.getInt();
-  }
+  if (auto affineApply = current.getDefiningOp<affine::AffineApplyOp>())
+    return evaluateAffineApply(affineApply,
+                               [&](Value operand) { return evaluateProjectedOffsetForLane(operand, laneArg, lane); });
 
   return failure();
 }
@@ -3503,7 +3473,7 @@ Value createBatchRunFlatIndex(MaterializerState& state, MaterializedClass& targe
 
   int64_t laneCount = static_cast<int64_t>(targetClass.cpus.size());
   AffineMap map = AffineMap::get(/*dimCount=*/2, /*symbolCount=*/0, d0 * laneCount + d1);
-  return createAffineApplyOrFoldedConstant(state.rewriter, loc, map, ValueRange {slotIndex, *laneArg}, state.func);
+  return createOrFoldAffineApply(state.rewriter, loc, map, ValueRange {slotIndex, *laneArg}, state.func);
 }
 
 Value createBatchClassRunSourceLane(MaterializerState& state,

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp

@@ -21,6 +21,8 @@
 
 #include "ComputeGraph.hpp"
 #include "ComputeInstanceUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Support/TypeUtilities.hpp"
 
 namespace onnx_mlir {
@@ -148,57 +150,6 @@ static CrossbarWeight getOpaqueCrossbarWeight(Value value, std::optional<uint32_
   return weight;
 }
 
-static FailureOr<int64_t> evaluateAffineExpr(AffineExpr expr, ArrayRef<int64_t> dims, ArrayRef<int64_t> symbols) {
-  if (auto constant = dyn_cast<AffineConstantExpr>(expr))
-    return constant.getValue();
-  if (auto dim = dyn_cast<AffineDimExpr>(expr)) {
-    unsigned position = dim.getPosition();
-    if (position >= dims.size())
-      return failure();
-    return dims[position];
-  }
-  if (auto symbol = dyn_cast<AffineSymbolExpr>(expr)) {
-    unsigned position = symbol.getPosition();
-    if (position >= symbols.size())
-      return failure();
-    return symbols[position];
-  }
-
-  auto binary = dyn_cast<AffineBinaryOpExpr>(expr);
-  if (!binary)
-    return failure();
-
-  FailureOr<int64_t> lhs = evaluateAffineExpr(binary.getLHS(), dims, symbols);
-  FailureOr<int64_t> rhs = evaluateAffineExpr(binary.getRHS(), dims, symbols);
-  if (failed(lhs) || failed(rhs))
-    return failure();
-
-  auto floorDiv = [](int64_t value, int64_t divisor) -> FailureOr<int64_t> {
-    if (divisor <= 0)
-      return failure();
-    if (value >= 0)
-      return value / divisor;
-    return -((-value + divisor - 1) / divisor);
-  };
-
-  switch (binary.getKind()) {
-  case AffineExprKind::Add:      return *lhs + *rhs;
-  case AffineExprKind::Mul:      return *lhs * *rhs;
-  case AffineExprKind::FloorDiv: return floorDiv(*lhs, *rhs);
-  case AffineExprKind::CeilDiv:
-    if (*rhs <= 0)
-      return failure();
-    return (*lhs + *rhs - 1) / *rhs;
-  case AffineExprKind::Mod: {
-    FailureOr<int64_t> div = floorDiv(*lhs, *rhs);
-    if (failed(div))
-      return failure();
-    return *lhs - *div * *rhs;
-  }
-  default: return failure();
-  }
-}
-
 static FailureOr<int64_t>
 evaluateIndexLike(Value value, const DenseMap<Value, int64_t>& bindings, std::optional<uint32_t> lane, Value laneArg);
 
@@ -222,12 +173,8 @@ evaluateIndexLike(Value value, const DenseMap<Value, int64_t>& bindings, std::op
   if (auto it = bindings.find(value); it != bindings.end())
     return it->second;
 
-  if (auto constant = value.getDefiningOp<arith::ConstantIndexOp>())
-    return constant.value();
-
-  if (auto constant = value.getDefiningOp<arith::ConstantOp>())
-    if (auto intAttr = dyn_cast<IntegerAttr>(constant.getValue()))
-      return intAttr.getInt();
+  if (std::optional<int64_t> constant = matchConstantIndexValue(value))
+    return *constant;
 
   if (auto extract = value.getDefiningOp<tensor::ExtractOp>()) {
     auto constant = extract.getTensor().getDefiningOp<arith::ConstantOp>();
@@ -245,26 +192,11 @@ evaluateIndexLike(Value value, const DenseMap<Value, int64_t>& bindings, std::op
     return failure();
   }
 
-  auto affineApply = value.getDefiningOp<affine::AffineApplyOp>();
-  if (!affineApply)
-    return failure();
+  if (auto affineApply = value.getDefiningOp<affine::AffineApplyOp>())
+    return evaluateAffineApply(affineApply,
+                               [&](Value operand) { return evaluateIndexLike(operand, bindings, lane, laneArg); });
 
-  AffineMap map = affineApply.getAffineMap();
-  if (map.getNumResults() != 1)
-    return failure();
-
-  SmallVector<int64_t, 4> operands;
-  operands.reserve(affineApply.getMapOperands().size());
-  for (Value operand : affineApply.getMapOperands()) {
-    FailureOr<int64_t> folded = evaluateIndexLike(operand, bindings, lane, laneArg);
-    if (failed(folded))
-      return failure();
-    operands.push_back(*folded);
-  }
-
-  ArrayRef<int64_t> dims(operands.data(), map.getNumDims());
-  ArrayRef<int64_t> symbols(operands.data() + map.getNumDims(), map.getNumSymbols());
-  return evaluateAffineExpr(map.getResult(0), dims, symbols);
+  return failure();
 }
 
 static FailureOr<SmallVector<int64_t, 4>> evaluateIndexList(ArrayRef<OpFoldResult> values,