[mlir] Add loop bounds normalization pass

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

@@ -140,6 +140,7 @@ def ForOp : SCF_Op<"for",
         "getSingleUpperBound", "getYieldedValuesMutable",
         "promoteIfSingleIteration", "replaceWithAdditionalYields",
         "yieldTiledValuesAndReplace"]>,
+       LoopLikeWithInductionVarsOpInterface,
        AllTypesMatch<["lowerBound", "upperBound", "step"]>,
        ConditionallySpeculatable,
        DeclareOpInterfaceMethods<RegionBranchOpInterface,
@@ -267,6 +268,74 @@ def ForOp : SCF_Op<"for",
       return getBody()->getArguments().drop_front(getNumInductionVars())[index];
     }
 
+    /// Return the induction variables.
+    ::mlir::ValueRange getInductionVars() {
+      return getBody()->getArguments().take_front(getNumInductionVars());
+    }
+
+    /// Get lower bounds as `OpFoldResult`.
+    SmallVector<OpFoldResult> getMixedLowerBound() {
+      return {getAsOpFoldResult(getLowerBound())};
+    }
+
+    /// Get upper bounds as `OpFoldResult`.
+    SmallVector<OpFoldResult> getMixedUpperBound() {
+      return {getAsOpFoldResult(getUpperBound())};
+    }
+
+    // Get steps as `OpFoldResult`.
+    SmallVector<OpFoldResult> getMixedStep() {
+      return {getAsOpFoldResult(getStep())};
+    }
+
+    /// Get lower bounds as values.
+    SmallVector<Value> getLowerBound(OpBuilder &b) {
+      return ValueRange{getLowerBound()};
+    }
+
+    /// Get upper bounds as values.
+    SmallVector<Value> getUpperBound(OpBuilder &b) {
+      return ValueRange{getUpperBound()};
+    }
+
+    /// Get steps as values.
+    SmallVector<Value> getStep(OpBuilder &b) {
+      return ValueRange{getStep()};
+    }
+
+    /// Set the lower bounds from `OpFoldResult`.
+    void setMixedLowerBounds(OpBuilder &b, ArrayRef<OpFoldResult> lbs) {
+      setLowerBound(getValueOrCreateConstantIndexOp(b, getLoc(), lbs[0]));
+    }
+
+    /// Set the upper bounds from `OpFoldResult`.
+    void setMixedUpperBounds(OpBuilder &b, ArrayRef<OpFoldResult> ubs) {
+      setUpperBound(getValueOrCreateConstantIndexOp(b, getLoc(), ubs[0]));
+    }
+
+    /// Set the steps from `OpFoldResult`.
+    void setMixedSteps(OpBuilder &b, ArrayRef<OpFoldResult> steps) {
+      setStep(getValueOrCreateConstantIndexOp(b, getLoc(), steps[0]));
+    }
+
+    /// Set the lower bounds from values.
+    void setLowerBounds(ArrayRef<Value> lbs) {
+      assert(lbs.size() == 1 && "expected a single lower bound");
+      setLowerBound(lbs[0]);
+    }
+
+    /// Set the upper bounds from values.
+    void setUpperBounds(ArrayRef<Value> ubs) {
+      assert(ubs.size() == 1 && "expected a single upper bound");
+      setUpperBound(ubs[0]);
+    }
+
+    /// Set the steps from values.
+    void setSteps(ArrayRef<Value> steps) {
+      assert(steps.size() == 1 && "expected a single step");
+      setStep(steps[0]);
+    }
+
     void setLowerBound(Value bound) { getOperation()->setOperand(0, bound); }
     void setUpperBound(Value bound) { getOperation()->setOperand(1, bound); }
     void setStep(Value step) { getOperation()->setOperand(2, step); }
