AffineApplyOp index support

Author: Hardcode84

mlir/lib/Dialect/Vector/Transforms/SLPVectorizer.cpp

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/DataLayoutAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
@@ -142,6 +143,27 @@ static Type getElementType(Operation *op) {
   return {};
 }
 
+static bool isAdjacentAffineMapIndices(Value idx1, Value idx2) {
+  auto applyOp1 = idx1.getDefiningOp<affine::AffineApplyOp>();
+  if (!applyOp1)
+    return false;
+
+  auto applyOp2 = idx2.getDefiningOp<affine::AffineApplyOp>();
+  if (!applyOp2)
+    return false;
+
+  if (applyOp1.getOperands() != applyOp2.getOperands())
+    return false;
+
+  AffineExpr expr1 = applyOp1.getAffineMap().getResult(0);
+  AffineExpr expr2 = applyOp2.getAffineMap().getResult(0);
+  auto diff =
+      simplifyAffineExpr(expr2 - expr1, 0, applyOp1.getOperands().size());
+
+  auto diffConst = dyn_cast<AffineConstantExpr>(diff);
+  return diffConst && diffConst.getValue() == 1;
+}
+
 /// Check if two indices are consecutive, i.e index1 + 1 == index2.
 static bool isAdjacentIndices(Value idx1, Value idx2) {
   if (auto c1 = getConstantIntValue(idx1)) {
@@ -160,7 +182,9 @@ static bool isAdjacentIndices(Value idx1, Value idx2) {
     }
   }
 
-  // TODO: Handle affine.apply, etc
+  if (isAdjacentAffineMapIndices(idx1, idx2))
+    return true;
+
   return false;
 }
 

mlir/test/Dialect/Vector/slp-vectorize.mlir

@@ -143,6 +143,37 @@ func.func @read_write_add_index_interleaved(%arg0: memref<8xi32>, %arg1: memref<
 }
 
 
+#map0 = affine_map<()[s0, s1] -> (s1 * s0)>
+#map1 = affine_map<()[s0, s1] -> (s1 * s0 + 1)>
+#map2 = affine_map<()[s0, s1] -> (s1 * s0 + 2)>
+#map3 = affine_map<()[s0, s1] -> (s1 * s0 + 3)>
+
+// CHECK-LABEL: func @read_write_affine_apply
+//  CHECK-SAME: (%[[ARG0:.*]]: memref<8xi32>, %[[ARG1:.*]]: memref<8xi32>, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index)
+func.func @read_write_affine_apply(%arg0: memref<8xi32>, %arg1: memref<8xi32>, %arg2: index, %arg3: index) {
+  // CHECK:     %[[IDX:.*]] = affine.apply #{{.*}}()[%[[ARG2]], %[[ARG3]]]
+  // CHECK:     %[[RES:.*]] = vector.load %[[ARG0]][%[[IDX]]] : memref<8xi32>, vector<4xi32>
+  // CHECK:     vector.store %[[RES]], %[[ARG0]][%[[IDX]]] : memref<8xi32>, vector<4xi32>
+
+  %ind0 = affine.apply #map0()[%arg2, %arg3]
+  %ind1 = affine.apply #map1()[%arg2, %arg3]
+  %ind2 = affine.apply #map2()[%arg2, %arg3]
+  %ind3 = affine.apply #map3()[%arg2, %arg3]
+
+  %0 = memref.load %arg0[%ind0] : memref<8xi32>
+  %1 = memref.load %arg0[%ind1] : memref<8xi32>
+  %2 = memref.load %arg0[%ind2] : memref<8xi32>
+  %3 = memref.load %arg0[%ind3] : memref<8xi32>
+
+  memref.store %0, %arg0[%ind0] : memref<8xi32>
+  memref.store %1, %arg0[%ind1] : memref<8xi32>
+  memref.store %2, %arg0[%ind2] : memref<8xi32>
+  memref.store %3, %arg0[%ind3] : memref<8xi32>
+
+  return
+}
+
+
 // CHECK-LABEL: func @read_read_add
 //  CHECK-SAME: (%[[ARG0:.*]]: memref<8xi32>, %[[ARG1:.*]]: memref<8xi32>)
 func.func @read_read_add(%arg0: memref<8xi32>, %arg1: memref<8xi32>) -> (i32, i32, i32, i32) {