The Relax transform RewriteDataflowReshape identifies TIR functions that are equivalent to relax.op.reshape, and replaces them with calls to relax.op.reshape. This is used as a precursor for simplifications that rely on the high-level knowledge that an operator is a reshape, but also require the low-level knowledge of the adjacent TIR PrimFuncs.
Prior to this commit, the RewriteDataflowReshape pass would only recognize static shapes, or dynamic shapes that could be inferred from the shapes of tensor arguments. This commit updates RewriteDataflowReshape to recognize cases where an extra symbolic variable has been provided.
masahi
commented
3 months ago
The Relax transform
RewriteDataflowReshapeidentifies TIR functions that are equivalent torelax.op.reshape, and replaces them with calls torelax.op.reshape. This is used as a precursor for simplifications that rely on the high-level knowledge that an operator is a reshape, but also require the low-level knowledge of the adjacent TIR PrimFuncs.Prior to this commit, the
RewriteDataflowReshapepass would only recognize static shapes, or dynamic shapes that could be inferred from the shapes of tensor arguments. This commit updatesRewriteDataflowReshapeto recognize cases where an extra symbolic variable has been provided.