From 5de1152e2aefa732c895f29c53639c9511a035ac Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko@google.com>
Date: Thu, 13 Jun 2019 07:08:46 -0700
Subject: [PATCH] Disallow non-index operands and results in affine.apply

`affine.apply` is supposed to operate on values of index types in context of
affine loops.  It is possible to programmatically constuct an `affine.apply`
that takes values of other types as operands or returns them, but it would not
be parseable.  Disallow such cases in the verifier.

PiperOrigin-RevId: 253021704
---
 mlir/lib/AffineOps/AffineOps.cpp |  9 +++++++++
 mlir/test/AffineOps/invalid.mlir | 20 ++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/mlir/lib/AffineOps/AffineOps.cpp b/mlir/lib/AffineOps/AffineOps.cpp
index fc70807..26a35ad 100644
--- a/mlir/lib/AffineOps/AffineOps.cpp
+++ b/mlir/lib/AffineOps/AffineOps.cpp
@@ -174,6 +174,15 @@ LogicalResult AffineApplyOp::verify() {
     return emitOpError(
         "operand count and affine map dimension and symbol count must match");
 
+  // Verify that all operands are of `index` type.
+  for (Type t : getOperandTypes()) {
+    if (!t.isIndex())
+      return emitOpError("operands must be of type 'index'");
+  }
+
+  if (!getResult()->getType().isIndex())
+    return emitOpError("result must be of type 'index'");
+
   // Verify that the operands are valid dimension and symbol identifiers.
   if (failed(verifyDimAndSymbolIdentifiers(*this, getOperands(),
                                            map.getNumDims())))
diff --git a/mlir/test/AffineOps/invalid.mlir b/mlir/test/AffineOps/invalid.mlir
index 69260a7..748e0fe 100644
--- a/mlir/test/AffineOps/invalid.mlir
+++ b/mlir/test/AffineOps/invalid.mlir
@@ -28,6 +28,26 @@ func @affine_apply_invalid_sym() {
 
 // -----
 
+func @affine_apply_operand_non_index(%arg0 : i32) {
+  // Custom parser automatically assigns all arguments the `index` so we must
+  // use the generic syntax here to exercise the verifier.
+  // expected-error@+1 {{operands must be of type 'index'}}
+  %0 = "affine.apply"(%arg0) {map: (d0) -> (d0)} : (i32) -> (index)
+  return
+}
+
+// -----
+
+func @affine_apply_resul_non_index(%arg0 : index) {
+  // Custom parser automatically assigns `index` as the result type so we must
+  // use the generic syntax here to exercise the verifier.
+  // expected-error@+1 {{result must be of type 'index'}}
+  %0 = "affine.apply"(%arg0) {map: (d0) -> (d0)} : (index) -> (i32)
+  return
+}
+
+// -----
+
 #map = (d0)[s0] -> (d0 + s0)
 
 func @affine_for_lower_bound_invalid_dim(%arg : index) {
-- 
2.7.4