8000 [Backend][AIE] Experimental mlir-aie codegen using MLIR builder by Fangtangtang · Pull Request #360 · cornell-zhang/allo · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

[Backend][AIE] Experimental mlir-aie codegen using MLIR builder #360

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
Merged
merged 44 commits into from
May 19, 2025
Merged

[Backend][AIE] Experimental mlir-aie codegen using MLIR builder #360

merged 44 commits into from
May 19, 2025

Conversation

Fangtangtang
Copy link
Contributor
@Fangtangtang Fangtangtang commented May 4, 2025 •
edited
Loading

Description

This PR enables MLIR-AIE codegen in Allo by setting up a compatible environment and introducing a lowering method from Allo IR to MLIR-AIE dialects via MLIR context conversion.

Problems

Previously, the environment setup led to conflicts when importing both allo and mlir-aie simultaneously. MLIR-AIE code generation in Allo relied on string matching and substitution.

Proposed Solutions

  • Configure a compatible environment using the IRON library, mlir-aie, and llvm-aie compilers from the v1.0 release wheels.
  • Manually construct codegen logic to lower Allo IR into MLIR-AIE dialects, migrating the IR from Allo’s MLIR context into the MLIR-AIE context.

Examples

Ty = int32
M = 1024

@df.region()
def top():
    @df.kernel(mapping=[1])
    def core(A: Ty[M], B: Ty[M]):
        B[:] = allo.add(A, 1)

A = np.random.randint(0, 100, M).astype(np.int32)
mod = df.build(top, target="aie-mlir")
B = np.zeros(M).astype(np.int32)
mod(A, B)
np.testing.assert_allclose(B, A + 1)
print("PASSED!")

Please refer to tests/dataflow/aie/test_vector_add.py for details.

Checklist

Please make sure to review and check all of these items:

  • PR's title starts with a category (e.g. [Bugfix], [IR], [Builder], etc)
  • All changes have test coverage (It would be good to provide ~2 different test cases to test the robustness of your code)
  • Pass the formatting check locally
  • Code is well-documented

@Fangtangtang
working on context transformation
be91b88
@Fangtangtang
[feat] setting up params for aie codegen
6eda8f1
@Fangtangtang
[feat] append aie tiles into module body
cb22567 
parse top module with external functions
@Fangtangtang
[feat] compute core wrapper
d0b188a
@Fangtangtang
[feat] resolve simple function without parameter
a8b37b1
@Fangtangtang
[ckpt] prepare for refraction
c1faae9
@Fangtangtang
[refactor] modify code structure
eb103d9
@Fangtangtang
[feat] fifo and link
cb70a79
@Fangtangtang
[feat] replace arguments with fifo
37fda1e
@Fangtangtang
[feat] replace memref.alloc with buffer
dd4361c 
TODO:
- runtime sequence
- fix mismatched copy shape
@Fangtangtang
[fix] mismatched arg type
1aab7b0
@Fangtangtang
[feat] mlir module reconstruction under aie context
bdbf8ee
@Fangtangtang
[feat] external kernels
f9be996
@Fangtangtang
[feat] generate host code and call module.
4f25c53
@Fangtangtang
[fix] incompatible host code
1b551e9 
- skip using external kernels since compilation using aiecc would crash
@Fangtangtang
[testcase] vector add
< 8000 /div>
0f10979
@Fangtangtang
[docs] add setup and patch for MLIR-AIE codegen
b31be4b
@chhzh123
Copy link
Member
chhzh123 commented May 4, 2025

@Fangtangtang Thanks for contributing! Could you follow this guide to format your code?

@chhzh123
Copy link
Member
chhzh123 commented May 4, 2025 •
edited
Loading

You can just add #pylint disable=import-error at the top of the file

@chhzh123 chhzh123 changed the title [Backend][AIE] mlir-aie codegen [Backend][AIE] Experimental mlir-aie codegen using MLIR builder May 4, 2025
< 8000 /div>
@chhzh123
Copy link
Member
chhzh123 commented May 5, 2025

Let me know when you finish the changes

@Fangtangtang
Copy link
Contributor Author

I’ve finished all the non-streaming cases except for the one using bf16.

@Fangtangtang
Copy link
Contributor Author

I’ve finished all the non-streaming cases except for the one using bf16.

fixed

@chhzh123 chhzh123 requested a review from EthanMeng324 May 8, 2025 17:35
@chhzh123
Copy link
Member
chhzh123 commented May 8, 2025

@EthanMeng324 Could you help review this PR first?

@EthanMeng324
Copy link
Contributor

@EthanMeng324 Could you help review this PR first?

Sure.

EthanMeng324
EthanMeng324 reviewed May 9, 2025
View reviewed changes
Copy link
Contributor
@EthanMeng324 EthanMeng324 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for contributing! The code is very well-organized. Just some comments below.

@Fangtangtang
add comments and adjust build flow in AIE_MLIRModule
7ad7a5e
@Fangtangtang
[refactor]: remove experimental test folder and enable codegen select…
cc8b610 
…ion via environment variable

- Removed the `tests/dataflow/aie/experimental` folder based on review feedback
- Added support for selecting between `aie` and `aie-mlir` codegen via the `USE_AIE_MLIR_BUILDER` environment variable
@Fangtangtang
format
c2adf0e
@Fangtangtang
Copy link
Contributor Author

@EthanMeng324 Thanks for the feedback! I've made the changes accordingly — please take a look and let me know if it now meets your expectations.

@EthanMeng324
Copy link
Contributor

@EthanMeng324 Thanks for the feedback! I've made the changes accordingly — please take a look and let me know if it now meets your expectations.

LGTM, thanks! @chhzh123

chhzh123
chhzh123 reviewed May 16, 2025
View reviewed changes
Copy link
Member
@chhzh123 chhzh123 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for migrating the AIE codegen. This PR looks pretty nice! Just some clarification questions listed below

@Fangtangtang
Copy link
Contributor Author

Thanks for migrating the AIE codegen. This PR looks pretty nice! Just some clarification questions listed below

Thanks! I've made the updates based on your suggestions.
Let me know if there's anything else that needs to be changed.

@Fangtangtang
Copy link
Contributor Author

@chhzh123 I've merged AIE_DTensor and DTensor as discussed. Please let me know if anything else needs to be updated.

chhzh123
chhzh123 approved these changes May 19, 2025
View reviewed changes
Copy link
Member
@chhzh123 chhzh123 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thanks!

@chhzh123 chhzh123 merged commit 97ec69f into cornell-zhang:main May 19, 2025
1 check passed
@Fangtangtang Fangtangtang mentioned this pull request May 19, 2025
Merged
4 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@EthanMeng324 EthanMeng324 EthanMeng324 left review comments

@chhzh123 chhzh123 chhzh123 approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@Fangtangtang @chhzh123 @EthanMeng324
0