microsoft / triton-shared

Shared Middle-Layer for Triton Compilation
MIT License
134 stars 26 forks source link

Introduce TritonStructured dialect and triton-to-structured pass #82

Closed haishanzzzz closed 5 months ago

haishanzzzz commented 6 months ago

This PR introduces TritonStructured dialect and triton-to-structured pass. Please see #81 for background.

This PR is broken into 5 commits:

  1. Introduce TritonStructured dialect, which includes three ops: tts.make_tptr, tts.load, and tts.store
  2. Update the current MaskAnalysis, which removed unused functions that produces different flavors of subviews
  3. Minor updates to OpFoldResultsUtils
  4. Introduce triton-to-structured pass
  5. Adds LIT tests

A few notes:

A main follow up to this PR is to get consensus on how we'd like to handle block pointers in this pass.

There are a few other misc. items on the wishlist for the new code:

nhat-nguyen commented 6 months ago

Thank you @haishanzzzz, I'll start taking a look :)

haishanzzzz commented 6 months ago

@haishanzzzz please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.

@microsoft-github-policy-service agree [company="{your company}"]

Options:

  • (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.
@microsoft-github-policy-service agree
  • (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.
@microsoft-github-policy-service agree company="Microsoft"

Contributor License Agreement

@microsoft-github-policy-service agree company="Meta"

manbearian commented 6 months ago

@haishanzzzz please , can you tell me more about this:

Investigate why tt.reduce make --remove-dead-values pass fail

haishanzzzz commented 6 months ago

@haishanzzzz please , can you tell me more about this:

Investigate why tt.reduce make --remove-dead-values pass fail

This is technically not related to this PR, but if I just do triton-opt --remove-dead-values with tt.reduce things actually fail.

triton-opt --remove-dead-values reducemax_32_256_bf16.mlir                                                                                                                                                                    16:57:35
reducemax_32_256_bf16.mlir:35:7: error: null operand found
      tt.reduce.return %22 : bf16
      ^
reducemax_32_256_bf16.mlir:35:7: note: see current operation: "tt.reduce.return"(<<NULL VALUE>>) : (<<NULL TYPE>>) -> ()

Is this what you see too?

nhat-nguyen commented 5 months ago

@haishanzzzz I reviewed the changes and left some comments, overall, I think the approach is good and makes the code much simpler. Love it.

nhat-nguyen commented 5 months ago

@haishanzzzz Would you mind updating the description with the decision around keeping tt.addptr for scalars after this pass for structured-to-memref to process later?

haishanzzzz commented 5 months ago

@haishanzzzz Would you mind updating the description with the decision around keeping tt.addptr for scalars after this pass for structured-to-memref to process later?

Updated the description to include more info on relying on tt.addptr for scalar pointers, and that we should expect to see tt.addptr at the output of the pass if analysis fails or for scalar pointers.

haishanzzzz commented 5 months ago

@nhat-nguyen @manbearian Please let me know if there are other things I can do before closing this PR.