[Codegen] Change swizzle hint offset logic to use arith (iree-org#21237)

After the addition of PropagateConstantOffsets, the expected incoming
form for a constant offset is with `arith.addi ..., %c`, so change the
logic accordingly. This logic is overall more robust now that there is
a preceding pass giving some guarantee as to the form of the input.

Author: qedawkins

compiler/src/iree/compiler/Codegen/Common/ResolveSwizzleHints.cpp

@@ -26,6 +26,26 @@ struct ResolveSwizzleHintsPass final
 };
 } // namespace
 
+static Value createOrFoldNewStaticAdd(RewriterBase &rewriter, Value v,
+                                      int64_t offset) {
+  // Early exit for the common offset = 0 case.
+  if (offset == 0) {
+    return v;
+  }
+
+  if (auto add = v.getDefiningOp<arith::AddIOp>()) {
+    llvm::APInt constant;
+    if (matchPattern(add.getRhs(), m_ConstantInt(&constant))) {
+      Value combined = rewriter.create<arith::ConstantIndexOp>(
+          add.getLoc(), offset + constant.getSExtValue());
+      return rewriter.create<arith::AddIOp>(add.getLoc(), add.getLhs(),
+                                            combined, add.getOverflowFlags());
+    }
+  }
+  Value offsetVal = rewriter.create<arith::ConstantIndexOp>(v.getLoc(), offset);
+  return rewriter.create<arith::AddIOp>(v.getLoc(), v, offsetVal);
+}
+
 /// Swizzles vector.load(iree_codegen.swizzle_hint, offset). The
 /// SwizzleInterfaceAttr exposes two methods:
 ///   1. getAccessElementCount -> int64_t
