[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

2819

unsigned

2905

unsigned

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

22

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

It has users outside of this repository.

[flaub]

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

2819

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

2819

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

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

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

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.
+    OpBuilder<"OpBuilder &builder, OperationState &result, "
+              "VectorType resultType, AffineMap map, ValueRange operands">,
+    /// Builds an affine vector load op with an identity map and operands.
+    OpBuilder<"OpBuilder &builder, OperationState &result, "
+              "VectorType resultType, Value memref, ValueRange indices = {}">,
+    /// Builds an affine vector load op with the specified map and its operands.
+    OpBuilder<"OpBuilder &builder, OperationState &result, "
+              "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 = [
+    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 = extraClassDeclarationBase # [{
     VectorType getVectorType() {
       return value().getType().cast<VectorType>();
     }
   }];
 }
 
 #endif // AFFINE_OPS

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

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>();
+  auto 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>();
+  auto 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;

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

 //===- AffineParallelNormalize.cpp - AffineParallelNormalize Pass ---------===//
 //
 // 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 a normalizer for affine parallel loops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Affine/Passes.h"
+#include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/IR/PatternMatch.h"
 
 using namespace mlir;
 
-void normalizeAffineParallel(AffineParallelOp op) {
+void mlir::normalizeAffineParallel(AffineParallelOp op) {
   AffineMap lbMap = op.lowerBoundsMap();
   SmallVector<int64_t, 8> steps = op.getSteps();
   // No need to do any work if the parallel op is already normalized.
   bool isAlreadyNormalized =
       llvm::all_of(llvm::zip(steps, lbMap.getResults()), [](auto tuple) {
         int64_t step = std::get<0>(tuple);
         auto lbExpr =
             std::get<1>(tuple).template dyn_cast<AffineConstantExpr>();