[llvm-branch-commits] [mlir] 9ffba19 - [MLIR][Affine] Add custom builders for AffineVectorLoadOp/AffineVectorStoreOp

[Frank Laub via llvm-branch-commits]

Author: Frank Laub
Date: 2020-11-25T20:22:56Z
New Revision: 9ffba19e86ce234d7289b31b4a7b78cd8878f159

URL: https://github.com/llvm/llvm-project/commit/9ffba19e86ce234d7289b31b4a7b78cd8878f159
DIFF: https://github.com/llvm/llvm-project/commit/9ffba19e86ce234d7289b31b4a7b78cd8878f159.diff

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

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).

Differential Revision: https://reviews.llvm.org/D86380

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index e2fffc4ab591..5b1180889072 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -910,6 +910,18 @@ def AffineVectorLoadOp : AffineLoadOpBase<"vector_load"> {
 
   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>();
@@ -964,6 +976,14 @@ def AffineVectorStoreOp : AffineStoreOpBase<"vector_store"> {
       [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>();

diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 299569602ea1..b83807475f34 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2088,7 +2088,7 @@ void AffineLoadOp::build(OpBuilder &builder, OperationState &result,
 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 =
@@ -2202,7 +2202,7 @@ 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 =
@@ -2902,6 +2902,38 @@ static LogicalResult verify(AffineYieldOp op) {
 // 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();
@@ -2965,6 +2997,29 @@ static LogicalResult verify(AffineVectorLoadOp op) {
 // 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();