@@ -53,10 +73,6 @@ static void swizzleLoad(RewriterBase &rewriter, vector::LoadOp load,
   VectorType swizzledLoadType =
       VectorType::get({accessWidth}, type.getElementType());
 
-  AffineExpr s0, s1;
-  bindSymbols(rewriter.getContext(), s0, s1);
-  AffineMap sum = AffineMap::get(0, 2, s0 + s1);
-
   // ~ vector.undef, overwritten by unrolling.
   Value replacement = rewriter.create<arith::ConstantOp>(
       hintLoc, type, rewriter.getZeroAttr(type));
@@ -65,10 +81,7 @@ static void swizzleLoad(RewriterBase &rewriter, vector::LoadOp load,
   // i = 0 -> C += k is the offset into the vector of a contiguous group of
   // swizzled elements.
   for (int64_t i = 0; i < loadWidth; i += accessWidth) {
-    auto vecOffset = rewriter.getIndexAttr(i);
-    auto newBaseOffset = affine::makeComposedFoldedAffineApply(
-        rewriter, hintLoc, sum, {memrefOffset, vecOffset});
-
+    Value newBaseOffset = createOrFoldNewStaticAdd(rewriter, memrefOffset, i);
     Value newOffset = getValueOrCreateConstantIndexOp(
         rewriter, hintLoc,
         hintOp.getSwizzle().swizzleOffset(rewriter, hintOp.getLoc(),
@@ -103,20 +116,14 @@ static void swizzleStore(RewriterBase &rewriter, vector::StoreOp store,
   int64_t storeWidth = type.getShape()[0];
   Value memrefOffset = store.getIndices()[0];
 
-  AffineExpr s0, s1;
-  bindSymbols(rewriter.getContext(), s0, s1);
-  AffineMap sum = AffineMap::get(0, 2, s0 + s1);
-
   // Store type = vector<C>, k = accessWidth
   // i = 0 -> C += k is the offset into the vector of a contiguous group of
   // swizzled elements.
   for (int64_t i = 0; i < storeWidth; i += accessWidth) {
     Value subVec = rewriter.create<vector::ExtractStridedSliceOp>(
         store.getLoc(), store.getValueToStore(), ArrayRef<int64_t>{i},
         ArrayRef<int64_t>{accessWidth}, ArrayRef<int64_t>{1});
-    auto vecOffset = rewriter.getIndexAttr(i);
-    auto newBaseOffset = affine::makeComposedFoldedAffineApply(
-        rewriter, hintLoc, sum, {memrefOffset, vecOffset});
+    Value newBaseOffset = createOrFoldNewStaticAdd(rewriter, memrefOffset, i);
 
     Value newOffset = getValueOrCreateConstantIndexOp(
         rewriter, hintLoc,

compiler/src/iree/compiler/Codegen/Common/test/resolve_swizzle_hints.mlir

@@ -156,23 +156,26 @@ func.func @swizzle_dynamic(%src: memref<?xf32>, %vec: vector<4xf32>, %offset: in
 
 // -----
 
-func.func @swizzle_adjust_affine_offset(%src: memref<?xf32>, %vec: vector<4xf32>, %offset_base: index) -> vector<4xf32> {
+func.func @swizzle_adjust_add_offset(%src: memref<?xf32>, %vec: vector<4xf32>, %offset_base: index) -> vector<4xf32> {
   %0 = iree_codegen.swizzle_hint %src[#iree_codegen.rotate_rows<64, 4>] : memref<?xf32>
-  %load_offset = affine.apply affine_map<()[s0] -> (16 + s0)>()[%offset_base]
+  %c16 = arith.constant 16 : index
+  %c1040 = arith.constant 1040 : index
+  %load_offset = arith.addi %offset_base, %c16 overflow<nsw> : index
   %1 = vector.load %0[%load_offset] : memref<?xf32>, vector<4xf32>
-  %store_offset = affine.apply affine_map<()[s0] -> (1040 + s0)>()[%offset_base]
+  %store_offset = arith.addi %offset_base, %c1040 overflow<nsw> : index
   vector.store %vec, %0[%store_offset] : memref<?xf32>, vector<4xf32>
   return %1: vector<4xf32>
 }
 
-// CHECK-LABEL: func @swizzle_adjust_affine_offset
+// CHECK-LABEL: func @swizzle_adjust_add_offset
 //  CHECK-SAME:   %[[SRC:[A-Za-z0-9]+]]: memref<?xf32>
 //  CHECK-SAME:   %[[VEC:[A-Za-z0-9]+]]: vector<4xf32>
 //  CHECK-SAME:   %[[OFFSET:[A-Za-z0-9]+]]: index
 //   CHECK-DAG:   %[[ROW_WIDTH:.+]] = arith.constant 64 : index
 //   CHECK-DAG:   %[[GROUP_COUNT:.+]] = arith.constant 16 : index
 //   CHECK-DAG:   %[[GROUP_WIDTH:.+]] = arith.constant 4 : index
-//       CHECK:   %[[APPLY_BASE:.+]] = affine.apply affine_map<()[s0] -> (s0 + 16)>()[%[[OFFSET]]]
+//   CHECK-DAG:   %[[C1040:.+]] = arith.constant 1040 : index
+//       CHECK:   %[[APPLY_BASE:.+]] = arith.addi %[[OFFSET]], %[[GROUP_COUNT]] overflow<nsw> : index
 //       CHECK:   %[[I:.+]] = arith.divui %[[APPLY_BASE]], %[[ROW_WIDTH]] : index
 //       CHECK:   %[[JELEM:.+]] = arith.remui %[[APPLY_BASE]], %[[ROW_WIDTH]] : index
 //       CHECK:   %[[J:.+]] = arith.divui %[[JELEM]], %[[GROUP_WIDTH]] : index
@@ -184,7 +187,7 @@ func.func @swizzle_adjust_affine_offset(%src: memref<?xf32>, %vec: vector<4xf32>
 
 //       CHECK:   %[[VECTOR:.+]] = vector.load %[[SRC]][%[[SWOFF]]]
 
-//       CHECK:   %[[STORE_BASE:.+]] = affine.apply affine_map<()[s0] -> (s0 + 1040)>()[%[[OFFSET]]]
+//       CHECK:   %[[STORE_BASE:.+]] = arith.addi %[[OFFSET]], %[[C1040]] overflow<nsw> : index
 //       CHECK:   %[[OFFSET_DIFF:.+]] = arith.subi %[[SWOFF]], %[[APPLY_BASE]] : index
 //       CHECK:   %[[STORE_SWOFF:.+]] = arith.addi %[[STORE_BASE]], %[[OFFSET_DIFF]] : index
 //       CHECK:   vector.store %[[VEC]], %[[SRC]][%[[STORE_SWOFF]]]

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp

@@ -467,55 +467,41 @@ int64_t WorkgroupMappingAttr::getRelativeIndex() const {
 // iree_codegen.rotate_rows
 //===---------------------------------------------------------------------===//
 
-/// Given `r = |rotationInvariant|`, simplify affine maps of the following form:
+/// Given `r = |rotationInvariant|`, simplify additions of the following form:
 ///
-/// ```
-///   %offset = affine.apply (d0, ..., dn) -> (f(d0, ..., dn) + c)
-/// ```
+/// %offset = arith.addi %0, c
 ///
 /// Where `c` is a constant, to:
 ///
-/// ```
-///   %offset = affine.apply (d0, ..., dn) -> (f(d0, ..., dn) + c % r)
-/// ```
-static OpFoldResult getMinimumConstantOffsetMap(OpBuilder &b, Location loc,
-                                                OpFoldResult offset,
-                                                int64_t rotationInvariant) {
+/// %offset = arith.addi %0, c % r
+static OpFoldResult getMinimumConstantOffsetValue(OpBuilder &b, Location loc,
+                                                  OpFoldResult offset,
+                                                  int64_t rotationInvariant) {
   auto value = dyn_cast_if_present<Value>(offset);
   if (!value)
     return offset;
 
-  auto apply = value.getDefiningOp<affine::AffineApplyOp>();
-  if (!apply)
+  auto add = value.getDefiningOp<arith::AddIOp>();
+  if (!add)
     return offset;
 
-  AffineMap map = apply.getMap();
-  // Simplify the map to move `+ c` terms to the right most (first) expression
-  // in the tree.
-  map = simplifyAffineMap(map);
-  AffineExpr resultExpr = map.getResult(0);
-  auto addExpr = llvm::dyn_cast<AffineBinaryOpExpr>(resultExpr);
-
-  // After simplification, the add should be the first expression if present.
-  if (!addExpr || addExpr.getKind() != AffineExprKind::Add)
-    return offset;
-
-  // If RHS is not constant, nothing to do.
-  auto constantRhs = llvm::dyn_cast<AffineConstantExpr>(addExpr.getRHS());
-  if (!constantRhs)
+  llvm::APInt constant;
+  if (!matchPattern(add.getRhs(), m_ConstantInt(&constant)))
     return offset;
 
-  int64_t constantOffset = constantRhs.getValue();
+  int64_t constantOffset = constant.getSExtValue();
   int64_t baseMod = constantOffset % rotationInvariant;
 
   // Skip constructing the new apply if it's not needed (c < rotationInvariant).
   if (baseMod == constantOffset)
     return offset;
 
-  AffineExpr newExpr =
-      addExpr.getLHS() + getAffineConstantExpr(baseMod, b.getContext());
-  map = AffineMap::get(map.getNumDims(), map.getNumSymbols(), newExpr);
-  return b.create<affine::AffineApplyOp>(loc, map, apply.getOperands())
+  Value modOffset = b.create<arith::ConstantIndexOp>(loc, baseMod);
+  // If the original add is nsw/nuw, then the new add must also be given we're
+  // adding a smaller value.
+  return b
+      .create<arith::AddIOp>(loc, add.getLhs(), modOffset,
+                             add.getOverflowFlags())
       .getResult();
 }
 
@@ -531,7 +517,7 @@ OpFoldResult RotateRowsAttr::swizzleOffset(OpBuilder &b, Location loc,
   int64_t rotationInvariant =
       getRowWidth() * (getRowWidth() / getAccessWidth());
   OpFoldResult id =
-      getMinimumConstantOffsetMap(b, loc, offset, rotationInvariant);
+      getMinimumConstantOffsetValue(b, loc, offset, rotationInvariant);
 
   // Number of elements per row.
   Value rowAlignmentVal = b.create<arith::ConstantIndexOp>(loc, getRowWidth());