[mlir][scf] Support simple symbolic expression without depending on AffineDialect to simply trivial loops.

[Peiming]

Summary

Remove dependence of AffineDialect

mlir/lib/Dialect/SCF/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRSCFDialect
   SCF.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/SCF
 
   DEPENDS
   MLIRSCFOpsIncGen
 
   LINK_LIBS PUBLIC
-  MLIRAffineDialect
   MLIRArithmeticDialect
   MLIRBufferizationDialect
   MLIRControlFlowDialect
   MLIRIR
   MLIRLoopLikeInterface
   MLIRSideEffectInterfaces
   )
 

mlir/lib/Dialect/SCF/IR/SCF.cpp

 //===- SCF.cpp - Structured Control Flow Operations -----------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/FunctionInterfaces.h"
LogicalResult matchAndRewrite(scf::ForOp forOp,
     return success();
   }
 };
 
 /// Util function that tries to compute a constant diff between u and l.
 /// Returns llvm::None when the difference between two AffineValueMap is
 /// dynamic.
 static Optional<int64_t> computeConstDiff(Value l, Value u) {
-  auto alb = l.getDefiningOp<AffineApplyOp>();
-  auto aub = u.getDefiningOp<AffineApplyOp>();
-  // ID map: (d0)->d0
-  auto id = AffineMap::getMultiDimIdentityMap(1, l.getContext());
-  auto lb = alb ? alb.getAffineValueMap() : AffineValueMap(id, l);
-  auto ub = aub ? aub.getAffineValueMap() : AffineValueMap(id, u);
-
-  AffineValueMap diffMap;
-  AffineValueMap::difference(ub, lb, &diffMap);
-  if (auto constDiff = diffMap.getResult(0).dyn_cast<AffineConstantExpr>())
-    return constDiff.getValue();
+  auto clb = l.getDefiningOp<arith::ConstantOp>();
+  auto cub = u.getDefiningOp<arith::ConstantOp>();
+  if (cub && clb) {
+    llvm::APInt lbValue = clb.getValue().cast<IntegerAttr>().getValue();
+    llvm::APInt ubValue = cub.getValue().cast<IntegerAttr>().getValue();
+    return (ubValue - lbValue).getSExtValue();
+  }
+
+  // Else a simple pattern match for x + c or c + x
+  llvm::APInt diff;
+  if (matchPattern(
+          u, m_Op<arith::AddIOp>(matchers::m_Val(l), m_ConstantInt(&diff))) ||
+      matchPattern(
+          u, m_Op<arith::AddIOp>(m_ConstantInt(&diff), matchers::m_Val(l))))
+    return diff.getSExtValue();
   return llvm::None;
 }
 
 /// Rewriting pattern that erases loops that are known not to iterate, replaces
 /// single-iteration loops with their bodies, and removes empty loops that
 /// iterate at least once and only return values defined outside of the loop.
 struct SimplifyTrivialLoops : public OpRewritePattern<ForOp> {
   using OpRewritePattern<ForOp>::OpRewritePattern;

mlir/test/Dialect/SCF/canonicalize.mlir

func.func @replace_single_iteration_loop_non_unit_step() {
   // CHECK: "test.consume"(%[[VAL]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
 
 // -----
 
-#map = affine_map<(d0) -> (d0 + 1)>
 // CHECK-LABEL: func @replace_single_iteration_const_diff(
 //  CHECK-SAME: %[[A0:.*]]: index)
 func.func @replace_single_iteration_const_diff(%arg0 : index) {
   // CHECK-NEXT: %[[CST:.*]] = arith.constant 2
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
-  %5 = affine.apply #map(%arg0)
+  %5 = arith.addi %arg0, %c1 : index
   // CHECK-NOT: scf.for
   scf.for %arg2 = %arg0 to %5 step %c1 {
     // CHECK-NEXT: %[[MUL:.*]] = arith.muli %[[A0]], %[[CST]]
     %7 = arith.muli %c2, %arg2 : index
     // CHECK-NEXT: "test.consume"(%[[MUL]])
     "test.consume"(%7) : (index) -> ()
   }
   return

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     ),
     hdrs = glob(
         [
             "include/mlir/Dialect/SCF/IR/*.h",
         ],
     ),
     includes = ["include"],
     deps = [
-        ":AffineDialect",
         ":ArithmeticDialect",
         ":ArithmeticUtils",
         ":BufferizationDialect",
         ":ControlFlowDialect",
         ":ControlFlowInterfaces",
         ":FuncDialect",
         ":IR",
         ":LoopLikeInterface",