Symbolic shape stride calculation

[chungs31]

Hi there, I was hoping to be pointed to which files decide that when there is a te.var instead of an integer in a shape, create additional int32 stride arguments. For example I am using topi.nn.conv2d and the default nn schedule, and when I print the IR using tvm.lower I get this:

primfn(in_1: handle, w_1: handle) -> ()                                                                               
  attr = {"global_symbol": "main", "tir.noalias": True}                                                               
  buffers = {w: Buffer(w_2: Pointer(float32), float32, [ff: int32, rc: int32, 1, 1], [stride: int32, stride_1: int32, 
stride_2: int32, stride_3: int32], type="auto"),                                                                      
             in: Buffer(in_2: Pointer(float32), float32, [1, rc, xx: int32, xx], [stride_4: int32, stride_5: int32, st
ride_6: int32, stride_7: int32], type="auto")}                                                                        
  buffer_map = {in_1: in, w_1: w} {                                                                                   
  attr [compute: Pointer(float32)] "storage_scope" = "global";                                                        
  allocate(compute, float32, [((ff*xx)*xx)]);                                                                         
  attr [compute_1: Pointer(float32)] "storage_scope" = "global";                                                      
  allocate(compute_1, float32, [(xx*xx)]);                                                                            
  attr [IterVar(pipeline: int32, (nullptr), "ThreadIndex", "pipeline")] "pipeline_exec_scope" = 1;                    
  for (i1: int32, 0, ff) {                                                                                            
    for (yy: int32, 0, xx) {                                                                                          
      for (xx_1: int32, 0, xx) {                                                                                      
        compute_1[((yy*xx) + xx_1)] = 0f32                                                                            
        for (rc_1: int32, 0, rc) {                                                                                    
          compute_1[((yy*xx) + xx_1)] = ((float32*)compute_1[((yy*xx) + xx_1)] + ((float32*)in_2[(((rc_1*stride_5) + (yy*stride_6)) + (xx_1*stride_7))]*(float32*)w_2[((i1*stride) + (rc_1*stride_1))]))                                    
        }                                                                                                             
      }                                                                                                               
    }                                                                                                                 
    for (i2: int32, 0, xx) {                                                                                          
      for (i3: int32, 0, xx) {                                                                                        
        compute[((((i1*xx) + i2)*xx) + i3)] = max((float32*)compute_1[((i2*xx) + i3)], 0f32)                          
      }                                                                                                               
    }                                                                                                                 
  }                                                                                                                   
}

Where are those strides (stride_2, stride_3, stride_6, etc.) calculated? At which stage are these arguments created?

[giuseros]

Hi @chungs31,

I was looking into this a couple of days ago. This is my understanding (more experienced folks, please correct me if I am wrong).

There is a pass named MakePackedAPI, which has the task of binding the variables of DLTensors to TVM buffers.

Lines 177-179 of make_packed_api.cc are like:

  for (const auto& kv : buffer_def) {
    binder.BindDLTensor(kv.second, device_type, device_id, kv.first, kv.first->name_hint);
  }

BindDLTensor adds LetStmt for every variable related to the tensor (size, stride, etc…) calling the function Bind_. For stride, you should be looking at lines 251-258 of arg_binder.cc.

What Bind_ does is to look for the requested variable, creating the LetStmt.

Hope this helps,