-
Notifications
You must be signed in to change notification settings - Fork 11.8k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[mlir] Change tensor.extract/insert
to take static/dynamic indices.
#104488
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -332,12 +332,37 @@ def Tensor_ExtractOp : Tensor_Op<"extract", [ | |
```mlir | ||
%4 = tensor.extract %t[%1, %2] : tensor<4x4xi32> | ||
%5 = tensor.extract %rt[%1, %2] : tensor<?x?xi32> | ||
%6 = tensor.extract %rt[3, 4] : tensor<?x?xi32> | ||
%7 = tensor.extract %rt[%1, 4] : tensor<?x?xi32> | ||
``` | ||
}]; | ||
|
||
let arguments = (ins AnyRankedTensor:$tensor, Variadic<Index>:$indices); | ||
let arguments = (ins | ||
AnyRankedTensor:$tensor, | ||
Variadic<Index>:$indices, | ||
DenseI64ArrayAttr:$static_indices | ||
); | ||
let results = (outs AnyType:$result); | ||
let assemblyFormat = "$tensor `[` $indices `]` attr-dict `:` type($tensor)"; | ||
let assemblyFormat = [{ | ||
$tensor `` | ||
custom<DynamicIndexList>($indices, $static_indices) | ||
attr-dict `:` type($tensor) | ||
}]; | ||
|
||
let builders = [ | ||
// Build an ExtractOp with mixed static and dynamic indexes. | ||
OpBuilder<(ins "Value":$tensor, "ArrayRef<OpFoldResult>":$indexes, | ||
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>, | ||
// Build an ExtractOp with mixed static, dynamic indexes and inferred result type. | ||
OpBuilder<(ins "Type":$resultType, "Value":$tensor, "ArrayRef<OpFoldResult>":$indexes, | ||
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>, | ||
// Build an ExtractOp with dynamic indexes. | ||
OpBuilder<(ins "Value":$source, CArg<"ValueRange", "{}">:$indexes, | ||
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>, | ||
// Build an ExtractOp with dynamic indexes and inferred result type. | ||
OpBuilder<(ins "Type":$resultType, "Value":$source, CArg<"ValueRange", "{}">:$indexes, | ||
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>, | ||
]; | ||
|
||
let hasCanonicalizer = 1; | ||
let hasFolder = 1; | ||
|
@@ -808,16 +833,35 @@ def Tensor_InsertOp : Tensor_Op<"insert", [ | |
|
||
let arguments = (ins AnyType:$scalar, | ||
AnyRankedTensor:$dest, | ||
Variadic<Index>:$indices); | ||
Variadic<Index>:$indices, | ||
DenseI64ArrayAttr:$static_indices | ||
); | ||
let results = (outs AnyRankedTensor:$result); | ||
let assemblyFormat = [{ | ||
$scalar `into` $dest `[` $indices `]` attr-dict `:` type($dest) | ||
$scalar `into` | ||
$dest `` custom<DynamicIndexList>($indices, $static_indices) | ||
attr-dict `:` type($dest) | ||
}]; | ||
|
||
let extraClassDeclaration = [{ | ||
MutableOperandRange getDpsInitsMutable() { return getDestMutable(); } | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Both ops should have a There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks for the suggestion! Do you think a new interface like MixedIndicesInterface is needed for this? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. That could be useful. (But can also be done without.) I tried something like that in the past (https://reviews.llvm.org/D156899), but I didn't land it for some reason... Don't really remember why. There was also an RFC (https://discourse.llvm.org/t/rfc-more-opfoldresult-and-mixed-indices-in-ops-that-deal-with-shaped-values/72510). I would read through that discussion before adding a new interface. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. +1 I think there is something more general to be done here, I think this is a good starting point and can see other pain points and think about generalizing. |
||
}]; | ||
|
||
let builders = [ | ||
// Build an InsertOp with mixed static and dynamic indexes. | ||
OpBuilder<(ins "Value":$scalar, "Value":$dest, "ArrayRef<OpFoldResult>":$indexes, | ||
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>, | ||
// Build an InsertOp with mixed static, dynamic indexes and inferred result type. | ||
OpBuilder<(ins "Type":$resultType, "Value":$scalar, "Value":$dest, "ArrayRef<OpFoldResult>":$indexes, | ||
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>, | ||
// Build an InsertOp with dynamic indexes. | ||
OpBuilder<(ins "Value":$scalar, "Value":$dest, CArg<"ValueRange", "{}">:$indexes, | ||
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>, | ||
// Build an InsertOp with dynamic indexes and inferred result type. | ||
OpBuilder<(ins "Type":$resultType, "Value":$scalar, "Value":$dest, CArg<"ValueRange", "{}">:$indexes, | ||
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>, | ||
]; | ||
|
||
let hasFolder = 1; | ||
let hasVerifier = 1; | ||
} | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -1736,6 +1736,32 @@ struct ShapeOfFromReshape : public OpRewritePattern<shape::ShapeOfOp> { | |
} | ||
}; | ||
|
||
struct ExtractFromShapeOfExtentTensor | ||
: public OpRewritePattern<tensor::ExtractOp> { | ||
using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern; | ||
|
||
LogicalResult matchAndRewrite(tensor::ExtractOp op, | ||
PatternRewriter &rewriter) const override { | ||
auto tensorShapeOfOp = op.getTensor().getDefiningOp<shape::ShapeOfOp>(); | ||
if (!tensorShapeOfOp) | ||
return rewriter.notifyMatchFailure(op, "producer is not shape.shape_of"); | ||
|
||
int64_t staticIndice = op.getStaticIndices()[0]; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Its sort of weird that static index could be dynamic ... I seem to recall poking about this on a previous review, why not just store only static in one and only dynamic in the other and then using the type to differentiate - that would result in more operations for indexing. Not something to address here as this is keeping the form. |
||
Type indexType = rewriter.getIndexType(); | ||
Value indice = | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. typo: There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Rather index ? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. (funnily enough it seems indice is index in Spanish) |
||
staticIndice != ShapedType::kDynamic | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Prefer using isDynamic(staticIndex) |
||
? tensorShapeOfOp->getDialect() | ||
->materializeConstant( | ||
rewriter, IntegerAttr::get(indexType, staticIndice), | ||
indexType, op.getLoc()) | ||
->getResult(0) | ||
: op.getIndices()[0]; | ||
rewriter.replaceOpWithNewOp<tensor::DimOp>(op, tensorShapeOfOp.getArg(), | ||
indice); | ||
return success(); | ||
} | ||
}; | ||
|
||
// Canonicalize | ||
// ``` | ||
// %0 = shape.shape_of %arg : tensor<?x?x?xf32> -> tensor<3xindex> | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -39,6 +39,19 @@ using llvm::divideCeilSigned; | |
using llvm::divideFloorSigned; | ||
using llvm::mod; | ||
|
||
static LogicalResult | ||
checkTensorRankMatchIndices(Value tensor, ValueRange dynamicIndices, | ||
ArrayRef<int64_t> staticIndices) { | ||
auto tensorType = llvm::cast<RankedTensorType>(tensor.getType()); | ||
int64_t dynamicDimCount = llvm::count_if(staticIndices, [](int64_t element) { | ||
return element == ShapedType::kDynamic; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Same |
||
}); | ||
if (tensorType.getRank() != staticIndices.size() || | ||
dynamicDimCount != static_cast<int64_t>(dynamicIndices.size())) | ||
return LogicalResult::failure(); | ||
return LogicalResult::success(); | ||
} | ||
|
||
/// Materialize a single constant operation from a given attribute value with | ||
/// the desired resultant type. | ||
Operation *TensorDialect::materializeConstant(OpBuilder &builder, | ||
|
@@ -1120,10 +1133,49 @@ void ExtractOp::getAsmResultNames( | |
setNameFn(getResult(), "extracted"); | ||
} | ||
|
||
// Build an ExtractOp with mixed static and dynamic indexes. | ||
void ExtractOp::build(OpBuilder &b, OperationState &result, Value tensor, | ||
ArrayRef<OpFoldResult> indices, | ||
ArrayRef<NamedAttribute> attrs) { | ||
Type resultType = llvm::cast<TensorType>(tensor.getType()).getElementType(); | ||
build(b, result, resultType, tensor, indices, attrs); | ||
} | ||
|
||
// Build an ExtractOp with mixed static, dynamic indexes and inferred result | ||
// Type. | ||
void ExtractOp::build(OpBuilder &b, OperationState &result, Type resultType, | ||
Value tensor, ArrayRef<OpFoldResult> indices, | ||
ArrayRef<NamedAttribute> attrs) { | ||
SmallVector<int64_t> staticIndices; | ||
SmallVector<Value> dynamicIndices; | ||
dispatchIndexOpFoldResults(indices, dynamicIndices, staticIndices); | ||
result.addAttributes(attrs); | ||
build(b, result, resultType, tensor, dynamicIndices, | ||
b.getDenseI64ArrayAttr(staticIndices)); | ||
} | ||
|
||
// Build an ExtractOp with dynamic indexes and inferred result type. | ||
void ExtractOp::build(OpBuilder &b, OperationState &result, Type resultType, | ||
Value tensor, ValueRange indices, | ||
ArrayRef<NamedAttribute> attrs) { | ||
SmallVector<OpFoldResult> indicesValues = llvm::to_vector<4>( | ||
llvm::map_range(indices, [](Value v) -> OpFoldResult { return v; })); | ||
build(b, result, resultType, tensor, indicesValues, attrs); | ||
} | ||
|
||
// Build an ExtractOp with dynamic indexes. | ||
void ExtractOp::build(OpBuilder &b, OperationState &result, Value tensor, | ||
ValueRange indices, ArrayRef<NamedAttribute> attrs) { | ||
Type resultType = llvm::cast<TensorType>(tensor.getType()).getElementType(); | ||
SmallVector<OpFoldResult> indicesValues = llvm::to_vector<4>( | ||
llvm::map_range(indices, [](Value v) -> OpFoldResult { return v; })); | ||
build(b, result, resultType, tensor, indicesValues, attrs); | ||
} | ||
|
||
LogicalResult ExtractOp::verify() { | ||
// Verify the # indices match if we have a ranked type. | ||
auto tensorType = llvm::cast<RankedTensorType>(getTensor().getType()); | ||
if (tensorType.getRank() != static_cast<int64_t>(getIndices().size())) | ||
if (failed(checkTensorRankMatchIndices(getTensor(), getIndices(), | ||
getStaticIndices()))) | ||
return emitOpError("incorrect number of indices for extract_element"); | ||
return success(); | ||
} | ||
|
@@ -1137,12 +1189,18 @@ OpFoldResult ExtractOp::fold(FoldAdaptor adaptor) { | |
|
||
// Collect the constant indices into the tensor. | ||
SmallVector<uint64_t, 8> indices; | ||
for (Attribute indice : adaptor.getIndices()) { | ||
if (!indice || !llvm::isa<IntegerAttr>(indice)) | ||
return {}; | ||
indices.push_back(llvm::cast<IntegerAttr>(indice).getInt()); | ||
auto dynamicIndicesIt = adaptor.getIndices().begin(); | ||
for (int64_t i : getStaticIndices()) { | ||
if (i != ShapedType::kDynamic) { | ||
indices.push_back(i); | ||
} else { | ||
Attribute indice = *dynamicIndicesIt; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. typo: |
||
if (!indice || !llvm::isa<IntegerAttr>(indice)) | ||
return {}; | ||
indices.push_back(llvm::cast<IntegerAttr>(indice).getInt()); | ||
dynamicIndicesIt++; | ||
} | ||
} | ||
|
||
// Fold extract(from_elements(...)). | ||
if (auto fromElementsOp = getTensor().getDefiningOp<FromElementsOp>()) { | ||
auto tensorType = llvm::cast<RankedTensorType>(fromElementsOp.getType()); | ||
|
@@ -1354,10 +1412,48 @@ void InsertOp::getAsmResultNames( | |
setNameFn(getResult(), "inserted"); | ||
} | ||
|
||
// Build an ExtractOp with mixed static and dynamic indexes. | ||
void InsertOp::build(OpBuilder &b, OperationState &result, Value scalar, | ||
Value dest, ArrayRef<OpFoldResult> indices, | ||
ArrayRef<NamedAttribute> attrs) { | ||
build(b, result, dest.getType(), scalar, dest, indices, attrs); | ||
} | ||
|
||
// Build an InsertOp with mixed static, dynamic indexes and inferred result | ||
// Type. | ||
void InsertOp::build(OpBuilder &b, OperationState &result, Type resultType, | ||
Value scalar, Value dest, ArrayRef<OpFoldResult> indices, | ||
ArrayRef<NamedAttribute> attrs) { | ||
SmallVector<int64_t> staticIndices; | ||
SmallVector<Value> dynamicIndices; | ||
dispatchIndexOpFoldResults(indices, dynamicIndices, staticIndices); | ||
result.addAttributes(attrs); | ||
build(b, result, resultType, scalar, dest, dynamicIndices, | ||
b.getDenseI64ArrayAttr(staticIndices)); | ||
} | ||
|
||
// Build an ExtractOp with dynamic indexes and inferred result type. | ||
void InsertOp::build(OpBuilder &b, OperationState &result, Type resultType, | ||
Value scalar, Value dest, ValueRange indices, | ||
ArrayRef<NamedAttribute> attrs) { | ||
SmallVector<OpFoldResult> indicesValues = llvm::to_vector<4>( | ||
llvm::map_range(indices, [](Value v) -> OpFoldResult { return v; })); | ||
build(b, result, resultType, scalar, dest, indicesValues, attrs); | ||
} | ||
|
||
// Build an InsertOp with dynamic indexes. | ||
void InsertOp::build(OpBuilder &b, OperationState &result, Value scalar, | ||
Value dest, ValueRange indices, | ||
ArrayRef<NamedAttribute> attrs) { | ||
SmallVector<OpFoldResult> indicesValues = llvm::to_vector<4>( | ||
llvm::map_range(indices, [](Value v) -> OpFoldResult { return v; })); | ||
build(b, result, dest.getType(), scalar, dest, indicesValues, attrs); | ||
} | ||
|
||
LogicalResult InsertOp::verify() { | ||
// Verify the # indices match if we have a ranked type. | ||
auto destType = llvm::cast<RankedTensorType>(getDest().getType()); | ||
if (destType.getRank() != static_cast<int64_t>(getIndices().size())) | ||
if (failed(checkTensorRankMatchIndices(getDest(), getIndices(), | ||
getStaticIndices()))) | ||
return emitOpError("incorrect number of indices"); | ||
return success(); | ||
} | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are some spelling inconsistences:
indexes
,indices