[mlir][Affine] Introduce affine.vector_load and affine.vector_store

[dcaballe]

Summary

This patch is the first of a short series that aims to extend Affine dialect
with vector memory semantics. Further context and details can be found here:
https://llvm.discourse.group/t/understanding-the-vector-abstraction-in-mlir/

In particular, this patch adds affine.vector_load and affine.vector_store ops to
the Affine dialect and lowers them to vector.transfer_read and
vector.transfer_write, respectively, in the Vector dialect.

Subsequent patches will add new interfaces for affine load-like and
store-like ops (already implemented locally), port affine fusion to
these interfaces as an example (already implemented locally), and
extend affine analyses to deal with vector semantics (WIP, may need
some help :)). Some refactoring may also be needed.

Thanks!
Diego

[dcaballe]

A few things I found along the road...

[dcaballe]

mlir/include/mlir/Conversion/Passes.td

18 ↗

(On Diff #262985)

Should we rename this to convert-affine-to-loops?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

811

Not sure if it's feasible at all in tablegen but it would be great if we could make "code" variables (extraClassDeclaration, builders, etc.) "additive" in a hierarchy instead of overriding the variable with the value from the subclass.

mlir/lib/Conversion/AffineToVector/AffineToVector.cpp

55 ↗

(On Diff #262985)

I would be great if we could have a builder that sets getMinorIdentityMap as the default value of the permutation map for those cases where no permutation is needed.

[bondhugula]

Thanks for initiating this - this would be a great addition to the affine dialect.

mlir/include/mlir/Conversion/Passes.td

18 ↗

(On Diff #262985)

The reasons this is still called lower-affine is that it's actually converting affine not just to SCF (formerly loop) but SCF + std. Because the affine.apply's as well as the expressions in the affine for bounds, affine if conditionals, and load/store subscripts are lowered to add/mul/div etc. -convert-affine-to-<something> is perhaps more consistent / better.

mlir/lib/Conversion/AffineToVector/AffineToVector.cpp

55 ↗

(On Diff #262985)

+1.

[bondhugula]

mlir/include/mlir/Conversion/AffineToVector/AffineToVector.h

1 ↗

(On Diff #262985)

Instead of another conversion pass for this, we could just do this as part of -lower-affine? Because lower-affine already converts to scf + std, and vector is already at the same "level" of abstraction as std, so, -lower-affine could as well just convert to scf + std + vector.

mlir/test/Dialect/Affine/load-store.mlir

228

Do you want to switch the type order in the type list to keep it consistent with affine.load? Just makes it easier to copy / paste while writing test cases and a bit consistent. The order of the operands need not be the same as in the colon type list, but it's the semantics "store to a memref of type ..., a vector value of type ..." that we'd read.

[dcaballe]

mlir/include/mlir/Conversion/AffineToVector/AffineToVector.h

1 ↗

(On Diff #262985)

Ok! I did it this way for composability reasons since I need the vector part but not necessarily the scf + std. I'll move it to -lower-affine but leave the populateAffineToVectorConversionPatterns around so that I can create my own lowering pass locally.

mlir/lib/Conversion/AffineToVector/AffineToVector.cpp

55 ↗

(On Diff #262985)

Ok, I'll address this separately

mlir/test/Dialect/Affine/load-store.mlir

228

Not sure about this. I tried to keep it consistent with vector.transfer_write so that we don't have to do a different interpretation of a vector store depending on the dialect we are. It could be confusing if we switch the type order since one would initially expect the types to match the order of the operands when the number of operands and types are the same (at least I would :)). The way I read it is more like: "store a vector value of type ... to a memref of type...".

[nicolasvasilache]

Great, thanks @dcaballe !

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

811

In base:

let extraClassDeclarationBase = " ... "

In derived:

let extraClassDeclaration =  extraClassDeclarationBase # "..."

?

mlir/lib/Conversion/AffineToVector/AffineToVector.cpp

55 ↗

(On Diff #262985)

Feel free to also fold it into this revision if you want.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2502

Not for this revision but generally, should printer/parser based on parseAffineMapOfSSAIds have a special syntax in the declarative assembly, e.g. "$affine_indexing" ?

mlir/test/Dialect/Affine/load-store.mlir

228

feel free to change this syntax, parser and printer + the vector_transfer syntax to whichever combination of you prefer:
vector from memref / vector to memref or memref to vector / memref from vector

[dcaballe]

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

811

Something like that should work! However, tablegen complains about extraClassDeclarationBase not being defined. I can file a bugzilla for this. I wouldn't like to go into extending ODS in this patch. I could keep getVectorType in the base class for now and point to bugzilla.

[dcaballe]

Addressing feedback. Thanks!

[dcaballe]

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

811

Also note that part of the base extraClassDeclaration, if not all, will go to the interfaces for affine loads/stores (next patch).

mlir/test/Dialect/Affine/load-store.mlir

228

No strong opinion here but the current combination in vector.transfer_* seems ok to me. In any case, if we want to change it, probably better to do it in a separate patch since we would have to change vector tests, etc.

[mehdi_amini]

Just a nit but the one letter difference between affine.vstore and affine.store does not seem ideal to me from a readability point of view. Also one-letter abbreviation isn't great in general, "v" could be for "volatile" for example.
What about affine.vector_store (or affine.vec_store)?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

803

Typo: affine.load -> affine.vload

817

Do we have another term than "blocking"? It can be ambiguous with respect "synchronous" for example, and it isn't clear to me in general.
Is this "store a vector to a contiguous memory segment"?

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

998

Typo permunation

1069

Typo permunation

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2502

I thought the same! Can you file a bug for this?

mlir/test/Conversion/AffineToStandard/lower-affine-to-vector.mlir

9

Isn't this reading out-of-bound?

[dcaballe]

Thanks, Mehdi!

What about affine.vector_store (or affine.vec_store)?

affine.vector_store sounds good to me. Anybody else?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

817

This doc is a copy-paste mix of AffineStoreOp doc and VectorTransferWriteOp doc, including the "blocking" terminology. I decided to do that for consistent with the vector dialect. I can replace "blocking write" by something along the lines of "store to a vector to a contiguous memory segment" in both dialects if that is clearer. We may have to revisit this doc anyways if/when we introduce strided or gather/scatter accesses.

mlir/test/Conversion/AffineToStandard/lower-affine-to-vector.mlir

9

It might anyways due to arg0 being an arbitrary input but I'll make it less obvious :)

[bondhugula]

mlir/include/mlir/Conversion/AffineToVector/AffineToVector.h

1 ↗

(On Diff #262985)

Yes, that sounds good. You don't need a separate pass. Curious: why do you need just the conversion to vector read/write transfer. What will the output give you that you won't get by looking at affine.vload/vstore?

mlir/test/Dialect/Affine/load-store.mlir

228

IMHO, it's more consistent and readable to always find the memref first and the vector type trailing it. The order of operands is secondary - for example, the load doesn't even have a vector type operand. The type list is all custom in the custom form, and so consistency for all memory ops would be good.

[bondhugula]

Please update the commit summary that there is no new conversion pass.

[ftynse]

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

809

This can be omitted.

811

However, tablegen complains about extraClassDeclarationBase not being defined.

let _redefines_ a variable, <typename> defines a new one. Try code extraClassDeclarationBase = [{...}]; instead.

817

I think the "blocking" terminology in vector transfers indicates that it's not a DMA operation and it will block execution until the result is ready. It doesn't make sense for load/store.

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

1002

This looks complex enough to deserve being in a .cpp file instead.

[dcaballe]

Addressing all the feedback. Thanks!

mlir/include/mlir/Conversion/AffineToVector/AffineToVector.h

1 ↗

(On Diff #262985)

I meant that I may want to compose the affine->vector lowering with the lowering to other dialects or a customized scf+std lowering.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

811

Awesome! It works, thanks!

817

I removed "blocking" and rephrase it a bit, describing the read/write starting position and the shape of the slice being read/written. Please, let me know if you have more comments. We can iterate on this.

mlir/test/Dialect/Affine/load-store.mlir

228

Ok, I changed that for affine.vector_load and affine.vector_store. We can change vector.transfer_* separately.

[dcaballe]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2502

https://bugs.llvm.org/show_bug.cgi?id=45891

[mehdi_amini]

Thanks, LGTM but probably worth having Uday or Alex approval on this.

Also: you did mention in the past that you were interested to look into gather/scatter load/store with some 2D native memory unit in mind, do you still plan to look into this later?

[dcaballe]

Thanks, Mehdi!

Also: you did mention in the past that you were interested to look into gather/scatter load/store with some 2D native memory unit in mind, do you still plan to look into this later?

The 2D native memory unit was just an example. I'm interested in having a more generic vector representation, including strided loads/stores, generic gathers/scatters, and multi-dimensional vector loads/stores in general. However, it depends on priorities and my workload. I think there are a few things that I may need to look into first around the current affine vector ops before going any further. I'll be happy to help, though, if somebody else wants to champion this.

[bondhugula]

This overall looks great. Please also update the commit title for vload/vstore -> vector_load/store. (since arc won't do it by itself as you know.)

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

783

Nit: I assume you also want strides in the future. You could document somewhere that this slice/vector is currently contiguous along the respective dimensions, but will be extended to strided vectors.

826

Nit: "... affine function of ... " or "... affine expression in ..."

827

Nit: ... the start position of the write within ...

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

1002

Was a build method moved into AffineOps.cpp? I don't see anything there.

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

579

Argument comment for 1 (i.e., /*numResults=*/1). I think we should also improve the doc comment for getMinorIdentityMap. It doesn't say what 'results' is, and it should be numResults.

607

Comment update: std.store ->

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2572

We are missing a check to verify that this op's map is indeed a permutation map. Is vector.read/write_transfer also missing it? If yes, please add a TODO for it on AffineVector/Load/Store Op's verifiers. In future, I assume you plan to support strides as well - so that could be marked as a TODO as well.

mlir/test/Conversion/AffineToStandard/lower-affine-to-vector.mlir

9

-memref-bound-check can catch such out-of-bound errors (but for affine.load/store ops).

mlir/test/Dialect/Affine/load-store.mlir

220

Nit: MAP0 -> MAP_ID or IDENTITY for readability.

257

We could add an extra negative test case to test/Dialect/Affine/invalid.mlir where you use a different element type for the vector and for the memref, and check for errors.
affine.vector_load %A[%i] : memref<100xvector<8xf32>>, vector<16xf32>
affine.vector_store %v, %A : memref<100xf32>, vector<8xf64>
affine.vector_load %A[%i] : memref<100xvector<8xf32>>, vector<8xf32>

All of these should fail verification. Very likely that users will try these. :)

[dcaballe]

Addressing feedback.

• Added specific verifiers for affine.vector_load and affine.vector_store.
• Refactored affine.load and affine.store verifiers to avoid code duplication.
• Added invalid tests.

[dcaballe]

This overall looks great. Please also update the commit title for vload/vstore -> vector_load/store. (since arc won't do it by itself as you know.)

Thanks! Yes, I'll clean up the message before committing to remove irrelevant information.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

826

I think 'expression' is referring to the AffineExpr class here (note that this is copy-pasted from AffineStoreOp). "affine expression _in_ loop induction variables and symbols" sounds odd to me but I can be wrong. Is that what you want?

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

1002

It was moved to VectorOps.cpp since this is VectorOps.td

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

579

We can actually remove all of this and use the new build with default permutation and padding.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2572

The current map in affine.vector_load/store is not a permutation map (at least a permutation map in the way is defined in vector.transfer_load/store), right? IIUC, the one in affine.vector_load/store is the same as in affine.load/store, i.e., it's just for indexing (computing the start position of the read/write), not for permuting the slice we are reading/writing from/to memory. I think we would need a separate map for that in affine, if we are interested in modeling that.

Added TODOs for both things to the doc in the td file but, please, let me know if I'm wrong.

mlir/test/Dialect/Affine/load-store.mlir

257

Done. I added specific verifiers since this was not covered by the auto-generated ones. This required some refactoring of the scalar verifiers since I didn't want to have the same code replicated 4 times.

[bondhugula]

Looks good. All your examples/test cases are with 1-d vectors, but the op documentation isn't restricted. You may just want to either specifically add that the vector could be multi-dimensional or include one test case in load-store.mlir (or example in the op doc). I assume such multi-dimensional affine vector load/stores correctly lower to vector read/write transfer as is with what you have.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

826

Never mind - this is too minor. Sounds good as is.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2572

Oh, I see. Looks like I got confused. There is no permutation map here.

[dcaballe]

Enabled multi-dim vector loads/stores (minor fixes)
Added tests and doc examples.

I'll commit it tomorrow if no more comments.
Thanks!

mlir/include/mlir/Conversion/AffineToStandard/AffineToStandard.h

Optional<SmallVector<Value, 8>> expandAffineMap(OpBuilder &builder,
                                                 ValueRange operands);
 
 /// Collect a set of patterns to convert from the Affine dialect to the Standard
 /// dialect, in particular convert structured affine control flow into CFG
 /// branch-based control flow.
 void populateAffineToStdConversionPatterns(OwningRewritePatternList &patterns,
                                            MLIRContext *ctx);
 
+/// Collect a set of patterns to convert vector-related Affine ops to the Vector
+/// dialect.
+void populateAffineToVectorConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx);
+
 /// Emit code that computes the lower bound of the given affine loop using
 /// standard arithmetic operations.
 Value lowerAffineLowerBound(AffineForOp op, OpBuilder &builder);
 
 /// Emit code that computes the upper bound of the given affine loop using
 /// standard arithmetic operations.
 Value lowerAffineUpperBound(AffineForOp op, OpBuilder &builder);
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_AFFINETOSTANDARD_AFFINETOSTANDARD_H

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

OpBuilder getElseBodyBuilder() {
       return OpBuilder(&body, std::prev(body.end()));
     }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
-def AffineLoadOp : Affine_Op<"load", []> {
+class AffineLoadOpBase<string mnemonic, list<OpTrait> traits = []> :
+    Affine_Op<mnemonic, traits> {
+  let arguments = (ins Arg<AnyMemRef, "the reference to load from",
+      [MemRead]>:$memref,
+      Variadic<Index>:$indices);
+
+  code extraClassDeclarationBase = [{
+    /// Returns the operand index of the memref.
+    unsigned getMemRefOperandIndex() { return 0; }
+
+    /// Get memref operand.
+    Value getMemRef() { return getOperand(getMemRefOperandIndex()); }
+    void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
+    MemRefType getMemRefType() {
+      return getMemRef().getType().cast<MemRefType>();
+    }
+
+    /// Get affine map operands.
+    operand_range getMapOperands() { return llvm::drop_begin(getOperands(), 1); }
+
+    /// Returns the affine map used to index the memref for this operation.
+    AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
+    AffineMapAttr getAffineMapAttr() {
+      return getAttr(getMapAttrName()).cast<AffineMapAttr>();
+    }
+
+    /// Returns the AffineMapAttr associated with 'memref'.
+    NamedAttribute getAffineMapAttrForMemRef(Value memref) {
+      assert(memref == getMemRef());
+      return {Identifier::get(getMapAttrName(), getContext()),
+              getAffineMapAttr()};
+    }
+
+    static StringRef getMapAttrName() { return "map"; }
+  }];
+}
+
+def AffineLoadOp : AffineLoadOpBase<"load", []> {
   let summary = "affine load operation";
   let description = [{
     The "affine.load" op reads an element from a memref, where the index
     for each memref dimension is an affine expression of loop induction
     variables and symbols. The output of 'affine.load' is a new value with the
     same type as the elements of the memref. An affine expression of loop IVs
     and symbols must be specified for each dimension of the memref. The keyword
     'symbol' can be used to indicate SSA identifiers which are symbolic.
 
     Example 1:
 
     ```mlir
     %1 = affine.load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>
     ```
 
     Example 2: Uses 'symbol' keyword for symbols '%n' and '%m'.
 
     ```mlir
     %1 = affine.load %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32>
     ```
   }];
 
-  let arguments = (ins Arg<AnyMemRef, "the reference to load from",
-      [MemRead]>:$memref,
-      Variadic<Index>:$indices);
   let results = (outs AnyType:$result);
 
   let builders = [
     /// Builds an affine load op with the specified map and operands.
     OpBuilder<"OpBuilder &builder, OperationState &result, AffineMap map, "
                       "ValueRange operands">,
     /// Builds an affine load op with an identity map and operands.
     OpBuilder<"OpBuilder &builder, OperationState &result, Value memref, "
                       "ValueRange indices = {}">,
     /// Builds an affine load op with the specified map and its operands.
     OpBuilder<"OpBuilder &builder, OperationState &result, Value memref, "
                       "AffineMap map, ValueRange mapOperands">
   ];
 
-  let extraClassDeclaration = [{
-    /// Returns the operand index of the memref.
-    unsigned getMemRefOperandIndex() { return 0; }
-
-    /// Get memref operand.
-    Value getMemRef() { return getOperand(getMemRefOperandIndex()); }
-    void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
-    MemRefType getMemRefType() {
-      return getMemRef().getType().cast<MemRefType>();
-    }
-
-    /// Get affine map operands.
-    operand_range getMapOperands() { return llvm::drop_begin(getOperands(), 1); }
-
-    /// Returns the affine map used to index the memref for this operation.
-    AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
-    AffineMapAttr getAffineMapAttr() {
-      return getAttr(getMapAttrName()).cast<AffineMapAttr>();
-    }
-
-    /// Returns the AffineMapAttr associated with 'memref'.
-    NamedAttribute getAffineMapAttrForMemRef(Value memref) {
-      assert(memref == getMemRef());
-      return {Identifier::get(getMapAttrName(), getContext()),
-              getAffineMapAttr()};
-    }
-
-    static StringRef getMapAttrName() { return "map"; }
-  }];
+  let extraClassDeclaration = extraClassDeclarationBase;
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 class AffineMinMaxOpBase<string mnemonic, list<OpTrait> traits = []> :
     Op<Affine_Dialect, mnemonic, traits> {
   let arguments = (ins AffineMapAttr:$map, Variadic<Index>:$operands);
let extraClassDeclaration = [{
     static StringRef getIsWriteAttrName() { return "isWrite"; }
     static StringRef getIsDataCacheAttrName() { return "isDataCache"; }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
-def AffineStoreOp : Affine_Op<"store", []> {
+class AffineStoreOpBase<string mnemonic, list<OpTrait> traits = []> :
+    Affine_Op<mnemonic, traits> {
+
+  code extraClassDeclarationBase = [{
+    /// Get value to be stored by store operation.
+    Value getValueToStore() { return getOperand(0); }
+
+    /// Returns the operand index of the memref.
+    unsigned getMemRefOperandIndex() { return 1; }
+
+    /// Get memref operand.
+    Value getMemRef() { return getOperand(getMemRefOperandIndex()); }
+    void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
+
+    MemRefType getMemRefType() {
+      return getMemRef().getType().cast<MemRefType>();
+    }
+
+    /// Get affine map operands.
+    operand_range getMapOperands() { return llvm::drop_begin(getOperands(), 2); }
+
+    /// Returns the affine map used to index the memref for this operation.
+    AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
+    AffineMapAttr getAffineMapAttr() {
+      return getAttr(getMapAttrName()).cast<AffineMapAttr>();
+    }
+
+    /// Returns the AffineMapAttr associated with 'memref'.
+    NamedAttribute getAffineMapAttrForMemRef(Value memref) {
+      assert(memref == getMemRef());
+      return {Identifier::get(getMapAttrName(), getContext()),
+              getAffineMapAttr()};
+    }
+
+    static StringRef getMapAttrName() { return "map"; }
+  }];
+}
+
+def AffineStoreOp : AffineStoreOpBase<"store", []> {
   let summary = "affine store operation";
   let description = [{
     The "affine.store" op writes an element to a memref, where the index
     for each memref dimension is an affine expression of loop induction
     variables and symbols. The 'affine.store' op stores a new value which is the
     same type as the elements of the memref. An affine expression of loop IVs
     and symbols must be specified for each dimension of the memref. The keyword
     'symbol' can be used to indicate SSA identifiers which are symbolic.
let description = [{
     affine.store %v0, %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32>
     ```
   }];
   let arguments = (ins AnyType:$value,
       Arg<AnyMemRef, "the reference to store to",
       [MemWrite]>:$memref,
       Variadic<Index>:$indices);
 
-
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
                       "Value valueToStore, Value memref, ValueRange indices">,
     OpBuilder<"OpBuilder &builder, OperationState &result, "
                       "Value valueToStore, Value memref, AffineMap map, "
                       "ValueRange mapOperands">
   ];
 
-  let extraClassDeclaration = [{
-    /// Get value to be stored by store operation.
-    Value getValueToStore() { return getOperand(0); }
-
-    /// Returns the operand index of the memref.
-    unsigned getMemRefOperandIndex() { return 1; }
-
-    /// Get memref operand.
-    Value getMemRef() { return getOperand(getMemRefOperandIndex()); }
-    void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
-
-    MemRefType getMemRefType() {
-      return getMemRef().getType().cast<MemRefType>();
-    }
-
-    /// Get affine map operands.
-    operand_range getMapOperands() { return llvm::drop_begin(getOperands(), 2); }
-
-    /// Returns the affine map used to index the memref for this operation.
-    AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
-    AffineMapAttr getAffineMapAttr() {
-      return getAttr(getMapAttrName()).cast<AffineMapAttr>();
-    }
-
-    /// Returns the AffineMapAttr associated with 'memref'.
-    NamedAttribute getAffineMapAttrForMemRef(Value memref) {
-      assert(memref == getMemRef());
-      return {Identifier::get(getMapAttrName(), getContext()),
-              getAffineMapAttr()};
-    }
-
-    static StringRef getMapAttrName() { return "map"; }
-  }];
+  let extraClassDeclaration = extraClassDeclarationBase;
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 def AffineTerminatorOp :
     Affine_Op<"terminator", [NoSideEffect, Terminator]> {
   let summary = "affine terminator operation";
def AffineTerminatorOp :
   // No custom parsing/printing form.
   let parser = ?;
   let printer = ?;
 
   // Fully specified by traits.
   let verifier = ?;
 }
 
+def AffineVectorLoadOp : AffineLoadOpBase<"vector_load", []> {
+  let summary = "affine vector load operation";
+  let description = [{
+    The "affine.vector_load" is the vector counterpart of
+    [affine.load](#affineload-operation). It reads a slice from a
+    [MemRef](../LangRef.md#memref-type), supplied as its first operand,
+    into a [vector](../LangRef.md#vector-type) of the same base elemental type.
+    The index for each memref dimension is an affine expression of loop induction
+    variables and symbols. These indices determine the start position of the read
+    within the memref. The shape of the return vector type determines the shape of
+    the slice read from the memref. This slice is contiguous along the respective
+    dimensions of the shape. Strided vector loads will be supported in the future.
+    An affine expression of loop IVs and symbols must be specified for each
+    dimension of the memref. The keyword 'symbol' can be used to indicate SSA
+    identifiers which are symbolic.
+
+    Example 1: 8-wide f32 vector load.
+
+    ```mlir
+    %1 = affine.vector_load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<8xf32>
+    ```
+
+    Example 2: 4-wide f32 vector load. Uses 'symbol' keyword for symbols '%n' and '%m'.
+
+    ```mlir
+    %1 = affine.vector_load %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32>, vector<4xf32>
+    ```
+
+    Example 3: 2-dim f32 vector load.
+
+    ```mlir
+    %1 = affine.vector_load %0[%i0, %i1] : memref<100x100xf32>, vector<2x8xf32>
+    ```
+
+    TODOs:
+    * Add support for strided vector loads.
+    * Consider adding a permutation map to permute the slice that is read from memory
+    (see [vector.transfer_read](../Vector/#vectortransfer_read-vectortransferreadop)).
+  }];
+
+  let results = (outs AnyVector:$result);
+
+  let extraClassDeclaration = extraClassDeclarationBase # [{
+    VectorType getVectorType() {
+      return result().getType().cast<VectorType>();
+    }
+  }];
+}
+
+def AffineVectorStoreOp : AffineStoreOpBase<"vector_store", []> {
+  let summary = "affine vector store operation";
+  let description = [{
+    The "affine.vector_store" is the vector counterpart of
+    [affine.store](#affinestore-affinestoreop). It writes a
+    [vector](../LangRef.md#vector-type), supplied as its first operand,
+    into a slice within a [MemRef](../LangRef.md#memref-type) of the same base
+    elemental type, supplied as its second operand.
+    The index for each memref dimension is an affine expression of loop
+    induction variables and symbols. These indices determine the start position
+    of the write within the memref. The shape of th input vector determines the
+    shape of the slice written to the memref. This slice is contiguous along the
+    respective dimensions of the shape. Strided vector stores will be supported
+    in the future.
+    An affine expression of loop IVs and symbols must be specified for each
+    dimension of the memref. The keyword 'symbol' can be used to indicate SSA
+    identifiers which are symbolic.
+
+    Example 1: 8-wide f32 vector store.
+
+    ```mlir
+    affine.vector_store %v0, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<8xf32>
+    ```
+
+    Example 2: 4-wide f32 vector store. Uses 'symbol' keyword for symbols '%n' and '%m'.
+
+    ```mlir
+    affine.vector_store %v0, %0[%i0 + symbol(%n), %i1 + symbol(%m)] : memref<100x100xf32>, vector<4xf32>
+    ```
+
+    Example 3: 2-dim f32 vector store.
+
+    ```mlir
+    affine.vector_store %v0, %0[%i0, %i1] : memref<100x100xf32>, vector<2x8xf32>
+    ```
+
+    TODOs:
+    * Add support for strided vector stores.
+    * Consider adding a permutation map to permute the slice that is written to memory
+    (see [vector.transfer_write](../Vector/#vectortransfer_write-vectortransferwriteop)).
+  }];
+
+  let arguments = (ins AnyVector:$value,
+      Arg<AnyMemRef, "the reference to store to",
+      [MemWrite]>:$memref,
+      Variadic<Index>:$indices);
+
+  let extraClassDeclaration = extraClassDeclarationBase # [{
+    VectorType getVectorType() {
+      return value().getType().cast<VectorType>();
+    }
+  }];
+}
+
 #endif // AFFINE_OPS

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

let description = [{
 
     // Read from a memref with vector element type.
     %4 = vector.transfer_read %arg1[%c3, %c3], %vf0
       {permutation_map = (d0, d1)->(d0, d1)}
         : memref<?x?xvector<4x3xf32>>, vector<1x1x4x3xf32>
     ```
   }];
 
+  let builders = [
+    // Builder that sets permutation map and padding to 'getMinorIdentityMap'
+    // and zero, respectively, by default.
+    OpBuilder<"OpBuilder &builder, OperationState &result, VectorType vector, "
+              "Value memref, ValueRange indices">
+  ];
+
   let extraClassDeclaration = [{
     MemRefType getMemRefType() {
       return memref().getType().cast<MemRefType>();
     }
     VectorType getVectorType() {
       return vector().getType().cast<VectorType>();
     }
   }];
let description = [{
 
     // write to a memref with vector element type.
     vector.transfer_write %4, %arg1[%c3, %c3]
       {permutation_map = (d0, d1)->(d0, d1)}
         : vector<1x1x4x3xf32>, memref<?x?xvector<4x3xf32>>
     ```
   }];
 
+  let builders = [
+    // Builder that sets permutation map and padding to 'getMinorIdentityMap'
+    // by default.
+    OpBuilder<"OpBuilder &builder, OperationState &result, Value vector, "
+              "Value memref, ValueRange indices">
+  ];
+
   let extraClassDeclaration = [{
     VectorType getVectorType() {
       return vector().getType().cast<VectorType>();
     }
     MemRefType getMemRefType() {
       return memref().getType().cast<MemRefType>();
     }
   }];

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 
 #include "../PassDetail.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
+using namespace mlir::vector;
 
 namespace {
 /// Visit affine expressions recursively and build the sequence of operations
 /// that correspond to it.  Visitation functions return an Value of the
 /// expression subtree they visited or `nullptr` on error.
 class AffineApplyExpander
     : public AffineExprVisitor<AffineApplyExpander, Value> {
 public:
LogicalResult matchAndRewrite(AffineDmaWaitOp op,
 
     // Build std.dma_wait operation with affine map results.
     rewriter.replaceOpWithNewOp<DmaWaitOp>(
         op, op.getTagMemRef(), *maybeExpandedTagMap, op.getNumElements());
     return success();
   }
 };
 
+/// Apply the affine map from an 'affine.vector_load' operation to its operands,
+/// and feed the results to a newly created 'vector.transfer_read' operation
+/// (which replaces the original 'affine.vector_load').
+class AffineVectorLoadLowering : public OpRewritePattern<AffineVectorLoadOp> {
+public:
+  using OpRewritePattern<AffineVectorLoadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AffineVectorLoadOp op,
+                                PatternRewriter &rewriter) const override {
+    // Expand affine map from 'affineVectorLoadOp'.
+    SmallVector<Value, 8> indices(op.getMapOperands());
+    auto resultOperands =
+        expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);
+    if (!resultOperands)
+      return failure();
+
+    // Build vector.transfer_read memref[expandedMap.results].
+    rewriter.replaceOpWithNewOp<TransferReadOp>(
+        op, op.getVectorType(), op.getMemRef(), *resultOperands);
+    return success();
+  }
+};
+
+/// Apply the affine map from an 'affine.vector_store' operation to its
+/// operands, and feed the results to a newly created 'vector.transfer_write'
+/// operation (which replaces the original 'affine.vector_store').
+class AffineVectorStoreLowering : public OpRewritePattern<AffineVectorStoreOp> {
+public:
+  using OpRewritePattern<AffineVectorStoreOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AffineVectorStoreOp op,
+                                PatternRewriter &rewriter) const override {
+    // Expand affine map from 'affineVectorStoreOp'.
+    SmallVector<Value, 8> indices(op.getMapOperands());
+    auto maybeExpandedMap =
+        expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);
+    if (!maybeExpandedMap)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<TransferWriteOp>(
+        op, op.getValueToStore(), op.getMemRef(), *maybeExpandedMap);
+    return success();
+  }
+};
+
 } // end namespace
 
 void mlir::populateAffineToStdConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   // clang-format off
   patterns.insert<
       AffineApplyLowering,
       AffineDmaStartLowering,
       AffineDmaWaitLowering,
       AffineLoadLowering,
       AffineMinLowering,
       AffineMaxLowering,
       AffinePrefetchLowering,
       AffineStoreLowering,
       AffineForLowering,
       AffineIfLowering,
       AffineTerminatorLowering>(ctx);
   // clang-format on
 }
 
+void mlir::populateAffineToVectorConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx) {
+  // clang-format off
+  patterns.insert<
+      AffineVectorLoadLowering,
+      AffineVectorStoreLowering>(ctx);
+  // clang-format on
+}
+
 namespace {
 class LowerAffinePass : public ConvertAffineToStandardBase<LowerAffinePass> {
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     populateAffineToStdConversionPatterns(patterns, &getContext());
+    populateAffineToVectorConversionPatterns(patterns, &getContext());
     ConversionTarget target(getContext());
-    target.addLegalDialect<scf::SCFDialect, StandardOpsDialect>();
+    target
+        .addLegalDialect<scf::SCFDialect, StandardOpsDialect, VectorDialect>();
     if (failed(applyPartialConversion(getFunction(), target, patterns)))
       signalPassFailure();
   }
 };
 } // namespace
 
 /// Lowers If and For operations within a function into their lower level CFG
 /// equivalent blocks.
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLowerAffinePass() {
   return std::make_unique<LowerAffinePass>();
 }

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

void print(OpAsmPrinter &p, AffineLoadOp op) {
   if (AffineMapAttr mapAttr =
           op.getAttrOfType<AffineMapAttr>(op.getMapAttrName()))
     p.printAffineMapOfSSAIds(mapAttr, op.getMapOperands());
   p << ']';
   p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{op.getMapAttrName()});
   p << " : " << op.getMemRefType();
 }
 
-LogicalResult verify(AffineLoadOp op) {
-  if (op.getType() != op.getMemRefType().getElementType())
-    return op.emitOpError("result type must match element type of memref");
-
-  auto mapAttr = op.getAttrOfType<AffineMapAttr>(op.getMapAttrName());
+/// Verify common indexing invariants of affine.load, affine.store,
+/// affine.vector_load and affine.vector_store.
+static LogicalResult
+verifyMemoryOpIndexing(Operation *op, AffineMapAttr mapAttr,
+                       Operation::operand_range mapOperands,
+                       MemRefType memrefType, unsigned numIndexOperands) {
   if (mapAttr) {
-    AffineMap map =
-        op.getAttrOfType<AffineMapAttr>(op.getMapAttrName()).getValue();
-    if (map.getNumResults() != op.getMemRefType().getRank())
-      return op.emitOpError("affine.load affine map num results must equal"
-                            " memref rank");
-    if (map.getNumInputs() != op.getNumOperands() - 1)
-      return op.emitOpError("expects as many subscripts as affine map inputs");
+    AffineMap map = mapAttr.getValue();
+    if (map.getNumResults() != memrefType.getRank())
+      return op->emitOpError("affine map num results must equal memref rank");
+    if (map.getNumInputs() != numIndexOperands)
+      return op->emitOpError("expects as many subscripts as affine map inputs");
   } else {
-    if (op.getMemRefType().getRank() != op.getNumOperands() - 1)
-      return op.emitOpError(
+    if (memrefType.getRank() != numIndexOperands)
+      return op->emitOpError(
           "expects the number of subscripts to be equal to memref rank");
   }
 
   Region *scope = getAffineScope(op);
-  for (auto idx : op.getMapOperands()) {
+  for (auto idx : mapOperands) {
     if (!idx.getType().isIndex())
-      return op.emitOpError("index to load must have 'index' type");
+      return op->emitOpError("index to load must have 'index' type");
     if (!isValidAffineIndexOperand(idx, scope))
-      return op.emitOpError("index must be a dimension or symbol identifier");
+      return op->emitOpError("index must be a dimension or symbol identifier");
   }
+
+  return success();
+}
+
+LogicalResult verify(AffineLoadOp op) {
+  auto memrefType = op.getMemRefType();
+  if (op.getType() != memrefType.getElementType())
+    return op.emitOpError("result type must match element type of memref");
+
+  if (failed(verifyMemoryOpIndexing(
+          op.getOperation(),
+          op.getAttrOfType<AffineMapAttr>(op.getMapAttrName()),
+          op.getMapOperands(), memrefType,
+          /*numIndexOperands=*/op.getNumOperands() - 1)))
+    return failure();
+
   return success();
 }
 
 void AffineLoadOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<SimplifyAffineOp<AffineLoadOp>>(context);
 }
 
if (AffineMapAttr mapAttr =
     p.printAffineMapOfSSAIds(mapAttr, op.getMapOperands());
   p << ']';
   p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{op.getMapAttrName()});
   p << " : " << op.getMemRefType();
 }
 
 LogicalResult verify(AffineStoreOp op) {
   // First operand must have same type as memref element type.
-  if (op.getValueToStore().getType() != op.getMemRefType().getElementType())
+  auto memrefType = op.getMemRefType();
+  if (op.getValueToStore().getType() != memrefType.getElementType())
     return op.emitOpError(
         "first operand must have same type memref element type");
 
-  auto mapAttr = op.getAttrOfType<AffineMapAttr>(op.getMapAttrName());
-  if (mapAttr) {
-    AffineMap map = mapAttr.getValue();
-    if (map.getNumResults() != op.getMemRefType().getRank())
-      return op.emitOpError("affine.store affine map num results must equal"
-                            " memref rank");
-    if (map.getNumInputs() != op.getNumOperands() - 2)
-      return op.emitOpError("expects as many subscripts as affine map inputs");
-  } else {
-    if (op.getMemRefType().getRank() != op.getNumOperands() - 2)
-      return op.emitOpError(
-          "expects the number of subscripts to be equal to memref rank");
-  }
+  if (failed(verifyMemoryOpIndexing(
+          op.getOperation(),
+          op.getAttrOfType<AffineMapAttr>(op.getMapAttrName()),
+          op.getMapOperands(), memrefType,
+          /*numIndexOperands=*/op.getNumOperands() - 2)))
+    return failure();
 
-  Region *scope = getAffineScope(op);
-  for (auto idx : op.getMapOperands()) {
-    if (!idx.getType().isIndex())
-      return op.emitOpError("index to store must have 'index' type");
-    if (!isValidAffineIndexOperand(idx, scope))
-      return op.emitOpError("index must be a dimension or symbol identifier");
-  }
   return success();
 }
 
 void AffineStoreOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<SimplifyAffineOp<AffineStoreOp>>(context);
 }
 
if (parser.parseRegion(*body, ivs, types) ||
     return failure();
 
   // Add a terminator if none was parsed.
   AffineParallelOp::ensureTerminator(*body, builder, result.location);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// AffineVectorLoadOp
+//===----------------------------------------------------------------------===//
+
+ParseResult parseAffineVectorLoadOp(OpAsmParser &parser,
+                                    OperationState &result) {
+  auto &builder = parser.getBuilder();
+  auto indexTy = builder.getIndexType();
+
+  MemRefType memrefType;
+  VectorType resultType;
+  OpAsmParser::OperandType memrefInfo;
+  AffineMapAttr mapAttr;
+  SmallVector<OpAsmParser::OperandType, 1> mapOperands;
+  return failure(
+      parser.parseOperand(memrefInfo) ||
+      parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
+                                    AffineVectorLoadOp::getMapAttrName(),
+                                    result.attributes) ||
+      parser.parseOptionalAttrDict(result.attributes) ||
+      parser.parseColonType(memrefType) || parser.parseComma() ||
+      parser.parseType(resultType) ||
+      parser.resolveOperand(memrefInfo, memrefType, result.operands) ||
+      parser.resolveOperands(mapOperands, indexTy, result.operands) ||
+      parser.addTypeToList(resultType, result.types));
+}
+
+void print(OpAsmPrinter &p, AffineVectorLoadOp op) {
+  p << "affine.vector_load " << op.getMemRef() << '[';
+  if (AffineMapAttr mapAttr =
+          op.getAttrOfType<AffineMapAttr>(op.getMapAttrName()))
+    p.printAffineMapOfSSAIds(mapAttr, op.getMapOperands());
+  p << ']';
+  p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{op.getMapAttrName()});
+  p << " : " << op.getMemRefType() << ", " << op.getType();
+}
+
+/// Verify common invariants of affine.vector_load and affine.vector_store.
+static LogicalResult verifyVectorMemoryOp(Operation *op, MemRefType memrefType,
+                                          VectorType vectorType) {
+  // Check that memref and vector element types match.
+  if (memrefType.getElementType() != vectorType.getElementType())
+    return op->emitOpError(
+        "requires memref and vector types of the same elemental type");
+
+  return success();
+}
+
+static LogicalResult verify(AffineVectorLoadOp op) {
+  MemRefType memrefType = op.getMemRefType();
+  if (failed(verifyMemoryOpIndexing(
+          op.getOperation(),
+          op.getAttrOfType<AffineMapAttr>(op.getMapAttrName()),
+          op.getMapOperands(), memrefType,
+          /*numIndexOperands=*/op.getNumOperands() - 1)))
+    return failure();
+
+  if (failed(verifyVectorMemoryOp(op.getOperation(), memrefType,
+                                  op.getVectorType())))
+    return failure();
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// AffineVectorStoreOp
+//===----------------------------------------------------------------------===//
+
+ParseResult parseAffineVectorStoreOp(OpAsmParser &parser,
+                                     OperationState &result) {
+  auto indexTy = parser.getBuilder().getIndexType();
+
+  MemRefType memrefType;
+  VectorType resultType;
+  OpAsmParser::OperandType storeValueInfo;
+  OpAsmParser::OperandType memrefInfo;
+  AffineMapAttr mapAttr;
+  SmallVector<OpAsmParser::OperandType, 1> mapOperands;
+  return failure(
+      parser.parseOperand(storeValueInfo) || parser.parseComma() ||
+      parser.parseOperand(memrefInfo) ||
+      parser.parseAffineMapOfSSAIds(mapOperands, mapAttr,
+                                    AffineVectorStoreOp::getMapAttrName(),
+                                    result.attributes) ||
+      parser.parseOptionalAttrDict(result.attributes) ||
+      parser.parseColonType(memrefType) || parser.parseComma() ||
+      parser.parseType(resultType) ||
+      parser.resolveOperand(storeValueInfo, resultType, result.operands) ||
+      parser.resolveOperand(memrefInfo, memrefType, result.operands) ||
+      parser.resolveOperands(mapOperands, indexTy, result.operands));
+}
+
+void print(OpAsmPrinter &p, AffineVectorStoreOp op) {
+  p << "affine.vector_store " << op.getValueToStore();
+  p << ", " << op.getMemRef() << '[';
+  if (AffineMapAttr mapAttr =
+          op.getAttrOfType<AffineMapAttr>(op.getMapAttrName()))
+    p.printAffineMapOfSSAIds(mapAttr, op.getMapOperands());
+  p << ']';
+  p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{op.getMapAttrName()});
+  p << " : " << op.getMemRefType() << ", " << op.getValueToStore().getType();
+}
+
+static LogicalResult verify(AffineVectorStoreOp op) {
+  MemRefType memrefType = op.getMemRefType();
+  if (failed(verifyMemoryOpIndexing(
+          op.getOperation(),
+          op.getAttrOfType<AffineMapAttr>(op.getMapAttrName()),
+          op.getMapOperands(), memrefType,
+          /*numIndexOperands=*/op.getNumOperands() - 2)))
+    return failure();
+
+  if (failed(verifyVectorMemoryOp(op.getOperation(), memrefType,
+                                  op.getVectorType())))
+    return failure();
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Affine/IR/AffineOps.cpp.inc"

mlir/lib/Dialect/Vector/VectorOps.cpp

static LogicalResult verifyTransferOp(Operation *op, MemRefType memrefType,
   if (permutationMap.getNumSymbols() != 0)
     return op->emitOpError("requires permutation_map without symbols");
   if (permutationMap.getNumInputs() != memrefType.getRank())
     return op->emitOpError("requires a permutation_map with input dims of the "
                            "same rank as the memref type");
   return success();
 }
 
+/// Builder that sets permutation map and padding to 'getMinorIdentityMap' and
+/// zero, respectively, by default.
+void TransferReadOp::build(OpBuilder &builder, OperationState &result,
+                           VectorType vector, Value memref,
+                           ValueRange indices) {
+  auto permMap = AffineMap::getMinorIdentityMap(
+      memref.getType().cast<MemRefType>().getRank(), vector.getRank(),
+      builder.getContext());
+  Type elemType = vector.cast<VectorType>().getElementType();
+  Value padding = builder.create<ConstantOp>(result.location, elemType,
+                                             builder.getZeroAttr(elemType));
+
+  build(builder, result, vector, memref, indices, permMap, padding);
+}
+
 static void print(OpAsmPrinter &p, TransferReadOp op) {
   p << op.getOperationName() << " " << op.memref() << "[" << op.indices()
     << "], " << op.padding() << " ";
   p.printOptionalAttrDict(op.getAttrs());
   p << " : " << op.getMemRefType() << ", " << op.getVectorType();
 }
 
 static ParseResult parseTransferReadOp(OpAsmParser &parser,
static LogicalResult verify(TransferReadOp op) {
   return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
 //===----------------------------------------------------------------------===//
 // TransferWriteOp
 //===----------------------------------------------------------------------===//
 
+/// Builder that sets permutation map and padding to 'getMinorIdentityMap' by
+/// default.
+void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
+                            Value vector, Value memref, ValueRange indices) {
+  auto vectorType = vector.getType().cast<VectorType>();
+  auto permMap = AffineMap::getMinorIdentityMap(
+      memref.getType().cast<MemRefType>().getRank(), vectorType.getRank(),
+      builder.getContext());
+  build(builder, result, vector, memref, indices, permMap);
+}
+
 static LogicalResult verify(TransferWriteOp op) {
   // Consistency of elemental types in memref and vector.
   MemRefType memrefType = op.getMemRefType();
   VectorType vectorType = op.getVectorType();
   auto permutationMap = op.permutation_map();
 
   if (llvm::size(op.indices()) != memrefType.getRank())
     return op.emitOpError("requires ") << memrefType.getRank() << " indices";

mlir/test/Conversion/AffineToStandard/lower-affine-to-vector.mlir

+// RUN: mlir-opt -lower-affine --split-input-file %s | FileCheck %s
+
+// CHECK: #[[perm_map:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: func @affine_vector_load
+func @affine_vector_load(%arg0 : index) {
+  %0 = alloc() : memref<100xf32>
+  affine.for %i0 = 0 to 16 {
+    %1 = affine.vector_load %0[%i0 + symbol(%arg0) + 7] : memref<100xf32>, vector<8xf32>
+  }
+// CHECK:       %[[buf:.*]] = alloc
+// CHECK:       %[[a:.*]] = addi %{{.*}}, %{{.*}} : index
+// CHECK-NEXT:  %[[c7:.*]] = constant 7 : index
+// CHECK-NEXT:  %[[b:.*]] = addi %[[a]], %[[c7]] : index
+// CHECK-NEXT:  %[[pad:.*]] = constant 0.0
+// CHECK-NEXT:  vector.transfer_read %[[buf]][%[[b]]], %[[pad]] {permutation_map = #[[perm_map]]} : memref<100xf32>, vector<8xf32>
+  return
+}
+
+// -----
+
+// CHECK: #[[perm_map:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: func @affine_vector_store
+func @affine_vector_store(%arg0 : index) {
+  %0 = alloc() : memref<100xf32>
+  %1 = constant dense<11.0> : vector<4xf32>
+  affine.for %i0 = 0 to 16 {
+    affine.vector_store %1, %0[%i0 - symbol(%arg0) + 7] : memref<100xf32>, vector<4xf32>
+}
+// CHECK:       %[[buf:.*]] = alloc
+// CHECK:       %[[val:.*]] = constant dense
+// CHECK:       %[[c_1:.*]] = constant -1 : index
+// CHECK-NEXT:  %[[a:.*]] = muli %arg0, %[[c_1]] : index
+// CHECK-NEXT:  %[[b:.*]] = addi %{{.*}}, %[[a]] : index
+// CHECK-NEXT:  %[[c7:.*]] = constant 7 : index
+// CHECK-NEXT:  %[[c:.*]] = addi %[[b]], %[[c7]] : index
+// CHECK-NEXT:  vector.transfer_write  %[[val]], %[[buf]][%[[c]]] {permutation_map = #[[perm_map]]} : vector<4xf32>, memref<100xf32>
+  return
+}
+
+// -----
+
+// CHECK: #[[perm_map:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: func @affine_vector_load
+func @affine_vector_load(%arg0 : index) {
+  %0 = alloc() : memref<100xf32>
+  affine.for %i0 = 0 to 16 {
+    %1 = affine.vector_load %0[%i0 + symbol(%arg0) + 7] : memref<100xf32>, vector<8xf32>
+  }
+// CHECK:       %[[buf:.*]] = alloc
+// CHECK:       %[[a:.*]] = addi %{{.*}}, %{{.*}} : index
+// CHECK-NEXT:  %[[c7:.*]] = constant 7 : index
+// CHECK-NEXT:  %[[b:.*]] = addi %[[a]], %[[c7]] : index
+// CHECK-NEXT:  %[[pad:.*]] = constant 0.0
+// CHECK-NEXT:  vector.transfer_read %[[buf]][%[[b]]], %[[pad]] {permutation_map = #[[perm_map]]} : memref<100xf32>, vector<8xf32>
+  return
+}
+
+// -----
+
+// CHECK: #[[perm_map:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: func @affine_vector_store
+func @affine_vector_store(%arg0 : index) {
+  %0 = alloc() : memref<100xf32>
+  %1 = constant dense<11.0> : vector<4xf32>
+  affine.for %i0 = 0 to 16 {
+    affine.vector_store %1, %0[%i0 - symbol(%arg0) + 7] : memref<100xf32>, vector<4xf32>
+}
+// CHECK:       %[[buf:.*]] = alloc
+// CHECK:       %[[val:.*]] = constant dense
+// CHECK:       %[[c_1:.*]] = constant -1 : index
+// CHECK-NEXT:  %[[a:.*]] = muli %arg0, %[[c_1]] : index
+// CHECK-NEXT:  %[[b:.*]] = addi %{{.*}}, %[[a]] : index
+// CHECK-NEXT:  %[[c7:.*]] = constant 7 : index
+// CHECK-NEXT:  %[[c:.*]] = addi %[[b]], %[[c7]] : index
+// CHECK-NEXT:  vector.transfer_write  %[[val]], %[[buf]][%[[c]]] {permutation_map = #[[perm_map]]} : vector<4xf32>, memref<100xf32>
+  return
+}
+
+// -----
+
+// CHECK: #[[perm_map:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @vector_load_2d
+func @vector_load_2d() {
+  %0 = alloc() : memref<100x100xf32>
+  affine.for %i0 = 0 to 16 step 2{
+    affine.for %i1 = 0 to 16 step 8 {
+      %1 = affine.vector_load %0[%i0, %i1] : memref<100x100xf32>, vector<2x8xf32>
+// CHECK:      %[[buf:.*]] = alloc
+// CHECK:      scf.for %[[i0:.*]] =
+// CHECK:        scf.for %[[i1:.*]] =
+// CHECK-NEXT:     %[[pad:.*]] = constant 0.0
+// CHECK-NEXT:     vector.transfer_read %[[buf]][%[[i0]], %[[i1]]], %[[pad]] {permutation_map = #[[perm_map]]} : memref<100x100xf32>, vector<2x8xf32>
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK: #[[perm_map:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @vector_store_2d
+func @vector_store_2d() {
+  %0 = alloc() : memref<100x100xf32>
+  %1 = constant dense<11.0> : vector<2x8xf32>
+  affine.for %i0 = 0 to 16 step 2{
+    affine.for %i1 = 0 to 16 step 8 {
+      affine.vector_store %1, %0[%i0, %i1] : memref<100x100xf32>, vector<2x8xf32>
+// CHECK:      %[[buf:.*]] = alloc
+// CHECK:      %[[val:.*]] = constant dense
+// CHECK:      scf.for %[[i0:.*]] =
+// CHECK:        scf.for %[[i1:.*]] =
+// CHECK-NEXT:     vector.transfer_write  %[[val]], %[[buf]][%[[i0]], %[[i1]]] {permutation_map = #[[perm_map]]} : vector<2x8xf32>, memref<100x100xf32>
+    }
+  }
+  return
+}
+

mlir/test/Dialect/Affine/invalid.mlir

func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) {
   affine.for %x = 0 to 7 {
     %y = addi %x, %x : index
     // expected-error@+1 {{operand cannot be used as a symbol}}
     affine.parallel (%i, %j) = (0, 0) to (symbol(%y), 100) step (10, 10) {
     }
   }
   return
 }
+
+// -----
+
+func @vector_load_invalid_vector_type() {
+  %0 = alloc() : memref<100xf32>
+  affine.for %i0 = 0 to 16 step 8 {
+    // expected-error@+1 {{requires memref and vector types of the same elemental type}}
+    %1 = affine.vector_load %0[%i0] : memref<100xf32>, vector<8xf64>
+  }
+  return
+}
+
+// -----
+
+func @vector_store_invalid_vector_type() {
+  %0 = alloc() : memref<100xf32>
+  %1 = constant dense<7.0> : vector<8xf64>
+  affine.for %i0 = 0 to 16 step 8 {
+    // expected-error@+1 {{requires memref and vector types of the same elemental type}}
+    affine.vector_store %1, %0[%i0] : memref<100xf32>, vector<8xf64>
+  }
+  return
+}
+
+// -----
+
+func @vector_load_vector_memref() {
+  %0 = alloc() : memref<100xvector<8xf32>>
+  affine.for %i0 = 0 to 4 {
+    // expected-error@+1 {{requires memref and vector types of the same elemental type}}
+    %1 = affine.vector_load %0[%i0] : memref<100xvector<8xf32>>, vector<8xf32>
+  }
+  return
+}
+
+// -----
+
+func @vector_store_vector_memref() {
+  %0 = alloc() : memref<100xvector<8xf32>>
+  %1 = constant dense<7.0> : vector<8xf32>
+  affine.for %i0 = 0 to 4 {
+    // expected-error@+1 {{requires memref and vector types of the same elemental type}}
+    affine.vector_store %1, %0[%i0] : memref<100xvector<8xf32>>, vector<8xf32>
+  }
+  return
+}

mlir/test/Dialect/Affine/load-store.mlir

affine.for %i1 = 0 to 10 {
       affine.prefetch %0[%i0 + 3, %i1 + 11], write, locality<0>, data : memref<100x100xf32>
       // CHECK: affine.prefetch %{{.*}}[%{{.*}} + 3, %{{.*}} + 11], write, locality<0>, data : memref<100x100xf32>
       affine.prefetch %0[%i0, %i1 + 1], read, locality<3>, instr : memref<100x100xf32>
       // CHECK: affine.prefetch %{{.*}}[%{{.*}}, %{{.*}} + 1], read, locality<3>, instr : memref<100x100xf32>
     }
   }
   return
 }
+
+// -----
+
+// CHECK: [[MAP_ID:#map[0-9]+]] = affine_map<(d0, d1) -> (d0, d1)>
+
+// Test with just loop IVs.
+func @vector_load_vector_store_iv() {
+  %0 = alloc() : memref<100x100xf32>
+  affine.for %i0 = 0 to 16 {
+    affine.for %i1 = 0 to 16 step 8 {
+      %1 = affine.vector_load %0[%i0, %i1] : memref<100x100xf32>, vector<8xf32>
+      affine.vector_store %1, %0[%i0, %i1] : memref<100x100xf32>, vector<8xf32>
+// CHECK:      %[[buf:.*]] = alloc
+// CHECK-NEXT: affine.for %[[i0:.*]] = 0
+// CHECK-NEXT:   affine.for %[[i1:.*]] = 0
+// CHECK-NEXT:     %[[val:.*]] = affine.vector_load %[[buf]][%[[i0]], %[[i1]]] : memref<100x100xf32>, vector<8xf32>
+// CHECK-NEXT:     affine.vector_store %[[val]], %[[buf]][%[[i0]], %[[i1]]] : memref<100x100xf32>, vector<8xf32>
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK: [[MAP0:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 + 3, d1 + 7)>
+
+// Test with loop IVs and constants.
+func @vector_load_vector_store_iv_constant() {
+  %0 = alloc() : memref<100x100xf32>
+  affine.for %i0 = 0 to 10 {
+    affine.for %i1 = 0 to 16 step 4 {
+      %1 = affine.vector_load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<4xf32>
+      affine.vector_store %1, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<4xf32>
+// CHECK:      %[[buf:.*]] = alloc
+// CHECK-NEXT: affine.for %[[i0:.*]] = 0
+// CHECK-NEXT:   affine.for %[[i1:.*]] = 0
+// CHECK-NEXT:     %[[val:.*]] = affine.vector_load %{{.*}}[%{{.*}} + 3, %{{.*}} + 7] : memref<100x100xf32>, vector<4xf32>
+// CHECK-NEXT:     affine.vector_store %[[val]], %[[buf]][%[[i0]] + 3, %[[i1]] + 7] : memref<100x100xf32>, vector<4xf32>
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK: [[MAP0:#map[0-9]+]] = affine_map<(d0, d1) -> (d0, d1)>
+
+func @vector_load_vector_store_2d() {
+  %0 = alloc() : memref<100x100xf32>
+  affine.for %i0 = 0 to 16 step 2{
+    affine.for %i1 = 0 to 16 step 8 {
+      %1 = affine.vector_load %0[%i0, %i1] : memref<100x100xf32>, vector<2x8xf32>
+      affine.vector_store %1, %0[%i0, %i1] : memref<100x100xf32>, vector<2x8xf32>
+// CHECK:      %[[buf:.*]] = alloc
+// CHECK-NEXT: affine.for %[[i0:.*]] = 0
+// CHECK-NEXT:   affine.for %[[i1:.*]] = 0
+// CHECK-NEXT:     %[[val:.*]] = affine.vector_load %[[buf]][%[[i0]], %[[i1]]] : memref<100x100xf32>, vector<2x8xf32>
+// CHECK-NEXT:     affine.vector_store %[[val]], %[[buf]][%[[i0]], %[[i1]]] : memref<100x100xf32>, vector<2x8xf32>
+    }
+  }
+  return
+}