POC for Buffer Descriptor Analysis

[jamestclin]

This draft shows an implementation of an analysis for generating Buffer Descriptors.

• Problem formulation:
  • Input:
    • LinAlg IR converted from Triton IR by --triton-to-linalg conversion pass
    • Make dynamic strides/sizes from function argument to constant strides/sizes in arith.constant (task 3.2.1., this is done manually at the moment)
  • Output:
    • Buffer descriptor with constant length/steps/wraps, and dynamic base (base pointer address, a function of program_id)
• Proposed method in high-level:
  • Analysis is done by visiting all the memref.copy/memref.tensor_store in the function and recursively call visitor functions for operands until reaching memref.reinterpretcast.
  • Buffer descriptor is propagated from memref.reinterpretcast to memref.copy/memref.tensor_store by recursion of visitor functions.
• Able to handle tasks 3.2.2. and 3.2.3. for triton tutorials:
  • add vector, softmax, and layernorm
  • matmul (works for various group sizes)
  • fused attention (only works for full attention that comes with constant for-loop tripcount. The original tutorial example is masked attention where the tripcount is dynamic to the program_id)
• Known limitation:
  • Memory block size masking/padding is skipped
  • Multiple load of a tensor in a single for loop iteration isn't supported
  • For-loop has to have constant tripcount to infer the BD for task 3.2.3.
  • Haven't looked into the case of nested for-loop

[stephenneuendorffer]

Advanced Micro Devices, Inc.

[jamestclin]

Will update the name here. Also, git checks show errors regarding clang formatting, and compilation warning. Will fix them as well.

[stephenneuendorffer]

What does this do? Is it really specific to your target, or is it more general?

[jamestclin]

It's more to the target. Any suggestion on where this header file should live?

[stephenneuendorffer]

@jamestclin What's the status with this?

[jamestclin]

Close this PR for now, as the POC isn't in the current pursuit.