[mlir][Vector] Provide progressive lowering of masked n-D vector transfers

[nicolasvasilache]

Summary

This revision allows masked vector transfers with m-D buffers and n-D vectors to
progressively lower to m-D buffer and 1-D vector transfers.

For a vector.transfer_read, assuming a memref<leading_dims) x (major_dims) x (minor_dims) x type> and a vector<(minor_dims) x type> are involved in the transfer, this generates pseudo-IR resembling:

if (any_of(%ivs_major + %offsets, <, major_dims)) {
  %v = vector_transfer_read(
    {%offsets_leading, %ivs_major + %offsets_major, %offsets_minor},
     %ivs_minor):
    memref<leading_dims) x (major_dims) x (minor_dims) x type>,
    vector<(minor_dims) x type>;
} else {
  %v = splat(vector<(minor_dims) x type>, %fill)
}

[rriddle]

mlir/lib/Conversion/VectorToLoops/ConvertVectorToLoops.cpp

71

///?

137

Why are you building a vector here? Is it because AffineLoopNestBuilder takes an ArrayRef<Value> instead of ValueRange?

159

inbounds &= inBounds2;?

[nicolasvasilache]

Format

[ftynse]

mlir/include/mlir/EDSC/Builders.h

122–123

This comment needs an update

129

This will crash if a block has an operation, which is not known to be a terminator (which may happen in rewrite patterns). I would prefer a more cautious way: get the last operation, check if it is a known terminator, and if so, update the insertion point.

205–206

Typo: "thes"

mlir/lib/Conversion/VectorToLoops/ConvertVectorToLoops.cpp

62

Here and below + in the commit description, the type spelling memref<leading_dims) x (major_dims) x (minor_dims) x type> looks broken to me.

94

Is this memref always in the default address space? What if the memref we are transferring to/from is in the non-default address space?

98

plz document here and below

101

Nit: this "1." without any continuation in the list looks fishy

104

ValueRange?

153

Nit: please add /*name=*/ annotations to differentiate which 1 corresponds to what

184

Maybe use minorRank instead of hardcoded 1 here?

586–587

Why rename Conversion to Lowering? We decided to use the former instead of the latter....

mlir/lib/Dialect/Vector/VectorOps.cpp

1491

Should we also carry over the address space of the original memref type?

[nicolasvasilache]

mlir/lib/Conversion/VectorToLoops/ConvertVectorToLoops.cpp

137

It's because the first argument of AffineLoopNestBuilder is a MutableArrayRef<Value> which captures the induction variables so they can be reused.

159

This would require Value::operator&=(Value) I believe ?

[nicolasvasilache]

Address review comments.

