[Library] Add RoPE and modulate_fused operators to the nn.py library

[silvenachen]

Description

Add two new operator functions to the nn.py library for better support of diffusion transformers and LLM-style workloads.

Problems

Currently, the Allo nn.py library lacks support for:

• RoPE (Rotary Positional Embedding), which is widely used in attention layers in LLMs
• modulate_fused, a fused operator for scaling and shifting used in diffusion models (e.g., timestep conditioning)

Adding these operators could improve coverage for transformer-based and diffusion model architectures.

Proposed Solutions

New operators in /allo/library/nn.py

• Implemented RoPE operator.
• Implemented modulate_fused operator, which combines scale and shift into a single operator.
• Added accompanying schedule_modulate_fused function.

New test cases in tests/test_nn.py

• test_RoPE and test_modulate_fused for the above mentioned two operators.

Examples

(Please provide an example of the input program and the expected behavior, e.g., the generated IR.)

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

[chhzh123]

Thanks for contributing to Allo! As these two are not commonly used operations, could you add the corresponding paper reference to the operator definition? If they have PyTorch implementation, please also provide a link

[chhzh123]

Can you add the reference paper here?

[silvenachen]

Thanks so much for the suggestion! I've added a comment of ref to the original paper. You can also find more details here:
Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding” — arXiv:2104.09864

[chhzh123]

I suppose the RoPE kernel also needs the corresponding scheduling function?

[silvenachen]

I originally planned to add a schedule_RoPE function, but it seems Allo raises error when trying to access the loop labels. This is the schedule function I tried to pipeline the inner loop j:

def schedule_RoPE(s):
    lj1 = s.get_loops(s.top_func_name)["rope_split_1"]["j"]
    s.pipeline(lj1)
    lj2 = s.get_loops(s.top_func_name)["rope_split_2"]["j"]
    s.pipeline(lj2)
    lj3 = s.get_loops(s.top_func_name)["rotary_1"]["j"]
    s.pipeline(lj3)
    lj4 = s.get_loops(s.top_func_name)["rotary_2"]["j"]
    s.pipeline(lj4)
    lj5 = s.get_loops(s.top_func_name)["rotary_merge_1"]["j"]
    s.pipeline(lj5)
    lj6 = s.get_loops(s.top_func_name)["rotary_merge_2"]["j"]
    s.pipeline(lj6)
    return s

However, I get the following error message. A potential cause is that the main body of the RoPE operator is nested inside a standard Python for loop, so the nested dsl.grid() loops are not visible to s.get_loops() at the top level. This seems similar to the structure of scaled_dot_product_attention, which I suppose also doesn't have a schedule_* function for the same reason.

(allo) root@2a2f74889c68:~/allo/allo/library# python3 tests.py 
Traceback (most recent call last):
  File "/root/allo/allo/library/tests.py", line 66, in <module>
    test_RoPE()
  File "/root/allo/allo/library/tests.py", line 26, in test_RoPE
    schedule_RoPE(s)
  File "/root/allo/allo/library/nn.py", line 231, in schedule_RoPE
    lj1 = s.get_loops(s.top_func_name)["rope_split_1"]["j"]
          ~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^
  File "/root/allo/allo/ir/transform.py", line 66, in __getitem__
    raise AttributeError(f"No such loop {name}")
AttributeError: No such loop rope_split_1

Could you please kindly tell me how to add a schedule in this case, or it's fine to leave this function unscheduled for now?

[chhzh123]

I see. Can you add a name for the h loop so that you can access h first? Also, you can print out the loops by print(s.get_loops(s.top_func_name)), so you can know what kinds of loops are accesible