Skip to content

[IR] Add explicit unroll#184

Merged
yaoyaoding merged 5 commits intohidet-org:mainfrom
yaoyaoding:hidet-unroll
Apr 21, 2023
Merged

[IR] Add explicit unroll#184
yaoyaoding merged 5 commits intohidet-org:mainfrom
yaoyaoding:hidet-unroll

Conversation

@yaoyaoding
Copy link
Copy Markdown
Member

@yaoyaoding yaoyaoding commented Apr 21, 2023

  1. update the API to specify unroll.
  2. add support for explicit unroll in hidet (expand the loop in hidet instead of using pragma unroll)

Syntax for the attr string:

        attr-string: attr+
        attr:
             | unroll
             | default
        unroll:
             | 'u'          # unroll
             | 'u' INT+     # unroll with factor, e.g., u1 u2 u3. u1 indicates unroll with factor 1 (i.e., no unroll)
             | 'u' '+'      # explicit unroll, will be unrolled by hidet instead of underlying compiler
        default: '.'

    with hidet.script_module() as script_module:
        @hidet.script
        def example():
            attr.func_kind = 'host_kernel'

            for i in grid(10, attrs='u'):   # unroll
                printf("i = %d\n", i)

            for i in grid(10, attrs='u+'):   # unroll explicitly
                printf("i = %d\n", i)

            for i in grid(10, attrs='u2'):  # unroll with factor 2
                printf("i = %d\n", i)

            for i, j in grid(2, 5, attrs='u.'):   # unroll the first loop while keep the second loop unchanged
                printf("i = %d, j = %d\n", i, j)

            for w in range(32):
                # unroll the first loop while keep the second loop unchanged in the repeat task mapping
                for i, j in repeat(2, 8, attrs='u.').spatial(4, 8).on(w):
                    printf("i = %d, j = %d\n", i, j)
__host__ void hidet_example() {
  #pragma unroll
  for (int32_t i = 0; (i < 10); i = (i + 1)) {
    printf("i = %d\n", i);
  } 
  printf("i = %d\n", 0);
  printf("i = %d\n", 1);
  printf("i = %d\n", 2);
  printf("i = %d\n", 3);
  printf("i = %d\n", 4);
  printf("i = %d\n", 5);
  printf("i = %d\n", 6);
  printf("i = %d\n", 7);
  printf("i = %d\n", 8);
  printf("i = %d\n", 9);
  #pragma unroll 2
  for (int32_t i_1 = 0; (i_1 < 10); i_1 = (i_1 + 1)) {
    printf("i = %d\n", i_1);
  } 
  #pragma unroll
  for (int32_t i_2 = 0; (i_2 < 2); i_2 = (i_2 + 1)) {
    for (int32_t j = 0; (j < 5); j = (j + 1)) {
      printf("i = %d, j = %d\n", i_2, j);
    } 
  } 
  for (int32_t w = 0; (w < 32); w = (w + 1)) {
    #pragma unroll
    for (int32_t i_3 = 0; (i_3 < 2); i_3 = (i_3 + 1)) {
      for (int32_t i_4 = 0; (i_4 < 8); i_4 = (i_4 + 1)) {
        printf("i = %d, j = %d\n", ((i_3 * 4) + (w / 8)), ((i_4 * 8) + (w % 8)));
      } 
    } 
  } 
}

@yaoyaoding yaoyaoding changed the title [IR] Add explicit unroll [WIP][IR] Add explicit unroll Apr 21, 2023
@yaoyaoding yaoyaoding changed the title [WIP][IR] Add explicit unroll [IR] Add explicit unroll Apr 21, 2023
@yaoyaoding yaoyaoding merged commit f15af4d into hidet-org:main Apr 21, 2023
@yaoyaoding yaoyaoding deleted the hidet-unroll branch April 21, 2023 15:34
AndreSlavescu pushed a commit to AndreSlavescu/hidet that referenced this pull request Apr 25, 2023
@yaoyaoding @AndreSlavescu
[IR] Add explicit unroll (hidet-org#184)
9da8a80
AndreSlavescu pushed a commit to AndreSlavescu/hidet that referenced this pull request Apr 25, 2023
@yaoyaoding @AndreSlavescu
[IR] Add explicit unroll (hidet-org#184)
d257451
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@yaoyaoding