[MLIR][Dataflow] Add Allo operations for SPMW and setup support (#555)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: Fangtangtang

allo/customize.py

@@ -5,7 +5,6 @@
 import sys
 import re
 import inspect
-import textwrap
 import traceback
 import copy
 from dataclasses import dataclass
@@ -1350,15 +1349,8 @@ def customize(
             Defaults to `"default"`.
     """
     # Get Python AST
-    if isinstance(fn, str):
-        src, starting_line_no = fn, 1
-        file_name = None
-    else:
-        src, starting_line_no = inspect.getsourcelines(fn)
-        src = [textwrap.fill(line, tabsize=4, width=9999) for line in src]
-        src = textwrap.dedent("\n".join(src))
-        file_name = inspect.getfile(fn)
-    tree = parse_ast(src, starting_line_no=starting_line_no, verbose=verbose)
+    file_name = None if isinstance(fn, str) else inspect.getfile(fn)
+    tree = parse_ast(fn, verbose=verbose)
     if instantiate is None:
         instantiate = []
     if global_vars is None:

allo/ir/infer.py

@@ -5,8 +5,6 @@
 import ast
 import copy
 import os
-import inspect
-import textwrap
 import warnings
 import sympy
 import numpy as np
@@ -29,7 +27,7 @@
     float64,
     Struct,
     Stream,
-    stateful,
+    Stateful,
     ConstExpr,
 )
 from .typing_rule import get_typing_rule
@@ -144,7 +142,7 @@ def visit_type_hint(ctx: ASTContext, node: ast.AST):
             spec = ASTResolver.resolve(node.right, ctx.global_vars)
             if isinstance(spec, list):
                 spec = Layout(spec)
-            if spec is stateful:
+            if spec is Stateful:
                 # Create a copy with stateful=True
                 stateful_dtype = copy.deepcopy(dtype)
                 stateful_dtype.stateful = True
@@ -1249,12 +1247,7 @@ class ExternalModule:
         else:
             # Visit arguments in the top-level
             visit_stmts(ctx, node.args)
-            src, starting_line_no = inspect.getsourcelines(func)
-            src = [textwrap.fill(line, tabsize=4, width=9999) for line in src]
-            src = textwrap.dedent("\n".join(src))
-            tree = parse_ast(
-                src, starting_line_no=starting_line_no, verbose=ctx.verbose
-            )
+            tree = parse_ast(func, verbose=ctx.verbose)
             # Create a new context to avoid name collision
             func_ctx = ctx.copy()
             stmts = visit_stmts(func_ctx, tree.body)

allo/ir/types.py

@@ -35,7 +35,7 @@ def __init__(self, dtype, shape):
         self.shape = shape
 
     def __matmul__(self, other):
-        """Support the @ operator for memory/layout annotations."""
+        """Support the @ operator for memory/layout/stateful annotations."""
         # Return self to allow chaining, the actual spec is extracted from AST
         return self
 
@@ -91,27 +91,14 @@ def __hash__(self):
         return hash((self.name, self.stateful))
 
 
-def stateful(dtype: AlloType):
+class Stateful:
     """
-    Marks a type as stateful, making it persistent across kernel invocations.
+    Refinement type, marks a type as stateful, making it persistent across kernel invocations.
 
     Usage:
-        state: stateful(int32)           # Stateful scalar
-        counter: stateful(int32[10])     # Stateful array
-
-    Parameters
-    ----------
-    dtype : AlloType
-        The underlying type to mark as stateful
-
-    Returns
-    -------
-    tuple
-        A marker tuple that will be processed during type inference
+        acc: Int(4) @ Stateful = 0          # Stateful scalar
+        window: float32[4] @ Stateful       # Stateful array
     """
-    # TODO: Return type should be AlloType with stateful attribute set,
-    #       but currently returns tuple for AST
-    return ("stateful", dtype)
 
 
 class Index(AlloType):

allo/ir/utils.py

@@ -4,6 +4,7 @@
 
 import ast
 import inspect
+import textwrap
 from collections.abc import Callable
 from types import FunctionType as PyFunctionType
 from .._mlir.ir import (
@@ -27,19 +28,43 @@
 from .symbol_resolver import ASTResolver
 
 
-def _get_global_vars(_func):
+def _get_global_vars(_func, skip: set[str] = None, stop: set[str] = None):
+    """
+    Collect global variables from the call stack of a Python function.
+
+    Args:
+        skip: Set of frame names to skip over when walking the call stack.
+              Frames whose co_name is in `skip` are ignored (no variables collected), and the walk continues to the next outer frame.
+              This is mainly used to skip compiler internal functions when collecting global variables used in source code.
+        stop: Set of frame names that act as boundaries for the stack walk.
+              When a frame whose co_name is in `stop` is reached, its variables are collected and then the walk terminates.
+    """
+    if skip is None:
+        skip = {"get_global_vars", "customize", "build"}
+    if stop is None:
+        stop = {"<module>"}
     if isinstance(_func, Callable):
         # Discussions: https://github.com/taichi-dev/taichi/issues/282
         global_vars = _func.__globals__.copy()
     else:
         global_vars = {}
 
-    # Get back to the outer-most scope (user-defined function)
+    # Get back to outer scopes
     # Mainly used to get the annotation definitions (shape and type),
     # which are probably not defined in __globals__
-    for name, var in inspect.stack()[3][0].f_locals.items():
-        if isinstance(var, (int, float, AlloType)) or inspect.isfunction(var):
-            global_vars[name] = var
+    frame = inspect.currentframe().f_back
+    while frame:
+        if frame.f_code.co_name in skip:
+            frame = frame.f_back
+            continue
+        # collect allowed types
+        for name, var in frame.f_locals.items():
+            if isinstance(var, (int, float, AlloType)) or inspect.isfunction(var):
+                global_vars[name] = var
+        # boundary
+        if frame.f_code.co_name in stop:
+            break
+        frame = frame.f_back
 
     if isinstance(_func, Callable):
         freevar_names = _func.__code__.co_freevars
@@ -52,13 +77,25 @@ def _get_global_vars(_func):
 
 
 def get_global_vars(func):
-    global_vars = _get_global_vars(func)
-    new_global_vars = global_vars.copy()
-    for var in global_vars.values():
-        # import functions from other files
-        if isinstance(var, PyFunctionType):
-            new_global_vars.update(_get_global_vars(var))
-    return new_global_vars
+    all_globals = {}
+    worklist = [func]
+    visited_funcs = set()
+
+    while worklist:
+        f = worklist.pop()
+        if f in visited_funcs:
+            continue
+        visited_funcs.add(f)
+
+        gv = _get_global_vars(f)
+        for name, val in gv.items():
+            if name not in all_globals:
+                all_globals[name] = val
+                # import functions from other files
+                if isinstance(val, PyFunctionType):
+                    worklist.append(val)
+
+    return all_globals
 
 
 def get_extra_type_hints(dtype: AlloType):
@@ -104,7 +141,13 @@ def _adjust_line_numbers(node, offset):
             child.end_lineno += offset
 
 
-def parse_ast(src, starting_line_no=1, verbose=False):
+def parse_ast(src, verbose=False):
+    if isinstance(src, str):
+        starting_line_no = 1
+    else:
+        src, starting_line_no = inspect.getsourcelines(src)
+        src = [textwrap.fill(line, tabsize=4, width=9999) for line in src]
+        src = textwrap.dedent("\n".join(src))
     tree = ast.parse(src)
     _adjust_line_numbers(tree, starting_line_no - 1)
     if verbose:

allo/ir/visitor.py

@@ -5,6 +5,7 @@
 import ast
 from .._mlir import InsertionPoint
 from .._mlir.dialects import allo as allo_d
+from ..utils import register_dialect
 
 
 class BlockScopeGuard:
@@ -69,7 +70,7 @@ def __init__(
         self.global_vars = global_vars
         self.mlir_ctx = mlir_ctx
         self.file_name = None
-        allo_d.register_dialect(mlir_ctx)
+        register_dialect(mlir_ctx)
         # map from function name to function arguments
         self.func_args = {} if func_args is None else func_args
         self.func_id = None
@@ -413,8 +414,6 @@ def __init__(self, symbolic_mapping, var_map, variables):
         self.special_symbol = set()
 
     def visit_Name(self, node):
-        if node.id in self.variables:
-            raise ValueError("Fail to resolve the expression as symbolic expression.")
         if node.id in self.symbolic_mapping:
             symbol_var = self.symbolic_mapping[node.id]
             if isinstance(symbol_var, str):
@@ -428,6 +427,8 @@ def visit_Name(self, node):
             return new_node
         if node.id in self.var_map:
             return ast.Constant(self.var_map[node.id])
+        if node.id in self.variables:
+            raise ValueError("Fail to resolve the expression as symbolic expression.")
         return node
 
 

allo/utils.py

@@ -8,6 +8,7 @@
 import numpy.typing as npt
 import ml_dtypes
 from ._mlir.ir import (
+    Context,
     MemRefType,
     RankedTensorType,
     IntegerType,
@@ -580,3 +581,7 @@ def allo_to_numpy_dtype(allo_type: AlloType) -> npt.DTypeLike:
             dtype = np.int64 if isinstance(allo_type, Fixed) else np.uint64
 
     return dtype
+
+
+def register_dialect(ctx: Context):
+    allo_d.register_dialect(ctx)

mlir/include/allo/Conversion/Passes.h

@@ -21,13 +21,15 @@ std::unique_ptr<OperationPass<ModuleOp>> createFixedPointToIntegerPass();
 std::unique_ptr<OperationPass<ModuleOp>> createLowerCompositeTypePass();
 std::unique_ptr<OperationPass<ModuleOp>> createLowerBitOpsPass();
 std::unique_ptr<OperationPass<ModuleOp>> createLowerTransformLayoutOpsPass();
+std::unique_ptr<OperationPass<ModuleOp>> createLowerMemCopyOpsPass();
 std::unique_ptr<OperationPass<ModuleOp>> createLowerPrintOpsPass();
 
 bool applyAlloToLLVMLoweringPass(ModuleOp &module, MLIRContext &context);
 bool applyFixedPointToInteger(ModuleOp &module);
 bool applyLowerCompositeType(ModuleOp &module);
 bool applyLowerBitOps(ModuleOp &module);
 bool applyLowerTransformLayoutOps(ModuleOp &module);
+bool applyLowerMemCopyOps(ModuleOp &module);
 bool applyLowerPrintOps(ModuleOp &module);
 
 /// Registers all Allo conversion passes

mlir/include/allo/Conversion/Passes.td

@@ -33,6 +33,11 @@ def LowerTransformLayoutOps : Pass<"lower-transform-layout-ops", "ModuleOp"> {
   let constructor = "mlir::allo::createLowerTransformLayoutOpsPass()";
 }
 
+def LowerMemCopyOps : Pass<"lower-memcopy-ops", "ModuleOp"> {
+  let summary = "Lower memcopy operations";
+  let constructor = "mlir::allo::createLowerMemCopyOpsPass()";
+}
+
 def LowerPrintOps : Pass<"lower-print-ops", "ModuleOp"> {
   let summary = "Lower print operations";
   let constructor = "mlir::allo::createLowerPrintOpsPass()";

mlir/include/allo/Dialect/AlloOps.td

@@ -935,7 +935,7 @@ def OrOp : Allo_Op<"or"> {
 //===----------------------------------------------------------------------===//
 
 def YieldOp : Allo_Op<"yield", [NoMemoryEffect, Terminator,
-                               ParentOneOf<["AndOp, OrOp"]>]> {
+                               ParentOneOf<["AndOp", "OrOp", "GridMapOp"]>]> {
   let summary = "yield and termination operation";
   let description = [{
     "allo.yield" yields an SSA value from the Allo dialect op region and
@@ -1048,6 +1048,56 @@ def StreamPutOp : Allo_Op<"stream_put"> {
   }];
 }
 
+def StreamGlobalOp : Allo_Op<"stream_global", [Symbol]> {
+  let summary = "Create a global stream object";
+  let arguments = (ins SymbolNameAttr:$sym_name, TypeAttr:$element_type, DenseI64ArrayAttr:$shape);
+  let assemblyFormat = [{
+      $sym_name `:` $element_type $shape  attr-dict
+  }];
+  let hasVerifier = 1;
+}
+
+def GlobalStreamGetOp : Allo_Op<"get_stream_global"> {
+  let summary = "Get an object from a global stream";
+  // arguments:
+  // - $global: symbol name of the global stream
+  // - $indices: variadic index operands
+  // - $map: affine map
+  let arguments = (
+    ins FlatSymbolRefAttr:$global,
+        Variadic<Index>:$indices,
+        AffineMapAttr:$map
+  );
+  let results = (outs AnyType:$result);
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDeclaration = [{
+    void simplifyAffineMap();
+  }];
+}
+
+def GlobalStreamPutOp : Allo_Op<"put_stream_global"> {
+  let summary = "Put an object to a global stream";
+  // arguments:
+  // - $global: symbol name of the global stream
+  // - $indices: variadic index operands
+  // - $data: value to put
+  // - $map: affine map to compute indices
+  let arguments = (
+      ins FlatSymbolRefAttr:$global,
+          Variadic<Index>:$indices,
+          AnyType:$data,
+          AffineMapAttr:$map
+  );
+  let results = (outs);
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDeclaration = [{
+    void simplifyAffineMap();
+  }];
+}
+
+
 //===----------------------------------------------------------------------===//
 // Layout operations
 //===----------------------------------------------------------------------===//
@@ -1096,4 +1146,35 @@ def TransformLayoutOp : Allo_Op<"transform_layout"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// SPMW operations
+//===----------------------------------------------------------------------===//
+def GridMapOp : Allo_Op<"grid_map",
+  [RecursiveMemoryEffects, SingleBlockImplicitTerminator<"YieldOp">]> {
+  let summary = "Grid map operation";
+  let description = [{
+    The grid_map operation distributes a computation over a logical grid.
+    It takes a list of memrefs with static shape (`tensors`) as input arguments.
+    It contains a single block whose arguments are the sharded memrefs. The operations in the block can access variables in parent regions.
+  }];
+
+  let arguments = (ins
+    Variadic<AnyStaticShapeMemRef>:$tensors,
+    ArrayAttr:$sharding,
+    DenseI64ArrayAttr:$grid
+  );
+  let regions = (region SizedRegion<1>:$body);
+
+  let assemblyFormat = [{
+    `(` $tensors `)`
+    `sharding` `=` $sharding
+    `grid` `=` $grid
+    $body
+    attr-dict
+    `:` type($tensors)
+  }];
+
+  let hasVerifier = 1;
+}
+
 #endif // ALLO_OPS