[mlir][VectorOps] Redo the scalar loop emission in VectoToSCF to pad instead of clipping

[bkramer]

Summary

This replaces the select chain for edge-padding with an scf.if that
performs the memory operation when the index is in bounds and uses the
pad value when it's not. For transfer_write the same mechanism is used,
skipping the store when the index is out of bounds.

The integration test has a bunch of cases of how I believe this should
work.

[nicolasvasilache]

Thanks much for updating this older piece!

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

192

Return a Value of type i1 ?

228

would it be simpler to cast at call site and get rid of this template altogether ?

424–425

The artificial split was due to C++ comma operator undefined ordering.
Now that you have simplified this we could just write:

scalarAccessExprs[memRefDim] = (loopIndex < 0) ? i : i + ivs[loopIndex];

442

Any way to make this look a little more like the actual if-then-else that it builds?
Maybe with some clang-format ?

443

Could we rename to inBoundsFun/outOfBoundsFun or s/Fun/Lambda or something along those lines ?
When I looked at the use at the end of the function I had trouble figuring out what was what.

589

Better, thanks!

[rriddle]

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

444

nit: Static functions should be placed in the top-level namespace.

mlir/integration_test/Dialect/Vector/CPU/test-transfer-to-loops.mlir

 // RUN: mlir-opt %s -convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void  \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext,%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 #map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1) -> (d1)>
 
 func @print_memref_f32(memref<*xf32>)
 
 func @alloc_2d_filled_f32(%arg0: index, %arg1: index) -> memref<?x?xf32> {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c10 = constant 10 : index
   %c100 = constant 100 : index
func @alloc_2d_filled_f32(%arg0: index, %arg1: index) -> memref<?x?xf32> {
   }
   return %0 : memref<?x?xf32>
 }
 
 func @main() {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index
+  %c3 = constant 3 : index
   %c6 = constant 6 : index
   %cst = constant -4.2e+01 : f32
   %0 = call @alloc_2d_filled_f32(%c6, %c6) : (index, index) -> memref<?x?xf32>
   %converted = memref_cast %0 : memref<?x?xf32> to memref<*xf32>
   call @print_memref_f32(%converted): (memref<*xf32>) -> ()
   // CHECK:      Unranked{{.*}}data =
   // CHECK:      [
   // CHECK-SAME:  [0,   100,   200,   300,   400,   500],
func @main() {
   //   (i.e. last row: 105 .. 505, last col: 501 .. 505)
   // and the 4x4 subblock 202 .. 505 rooted @{0, 0} in the vector
   // CHECK-NEXT:  ( ( 202, 302, 402, 502, 501 ),
   // CHECK-SAME:    ( 203, 303, 403, 503, 502 ),
   // CHECK-SAME:    ( 204, 304, 404, 504, 503 ),
   // CHECK-SAME:    ( 205, 305, 405, 505, 504 ),
   // CHECK-SAME:    ( 105, 205, 305, 405, 505 ) )
 
+  %3 = vector.transfer_read %0[%c2, %c3], %cst : memref<?x?xf32>, vector<5x5xf32>
+  vector.print %3 : vector<5x5xf32>
+  // New 5x5 block rooted @{2, 3} in memory.
+  // CHECK-NEXT: ( ( 403, 503, 502, -42, -42 ),
+  // CHECK-SAME:   ( 404, 504, 503, -42, -42 ),
+  // CHECK-SAME:   ( 405, 505, 504, -42, -42 ),
+  // CHECK-SAME:   ( 305, 405, 505, -42, -42 ),
+  // CHECK-SAME:   ( -42, -42, -42, -42, -42 ) )
+
+  %4 = vector.transfer_read %0[%c2, %c3], %cst {permutation_map = #map0} : memref<?x?xf32>, vector<5x5xf32>
+  vector.print %4 : vector<5x5xf32>
+  // Transposed 5x5 block rooted @{2, 3} in memory.
+  // CHECK-NEXT: ( ( 403, 404, 405, 305, -42 ),
+  // CHECK-SAME:   ( 503, 504, 505, 405, -42 ),
+  // CHECK-SAME:   ( 502, 503, 504, 505, -42 ),
+  // CHECK-SAME:   ( -42, -42, -42, -42, -42 ),
+  // CHECK-SAME:   ( -42, -42, -42, -42, -42 ) )
+
+  %5 = vector.transfer_read %0[%c2, %c3], %cst {permutation_map = #map1} : memref<?x?xf32>, vector<5xf32>
+  vector.print %5 : vector<5xf32>
+  // CHECK-NEXT: ( 403, 503, 502, -42, -42 )
+
   dealloc %0 : memref<?x?xf32>
   return
 }

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 };
 
 /// Returns true if the memory underlying `memRefType` has a contiguous layout.
 /// Strides are written to `strides`.
 static bool isContiguous(MemRefType memRefType,
                          SmallVectorImpl<int64_t> &strides) {
   int64_t offset;
   auto successStrides = getStridesAndOffset(memRefType, strides, offset);
-  bool isContiguous = (strides.back() == 1);
+  bool isContiguous = strides.empty() || strides.back() == 1;
   if (isContiguous) {
     auto sizes = memRefType.getShape();
     for (int index = 0, e = strides.size() - 2; index < e; ++index) {
       if (strides[index] != strides[index + 1] * sizes[index + 1]) {
         isContiguous = false;
         break;
       }
     }

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

public:
   LogicalResult doReplace();
 
 private:
   /// Creates the loop nest on the "major" dimensions and calls the
   /// `loopBodyBuilder` lambda in the context of the loop nest.
   template <typename Lambda>
   void emitLoops(Lambda loopBodyBuilder);
 
-  /// Operate within the body of `emitLoops` to:
-  ///   1. Compute the indexings `majorIvs + majorOffsets` and save them in
-  ///      `majorIvsPlusOffsets`.
-  ///   2. Return a boolean that determines whether the first `majorIvs.rank()`
-  ///      dimensions `majorIvs + majorOffsets` are all within `memrefBounds`.
-  Value emitInBoundsCondition(ValueRange majorIvs, ValueRange majorOffsets,
-                              MemRefBoundsCapture &memrefBounds,
-                              SmallVectorImpl<Value> &majorIvsPlusOffsets);
-
   /// Common state to lower vector transfer ops.
   PatternRewriter &rewriter;
   const VectorTransferToSCFOptions &options;
   Location loc;
   std::unique_ptr<ScopedContext> scope;
   ConcreteOp xferOp;
   Operation *op;
   // A vector transfer copies data between:
static Value onTheFlyFoldSLT(Value v, Value ub) {
   using namespace mlir::edsc::op;
   auto maybeCstV = extractConstantIndex(v);
   auto maybeCstUb = extractConstantIndex(ub);
   if (maybeCstV && maybeCstUb && *maybeCstV < *maybeCstUb)
     return Value();
   return slt(v, ub);
 }
 
-template <typename ConcreteOp>
-Value NDTransferOpHelper<ConcreteOp>::emitInBoundsCondition(
+///   1. Compute the indexings `majorIvs + majorOffsets` and save them in
+///      `majorIvsPlusOffsets`.
+///   2. Return a value of i1 that determines whether the first `majorIvs.rank()`
+///      dimensions `majorIvs + majorOffsets` are all within `memrefBounds`.
+static Value
+emitInBoundsCondition(PatternRewriter &rewriter,
+                      VectorTransferOpInterface xferOp, unsigned leadingRank,
                       ValueRange majorIvs, ValueRange majorOffsets,
                       MemRefBoundsCapture &memrefBounds,
                       SmallVectorImpl<Value> &majorIvsPlusOffsets) {
   Value inBoundsCondition;
   majorIvsPlusOffsets.reserve(majorIvs.size());
   unsigned idx = 0;
   SmallVector<Value, 4> bounds =
       linalg::applyMapToValues(rewriter, xferOp.getLoc(),
                                xferOp.permutation_map(), memrefBounds.getUbs());
   for (auto it : llvm::zip(majorIvs, majorOffsets, bounds)) {
     Value iv = std::get<0>(it), off = std::get<1>(it), ub = std::get<2>(it);

 }
 
 // TODO: Parallelism and threadlocal considerations.
 static Value setAllocAtFunctionEntry(MemRefType memRefMinorVectorType,
                                      Operation *op) {
   auto &b = ScopedContext::getBuilderRef();
   OpBuilder::InsertionGuard guard(b);
   Operation *scope =
       op->getParentWithTrait<OpTrait::AutomaticAllocationScope>();
   assert(scope && "Expected op to be inside automatic allocation scope");
   b.setInsertionPointToStart(&scope->getRegion(0).front());
   Value res = std_alloca(memRefMinorVectorType);
   return res;
 }
 
 template <>
 LogicalResult NDTransferOpHelper<TransferReadOp>::doReplace() {
auto load1DVector = [&](ValueRange majorIvsPlusOffsets) -> Value {
                                   AffineMapAttr::get(map), xferOp.padding(),
                                   masked);
     };
 
     // 1. Compute the inBoundsCondition in the current loops ivs + offset
     // context.
     SmallVector<Value, 4> majorIvsPlusOffsets;
     Value inBoundsCondition = emitInBoundsCondition(
-        majorIvs, majorOffsets, memrefBounds, majorIvsPlusOffsets);
+        rewriter, cast<VectorTransferOpInterface>(xferOp.getOperation()),
+        leadingRank, majorIvs, majorOffsets, memrefBounds, majorIvsPlusOffsets);
 
     if (inBoundsCondition) {
       // 2. If the condition is not null, we need an IfOp, which may yield
       // if `options.unroll` is true.
       SmallVector<Type, 1> resultType;
       if (options.unroll)
         resultType.push_back(vectorType);
 
auto emitTransferWrite = [&](ValueRange majorIvsPlusOffsets) {
       vector_transfer_write(result, xferOp.memref(), indexing,
                             AffineMapAttr::get(map), masked);
     };
 
     // 1. Compute the inBoundsCondition in the current loops ivs + offset
     // context.
     SmallVector<Value, 4> majorIvsPlusOffsets;
     Value inBoundsCondition = emitInBoundsCondition(
-        majorIvs, majorOffsets, memrefBounds, majorIvsPlusOffsets);
+        rewriter, cast<VectorTransferOpInterface>(xferOp.getOperation()),
+        leadingRank, majorIvs, majorOffsets, memrefBounds, majorIvsPlusOffsets);
 
     if (inBoundsCondition) {
       // 2.a. If the condition is not null, we need an IfOp, to write
       // conditionally. Progressively lower to a 1-D transfer write.
       conditionBuilder(inBoundsCondition,
                        [&] { emitTransferWrite(majorIvsPlusOffsets); });
     } else {
       // 2.b. Guaranteed in-bounds. Progressively lower to a 1-D transfer write.
if (memRefDim == remoteRank - 1) {
       // memRefDim has coalescing properties, it should be swapped in the last
       // position.
       assert(coalescedIdx == -1 && "Unexpected > 1 coalesced indices");
       coalescedIdx = en.index();
     }
   }
   return coalescedIdx;
 }
 
-/// Emits remote memory accesses that are clipped to the boundaries of the
-/// MemRef.
+namespace mlir {
+
 template <typename TransferOpTy>
-static SmallVector<Value, 8>
-clip(TransferOpTy transfer, MemRefBoundsCapture &bounds, ArrayRef<Value> ivs) {
-  using namespace mlir::edsc;
+VectorTransferRewriter<TransferOpTy>::VectorTransferRewriter(
+    VectorTransferToSCFOptions options, MLIRContext *context)
+    : RewritePattern(TransferOpTy::getOperationName(), 1, context),
+      options(options) {}
+
+/// Used for staging the transfer in a local buffer.
+template <typename TransferOpTy>
+MemRefType VectorTransferRewriter<TransferOpTy>::tmpMemRefType(
+    TransferOpTy transfer) const {
+  auto vectorType = transfer.getVectorType();
+  return MemRefType::get(vectorType.getShape().drop_back(),
+                         VectorType::get(vectorType.getShape().take_back(),
+                                         vectorType.getElementType()),
+                         {}, 0);
+}
+
+static void emitWithBoundsChecks(
+    PatternRewriter &rewriter, VectorTransferOpInterface transfer,
+    ValueRange ivs, MemRefBoundsCapture &memRefBoundsCapture,
+    function_ref<void(ArrayRef<Value>)> inBoundsFun,
+    function_ref<void(ArrayRef<Value>)> outOfBoundsFun = nullptr) {
+  // Permute the incoming indices according to the permutation map.
+  SmallVector<Value, 4> indices =
+      linalg::applyMapToValues(rewriter, transfer.getLoc(),
+                               transfer.permutation_map(), transfer.indices());
 
-  Value zero(std_constant_index(0)), one(std_constant_index(1));
-  SmallVector<Value, 8> memRefAccess(transfer.indices());
-  SmallVector<Value, 8> clippedScalarAccessExprs(memRefAccess.size());
-  // Indices accessing to remote memory are clipped and their expressions are
-  // returned in clippedScalarAccessExprs.
-  for (unsigned memRefDim = 0; memRefDim < clippedScalarAccessExprs.size();
+  // Generate a bounds check if necessary.
+  SmallVector<Value, 4> majorIvsPlusOffsets;
+  Value inBoundsCondition =
+      emitInBoundsCondition(rewriter, transfer, 0, ivs, indices,
+                            memRefBoundsCapture, majorIvsPlusOffsets);
+
+  // Apply the permutation map to the ivs. The permutation map may not use all
+  // the inputs.
+  SmallVector<Value, 4> scalarAccessExprs(transfer.indices().size());
+  for (unsigned memRefDim = 0; memRefDim < transfer.indices().size();
        ++memRefDim) {
     // Linear search on a small number of entries.
     int loopIndex = -1;
     auto exprs = transfer.permutation_map().getResults();
     for (auto en : llvm::enumerate(exprs)) {
       auto expr = en.value();
-      auto dim = expr.template dyn_cast<AffineDimExpr>();
+      auto dim = expr.dyn_cast<AffineDimExpr>();
       // Sanity check.
-      assert(
-          (dim || expr.template cast<AffineConstantExpr>().getValue() == 0) &&
+      assert((dim || expr.cast<AffineConstantExpr>().getValue() == 0) &&
              "Expected dim or 0 in permutationMap");
       if (dim && memRefDim == dim.getPosition()) {
         loopIndex = en.index();
         break;
       }
     }
 
-    // We cannot distinguish atm between unrolled dimensions that implement
-    // the "always full" tile abstraction and need clipping from the other
-    // ones. So we conservatively clip everything.
     using namespace edsc::op;
-    auto N = bounds.ub(memRefDim);
-    auto i = memRefAccess[memRefDim];
-    if (loopIndex < 0) {
-      auto N_minus_1 = N - one;
-      auto select_1 = std_select(slt(i, N), i, N_minus_1);
-      clippedScalarAccessExprs[memRefDim] =
-          std_select(slt(i, zero), zero, select_1);
-    } else {
-      auto ii = ivs[loopIndex];
-      auto i_plus_ii = i + ii;
-      auto N_minus_1 = N - one;
-      auto select_1 = std_select(slt(i_plus_ii, N), i_plus_ii, N_minus_1);
-      clippedScalarAccessExprs[memRefDim] =
-          std_select(slt(i_plus_ii, zero), zero, select_1);
-    }
-  }
-
-  return clippedScalarAccessExprs;
+    auto i = transfer.indices()[memRefDim];
+    scalarAccessExprs[memRefDim] = loopIndex < 0 ? i : i + ivs[loopIndex];
   }
 
-namespace mlir {
-
-template <typename TransferOpTy>
-VectorTransferRewriter<TransferOpTy>::VectorTransferRewriter(
-    VectorTransferToSCFOptions options, MLIRContext *context)
-    : RewritePattern(TransferOpTy::getOperationName(), 1, context),
-      options(options) {}
-
-/// Used for staging the transfer in a local buffer.
-template <typename TransferOpTy>
-MemRefType VectorTransferRewriter<TransferOpTy>::tmpMemRefType(
-    TransferOpTy transfer) const {
-  auto vectorType = transfer.getVectorType();
-  return MemRefType::get(vectorType.getShape().drop_back(),
-                         VectorType::get(vectorType.getShape().take_back(),
-                                         vectorType.getElementType()),
-                         {}, 0);
+  if (inBoundsCondition)
+    conditionBuilder(
+        /* scf.if */ inBoundsCondition, // {
+        [&] { inBoundsFun(scalarAccessExprs); },
+        // } else {
+        outOfBoundsFun ? [&] { outOfBoundsFun(scalarAccessExprs); }
+                       : function_ref<void()>()
+        // }
+    );
+  else
+    inBoundsFun(scalarAccessExprs);
 }
 
 /// Lowers TransferReadOp into a combination of:
 ///   1. local memory allocation;
 ///   2. perfect loop nest over:
 ///      a. scalar load from local buffers (viewed as a scalar memref);
 ///      a. scalar store to original memref (with clipping).
 ///   3. vector_load from local buffer (viewed as a memref<1 x vector>);
LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite(
   for (auto step : vectorBoundsCapture.getSteps())
     steps.push_back(std_constant_index(step));
 
   // 2. Emit alloc-copy-load-dealloc.
   MLIRContext *ctx = op->getContext();
   Value tmp = setAllocAtFunctionEntry(tmpMemRefType(transfer), transfer);
   StdIndexedValue local(tmp);
   loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
-    auto ivs = llvm::to_vector<8>(loopIvs);
+    auto ivsStorage = llvm::to_vector<8>(loopIvs);
     // Swap the ivs which will reorder memory accesses.
     if (coalescedIdx >= 0)
-      std::swap(ivs.back(), ivs[coalescedIdx]);
-    // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist).
-    SmallVector<Value, 8> indices = clip(transfer, memRefBoundsCapture, ivs);
-    ArrayRef<Value> indicesRef(indices), ivsRef(ivs);
-    Value pos = std_index_cast(IntegerType::get(32, ctx), ivsRef.back());
-    Value scal = remote(indicesRef);
-    Value vector = vector_insert_element(scal, local(ivsRef.drop_back()), pos);
-    local(ivsRef.drop_back()) = vector;
+      std::swap(ivsStorage.back(), ivsStorage[coalescedIdx]);
+
+    ArrayRef<Value> ivs(ivsStorage);
+    Value pos = std_index_cast(IntegerType::get(32, ctx), ivs.back());
+    Value inVector = local(ivs.drop_back());
+    auto loadValue = [&](ArrayRef<Value> indices) {
+      Value vector = vector_insert_element(remote(indices), inVector, pos);
+      local(ivs.drop_back()) = vector;
+    };
+    auto loadPadding = [&](ArrayRef<Value>) {
+      Value vector = vector_insert_element(transfer.padding(), inVector, pos);
+      local(ivs.drop_back()) = vector;
+    };
+    emitWithBoundsChecks(
+        rewriter, cast<VectorTransferOpInterface>(transfer.getOperation()), ivs,
+        memRefBoundsCapture, loadValue, loadPadding);
   });
   Value vectorValue = std_load(vector_type_cast(tmp));
 
   // 3. Propagate.
   rewriter.replaceOp(op, vectorValue);
   return success();
 }
 
for (auto step : vectorBoundsCapture.getSteps())
     steps.push_back(std_constant_index(step));
 
   // 2. Emit alloc-store-copy-dealloc.
   Value tmp = setAllocAtFunctionEntry(tmpMemRefType(transfer), transfer);
   StdIndexedValue local(tmp);
   Value vec = vector_type_cast(tmp);
   std_store(vectorValue, vec);
   loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
-    auto ivs = llvm::to_vector<8>(loopIvs);
-    // Swap the ivs which will reorder memory accesses.
+    auto ivsStorage = llvm::to_vector<8>(loopIvs);
+    // Swap the ivsStorage which will reorder memory accesses.
     if (coalescedIdx >= 0)
-      std::swap(ivs.back(), ivs[coalescedIdx]);
-    // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist).
-    SmallVector<Value, 8> indices = clip(transfer, memRefBoundsCapture, ivs);
-    ArrayRef<Value> indicesRef(indices), ivsRef(ivs);
+      std::swap(ivsStorage.back(), ivsStorage[coalescedIdx]);
+
+    ArrayRef<Value> ivs(ivsStorage);
     Value pos =
-        std_index_cast(IntegerType::get(32, op->getContext()), ivsRef.back());
-    Value scalar = vector_extract_element(local(ivsRef.drop_back()), pos);
+        std_index_cast(IntegerType::get(32, op->getContext()), ivs.back());
+    auto storeValue = [&](ArrayRef<Value> indices) {
+      Value scalar = vector_extract_element(local(ivs.drop_back()), pos);
       remote(indices) = scalar;
+    };
+    emitWithBoundsChecks(
+        rewriter, cast<VectorTransferOpInterface>(transfer.getOperation()), ivs,
+        memRefBoundsCapture, storeValue);
   });
 
   // 3. Erase.
   rewriter.eraseOp(op);
   return success();
 }
 
 void populateVectorToSCFConversionPatterns(

mlir/test/Conversion/VectorToSCF/vector-to-loops.mlir

affine.for %i1 = 0 to 42 step 4 {
       %f1 = vector.transfer_read %A[%i0, %i1], %f0 {permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<7x42xf32>, vector<4xf32>
       %ip1 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i1)
       %f2 = vector.transfer_read %A[%i0, %ip1], %f0 {permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<7x42xf32>, vector<4xf32>
       %ip2 = affine.apply affine_map<(d0) -> (d0 + 2)> (%i1)
       %f3 = vector.transfer_read %A[%i0, %ip2], %f0 {permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<7x42xf32>, vector<4xf32>
       %ip3 = affine.apply affine_map<(d0) -> (d0 + 3)> (%i1)
       %f4 = vector.transfer_read %A[%i0, %ip3], %f0 {permutation_map = affine_map<(d0, d1) -> (d0)>} : memref<7x42xf32>, vector<4xf32>
       // Both accesses in the load must be clipped otherwise %i1 + 2 and %i1 + 3 will go out of bounds.
-      // CHECK: {{.*}} = select
-      // CHECK: %[[FILTERED1:.*]] = select
-      // CHECK: {{.*}} = select
-      // CHECK: %[[FILTERED2:.*]] = select
-      // CHECK: %{{.*}} = load {{.*}}[%[[FILTERED1]], %[[FILTERED2]]] : memref<7x42xf32>
+      // CHECK: scf.if
+      // CHECK-NEXT: load
+      // CHECK-NEXT: vector.insertelement
+      // CHECK-NEXT: store
+      // CHECK-NEXT: else
+      // CHECK-NEXT: vector.insertelement
+      // CHECK-NEXT: store
     }
   }
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @materialize_read_1d_partially_specialized
func @materialize_read_1d_partially_specialized(%dyn1 : index, %dyn2 : index, %dyn4 : index) {
   // CHECK-NOT: {{.*}} dim %[[tensor]], %c0
   // CHECK-NOT: {{.*}} dim %[[tensor]], %c3
   return
 }
 
 // -----
 
 // CHECK: #[[$ADD:map[0-9]+]] = affine_map<(d0, d1) -> (d0 + d1)>
-// CHECK: #[[$SUB:map[0-9]+]] = affine_map<()[s0] -> (s0 - 1)>
 
 // CHECK-LABEL: func @materialize_read(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
 func @materialize_read(%M: index, %N: index, %O: index, %P: index) {
   %f0 = constant 0.0: f32
   // CHECK-DAG:  %[[ALLOC:.*]] = alloca() : memref<5x4xvector<3xf32>>
   // CHECK-DAG:  %[[C0:.*]] = constant 0 : index
   // CHECK-DAG:  %[[C1:.*]] = constant 1 : index
   // CHECK-DAG:  %[[C3:.*]] = constant 3 : index
   // CHECK-DAG:  %[[C4:.*]] = constant 4 : index
   // CHECK-DAG:  %[[C5:.*]] = constant 5 : index
   //     CHECK:  %{{.*}} = alloc(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : memref<?x?x?x?xf32>
   // CHECK-NEXT:  affine.for %[[I0:.*]] = 0 to %{{.*}} step 3 {
   // CHECK-NEXT:    affine.for %[[I1:.*]] = 0 to %{{.*}} {
   // CHECK-NEXT:      affine.for %[[I2:.*]] = 0 to %{{.*}} {
   // CHECK-NEXT:        affine.for %[[I3:.*]] = 0 to %{{.*}} step 5 {
   // CHECK-NEXT:          scf.for %[[I4:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
   // CHECK-NEXT:            scf.for %[[I5:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
   // CHECK-NEXT:              scf.for %[[I6:.*]] = %[[C0]] to %[[C5]] step %[[C1]] {
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$ADD]](%[[I0]], %[[I4]])
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$SUB]]()[%{{.*}}]
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = select
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, %[[C0]] : index
-  // CHECK-NEXT:                %[[L0:.*]] = select
-  //
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$SUB]]()[%{{.*}}]
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = select
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, %[[C0]] : index
-  // CHECK-NEXT:                %[[L1:.*]] = select
-  //
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$SUB]]()[%{{.*}}]
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = select
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, %[[C0]] : index
-  // CHECK-NEXT:                %[[L2:.*]] = select
-  //
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$ADD]](%[[I3]], %[[I6]])
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$SUB]]()[%{{.*}}]
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = select
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, %[[C0]] : index
-  // CHECK-NEXT:                %[[L3:.*]] = select
-  // CHECK-NEXT:                %[[VIDX:.*]] = index_cast %[[I4]]
-  //
-  // CHECK-NEXT:                %[[SCAL:.*]] = load %{{.*}}[%[[L0]], %[[L1]], %[[L2]], %[[L3]]] : memref<?x?x?x?xf32>
-  // CHECK-NEXT:                %[[VEC:.*]] = load %[[ALLOC]][%[[I6]], %[[I5]]] : memref<5x4xvector<3xf32>>
-  // CHECK-NEXT:                %[[RVEC:.*]] = vector.insertelement %[[SCAL]], %[[VEC]][%[[VIDX]] : i32] : vector<3xf32>
-  // CHECK-NEXT:                store %[[RVEC]], %[[ALLOC]][%[[I6]], %[[I5]]] : memref<5x4xvector<3xf32>>
+  // CHECK:                     %[[VIDX:.*]] = index_cast %[[I4]]
+  // CHECK:                     %[[VEC:.*]] = load %[[ALLOC]][%[[I6]], %[[I5]]] : memref<5x4xvector<3xf32>>
+  // CHECK:                     %[[L0:.*]] = affine.apply #[[$ADD]](%[[I0]], %[[I4]])
+  // CHECK:                     %[[L3:.*]] = affine.apply #[[$ADD]](%[[I3]], %[[I6]])
+  // CHECK-NEXT:                scf.if
+  // CHECK-NEXT:                  %[[SCAL:.*]] = load %{{.*}}[%[[L0]], %[[I1]], %[[I2]], %[[L3]]] : memref<?x?x?x?xf32>
+  // CHECK-NEXT:                  %[[RVEC:.*]] = vector.insertelement %[[SCAL]], %[[VEC]][%[[VIDX]] : i32] : vector<3xf32>
+  // CHECK-NEXT:                  store %[[RVEC]], %[[ALLOC]][%[[I6]], %[[I5]]] : memref<5x4xvector<3xf32>>
+  // CHECK-NEXT:                } else {
+  // CHECK-NEXT:                  %[[CVEC:.*]] = vector.insertelement
+  // CHECK-NEXT:                  store %[[CVEC]], %[[ALLOC]][%[[I6]], %[[I5]]] : memref<5x4xvector<3xf32>>
+  // CHECK-NEXT:                }
   // CHECK-NEXT:              }
   // CHECK-NEXT:            }
   // CHECK-NEXT:          }
   // CHECK-NEXT:        }
   // CHECK-NEXT:      }
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
affine.for %i0 = 0 to %M step 3 {
     }
   }
   return
 }
 
 // -----
 
 // CHECK: #[[$ADD:map[0-9]+]] = affine_map<(d0, d1) -> (d0 + d1)>
