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.

Sign up for GitHub

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

Jump to bottom

Refactor softmax templates to use outer dims #844

Closed
wants to merge 2 commits into from
Closed

Refactor softmax templates to use outer dims #844

wants to merge 2 commits into from

Commits on Jul 25, 2023

  1. Slightly improve softmax's codegen formatting 

    Summary:
Very minor cleanups I did while familiarizing myself with the code

Aside from whitespace changes I also removed a few unnecessary automatic variables

Differential Revision: D47732846

fbshipit-source-id: 88635358cf795f21ce1c582dea85185db6c5ac07
    @int3 @facebook-github-bot
    int3 authored and facebook-github-bot committed Jul 25, 2023
    Configuration menu
    Copy the full SHA
    09c988a View commit details
    Browse the repository at this point in the history

  2. Refactor softmax templates to use outer dims (facebookincubator#844) 

    Summary:
Pull Request resolved: facebookincubator#844

Previously, the softmax templates assumed that reduction would always be done over the last dim, so the only parameter passed to the templates was the rank of the tensor. To set the stage for generalizing softmax, we pass the reduction dim instead.

The output is functionally identical, though the codegen changes slightly in the case where all the inner dimensions are 1: we now pass only the outer dimensions to the function call, dropping the redundant inner dimension parameters.

For the `tail_shapes_all_1_bf16` softmax test case, we have

Before:
```
      softmax_0(
         X,
         Y,
         &input_batch,
         &X_dim_1,
         &X_dim_2,
         stream
      );
```

After:
```
      softmax_0(
         X,
         Y,
         &input_batch,
         stream
      );
```

Reviewed By: aakhundov

Differential Revision: D47732859

fbshipit-source-id: 9c081a21338ba5d6cb6386b06b113171611746e4
    @int3 @facebook-github-bot
    int3 authored and facebook-github-bot committed Jul 25, 2023
    Configuration menu
    Copy the full SHA
    11c86f7 View commit details
    Browse the repository at this point in the history