[AutoScheduler] New layout rewrite option: Weight pre-transpose

[jcf94]

In Ansor, we have an optimization called "Layout rewrite", which modifies the input weight of a specific op according to its schedule to get better performance, a previous PR about this is #6297 .

This PR brings another option for this feature, besides directly modify the input placeholder, we now support to insert a transpose stage between the placeholder and compute op.

Others:

• Add a CopyOnWrite methods for transform_steps
• Some small fix

cc @merrymercy @comaniac @minminsun

[junrushao]

Hey, thank you for the contribution! May I know the difference between layout rewrite and this weight pre-transpose? Looks like weight pre-transpose can be done in compile time, so why do we insert a stage instead? Thanks a lot!

[FrozenGene]

enum class LayoutRewriteOption : uint8 should be enough.

[FrozenGene]

Hey, thank you for the contribution! May I know the difference between layout rewrite and this weight pre-transpose? Looks like weight pre-transpose can be done in compile time, so why do we insert a stage instead? Thanks a lot!

This could keep the weight shape is the same as before. However, I am curious too what benefit this will bring in.

[merrymercy]

The code looks good to me. But the naming is not intuitive. For example, the "RewriteWithPlaceholder" option in your code actually accept "PreTransposed" tensors as inputs.
The key difference between the two options is whether adding new stages. We can make the names much clearer as follows.

/*!
 * \brief Options for applying layout rewrite.
 * This is an optimization to rewrite the layout of input tensors according to the schedule we get.
 */
enum class LayoutRewriteOption : int {
  /*! \brief Do not process layout rewrite. */
  NoRewrite = 0,
  /*! \brief Insert layout transformation stages for input placeholders in the compute DAG */
  InsertTransformStage = 1,
  /*!
   * \brief Do not insert layout transformation stages and assume the input placeholders
   * are pre-transformed.
   * \note The lowered function with this option does not accept the origial input shapes,
   * so this option must be used along with a layout conversion pass in Relay.
   */
  RewriteForPreTransformed = 2,
};

In addition, I prefer "transform" over "transpose" because we can support other kinds of rewriting besides the current simple "transpose".

[jcf94]

@junrushao1994 @FrozenGene Thanks!
Our former implementation of layout rewrite is used for end to end relay integration, and we'll modify the input tensor directly since all these can be seen as parts of the pre-process.
Currently we got a simpler demand, we would like to get the well scheduled kernel from TVM and apply it to some where else(out of TVM runtime). We want to keep the input/output size the same, so it can be used as a out-of-the-box kernel, while at the same time we can still benefit from the layout rewrite optimization.

[comaniac]

I don't have other comments. Thanks!

[merrymercy]

@jcf94 Please fix the test cases.

[merrymercy]

The problem is not about random seed. It is because the condition of np.testing.assert_allclose is too strict.

[jcf94]

The problem is not about random seed. It is because the condition of np.testing.assert_allclose is too strict.

Emm ... That may also be a problem, I'm thinking that these float operations should be in theory absolutely the same, even with a very strict precision.

At least the schedules generated with fixed random seed 0 can pass these check.

[comaniac]

@merrymercy @jcf94 the tests added by this PR seem flaky. Please see https://ci.tlcpack.ai/blue/organizations/jenkins/tvm/detail/main/124/pipeline

[comaniac]

It turns out that #6828 didn't really disable the flaky tests. It simply comments out the function call in __main__, but pytest will look for all functions with test_ as the prefix. I'll file another PR to disable them.

[comaniac]

#6841 filed.