-// CHECK: #[[$SUB:map[0-9]+]] = affine_map<()[s0] -> (s0 - 1)>
 
 // CHECK-LABEL:func @materialize_write(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
 func @materialize_write(%M: index, %N: index, %O: index, %P: index) {
   // CHECK-DAG:  %[[ALLOC:.*]] = alloca() : memref<5x4xvector<3xf32>>
   // CHECK-DAG:  %{{.*}} = constant dense<1.000000e+00> : vector<5x4x3xf32>
   // CHECK-DAG:  %[[C0:.*]] = constant 0 : index
   // CHECK-DAG:  %[[C1:.*]] = constant 1 : index
   // CHECK-DAG:  %[[C3:.*]] = constant 3 : index
   // CHECK-DAG:  %[[C4:.*]] = constant 4 : index
   // CHECK-DAG:  %[[C5:.*]] = constant 5 : index
   //     CHECK:  %{{.*}} = alloc(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : memref<?x?x?x?xf32>
   // CHECK-NEXT:  affine.for %[[I0:.*]] = 0 to %{{.*}} step 3 {
   // CHECK-NEXT:    affine.for %[[I1:.*]] = 0 to %{{.*}} step 4 {
   // CHECK-NEXT:      affine.for %[[I2:.*]] = 0 to %{{.*}} {
   // CHECK-NEXT:        affine.for %[[I3:.*]] = 0 to %{{.*}} step 5 {
   // CHECK-NEXT:          %[[VECTOR_VIEW:.*]] = vector.type_cast {{.*}} : memref<5x4xvector<3xf32>>
   //      CHECK:          store %{{.*}}, {{.*}} : memref<vector<5x4x3xf32>>
   // CHECK-NEXT:          scf.for %[[I4:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
   // CHECK-NEXT:            scf.for %[[I5:.*]] = %[[C0]] to %[[C4]] step %[[C1]] {
   // CHECK-NEXT:              scf.for %[[I6:.*]] = %[[C0]] to %[[C5]] step %[[C1]] {
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$ADD]](%[[I0]], %[[I4]])
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$SUB]]()[%{{.*}}]
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = select {{.*}}, {{.*}}, {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, %[[C0]] : index
-  // CHECK-NEXT:                %[[S0:.*]] = select {{.*}}, %[[C0]], {{.*}} : index
-  //
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$ADD]](%[[I1]], %[[I5]])
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$SUB]]()[%{{.*}}]
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = select {{.*}}, {{.*}}, {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, %[[C0]] : index
-  // CHECK-NEXT:                %[[S1:.*]] = select {{.*}}, %[[C0]], {{.*}} : index
-  //
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$SUB]]()[%{{.*}}]
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", %[[I2]], %{{.*}} : index
-  // CHECK-NEXT:                {{.*}} = select {{.*}}, %[[I2]], {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", %[[I2]], %[[C0]] : index
-  // CHECK-NEXT:                %[[S2:.*]] = select {{.*}}, %[[C0]], {{.*}} : index
-  //
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$ADD]](%[[I3]], %[[I6]])
-  // CHECK-NEXT:                {{.*}} = affine.apply #[[$SUB]]()[%{{.*}}]
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = select {{.*}}, {{.*}}, {{.*}} : index
-  // CHECK-NEXT:                {{.*}} = cmpi "slt", {{.*}}, %[[C0]] : index
-  // CHECK-NEXT:                %[[S3:.*]] = select {{.*}}, %[[C0]], {{.*}} : index
-  // CHECK-NEXT:                %[[VIDX:.*]] = index_cast %[[I4]]
-  //
-  // CHECK-NEXT:                %[[VEC:.*]] = load {{.*}}[%[[I6]], %[[I5]]] : memref<5x4xvector<3xf32>>
-  // CHECK-NEXT:                %[[SCAL:.*]] = vector.extractelement %[[VEC]][%[[VIDX]] : i32] : vector<3xf32>
-  // CHECK-NEXT:                store %[[SCAL]], {{.*}}[%[[S0]], %[[S1]], %[[S2]], %[[S3]]] : memref<?x?x?x?xf32>
+  // CHECK:                     %[[VIDX:.*]] = index_cast %[[I4]]
+  // CHECK:                     %[[S0:.*]] = affine.apply #[[$ADD]](%[[I0]], %[[I4]])
+  // CHECK:                     %[[S1:.*]] = affine.apply #[[$ADD]](%[[I1]], %[[I5]])
+  // CHECK:                     %[[S3:.*]] = affine.apply #[[$ADD]](%[[I3]], %[[I6]])
+  // CHECK-NEXT:                scf.if
+  // CHECK-NEXT:                  %[[VEC:.*]] = load {{.*}}[%[[I6]], %[[I5]]] : memref<5x4xvector<3xf32>>
+  // CHECK-NEXT:                  %[[SCAL:.*]] = vector.extractelement %[[VEC]][%[[VIDX]] : i32] : vector<3xf32>
+  //      CHECK:                  store %[[SCAL]], {{.*}}[%[[S0]], %[[S1]], %[[I2]], %[[S3]]] : memref<?x?x?x?xf32>
+  // CHECK-NEXT:                }
   // CHECK-NEXT:              }
   // CHECK-NEXT:            }
   // CHECK-NEXT:          }
   // CHECK-NEXT:        }
   // CHECK-NEXT:      }
   // CHECK-NEXT:    }
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return