feat(language): Add Scalar type, decorator function type, and SSA passes

[lyfne123]

• Add pl.Scalar[dtype] type and scalar.py with parser/type_resolver support
• Extract function type from @pl.function(type=pl.FunctionType.*) for program build
• Add ConvertToSSA and VerifySSA to default optimization strategy in pass manager
• Move test_pto_codegen from ir/transforms to codegen; update orchestration example and tests

[gemini-code-assist]

Summary of Changes

Hello @lyfne123, 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 significantly enhances the PyPTO Language DSL by introducing a dedicated Scalar type for explicit scalar handling and improving the declarative programming model. It integrates Static Single Assignment (SSA) form into the default optimization pipeline, which is crucial for robust IR analysis. Additionally, the changes streamline the definition of program functions by allowing direct type specification via decorators and refactor an existing example to showcase these new capabilities, making the DSL more expressive and easier to use.

Highlights

• New Scalar Type: Introduced pl.Scalar[dtype] for representing scalar values in the PyPTO Language DSL, complete with parser and type resolver support. This allows for explicit scalar type annotations in function signatures and operations.
• Declarative Function Type Specification: Enhanced the @pl.function decorator to allow explicit specification of function types (e.g., pl.FunctionType.InCore, pl.FunctionType.Orchestration) directly in the Python DSL, improving clarity and program construction.
• SSA Passes Integration: Added ConvertToSSA and VerifySSA passes to the default optimization strategy in the IR pass manager, ensuring that the Intermediate Representation adheres to Static Single Assignment form for improved analysis and optimization.
• Orchestration Example Refactoring: The orchestration_example.py was refactored to utilize the new @pl.program and @pl.function decorators, demonstrating a more declarative and Python-native way to define programs and kernels.
• Expanded Operation Support: Updated block_ops and tensor_ops to accept the new pl.Scalar type as an operand, enabling more flexible arithmetic operations between tensors/tiles and scalars.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 several significant features to the PyPTO language and compiler infrastructure. It adds a new pl.Scalar type for representing scalar values in the DSL, complete with parser and type resolver support. It also introduces the ability to specify a function's type (e.g., InCore, Orchestration) via the @pl.function decorator, which is a great enhancement for program construction. Additionally, ConvertToSSA and VerifySSA passes have been added to the default optimization strategy, strengthening the compiler's IR processing capabilities. The orchestration example has been nicely updated to use the new language features.

My review focuses on improving code maintainability, clarity, and type safety. I've identified a few areas where the code can be refactored for better readability and to reduce duplication. I also noted a case where static type hints don't align with the DSL's runtime behavior, leading to the use of type: ignore, and suggested potential improvements. Finally, I've recommended reinstating an updated test case to ensure pass manager strategies remain correctly configured.

I am having trouble creating individual review comments. Click here to see my feedback.

examples/ir_builder/orchestration_example.py (110)

The use of # type: ignore[reportArgumentType] here highlights a discrepancy between the static type hints and the DSL's runtime behavior. The kernel_add_scalar function's signature expects a pl.Scalar[pl.FP32], but a float literal is passed. While the parser correctly handles this implicit conversion, it makes the code less type-safe from a static analysis perspective and can be confusing for developers.

To improve type safety and clarity, you could consider one of the following approaches:

1. Update the kernel_add_scalar signature to scalar: Union[pl.Scalar[pl.FP32], float] if the parser can handle this without breaking type resolution.
2. If this implicit conversion is an intended feature of the DSL, it would be beneficial to document this behavior in the kernel_add_scalar docstring to clarify why the type: ignore is necessary.

python/pypto/language/op/tensor_ops.py (97-102)

The logic to unwrap Tensor or Scalar types can be simplified. Since both types have an unwrap method, you can check for isinstance(rhs, (Tensor, Scalar)) to reduce the number of conditional branches and make the code more concise. This same improvement can be applied to the add function below.

    if isinstance(rhs, (Tensor, Scalar)):
        rhs_expr = rhs.unwrap()
    else:
        rhs_expr = rhs

python/pypto/language/parser/decorator.py (168-174)

This if/elif chain for mapping the function type name to the enum can be made more concise and maintainable by using a dictionary lookup. This avoids a growing chain of conditions if more function types are added in the future and improves readability.

                            type_name = value.attr
                            type_map = {
                                "Orchestration": ir.FunctionType.Orchestration,
                                "InCore": ir.FunctionType.InCore,
                                "Opaque": ir.FunctionType.Opaque,
                            }
                            if type_name in type_map:
                                return type_map[type_name]

tests/ut/ir/transforms/test_pass_manager.py (86-100)

The test case test_multiple_instances_different_strategies was removed. While it was outdated due to changes in the Default strategy, the test itself is valuable for ensuring that different optimization strategies are configured correctly and are distinct. It would be better to update this test with the new pass lists rather than removing it. This helps prevent regressions in pass manager configuration.

Here is an example of an updated test:

def test_multiple_instances_different_strategies(self):
    """Test creating instances of different strategies."""
    pm_default = ir.PassManager.get_strategy(ir.OptimizationStrategy.Default)
    pm_ptoa = ir.PassManager.get_strategy(ir.OptimizationStrategy.PTOAS)

    # Should have different strategies
    assert pm_default.strategy != pm_ptoa.strategy

    # Check pass counts
    assert len(pm_default.passes) == 6
    assert len(pm_ptoa.passes) == 3

    # Verify pass names are properly configured
    assert pm_default.get_pass_names() == [
        "ConvertToSSA",
        "VerifySSA",
        "InitMemRef",
        "MemoryReuse",
        "InsertSync",
        "AddAlloc",
    ]
    assert pm_ptoa.get_pass_names() == ["InitMemRef", "MemoryReuse", "AddAlloc"]