Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 1 commit into from
Jun 27, 2022

Conversation

crazydemo
Copy link
Contributor

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
Contributor Author

@masahi Could you please help review this PR?

@masahi masahi merged commit 1115fd9 into apache:main Jun 27, 2022
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