[MLIR][Affine] Add custom builders for AffineVectorLoadOp/AffineVectorStoreOp

[flaub]

Summary

Adding missing custom builders for AffineVectorLoadOp & AffineVectorStoreOp. In practice, it is difficult to correctly construct these ops without these builders (because the AffineMap is not included at construction time).

[bondhugula]

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

2929

unsigned

3015

unsigned

mlir/lib/Dialect/Affine/Transforms/AffineParallelNormalize.cpp

22 ↗

(On Diff #287125)

This doesn't have any users? If yes, how was this being linked earlier without being in the namespace?

[bondhugula]

Can you please separate out the fix to normalizeAffineParallel? It has really nothing to do with the builder additions.

[flaub]

mlir/lib/Dialect/Affine/Transforms/AffineParallelNormalize.cpp

22 ↗

(On Diff #287125)

It has users outside of this repository.

[flaub]

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

2929

I can see that you really don't like auto :)

In general, this code was copy/pasted from existing MLIR code.

[flaub]

In D86380#2232369, @bondhugula wrote:

Can you please separate out the fix to normalizeAffineParallel? It has really nothing to do with the builder additions.

Of course, I can. I was hoping to get this in quickly since downstream users are broken without it.

[bondhugula]

In D86380#2232376, @flaub wrote:

In D86380#2232369, @bondhugula wrote:

Can you please separate out the fix to normalizeAffineParallel? It has really nothing to do with the builder additions.

Of course, I can. I was hoping to get this in quickly since downstream users are broken without it.

It would have been quicker with it in its own patch. :)

[dcaballe]

Thanks Frank! Sorry, I was OOO last week. Quick comments for now, I'll come back later.

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

2929

I think in general the rule is to only use auto when it help readability. I personally *try* to only use it for cast/isa/dyn_cast, for other cases where the type is obvious or redundant and for cases where the type is too large (e.g., iterators), but I sometimes abuse it as well :)

mlir/lib/Dialect/Affine/Transforms/AffineParallelNormalize.cpp

22 ↗

(On Diff #287125)

These changes seem unrelated to changes on affine.vector_load and affine.vector_store. Please, create a independent differential for them.

[dcaballe]

Hey @flaub! What happened to this patch? Do you plan to finish it and land it?

Thanks,
Diego

[flaub]

mlir/lib/Dialect/Affine/Transforms/AffineParallelNormalize.cpp

22 ↗

(On Diff #287125)

Did this get updated in your recent commit @dcaballe? If so, is there anything blocking this diff from landing pretty much as is?

[dcaballe]

mlir/lib/Dialect/Affine/Transforms/AffineParallelNormalize.cpp

22 ↗

(On Diff #287125)

Yeah, I think just a rebase would work. Not sure, though, if @bondhugula has any other comments.

[flaub]

Addressing CR comments

[flaub]

@bondhugula Ready when you are :)

[bondhugula]

It's been exactly two months since I submitted a review!! - this was I think a simple patch. I've forgotten all context! @dcaballe - could you please review this as it's directly on top of your work? @flaub please go with @dcaballe's review.

[dcaballe]

Thanks @bondhugula. Yeah, LGTM! It just needed a rebase. Thanks @flaub!

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

let description = [{
     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 builders = [
+    /// Builds an affine vector load op with the specified map and operands.
+    OpBuilderDAG<(ins "VectorType":$resultType, "AffineMap":$map,
+      "ValueRange":$operands)>,
+    /// Builds an affine vector load op with an identity map and operands.
+    OpBuilderDAG<(ins "VectorType":$resultType, "Value":$memref,
+      CArg<"ValueRange", "{}">:$indices)>,
+    /// Builds an affine vector load op with the specified map and its operands.
+    OpBuilderDAG<(ins "VectorType":$resultType, "Value":$memref,
+      "AffineMap":$map, "ValueRange":$mapOperands)>
+  ];
+
   let extraClassDeclaration = extraClassDeclarationBase # [{
     VectorType getVectorType() {
       return result().getType().cast<VectorType>();
     }
   }];
 }
 
 def AffineVectorStoreOp : AffineStoreOpBase<"vector_store"> {
let description = [{
     (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 skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilderDAG<(ins "Value":$valueToStore, "Value":$memref,
+      "ValueRange":$indices)>,
+    OpBuilderDAG<(ins "Value":$valueToStore, "Value":$memref, "AffineMap":$map,
+      "ValueRange":$mapOperands)>
+  ];
+
   let extraClassDeclaration = extraClassDeclarationBase # [{
     VectorType getVectorType() {
       return value().getType().cast<VectorType>();
     }
   }];
 }
 
 #endif // AFFINE_OPS

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

void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
   auto memrefType = memref.getType().cast<MemRefType>();
   result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
   result.types.push_back(memrefType.getElementType());
 }
 
 void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
                          Value memref, ValueRange indices) {
   auto memrefType = memref.getType().cast<MemRefType>();
-  auto rank = memrefType.getRank();
+  int64_t rank = memrefType.getRank();
   // Create identity map for memrefs with at least one dimension or () -> ()
   // for zero-dimensional memrefs.
   auto map =
       rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
   build(builder, result, memref, map, indices);
 }
 
 static ParseResult parseAffineLoadOp(OpAsmParser &parser,
void AffineStoreOp::build(OpBuilder &builder, OperationState &result,
   result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
 }
 
 // Use identity map.
 void AffineStoreOp::build(OpBuilder &builder, OperationState &result,
                           Value valueToStore, Value memref,
                           ValueRange indices) {
   auto memrefType = memref.getType().cast<MemRefType>();
-  auto rank = memrefType.getRank();
+  int64_t rank = memrefType.getRank();
   // Create identity map for memrefs with at least one dimension or () -> ()
   // for zero-dimensional memrefs.
   auto map =
       rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
   build(builder, result, valueToStore, memref, map, indices);
 }
 
 static ParseResult parseAffineStoreOp(OpAsmParser &parser,
static LogicalResult verify(AffineYieldOp op) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // AffineVectorLoadOp
 //===----------------------------------------------------------------------===//
 
+void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
+                               VectorType resultType, AffineMap map,
+                               ValueRange operands) {
+  assert(operands.size() == 1 + map.getNumInputs() && "inconsistent operands");
+  result.addOperands(operands);
+  if (map)
+    result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
+  result.types.push_back(resultType);
+}
+
+void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
+                               VectorType resultType, Value memref,
+                               AffineMap map, ValueRange mapOperands) {
+  assert(map.getNumInputs() == mapOperands.size() && "inconsistent index info");
+  result.addOperands(memref);
+  result.addOperands(mapOperands);
+  result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
+  result.types.push_back(resultType);
+}
+
+void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
+                               VectorType resultType, Value memref,
+                               ValueRange indices) {
+  auto memrefType = memref.getType().cast<MemRefType>();
+  int64_t rank = memrefType.getRank();
+  // Create identity map for memrefs with at least one dimension or () -> ()
+  // for zero-dimensional memrefs.
+  auto map =
+      rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
+  build(builder, result, resultType, memref, map, indices);
+}
+
 static ParseResult parseAffineVectorLoadOp(OpAsmParser &parser,
                                            OperationState &result) {
   auto &builder = parser.getBuilder();
   auto indexTy = builder.getIndexType();
 
   MemRefType memrefType;
   VectorType resultType;
   OpAsmParser::OperandType memrefInfo;
static LogicalResult verify(AffineVectorLoadOp op) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // AffineVectorStoreOp
 //===----------------------------------------------------------------------===//
 
+void AffineVectorStoreOp::build(OpBuilder &builder, OperationState &result,
+                                Value valueToStore, Value memref, AffineMap map,
+                                ValueRange mapOperands) {
+  assert(map.getNumInputs() == mapOperands.size() && "inconsistent index info");
+  result.addOperands(valueToStore);
+  result.addOperands(memref);
+  result.addOperands(mapOperands);
+  result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
+}
+
+// Use identity map.
+void AffineVectorStoreOp::build(OpBuilder &builder, OperationState &result,
+                                Value valueToStore, Value memref,
+                                ValueRange indices) {
+  auto memrefType = memref.getType().cast<MemRefType>();
+  int64_t rank = memrefType.getRank();
+  // Create identity map for memrefs with at least one dimension or () -> ()
+  // for zero-dimensional memrefs.
+  auto map =
+      rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
+  build(builder, result, valueToStore, memref, map, indices);
+}
+
 static ParseResult parseAffineVectorStoreOp(OpAsmParser &parser,
                                             OperationState &result) {
   auto indexTy = parser.getBuilder().getIndexType();
 
   MemRefType memrefType;
   VectorType resultType;
   OpAsmParser::OperandType storeValueInfo;
   OpAsmParser::OperandType memrefInfo;