@@ -304,6 +373,7 @@ def ForallOp : SCF_Op<"forall", [
           ["getInitsMutable", "getRegionIterArgs", "getSingleInductionVar", 
            "getSingleLowerBound", "getSingleUpperBound", "getSingleStep",
            "promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
+       LoopLikeWithInductionVarsOpInterface,
        RecursiveMemoryEffects,
        SingleBlockImplicitTerminator<"scf::InParallelOp">,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
@@ -543,6 +613,33 @@ def ForallOp : SCF_Op<"forall", [
       return getValueOrCreateConstantIndexOp(b, getLoc(), getMixedStep());
     }
 
+    /// Set the lower bounds from `OpFoldResult`.
+    void setMixedLowerBounds(OpBuilder &b, ArrayRef<OpFoldResult> lbs);
+
+    /// Set the upper bounds from `OpFoldResult`.
+    void setMixedUpperBounds(OpBuilder &b, ArrayRef<OpFoldResult> ubs);
+
+    /// Set the steps from `OpFoldResult`.
+    void setMixedSteps(OpBuilder &b, ArrayRef<OpFoldResult> steps);
+
+    /// Set the lower bounds from values.
+    void setLowerBounds(ArrayRef<Value> lbs) {
+      OpBuilder b(getOperation()->getContext());
+      return setMixedLowerBounds(b, getAsOpFoldResult(lbs));
+    }
+
+    /// Set the upper bounds from values.
+    void setUpperBounds(ArrayRef<Value> ubs) {
+      OpBuilder b(getOperation()->getContext());
+      return setMixedUpperBounds(b, getAsOpFoldResult(ubs));
+    }
+
+    /// Set the steps from values.
+    void setSteps(ArrayRef<Value> steps) {
+      OpBuilder b(getOperation()->getContext());
+      return setMixedSteps(b, getAsOpFoldResult(steps));
+    }
+
     int64_t getRank() { return getStaticLowerBound().size(); }
 
     /// Number of operands controlling the loop: lbs, ubs, steps

mlir/include/mlir/Dialect/Utils/LoopUtils.h

@@ -0,0 +1,30 @@
+//===- LoopUtils.h - Helpers related to loop operations ---------*- C++ -*-===//
+//
+// 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 header file defines utilities for loop operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/PatternMatch.h"
+
+namespace mlir {
+
+// This structure is to pass and return sets of loop parameters without
+// confusing the order.
+struct LoopParams {
+  Value lowerBound;
+  Value upperBound;
+  Value step;
+};
+
+/// Calculate the normalized loop upper bounds with lower bound equal to zero
+/// and step equal to one.
+LoopParams emitNormalizedLoopBounds(RewriterBase &rewriter, Location loc,
+                                    Value lb, Value ub, Value step);
+
+} // namespace mlir

mlir/include/mlir/Interfaces/LoopLikeInterface.h

@@ -13,6 +13,7 @@
 #ifndef MLIR_INTERFACES_LOOPLIKEINTERFACE_H_
 #define MLIR_INTERFACES_LOOPLIKEINTERFACE_H_
 
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/OpDefinition.h"
 
 namespace mlir {
@@ -28,6 +29,9 @@ using NewYieldValuesFn = std::function<SmallVector<Value>(
 namespace detail {
 /// Verify invariants of the LoopLikeOpInterface.
 LogicalResult verifyLoopLikeOpInterface(Operation *op);
+
+/// Verify invariants of the LoopLikeWithInductionVarsOpInterface.
+LogicalResult verifyLoopLikeWithInductionVarsOpInterface(Operation *op);
 } // namespace detail
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Interfaces/LoopLikeInterface.td

@@ -375,6 +375,132 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
   }];
 }
 
+def LoopLikeWithInductionVarsOpInterface 
+    : OpInterface<"LoopLikeWithInductionVarsOpInterface", [LoopLikeOpInterface]> {
+  let description = [{
+    Interface for loop-like operations with one or more induction variables.
+    This interface contains helper functions for retrieving and updating the
+    lower bound, upper bound and step size for each induction variable and
+    provides a utility function to check whether the loop is normalized., i.e.
+    all lower bounds are equal to zero and steps are equal to one.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Return the induction variables if they exist, otherwise return
+        std::nullopt.
+      }],
+      /*retTy=*/"::mlir::ValueRange",
+      /*methodName=*/"getInductionVars"
+    >,
+    InterfaceMethod<[{
+        Return the lower bound values or attributes as OpFoldResult.
+      }],
+      /*retTy=*/"SmallVector<::mlir::OpFoldResult>",
+      /*methodName=*/"getMixedLowerBound"
+    >,
+    InterfaceMethod<[{
+        Return the step values or attributes if they exist as OpFoldResult.
+      }],
+      /*retTy=*/"SmallVector<::mlir::OpFoldResult>",
+      /*methodName=*/"getMixedStep"
+    >,
+    InterfaceMethod<[{
+        Return the upper bound values or attributes as OpFoldResult.
+      }],
+      /*retTy=*/"SmallVector<::mlir::OpFoldResult>",
+      /*methodName=*/"getMixedUpperBound"
+    >,
+    InterfaceMethod<[{
+        Return the lower bounds as values.
+      }],
+      /*retTy=*/"SmallVector<Value>",
+      /*methodName=*/"getLowerBound",
+      /*args=*/(ins "OpBuilder &":$b)
+    >,
+    InterfaceMethod<[{
+        Return the steps as values.
+      }],
+      /*retTy=*/"SmallVector<Value>",
+      /*methodName=*/"getStep",
+      /*args=*/(ins "OpBuilder &":$b)
+    >,
+    InterfaceMethod<[{
+        Return the upper bounds as values.
+      }],
+      /*retTy=*/"SmallVector<Value>",
+      /*methodName=*/"getUpperBound",
+      /*args=*/(ins "OpBuilder &":$b)
+    >,
+    InterfaceMethod<[{
+        Set the lower bounds from an array of `OpFoldResult`.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setMixedLowerBounds",
+      /*args=*/(ins "OpBuilder &":$b, "ArrayRef<OpFoldResult>":$lbs)
+    >,
+    InterfaceMethod<[{
+        Set the steps from an array of `OpFoldResult`.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setMixedSteps",
+      /*args=*/(ins "OpBuilder &":$b, "ArrayRef<OpFoldResult>":$lbs)
+    >,
+    InterfaceMethod<[{
+        Set the upper bounds from an array of `OpFoldResult`.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setMixedUpperBounds",
+      /*args=*/(ins "OpBuilder &":$b, "ArrayRef<OpFoldResult>":$lbs)
+    >,
+    InterfaceMethod<[{
+        Set the lower bounds from an array of values.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setLowerBounds",
+      /*args=*/(ins "ArrayRef<Value>":$lbs)
+    >,
+    InterfaceMethod<[{
+        Set the steps from an array of values.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setSteps",
+      /*args=*/(ins "ArrayRef<Value>":$lbs)
+    >,
+    InterfaceMethod<[{
+        Set the upper bounds from an array of values.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setUpperBounds",
+      /*args=*/(ins "ArrayRef<Value>":$lbs)
+    >,
+    InterfaceMethod<[{
+        Checks if the lower bounds are zeros and steps are ones.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"isNormalized",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        auto allEqual = [](ArrayRef<OpFoldResult> results, int64_t val) {
+          return llvm::all_of(results, [&](OpFoldResult ofr) {
+            auto intValue = getConstantIntValue(ofr);
+            return intValue.has_value() && intValue == val;
+          });
+        };
+        SmallVector<::mlir::OpFoldResult> lbs = $_op.getMixedLowerBound();
+        SmallVector<::mlir::OpFoldResult> steps = $_op.getMixedStep();
+        return allEqual(lbs, 0) && allEqual(steps, 1);
+      }]
+    >
+  ];
+
+  let verify = [{
+    return detail::verifyLoopLikeWithInductionVarsOpInterface($_op);
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Traits
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Transforms/Passes.h

@@ -82,6 +82,10 @@ std::unique_ptr<Pass> createLoopInvariantCodeMotionPass();
 /// Creates a pass that hoists loop-invariant subset ops.
 std::unique_ptr<Pass> createLoopInvariantSubsetHoistingPass();
 
+/// Create a pass that normalizes the loop bounds of loop-like operations with
+/// induction variables.
+std::unique_ptr<Pass> createNormalizeLoopBoundsPass();
+
 /// Creates a pass to strip debug information from a function.
 std::unique_ptr<Pass> createStripDebugInfoPass();
 

mlir/include/mlir/Transforms/Passes.td

@@ -377,6 +377,12 @@ def Mem2Reg : Pass<"mem2reg"> {
   ];
 }
 
+def NormalizeLoopBounds : Pass<"normalize-loop-bounds"> {
+  let summary = "Normalize the loop bounds of loop-like operations with "
+                "induction variables.";
+  let constructor = "mlir::createNormalizeLoopBoundsPass()";
+}
+
 def PrintOpStats : Pass<"print-op-stats"> {
   let summary = "Print statistics of operations";
   let constructor = "mlir::createPrintOpStatsPass()";

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

@@ -1387,6 +1387,66 @@ void ForallOp::build(
   build(b, result, lbs, ubs, steps, outputs, mapping, bodyBuilderFn);
 }
 
+/// Set the lower bounds from `OpFoldResult`.
+void ForallOp::setMixedLowerBounds(OpBuilder &b, ArrayRef<OpFoldResult> lbs) {
+  SmallVector<int64_t> staticLbs;
+  SmallVector<Value> dynamicLbs;
+  dispatchIndexOpFoldResults(lbs, dynamicLbs, staticLbs);
+  getOperation()->setOperands(0, getDynamicLowerBound().size(), dynamicLbs);
+  (*this)->setAttr(getStaticLowerBoundAttrName(),
+                   b.getDenseI64ArrayAttr(staticLbs));
+  ArrayRef<int32_t> segmentSizes =
+      (*this)
+          ->getAttrOfType<DenseI32ArrayAttr>("operandSegmentSizes")
+          .asArrayRef();
+  SmallVector<int32_t> newSegmentSizes(segmentSizes.begin(),
+                                       segmentSizes.end());
+  newSegmentSizes[0] = dynamicLbs.size();
+  (*this)->setAttr("operandSegmentSizes",
+                   b.getDenseI32ArrayAttr(newSegmentSizes));
+}
+
+/// Set the upper bounds from `OpFoldResult`.
+void ForallOp::setMixedUpperBounds(OpBuilder &b, ArrayRef<OpFoldResult> ubs) {
+  SmallVector<int64_t> staticUbs;
+  SmallVector<Value> dynamicUbs;
+  dispatchIndexOpFoldResults(ubs, dynamicUbs, staticUbs);
+  size_t offset = getDynamicLowerBound().size();
+  getOperation()->setOperands(offset, getDynamicUpperBound().size(),
+                              dynamicUbs);
+  (*this)->setAttr(getStaticUpperBoundAttrName(),
+                   b.getDenseI64ArrayAttr(staticUbs));
+  ArrayRef<int32_t> segmentSizes =
+      (*this)
+          ->getAttrOfType<DenseI32ArrayAttr>("operandSegmentSizes")
+          .asArrayRef();
+  SmallVector<int32_t> newSegmentSizes(segmentSizes.begin(),
+                                       segmentSizes.end());
+  newSegmentSizes[1] = dynamicUbs.size();
+  (*this)->setAttr("operandSegmentSizes",
+                   b.getDenseI32ArrayAttr(newSegmentSizes));
+}
+
+/// Set the steps from `OpFoldResult`.
+void ForallOp::setMixedSteps(OpBuilder &b, ArrayRef<OpFoldResult> steps) {
+  SmallVector<int64_t> staticSteps;
+  SmallVector<Value> dynamicSteps;
+  dispatchIndexOpFoldResults(steps, dynamicSteps, staticSteps);
+  size_t offset = getDynamicLowerBound().size() + getDynamicUpperBound().size();
+  getOperation()->setOperands(offset, getDynamicStep().size(), dynamicSteps);
+  (*this)->setAttr(getStaticStepAttrName(),
+                   b.getDenseI64ArrayAttr(staticSteps));
+  ArrayRef<int32_t> segmentSizes =
+      (*this)
+          ->getAttrOfType<DenseI32ArrayAttr>("operandSegmentSizes")
+          .asArrayRef();
+  SmallVector<int32_t> newSegmentSizes(segmentSizes.begin(),
+                                       segmentSizes.end());
+  newSegmentSizes[2] = dynamicSteps.size();
+  (*this)->setAttr("operandSegmentSizes",
+                   b.getDenseI32ArrayAttr(newSegmentSizes));
+}
+
 // Checks if the lbs are zeros and steps are ones.
 bool ForallOp::isNormalized() {
   auto allEqual = [](ArrayRef<OpFoldResult> results, int64_t val) {