| 2760 | } |
| 2761 | |
| 2762 | bool StridedSliceOp::GetSlicedBoundRanges( |
| 2763 | SmallVectorImpl<int64_t> *begin_indices, |
| 2764 | SmallVectorImpl<int64_t> *end_indices, SmallVectorImpl<int64_t> *strides) { |
| 2765 | // TODO(hinsu): Support lowering for ops with dynamic begin and end values |
| 2766 | // when it is possible to derive indices based on mask attributes. |
| 2767 | DenseIntElementsAttr sparse_begin_attr, sparse_end_attr, sparse_strides_attr; |
| 2768 | if (!matchPattern(this->begin(), m_Constant(&sparse_begin_attr)) || |
| 2769 | !matchPattern(this->end(), m_Constant(&sparse_end_attr)) || |
| 2770 | !matchPattern(this->strides(), m_Constant(&sparse_strides_attr))) |
| 2771 | return false; |
| 2772 | |
| 2773 | auto input_ty = this->input().getType().dyn_cast<RankedTensorType>(); |
| 2774 | if (!input_ty || !input_ty.hasStaticShape()) return false; |
| 2775 | auto input_shape = llvm::to_vector<4>(input_ty.getShape()); |
| 2776 | int rank = input_shape.size(); |
| 2777 | |
| 2778 | SmallVector<int64_t, 4> sparse_begin, sparse_end, sparse_strides; |
| 2779 | |
| 2780 | for (const APInt &index : sparse_begin_attr) |
| 2781 | sparse_begin.push_back(index.getSExtValue()); |
| 2782 | for (const APInt &index : sparse_end_attr) |
| 2783 | sparse_end.push_back(index.getSExtValue()); |
| 2784 | for (const APInt &stride : sparse_strides_attr) |
| 2785 | sparse_strides.push_back(stride.getSExtValue()); |
| 2786 | |
| 2787 | auto num_sparse_indices = sparse_begin_attr.getNumElements(); |
| 2788 | SparseSliceSpec sparse = {num_sparse_indices, |
| 2789 | this->begin_mask().getZExtValue(), |
| 2790 | this->end_mask().getZExtValue(), |
| 2791 | this->ellipsis_mask().getZExtValue(), |
| 2792 | this->new_axis_mask().getZExtValue(), |
| 2793 | this->shrink_axis_mask().getZExtValue(), |
| 2794 | sparse_begin, |
| 2795 | sparse_end, |
| 2796 | sparse_strides}; |
| 2797 | |
| 2798 | DenseSliceSpec dense = {rank, |
| 2799 | /*begin_mask = */ 0, |
| 2800 | /*end_mask = */ 0, |
| 2801 | /*shrink_axis_mask = */ 0, |
| 2802 | *begin_indices, |
| 2803 | *end_indices, |
| 2804 | *strides}; |
| 2805 | |
| 2806 | if (failed(BuildDenseSliceSpec(sparse, &dense))) return false; |
| 2807 | |
| 2808 | CalculateSlicedShapeAndBoundRanges(input_shape, dense.begin_mask, |
| 2809 | dense.end_mask, dense.shrink_axis_mask, |
| 2810 | *begin_indices, *end_indices, *strides); |
| 2811 | return true; |
| 2812 | } |
| 2813 | |
| 2814 | //===----------------------------------------------------------------------===// |
| 2815 | // StridedSliceGradOp |
no test coverage detected