[Cuda] Slight optimize the default injective schedule

[ANSHUMAN87]

We do not need additional thread extent when the tensor size is known.
Below output shows the impact of the change, taken Transpose Op for illustration.

Before:

producer_realize T_transpose([0, 3], [0, 4]) {
  // attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = 1
  // attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 1024
  if (tir.likely((floordiv(threadIdx.x, 4) < 3))) {
    if (tir.likely((threadIdx.x < 12))) {
      T_transpose[floordiv(threadIdx.x, 4), floormod(threadIdx.x, 4)] =placeholder[floormod(threadIdx.x, 4), floordiv(threadIdx.x, 4)]
    }
  }
}

After:

producer_realize T_transpose([0, 3], [0, 4]) {
  // attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = 1
  // attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 12
  T_transpose[floordiv(threadIdx.x, 4), floormod(threadIdx.x, 4)] =placeholder[floormod(threadIdx.x, 4), floordiv(threadIdx.x, 4)]
}

[junrushao]

LGTM. Thanks!

[ANSHUMAN87]

Thanks @junrushao1994 ! Looks like the test case is flaky. Does not depend on this PR change. Will re-trigger CI.

[junrushao]

Thanks @ANSHUMAN87! I don’t have other comments. @comaniac would you like to take a look?

[comaniac]

Thanks @ANSHUMAN87 @junrushao1994