diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td @@ -910,6 +910,18 @@ 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(); @@ -964,6 +976,14 @@ [MemWrite]>:$memref, Variadic:$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(); diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp --- 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, Value memref, ValueRange indices) { auto memrefType = memref.getType().cast(); - 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 @@ Value valueToStore, Value memref, ValueRange indices) { auto memrefType = memref.getType().cast(); - 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 @@ // 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(); + 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 @@ // 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(); + 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();