PR Category

CINN

PR Types

Improvements

Description

This PR introduces an optimization for the pe::Reshape index calculation. The previous implementation involved complex calculations to determine the input tensor indices from the output tensor indices, by computing the absolute offset in the output and then back-calculating each dimension of the input tensor. For example, transforming an input tensor x of shape [S0, S1, S2, 128] to an output tensor y of shape [S0, S1*S2, 32, 4] required calculating the offset as offset=((i*S1*S2+j)*32+k)*4+m, and then converting it to the input indices using [offset/(S1*S2*128), offset/(S2*128)%S1, offset/128%S2, offset%128].

The optimization leverages the correspondence between input and output axes, for instance, [32, 4] corresponding to [128], and [S1*S2] corresponding to [S1, S2]. The new approach generates each axis index directly based on this correspondence, resulting in simpler and more efficient calculations. For the given example, the optimized indices are computed as [i, j/S2, j%S2, k*4+m].

pcard-85711