[mlir][Vector] Add custom slt / SCF.if folding to VectorToSCF

[nicolasvasilache]

Summary

scf.if currently lacks folding on true / false conditionals.
Such foldings are a bit more involved than can be addressed immediately.
This revision introduces an eager folding for lowering vector.transfer operations in the presence of unrolling.

[nicolasvasilache]

Fix test

[nicolasvasilache]

Fix typo in test

[ThomasRaoux]

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

193

Should it be *maybeCstV < *maybeCstUb since we are trying to fold less than?

217

Don't we need to increment idx even if the condition was folded?

[nicolasvasilache]

Address review

[nicolasvasilache]

Thanks @ThomasRaoux

[ftynse]

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

215

Nit: I'd give this a different name. if (inBounds) reads like "if the transfer is in bounds", but in fact means "if there exists a dynamic condition for the transfer to be in bounds" and if (!inBounds) will actually give a static guarantee that the transfer is in bounds.

[nicolasvasilache]

Rename variable

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

affineLoopNestBuilder(
           ValueRange indices(xferOp.indices());
           loopBodyBuilder(majorIvs, indices.take_front(leadingRank),
                           indices.drop_front(leadingRank).take_front(majorRank),
                           indices.take_back(minorRank), memrefBoundsCapture);
         });
   }
 }
 
+static Optional<int64_t> extractConstantIndex(Value v) {
+  if (auto cstOp = v.getDefiningOp<ConstantIndexOp>())
+    return cstOp.getValue();
+  if (auto affineApplyOp = v.getDefiningOp<AffineApplyOp>())
+    if (affineApplyOp.getAffineMap().isSingleConstant())
+      return affineApplyOp.getAffineMap().getSingleConstantResult();
+  return None;
+}
+
+// Missing foldings of scf.if make it necessary to perform poor man's folding
+// eagerly, especially in the case of unrolling. In the future, this should go
+// away once scf.if folds properly.
+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(
     ValueRange majorIvs, ValueRange majorOffsets,
     MemRefBoundsCapture &memrefBounds,
     SmallVectorImpl<Value> &majorIvsPlusOffsets) {
   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 = slt(majorIvsPlusOffsets.back(), ub);
+      Value inBounds = onTheFlyFoldSLT(majorIvsPlusOffsets.back(), ub);
+      if (!inBounds)
+        continue;
       inBoundsCondition =
           (inBoundsCondition) ? (inBoundsCondition && inBounds) : inBounds;
     }
     ++idx;
   }
   return inBoundsCondition;
 }
 
 // TODO: Parallelism and threadlocal considerations.
 static Value setAllocAtFunctionEntry(MemRefType memRefMinorVectorType,
                                      Operation *op) {
   auto &b = ScopedContext::getBuilderRef();

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

 // RUN: mlir-opt %s -convert-vector-to-scf -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -convert-vector-to-scf=full-unroll=true -split-input-file | FileCheck %s --check-prefix=FULL-UNROLL
+// RUN: mlir-opt %s -convert-vector-to-scf=full-unroll=true -split-input-file
+//| FileCheck %s --check-prefix=FULL-UNROLL
 
 // CHECK-LABEL: func @materialize_read_1d() {
 func @materialize_read_1d() {
   %f0 = constant 0.0: f32
   %A = alloc () : memref<7x42xf32>
   affine.for %i0 = 0 to 7 step 4 {
     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>
func @transfer_write_progressive_not_masked(%A : memref<?x?xf32>, %base: index, %vec: vector<3x15xf32>) {
   // FULL-UNROLL: vector.transfer_write %2, %[[A]][%[[I1]], %[[base]]] : vector<15xf32>, memref<?x?xf32>
   // FULL-UNROLL: %[[I2:.*]] = affine.apply #[[$MAP2]]()[%[[base]]]
   // FULL-UNROLL: %[[VEC2:.*]] = vector.extract %[[vec]][2] : vector<3x15xf32>
   // FULL-UNROLL: vector.transfer_write %[[VEC2:.*]], %[[A]][%[[I2]], %[[base]]] : vector<15xf32>, memref<?x?xf32>
   vector.transfer_write %vec, %A[%base, %base] {masked = [false, false]} :
     vector<3x15xf32>, memref<?x?xf32>
   return
 }
+
+// -----
+
+func @transfer_read_simple(%A : memref<2x2xf32>) -> vector<2x2xf32> {
+  %c0 = constant 0 : index
+  %f0 = constant 0.0 : f32
+  // FULL-UNROLL: %[[VC0]] = constant dense<0.000000e+00> : vector<2x2xf32>
+  // FULL-UNROLL: %[[C0:.*]] = constant 0 : index
+  // FULL-UNROLL: %[[C1:.*]] = constant 1 : index
+  // FULL-UNROLL: %[[V0:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]]]
+  // FULL-UNROLL: %[[RES0:.*]] = vector.insert %[[V0]], %[[VC0]] [0] : vector<2xf32> into vector<2x2xf32>
+  // FULL-UNROLL: %[[V1:.*]] = vector.transfer_read %{{.*}}[%[[C1]], %[[C0]]]
+  // FULL-UNROLL: %[[RES1:.*]] = vector.insert %[[C1]], %[[RES0]] [1] : vector<2xf32> into vector<2x2xf32>
+  %0 = vector.transfer_read %A[%c0, %c0], %f0 : memref<2x2xf32>, vector<2x2xf32>
+  return %0 : vector<2x2xf32>
+}