已合并
fix(random): InitTensorSliceState增加coalesce对齐PyTorch切分 #3980
liangtongxue创建于 7月10日
fix(random): InitTensorSliceState增加coalesce对齐PyTorch切分 #3980
已合并
共 1 个文件变更+42-7
| @@ -88,20 +88,55 @@ void TensorSliceState::PartitionDim(int64_t dim, TensorSliceState& other) | |||
| 88 | ge::graphStatus InitTensorSliceState( | 88 | ge::graphStatus InitTensorSliceState( |
| 89 | TensorSliceState& state, const gert::Shape& outputTensor, int64_t outputSize, ge::DataType outputDtype) | 89 | TensorSliceState& state, const gert::Shape& outputTensor, int64_t outputSize, ge::DataType outputDtype) |
| 90 | { | 90 | { |
| 91 | - state.ndim = static_cast<int64_t>(outputTensor.GetDimNum()); | ||
| 92 | state.numel = outputSize; | 91 | state.numel = outputSize; |
| 93 | state.gmOffset = 0; | 92 | state.gmOffset = 0; |
| 94 | state.elementSize = ge::GetSizeByDataType(outputDtype); | 93 | state.elementSize = ge::GetSizeByDataType(outputDtype); |
| 95 | 94 | ||
| 96 | - for (int64_t dim = 0; dim < state.ndim && dim < MAX_TENSOR_DIMS; dim++) { | 95 | + int64_t rawShape[MAX_TENSOR_DIMS] = {0}; |
| 97 | - state.shape[dim] = outputTensor.GetDim(static_cast<size_t>(dim)); | 96 | + int64_t rawStrides[MAX_TENSOR_DIMS] = {0}; |
| 97 | + int64_t rawNdim = static_cast<int64_t>(outputTensor.GetDimNum()); | ||
| 98 | + for (int64_t dim = 0; dim < rawNdim && dim < MAX_TENSOR_DIMS; dim++) { | ||
| 99 | + rawShape[dim] = outputTensor.GetDim(static_cast<size_t>(dim)); | ||
| 100 | + } | ||
| 101 | + if (rawNdim > 0) { | ||
| 102 | + rawStrides[rawNdim - 1] = 1; | ||
| 103 | + for (int64_t dim = rawNdim - 2; dim >= 0; dim--) { | ||
| 104 | + rawStrides[dim] = rawShape[dim + 1] * rawStrides[dim + 1]; | ||
| 105 | + } | ||
| 98 | } | 106 | } |
| 99 | 107 | ||
| 100 | - if (state.ndim > 0) { | 108 | + int64_t coalescedShape[MAX_TENSOR_DIMS] = {0}; |
| 101 | - state.strides[state.ndim - 1] = 1; | 109 | + int64_t coalescedStrides[MAX_TENSOR_DIMS] = {0}; |
| 102 | - for (int64_t dim = state.ndim - 2; dim >= 0; dim--) { | 110 | + int64_t coalescedNdim = 0; |
| 103 | - state.strides[dim] = state.shape[dim + 1] * state.strides[dim + 1]; | 111 | + for (int64_t dim = 0; dim < rawNdim; dim++) { |
| 112 | + if (rawShape[dim] == 1) { | ||
| 113 | + continue; | ||
| 104 | } | 114 | } |
| 115 | + coalescedShape[coalescedNdim] = rawShape[dim]; | ||
| 116 | + coalescedStrides[coalescedNdim] = rawStrides[dim]; | ||
| 117 | + coalescedNdim++; | ||
| 118 | + } | ||
| 119 | + if (coalescedNdim == 0) { | ||
| 120 | + coalescedShape[0] = 1; | ||
| 121 | + coalescedStrides[0] = 1; | ||
| 122 | + coalescedNdim = 1; | ||
| 123 | + } | ||
| 124 | + for (int64_t dim = coalescedNdim - 1; dim > 0; dim--) { | ||
| 125 | + if (coalescedShape[dim] * coalescedStrides[dim] == coalescedStrides[dim - 1]) { | ||
| 126 | + coalescedShape[dim - 1] *= coalescedShape[dim]; | ||
| 127 | + coalescedStrides[dim - 1] = coalescedStrides[dim]; | ||
| 128 | + for (int64_t d = dim; d < coalescedNdim - 1; d++) { | ||
| 129 | + coalescedShape[d] = coalescedShape[d + 1]; | ||
| 130 | + coalescedStrides[d] = coalescedStrides[d + 1]; | ||
| 131 | + } | ||
| 132 | + coalescedNdim--; | ||
| 133 | + } | ||
| 134 | + } | ||
| 135 | + | ||
| 136 | + state.ndim = coalescedNdim; | ||
| 137 | + for (int64_t dim = 0; dim < coalescedNdim && dim < MAX_TENSOR_DIMS; dim++) { | ||
| 138 | + state.shape[dim] = coalescedShape[dim]; | ||
| 139 | + state.strides[dim] = coalescedStrides[dim]; | ||
| 105 | } | 140 | } |
| 106 | 141 | ||
| 107 | return ge::GRAPH_SUCCESS; | 142 | return ge::GRAPH_SUCCESS; |