add unsigned comparison builders to Affine EDSC

[adstraw]

Summary

Current Affine comparison builders, which use operator overload, default to signed comparison. This creates the possibility of misuse of these builders and potential correctness issues when dealing with unsigned integers. This change makes the distinction between signed and unsigned comparison builders and forces the caller to make a choice between the two.

[adstraw]

@ftynse I do not have push privileges. Can you submit this patch for me?

[ftynse]

@adstraw could you please rebase on head? This does not seem to apply anymore

[adstraw]

@ftynse patch updated.

mlir/include/mlir/Dialect/Affine/EDSC/Builders.h

 Value negate(Value value);
 Value operator&&(Value lhs, Value rhs);
 Value operator||(Value lhs, Value rhs);
 Value operator^(Value lhs, Value rhs);
 
 /// Comparison operator overloadings.
 Value eq(Value lhs, Value rhs);
 Value ne(Value lhs, Value rhs);
-Value operator<(Value lhs, Value rhs);
-Value operator<=(Value lhs, Value rhs);
-Value operator>(Value lhs, Value rhs);
-Value operator>=(Value lhs, Value rhs);
+Value slt(Value lhs, Value rhs);
+Value sle(Value lhs, Value rhs);
+Value sgt(Value lhs, Value rhs);
+Value sge(Value lhs, Value rhs);
+Value ult(Value lhs, Value rhs);
+Value ule(Value lhs, Value rhs);
+Value ugt(Value lhs, Value rhs);
+Value uge(Value lhs, Value rhs);
 
 } // namespace op
 
 /// Arithmetic operator overloadings.
 template <typename Load, typename Store>
 Value TemplatedIndexedValue<Load, Store>::operator+(Value e) {
   using op::operator+;
   return static_cast<Value>(*this) + e;

 Value TemplatedIndexedValue<Load, Store>::eq(Value e) {
   return eq(value, e);
 }
 template <typename Load, typename Store>
 Value TemplatedIndexedValue<Load, Store>::ne(Value e) {
   return ne(value, e);
 }
 template <typename Load, typename Store>
-Value TemplatedIndexedValue<Load, Store>::operator<(Value e) {
-  using op::operator<;
-  return static_cast<Value>(*this) < e;
+Value TemplatedIndexedValue<Load, Store>::slt(Value e) {
+  using op::slt;
+  return slt(static_cast<Value>(*this), e);
 }
 template <typename Load, typename Store>
-Value TemplatedIndexedValue<Load, Store>::operator<=(Value e) {
-  using op::operator<=;
-  return static_cast<Value>(*this) <= e;
+Value TemplatedIndexedValue<Load, Store>::sle(Value e) {
+  using op::sle;
+  return sle(static_cast<Value>(*this), e);
 }
 template <typename Load, typename Store>
-Value TemplatedIndexedValue<Load, Store>::operator>(Value e) {
-  using op::operator>;
-  return static_cast<Value>(*this) > e;
+Value TemplatedIndexedValue<Load, Store>::sgt(Value e) {
+  using op::sgt;
+  return sgt(static_cast<Value>(*this), e);
 }
 template <typename Load, typename Store>
-Value TemplatedIndexedValue<Load, Store>::operator>=(Value e) {
-  using op::operator>=;
-  return static_cast<Value>(*this) >= e;
+Value TemplatedIndexedValue<Load, Store>::sge(Value e) {
+  using op::sge;
+  return sge(static_cast<Value>(*this), e);
+}
+template <typename Load, typename Store>
+Value TemplatedIndexedValue<Load, Store>::ult(Value e) {
+  using op::ult;
+  return ult(static_cast<Value>(*this), e);
+}
+template <typename Load, typename Store>
+Value TemplatedIndexedValue<Load, Store>::ule(Value e) {
+  using op::ule;
+  return ule(static_cast<Value>(*this), e);
+}
+template <typename Load, typename Store>
+Value TemplatedIndexedValue<Load, Store>::ugt(Value e) {
+  using op::ugt;
+  return ugt(static_cast<Value>(*this), e);
+}
+template <typename Load, typename Store>
+Value TemplatedIndexedValue<Load, Store>::uge(Value e) {
+  using op::uge;
+  return uge(static_cast<Value>(*this), e);
 }
 
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_EDSC_BUILDERS_H_

mlir/include/mlir/EDSC/Builders.h

public:
   }
   Value operator||(TemplatedIndexedValue e) {
     return *this || static_cast<Value>(e);
   }
 
   /// Comparison operator overloadings.
   Value eq(Value e);
   Value ne(Value e);
-  Value operator<(Value e);
-  Value operator<=(Value e);
-  Value operator>(Value e);
-  Value operator>=(Value e);
-  Value operator<(TemplatedIndexedValue e) {
-    return *this < static_cast<Value>(e);
+  Value slt(Value e);
+  Value sle(Value e);
+  Value sgt(Value e);
+  Value sge(Value e);
+  Value ult(Value e);
+  Value ule(Value e);
+  Value ugt(Value e);
+  Value uge(Value e);
+  Value slt(TemplatedIndexedValue e) {
+    return slt(*this, static_cast<Value>(e));
+  }
+  Value sle(TemplatedIndexedValue e) {
+    return sle(*this, static_cast<Value>(e));
+  }
+  Value sgt(TemplatedIndexedValue e) {
+    return sgt(*this, static_cast<Value>(e));
+  }
+  Value sge(TemplatedIndexedValue e) {
+    return sge(*this, static_cast<Value>(e));
+  }
+  Value ult(TemplatedIndexedValue e) {
+    return ult(*this, static_cast<Value>(e));
   }
-  Value operator<=(TemplatedIndexedValue e) {
-    return *this <= static_cast<Value>(e);
+  Value ule(TemplatedIndexedValue e) {
+    return ule(*this, static_cast<Value>(e));
   }
-  Value operator>(TemplatedIndexedValue e) {
-    return *this > static_cast<Value>(e);
+  Value ugt(TemplatedIndexedValue e) {
+    return ugt(*this, static_cast<Value>(e));
   }
-  Value operator>=(TemplatedIndexedValue e) {
-    return *this >= static_cast<Value>(e);
+  Value uge(TemplatedIndexedValue e) {
+    return uge(*this, static_cast<Value>(e));
   }
 
 private:
   TemplatedIndexedValue(Value value, ArrayRef<Value> indices)
       : value(value), indices(indices.begin(), indices.end()) {}
 
   TemplatedIndexedValue &append() { return *this; }
 

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

Value NDTransferOpHelper<ConcreteOp>::emitInBoundsCondition(
   Value inBoundsCondition;
   majorIvsPlusOffsets.reserve(majorIvs.size());
   unsigned idx = 0;
   for (auto it : llvm::zip(majorIvs, majorOffsets, memrefBounds.getUbs())) {
     Value iv = std::get<0>(it), off = std::get<1>(it), ub = std::get<2>(it);
     using namespace mlir::edsc::op;
     majorIvsPlusOffsets.push_back(iv + off);
     if (xferOp.isMaskedDim(leadingRank + idx)) {
-      Value inBounds = majorIvsPlusOffsets.back() < ub;
+      Value inBounds = slt(majorIvsPlusOffsets.back(), ub);
       inBoundsCondition =
           (inBoundsCondition) ? (inBoundsCondition && inBounds) : inBounds;
     }
     ++idx;
   }
   return inBoundsCondition;
 }
 
for (unsigned memRefDim = 0; memRefDim < clippedScalarAccessExprs.size();
     // 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(i < N, i, N_minus_1);
+      auto select_1 = std_select(slt(i, N), i, N_minus_1);
       clippedScalarAccessExprs[memRefDim] =
-          std_select(i < zero, zero, select_1);
+          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(i_plus_ii < N, i_plus_ii, N_minus_1);
+      auto select_1 = std_select(slt(i_plus_ii, N), i_plus_ii, N_minus_1);
       clippedScalarAccessExprs[memRefDim] =
-          std_select(i_plus_ii < zero, zero, select_1);
+          std_select(slt(i_plus_ii, zero), zero, select_1);
     }
   }
 
   return clippedScalarAccessExprs;
 }
 
 namespace mlir {
 

mlir/lib/Dialect/Affine/EDSC/Builders.cpp

return type.isa<FloatType>()
              : createIComparisonExpr(CmpIPredicate::eq, lhs, rhs);
 }
 Value mlir::edsc::op::ne(Value lhs, Value rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::ONE, lhs, rhs)
              : createIComparisonExpr(CmpIPredicate::ne, lhs, rhs);
 }
-Value mlir::edsc::op::operator<(Value lhs, Value rhs) {
+Value mlir::edsc::op::slt(Value lhs, Value rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::OLT, lhs, rhs)
-             :
-             // TODO(ntv,zinenko): signed by default, how about unsigned?
-             createIComparisonExpr(CmpIPredicate::slt, lhs, rhs);
+             : createIComparisonExpr(CmpIPredicate::slt, lhs, rhs);
 }
-Value mlir::edsc::op::operator<=(Value lhs, Value rhs) {
+Value mlir::edsc::op::sle(Value lhs, Value rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::OLE, lhs, rhs)
              : createIComparisonExpr(CmpIPredicate::sle, lhs, rhs);
 }
-Value mlir::edsc::op::operator>(Value lhs, Value rhs) {
+Value mlir::edsc::op::sgt(Value lhs, Value rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::OGT, lhs, rhs)
              : createIComparisonExpr(CmpIPredicate::sgt, lhs, rhs);
 }
-Value mlir::edsc::op::operator>=(Value lhs, Value rhs) {
+Value mlir::edsc::op::sge(Value lhs, Value rhs) {
   auto type = lhs.getType();
   return type.isa<FloatType>()
              ? createFComparisonExpr(CmpFPredicate::OGE, lhs, rhs)
              : createIComparisonExpr(CmpIPredicate::sge, lhs, rhs);
 }
+Value mlir::edsc::op::ult(Value lhs, Value rhs) {
+  auto type = lhs.getType();
+  return type.isa<FloatType>()
+             ? createFComparisonExpr(CmpFPredicate::OLT, lhs, rhs)
+             : createIComparisonExpr(CmpIPredicate::ult, lhs, rhs);
+}
+Value mlir::edsc::op::ule(Value lhs, Value rhs) {
+  auto type = lhs.getType();
+  return type.isa<FloatType>()
+             ? createFComparisonExpr(CmpFPredicate::OLE, lhs, rhs)
+             : createIComparisonExpr(CmpIPredicate::ule, lhs, rhs);
+}
+Value mlir::edsc::op::ugt(Value lhs, Value rhs) {
+  auto type = lhs.getType();
+  return type.isa<FloatType>()
+             ? createFComparisonExpr(CmpFPredicate::OGT, lhs, rhs)
+             : createIComparisonExpr(CmpIPredicate::ugt, lhs, rhs);
+}
+Value mlir::edsc::op::uge(Value lhs, Value rhs) {
+  auto type = lhs.getType();
+  return type.isa<FloatType>()
+             ? createFComparisonExpr(CmpFPredicate::OGE, lhs, rhs)
+             : createIComparisonExpr(CmpIPredicate::uge, lhs, rhs);
+}

mlir/lib/Dialect/Linalg/EDSC/Builders.cpp

BinaryPointwiseOpBuilder binOp(
       [](Value a, Value b) -> Value { return a + b; });
   return linalg_generic_pointwise(binOp, I1, I2, O);
 }
 
 Operation *mlir::edsc::ops::linalg_generic_pointwise_max(StructuredIndexed I1,
                                                          StructuredIndexed I2,
                                                          StructuredIndexed O) {
   BinaryPointwiseOpBuilder binOp([](Value a, Value b) -> Value {
-    using edsc::op::operator>;
-    return std_select(a > b, a, b);
+    using edsc::op::sgt;
+    return std_select(sgt(a, b), a, b);
   });
   return linalg_generic_pointwise(binOp, I1, I2, O);
 }
 
 Operation *
 mlir::edsc::ops::linalg_generic_matmul(Value vA, Value vB, Value vC,
                                        MatmulRegionBuilder regionBuilder) {
   // clang-format off

mlir/lib/Dialect/Linalg/Transforms/Loops.cpp

for (auto iter : llvm::enumerate(imIdx)) {
     int idx = iter.index();
     auto dim = iter.value();
     // Only need to iterate over the window dimensions.
     if (idx == 0 || idx == static_cast<int>(imIdx.size()) - 1) {
       clampedImIdx.push_back(dim);
       continue;
     }
 
-    using edsc::op::operator<;
-    using edsc::op::operator>=;
+    using edsc::op::sge;
+    using edsc::op::slt;
     using edsc::op::operator||;
-    Value leftOutOfBound = dim < zeroIndex;
+    Value leftOutOfBound = slt(dim, zeroIndex);
     if (conds.empty())
       conds.push_back(leftOutOfBound);
     else
       conds.push_back(conds.back() || leftOutOfBound);
     Value rightBound = std_dim(convOp.input(), idx);
-    conds.push_back(conds.back() || (dim >= rightBound));
+    conds.push_back(conds.back() || (sge(dim, rightBound)));
 
     // When padding is involved, the indices will only be shifted to negative,
     // so having a max op is enough.
     auto maxMap = AffineMap::get(/*dimCount=*/1, 0,
                                  {getAffineDimExpr(/*position=*/0, context),
                                   getAffineConstantExpr(0, context)},
                                  context);
     clampedImIdx.push_back(affine_max(dim.getType(), maxMap, ValueRange{dim}));

 }
 
 template <typename IndexedValueType>
 void emitScalarImplementation(ArrayRef<Value> allIvs, PoolingMaxOp op) {
   auto indices = getInputAndOutputIndices(allIvs, op);
   // Emit scalar form.
   Value lhs = std_load(op.output(), indices.outputs);
   Value rhs = std_load(op.input(), indices.inputs);
-  using edsc::op::operator>;
-  Value maxValue = std_select(lhs > rhs, lhs, rhs);
+  using edsc::op::sgt;
+  Value maxValue = std_select(sgt(lhs, rhs), lhs, rhs);
   std_store(maxValue, op.output(), indices.outputs);
 }
 template <typename IndexedValueType>
 void emitScalarImplementation(ArrayRef<Value> allIvs, PoolingMinOp op) {
   auto indices = getInputAndOutputIndices(allIvs, op);
   // Emit scalar form.
   Value lhs = std_load(op.output(), indices.outputs);
   Value rhs = std_load(op.input(), indices.inputs);
-  using edsc::op::operator<;
-  Value minValue = std_select(lhs < rhs, lhs, rhs);
+  using edsc::op::slt;
+  Value minValue = std_select(slt(lhs, rhs), lhs, rhs);
   std_store(minValue, op.output(), indices.outputs);
 }
 template <typename IndexedValueType>
 void emitScalarImplementation(ArrayRef<Value> allIvs, PoolingSumOp op) {
   auto indices = getInputAndOutputIndices(allIvs, op);
   IndexedValueType input(op.input()), output(op.output());
 
   // Emit scalar form.

mlir/test/EDSC/builder-api-test.cpp

TEST_FUNC(diviu_op_i32) {
   //  CHECK-NEXT:     {{.*}} = divi_unsigned
   // clang-format on
   f.print(llvm::outs());
   f.erase();
 }
 
 TEST_FUNC(select_op_i32) {
   using namespace edsc::op;
-  auto f32Type = FloatType::getF32(&globalContext());
+  auto i32Type = IntegerType::get(32, &globalContext());
   auto memrefType = MemRefType::get(
-      {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
+      {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, i32Type, {}, 0);
   auto f = makeFunction("select_op", {}, {memrefType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   Value zero = std_constant_index(0), one = std_constant_index(1);
   MemRefBoundsCapture vA(f.getArgument(0));
   AffineIndexedValue A(f.getArgument(0));
   affineLoopNestBuilder({zero, zero}, {one, one}, {1, 1}, [&](ValueRange ivs) {
-    std_select(eq(ivs[0], zero), A(zero, zero), A(ivs[0], ivs[1]));
+    using namespace edsc::op;
+    Value i = ivs[0], j = ivs[1];
+    std_select(eq(i, zero), A(zero, zero), A(i, j));
+    std_select(ne(i, zero), A(zero, zero), A(i, j));
+    std_select(slt(i, zero), A(zero, zero), A(i, j));
+    std_select(sle(i, zero), A(zero, zero), A(i, j));
+    std_select(sgt(i, zero), A(zero, zero), A(i, j));
+    std_select(sge(i, zero), A(zero, zero), A(i, j));
+    std_select(ult(i, zero), A(zero, zero), A(i, j));
+    std_select(ule(i, zero), A(zero, zero), A(i, j));
+    std_select(ugt(i, zero), A(zero, zero), A(i, j));
+    std_select(uge(i, zero), A(zero, zero), A(i, j));
   });
 
   // clang-format off
   // CHECK-LABEL: @select_op
   //      CHECK: affine.for %{{.*}} = 0 to 1 {
   // CHECK-NEXT:   affine.for %{{.*}} = 0 to 1 {
   //  CHECK-DAG:     {{.*}} = cmpi "eq"
   //  CHECK-DAG:     {{.*}} = affine.load
   //  CHECK-DAG:     {{.*}} = affine.load
   // CHECK-NEXT:     {{.*}} = select
+  //  CHECK-DAG:     {{.*}} = cmpi "ne"
+  //  CHECK-DAG:     {{.*}} = affine.load
+  //  CHECK-DAG:     {{.*}} = affine.load
+  // CHECK-NEXT:     {{.*}} = select
+  //  CHECK-DAG:     {{.*}} = cmpi "slt"
+  //  CHECK-DAG:     {{.*}} = affine.load
+  //  CHECK-DAG:     {{.*}} = affine.load
+  // CHECK-NEXT:     {{.*}} = select
+  //  CHECK-DAG:     {{.*}} = cmpi "sle"
+  //  CHECK-DAG:     {{.*}} = affine.load
+  //  CHECK-DAG:     {{.*}} = affine.load
+  // CHECK-NEXT:     {{.*}} = select
+  //  CHECK-DAG:     {{.*}} = cmpi "sgt"
+  //  CHECK-DAG:     {{.*}} = affine.load
+  //  CHECK-DAG:     {{.*}} = affine.load
+  // CHECK-NEXT:     {{.*}} = select
+  //  CHECK-DAG:     {{.*}} = cmpi "sge"
+  //  CHECK-DAG:     {{.*}} = affine.load
+  //  CHECK-DAG:     {{.*}} = affine.load
+  // CHECK-NEXT:     {{.*}} = select
+  //  CHECK-DAG:     {{.*}} = cmpi "ult"
+  //  CHECK-DAG:     {{.*}} = affine.load
+  //  CHECK-DAG:     {{.*}} = affine.load
+  // CHECK-NEXT:     {{.*}} = select
+  //  CHECK-DAG:     {{.*}} = cmpi "ule"
+  //  CHECK-DAG:     {{.*}} = affine.load
+  //  CHECK-DAG:     {{.*}} = affine.load
+  // CHECK-NEXT:     {{.*}} = select
+  //  CHECK-DAG:     {{.*}} = cmpi "ugt"
+  //  CHECK-DAG:     {{.*}} = affine.load
+  //  CHECK-DAG:     {{.*}} = affine.load
+  // CHECK-NEXT:     {{.*}} = select
+  //  CHECK-DAG:     {{.*}} = cmpi "uge"
+  //  CHECK-DAG:     {{.*}} = affine.load
+  //  CHECK-DAG:     {{.*}} = affine.load
+  // CHECK-NEXT:     {{.*}} = select
   // clang-format on
   f.print(llvm::outs());
   f.erase();
 }
 
 TEST_FUNC(select_op_f32) {
   auto f32Type = FloatType::getF32(&globalContext());
   auto memrefType = MemRefType::get(
       {ShapedType::kDynamicSize, ShapedType::kDynamicSize}, f32Type, {}, 0);
   auto f = makeFunction("select_op", {}, {memrefType, memrefType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   // clang-format off
   Value zero = std_constant_index(0), one = std_constant_index(1);
   MemRefBoundsCapture vA(f.getArgument(0)), vB(f.getArgument(1));
   AffineIndexedValue A(f.getArgument(0)), B(f.getArgument(1));
   affineLoopNestBuilder({zero, zero}, {one, one}, {1, 1}, [&](ValueRange ivs) {
     using namespace edsc::op;
     Value i = ivs[0], j = ivs[1];
     std_select(eq(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
     std_select(ne(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
-    std_select(B(i, j) >= B(i + one, j), A(zero, zero), A(i, j));
-    std_select(B(i, j) <= B(i + one, j), A(zero, zero), A(i, j));
-    std_select(B(i, j) < B(i + one, j), A(zero, zero), A(i, j));
-    std_select(B(i, j) > B(i + one, j), A(zero, zero), A(i, j));
+    std_select(sge(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
+    std_select(sle(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
+    std_select(slt(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
+    std_select(sgt(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
+    std_select(uge(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
+    std_select(ule(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
+    std_select(ult(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
+    std_select(ugt(B(i, j), B(i + one, j)), A(zero, zero), A(i, j));
   });
 
   // CHECK-LABEL: @select_op
   //      CHECK: affine.for %{{.*}} = 0 to 1 {
   // CHECK-NEXT:   affine.for %{{.*}} = 0 to 1 {
   //  CHECK-DAG:     cmpf "oeq"
   //  CHECK-DAG:     affine.load
   //  CHECK-DAG:     affine.load
TEST_FUNC(select_op_f32) {
   // CHECK-NEXT:     select
   //  CHECK-DAG:     cmpf "ogt"
   //  CHECK-DAG:     affine.load
   //  CHECK-DAG:     affine.load
   //  CHECK-DAG:     affine.load
   //  CHECK-DAG:     affine.load
   //  CHECK-DAG:     affine.apply
   // CHECK-NEXT:     select
+  //  CHECK-DAG:     cmpf "oge"
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.apply
+  // CHECK-NEXT:     select
+  //  CHECK-DAG:     cmpf "ole"
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.apply
+  // CHECK-NEXT:     select
+  //  CHECK-DAG:     cmpf "olt"
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.apply
+  // CHECK-NEXT:     select
+  //  CHECK-DAG:     cmpf "ogt"
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.load
+  //  CHECK-DAG:     affine.apply
+  // CHECK-NEXT:     select
   // clang-format on
   f.print(llvm::outs());
   f.erase();
 }
 
 // Inject an EDSC-constructed computation to exercise imperfectly nested 2-d
 // tiling.
 TEST_FUNC(tile_2d) {