【Question】How to use storage_align?

Problem Description: Example: The axis of s[conv2d_bias_stage].op is dim= [1,1,16,16] , But, In the process of ScheduleOps, dim[2]=16 is optimized by GatherBound to dim[2] = 1(Only one of the 16 data of dim[2] is valid. I want to keep dim[2] = 16 not optimized. I found that there is a storage_align method in schedule. So I tried to do this: nns, mm, m16, n16= s[conv2d_bias_stage].op.axis s[conv2d_bias_stage].storage_align(m16, 16,0) But the actual effect is still optimized.

IR:
// attr [conv2d_bias] storage_scope = "local.out_buffer"
     allocate conv2d_bias[float16 * 16]
     produce conv2d_bias {
       for (m16.inner, 0, 1) {
         for (n16, 0, 16) {
           xxxxxx
           }
         }
       }
IR wanted:
// attr [conv2d_bias] storage_scope = "local.out_buffer"
     allocate conv2d_bias[float16 * 16]
     produce conv2d_bias {
       for (m16.inner, 0, 16) {
         for (n16, 0, 16) {
           xxxxxx
           }
         }
       }

Question : How can I use storage_align? Is it possible to achieve this goal?

1 Like

Hi, as far as I know, storage_align cannot achieve your goal to keep the range of dim[2], and it only updates the strides in compute body and the extents in allocate node, but not modify the extent of dim[2]. See tvm/storage_flatten.cc at d425c144adb6e7a840da071b044773eb525e67c1 · apache/tvm · GitHub

Thank you very much for your support! Except for storage_alig, is there any other way to achieve the goal?

@lcjdd I think there is no way. But I saw other methods before, see this project based on TVM Search · buffer_align (github.com)

1 Like

OK, Thanks!! I try other ways to achieve。