Skip to content

Context

context

FuncInfo

Bases: NamedTuple

Information about a decorated function being generated into IR.

Source code in xdsl/frontend/pyast/context.py 
23
24
25
26
27
28
29
30
31
32
33
class FuncInfo(NamedTuple):
    """Information about a decorated function being generated into IR."""

    file: str
    """The path of the file containing the function."""

    globals: dict[str, Any]
    """The globals defined in that file up to the point of function definition."""

    ast: ast.FunctionDef
    """The Python AST representation of the function."""

file: str instance-attribute

The path of the file containing the function.

globals: dict[str, Any] instance-attribute

The globals defined in that file up to the point of function definition.

ast: ast.FunctionDef instance-attribute

The Python AST representation of the function.

PyASTContext dataclass

Encapsulate the mapping between Python and IR types and operations.

Source code in xdsl/frontend/pyast/context.py 
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
@dataclass
class PyASTContext:
    """Encapsulate the mapping between Python and IR types and operations."""

    type_registry: TypeRegistry = field(default_factory=TypeRegistry)
    """Mappings between source code and IR type."""

    function_registry: FunctionRegistry = field(default_factory=FunctionRegistry)
    """Mappings between functions and their operation types."""

    literal_registry: LiteralRegistry = field(default_factory=LiteralRegistry)
    """Mappings between literal types and their operation constructors."""

    post_transforms: list[ModulePass] = field(
        default_factory=lambda: [FrontendDesymrefyPass()]
    )
    """An ordered list of passes to apply to the built module."""

    post_callback: PassPipelineCallbackType | None = default_pipeline_callback
    """Callback to run between post transforms."""

    ir_context: Context = field(
        default_factory=lambda: Context(allow_unregistered=True)
    )
    """The xDSL context to use when applying transformations to the built module."""

    def register_type(
        self,
        source_type: type,
        ir_type: TypeAttribute,
    ) -> None:
        """Associate a type in the source code with its type in the IR."""
        self.type_registry.insert(source_type, ir_type)

    def register_function(
        self, function: Callable[..., Any], ir_constructor: Callable[..., Operation]
    ) -> None:
        """Associate a method on an object in the source code with its IR implementation."""
        self.function_registry.insert(function, ir_constructor)

    def register_literal(
        self, value_type: type[object], ir_constructor: Callable[[Any], Operation]
    ) -> None:
        """Associate a Python literal type with an IR constructor."""
        self.literal_registry.insert(value_type, ir_constructor)

    def register_post_transform(self, transform: ModulePass) -> None:
        """Add a module pass to be run on the generated IR."""
        self.post_transforms.append(transform)

    def register_dialect(self, dialect: Dialect) -> None:
        """Add a dialect to the context used for transformation."""
        self.ir_context.load_dialect(dialect)

    @property
    def pass_pipeline(self) -> PassPipeline:
        """Get a pass pipeline from the context state."""
        return PassPipeline(tuple(self.post_transforms), self.post_callback)

    @classmethod
    def _get_func_info(cls, func: Callable[P, R]) -> FuncInfo:
        """Get information about the decorated function."""
        func_file = func.__code__.co_filename
        func_globals = func.__globals__

        # Remove leading indentation from the source code to avoid parsing errors
        source = getsource(func.__code__)
        source = textwrap.dedent(source)

        # Retrieve the AST for the function body, without the decorator
        func_ast = ast.parse(source).body[0]
        assert isinstance(func_ast, ast.FunctionDef)
        assert func_ast.name == func.__name__
        assert len(func_ast.decorator_list) == 1
        func_ast.decorator_list = []

        # Return the information about the function
        return FuncInfo(func_file, func_globals, func_ast)

    @classmethod
    def _get_wrapped_program(
        cls, func: Callable[P, R], builder: PyASTBuilder
    ) -> PyASTProgram[P, R]:
        """Return a PyAST program for this function with the builder."""
        program = PyASTProgram[P, R](
            name=func.__name__,
            func=func,
            _builder=builder,
        )
        functools.update_wrapper(program, func)
        assert program.__doc__ == func.__doc__
        return program

    def parse_program(self, func: Callable[P, R]) -> PyASTProgram[P, R]:
        """Get a program wrapper by decorating a function."""
        func_file, func_globals, func_ast = self._get_func_info(func)
        builder = PyASTBuilder(
            type_registry=self.type_registry,
            function_registry=self.function_registry,
            literal_registry=self.literal_registry,
            file=func_file,
            globals=func_globals,
            function_ast=func_ast,
            build_context=self.ir_context,
            post_transforms=self.pass_pipeline,
        )
        return self._get_wrapped_program(func, builder)

