[Ansor][AutoTVM v2.0] Phase 1: feature extraction for cost models

[merrymercy]

For the full upstream plan, see Ansor RFC.

This PR adds feature extraction for the cost models.
It is similar to the existing feature extraction in autotvm but

1. Uses a general format for all loop structures, so all workloads can share the same cost model.
2. Includes more analysis results (e.g., reuse distance) into the feature vector.

[merrymercy]

cc @tqchen @junrushao1994 @jcf94 @comaniac @whbldhwj

[comaniac]

Reviewed all but feature.cc. Will finish it by tomorrow.

[comaniac]

Tried to review the features but they are too specific. Do you think if we can make up a registration mechanism like register("feature_name", extractor) to better manage them?

[merrymercy]

@comaniac It is not easy to do so. Now we share a lot of computation for different features. If we compute them separately, redundant computation will be introduced. In addition, the global context is also shared for these features. How to register the required global context for a specific feature is also a problem.

[junrushao]

nitpick on the doc

Suggested change

	# int[n+2] sizes
	# int sizes[0]
	# ...
	# int sizes[n + 1]

[yangjunpro]

Personally, I think junrushao1994's suggestion is clearer than current comment convention.Since I think "float sizes[n + 1]" illustrates the format semantics more precisely than "float[sizes[0]]".

[merrymercy]

I won't take either of your suggestions. I use int[x] variable to denote an array of x integer values with the name variable.

@junrushao1994 's suggestion is wrong because int size[n] has another meaning of n+1 integers, but here we actually mean nth integer.
@yangjunpro 's suggestion is also wrong. It makes sizes become the name of the filed, but actually size only specifies the size.

[junrushao]

perhaps add a comment here, like "for avoiding rounding error"? btw, why we use a float for integer?

[merrymercy]

Some of them are int while the others are float. I want to store all of them in a single array, but we do not have union in tvm::Object. So I use a single float array to store both int and float

[yangjunpro]

+1, how about store the float array and the n_stmts separately? By doing this we may need to add extra code, but I think the program semantic is clearer.

[tqchen]

per discussion with @merrymercy :

• split the feature extraction logic into FeatureExtractor(callback by visitor) and visitor.
• Provide FeatureExtractor for each group of features

[merrymercy]

I organized the features into 5 groups.

  // Group 1: Computation related features
  // Group 2: Buffer access related features (per buffer)
  // Group 3: Arithmetic intensity related features
  // Group 4: Allocation related features
  // Group 5: Outer scope related features

The specification can be found in src/auto_scheduler/feature.cc::FeatureSet.

Each group has one corresponding extraction function, they are called in the main visitor (PerStoreFeatureExtractor).

  void VisitStmt_(const BufferStoreNode* node) final {
    ...
    // Group 1: Computation related features
    ExtractComputationFeature(node, math_op_counter);

    // Group 2: Buffer access related features (per buffer)
    ExtractBufferAccessFeature(node, math_op_counter, &cur_compute_ops, &compute_ops_list,
                               &mem_bytes_list);

    // Group 3: Arithmetic intensity related features
    ExtractArithmeticIntensityFeature(node, cur_compute_ops, compute_ops_list, mem_bytes_list);

    // Group 4: Allocation related features
    ExtractOuterScopeFeature(node);
  }

  void VisitStmt_(const BufferRealizeNode* node) final {
    StmtExprVisitor::VisitStmt_(node);

    // Group 5: Outer scope related features
    ExtractAllocationFeature(node);
  }

I think the code is very clean and much better than the old autotvm now.
I don't like registration or adding an extra layer of callback. It is over design and make things more complicated.

@tqchen @comaniac @FrozenGene @jroesch @junrushao1994 Your comments are all addressed. This PR is ready to be merged.

[comaniac]

Much cleaner now. LGTM

[tqchen]

@merrymercy please fix the conflict

[yangjunpro]

Just some nitpicking comments.

[yangjunpro]

Personally, I think junrushao1994's suggestion is clearer than current comment convention.Since I think "float sizes[n + 1]" illustrates the format semantics more precisely than "float[sizes[0]]".

[yangjunpro]

+1, how about store the float array and the n_stmts separately? By doing this we may need to add extra code, but I think the program semantic is clearer.