Enhance ConstantIntOp to work with AffineInt, move use/def list processing code

into a more logical header.

PiperOrigin-RevId: 206175435

mlir/include/mlir/IR/CFGValue.h

@@ -22,7 +22,6 @@
 #ifndef MLIR_IR_CFGOPERAND_H
 #define MLIR_IR_CFGOPERAND_H
 
-#include "mlir/IR/SSAOperand.h"
 #include "mlir/IR/SSAValue.h"
 
 namespace mlir {

mlir/include/mlir/IR/SSAValue.h

@@ -23,48 +23,12 @@
 #define MLIR_IR_SSAVALUE_H
 
 #include "mlir/IR/Types.h"
+#include "mlir/IR/UseDefLists.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/PointerIntPair.h"
-#include "llvm/ADT/iterator_range.h"
 
 namespace mlir {
 class OperationInst;
-class IROperand;
-template <typename OperandType, typename OwnerType> class SSAValueUseIterator;
-
-class IRObjectWithUseList {
-public:
-  ~IRObjectWithUseList() {
-    assert(use_empty() && "Cannot destroy a value that still has uses!");
-  }
-
-  /// Returns true if this value has no uses.
-  bool use_empty() const { return firstUse == nullptr; }
-
-  /// Returns true if this value has exactly one use.
-  inline bool hasOneUse() const;
-
-  using use_iterator = SSAValueUseIterator<IROperand, void>;
-  using use_range = llvm::iterator_range<use_iterator>;
-
-  inline use_iterator use_begin() const;
-  inline use_iterator use_end() const;
-
-  /// Returns a range of all uses, which is useful for iterating over all uses.
-  inline use_range getUses() const;
-
-  /// Replace all uses of 'this' value with the new value, updating anything in
-  /// the IR that uses 'this' to use the other value instead.  When this returns
-  /// there are zero uses of 'this'.
-  void replaceAllUsesWith(IRObjectWithUseList *newValue);
-
-protected:
-  IRObjectWithUseList() {}
-
-private:
-  friend class IROperand;
-  IROperand *firstUse = nullptr;
-};
 
 /// This enumerates all of the SSA value kinds in the MLIR system.
 enum class SSAValueKind {
@@ -133,48 +97,6 @@ protected:
   SSAValueImpl(KindTy kind, Type *type) : SSAValue((SSAValueKind)kind, type) {}
 };
 
-// FIXME: Implement SSAValueUseIterator here.
-
-/// An iterator over all uses of a ValueBase.
-template <typename OperandType, typename OwnerType>
-class SSAValueUseIterator
-    : public std::iterator<std::forward_iterator_tag, IROperand> {
-public:
-  SSAValueUseIterator() = default;
-  explicit SSAValueUseIterator(IROperand *current) : current(current) {}
-  OperandType *operator->() const { return current; }
-  OperandType &operator*() const { return current; }
-
-  template<typename SFINAE_Owner = OwnerType>
-  typename std::enable_if<!std::is_void<OwnerType>::value, SFINAE_Owner>::type
-  getUser() const {
-    return current->getOwner();
-  }
-
-  SSAValueUseIterator &operator++() {
-    assert(current && "incrementing past end()!");
-    current = (OperandType *)current->getNextOperandUsingThisValue();
-    return *this;
-  }
-
-  SSAValueUseIterator operator++(int unused) {
-    SSAValueUseIterator copy = *this;
-    ++*this;
-    return copy;
-  }
-
-  friend bool operator==(SSAValueUseIterator lhs, SSAValueUseIterator rhs) {
-    return lhs.current == rhs.current;
-  }
-
-  friend bool operator!=(SSAValueUseIterator lhs, SSAValueUseIterator rhs) {
-    return !(lhs == rhs);
-  }
-
-private:
-  OperandType *current;
-};
-
 } // namespace mlir
 
 #endif

mlir/include/mlir/IR/StandardOps.h

@@ -74,7 +74,7 @@ protected:
 };
 
 /// This is a refinement of the "constant" op for the case where it is
-/// returning an integer value.
+/// returning an integer value (either an IntegerType or AffineInt).
 ///
 ///   %1 = "constant"(){value: 42}
 ///

mlir/include/mlir/IR/UseDefLists.h

@@ -1,4 +1,4 @@
-//===- SSAOperand.h ---------------------------------------------*- C++ -*-===//
+//===- UseDefLists.h --------------------------------------------*- C++ -*-===//
 //
 // Copyright 2019 The MLIR Authors.
 //
