1 files changed, 37 insertions, 1 deletions
diff --git a/mlir/lib/IR/Statement.cpp b/mlir/lib/IR/Statement.cpp
index e4cb3e5bd87..9dd22852bac 100644
--- a/mlir/lib/IR/Statement.cpp
+++ b/mlir/lib/IR/Statement.cpp
@@ -17,10 +17,10 @@
 
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLFunction.h"
 #include "mlir/IR/MLIRContext.h"
-#include "mlir/IR/StandardOps.h"
 #include "mlir/IR/Statements.h"
 #include "mlir/IR/StmtVisitor.h"
 #include "llvm/ADT/DenseMap.h"
@@ -92,6 +92,42 @@ const MLValue *Statement::getOperand(unsigned idx) const {
   return getStmtOperand(idx).get();
 }
 
+// MLValue can be used as a dimension id if it is valid as a symbol, or
+// it is an induction variable, or it is a result of affine apply operation
+// with dimension id arguments.
+bool MLValue::isValidDim() const {
+  if (auto *stmt = getDefiningStmt()) {
+    // Top level statement or constant operation is ok.
+    if (stmt->getParentStmt() == nullptr || stmt->is<ConstantOp>())
+      return true;
+    // Affine apply operation is ok if all of its operands are ok.
+    if (auto op = stmt->getAs<AffineApplyOp>())
+      return op->isValidDim();
+    return false;
+  }
+  // This value is either a function argument or an induction variable. Both
+  // are ok.
+  return true;
+}
+
+// MLValue can be used as a symbol if it is a constant, or it is defined at
+// the top level, or it is a result of affine apply operation with symbol
+// arguments.
+bool MLValue::isValidSymbol() const {
+  if (auto *stmt = getDefiningStmt()) {
+    // Top level statement or constant operation is ok.
+    if (stmt->getParentStmt() == nullptr || stmt->is<ConstantOp>())
+      return true;
+    // Affine apply operation is ok if all of its operands are ok.
+    if (auto op = stmt->getAs<AffineApplyOp>())
+      return op->isValidSymbol();
+    return false;
+  }
+  // This value is either a function argument or an induction variable.
+  // Function argument is ok, induction variable is not.
+  return isa<MLFuncArgument>(this);
+}
+
 void Statement::setOperand(unsigned idx, MLValue *value) {
   getStmtOperand(idx).set(value);
 }