From e5691c512f7335fe90093bb9989da1760d80701b Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday@polymagelabs.com>
Date: Fri, 20 Dec 2019 09:43:34 -0800
Subject: [PATCH] fix isValidDim for block arg case

- a block argument associated with an arbitrary op can't be a valid
  dimensional identifier; it has to be the block argument of either
  a function op or an affine.for.

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Closes tensorflow/mlir#331

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/mlir/pull/331 from bondhugula:valid_dim 3273b4fcbaa31fb7b6671d93c9e42a6b2a6a4e4c
PiperOrigin-RevId: 286593693
---
 mlir/lib/Dialect/AffineOps/AffineOps.cpp | 14 +++++++-------
 mlir/test/AffineOps/invalid.mlir         | 13 +++++++++++++
 mlir/test/Transforms/slicing-utils.mlir  |  7 +++----
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/mlir/lib/Dialect/AffineOps/AffineOps.cpp
index 6768aa5..ef4060d4 100644
--- a/mlir/lib/Dialect/AffineOps/AffineOps.cpp
+++ b/mlir/lib/Dialect/AffineOps/AffineOps.cpp
@@ -142,8 +142,9 @@ bool mlir::isValidDim(Value *value) {
       return isTopLevelValue(dimOp.getOperand());
     return false;
   }
-  // This value is a block argument (which also includes 'affine.for' loop IVs).
-  return true;
+  // This value has to be a block argument for a FuncOp or an affine.for.
+  auto *parentOp = cast<BlockArgument>(value)->getOwner()->getParentOp();
+  return isa<FuncOp>(parentOp) || isa<AffineForOp>(parentOp);
 }
 
 /// Returns true if the 'index' dimension of the `memref` defined by
@@ -301,16 +302,15 @@ LogicalResult AffineApplyOp::verify() {
   return success();
 }
 
-// The result of the affine apply operation can be used as a dimension id if it
-// is a CFG value or if it is an Value, and all the operands are valid
-// dimension ids.
+// The result of the affine apply operation can be used as a dimension id if all
+// its operands are valid dimension ids.
 bool AffineApplyOp::isValidDim() {
   return llvm::all_of(getOperands(),
                       [](Value *op) { return mlir::isValidDim(op); });
 }
 
-// The result of the affine apply operation can be used as a symbol if it is
-// a CFG value or if it is an Value, and all the operands are symbols.
+// The result of the affine apply operation can be used as a symbol if all its
+// operands are symbols.
 bool AffineApplyOp::isValidSymbol() {
   return llvm::all_of(getOperands(),
                       [](Value *op) { return mlir::isValidSymbol(op); });
diff --git a/mlir/test/AffineOps/invalid.mlir b/mlir/test/AffineOps/invalid.mlir
index 4534d70..2ade76c 100644
--- a/mlir/test/AffineOps/invalid.mlir
+++ b/mlir/test/AffineOps/invalid.mlir
@@ -51,6 +51,19 @@ func @affine_for_upper_bound_invalid_dim(%arg : index) {
 }
 
 // -----
+func @affine_load_invalid_dim(%M : memref<10xi32>) {
+  "unknown"() ({
+  ^bb0(%arg: index):
+    affine.load %M[%arg] : memref<10xi32>
+    // expected-error@-1 {{index must be a dimension or symbol identifier}}
+    br ^bb1
+  ^bb1:
+    br ^bb1
+  }) : () -> ()
+  return
+}
+
+// -----
 
 #map0 = (d0)[s0] -> (d0 + s0)
 
diff --git a/mlir/test/Transforms/slicing-utils.mlir b/mlir/test/Transforms/slicing-utils.mlir
index 49410db..8c6fb01 100644
--- a/mlir/test/Transforms/slicing-utils.mlir
+++ b/mlir/test/Transforms/slicing-utils.mlir
@@ -222,13 +222,12 @@ func @slicing_test() {
 // FWDBWD-LABEL: slicing_test_2
 func @slicing_test_2() {
   %c0 = constant 0 : index
-  %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c16 = constant 16 : index
-  loop.for %i0 = %c0 to %c16 step %c1 {
+  affine.for %i0 = %c0 to %c16 {
     affine.for %i1 = (i)[] -> (i)(%i0) to 10 {
       // BWD: matched: %[[b:.*]] {{.*}} backward static slice:
-      // BWD: loop.for {{.*}}
+      // BWD: affine.for {{.*}}
 
       // affine.for appears in the body of loop.for
       // BWD: affine.for {{.*}}
@@ -238,7 +237,7 @@ func @slicing_test_2() {
       %b = "slicing-test-op"(%i1): (index) -> index
 
       // BWD: matched: %[[c:.*]] {{.*}} backward static slice:
-      // BWD: loop.for {{.*}}
+      // BWD: affine.for {{.*}}
 
       // affine.for appears in the body of loop.for
       // BWD-NEXT: affine.for {{.*}}
-- 
2.7.4