[WIP] feat(ir): Add ub_copy operation and fix memory reuse for view ops

[YunjiQin]

Summary

• Add block.ub_copy operation for UB-to-UB memory copies with TMOV
• Fix memory reuse pass to handle view operations (reshape) correctly
• Track MemRef sharing groups to merge lifetimes for shared memory
• Move create_tile from unified_ops to block_ops namespace
• Add comprehensive tests for view operation memory reuse

[gemini-code-assist]

Summary of Changes

Hello @YunjiQin, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new ub_copy operation for direct Unified Buffer memory transfers and significantly improves the memory reuse optimization pass. It addresses a critical issue where view operations like reshape were not correctly handled, leading to inefficient memory allocation. By tracking shared memory references and merging their lifetimes, the pass can now apply memory reuse more effectively across variables that logically share the same buffer, ultimately reducing memory footprint and improving performance.

Highlights

• New ub_copy operation: Introduced block.ub_copy for efficient Unified Buffer (UB) to UB memory copies, leveraging the TMOV instruction.
• Memory Reuse for View Operations: Enhanced the memory reuse pass to correctly handle view operations (e.g., reshape) by ensuring they share the underlying MemRef.
• MemRef Sharing Group Tracking: Implemented tracking of MemRef sharing groups to merge lifetimes for variables that share the same memory reference, improving reuse accuracy.
• create_tile namespace change: Relocated the create_tile operation from the unified_ops namespace to block_ops for better organization.
• Comprehensive Testing: Added extensive unit tests to validate the correct memory reuse behavior for view operations.

Changelog

• python/pypto/ir/op/block_ops.py
  • Added ub_copy function for UB-to-UB tile copies.
• python/pypto/language/init.py
  • Imported create_tile and ub_copy operations.
  • Added create_tile and ub_copy to the list of promoted block-only operations.
• python/pypto/language/op/init.py
  • Imported create_tile and ub_copy operations.
  • Removed create_tile from unified_ops import.
  • Added ub_copy to the list of promoted block-only operations.
• python/pypto/language/op/block_ops.py
  • Added ub_copy function with documentation for UB-to-UB tile copies.
  • Exported ub_copy in the __all__ list.
• src/backend/910B_CCE/backend_910b_cce_ops.cpp
  • Added MakeBlockUbCopyCodegenCCE to generate TMOV instruction for block.ub_copy.
  • Introduced validation in MakeBlockMoveCodegenCCE to prevent block.move from performing UB-to-UB copies, directing users to block.ub_copy.
  • Registered the block.ub_copy backend operation.
• src/ir/op/block_ops/memory.cpp
  • Added DeduceBlockUbCopyType function to deduce the type for block.ub_copy.
  • Registered the block.ub_copy operation with its description and argument.
• src/ir/transforms/basic_memory_reuse_pass.cpp
  • Included <climits> header.
  • Modified ComputeLifetimesFromDependencies to identify and group variables sharing the same MemRef, computing a merged lifetime for the group.
  • Updated ApplyMemRefSharing to propagate MemRef reuse to all variables within a sharing group when one variable is reused.
  • Modified MemRefSharingMutator constructor to accept func_body and build sharing groups.
• src/ir/transforms/init_memref.cpp
  • Added IsViewOperation helper function to identify view operations like block.reshape.
  • Modified InitMemRefMutator to ensure that view operations share the MemRef of their input tile.
• tests/ut/ir/transforms/test_basic_memory_reuse.py
  • Added TestViewOperationsMemoryReuse class with tests for reshape operations, including single reshape, chained reshapes, and ensuring correct propagation of MemRef reuse.

Activity

• The pull request is currently marked as Work In Progress (WIP).

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a new block.ub_copy operation for efficient UB-to-UB memory copies and significantly improves the memory reuse pass to correctly handle view operations like reshape. The changes to track MemRef sharing groups and merge their lifetimes are well-implemented and crucial for correctness. The addition of comprehensive tests for view operation memory reuse is also a great inclusion. I've identified one potential issue regarding validation and one opportunity for a performance improvement in the memory reuse pass.