@@ -15,17 +15,53 @@
 // limitations under the License.
 // =============================================================================
 //
-// This file defines generic SSA operand and manipulation utilities.
+// This file defines generic use/def list machinery and manipulation utilities.
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_IR_SSAOPERAND_H
-#define MLIR_IR_SSAOPERAND_H
+#ifndef MLIR_IR_USEDEFLISTS_H
+#define MLIR_IR_USEDEFLISTS_H
 
-#include "mlir/IR/SSAValue.h"
+#include "llvm/ADT/iterator_range.h"
 
 namespace mlir {
-class SSAValue;
+
+class IROperand;
+template <typename OperandType, typename OwnerType> class SSAValueUseIterator;
+
+class IRObjectWithUseList {
+public:
+  ~IRObjectWithUseList() {
+    assert(use_empty() && "Cannot destroy a value that still has uses!");
+  }
+
+  /// Returns true if this value has no uses.
+  bool use_empty() const { return firstUse == nullptr; }
+
+  /// Returns true if this value has exactly one use.
+  inline bool hasOneUse() const;
+
+  using use_iterator = SSAValueUseIterator<IROperand, void>;
+  using use_range = llvm::iterator_range<use_iterator>;
+
+  inline use_iterator use_begin() const;
+  inline use_iterator use_end() const;
+
+  /// Returns a range of all uses, which is useful for iterating over all uses.
+  inline use_range getUses() const;
+
+  /// Replace all uses of 'this' value with the new value, updating anything in
+  /// the IR that uses 'this' to use the other value instead.  When this returns
+  /// there are zero uses of 'this'.
+  void replaceAllUsesWith(IRObjectWithUseList *newValue);
+
+protected:
+  IRObjectWithUseList() {}
+
+private:
+  friend class IROperand;
+  IROperand *firstUse = nullptr;
+};
 
 /// A reference to a value, suitable for use as an operand of an instruction,
 /// statement, etc.
@@ -137,6 +173,46 @@ private:
   IROwnerTy *const owner;
 };
 
+/// An iterator over all uses of a ValueBase.
+template <typename OperandType, typename OwnerType>
+class SSAValueUseIterator
+    : public std::iterator<std::forward_iterator_tag, IROperand> {
+public:
+  SSAValueUseIterator() = default;
+  explicit SSAValueUseIterator(IROperand *current) : current(current) {}
+  OperandType *operator->() const { return current; }
+  OperandType &operator*() const { return current; }
+
+  template <typename SFINAE_Owner = OwnerType>
+  typename std::enable_if<!std::is_void<OwnerType>::value, SFINAE_Owner>::type
+  getUser() const {
+    return current->getOwner();
+  }
+
+  SSAValueUseIterator &operator++() {
+    assert(current && "incrementing past end()!");
+    current = (OperandType *)current->getNextOperandUsingThisValue();
+    return *this;
+  }
+
+  SSAValueUseIterator operator++(int unused) {
+    SSAValueUseIterator copy = *this;
+    ++*this;
+    return copy;
+  }
+
+  friend bool operator==(SSAValueUseIterator lhs, SSAValueUseIterator rhs) {
+    return lhs.current == rhs.current;
+  }
+
+  friend bool operator!=(SSAValueUseIterator lhs, SSAValueUseIterator rhs) {
+    return !(lhs == rhs);
+  }
+
+private:
+  OperandType *current;
+};
+
 inline auto IRObjectWithUseList::use_begin() const -> use_iterator {
   return use_iterator(firstUse);
 }

mlir/lib/IR/StandardOps.cpp

@@ -73,10 +73,12 @@ const char *ConstantOp::verify() const {
   return "requires a result type that aligns with the 'value' attribute";
 }
 
-/// ConstantIntOp only matches values whose result type is an IntegerType.
+/// ConstantIntOp only matches values whose result type is an IntegerType or
+/// AffineInt.
 bool ConstantIntOp::isClassFor(const Operation *op) {
   return ConstantOp::isClassFor(op) &&
-         isa<IntegerType>(op->getResult(0)->getType());
+         (isa<IntegerType>(op->getResult(0)->getType()) ||
+          op->getResult(0)->getType()->isAffineInt());
 }
 
 void DimOp::print(OpAsmPrinter *p) const {