type_registry: TypeRegistry = field(default_factory=TypeRegistry) class-attribute instance-attribute

Mappings between source code and IR type.

function_registry: FunctionRegistry = field(default_factory=FunctionRegistry) class-attribute instance-attribute

Mappings between functions and their operation types.

literal_registry: LiteralRegistry = field(default_factory=LiteralRegistry) class-attribute instance-attribute

Mappings between literal types and their operation constructors.

post_transforms: list[ModulePass] = field(default_factory=lambda: [FrontendDesymrefyPass()]) class-attribute instance-attribute

An ordered list of passes to apply to the built module.

post_callback: PassPipelineCallbackType | None = default_pipeline_callback class-attribute instance-attribute

Callback to run between post transforms.

ir_context: Context = field(default_factory=lambda: Context(allow_unregistered=True)) class-attribute instance-attribute

The xDSL context to use when applying transformations to the built module.

pass_pipeline: PassPipeline property

Get a pass pipeline from the context state.

__init__(type_registry: TypeRegistry = TypeRegistry(), function_registry: FunctionRegistry = FunctionRegistry(), literal_registry: LiteralRegistry = LiteralRegistry(), post_transforms: list[ModulePass] = (lambda: [FrontendDesymrefyPass()])(), post_callback: PassPipelineCallbackType | None = default_pipeline_callback, ir_context: Context = (lambda: Context(allow_unregistered=True))()) -> None

register_type(source_type: type, ir_type: TypeAttribute) -> None

Associate a type in the source code with its type in the IR.

Source code in xdsl/frontend/pyast/context.py 
70
71
72
73
74
75
76
def register_type(
    self,
    source_type: type,
    ir_type: TypeAttribute,
) -> None:
    """Associate a type in the source code with its type in the IR."""
    self.type_registry.insert(source_type, ir_type)

register_function(function: Callable[..., Any], ir_constructor: Callable[..., Operation]) -> None

Associate a method on an object in the source code with its IR implementation.

Source code in xdsl/frontend/pyast/context.py 
78
79
80
81
82
def register_function(
    self, function: Callable[..., Any], ir_constructor: Callable[..., Operation]
) -> None:
    """Associate a method on an object in the source code with its IR implementation."""
    self.function_registry.insert(function, ir_constructor)

register_literal(value_type: type[object], ir_constructor: Callable[[Any], Operation]) -> None

Associate a Python literal type with an IR constructor.

Source code in xdsl/frontend/pyast/context.py 
84
85
86
87
88
def register_literal(
    self, value_type: type[object], ir_constructor: Callable[[Any], Operation]
) -> None:
    """Associate a Python literal type with an IR constructor."""
    self.literal_registry.insert(value_type, ir_constructor)

register_post_transform(transform: ModulePass) -> None

Add a module pass to be run on the generated IR.

Source code in xdsl/frontend/pyast/context.py 
90
91
92
def register_post_transform(self, transform: ModulePass) -> None:
    """Add a module pass to be run on the generated IR."""
    self.post_transforms.append(transform)

register_dialect(dialect: Dialect) -> None

Add a dialect to the context used for transformation.

Source code in xdsl/frontend/pyast/context.py 
94
95
96
def register_dialect(self, dialect: Dialect) -> None:
    """Add a dialect to the context used for transformation."""
    self.ir_context.load_dialect(dialect)

parse_program(func: Callable[P, R]) -> PyASTProgram[P, R]

Get a program wrapper by decorating a function.

Source code in xdsl/frontend/pyast/context.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def parse_program(self, func: Callable[P, R]) -> PyASTProgram[P, R]:
    """Get a program wrapper by decorating a function."""
    func_file, func_globals, func_ast = self._get_func_info(func)
    builder = PyASTBuilder(
        type_registry=self.type_registry,
        function_registry=self.function_registry,
        literal_registry=self.literal_registry,
        file=func_file,
        globals=func_globals,
        function_ast=func_ast,
        build_context=self.ir_context,
        post_transforms=self.pass_pipeline,
    )
    return self._get_wrapped_program(func, builder)

default_pipeline_callback(previous_pass: ModulePass | None, module: ModuleOp, next_pass: ModulePass | None) -> None

Default callback to verify the module after each transformation pass.

Source code in xdsl/frontend/pyast/context.py 
36
37
38
39
40
41
def default_pipeline_callback(
    previous_pass: ModulePass | None, module: ModuleOp, next_pass: ModulePass | None
) -> None:
    """Default callback to verify the module after each transformation pass."""
    if previous_pass and next_pass:
        module.verify()