Skip to content

[Injective Schedule] make injective ops's opt schedule applied to every output tensor#11820

Merged
masahi merged 1 commit intoapache:mainfrom
crazydemo:upstream-schedule
Jun 27, 2022
Merged

[Injective Schedule] make injective ops's opt schedule applied to every output tensor#11820
masahi merged 1 commit intoapache:mainfrom
crazydemo:upstream-schedule

Conversation

@crazydemo
Copy link

Modify the schedule for injective ops. All outputs of injective ops should share the same schedule optimization, e.g. outputs of OP split are all supposed to be lowered into parallel for.

new ir:
IRModule({GlobalVar(tvmgen_default_fused_split): PrimFunc([placeholder, T_split_sections, T_split_sections]) attrs={"from_legacy_te_schedule": (bool)1, "global_symbol": "tvmgen_default_fused_split", "tir.noalias": (bool)1} {
  buffer_realize T_split_sections([0, 1], [0, 64], [0, 56], [0, 56]) {
    parallel (ax0.ax1.fused, 0, 64) {
      for (ax2, 0, 56) {
        for (ax3.outer, 0, 4) {
          vectorized (ax3.inner, 0, 16) {
            if (tir.likely(((ax3.inner + (ax3.outer*16)) < 56))) {
              T_split_sections[floordiv(ax0.ax1.fused, 64), floormod(ax0.ax1.fused, 64), ax2, (ax3.inner + (ax3.outer*16))] = placeholder[floordiv(ax0.ax1.fused, 64), floormod(ax0.ax1.fused, 64), ax2, (ax3.inner + (ax3.outer*16))]
            }
          }
        }
      }
    }
    buffer_realize T_split_sections([0, 1], [0, 64], [0, 56], [0, 56]) {
      parallel (ax0.ax1.fused, 0, 64) {
        for (ax2, 0, 56) {
          for (ax3.outer, 0, 4) {
            vectorized (ax3.inner, 0, 16) {
              if (tir.likely(((ax3.inner + (ax3.outer*16)) < 56))) {
                T_split_sections[floordiv(ax0.ax1.fused, 64), floormod(ax0.ax1.fused, 64), ax2, (ax3.inner + (ax3.outer*16))] = placeholder[floordiv(ax0.ax1.fused, 64), (floormod(ax0.ax1.fused, 64) + 64), ax2, (ax3.inner + (ax3.outer*16))]
              }
            }
          }
        }
      }
    }
  }
}
})

old ir:
IRModule({GlobalVar(tvmgen_default_fused_split): PrimFunc([placeholder, T_split_sections, T_split_sections]) attrs={"from_legacy_te_schedule": (bool)1, "global_symbol": "tvmgen_default_fused_split", "tir.noalias": (bool)1} {
  buffer_realize T_split_sections([0, 1], [0, 64], [0, 56], [0, 56]) {
    parallel (ax0.ax1.fused, 0, 64) {
      for (ax2, 0, 56) {
        for (ax3.outer, 0, 4) {
          vectorized (ax3.inner, 0, 16) {
            if (tir.likely(((ax3.inner + (ax3.outer*16)) < 56))) {
              T_split_sections[floordiv(ax0.ax1.fused, 64), floormod(ax0.ax1.fused, 64), ax2, (ax3.inner + (ax3.outer*16))] = placeholder[floordiv(ax0.ax1.fused, 64), floormod(ax0.ax1.fused, 64), ax2, (ax3.inner + (ax3.outer*16))]
            }
          }
        }
      }
    }
    buffer_realize T_split_sections([0, 1], [0, 64], [0, 56], [0, 56]) {
      for (ax1, 0, 64) {
        for (ax2, 0, 56) {
          for (ax3, 0, 56) {
            T_split_sections[0, ax1, ax2, ax3] = placeholder[0, (ax1 + 64), ax2, ax3]
          }
        }
      }
    }
  }
}
})

@crazydemo
Copy link
Author

@masahi Could you please help review this PR?

@masahi masahi merged commit 1115fd9 into apache:main Jun 27, 2022
mikeseven pushed a commit to mikeseven/tvm that referenced this pull request Sep 27, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants