Support ingesting Opset 15 Shape Op #1684
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -9,60 +9,83 @@ | |
| //===----------------------------------------------------------------------===// | ||
|
|
||
| #include "src/Dialect/ONNX/ShapeInference/ONNXShapeHelper.hpp" | ||
| #include <tuple> | ||
| #include <utility> | ||
|
|
||
| using namespace mlir; | ||
|
|
||
| namespace onnx_mlir { | ||
|
|
||
| namespace { | ||
|
|
||
| // The Shape op spec says: | ||
| // | ||
| // "Note that axes will be clipped to the range [0, r-1], where r is the | ||
| // rank of the input tensor if they are out-of-range (after adding r in the case | ||
| // of negative axis). Thus, specifying any end value > r is equivalent to | ||
| // specifying an end value of r, and specifying any start value < -r is | ||
| // equivalent to specifying a start value of 0." | ||
| int64_t normalize(int64_t axis, int64_t rank) { | ||
|
There was a problem hiding this comment. I see what you do here, it sounds reasonable but is it the standard? Or should we instead report an error? @chentong319 do you have an opinion on this issue? There was a problem hiding this comment. Sorry if this was confusing. This was actually lifted for the ONNX specification
It seems weird, but I have based it off the ONNX spec of Shape. There was a problem hiding this comment. Great, I would then just reflect this policy in the function name, so that others may not be tempted to use it when it is not applicable. Add something of the sort "ClampedPerSpec" or something similar. |
||
| if (axis < 0) | ||
| axis += rank; | ||
|
|
||
| if (axis < 0) | ||
| axis = 0; | ||
|
|
||
| if (axis > rank) | ||
| axis = rank; | ||
|
|
||
| return axis; | ||
| } | ||
|
|
||
| } // namespace | ||
|
|
||
| // Compute a slice of the input tensor's shape. The slice starts from axis 0. | ||
| // The axes up to the last one will be included. Negative axes indicate counting | ||
| // back from the last axis. | ||
| std::pair<int64_t, int64_t> getDataShapeBounds( | ||
| ONNXShapeOpAdaptor &operandAdaptor) { | ||
| Value data = operandAdaptor.data(); | ||
| MemRefBoundsIndexCapture dataBounds(data); | ||
| int64_t rank = dataBounds.getRank(); | ||
|
|
||
| // Compute the normalized start/end. Negative value means counting | ||
| // dimensions from the back. | ||
| int64_t start = operandAdaptor.start(); | ||
| int64_t end = rank; | ||
| if (operandAdaptor.end().has_value()) { | ||
| end = operandAdaptor.end().value(); | ||
| } | ||
|
|
||
| return std::make_pair(normalize(start, rank), normalize(end, rank)); | ||
| } | ||
|
|
||
| LogicalResult ONNXShapeOpShapeHelper::computeShape( | ||
| ONNXShapeOpAdaptor operandAdaptor) { | ||
| Value data = operandAdaptor.data(); | ||
| MemRefBoundsIndexCapture dataBounds(data); | ||
|
|
||
| int64_t start; | ||
| int64_t end; | ||
| std::tie(start, end) = getDataShapeBounds(operandAdaptor); | ||
|
|
||
| // Output is the actual number of values (1D) | ||
| dimsForOutput().emplace_back(LiteralIndexExpr(end - start)); | ||
|
|
||
| return success(); | ||
| } | ||
|
|
||
| // Compute the data selected by the Shape operator. | ||
| DimsExpr computeSelectedData(ONNXShapeOpAdaptor &operandAdaptor) { | ||
| MemRefBoundsIndexCapture dataBounds(operandAdaptor.data()); | ||
| int64_t start; | ||
| int64_t end; | ||
| std::tie(start, end) = getDataShapeBounds(operandAdaptor); | ||
| assert(start >= 0 && start <= end && end <= (int64_t)dataBounds.getRank() && | ||
| "Unexpected bounds"); | ||
|
|
||
| DimsExpr selectedData; | ||
| for (int64_t i = start; i < end; ++i) | ||
| selectedData.emplace_back(dataBounds.getDim(i)); | ||
|
|
||
| return selectedData; | ||
|
|
||
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.
would it make sense to have an assert here for sizes, or better, it could be added in the verifier (apologies if it is already there)