mplang-nightly 0.1.dev269__py3-none-any.whl → 0.1.dev271__py3-none-any.whl

mplang/v1/core/expr/printer.py

@@ -1,285 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Expression printer for debugging and visualization.
-"""
-
-from __future__ import annotations
-
-from typing import Any
-
-from mplang.v1.core.dtypes import DType
-from mplang.v1.core.expr.ast import (
-    AccessExpr,
-    CallExpr,
-    CondExpr,
-    ConvExpr,
-    EvalExpr,
-    Expr,
-    FuncDefExpr,
-    ShflExpr,
-    ShflSExpr,
-    TupleExpr,
-    VariableExpr,
-    WhileExpr,
-)
-from mplang.v1.core.expr.visitor import ExprVisitor
-from mplang.v1.core.mptype import MPType
-from mplang.v1.core.pfunc import PFunction
-from mplang.v1.core.tensor import Shape, TensorType
-
-
-class Printer(ExprVisitor):
-    """Printer that prints Expression DAG in IR style."""
-
-    def __init__(
-        self,
-        indent_size: int = 2,
-        compact_format: bool = True,
-        *,
-        verbose_peval: bool = False,
-        inline_pcall: bool = True,
-    ):
-        super().__init__()  # Initialize MemorizedVisitor
-        self.indent_size = indent_size
-        self.compact_format = compact_format
-        self.verbose_peval = verbose_peval
-        self.inline_pcall = inline_pcall
-        self._cur_indent = 0
-        self._output: list[str] = []
-        self._visited: dict[Expr, str] = {}
-        self._counter = 0
-
-    def _write(self, text: str) -> None:
-        """Write a line with current indentation."""
-        indent = " " * (self._cur_indent * self.indent_size)
-        for line in text.split("\n"):
-            self._output.append(f"{indent}{line}")
-
-    def _do_print(
-        self,
-        op_name: str,
-        op_args: list[str],
-        attrs: dict | None = None,
-        regions: dict[str, FuncDefExpr] | None = None,
-        mptypes: list | None = None,
-    ) -> str:
-        """A generic node printer that prints in the MLIR style."""
-        ret_name = f"%{self._counter}"
-        self._counter += 1
-
-        args_str = f"({', '.join(op_args)})"
-        attrs_str = ""
-        if attrs:
-            attr_parts = [f"{k}={v}" for k, v in attrs.items()]
-            attrs_str = f" {{{', '.join(attr_parts)}}}"
-
-        regions_str = ""
-        if regions:
-            regions_str = " {\n"
-            indent = " " * self.indent_size
-            for r_name, func_def_expr in regions.items():
-                body_printer = Printer(
-                    indent_size=self.indent_size,
-                    compact_format=self.compact_format,
-                    inline_pcall=self.inline_pcall,
-                )
-                func_def_expr.accept(body_printer)
-                regions_str += f"{indent}{r_name}: "
-                body_content = ("\n" + indent).join(body_printer._output)
-                regions_str += f"{body_content}\n"
-            regions_str += "}"
-
-        type_str = ""
-        if mptypes:
-            type_parts = [str(mptype) for mptype in mptypes]
-            if len(type_parts) == 1:
-                type_str = f" : {type_parts[0]}"
-            else:
-                type_str = f" : ({', '.join(type_parts)})"
-
-        self._write(
-            f"{ret_name} = {op_name}{args_str}{attrs_str}{regions_str}{type_str}"
-        )
-        return ret_name
-
-    def _var_name(self, expr: Expr) -> str:
-        key = expr
-        if key not in self._visited:
-            self._visited[key] = expr.accept(self)
-        return self._visited[key]
-
-    def print_expr(self, expr: Expr) -> str:
-        """Print an expression and return the formatted string."""
-        self._output = []
-        self._visited = {}
-        self._cache: dict[str, Any] = {}  # Reset memorized visitor cache
-        self._counter = 0
-        expr.accept(self)
-        return "\n".join(self._output)
-
-    def _get_const_data(self, dtype: DType, shape: Shape, data_bytes: bytes) -> str:
-        # Get dtype and shape from output info (following party.py implementation)
-        import numpy as np
-
-        np_array = np.frombuffer(data_bytes, dtype=dtype.to_numpy()).reshape(shape)
-
-        # Format the display based on array size
-        if np_array.size <= 10:
-            # Small arrays - show full content
-            if np_array.size == 1:
-                # Scalar value
-                value_str = str(np_array.item())
-            else:
-                value_str = str(np_array.tolist())
-        else:
-            # Large arrays - use numpy's default string representation which handles truncation
-            value_str = str(np_array)
-        return value_str
-
-    def _print_const(self, pfunc: PFunction, mptypes: list[MPType]) -> str:
-        assert len(pfunc.outs_info) == 1
-        out_type = pfunc.outs_info[0]
-        assert isinstance(out_type, TensorType)
-        attrs = {
-            "data": self._get_const_data(
-                out_type.dtype, out_type.shape, pfunc.attrs["data_bytes"]
-            )
-        }
-        return self._do_print("pconst", [], attrs=attrs, mptypes=mptypes)
-
-    def visit_eval(self, expr: EvalExpr) -> str:
-        arg_names = [self._var_name(arg) for arg in expr.args]
-        fn_type = expr.pfunc.fn_type
-
-        # for well known basic functions
-        if fn_type == "basic.constant":
-            return self._print_const(expr.pfunc, expr.mptypes)
-
-        attrs = {"fn_type": fn_type}
-        if expr.pfunc.fn_name:
-            attrs["fn_name"] = str(expr.pfunc.fn_name)
-        if self.verbose_peval:
-            attrs["fn_text"] = str(expr.pfunc.fn_text)
-
-        if expr.rmask is not None:
-            attrs["rmask"] = f"0x{expr.rmask.value:x}"
-        return self._do_print("peval", arg_names, attrs=attrs, mptypes=expr.mptypes)
-
-    def visit_variable(self, expr: VariableExpr) -> str:
-        if self.compact_format:
-            # Use $param format and don't print the variable definition
-            return f"{expr.name}"
-        else:
-            return self._do_print(
-                "pname", [f'"{expr.name}"'], attrs={}, mptypes=expr.mptypes
-            )
-
-    def visit_tuple(self, expr: TupleExpr) -> str:
-        arg_names = [self._var_name(arg) for arg in expr.args]
-        return self._do_print("tuple", arg_names, mptypes=expr.mptypes)
-
-    def visit_cond(self, expr: CondExpr) -> str:
-        pred_name = self._var_name(expr.pred)
-        arg_names = [self._var_name(arg) for arg in expr.args]
-
-        # Directly pass FuncDefExpr objects
-        return self._do_print(
-            "pcond",
-            [pred_name, *arg_names],
-            regions={
-                "then_fn": expr.then_fn,
-                "else_fn": expr.else_fn,
-            },
-            mptypes=expr.mptypes,
-        )
-
-    def visit_call(self, expr: CallExpr) -> str:
-        arg_names = [self._var_name(arg) for arg in expr.args]
-        if self.inline_pcall:
-            return self._do_print(
-                expr.name,
-                arg_names,
-                mptypes=expr.mptypes,
-            )
-        else:
-            return self._do_print(
-                "pcall",
-                arg_names,
-                regions={"fn": expr.fn},
-                mptypes=expr.mptypes,
-            )
-
-    def visit_while(self, expr: WhileExpr) -> str:
-        arg_names = [self._var_name(arg) for arg in expr.args]
-
-        return self._do_print(
-            "pwhile",
-            arg_names,
-            regions={
-                "cond_fn": expr.cond_fn,
-                "body_fn": expr.body_fn,
-            },
-            mptypes=expr.mptypes,
-        )
-
-    def visit_conv(self, expr: ConvExpr) -> str:
-        var_names = [self._var_name(var) for var in expr.vars]
-        return self._do_print("pconv", var_names, mptypes=expr.mptypes)
-
-    def visit_shfl_s(self, expr: ShflSExpr) -> str:
-        src_val_name = self._var_name(expr.src_val)
-        attrs = {"pmask": expr.pmask, "src_ranks": expr.src_ranks}
-        return self._do_print(
-            "pshfl_s", [src_val_name], attrs=attrs, mptypes=expr.mptypes
-        )
-
-    def visit_shfl(self, expr: ShflExpr) -> str:
-        src_name = self._var_name(expr.src)
-        index_name = self._var_name(expr.index)
-        return self._do_print("pshfl", [src_name, index_name], mptypes=expr.mptypes)
-
-    def visit_access(self, expr: AccessExpr) -> str:
-        expr_name = self._var_name(expr.src)
-        if self.compact_format:
-            # Original:
-            #   %x = ...
-            #   %y = %x[0]
-            #   %z = some_fn(%y)
-            # Single output(optimized):
-            #   %x = ...
-            #   %z = some_fn(%x)
-            # Multiple outputs (optimized):
-            #   %x = ...
-            #   %z = some_fn(%x:0, %x:1)
-            if len(expr.src.mptypes) > 1:
-                return f"{expr_name}:{expr.index}"
-            else:
-                return expr_name
-        else:
-            attrs = {"index": str(expr.index)}
-            return self._do_print(
-                "access", [expr_name], attrs=attrs, mptypes=expr.mptypes
-            )
-
-    def visit_func_def(self, expr: FuncDefExpr) -> str:
-        param_names = expr.params
-        self._write(f"({', '.join(param_names)}) {{")
-        self._cur_indent += 1
-        body_name = expr.body.accept(self)
-        self._write(f"return {body_name}")
-        self._cur_indent -= 1
-        self._write("}")
-        return ""

mplang/v1/core/expr/transformer.py

@@ -1,141 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Expression transformer based on visitor pattern.
-"""
-
-from collections.abc import Callable
-
-from mplang.v1.core.expr.ast import (
-    AccessExpr,
-    CallExpr,
-    CondExpr,
-    ConvExpr,
-    EvalExpr,
-    Expr,
-    FuncDefExpr,
-    ShflExpr,
-    ShflSExpr,
-    TupleExpr,
-    VariableExpr,
-    WhileExpr,
-)
-from mplang.v1.core.expr.visitor import ExprVisitor
-
-
-class ExprTransformer(ExprVisitor):
-    """Transformer that transforms expressions based on rules."""
-
-    def __init__(self, trans_rules: dict[str, Callable[[Expr], Expr]] | None = None):
-        self.trans_rules = trans_rules or {}
-
-    def visit_eval(self, expr: EvalExpr) -> Expr:
-        # Transform child expressions first
-        transformed_args = [arg.accept(self) for arg in expr.args]
-        new_expr = EvalExpr(expr.pfunc, transformed_args, expr.rmask)
-
-        if "eval" in self.trans_rules:
-            return self.trans_rules["eval"](new_expr)
-        return new_expr
-
-    def visit_variable(self, expr: VariableExpr) -> Expr:
-        if "name" in self.trans_rules:
-            return self.trans_rules["name"](expr)
-        return expr
-
-    def visit_tuple(self, expr: TupleExpr) -> Expr:
-        # Transform child expressions first
-        transformed_args = [arg.accept(self) for arg in expr.args]
-        new_expr = TupleExpr(transformed_args)
-
-        if "tuple" in self.trans_rules:
-            return self.trans_rules["tuple"](new_expr)
-        return new_expr
-
-    def visit_cond(self, expr: CondExpr) -> Expr:
-        # Transform child expressions first
-        transformed_pred = expr.pred.accept(self)
-        transformed_args = [arg.accept(self) for arg in expr.args]
-        new_expr = CondExpr(
-            transformed_pred, expr.then_fn, expr.else_fn, transformed_args
-        )
-
-        if "cond" in self.trans_rules:
-            return self.trans_rules["cond"](new_expr)
-        return new_expr
-
-    def visit_call(self, expr: CallExpr) -> Expr:
-        # Transform child expressions first
-        transformed_args = [arg.accept(self) for arg in expr.args]
-        new_expr = CallExpr(expr.name, expr.fn, transformed_args)
-
-        if "call" in self.trans_rules:
-            return self.trans_rules["call"](new_expr)
-        return new_expr
-
-    def visit_while(self, expr: WhileExpr) -> Expr:
-        # Transform all arguments
-        transformed_args = [arg.accept(self) for arg in expr.args]
-        new_expr = WhileExpr(expr.cond_fn, expr.body_fn, transformed_args)
-
-        if "while" in self.trans_rules:
-            return self.trans_rules["while"](new_expr)
-        return new_expr
-
-    def visit_conv(self, expr: ConvExpr) -> Expr:
-        # Transform child expressions first
-        transformed_vars = [var.accept(self) for var in expr.vars]
-        new_expr = ConvExpr(transformed_vars)
-
-        if "conv" in self.trans_rules:
-            return self.trans_rules["conv"](new_expr)
-        return new_expr
-
-    def visit_shfl_s(self, expr: ShflSExpr) -> Expr:
-        # Transform child expression first
-        transformed_src_val = expr.src_val.accept(self)
-        new_expr = ShflSExpr(transformed_src_val, expr.pmask, expr.src_ranks)
-
-        if "shfl_s" in self.trans_rules:
-            return self.trans_rules["shfl_s"](new_expr)
-        return new_expr
-
-    def visit_shfl(self, expr: ShflExpr) -> Expr:
-        # Transform child expressions first
-        transformed_src = expr.src.accept(self)
-        transformed_index = expr.index.accept(self)
-        new_expr = ShflExpr(transformed_src, transformed_index)
-
-        if "shfl" in self.trans_rules:
-            return self.trans_rules["shfl"](new_expr)
-        return new_expr
-
-    def visit_access(self, expr: AccessExpr) -> Expr:
-        # Transform child expression first
-        transformed_expr = expr.src.accept(self)
-        new_expr = AccessExpr(transformed_expr, expr.index)
-
-        if "access" in self.trans_rules:
-            return self.trans_rules["access"](new_expr)
-        return new_expr
-
-    def visit_func_def(self, expr: FuncDefExpr) -> Expr:
-        # Transform body only, params are just strings now
-        transformed_body = expr.body.accept(self)
-        new_expr = FuncDefExpr(expr.params, transformed_body)
-
-        if "func_def" in self.trans_rules:
-            return self.trans_rules["func_def"](new_expr)
-        return new_expr