mlir/include/mlir/Conversion/VectorToLoops/ConvertVectorToLoops.h

 //===- ConvertVectorToLoops.h - Utils to convert from the vector dialect --===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_VECTORTOLLVM_CONVERTVECTORTOLOOPS_H_
 #define MLIR_CONVERSION_VECTORTOLLVM_CONVERTVECTORTOLOOPS_H_
 
 namespace mlir {
 class MLIRContext;
 class OwningRewritePatternList;
 
 /// Collect a set of patterns to convert from the Vector dialect to loops + std.
-void populateVectorToAffineLoopsConversionPatterns(
-    MLIRContext *context, OwningRewritePatternList &patterns);
+void populateVectorToLoopsConversionPatterns(OwningRewritePatternList &patterns,
+                                             MLIRContext *context);
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_VECTORTOLLVM_CONVERTVECTORTOLOOPS_H_

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

 
 /// A VectorBoundsCapture represents the information required to step through a
 /// Vector accessing each scalar element at a time. It is the counterpart of
 /// a MemRefBoundsCapture but for vectors. This exists purely for boilerplate
 /// avoidance.
 class VectorBoundsCapture : public BoundsCapture {
 public:
   explicit VectorBoundsCapture(Value v);
+  explicit VectorBoundsCapture(VectorType t);
   VectorBoundsCapture(const VectorBoundsCapture &) = default;
   VectorBoundsCapture &operator=(const VectorBoundsCapture &) = default;
 
 private:
   Value base;
 };
 
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_STANDARDOPS_EDSC_BUILDERS_H_

mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h

 //===- Intrinsics.h - MLIR EDSC Intrinsics for StandardOps ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_DIALECT_STANDARDOPS_EDSC_INTRINSICS_H_
 #define MLIR_DIALECT_STANDARDOPS_EDSC_INTRINSICS_H_
 
 #include "mlir/Dialect/StandardOps/EDSC/Builders.h"
 
 namespace mlir {
 namespace edsc {
 namespace intrinsics {
 
+using std_addi = ValueBuilder<AddIOp>;
 using std_addf = ValueBuilder<AddFOp>;
 using std_alloc = ValueBuilder<AllocOp>;
+using std_alloca = ValueBuilder<AllocaOp>;
 using std_call = OperationBuilder<CallOp>;
 using std_constant = ValueBuilder<ConstantOp>;
 using std_constant_float = ValueBuilder<ConstantFloatOp>;
 using std_constant_index = ValueBuilder<ConstantIndexOp>;
 using std_constant_int = ValueBuilder<ConstantIntOp>;
 using std_dealloc = OperationBuilder<DeallocOp>;
 using std_dim = ValueBuilder<DimOp>;
 using std_extract_element = ValueBuilder<ExtractElementOp>;
 using std_index_cast = ValueBuilder<IndexCastOp>;
 using std_muli = ValueBuilder<MulIOp>;
 using std_mulf = ValueBuilder<MulFOp>;
 using std_memref_cast = ValueBuilder<MemRefCastOp>;
 using std_ret = OperationBuilder<ReturnOp>;
 using std_select = ValueBuilder<SelectOp>;
 using std_load = ValueBuilder<LoadOp>;
+using std_splat = ValueBuilder<SplatOp>;
 using std_store = OperationBuilder<StoreOp>;
 using std_subi = ValueBuilder<SubIOp>;
 using std_sub_view = ValueBuilder<SubViewOp>;
 using std_tanh = ValueBuilder<TanhOp>;
 using std_tensor_load = ValueBuilder<TensorLoadOp>;
 using std_tensor_store = OperationBuilder<TensorStoreOp>;
 using std_view = ValueBuilder<ViewOp>;
 using std_zero_extendi = ValueBuilder<ZeroExtendIOp>;

mlir/include/mlir/Dialect/Vector/EDSC/Intrinsics.h

 namespace mlir {
 namespace edsc {
 namespace intrinsics {
 
 using vector_broadcast = ValueBuilder<vector::BroadcastOp>;
 using vector_contract = ValueBuilder<vector::ContractionOp>;
 using vector_matmul = ValueBuilder<vector::MatmulOp>;
 using vector_print = OperationBuilder<vector::PrintOp>;
+using vector_transfer_read = ValueBuilder<vector::TransferReadOp>;
+using vector_transfer_write = OperationBuilder<vector::TransferWriteOp>;
 using vector_type_cast = ValueBuilder<vector::TypeCastOp>;
 
 } // namespace intrinsics
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_VECTOR_EDSC_INTRINSICS_H_

mlir/include/mlir/Dialect/Vector/VectorOps.td

let extraClassDeclaration = [{
     }
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
 }
 
 def Vector_TypeCastOp :
   Vector_Op<"type_cast", [NoSideEffect]>,
     Arguments<(ins StaticShapeMemRefOf<[AnyType]>:$memref)>,
-    Results<(outs AnyMemRef)> {
+    Results<(outs AnyMemRef:$result)> {
   let summary = "type_cast op converts a scalar memref to a vector memref";
   let description = [{
     Performs a conversion from a memref with scalar element to a memref with a
     *single* vector element, copying the shape of the memref to the vector. This
     is the minimal viable operation that is required to makeke
     super-vectorization operational. It can be seen as a special case of the
     `view` operation but scoped in the super-vectorization context.
 
     Syntax:
 
     ```
     operation ::= `vector.type_cast` ssa-use : memref-type to memref-type
     ```
 
     Example:
 
     ```mlir
     %A  = alloc() : memref<5x4x3xf32>
     %VA = vector.type_cast %A : memref<5x4x3xf32> to memref<vector<5x4x3xf32>>
     ```
   }];
 
+  /// Build the canonical memRefType with a single vector.
+  /// E.g. memref<4 x 5 x vector<6 x f32>> -> memref<vector<4 x 5 x 6 x f32>>.
   let builders = [OpBuilder<
     "OpBuilder &builder, OperationState &result, Value source">];
 
-  let parser = [{
-    return impl::parseCastOp(parser, result);
-  }];
-
   let extraClassDeclaration = [{
     MemRefType getMemRefType() {
       return memref().getType().cast<MemRefType>();
     }
     MemRefType getResultMemRefType() {
       return getResult().getType().cast<MemRefType>();
     }
   }];
+
+  let assemblyFormat = [{
+    $memref attr-dict `:` type($memref) `to` type($result)
+  }];
 }
 
 def Vector_ConstantMaskOp :
   Vector_Op<"constant_mask", [NoSideEffect]>,
     Arguments<(ins I64ArrayAttr:$mask_dim_sizes)>,
     Results<(outs VectorOf<[I1]>)> {
   let summary = "creates a constant vector mask";
   let description = [{

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

 
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 class MLIRContext;
 class OwningRewritePatternList;
 
 /// Collect a set of patterns to convert from the Vector dialect to itself.
-/// Should be merged with populateVectorToAffineLoopsConversionPatterns.
+/// Should be merged with populateVectorToLoopsLoweringPattern.
 void populateVectorToVectorConversionPatterns(
     MLIRContext *context, OwningRewritePatternList &patterns,
     ArrayRef<int64_t> coarseVectorShape = {},
     ArrayRef<int64_t> fineVectorShape = {});
 
 ////////////////////////////////////////////////////////////////////////////////
 // The following Declarative Rewrite Rule (DRR) helpers are used in rewrite
 // patterns. As such, they must not call into `rewriter.erase/replace` APIs and

mlir/include/mlir/EDSC/Builders.h

NestedBuilder &operator=(NestedBuilder &&other) {
     std::swap(bodyScope, other.bodyScope);
     return *this;
   }
 
   /// Enter an mlir::Block and setup a ScopedContext to insert operations at
   /// the end of it. Since we cannot use c++ language-level scoping to implement
   /// scoping itself, we use enter/exit pairs of operations.
   /// As a consequence we must allocate a new OpBuilder + ScopedContext and
   /// let the escape.
-  /// Step back "prev" times from the end of the block to set up the insertion
-  /// point, which is useful for non-empty blocks.
-  void enter(mlir::Block *block, int prev = 0) {
+  void enter(mlir::Block *block) {
     bodyScope = new ScopedContext(
         ScopedContext::getBuilder(),
-        OpBuilder::InsertPoint(block, std::prev(block->end(), prev)),
+        OpBuilder::InsertPoint(block, std::prev(block->end())),
         ScopedContext::getLocation());
+    if (!block->empty()) {
+      auto &termOp = block->back();
+      if (termOp.isKnownTerminator())
+        ScopedContext::getBuilder().setInsertionPoint(&termOp);
+    }
     bodyScope->nestedBuilder = this;
   }
 
   /// Exit the current mlir::Block by explicitly deleting the dynamically
   /// allocated OpBuilder and ScopedContext.
   void exit() {
     // Reclaim now to exit the scope.
     bodyScope->nestedBuilder = nullptr;

 /// A BlockBuilder is a NestedBuilder for mlir::Block*.
 /// This exists by opposition to LoopBuilder which is not related to an
 /// mlir::Block* but to a mlir::Value.
 /// It is meant to be used as a temporary object for representing any nested
 /// MLIR construct that is "related to" an mlir::Block*.
 class BlockBuilder : public NestedBuilder {
 public:
   /// Enters the mlir::Block* previously captured by `bh` and sets the insertion
-  /// point to its end.
+  /// point to its end. If the block already contains a terminator, set the
+  /// insertion point before the terminator.
   BlockBuilder(BlockHandle bh, Append);
 
   /// Constructs a new mlir::Block with argument types derived from `args`.
   /// Captures the new block in `bh` and its arguments into `args`.
   /// Enters the new mlir::Block* and sets the insertion point to its end.
   ///
   /// Prerequisites:
   ///   The Value `args` are typed delayed Values; i.e. they are

mlir/include/mlir/IR/AffineMap.h

public:
 
   /// Returns a single constant result affine map.
   static AffineMap getConstantMap(int64_t val, MLIRContext *context);
 
   /// Returns an AffineMap with 'numDims' identity result dim exprs.
   static AffineMap getMultiDimIdentityMap(unsigned numDims,
                                           MLIRContext *context);
 
+  /// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most
+  /// minor dimensions.
+  static AffineMap getMinorIdentityMap(unsigned dims, unsigned results,
+                                       MLIRContext *context);
+
   /// Returns an AffineMap representing a permutation.
   /// The permutation is expressed as a non-empty vector of integers.
   /// E.g. the permutation `(i,j,k) -> (j,k,i)` will be expressed with
   /// `permutation = [1,2,0]`. All values in `permutation` must be
   /// integers, in the range 0..`permutation.size()-1` without duplications
   /// (i.e. `[1,1,2]` is an invalid permutation).
   static AffineMap getPermutationMap(ArrayRef<unsigned> permutation,
                                      MLIRContext *context);
public:
   bool operator==(AffineMap other) const { return other.map == map; }
   bool operator!=(AffineMap other) const { return !(other.map == map); }
 
   /// Returns true if this affine map is an identity affine map.
   /// An identity affine map corresponds to an identity affine function on the
   /// dimensional identifiers.
   bool isIdentity() const;
 
+  /// Returns true if the map is a minor identity map, i.e. an identity affine
+  /// map (d0, ..., dn) -> (dp, ..., dn) on the most minor dimensions.
+  static bool isMinorIdentity(AffineMap map);
+
   /// Returns true if this affine map is an empty map, i.e., () -> ().
   bool isEmpty() const;
 
   /// Returns true if this affine map is a single result constant function.
   bool isSingleConstant() const;
 
   /// Returns the constant result of this map. This methods asserts that the map
   /// has a single constant result.

mlir/lib/Conversion/LinalgToLLVM/CMakeLists.txt

target_link_libraries(MLIRLinalgToLLVM
   MLIRAffineToStandard
   MLIREDSC
   MLIRIR
   MLIRLinalgOps
   MLIRLLVMIR
   MLIRLoopToStandard
   MLIRStandardToLLVM
   MLIRVectorToLLVM
+  MLIRVectorToLoops
   MLIRTransforms
   LLVMCore
   LLVMSupport
   )

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

 //===- LinalgToLLVM.cpp - conversion from Linalg to LLVM dialect ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/LoopToStandard/ConvertLoopToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
+#include "mlir/Conversion/VectorToLoops/ConvertVectorToLoops.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
void ConvertLinalgToLLVMPass::runOnOperation() {
   auto module = getOperation();
 
   // Convert to the LLVM IR dialect using the converter defined above.
   OwningRewritePatternList patterns;
   LLVMTypeConverter converter(&getContext());
   populateAffineToStdConversionPatterns(patterns, &getContext());
   populateLoopToStdConversionPatterns(patterns, &getContext());
   populateStdToLLVMConversionPatterns(converter, patterns);
+  populateVectorToLoopsConversionPatterns(patterns, &getContext());
   populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
   populateVectorToLLVMConversionPatterns(converter, patterns);
   populateLinalgToStandardConversionPatterns(patterns, &getContext());
   populateLinalgToLLVMConversionPatterns(converter, patterns, &getContext());
 
   LLVMConversionTarget target(getContext());
   target.addDynamicallyLegalOp<FuncOp>(
       [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
   target.addLegalOp<ModuleOp, ModuleTerminatorOp>();
   if (failed(applyFullConversion(module, target, patterns, &converter)))
     signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertLinalgToLLVMPass() {
   return std::make_unique<ConvertLinalgToLLVMPass>();
 }

mlir/lib/Conversion/VectorToLoops/ConvertVectorToLoops.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include <type_traits>
 
 #include "mlir/Conversion/VectorToLoops/ConvertVectorToLoops.h"
 #include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
 #include "mlir/Dialect/LoopOps/EDSC/Builders.h"
+#include "mlir/Dialect/LoopOps/EDSC/Intrinsics.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Types.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using vector::TransferReadOp;
 using vector::TransferWriteOp;
 
+/// Helper class captures the common information needed to lower N>1-D vector
+/// transfer operations (read and write).
+/// On construction, this class opens an edsc::ScopedContext for simpler IR
+/// manipulation.
+/// In pseudo-IR, for an n-D vector_transfer_read such as:
+///
+/// ```
+///   vector_transfer_read(%m, %offsets, identity_map, %fill) :
+///     memref<(leading_dims) x (major_dims) x (minor_dims) x type>,
+///     vector<(major_dims) x (minor_dims) x type>
+/// ```
+///
+/// where rank(minor_dims) is the lower-level vector rank (e.g. 1 for LLVM or
+/// higher).
+///
+/// This is the entry point to emitting pseudo-IR resembling:
+///
+/// ```
+///   %tmp = alloc(): memref<(major_dims) x vector<minor_dim x type>>
+///   for (%ivs_major, {0}, {vector_shape}, {1}) { // (N-1)-D loop nest
+///     if (any_of(%ivs_major + %offsets, <, major_dims)) {
+///       %v = vector_transfer_read(
+///         {%offsets_leading, %ivs_major + %offsets_major, %offsets_minor},
+///          %ivs_minor):
+///         memref<(leading_dims) x (major_dims) x (minor_dims) x type>,
+///         vector<(minor_dims) x type>;
+///       store(%v, %tmp);
+///     } else {
+///       %v = splat(vector<(minor_dims) x type>, %fill)
+///       store(%v, %tmp, %ivs_major);
+///     }
+///   }
+///   %res = load(%tmp, %0): memref<(major_dims) x vector<minor_dim x type>>):
+//      vector<(major_dims) x (minor_dims) x type>
+/// ```
+///
+template <typename ConcreteOp>
+class NDTransferOpHelper {
+public:
+  NDTransferOpHelper(PatternRewriter &rewriter, ConcreteOp xferOp)
+      : rewriter(rewriter), loc(xferOp.getLoc()),
+        scope(std::make_unique<ScopedContext>(rewriter, loc)), xferOp(xferOp),
+        op(xferOp.getOperation()) {
+    vectorType = xferOp.getVectorType();
+    // TODO(ntv, ajcbik): when we go to k > 1-D vectors adapt minorRank.
+    minorRank = 1;
+    majorRank = vectorType.getRank() - minorRank;
+    leadingRank = xferOp.getMemRefType().getRank() - (majorRank + minorRank);
+    majorVectorType =
+        VectorType::get(vectorType.getShape().take_front(majorRank),
+                        vectorType.getElementType());
+    minorVectorType =
+        VectorType::get(vectorType.getShape().take_back(minorRank),
+                        vectorType.getElementType());
+    /// Memref of minor vector type is used for individual transfers.
+    memRefMinorVectorType =
+        MemRefType::get(majorVectorType.getShape(), minorVectorType, {},
+                        xferOp.getMemRefType().getMemorySpace());
+  }
+
+  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`.
+  ///   2. Compute a boolean that determines whether the first `majorIvs.rank()`
+  ///      dimensions `majorIvs + majorOffsets` are all within `memrefBounds`.
+  ///   3. Create an IfOp conditioned on the boolean in step 2.
+  ///   4. Call a `thenBlockBuilder` and an `elseBlockBuilder` to append
+  ///      operations to the IfOp blocks as appropriate.
+  template <typename LambdaThen, typename LambdaElse>
+  void emitInBounds(ValueRange majorIvs, ValueRange majorOffsets,
+                    MemRefBoundsCapture &memrefBounds,
+                    LambdaThen thenBlockBuilder, LambdaElse elseBlockBuilder);
+
+  /// Common state to lower vector transfer ops.
+  PatternRewriter &rewriter;
+  Location loc;
+  std::unique_ptr<ScopedContext> scope;
+  ConcreteOp xferOp;
+  Operation *op;
+  // A vector transfer copies data between:
+  //   - memref<(leading_dims) x (major_dims) x (minor_dims) x type>
+  //   - vector<(major_dims) x (minor_dims) x type>
+  unsigned minorRank;         // for now always 1
+  unsigned majorRank;         // vector rank - minorRank
+  unsigned leadingRank;       // memref rank - vector rank
+  VectorType vectorType;      // vector<(major_dims) x (minor_dims) x type>
+  VectorType majorVectorType; // vector<(major_dims) x type>
+  VectorType minorVectorType; // vector<(minor_dims) x type>
+  MemRefType memRefMinorVectorType; // memref<vector<(minor_dims) x type>>
+};
+
+template <typename ConcreteOp>
+template <typename Lambda>
+void NDTransferOpHelper<ConcreteOp>::emitLoops(Lambda loopBodyBuilder) {
+  /// Loop nest operates on the major dimensions
+  MemRefBoundsCapture memrefBoundsCapture(xferOp.memref());
+  VectorBoundsCapture vectorBoundsCapture(majorVectorType);
+  auto majorLbs = vectorBoundsCapture.getLbs();
+  auto majorUbs = vectorBoundsCapture.getUbs();
+  auto majorSteps = vectorBoundsCapture.getSteps();
+  SmallVector<Value, 8> majorIvs(vectorBoundsCapture.rank());
+  AffineLoopNestBuilder(majorIvs, majorLbs, majorUbs, majorSteps)([&] {
+    ValueRange indices(xferOp.indices());
+    loopBodyBuilder(majorIvs, indices.take_front(leadingRank),
+                    indices.drop_front(leadingRank).take_front(majorRank),
+                    indices.take_back(minorRank), memrefBoundsCapture);
+  });
+}
+
+template <typename ConcreteOp>
+template <typename LambdaThen, typename LambdaElse>
+void NDTransferOpHelper<ConcreteOp>::emitInBounds(
+    ValueRange majorIvs, ValueRange majorOffsets,
+    MemRefBoundsCapture &memrefBounds, LambdaThen thenBlockBuilder,
+    LambdaElse elseBlockBuilder) {
+  Value inBounds = std_constant_int(/*value=*/1, /*width=*/1);
+  SmallVector<Value, 4> majorIvsPlusOffsets;
+  majorIvsPlusOffsets.reserve(majorIvs.size());
+  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);
+    Value inBounds2 = majorIvsPlusOffsets.back() < ub;
+    inBounds = inBounds && inBounds2;
+  }
+
+  auto ifOp = ScopedContext::getBuilder().create<loop::IfOp>(
+      ScopedContext::getLocation(), TypeRange{}, inBounds,
+      /*withElseRegion=*/std::is_same<ConcreteOp, TransferReadOp>());
+  BlockBuilder(&ifOp.thenRegion().front(),
+               Append())([&] { thenBlockBuilder(majorIvsPlusOffsets); });
+  if (std::is_same<ConcreteOp, TransferReadOp>())
+    BlockBuilder(&ifOp.elseRegion().front(),
+                 Append())([&] { elseBlockBuilder(majorIvsPlusOffsets); });
+}
+
+template <>
+LogicalResult NDTransferOpHelper<TransferReadOp>::doReplace() {
+  Value alloc = std_alloc(memRefMinorVectorType);
+
+  emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets,
+                ValueRange majorOffsets, ValueRange minorOffsets,
+                MemRefBoundsCapture &memrefBounds) {
+    // If in-bounds, index into memref and lower to 1-D transfer read.
+    auto thenBlockBuilder = [&](ValueRange majorIvsPlusOffsets) {
+      auto map = AffineMap::getMinorIdentityMap(
+          xferOp.getMemRefType().getRank(), minorRank, xferOp.getContext());
+      // Lower to 1-D vector_transfer_read and let recursion handle it.
+      Value memref = xferOp.memref();
+      SmallVector<Value, 8> indexing;
+      indexing.reserve(leadingRank + majorRank + minorRank);
+      indexing.append(leadingOffsets.begin(), leadingOffsets.end());
+      indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end());
+      indexing.append(minorOffsets.begin(), minorOffsets.end());
+      auto loaded1D =
+          vector_transfer_read(minorVectorType, memref, indexing,
+                               AffineMapAttr::get(map), xferOp.padding());
+      // Store the 1-D vector.
+      std_store(loaded1D, alloc, majorIvs);
+    };
+    // If out-of-bounds, just store a splatted vector.
+    auto elseBlockBuilder = [&](ValueRange majorIvsPlusOffsets) {
+      auto vector = std_splat(minorVectorType, xferOp.padding());
+      std_store(vector, alloc, majorIvs);
+    };
+    emitInBounds(majorIvs, majorOffsets, memrefBounds, thenBlockBuilder,
+                 elseBlockBuilder);
+  });
+
+  Value loaded =
+      std_load(vector_type_cast(MemRefType::get({}, vectorType), alloc));
+  rewriter.replaceOp(op, loaded);
+
+  return success();
+}
+
+template <>
+LogicalResult NDTransferOpHelper<TransferWriteOp>::doReplace() {
+  Value alloc = std_alloc(memRefMinorVectorType);
+
+  std_store(xferOp.vector(),
+            vector_type_cast(MemRefType::get({}, vectorType), alloc));
+
+  emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets,
+                ValueRange majorOffsets, ValueRange minorOffsets,
+                MemRefBoundsCapture &memrefBounds) {
+    auto thenBlockBuilder = [&](ValueRange majorIvsPlusOffsets) {
+      // Lower to 1-D vector_transfer_write and let recursion handle it.
+      Value loaded1D = std_load(alloc, majorIvs);
+      auto map = AffineMap::getMinorIdentityMap(
+          xferOp.getMemRefType().getRank(), minorRank, xferOp.getContext());
+      SmallVector<Value, 8> indexing;
+      indexing.reserve(leadingRank + majorRank + minorRank);
+      indexing.append(leadingOffsets.begin(), leadingOffsets.end());
+      indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end());
+      indexing.append(minorOffsets.begin(), minorOffsets.end());
+      vector_transfer_write(loaded1D, xferOp.memref(), indexing,
+                            AffineMapAttr::get(map));
+    };
+    // Don't write anything when out of bounds.
+    auto elseBlockBuilder = [&](ValueRange majorIvsPlusOffsets) {};
+    emitInBounds(majorIvs, majorOffsets, memrefBounds, thenBlockBuilder,
+                 elseBlockBuilder);
+  });
+
+  rewriter.eraseOp(op);
+
+  return success();
+}
+
 /// Analyzes the `transfer` to find an access dimension along the fastest remote
 /// MemRef dimension. If such a dimension with coalescing properties is found,
 /// `pivs` and `vectorBoundsCapture` are swapped so that the invocation of
 /// LoopNestBuilder captures it in the innermost loop.
 template <typename TransferOpTy>
 static int computeCoalescedIndex(TransferOpTy transfer) {
   // rank of the remote memory access, coalescing behavior occurs on the
   // innermost memory dimension.

 
 /// Performs the rewrite.
 template <>
 LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   using namespace mlir::edsc::op;
 
   TransferReadOp transfer = cast<TransferReadOp>(op);
+  if (AffineMap::isMinorIdentity(transfer.permutation_map())) {
+    // If > 1D, emit a bunch of loops around 1-D vector transfers.
+    if (transfer.getVectorType().getRank() > 1)
+      return NDTransferOpHelper<TransferReadOp>(rewriter, transfer).doReplace();
+    // If 1-D this is now handled by the target-specific lowering.
+    if (transfer.getVectorType().getRank() == 1)
+      return failure();
+  }
 
+  // Conservative lowering to scalar load / stores.
   // 1. Setup all the captures.
   ScopedContext scope(rewriter, transfer.getLoc());
   StdIndexedValue remote(transfer.memref());
   MemRefBoundsCapture memRefBoundsCapture(transfer.memref());
   VectorBoundsCapture vectorBoundsCapture(transfer.vector());
   int coalescedIdx = computeCoalescedIndex(transfer);
   // Swap the vectorBoundsCapture which will reorder loop bounds.
   if (coalescedIdx >= 0)

 /// TODO(ntv): implement alternatives to clipping.
 /// TODO(ntv): support non-data-parallel operations.
 template <>
 LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   using namespace edsc::op;
 
   TransferWriteOp transfer = cast<TransferWriteOp>(op);
+  if (AffineMap::isMinorIdentity(transfer.permutation_map())) {
+    // If > 1D, emit a bunch of loops around 1-D vector transfers.
+    if (transfer.getVectorType().getRank() > 1)
+      return NDTransferOpHelper<TransferWriteOp>(rewriter, transfer)
+          .doReplace();
+    // If 1-D this is now handled by the target-specific lowering.
+    if (transfer.getVectorType().getRank() == 1)
+      return failure();
+  }
 
   // 1. Setup all the captures.
   ScopedContext scope(rewriter, transfer.getLoc());
   StdIndexedValue remote(transfer.memref());
   MemRefBoundsCapture memRefBoundsCapture(transfer.memref());
   Value vectorValue(transfer.vector());
   VectorBoundsCapture vectorBoundsCapture(transfer.vector());
   int coalescedIdx = computeCoalescedIndex(transfer);
LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite(
   (std_dealloc(tmp)); // vexing parse...
 
   rewriter.eraseOp(op);
   return success();
 }
 
 } // namespace
 
-void mlir::populateVectorToAffineLoopsConversionPatterns(
-    MLIRContext *context, OwningRewritePatternList &patterns) {
+void mlir::populateVectorToLoopsConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<VectorTransferRewriter<vector::TransferReadOp>,
                   VectorTransferRewriter<vector::TransferWriteOp>>(context);
 }

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

if (auto staticForIv = emitStaticFor(lbs, ubs, step)) {
     *iv =
         Value(b.create<AffineForOp>(ScopedContext::getLocation(), lbs,
                                     b.getMultiDimIdentityMap(lbs.size()), ubs,
                                     b.getMultiDimIdentityMap(ubs.size()), step)
                   .getInductionVar());
   }
 
   auto *body = getForInductionVarOwner(*iv).getBody();
-  result.enter(body, /*prev=*/1);
+  result.enter(body);
   return result;
 }
 
 mlir::edsc::AffineLoopNestBuilder::AffineLoopNestBuilder(Value *iv,
                                                          ArrayRef<Value> lbs,
                                                          ArrayRef<Value> ubs,
                                                          int64_t step) {
   loops.emplace_back(makeAffineLoopBuilder(iv, lbs, ubs, step));

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

mlir::edsc::LoopRangeBuilder::LoopRangeBuilder(Value *iv, Value range) {
   assert(range.getDefiningOp() && "need operations to extract range parts");
   auto rangeOp = cast<RangeOp>(range.getDefiningOp());
   auto lb = rangeOp.min();
   auto ub = rangeOp.max();
   auto step = rangeOp.step();
   ForOp forOp = OperationBuilder<ForOp>(lb, ub, step);
   *iv = forOp.getInductionVar();
   auto *body = forOp.getBody();
-  enter(body, /*prev=*/1);
+  enter(body);
 }
 
 mlir::edsc::LoopRangeBuilder::LoopRangeBuilder(Value *iv,
                                                SubViewOp::Range range) {
   ForOp forOp = OperationBuilder<ForOp>(range.offset, range.size, range.stride);
   *iv = forOp.getInductionVar();
   auto *body = forOp.getBody();
-  enter(body, /*prev=*/1);
+  enter(body);
 }
 
 Value mlir::edsc::LoopRangeBuilder::operator()(std::function<void(void)> fun) {
   if (fun)
     fun();
   exit();
   return Value();
 }

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

LoopBuilder mlir::edsc::makeParallelLoopBuilder(MutableArrayRef<Value> ivs,
                                                 ArrayRef<Value> steps) {
   loop::ParallelOp parallelOp = OperationBuilder<loop::ParallelOp>(
       SmallVector<Value, 4>(lbs.begin(), lbs.end()),
       SmallVector<Value, 4>(ubs.begin(), ubs.end()),
       SmallVector<Value, 4>(steps.begin(), steps.end()));
   for (size_t i = 0, e = ivs.size(); i < e; ++i)
     ivs[i] = parallelOp.getBody()->getArgument(i);
   LoopBuilder result;
-  result.enter(parallelOp.getBody(), /*prev=*/1);
+  result.enter(parallelOp.getBody());
   return result;
 }
 
 mlir::edsc::LoopBuilder
 mlir::edsc::makeLoopBuilder(Value *iv, Value lb, Value ub, Value step,
                             MutableArrayRef<Value> iterArgsHandles,
                             ValueRange iterArgsInitValues) {
   mlir::edsc::LoopBuilder result;
   loop::ForOp forOp =
       OperationBuilder<loop::ForOp>(lb, ub, step, iterArgsInitValues);
   *iv = Value(forOp.getInductionVar());
   auto *body = loop::getForInductionVarOwner(*iv).getBody();
   for (size_t i = 0, e = iterArgsHandles.size(); i < e; ++i) {
     // Skipping the induction variable.
     iterArgsHandles[i] = body->getArgument(i + 1);
   }
   result.setOp(forOp);
-  result.enter(body, /*prev=*/1);
+  result.enter(body);
   return result;
 }

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

for (unsigned idx = 0, n = shape.size(); idx < n; ++idx) {
     if (shape[idx] == -1)
       res.push_back(std_dim(memRef, idx));
     else
       res.push_back(std_constant_index(shape[idx]));
   }
   return res;
 }
 
-mlir::edsc::MemRefBoundsCapture::MemRefBoundsCapture(Value v) : base(v) {
+mlir::edsc::MemRefBoundsCapture::MemRefBoundsCapture(Value v) {
   auto memrefSizeValues = getMemRefSizes(v);
   for (auto s : memrefSizeValues) {
     lbs.push_back(std_constant_index(0));
     ubs.push_back(s);
     steps.push_back(1);
   }
 }
 
-mlir::edsc::VectorBoundsCapture::VectorBoundsCapture(Value v) : base(v) {
-  auto vectorType = v.getType().cast<VectorType>();
-  for (auto s : vectorType.getShape()) {
+mlir::edsc::VectorBoundsCapture::VectorBoundsCapture(VectorType t) {
+  for (auto s : t.getShape()) {
     lbs.push_back(std_constant_index(0));
     ubs.push_back(std_constant_index(s));
     steps.push_back(1);
   }
 }
+
+mlir::edsc::VectorBoundsCapture::VectorBoundsCapture(Value v)
+    : VectorBoundsCapture(v.getType().cast<VectorType>()) {}

mlir/lib/Dialect/Vector/VectorOps.cpp

-//===- VectorOps.cpp - MLIR Super Vectorizer Operations -------------------===//
+//===- VectorOps.cpp - MLIR Vector Dialect Operations ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements convenience types for working with super-vectorization

 #define GET_OP_LIST
 #include "mlir/Dialect/Vector/VectorOps.cpp.inc"
       >();
 }
 
 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *VectorDialect::materializeConstant(OpBuilder &builder,
-                                                 Attribute value, Type type,
-                                                 Location loc) {
   return builder.create<ConstantOp>(loc, type, value);
 }
 
 IntegerType vector::getVectorSubscriptType(Builder &builder) {
   return builder.getIntegerType(64);
 }
 
 ArrayAttr vector::getVectorSubscriptAttr(Builder &builder,
static LogicalResult verify(ShapeCastOp op) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TypeCastOp
 //===----------------------------------------------------------------------===//
 
-static MemRefType inferVectorTypeCastResultType(MemRefType t) {
-  return MemRefType::get({}, VectorType::get(t.getShape(), t.getElementType()));
+static SmallVector<int64_t, 8> extractShape(MemRefType memRefType) {
+  auto vectorType = memRefType.getElementType().dyn_cast<VectorType>();
+  SmallVector<int64_t, 8> res(memRefType.getShape().begin(),
+                              memRefType.getShape().end());
+  if (vectorType) {
+    res.reserve(memRefType.getRank() + vectorType.getRank());
+    for (auto s : vectorType.getShape())
+      res.push_back(s);
+  }
+  return res;
 }
 
+/// Build the canonical memRefType with a single vector.
+/// E.g. memref<4 x 5 x vector<6 x f32>> -> memref<vector<4 x 5 x 6 x f32>>.
 void TypeCastOp::build(OpBuilder &builder, OperationState &result,
                        Value source) {
   result.addOperands(source);
+  MemRefType memRefType = source.getType().cast<MemRefType>();
+  VectorType vectorType =
+      VectorType::get(extractShape(memRefType),
+                      getElementTypeOrSelf(getElementTypeOrSelf(memRefType)));
   result.addTypes(
-      inferVectorTypeCastResultType(source.getType().cast<MemRefType>()));
-}
-
-static void print(OpAsmPrinter &p, TypeCastOp op) {
-  auto type = op.getOperand().getType().cast<MemRefType>();
-  p << op.getOperationName() << ' ' << op.memref() << " : " << type << " to "
-    << inferVectorTypeCastResultType(type);
+      MemRefType::get({}, vectorType, {}, memRefType.getMemorySpace()));
 }
 
 static LogicalResult verify(TypeCastOp op) {
   MemRefType canonicalType = canonicalizeStridedLayout(op.getMemRefType());
   if (!canonicalType.getAffineMaps().empty())
     return op.emitOpError("expects operand to be a memref with no layout");
-
-  auto resultType = inferVectorTypeCastResultType(op.getMemRefType());
-  if (op.getResultMemRefType() != resultType)
-    return op.emitOpError("expects result type to be: ") << resultType;
+  if (!op.getResultMemRefType().getAffineMaps().empty())
+    return op.emitOpError("expects result to be a memref with no layout");
+  if (op.getResultMemRefType().getMemorySpace() !=
+      op.getMemRefType().getMemorySpace())
+    return op.emitOpError("expects result in same memory space");
+
+  auto sourceType = op.getMemRefType();
+  auto resultType = op.getResultMemRefType();
+  if (getElementTypeOrSelf(getElementTypeOrSelf(sourceType)) !=
+      getElementTypeOrSelf(getElementTypeOrSelf(resultType)))
+    return op.emitOpError(
+               "expects result and operand with same underlying scalar type: ")
+           << resultType;
+  if (extractShape(sourceType) != extractShape(resultType))
+    return op.emitOpError(
+               "expects concatenated result and operand shapes to be equal: ")
+           << resultType;
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TupleOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseTupleOp(OpAsmParser &parser, OperationState &result) {

mlir/lib/IR/AffineMap.cpp

 } // end anonymous namespace
 
 /// Returns a single constant result affine map.
 AffineMap AffineMap::getConstantMap(int64_t val, MLIRContext *context) {
   return get(/*dimCount=*/0, /*symbolCount=*/0,
              {getAffineConstantExpr(val, context)});
 }
 
+/// Returns an identity affine map (d0, ..., dn) -> (dp, ..., dn) on the most
+/// minor dimensions.
+AffineMap AffineMap::getMinorIdentityMap(unsigned dims, unsigned results,
+                                         MLIRContext *context) {
+  assert(dims >= results && "Dimension mismatch");
+  auto id = AffineMap::getMultiDimIdentityMap(dims, context);
+  return AffineMap::get(dims, 0, id.getResults().take_back(results), context);
+}
+
+bool AffineMap::isMinorIdentity(AffineMap map) {
+  if (!map)
+    return false;
+  return map == getMinorIdentityMap(map.getNumDims(), map.getNumResults(),
+                                    map.getContext());
+};
+
 /// Returns an AffineMap representing a permutation.
 AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,
                                        MLIRContext *context) {
   assert(!permutation.empty() &&
          "Cannot create permutation map from empty permutation vector");
   SmallVector<AffineExpr, 4> affExprs;
   for (auto index : permutation)
     affExprs.push_back(getAffineDimExpr(index, context));

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

-// RUN: mlir-opt %s -test-convert-vector-to-loops | FileCheck %s
-
-// CHECK: #[[ADD:map[0-9]+]] = affine_map<(d0, d1) -> (d0 + d1)>
-// CHECK: #[[SUB:map[0-9]+]] = affine_map<()[s0] -> (s0 - 1)>
+// RUN: mlir-opt %s -test-convert-vector-to-loops -split-input-file | FileCheck %s
 
 // 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>
affine.for %i1 = 0 to 42 step 4 {
       // CHECK: {{.*}} = select
       // CHECK: %[[FILTERED2:.*]] = select
       // CHECK-NEXT: %{{.*}} = load {{.*}}[%[[FILTERED1]], %[[FILTERED2]]] : memref<7x42xf32>
     }
   }
   return
 }
 
+// -----
+
 // CHECK-LABEL: func @materialize_read_1d_partially_specialized
 func @materialize_read_1d_partially_specialized(%dyn1 : index, %dyn2 : index, %dyn4 : index) {
   %f0 = constant 0.0: f32
   %A = alloc (%dyn1, %dyn2, %dyn4) : memref<7x?x?x42x?xf32>
   affine.for %i0 = 0 to 7 {
     affine.for %i1 = 0 to %dyn1 {
       affine.for %i2 = 0 to %dyn2 {
         affine.for %i3 = 0 to 42 step 2 {
           affine.for %i4 = 0 to %dyn4 {
             %f1 = vector.transfer_read %A[%i0, %i1, %i2, %i3, %i4], %f0 {permutation_map = affine_map<(d0, d1, d2, d3, d4) -> (d3)>} : memref<7x?x?x42x?xf32>, vector<4xf32>
             %i3p1 = affine.apply affine_map<(d0) -> (d0 + 1)> (%i3)
             %f2 = vector.transfer_read %A[%i0, %i1, %i2, %i3p1, %i4], %f0 {permutation_map = affine_map<(d0, d1, d2, d3, d4) -> (d3)>} : memref<7x?x?x42x?xf32>, vector<4xf32>
           }
         }
       }
     }
   }
   // CHECK: %[[tensor:[0-9]+]] = alloc
   // CHECK-NOT: {{.*}} dim %[[tensor]], 0
   // CHECK-NOT: {{.*}} dim %[[tensor]], 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_read(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
 func @materialize_read(%M: index, %N: index, %O: index, %P: index) {
   %f0 = constant 0.0: f32
   // 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
affine.for %i1 = 0 to %N {
           %f = vector.transfer_read %A[%i0, %i1, %i2, %i3], %f0 {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, 0, d0)>} : memref<?x?x?x?xf32>, vector<5x4x3xf32>
         }
       }
     }
   }
   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:  %{{.*}} = 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
affine.for %i1 = 0 to %N step 4 {
         affine.for %i3 = 0 to %P step 5 {
           vector.transfer_write %f1, %A[%i0, %i1, %i2, %i3] {permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d1, d0)>} : vector<5x4x3xf32>, memref<?x?x?x?xf32>
         }
       }
     }
   }
   return
 }
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
+// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
+
+// CHECK-LABEL: transfer_read_progressive(
+//  CHECK-SAME:   %[[A:[a-zA-Z0-9]+]]: memref<?x?xf32>,
+//  CHECK-SAME:   %[[base:[a-zA-Z0-9]+]]: index
+func @transfer_read_progressive(%A : memref<?x?xf32>, %base: index) -> vector<17x15xf32> {
+  // CHECK: %[[cst:.*]] = constant 7.000000e+00 : f32
+  %f7 = constant 7.0: f32
+
+  // CHECK-DAG: %[[cond0:.*]] = constant 1 : i1
+  // CHECK-DAG: %[[splat:.*]] = constant dense<7.000000e+00> : vector<15xf32>
+  // CHECK-DAG: %[[alloc:.*]] = alloc() : memref<17xvector<15xf32>>
+  // CHECK-DAG: %[[dim:.*]] = dim %[[A]], 0 : memref<?x?xf32>
+  // CHECK: affine.for %[[I:.*]] = 0 to 17 {
+  // CHECK:   %[[add:.*]] = affine.apply #[[MAP0]](%[[I]])[%[[base]]]
+  // CHECK:   %[[cmp:.*]] = cmpi "slt", %[[add]], %[[dim]] : index
+  // CHECK:   %[[cond1:.*]] = and %[[cmp]], %[[cond0]] : i1
+  // CHECK:   loop.if %[[cond1]] {
+  // CHECK:     %[[vec_1d:.*]] = vector.transfer_read %[[A]][%[[add]], %[[base]]], %[[cst]]  {permutation_map = #[[MAP1]]} : memref<?x?xf32>, vector<15xf32>
+  // CHECK:     store %[[vec_1d]], %[[alloc]][%[[I]]] : memref<17xvector<15xf32>>
+  // CHECK:   } else {
+  // CHECK:     store %[[splat]], %[[alloc]][%[[I]]] : memref<17xvector<15xf32>>
+  // CHECK:   }
+  // CHECK: %[[vmemref:.*]] = vector.type_cast %[[alloc]] : memref<17xvector<15xf32>> to memref<vector<17x15xf32>>
+  // CHECK: %[[cst:.*]] = load %[[vmemref]][] : memref<vector<17x15xf32>>
+  %f = vector.transfer_read %A[%base, %base], %f7
+      {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
+    memref<?x?xf32>, vector<17x15xf32>
+
+  return %f: vector<17x15xf32>
+}
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
+// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
+
+// CHECK-LABEL: transfer_write_progressive(
+//  CHECK-SAME:   %[[A:[a-zA-Z0-9]+]]: memref<?x?xf32>,
+//  CHECK-SAME:   %[[base:[a-zA-Z0-9]+]]: index,
+//  CHECK-SAME:   %[[vec:[a-zA-Z0-9]+]]: vector<17x15xf32>
+func @transfer_write_progressive(%A : memref<?x?xf32>, %base: index, %vec: vector<17x15xf32>) {
+  // CHECK: %[[cond0:.*]] = constant 1 : i1
+  // CHECK: %[[alloc:.*]] = alloc() : memref<17xvector<15xf32>>
+  // CHECK: %[[vmemref:.*]] = vector.type_cast %[[alloc]] : memref<17xvector<15xf32>> to memref<vector<17x15xf32>>
+  // CHECK: store %[[vec]], %[[vmemref]][] : memref<vector<17x15xf32>>
+  // CHECK: %[[dim:.*]] = dim %[[A]], 0 : memref<?x?xf32>
+  // CHECK: affine.for %[[I:.*]] = 0 to 17 {
+  // CHECK:   %[[add:.*]] = affine.apply #[[MAP0]](%[[I]])[%[[base]]]
+  // CHECK:   %[[cmp:.*]] = cmpi "slt", %[[add]], %[[dim]] : index
+  // CHECK:   %[[cond1:.*]] = and %[[cmp]], %[[cond0]] : i1
+  // CHECK:   loop.if %[[cond1]] {
+  // CHECK:     %[[vec_1d:.*]] = load %0[%[[I]]] : memref<17xvector<15xf32>>
+  // CHECK:     vector.transfer_write %[[vec_1d]], %[[A]][%[[add]], %[[base]]] {permutation_map = #[[MAP1]]} : vector<15xf32>, memref<?x?xf32>
+  // CHECK:   }
+  vector.transfer_write %vec, %A[%base, %base]
+      {permutation_map = affine_map<(d0, d1) -> (d0, d1)>} :
+    vector<17x15xf32>, memref<?x?xf32>
+  return
+}

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

 
 TEST_FUNC(builder_blocks) {
   using namespace edsc::op;
   auto f = makeFunction("builder_blocks");
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   Value c1(std_constant_int(42, 32)), c2(std_constant_int(1234, 32));
+  ReturnOp ret = std_ret();
+
   Value r;
   Value args12[2];
   Value &arg1 = args12[0], &arg2 = args12[1];
   Value args34[2];
   Value &arg3 = args34[0], &arg4 = args34[1];
   BlockHandle b1, b2, functionBlock(&f.front());
   BlockBuilder(&b1, {c1.getType(), c1.getType()}, args12)(
       // b2 has not yet been constructed, need to come back later.
       // This is a byproduct of non-structured control-flow.
   );
   BlockBuilder(&b2, {c1.getType(), c1.getType()}, args34)([&] {
     std_br(b1, {arg3, arg4});
   });
   // The insertion point within the toplevel function is now past b2, we will
   // need to get back the entry block.
   // This is what happens with unstructured control-flow..
   BlockBuilder(b1, Append())([&] {
     r = arg1 + arg2;
     std_br(b2, {arg1, r});
   });
   // Get back to entry block and add a branch into b1
   BlockBuilder(functionBlock, Append())([&] { std_br(b1, {c1, c2}); });
+  ret.erase();
 
   // clang-format off
   // CHECK-LABEL: @builder_blocks
   // CHECK:        %{{.*}} = constant 42 : i32
   // CHECK-NEXT:   %{{.*}} = constant 1234 : i32
   // CHECK-NEXT:   br ^bb1(%{{.*}}, %{{.*}} : i32, i32)
   // CHECK-NEXT: ^bb1(%{{.*}}: i32, %{{.*}}: i32):   // 2 preds: ^bb0, ^bb2
   // CHECK-NEXT:   %{{.*}} = addi %{{.*}}, %{{.*}} : i32
TEST_FUNC(builder_cond_branch) {
   auto f = makeFunction("builder_cond_branch", {},
                         {IntegerType::get(1, &globalContext())});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   Value funcArg(f.getArgument(0));
   Value c32(std_constant_int(32, 32)), c64(std_constant_int(64, 64)),
       c42(std_constant_int(42, 32));
+  ReturnOp ret = std_ret();
+
   Value arg1;
   Value args23[2];
   BlockHandle b1, b2, functionBlock(&f.front());
   BlockBuilder(&b1, c32.getType(), arg1)([&] { std_ret(); });
   BlockBuilder(&b2, {c64.getType(), c32.getType()}, args23)([&] { std_ret(); });
   // Get back to entry block and add a conditional branch
   BlockBuilder(functionBlock, Append())([&] {
     std_cond_br(funcArg, b1, {c32}, b2, {c64, c42});
   });
+  ret.erase();
 
   // clang-format off
   // CHECK-LABEL: @builder_cond_branch
   // CHECK:   %{{.*}} = constant 32 : i32
   // CHECK-NEXT:   %{{.*}} = constant 64 : i64
   // CHECK-NEXT:   %{{.*}} = constant 42 : i32
   // CHECK-NEXT:   cond_br %{{.*}}, ^bb1(%{{.*}} : i32), ^bb2(%{{.*}}, %{{.*}} : i64, i32)
   // CHECK-NEXT: ^bb1(%{{.*}}: i32):   // pred: ^bb0

mlir/test/lib/Transforms/TestVectorToLoopsConversion.cpp

 
 namespace {
 
 struct TestVectorToLoopsPass
     : public PassWrapper<TestVectorToLoopsPass, FunctionPass> {
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     auto *context = &getContext();
-    populateVectorToAffineLoopsConversionPatterns(context, patterns);
+    populateVectorToLoopsConversionPatterns(patterns, context);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
   }
 };
 
 } // end anonymous namespace
 
 namespace mlir {
 void registerTestVectorToLoopsPass() {
   PassRegistration<TestVectorToLoopsPass> pass(
       "test-convert-vector-to-loops",
       "Converts vector transfer ops to loops over scalars and vector casts");
 }
 } // namespace mlir