diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td --- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td +++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td @@ -201,6 +201,201 @@ let hasCanonicalizer = 1; } +// Base class for ops with static/dynamic offset, sizes and strides +// attributes/arguments. +class BaseOpWithOffsetSizesAndStrides traits = []> : + Std_Op { + code extraBaseClassDeclaration = [{ + /// Returns the number of dynamic offset operands. + int64_t getNumOffsets() { return llvm::size(offsets()); } + + /// Returns the number of dynamic size operands. + int64_t getNumSizes() { return llvm::size(sizes()); } + + /// Returns the number of dynamic stride operands. + int64_t getNumStrides() { return llvm::size(strides()); } + + /// Returns the dynamic sizes for this subview operation if specified. + operand_range getDynamicSizes() { return sizes(); } + + /// Returns in `staticStrides` the static value of the stride + /// operands. Returns failure() if the static value of the stride + /// operands could not be retrieved. + LogicalResult getStaticStrides(SmallVectorImpl &staticStrides) { + if (!strides().empty()) + return failure(); + staticStrides.reserve(static_strides().size()); + for (auto s : static_strides().getAsValueRange()) + staticStrides.push_back(s.getZExtValue()); + return success(); + } + + /// Return the list of Range (i.e. offset, size, stride). Each + /// Range entry contains either the dynamic value or a ConstantIndexOp + /// constructed with `b` at location `loc`. + SmallVector getOrCreateRanges(OpBuilder &b, Location loc); + + /// Return the offsets as Values. Each Value is either the dynamic + /// value specified in the op or a ConstantIndexOp constructed + /// with `b` at location `loc` + SmallVector getOrCreateOffsets(OpBuilder &b, Location loc) { + unsigned dynamicIdx = 1; + return llvm::to_vector<4>(llvm::map_range( + static_offsets().cast(), [&](Attribute a) -> Value { + int64_t staticOffset = a.cast().getInt(); + if (ShapedType::isDynamicStrideOrOffset(staticOffset)) + return getOperand(dynamicIdx++); + else + return b.create( + loc, b.getIndexType(), b.getIndexAttr(staticOffset)); + })); + } + + /// Return the sizes as Values. Each Value is either the dynamic + /// value specified in the op or a ConstantIndexOp constructed + /// with `b` at location `loc` + SmallVector getOrCreateSizes(OpBuilder &b, Location loc) { + unsigned dynamicIdx = 1 + offsets().size(); + return llvm::to_vector<4>(llvm::map_range( + static_sizes().cast(), [&](Attribute a) -> Value { + int64_t staticSize = a.cast().getInt(); + if (ShapedType::isDynamic(staticSize)) + return getOperand(dynamicIdx++); + else + return b.create( + loc, b.getIndexType(), b.getIndexAttr(staticSize)); + })); + } + + /// Return the strides as Values. Each Value is either the dynamic + /// value specified in the op or a ConstantIndexOp constructed with + /// `b` at location `loc` + SmallVector getOrCreateStrides(OpBuilder &b, Location loc) { + unsigned dynamicIdx = 1 + offsets().size() + sizes().size(); + return llvm::to_vector<4>(llvm::map_range( + static_strides().cast(), [&](Attribute a) -> Value { + int64_t staticStride = a.cast().getInt(); + if (ShapedType::isDynamicStrideOrOffset(staticStride)) + return getOperand(dynamicIdx++); + else + return b.create( + loc, b.getIndexType(), b.getIndexAttr(staticStride)); + })); + } + + /// Return the rank of the source ShapedType. + unsigned getSourceRank() { + return source().getType().cast().getRank(); + } + + /// Return the rank of the result ShapedType. + unsigned getResultRank() { return getType().getRank(); } + + /// Return true if the offset `idx` is a static constant. + bool isDynamicOffset(unsigned idx) { + APInt v = *(static_offsets().getAsValueRange().begin() + idx); + return ShapedType::isDynamicStrideOrOffset(v.getSExtValue()); + } + /// Return true if the size `idx` is a static constant. + bool isDynamicSize(unsigned idx) { + APInt v = *(static_sizes().getAsValueRange().begin() + idx); + return ShapedType::isDynamic(v.getSExtValue()); + } + + /// Return true if the stride `idx` is a static constant. + bool isDynamicStride(unsigned idx) { + APInt v = *(static_strides().getAsValueRange().begin() + idx); + return ShapedType::isDynamicStrideOrOffset(v.getSExtValue()); + } + + /// Assert the offset `idx` is a static constant and return its value. + int64_t getStaticOffset(unsigned idx) { + assert(!isDynamicOffset(idx) && "expected static offset"); + APInt v = *(static_offsets().getAsValueRange().begin() + idx); + return v.getSExtValue(); + } + /// Assert the size `idx` is a static constant and return its value. + int64_t getStaticSize(unsigned idx) { + assert(!isDynamicSize(idx) && "expected static size"); + APInt v = *(static_sizes().getAsValueRange().begin() + idx); + return v.getSExtValue(); + } + /// Assert the stride `idx` is a static constant and return its value. + int64_t getStaticStride(unsigned idx) { + assert(!isDynamicStride(idx) && "expected static stride"); + APInt v = *(static_strides().getAsValueRange().begin() + idx); + return v.getSExtValue(); + } + + unsigned getNumDynamicEntriesUpToIdx(ArrayAttr attr, + llvm::function_ref isDynamic, unsigned idx) { + return std::count_if( + attr.getValue().begin(), attr.getValue().begin() + idx, + [&](Attribute attr) { + return isDynamic(attr.cast().getInt()); + }); + } + /// Assert the offset `idx` is dynamic and return the position of the + /// corresponding operand. + unsigned getIndexOfDynamicOffset(unsigned idx) { + assert(isDynamicOffset(idx) && "expected static offset"); + auto numDynamic = + getNumDynamicEntriesUpToIdx(static_offsets().cast(), + ShapedType::isDynamicStrideOrOffset, idx); + return 1 + numDynamic; + } + /// Assert the size `idx` is dynamic and return the position of the + /// corresponding operand. + unsigned getIndexOfDynamicSize(unsigned idx) { + assert(isDynamicSize(idx) && "expected static size"); + auto numDynamic = getNumDynamicEntriesUpToIdx( + static_sizes().cast(), ShapedType::isDynamic, idx); + return 1 + offsets().size() + numDynamic; + } + /// Assert the stride `idx` is dynamic and return the position of the + /// corresponding operand. + unsigned getIndexOfDynamicStride(unsigned idx) { + assert(isDynamicStride(idx) && "expected static stride"); + auto numDynamic = + getNumDynamicEntriesUpToIdx(static_strides().cast(), + ShapedType::isDynamicStrideOrOffset, idx); + return 1 + offsets().size() + sizes().size() + numDynamic; + } + + /// Assert the offset `idx` is dynamic and return its value. + Value getDynamicOffset(unsigned idx) { + return getOperand(getIndexOfDynamicOffset(idx)); + } + /// Assert the size `idx` is dynamic and return its value. + Value getDynamicSize(unsigned idx) { + return getOperand(getIndexOfDynamicSize(idx)); + } + /// Assert the stride `idx` is dynamic and return its value. + Value getDynamicStride(unsigned idx) { + return getOperand(getIndexOfDynamicStride(idx)); + } + + static StringRef getStaticOffsetsAttrName() { + return "static_offsets"; + } + static StringRef getStaticSizesAttrName() { + return "static_sizes"; + } + static StringRef getStaticStridesAttrName() { + return "static_strides"; + } + static ArrayRef getSpecialAttrNames() { + static SmallVector names{ + getStaticOffsetsAttrName(), + getStaticSizesAttrName(), + getStaticStridesAttrName(), + getOperandSegmentSizeAttr()}; + return names; + } + }]; +} + //===----------------------------------------------------------------------===// // AbsFOp //===----------------------------------------------------------------------===// @@ -2710,212 +2905,6 @@ // SubViewOp //===----------------------------------------------------------------------===// -class BaseOpWithOffsetSizesAndStrides traits = []> : - Std_Op { - let builders = [ - // Build a SubViewOp with mixed static and dynamic entries. - OpBuilder< - "Value source, ArrayRef staticOffsets, " - "ArrayRef staticSizes, ArrayRef staticStrides, " - "ValueRange offsets, ValueRange sizes, ValueRange strides, " - "ArrayRef attrs = {}">, - // Build a SubViewOp with all dynamic entries. - OpBuilder< - "Value source, ValueRange offsets, ValueRange sizes, ValueRange strides, " - "ArrayRef attrs = {}"> - ]; - - code extraBaseClassDeclaration = [{ - /// Returns the number of dynamic offset operands. - int64_t getNumOffsets() { return llvm::size(offsets()); } - - /// Returns the number of dynamic size operands. - int64_t getNumSizes() { return llvm::size(sizes()); } - - /// Returns the number of dynamic stride operands. - int64_t getNumStrides() { return llvm::size(strides()); } - - /// Returns the dynamic sizes for this subview operation if specified. - operand_range getDynamicSizes() { return sizes(); } - - /// Returns in `staticStrides` the static value of the stride - /// operands. Returns failure() if the static value of the stride - /// operands could not be retrieved. - LogicalResult getStaticStrides(SmallVectorImpl &staticStrides) { - if (!strides().empty()) - return failure(); - staticStrides.reserve(static_strides().size()); - for (auto s : static_strides().getAsValueRange()) - staticStrides.push_back(s.getZExtValue()); - return success(); - } - - /// Return the list of Range (i.e. offset, size, stride). Each - /// Range entry contains either the dynamic value or a ConstantIndexOp - /// constructed with `b` at location `loc`. - SmallVector getOrCreateRanges(OpBuilder &b, Location loc); - - /// Return the offsets as Values. Each Value is either the dynamic - /// value specified in the op or a ConstantIndexOp constructed - /// with `b` at location `loc` - SmallVector getOrCreateOffsets(OpBuilder &b, Location loc) { - unsigned dynamicIdx = 1; - return llvm::to_vector<4>(llvm::map_range( - static_offsets().cast(), [&](Attribute a) -> Value { - int64_t staticOffset = a.cast().getInt(); - if (ShapedType::isDynamicStrideOrOffset(staticOffset)) - return getOperand(dynamicIdx++); - else - return b.create( - loc, b.getIndexType(), b.getIndexAttr(staticOffset)); - })); - } - - /// Return the sizes as Values. Each Value is either the dynamic - /// value specified in the op or a ConstantIndexOp constructed - /// with `b` at location `loc` - SmallVector getOrCreateSizes(OpBuilder &b, Location loc) { - unsigned dynamicIdx = 1 + offsets().size(); - return llvm::to_vector<4>(llvm::map_range( - static_sizes().cast(), [&](Attribute a) -> Value { - int64_t staticSize = a.cast().getInt(); - if (ShapedType::isDynamic(staticSize)) - return getOperand(dynamicIdx++); - else - return b.create( - loc, b.getIndexType(), b.getIndexAttr(staticSize)); - })); - } - - /// Return the strides as Values. Each Value is either the dynamic - /// value specified in the op or a ConstantIndexOp constructed with - /// `b` at location `loc` - SmallVector getOrCreateStrides(OpBuilder &b, Location loc) { - unsigned dynamicIdx = 1 + offsets().size() + sizes().size(); - return llvm::to_vector<4>(llvm::map_range( - static_strides().cast(), [&](Attribute a) -> Value { - int64_t staticStride = a.cast().getInt(); - if (ShapedType::isDynamicStrideOrOffset(staticStride)) - return getOperand(dynamicIdx++); - else - return b.create( - loc, b.getIndexType(), b.getIndexAttr(staticStride)); - })); - } - - /// Return the rank of the source ShapedType. - unsigned getSourceRank() { - return source().getType().cast().getRank(); - } - - /// Return the rank of the result ShapedType. - unsigned getResultRank() { return getType().getRank(); } - - /// Return true if the offset `idx` is a static constant. - bool isDynamicOffset(unsigned idx) { - APInt v = *(static_offsets().getAsValueRange().begin() + idx); - return ShapedType::isDynamicStrideOrOffset(v.getSExtValue()); - } - /// Return true if the size `idx` is a static constant. - bool isDynamicSize(unsigned idx) { - APInt v = *(static_sizes().getAsValueRange().begin() + idx); - return ShapedType::isDynamic(v.getSExtValue()); - } - - /// Return true if the stride `idx` is a static constant. - bool isDynamicStride(unsigned idx) { - APInt v = *(static_strides().getAsValueRange().begin() + idx); - return ShapedType::isDynamicStrideOrOffset(v.getSExtValue()); - } - - /// Assert the offset `idx` is a static constant and return its value. - int64_t getStaticOffset(unsigned idx) { - assert(!isDynamicOffset(idx) && "expected static offset"); - APInt v = *(static_offsets().getAsValueRange().begin() + idx); - return v.getSExtValue(); - } - /// Assert the size `idx` is a static constant and return its value. - int64_t getStaticSize(unsigned idx) { - assert(!isDynamicSize(idx) && "expected static size"); - APInt v = *(static_sizes().getAsValueRange().begin() + idx); - return v.getSExtValue(); - } - /// Assert the stride `idx` is a static constant and return its value. - int64_t getStaticStride(unsigned idx) { - assert(!isDynamicStride(idx) && "expected static stride"); - APInt v = *(static_strides().getAsValueRange().begin() + idx); - return v.getSExtValue(); - } - - unsigned getNumDynamicEntriesUpToIdx(ArrayAttr attr, - llvm::function_ref isDynamic, unsigned idx) { - return std::count_if( - attr.getValue().begin(), attr.getValue().begin() + idx, - [&](Attribute attr) { - return isDynamic(attr.cast().getInt()); - }); - } - /// Assert the offset `idx` is dynamic and return the position of the - /// corresponding operand. - unsigned getIndexOfDynamicOffset(unsigned idx) { - assert(isDynamicOffset(idx) && "expected static offset"); - auto numDynamic = - getNumDynamicEntriesUpToIdx(static_offsets().cast(), - ShapedType::isDynamicStrideOrOffset, idx); - return 1 + numDynamic; - } - /// Assert the size `idx` is dynamic and return the position of the - /// corresponding operand. - unsigned getIndexOfDynamicSize(unsigned idx) { - assert(isDynamicSize(idx) && "expected static size"); - auto numDynamic = getNumDynamicEntriesUpToIdx( - static_sizes().cast(), ShapedType::isDynamic, idx); - return 1 + offsets().size() + numDynamic; - } - /// Assert the stride `idx` is dynamic and return the position of the - /// corresponding operand. - unsigned getIndexOfDynamicStride(unsigned idx) { - assert(isDynamicStride(idx) && "expected static stride"); - auto numDynamic = - getNumDynamicEntriesUpToIdx(static_strides().cast(), - ShapedType::isDynamicStrideOrOffset, idx); - return 1 + offsets().size() + sizes().size() + numDynamic; - } - - /// Assert the offset `idx` is dynamic and return its value. - Value getDynamicOffset(unsigned idx) { - return getOperand(getIndexOfDynamicOffset(idx)); - } - /// Assert the size `idx` is dynamic and return its value. - Value getDynamicSize(unsigned idx) { - return getOperand(getIndexOfDynamicSize(idx)); - } - /// Assert the stride `idx` is dynamic and return its value. - Value getDynamicStride(unsigned idx) { - return getOperand(getIndexOfDynamicStride(idx)); - } - - static StringRef getStaticOffsetsAttrName() { - return "static_offsets"; - } - static StringRef getStaticSizesAttrName() { - return "static_sizes"; - } - static StringRef getStaticStridesAttrName() { - return "static_strides"; - } - static ArrayRef getSpecialAttrNames() { - static SmallVector names{ - getStaticOffsetsAttrName(), - getStaticSizesAttrName(), - getStaticStridesAttrName(), - getOperandSegmentSizeAttr()}; - return names; - } - }]; -} - def SubViewOp : BaseOpWithOffsetSizesAndStrides< "subview", [DeclareOpInterfaceMethods] > { let summary = "memref subview operation";