mplang/v1/core/expr/utils.py

@@ -1,78 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Utility functions for expression system.
-"""
-
-from collections.abc import Sequence
-
-from mplang.v1.core.mask import Mask
-from mplang.v1.core.mptype import TensorLike
-
-
-def type_equal(*args: TensorLike) -> bool:
-    """Check if tensors have identical type properties (dtype, shape).
-
-    Args:
-        *args: Variable number of TensorLike objects to compare
-
-    Returns:
-        bool: True if all tensors have identical types, False otherwise
-    """
-    if len(args) <= 1:
-        return True
-    for i in range(1, len(args)):
-        if args[0].dtype != args[i].dtype or args[0].shape != args[i].shape:
-            return False
-    return True
-
-
-def ensure_scalar(obj: TensorLike) -> None:
-    """Ensure that a tensor is a scalar."""
-    if len(obj.shape) != 0:
-        raise TypeError(f"Expected a scalar, got {obj}.")
-
-
-def ensure_tensorlist_equal(*args: Sequence[TensorLike]) -> None:
-    """Ensure that multiple tensor lists have the same structure and types."""
-    if len(args) < 2:
-        raise ValueError(f"expect at least 2 args, got {len(args)}")
-    for i in range(1, len(args)):
-        if len(args[i]) != len(args[0]):
-            raise ValueError(f"Length mismatch: {len(args[i])} vs {len(args[0])}")
-        for j in range(len(args[0])):
-            if not type_equal(args[0][j], args[i][j]):
-                raise TypeError(f"Type mismatch: {args[0][j]} vs {args[i][j]}")
-
-
-def deduce_mask(*pmasks: Mask | None) -> Mask | None:
-    """Deduce the joint mask from multiple participant masks."""
-    if len(pmasks) == 0:
-        return None
-
-    if any(pmask is None for pmask in pmasks):
-        # If any pmask is None, we cannot deduce a specific mask.
-        return None
-
-    # return the joint mask of all provided pmasks.
-    # We already checked above, but add it here to make mypy happy
-    if pmasks[0] is None:
-        return None
-    result = Mask(pmasks[0])
-    for pmask in pmasks[1:]:
-        assert pmask is not None  # We already checked above
-        result = result.intersection(Mask(pmask))
-
-    return result

mplang/v1/core/expr/visitor.py

@@ -1,85 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Visitor pattern interface for expression system.
-"""
-
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    from mplang.v1.core.expr.ast import (
-        AccessExpr,
-        CallExpr,
-        CondExpr,
-        ConvExpr,
-        EvalExpr,
-        FuncDefExpr,
-        ShflExpr,
-        ShflSExpr,
-        TupleExpr,
-        VariableExpr,
-        WhileExpr,
-    )
-
-
-class ExprVisitor(ABC):
-    """Base visitor interface for expression types."""
-
-    @abstractmethod
-    def visit_eval(self, expr: EvalExpr) -> Any:
-        pass
-
-    @abstractmethod
-    def visit_variable(self, expr: VariableExpr) -> Any:
-        pass
-
-    @abstractmethod
-    def visit_tuple(self, expr: TupleExpr) -> Any:
-        pass
-
-    @abstractmethod
-    def visit_cond(self, expr: CondExpr) -> Any:
-        pass
-
-    @abstractmethod
-    def visit_call(self, expr: CallExpr) -> Any:
-        pass
-
-    @abstractmethod
-    def visit_while(self, expr: WhileExpr) -> Any:
-        pass
-
-    @abstractmethod
-    def visit_conv(self, expr: ConvExpr) -> Any:
-        pass
-
-    @abstractmethod
-    def visit_shfl_s(self, expr: ShflSExpr) -> Any:
-        pass
-
-    @abstractmethod
-    def visit_shfl(self, expr: ShflExpr) -> Any:
-        pass
-
-    @abstractmethod
-    def visit_access(self, expr: AccessExpr) -> Any:
-        pass
-
-    @abstractmethod
-    def visit_func_def(self, expr: FuncDefExpr) -> Any:
-        pass