tvm icon indicating copy to clipboard operation
tvm copied to clipboard

[Node][Metaschedule] Allow ignoring NDArray raw data in StructuralEqual / Hash

Open masahi opened this issue 3 years ago • 2 comments
trafficstars

Motivation: Currently, StructuralEqual/Hash take into account NDArray raw data to determine equality / compute hash. This is problematic for link-params = True case, because meta schedule task extraction and database look up will treat the following subgraphs as distinct, even though the only difference is the content of NDArray constants.

def @fused_nn.conv2d_add(%p0: Tensor[(1, 32, 32, 256), float32] /* ty=Tensor[(1, 32, 32, 256), float32] */, Primitive=1) -> Tensor[(1, 32, 32, 256), fl
oat32] {                                                                                                                      
  %0 = nn.conv2d(%p0, meta[relay.Constant][0] /* ty=Tensor[(3, 3, 256, 256), float32] */, padding=[1, 1, 1, 1], channels=256, kernel_size=[3, 3], data_
layout="NHWC", kernel_layout="HWIO") /* ty=Tensor[(1, 32, 32, 256), float32] */;                                                                       
  add(%0, meta[relay.Constant][1] /* ty=Tensor[(1, 1, 1, 256), float32] */) /* ty=Tensor[(1, 32, 32, 256), float32] */        
}

def @fused_nn.conv2d_add(%p0: Tensor[(1, 32, 32, 256), float32] /* ty=Tensor[(1, 32, 32, 256), float32] */, Primitive=1) -> Tensor[(1, 32, 32, 256), fl
oat32] {                                                                                                                      
  %0 = nn.conv2d(%p0, meta[relay.Constant][0] /* ty=Tensor[(3, 3, 256, 256), float32] */, padding=[1, 1, 1, 1], channels=256, kernel_size=[3, 3], data_
layout="NHWC", kernel_layout="HWIO") /* ty=Tensor[(1, 32, 32, 256), float32] */;                                                                       
  add(%0, meta[relay.Constant][1] /* ty=Tensor[(1, 1, 1, 256), float32] */) /* ty=Tensor[(1, 32, 32, 256), float32] */  
} 

This PR adds options to StructuralEqual/Hash to allow ignoring NDArray raw data when comparing / hashing. MS task extraction and database look up are now using these options, since NDArray raw data doesn't matter to them.

cc @Hzfengsy @junrushao1994

masahi avatar Sep 06 '22 01:09 masahi

thanks @masahi . I thin this belongs to a more general topic of canonicalizing the workload keys. While adding such option would indeed resolve this particular case, it would be good to think about if we can have a more generalized approach(for future cases that needs canonicalization, e.g. those that have followup elem-wise patterns which allows some form of reuse)

One possible approach is to have a CanonicalizeWorkloadKey which takes a TIR and Canonicalizes that into a different one(e.g. change alloc_const to a the function parameter which contains a buffer). Depends on how fast we can do this, we can choose to introduce the flag.

This also helps us to prevent possible future cases that can choose to schedule transform depending on the const value content

tqchen avatar Sep 06 '22 01:09 tqchen

@tqchen Thanks, I like the idea of the canonicalized key.

Last week I discussed with @zxybazh on the possibility of identifying two subgraphs with the identical anchor op workload (conv2d -> add and conv2d -> add -> add etc) as being equal, to reduce the tuning time like you suggested. I can imagine that, moving away from using the raw Relay / TIR mod as a workload key is a good direction toward that goal.

masahi avatar Sep 06 '22 01:09 masahi

Superceded by https://github.com/apache/tvm/pull/13091

masahi avatar Oct 17 '22 07:10 masahi