22023884668

Committed 14 Feb 2026 08:36PM UTC coverage: 64.903% (-1.4%) from 66.315%

Build # 22023884668

Build Type

Pull #525

github

Committed by

web-flow

Commit Message

Merge 1d47f8bf2 into 9d01cacd5

Pull Request Pull Request #525: Step 3 (Native Tensor Support): Refactor Python Frontend

Run Details

2522 of 3435 new or added lines in 32 files covered. (73.42%)

320 existing lines in 15 files now uncovered.

23204 of 35752 relevant lines covered (64.9%)

370.03 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

78.82

/python/docc/python/ast_parser.py

import ast
import copy
import inspect
import textwrap
from docc.sdfg import (
    Scalar,
    PrimitiveType,
    Pointer,
    TaskletCode,
    DebugInfo,
    Structure,
    CMathFunction,
)
from docc.python.ast_utils import (
    SliceRewriter,
    get_debug_info,
    contains_ufunc_outer,
    normalize_negative_index,
)
from docc.python.types import (
    sdfg_type_from_type,
    element_type_from_sdfg_type,
)
from docc.python.functions.numpy import NumPyHandler
from docc.python.functions.math import MathHandler
from docc.python.functions.python import PythonHandler
from docc.python.functions.scipy import SciPyHandler


class ASTParser(ast.NodeVisitor):
    def __init__(
        self,
        builder,
        array_info=None,
        symbol_table=None,
        filename="",
        function_name="",
        infer_return_type=False,
        globals_dict=None,
        unique_counter_ref=None,
        structure_member_info=None,
    ):
        self.builder = builder
        self.array_info = array_info if array_info is not None else {}
        self.symbol_table = symbol_table if symbol_table is not None else {}
        self.filename = filename
        self.function_name = function_name
        self.infer_return_type = infer_return_type
        self.globals_dict = globals_dict if globals_dict is not None else {}
        self._unique_counter_ref = (
            unique_counter_ref if unique_counter_ref is not None else [0]
        )
        self._access_cache = {}
        self.structure_member_info = (
            structure_member_info if structure_member_info is not None else {}
        )
        self.captured_return_shapes = {}  # Map param name to shape string list

        # Initialize handlers - they receive 'self' to access expression visitor methods
        self.numpy_visitor = NumPyHandler(self)
        self.math_handler = MathHandler(self)
        self.python_handler = PythonHandler(self)
        self.scipy_handler = SciPyHandler(self)

    def visit_Constant(self, node):
        if isinstance(node.value, bool):
            return "true" if node.value else "false"
        return str(node.value)

    def visit_Name(self, node):
        name = node.id
        if name not in self.symbol_table and self.globals_dict is not None:
            if name in self.globals_dict:
                val = self.globals_dict[name]
                if isinstance(val, (int, float)):
                    return str(val)
        return name

    def visit_Add(self, node):
        return "+"

    def visit_Sub(self, node):
        return "-"

    def visit_Mult(self, node):
        return "*"

    def visit_Div(self, node):
        return "/"

    def visit_FloorDiv(self, node):
        return "//"

    def visit_Mod(self, node):
        return "%"

    def visit_Pow(self, node):
        return "**"

    def visit_Eq(self, node):
        return "=="

    def visit_NotEq(self, node):
        return "!="

    def visit_Lt(self, node):
        return "<"

    def visit_LtE(self, node):
        return "<="

    def visit_Gt(self, node):
        return ">"

    def visit_GtE(self, node):
        return ">="

    def visit_And(self, node):
        return "&"

    def visit_Or(self, node):
        return "|"

    def visit_BitAnd(self, node):
        return "&"

    def visit_BitOr(self, node):
        return "|"

    def visit_BitXor(self, node):
        return "^"

    def visit_LShift(self, node):
        return "<<"

    def visit_RShift(self, node):
        return ">>"

    def visit_Not(self, node):
        return "!"

    def visit_USub(self, node):
        return "-"

    def visit_UAdd(self, node):
        return "+"

    def visit_Invert(self, node):
        return "~"

    def visit_BoolOp(self, node):
        op = self.visit(node.op)
        values = [f"({self.visit(v)} != 0)" for v in node.values]
        expr_str = f"{f' {op} '.join(values)}"

        tmp_name = self.builder.find_new_name()
        dtype = Scalar(PrimitiveType.Bool)
        self.builder.add_container(tmp_name, dtype, False)

        self.builder.begin_if(expr_str)
        self._add_assign_constant(tmp_name, "true", dtype)
        self.builder.begin_else()
        self._add_assign_constant(tmp_name, "false", dtype)
        self.builder.end_if()

        self.symbol_table[tmp_name] = dtype
        return tmp_name

    def visit_UnaryOp(self, node):
        if (
            isinstance(node.op, ast.USub)
            and isinstance(node.operand, ast.Constant)
            and isinstance(node.operand.value, (int, float))
        ):
            return f"-{node.operand.value}"

        op = self.visit(node.op)
        operand = self.visit(node.operand)

        if operand in self.array_info and op == "-":
            return self.numpy_visitor.handle_array_negate(operand)

        tmp_name = self.builder.find_new_name()
        dtype = Scalar(PrimitiveType.Double)
        if operand in self.symbol_table:
            dtype = self.symbol_table[operand]
            if isinstance(dtype, Pointer) and dtype.has_pointee_type():
                dtype = dtype.pointee_type
        elif self._is_int(operand):
            dtype = Scalar(PrimitiveType.Int64)
        elif isinstance(node.op, ast.Not):
            dtype = Scalar(PrimitiveType.Bool)

        self.builder.add_container(tmp_name, dtype, False)
        self.symbol_table[tmp_name] = dtype

        block = self.builder.add_block()
        t_src, src_sub = self._add_read(block, operand)
        t_dst = self.builder.add_access(block, tmp_name)

        if isinstance(node.op, ast.Not):
            t_const = self.builder.add_constant(
                block, "true", Scalar(PrimitiveType.Bool)
            )
            t_task = self.builder.add_tasklet(
                block, TaskletCode.int_xor, ["_in1", "_in2"], ["_out"]
            )
            self.builder.add_memlet(block, t_src, "void", t_task, "_in1", src_sub)
            self.builder.add_memlet(block, t_const, "void", t_task, "_in2", "")
            self.builder.add_memlet(block, t_task, "_out", t_dst, "void", "")

        elif op == "-":
            if dtype.primitive_type == PrimitiveType.Int64:
                t_const = self.builder.add_constant(block, "0", dtype)
                t_task = self.builder.add_tasklet(
                    block, TaskletCode.int_sub, ["_in1", "_in2"], ["_out"]
                )
                self.builder.add_memlet(block, t_const, "void", t_task, "_in1", "")
                self.builder.add_memlet(block, t_src, "void", t_task, "_in2", src_sub)
                self.builder.add_memlet(block, t_task, "_out", t_dst, "void", "")
            else:
                t_task = self.builder.add_tasklet(
                    block, TaskletCode.fp_neg, ["_in"], ["_out"]
                )
                self.builder.add_memlet(block, t_src, "void", t_task, "_in", src_sub)
                self.builder.add_memlet(block, t_task, "_out", t_dst, "void", "")

        elif op == "~":
            t_const = self.builder.add_constant(
                block, "-1", Scalar(PrimitiveType.Int64)
            )
            t_task = self.builder.add_tasklet(
                block, TaskletCode.int_xor, ["_in1", "_in2"], ["_out"]
            )
            self.builder.add_memlet(block, t_src, "void", t_task, "_in1", src_sub)
            self.builder.add_memlet(block, t_const, "void", t_task, "_in2", "")
            self.builder.add_memlet(block, t_task, "_out", t_dst, "void", "")

        else:
            t_task = self.builder.add_tasklet(
                block, TaskletCode.assign, ["_in"], ["_out"]
            )
            self.builder.add_memlet(block, t_src, "void", t_task, "_in", src_sub)
            self.builder.add_memlet(block, t_task, "_out", t_dst, "void", "")

        return tmp_name

    def visit_BinOp(self, node):
        if isinstance(node.op, ast.MatMult):
            return self.numpy_visitor.handle_numpy_matmul_op(node.left, node.right)

        left = self.visit(node.left)
        op = self.visit(node.op)
        right = self.visit(node.right)

        left_is_array = left in self.array_info
        right_is_array = right in self.array_info

        if left_is_array or right_is_array:
            op_map = {"+": "add", "-": "sub", "*": "mul", "/": "div", "**": "pow"}
            if op in op_map:
                return self.numpy_visitor.handle_array_binary_op(
                    op_map[op], left, right
                )
            else:
                raise NotImplementedError(f"Array operation {op} not supported")

        tmp_name = self.builder.find_new_name()

        dtype = Scalar(PrimitiveType.Double)

        left_is_int = self._is_int(left)
        right_is_int = self._is_int(right)

        if left_is_int and right_is_int and op not in ["/", "**"]:
            dtype = Scalar(PrimitiveType.Int64)

        if not self.builder.exists(tmp_name):
            self.builder.add_container(tmp_name, dtype, False)
            self.symbol_table[tmp_name] = dtype

        real_left = left
        real_right = right

        if dtype.primitive_type == PrimitiveType.Double:
            if left_is_int:
                left_cast = self.builder.find_new_name()
                self.builder.add_container(
                    left_cast, Scalar(PrimitiveType.Double), False
                )
                self.symbol_table[left_cast] = Scalar(PrimitiveType.Double)

                c_block = self.builder.add_block()
                t_src, src_sub = self._add_read(c_block, left)
                t_dst = self.builder.add_access(c_block, left_cast)
                t_task = self.builder.add_tasklet(
                    c_block, TaskletCode.assign, ["_in"], ["_out"]
                )
                self.builder.add_memlet(c_block, t_src, "void", t_task, "_in", src_sub)
                self.builder.add_memlet(c_block, t_task, "_out", t_dst, "void", "")

                real_left = left_cast

            if right_is_int:
                right_cast = self.builder.find_new_name()
                self.builder.add_container(
                    right_cast, Scalar(PrimitiveType.Double), False
                )
                self.symbol_table[right_cast] = Scalar(PrimitiveType.Double)

                c_block = self.builder.add_block()
                t_src, src_sub = self._add_read(c_block, right)
                t_dst = self.builder.add_access(c_block, right_cast)
                t_task = self.builder.add_tasklet(
                    c_block, TaskletCode.assign, ["_in"], ["_out"]
                )
                self.builder.add_memlet(c_block, t_src, "void", t_task, "_in", src_sub)
                self.builder.add_memlet(c_block, t_task, "_out", t_dst, "void", "")

                real_right = right_cast

        if op == "**":
            block = self.builder.add_block()
            t_left, left_sub = self._add_read(block, real_left)
            t_right, right_sub = self._add_read(block, real_right)
            t_out = self.builder.add_access(block, tmp_name)

            t_task = self.builder.add_cmath(block, CMathFunction.pow)
            self.builder.add_memlet(block, t_left, "void", t_task, "_in1", left_sub)
            self.builder.add_memlet(block, t_right, "void", t_task, "_in2", right_sub)
            self.builder.add_memlet(block, t_task, "_out", t_out, "void", "")

            return tmp_name
        elif op == "%":
            block = self.builder.add_block()
            t_left, left_sub = self._add_read(block, real_left)
            t_right, right_sub = self._add_read(block, real_right)
            t_out = self.builder.add_access(block, tmp_name)

            if dtype.primitive_type == PrimitiveType.Int64:
                t_rem1 = self.builder.add_tasklet(
                    block, TaskletCode.int_srem, ["_in1", "_in2"], ["_out"]
                )
                self.builder.add_memlet(block, t_left, "void", t_rem1, "_in1", left_sub)
                self.builder.add_memlet(
                    block, t_right, "void", t_rem1, "_in2", right_sub
                )

                rem1_name = self.builder.find_new_name()
                self.builder.add_container(rem1_name, dtype, False)
                t_rem1_out = self.builder.add_access(block, rem1_name)
                self.builder.add_memlet(block, t_rem1, "_out", t_rem1_out, "void", "")

                t_add = self.builder.add_tasklet(
                    block, TaskletCode.int_add, ["_in1", "_in2"], ["_out"]
                )
                self.builder.add_memlet(block, t_rem1_out, "void", t_add, "_in1", "")
                self.builder.add_memlet(
                    block, t_right, "void", t_add, "_in2", right_sub
                )

                add_name = self.builder.find_new_name()
                self.builder.add_container(add_name, dtype, False)
                t_add_out = self.builder.add_access(block, add_name)
                self.builder.add_memlet(block, t_add, "_out", t_add_out, "void", "")

                t_rem2 = self.builder.add_tasklet(
                    block, TaskletCode.int_srem, ["_in1", "_in2"], ["_out"]
                )
                self.builder.add_memlet(block, t_add_out, "void", t_rem2, "_in1", "")
                self.builder.add_memlet(
                    block, t_right, "void", t_rem2, "_in2", right_sub
                )
                self.builder.add_memlet(block, t_rem2, "_out", t_out, "void", "")

                return tmp_name
            else:
                t_rem1 = self.builder.add_tasklet(
                    block, TaskletCode.fp_rem, ["_in1", "_in2"], ["_out"]
                )
                self.builder.add_memlet(block, t_left, "void", t_rem1, "_in1", left_sub)
                self.builder.add_memlet(
                    block, t_right, "void", t_rem1, "_in2", right_sub
                )

                rem1_name = self.builder.find_new_name()
                self.builder.add_container(rem1_name, dtype, False)
                t_rem1_out = self.builder.add_access(block, rem1_name)
                self.builder.add_memlet(block, t_rem1, "_out", t_rem1_out, "void", "")

                t_add = self.builder.add_tasklet(
                    block, TaskletCode.fp_add, ["_in1", "_in2"], ["_out"]
                )
                self.builder.add_memlet(block, t_rem1_out, "void", t_add, "_in1", "")
                self.builder.add_memlet(
                    block, t_right, "void", t_add, "_in2", right_sub
                )

                add_name = self.builder.find_new_name()
                self.builder.add_container(add_name, dtype, False)
                t_add_out = self.builder.add_access(block, add_name)
                self.builder.add_memlet(block, t_add, "_out", t_add_out, "void", "")

                t_rem2 = self.builder.add_tasklet(
                    block, TaskletCode.fp_rem, ["_in1", "_in2"], ["_out"]
                )
                self.builder.add_memlet(block, t_add_out, "void", t_rem2, "_in1", "")
                self.builder.add_memlet(
                    block, t_right, "void", t_rem2, "_in2", right_sub
                )
                self.builder.add_memlet(block, t_rem2, "_out", t_out, "void", "")

                return tmp_name

        tasklet_code = None
        if dtype.primitive_type == PrimitiveType.Int64:
            if op == "+":
                tasklet_code = TaskletCode.int_add
            elif op == "-":
                tasklet_code = TaskletCode.int_sub
            elif op == "*":
                tasklet_code = TaskletCode.int_mul
            elif op == "/":
                tasklet_code = TaskletCode.int_sdiv
            elif op == "//":
                tasklet_code = TaskletCode.int_sdiv
            elif op == "&":
                tasklet_code = TaskletCode.int_and
            elif op == "|":
                tasklet_code = TaskletCode.int_or
            elif op == "^":
                tasklet_code = TaskletCode.int_xor
            elif op == "<<":
                tasklet_code = TaskletCode.int_shl
            elif op == ">>":
                tasklet_code = TaskletCode.int_lshr
        else:
            if op == "+":
                tasklet_code = TaskletCode.fp_add
            elif op == "-":
                tasklet_code = TaskletCode.fp_sub
            elif op == "*":
                tasklet_code = TaskletCode.fp_mul
            elif op == "/":
                tasklet_code = TaskletCode.fp_div
            elif op == "//":
                tasklet_code = TaskletCode.fp_div
            else:
                raise NotImplementedError(f"Operation {op} not supported for floats")

        block = self.builder.add_block()
        t_left, left_sub = self._add_read(block, real_left)
        t_right, right_sub = self._add_read(block, real_right)
        t_out = self.builder.add_access(block, tmp_name)

        t_task = self.builder.add_tasklet(
            block, tasklet_code, ["_in1", "_in2"], ["_out"]
        )

        self.builder.add_memlet(block, t_left, "void", t_task, "_in1", left_sub)
        self.builder.add_memlet(block, t_right, "void", t_task, "_in2", right_sub)
        self.builder.add_memlet(block, t_task, "_out", t_out, "void", "")

        return tmp_name

    def visit_Attribute(self, node):
        if node.attr == "shape":
            if isinstance(node.value, ast.Name) and node.value.id in self.array_info:
                return f"_shape_proxy_{node.value.id}"

        if node.attr == "T":
            value_name = None
            if isinstance(node.value, ast.Name):
                value_name = node.value.id
            elif isinstance(node.value, ast.Subscript):
                value_name = self.visit(node.value)

            if value_name and value_name in self.array_info:
                return self.numpy_visitor.handle_transpose_expr(node)

        if isinstance(node.value, ast.Name) and node.value.id == "math":
            val = ""
            if node.attr == "pi":
                val = "M_PI"
            elif node.attr == "e":
                val = "M_E"

            if val:
                tmp_name = self.builder.find_new_name()
                dtype = Scalar(PrimitiveType.Double)
                self.builder.add_container(tmp_name, dtype, False)
                self.symbol_table[tmp_name] = dtype
                self._add_assign_constant(tmp_name, val, dtype)
                return tmp_name

        if isinstance(node.value, ast.Name):
            obj_name = node.value.id
            attr_name = node.attr

            if obj_name in self.symbol_table:
                obj_type = self.symbol_table[obj_name]
                if isinstance(obj_type, Pointer) and obj_type.has_pointee_type():
                    pointee_type = obj_type.pointee_type
                    if isinstance(pointee_type, Structure):
                        struct_name = pointee_type.name

                        if (
                            struct_name in self.structure_member_info
                            and attr_name in self.structure_member_info[struct_name]
                        ):
                            member_index, member_type = self.structure_member_info[
                                struct_name
                            ][attr_name]
                        else:
                            raise RuntimeError(
                                f"Member '{attr_name}' not found in structure '{struct_name}'. "
                                f"Available members: {list(self.structure_member_info.get(struct_name, {}).keys())}"
                            )

                        tmp_name = self.builder.find_new_name()

                        self.builder.add_container(tmp_name, member_type, False)
                        self.symbol_table[tmp_name] = member_type

                        block = self.builder.add_block()
                        obj_access = self.builder.add_access(block, obj_name)
                        tmp_access = self.builder.add_access(block, tmp_name)

                        tasklet = self.builder.add_tasklet(
                            block, TaskletCode.assign, ["_in"], ["_out"]
                        )

                        subset = "0," + str(member_index)
                        self.builder.add_memlet(
                            block, obj_access, "", tasklet, "_in", subset
                        )
                        self.builder.add_memlet(block, tasklet, "_out", tmp_access, "")

                        return tmp_name

        raise NotImplementedError(f"Attribute access {node.attr} not supported")

    def visit_Compare(self, node):
        left = self.visit(node.left)
        if len(node.ops) > 1:
            raise NotImplementedError("Chained comparisons not supported yet")

        op = self.visit(node.ops[0])
        right = self.visit(node.comparators[0])

        left_is_array = left in self.array_info
        right_is_array = right in self.array_info

        if left_is_array or right_is_array:
            return self.numpy_visitor.handle_array_compare(
                left, op, right, left_is_array, right_is_array
            )

        expr_str = f"{left} {op} {right}"

        tmp_name = self.builder.find_new_name()
        dtype = Scalar(PrimitiveType.Bool)
        self.builder.add_container(tmp_name, dtype, False)

        self.builder.begin_if(expr_str)
        self.builder.add_transition(tmp_name, "true")
        self.builder.begin_else()
        self.builder.add_transition(tmp_name, "false")
        self.builder.end_if()

        self.symbol_table[tmp_name] = dtype
        return tmp_name

    def visit_Subscript(self, node):
        value_str = self.visit(node.value)

        if value_str.startswith("_shape_proxy_"):
            array_name = value_str[len("_shape_proxy_") :]
            if isinstance(node.slice, ast.Constant):
                idx = node.slice.value
            elif isinstance(node.slice, ast.Index):
                idx = node.slice.value.value
            else:
                try:
                    idx = int(self.visit(node.slice))
                except:
                    raise NotImplementedError(
                        "Dynamic shape indexing not fully supported yet"
                    )

            if (
                array_name in self.array_info
                and "shapes" in self.array_info[array_name]
            ):
                return self.array_info[array_name]["shapes"][idx]

            return f"_{array_name}_shape_{idx}"

        if value_str in self.array_info:
            ndim = self.array_info[value_str]["ndim"]
            shapes = self.array_info[value_str].get("shapes", [])

            if isinstance(node.slice, ast.Tuple):
                indices_nodes = node.slice.elts
            else:
                indices_nodes = [node.slice]

            all_full_slices = True
            for idx in indices_nodes:
                if isinstance(idx, ast.Slice):
                    if idx.lower is not None or idx.upper is not None:
                        all_full_slices = False
                        break
                else:
                    all_full_slices = False
                    break

            if all_full_slices:
                return value_str

            has_slices = any(isinstance(idx, ast.Slice) for idx in indices_nodes)
            if has_slices:
                return self._handle_expression_slicing(
                    node, value_str, indices_nodes, shapes, ndim
                )

            if len(indices_nodes) == 1 and self._is_array_index(indices_nodes[0]):
                if self.builder:
                    return self._handle_gather(value_str, indices_nodes[0])

            if isinstance(node.slice, ast.Tuple):
                indices = [self.visit(elt) for elt in node.slice.elts]
            else:
                indices = [self.visit(node.slice)]

            if len(indices) != ndim:
                raise ValueError(
                    f"Array {value_str} has {ndim} dimensions, but accessed with {len(indices)} indices"
                )

            normalized_indices = []
            for i, idx_str in enumerate(indices):
                shape_val = shapes[i] if i < len(shapes) else f"_{value_str}_shape_{i}"
                if isinstance(idx_str, str) and (
                    idx_str.startswith("-") or idx_str.startswith("(-")
                ):
                    normalized_indices.append(f"({shape_val} + {idx_str})")
                else:
                    normalized_indices.append(idx_str)

            linear_index = ""
            for i in range(ndim):
                term = normalized_indices[i]
                for j in range(i + 1, ndim):
                    shape_val = shapes[j] if j < len(shapes) else None
                    shape_sym = (
                        shape_val
                        if shape_val is not None
                        else f"_{value_str}_shape_{j}"
                    )
                    term = f"(({term}) * {shape_sym})"

                if i == 0:
                    linear_index = term
                else:
                    linear_index = f"({linear_index} + {term})"

            access_str = f"{value_str}({linear_index})"

            if self.builder and isinstance(node.ctx, ast.Load):
                dtype = Scalar(PrimitiveType.Double)
                if value_str in self.symbol_table:
                    t = self.symbol_table[value_str]
                    if type(t).__name__ == "Array" and hasattr(t, "element_type"):
                        et = t.element_type
                        if callable(et):
                            et = et()
                        dtype = et
                    elif type(t).__name__ == "Pointer" and hasattr(t, "pointee_type"):
                        et = t.pointee_type
                        if callable(et):
                            et = et()
                        dtype = et

                tmp_name = self.builder.find_new_name()
                self.builder.add_container(tmp_name, dtype, False)

                block = self.builder.add_block()
                t_src = self.builder.add_access(block, value_str)
                t_dst = self.builder.add_access(block, tmp_name)
                t_task = self.builder.add_tasklet(
                    block, TaskletCode.assign, ["_in"], ["_out"]
                )

                self.builder.add_memlet(
                    block, t_src, "void", t_task, "_in", linear_index
                )
                self.builder.add_memlet(block, t_task, "_out", t_dst, "void", "")

                self.symbol_table[tmp_name] = dtype
                return tmp_name

            return access_str

        slice_val = self.visit(node.slice)
        access_str = f"{value_str}({slice_val})"

        if (
            self.builder
            and isinstance(node.ctx, ast.Load)
            and value_str in self.array_info
        ):
            tmp_name = self.builder.find_new_name()
            self.builder.add_container(tmp_name, Scalar(PrimitiveType.Double), False)
            self.builder.add_assignment(tmp_name, access_str)
            self.symbol_table[tmp_name] = Scalar(PrimitiveType.Double)
            return tmp_name

        return access_str

    def visit_AugAssign(self, node):
        if isinstance(node.target, ast.Name) and node.target.id in self.array_info:
            # Convert to slice assignment: target[:] = target op value
            ndim = self.array_info[node.target.id]["ndim"]

            slices = []
            for _ in range(ndim):
                slices.append(ast.Slice(lower=None, upper=None, step=None))

            if ndim == 1:
                slice_arg = slices[0]
            else:
                slice_arg = ast.Tuple(elts=slices, ctx=ast.Load())

            slice_node = ast.Subscript(
                value=node.target, slice=slice_arg, ctx=ast.Store()
            )

            new_node = ast.Assign(
                targets=[slice_node],
                value=ast.BinOp(left=node.target, op=node.op, right=node.value),
            )
            self.visit_Assign(new_node)
        else:
            new_node = ast.Assign(
                targets=[node.target],
                value=ast.BinOp(left=node.target, op=node.op, right=node.value),
            )
            self.visit_Assign(new_node)

    def visit_Assign(self, node):
        if len(node.targets) > 1:
            tmp_name = self.builder.find_new_name()
            # Assign value to temporary
            val_assign = ast.Assign(
                targets=[ast.Name(id=tmp_name, ctx=ast.Store())], value=node.value
            )
            ast.copy_location(val_assign, node)
            self.visit_Assign(val_assign)

            # Assign temporary to targets
            for target in node.targets:
                assign = ast.Assign(
                    targets=[target], value=ast.Name(id=tmp_name, ctx=ast.Load())
                )
                ast.copy_location(assign, node)
                self.visit_Assign(assign)
            return

        target = node.targets[0]

        # Handle tuple unpacking: I, J, K = expr1, expr2, expr3
        if isinstance(target, ast.Tuple):
            if isinstance(node.value, ast.Tuple):
                # Unpacking tuple to tuple: a, b, c = x, y, z
                if len(target.elts) != len(node.value.elts):
                    raise ValueError("Tuple unpacking size mismatch")
                for tgt, val in zip(target.elts, node.value.elts):
                    assign = ast.Assign(targets=[tgt], value=val)
                    ast.copy_location(assign, node)
                    self.visit_Assign(assign)
            else:
                raise NotImplementedError(
                    "Tuple unpacking from non-tuple values not supported"
                )
            return

        # Special case: linear algebra functions (handled by NumPyHandler)
        if self.numpy_visitor.is_gemm(node.value):
            if self.numpy_visitor.handle_gemm(target, node.value):
                return
            if self.numpy_visitor.handle_dot(target, node.value):
                return

        # Special case: outer product (handled by NumPyHandler)
        if self.numpy_visitor.is_outer(node.value):
            if self.numpy_visitor.handle_outer(target, node.value):
                return

        # Special case: convolution (scipy.signal.correlate2d)
        if self.scipy_handler.is_correlate2d(node.value):
            if self.scipy_handler.handle_correlate2d(target, node.value):
                return

        # Special case: Transpose (handled by NumPyHandler)
        if self.numpy_visitor.is_transpose(node.value):
            if self.numpy_visitor.handle_transpose(target, node.value):
                return

        # Special case:
        if isinstance(target, ast.Subscript):
            target_name = self.visit(target.value)

            indices = []
            if isinstance(target.slice, ast.Tuple):
                indices = target.slice.elts
            else:
                indices = [target.slice]

            has_slice = False
            for idx in indices:
                if isinstance(idx, ast.Slice):
                    has_slice = True
                    break

            if has_slice:
                debug_info = get_debug_info(node, self.filename, self.function_name)
                self._handle_slice_assignment(
                    target, node.value, target_name, indices, debug_info
                )
                return

            target_name_full = self.visit(target)
            value_str = self.visit(node.value)
            debug_info = get_debug_info(node, self.filename, self.function_name)

            block = self.builder.add_block(debug_info)
            t_src, src_sub = self._add_read(block, value_str, debug_info)

            if "(" in target_name_full and target_name_full.endswith(")"):
                name = target_name_full.split("(")[0]
                subset = target_name_full[target_name_full.find("(") + 1 : -1]
                t_dst = self.builder.add_access(block, name, debug_info)
                dst_sub = subset
            else:
                t_dst = self.builder.add_access(block, target_name_full, debug_info)
                dst_sub = ""

            t_task = self.builder.add_tasklet(
                block, TaskletCode.assign, ["_in"], ["_out"], debug_info
            )

            self.builder.add_memlet(
                block, t_src, "void", t_task, "_in", src_sub, None, debug_info
            )
            self.builder.add_memlet(
                block, t_task, "_out", t_dst, "void", dst_sub, None, debug_info
            )
            return

        # Variable assignments
        if not isinstance(target, ast.Name):
            raise NotImplementedError("Only assignment to variables supported")

        target_name = target.id
        value_str = self.visit(node.value)
        debug_info = get_debug_info(node, self.filename, self.function_name)

        if not self.builder.exists(target_name):
            if isinstance(node.value, ast.Constant):
                val = node.value.value
                if isinstance(val, int):
                    dtype = Scalar(PrimitiveType.Int64)
                elif isinstance(val, float):
                    dtype = Scalar(PrimitiveType.Double)
                elif isinstance(val, bool):
                    dtype = Scalar(PrimitiveType.Bool)
                else:
                    raise NotImplementedError(f"Cannot infer type for {val}")

                self.builder.add_container(target_name, dtype, False)
                self.symbol_table[target_name] = dtype
            else:
                assert value_str in self.symbol_table
                self.builder.add_container(
                    target_name, self.symbol_table[value_str], False
                )
                self.symbol_table[target_name] = self.symbol_table[value_str]

        if value_str in self.array_info:
            self.array_info[target_name] = self.array_info[value_str]

        # Distinguish assignments: scalar -> tasklet, pointer -> reference_memlet
        src_type = self.symbol_table.get(value_str)
        dst_type = self.symbol_table[target_name]
        if src_type and isinstance(src_type, Pointer) and isinstance(dst_type, Pointer):
            block = self.builder.add_block(debug_info)
            t_src = self.builder.add_access(block, value_str, debug_info)
            t_dst = self.builder.add_access(block, target_name, debug_info)
            self.builder.add_reference_memlet(
                block, t_src, t_dst, "0", src_type, debug_info
            )
            return
        elif (src_type and isinstance(src_type, Scalar)) or isinstance(
            dst_type, Scalar
        ):
            block = self.builder.add_block(debug_info)
            t_dst = self.builder.add_access(block, target_name, debug_info)
            t_task = self.builder.add_tasklet(
                block, TaskletCode.assign, ["_in"], ["_out"], debug_info
            )

            if src_type:
                t_src = self.builder.add_access(block, value_str, debug_info)
            else:
                t_src = self.builder.add_constant(
                    block, value_str, dst_type, debug_info
                )

            self.builder.add_memlet(
                block, t_src, "void", t_task, "_in", "", None, debug_info
            )
            self.builder.add_memlet(
                block, t_task, "_out", t_dst, "void", "", None, debug_info
            )

            return

    def visit_Expr(self, node):
        self.visit(node.value)

    def visit_If(self, node):
        cond = self.visit(node.test)
        debug_info = get_debug_info(node, self.filename, self.function_name)
        self.builder.begin_if(f"{cond} != false", debug_info)

        for stmt in node.body:
            self.visit(stmt)

        if node.orelse:
            self.builder.begin_else(debug_info)
            for stmt in node.orelse:
                self.visit(stmt)

        self.builder.end_if()

    def visit_While(self, node):
        if node.orelse:
            raise NotImplementedError("while-else is not supported")

        debug_info = get_debug_info(node, self.filename, self.function_name)
        self.builder.begin_while(debug_info)

        # Evaluate condition inside the loop so it's re-evaluated each iteration
        cond = self.visit(node.test)

        # Create if-break pattern: if condition is false, break
        self.builder.begin_if(f"{cond} == false", debug_info)
        self.builder.add_break(debug_info)
        self.builder.end_if()

        for stmt in node.body:
            self.visit(stmt)

        self.builder.end_while()

    def visit_Break(self, node):
        debug_info = get_debug_info(node, self.filename, self.function_name)
        self.builder.add_break(debug_info)

    def visit_Continue(self, node):
        debug_info = get_debug_info(node, self.filename, self.function_name)
        self.builder.add_continue(debug_info)

    def visit_For(self, node):
        if node.orelse:
            raise NotImplementedError("while-else is not supported")
        if not isinstance(node.target, ast.Name):
            raise NotImplementedError("Only simple for loops supported")

        var = node.target.id
        debug_info = get_debug_info(node, self.filename, self.function_name)

        # Check if iterating over a range() call
        if (
            isinstance(node.iter, ast.Call)
            and isinstance(node.iter.func, ast.Name)
            and node.iter.func.id == "range"
        ):
            args = node.iter.args
            if len(args) == 1:
                start = "0"
                end = self.visit(args[0])
                step = "1"
            elif len(args) == 2:
                start = self.visit(args[0])
                end = self.visit(args[1])
                step = "1"
            elif len(args) == 3:
                start = self.visit(args[0])
                end = self.visit(args[1])

                # Special handling for step to avoid creating tasklets for constants
                step_node = args[2]
                if isinstance(step_node, ast.Constant):
                    step = str(step_node.value)
                elif (
                    isinstance(step_node, ast.UnaryOp)
                    and isinstance(step_node.op, ast.USub)
                    and isinstance(step_node.operand, ast.Constant)
                ):
                    step = f"-{step_node.operand.value}"
                else:
                    step = self.visit(step_node)
            else:
                raise ValueError("Invalid range arguments")

            if not self.builder.exists(var):
                self.builder.add_container(var, Scalar(PrimitiveType.Int64), False)
                self.symbol_table[var] = Scalar(PrimitiveType.Int64)

            self.builder.begin_for(var, start, end, step, debug_info)

            for stmt in node.body:
                self.visit(stmt)

            self.builder.end_for()
            return

        # Check if iterating over an ndarray (for x in array)
        if isinstance(node.iter, ast.Name):
            iter_name = node.iter.id
            if iter_name in self.array_info:
                arr_info = self.array_info[iter_name]
                if arr_info["ndim"] < 1:
                    raise NotImplementedError("Cannot iterate over 0-dimensional array")

                # Get the size of the first dimension
                arr_size = arr_info["shapes"][0]

                # Create a hidden index variable for the loop
                idx_var = self.builder.find_new_name()
                if not self.builder.exists(idx_var):
                    self.builder.add_container(
                        idx_var, Scalar(PrimitiveType.Int64), False
                    )
                    self.symbol_table[idx_var] = Scalar(PrimitiveType.Int64)

                # Determine the type of the loop variable (element type)
                # For a 1D array, it's a scalar; for ND array, it's a view of N-1 dimensions
                if arr_info["ndim"] == 1:
                    # Element is a scalar - get the element type from the array's type
                    arr_type = self.symbol_table.get(iter_name)
                    if isinstance(arr_type, Pointer):
                        elem_type = arr_type.pointee_type
                    else:
                        elem_type = Scalar(PrimitiveType.Double)  # Default fallback

                    if not self.builder.exists(var):
                        self.builder.add_container(var, elem_type, False)
                        self.symbol_table[var] = elem_type
                else:
                    # For multi-dimensional arrays, create a view/slice
                    # The loop variable becomes a pointer to the sub-array
                    inner_shapes = arr_info["shapes"][1:]
                    inner_ndim = arr_info["ndim"] - 1

                    arr_type = self.symbol_table.get(iter_name)
                    if isinstance(arr_type, Pointer):
                        elem_type = arr_type  # Keep as pointer type for views
                    else:
                        elem_type = Pointer(Scalar(PrimitiveType.Double))

                    if not self.builder.exists(var):
                        self.builder.add_container(var, elem_type, False)
                        self.symbol_table[var] = elem_type

                    # Register the view in array_info
                    self.array_info[var] = {"ndim": inner_ndim, "shapes": inner_shapes}

                # Begin the for loop
                self.builder.begin_for(idx_var, "0", str(arr_size), "1", debug_info)

                # Generate the assignment: var = array[idx_var]
                # Create an AST node for the assignment and visit it
                assign_node = ast.Assign(
                    targets=[ast.Name(id=var, ctx=ast.Store())],
                    value=ast.Subscript(
                        value=ast.Name(id=iter_name, ctx=ast.Load()),
                        slice=ast.Name(id=idx_var, ctx=ast.Load()),
                        ctx=ast.Load(),
                    ),
                )
                ast.copy_location(assign_node, node)
                self.visit_Assign(assign_node)

                # Visit the loop body
                for stmt in node.body:
                    self.visit(stmt)

                self.builder.end_for()
                return

        raise NotImplementedError(
            f"Only range() loops and iteration over ndarrays supported, got: {ast.dump(node.iter)}"
        )

    def visit_Return(self, node):
        if node.value is None:
            debug_info = get_debug_info(node, self.filename, self.function_name)
            self.builder.add_return("", debug_info)
            return

        if isinstance(node.value, ast.Tuple):
            values = node.value.elts
        else:
            values = [node.value]

        parsed_values = [self.visit(v) for v in values]
        debug_info = get_debug_info(node, self.filename, self.function_name)

        if self.infer_return_type:
            for i, res in enumerate(parsed_values):
                ret_name = f"_docc_ret_{i}"
                if not self.builder.exists(ret_name):
                    dtype = Scalar(PrimitiveType.Double)
                    if res in self.symbol_table:
                        dtype = self.symbol_table[res]
                    elif isinstance(values[i], ast.Constant):
                        val = values[i].value
                        if isinstance(val, int):
                            dtype = Scalar(PrimitiveType.Int64)
                        elif isinstance(val, float):
                            dtype = Scalar(PrimitiveType.Double)
                        elif isinstance(val, bool):
                            dtype = Scalar(PrimitiveType.Bool)

                    # Wrap Scalar in Pointer. Keep Arrays/Pointers as is.
                    arg_type = dtype
                    if isinstance(dtype, Scalar):
                        arg_type = Pointer(dtype)

                    self.builder.add_container(ret_name, arg_type, is_argument=True)
                    self.symbol_table[ret_name] = arg_type

                    if res in self.array_info:
                        self.array_info[ret_name] = self.array_info[res]

            self.infer_return_type = False

        for i, res in enumerate(parsed_values):
            ret_name = f"_docc_ret_{i}"
            typ = self.symbol_table.get(ret_name)

            is_array_return = False
            if res in self.array_info:
                # Only treat as array return if it has dimensions
                # 0-d arrays (scalars) should be handled by scalar assignment
                if self.array_info[res]["ndim"] > 0:
                    is_array_return = True
            elif res in self.symbol_table:
                if isinstance(self.symbol_table[res], Pointer):
                    is_array_return = True

            # Simple Scalar Assignment
            if not is_array_return:
                block = self.builder.add_block(debug_info)
                t_dst = self.builder.add_access(block, ret_name, debug_info)

                t_src, src_sub = self._add_read(block, res, debug_info)

                t_task = self.builder.add_tasklet(
                    block, TaskletCode.assign, ["_in"], ["_out"], debug_info
                )
                self.builder.add_memlet(
                    block, t_src, "void", t_task, "_in", src_sub, None, debug_info
                )
                self.builder.add_memlet(
                    block, t_task, "_out", t_dst, "void", "0", None, debug_info
                )

            # Array Assignment (Copy)
            else:
                # Record shape for metadata
                if res in self.array_info:
                    # Prefer runtime shapes if available (for indirect access patterns)
                    # Fall back to regular shapes otherwise
                    if "shapes_runtime" in self.array_info[res]:
                        shape = self.array_info[res]["shapes_runtime"]
                    else:
                        shape = self.array_info[res]["shapes"]
                    # Convert to string expressions
                    self.captured_return_shapes[ret_name] = [str(s) for s in shape]

                    # Ensure destination array info exists
                    if ret_name not in self.array_info:
                        self.array_info[ret_name] = self.array_info[res]

                # Copy Logic using visit_Assign
                ndim = 1
                if ret_name in self.array_info:
                    ndim = self.array_info[ret_name]["ndim"]

                slice_node = ast.Slice(lower=None, upper=None, step=None)
                if ndim > 1:
                    target_slice = ast.Tuple(elts=[slice_node] * ndim, ctx=ast.Load())
                else:
                    target_slice = slice_node

                target_sub = ast.Subscript(
                    value=ast.Name(id=ret_name, ctx=ast.Load()),
                    slice=target_slice,
                    ctx=ast.Store(),
                )

                # Value node reconstruction
                if isinstance(values[i], ast.Name):
                    val_node = values[i]
                else:
                    val_node = ast.Name(id=res, ctx=ast.Load())

                assign_node = ast.Assign(targets=[target_sub], value=val_node)
                self.visit_Assign(assign_node)

        # Add control flow return to exit the function/path
        self.builder.add_return("", debug_info)

    def visit_Call(self, node):
        func_name = ""
        module_name = ""
        submodule_name = ""
        if isinstance(node.func, ast.Attribute):
            if isinstance(node.func.value, ast.Name):
                if node.func.value.id == "math":
                    module_name = "math"
                    func_name = node.func.attr
                elif node.func.value.id in ["numpy", "np"]:
                    module_name = "numpy"
                    func_name = node.func.attr
                else:
                    array_name = node.func.value.id
                    method_name = node.func.attr
                    if array_name in self.array_info and method_name == "astype":
                        return self.numpy_visitor.handle_numpy_astype(node, array_name)
                    elif array_name in self.array_info and method_name == "copy":
                        return self.numpy_visitor.handle_numpy_copy(node, array_name)
            elif isinstance(node.func.value, ast.Attribute):
                if (
                    isinstance(node.func.value.value, ast.Name)
                    and node.func.value.value.id == "scipy"
                ):
                    module_name = "scipy"
                    submodule_name = node.func.value.attr
                    func_name = node.func.attr
                elif (
                    isinstance(node.func.value.value, ast.Name)
                    and node.func.value.value.id in ["numpy", "np"]
                    and node.func.attr == "outer"
                ):
                    ufunc_name = node.func.value.attr
                    return self.numpy_visitor.handle_ufunc_outer(node, ufunc_name)

        elif isinstance(node.func, ast.Name):
            func_name = node.func.id

        if module_name == "numpy":
            if self.numpy_visitor.has_handler(func_name):
                return self.numpy_visitor.handle_numpy_call(node, func_name)

        if module_name == "math":
            if self.math_handler.has_handler(func_name):
                return self.math_handler.handle_math_call(node, func_name)

        if module_name == "scipy":
            if self.scipy_handler.has_handler(submodule_name, func_name):
                return self.scipy_handler.handle_scipy_call(
                    node, submodule_name, func_name
                )

        if self.python_handler.has_handler(func_name):
            return self.python_handler.handle_python_call(node, func_name)

        if func_name in self.globals_dict:
            obj = self.globals_dict[func_name]
            if inspect.isfunction(obj):
                return self._handle_inline_call(node, obj)

        raise NotImplementedError(f"Function call {func_name} not supported")

    def _handle_inline_call(self, node, func_obj):
        try:
            source_lines, start_line = inspect.getsourcelines(func_obj)
            source = textwrap.dedent("".join(source_lines))
            tree = ast.parse(source)
            func_def = tree.body[0]
        except Exception as e:
            raise NotImplementedError(
                f"Could not parse function {func_obj.__name__}: {e}"
            )

        arg_vars = [self.visit(arg) for arg in node.args]

        if len(arg_vars) != len(func_def.args.args):
            raise NotImplementedError(
                f"Argument count mismatch for {func_obj.__name__}"
            )

        suffix = f"_{func_obj.__name__}_{self.builder.find_new_name()}"
        res_name = f"_res{suffix}"

        # Combine globals with closure variables of the inlined function
        combined_globals = dict(self.globals_dict)
        closure_constants = {}  # name -> value for numeric closure vars
        if func_obj.__closure__ is not None and func_obj.__code__.co_freevars:
            for name, cell in zip(func_obj.__code__.co_freevars, func_obj.__closure__):
                val = cell.cell_contents
                combined_globals[name] = val
                # Track numeric constants for injection
                if isinstance(val, (int, float)) and not isinstance(val, bool):
                    closure_constants[name] = val

        class VariableRenamer(ast.NodeTransformer):
            BUILTINS = {
                "range",
                "len",
                "int",
                "float",
                "bool",
                "str",
                "list",
                "dict",
                "tuple",
                "set",
                "print",
                "abs",
                "min",
                "max",
                "sum",
                "enumerate",
                "zip",
                "map",
                "filter",
                "sorted",
                "reversed",
                "True",
                "False",
                "None",
            }

            def __init__(self, suffix, globals_dict):
                self.suffix = suffix
                self.globals_dict = globals_dict

            def visit_Name(self, node):
                if node.id in self.globals_dict or node.id in self.BUILTINS:
                    return node
                return ast.Name(id=f"{node.id}{self.suffix}", ctx=node.ctx)

            def visit_Return(self, node):
                if node.value:
                    val = self.visit(node.value)
                    return ast.Assign(
                        targets=[ast.Name(id=res_name, ctx=ast.Store())],
                        value=val,
                    )
                return node

        renamer = VariableRenamer(suffix, combined_globals)
        new_body = [renamer.visit(stmt) for stmt in func_def.body]

        param_assignments = []

        # Inject closure constants as assignments
        for name, val in closure_constants.items():
            if isinstance(val, int):
                self.symbol_table[name] = Scalar(PrimitiveType.Int64)
                self.builder.add_container(name, Scalar(PrimitiveType.Int64), False)
                val_node = ast.Constant(value=val)
            else:
                self.symbol_table[name] = Scalar(PrimitiveType.Double)
                self.builder.add_container(name, Scalar(PrimitiveType.Double), False)
                val_node = ast.Constant(value=val)
            assign = ast.Assign(
                targets=[ast.Name(id=name, ctx=ast.Store())], value=val_node
            )
            param_assignments.append(assign)

        for arg_def, arg_val in zip(func_def.args.args, arg_vars):
            param_name = f"{arg_def.arg}{suffix}"

            if arg_val in self.symbol_table:
                self.symbol_table[param_name] = self.symbol_table[arg_val]
                self.builder.add_container(
                    param_name, self.symbol_table[arg_val], False
                )
                val_node = ast.Name(id=arg_val, ctx=ast.Load())
            elif self._is_int(arg_val):
                self.symbol_table[param_name] = Scalar(PrimitiveType.Int64)
                self.builder.add_container(
                    param_name, Scalar(PrimitiveType.Int64), False
                )
                val_node = ast.Constant(value=int(arg_val))
            else:
                try:
                    val = float(arg_val)
                    self.symbol_table[param_name] = Scalar(PrimitiveType.Double)
                    self.builder.add_container(
                        param_name, Scalar(PrimitiveType.Double), False
                    )
                    val_node = ast.Constant(value=val)
                except ValueError:
                    val_node = ast.Name(id=arg_val, ctx=ast.Load())

            assign = ast.Assign(
                targets=[ast.Name(id=param_name, ctx=ast.Store())], value=val_node
            )
            param_assignments.append(assign)

        final_body = param_assignments + new_body

        # Create a new parser instance for the inlined function
        parser = ASTParser(
            self.builder,
            self.array_info,
            self.symbol_table,
            globals_dict=combined_globals,
            unique_counter_ref=self._unique_counter_ref,
        )

        for stmt in final_body:
            parser.visit(stmt)

        return res_name

    def _add_assign_constant(self, target_name, value_str, dtype):
        block = self.builder.add_block()
        t_const = self.builder.add_constant(block, value_str, dtype)
        t_dst = self.builder.add_access(block, target_name)
        t_task = self.builder.add_tasklet(block, TaskletCode.assign, ["_in"], ["_out"])
        self.builder.add_memlet(block, t_const, "void", t_task, "_in", "")
        self.builder.add_memlet(block, t_task, "_out", t_dst, "void", "")

    def _handle_expression_slicing(self, node, value_str, indices_nodes, shapes, ndim):
        """Handle slicing in expressions (e.g., arr[1:, :, k+1])."""
        if not self.builder:
            raise ValueError("Builder required for expression slicing")

        dtype = Scalar(PrimitiveType.Double)
        if value_str in self.symbol_table:
            t = self.symbol_table[value_str]
            if isinstance(t, Pointer) and t.has_pointee_type():
                dtype = t.pointee_type

        result_shapes = []
        result_shapes_runtime = []
        slice_info = []
        index_info = []

        for i, idx in enumerate(indices_nodes):
            shape_val = shapes[i] if i < len(shapes) else f"_{value_str}_shape_{i}"

            if isinstance(idx, ast.Slice):
                start_str = "0"
                start_str_runtime = "0"
                if idx.lower is not None:
                    if self._contains_indirect_access(idx.lower):
                        start_str, start_str_runtime = (
                            self._materialize_indirect_access(
                                idx.lower, return_original_expr=True
                            )
                        )
                    else:
                        start_str = self.visit(idx.lower)
                        start_str_runtime = start_str
                    if isinstance(start_str, str) and (
                        start_str.startswith("-") or start_str.startswith("(-")
                    ):
                        start_str = f"({shape_val} + {start_str})"
                        start_str_runtime = f"({shape_val} + {start_str_runtime})"

                stop_str = str(shape_val)
                stop_str_runtime = str(shape_val)
                if idx.upper is not None:
                    if self._contains_indirect_access(idx.upper):
                        stop_str, stop_str_runtime = self._materialize_indirect_access(
                            idx.upper, return_original_expr=True
                        )
                    else:
                        stop_str = self.visit(idx.upper)
                        stop_str_runtime = stop_str
                    if isinstance(stop_str, str) and (
                        stop_str.startswith("-") or stop_str.startswith("(-")
                    ):
                        stop_str = f"({shape_val} + {stop_str})"
                        stop_str_runtime = f"({shape_val} + {stop_str_runtime})"

                step_str = "1"
                if idx.step is not None:
                    step_str = self.visit(idx.step)

                dim_size = f"({stop_str} - {start_str})"
                dim_size_runtime = f"({stop_str_runtime} - {start_str_runtime})"
                result_shapes.append(dim_size)
                result_shapes_runtime.append(dim_size_runtime)
                slice_info.append((i, start_str, stop_str, step_str))
            else:
                if self._contains_indirect_access(idx):
                    index_str = self._materialize_indirect_access(idx)
                else:
                    index_str = self.visit(idx)
                if isinstance(index_str, str) and (
                    index_str.startswith("-") or index_str.startswith("(-")
                ):
                    index_str = f"({shape_val} + {index_str})"
                index_info.append((i, index_str))

        tmp_name = self.builder.find_new_name("_slice_tmp_")
        result_ndim = len(result_shapes)

        if result_ndim == 0:
            self.builder.add_container(tmp_name, dtype, False)
            self.symbol_table[tmp_name] = dtype
        else:
            size_str = "1"
            for dim in result_shapes:
                size_str = f"({size_str} * {dim})"

            element_size = self.builder.get_sizeof(dtype)
            total_size = f"({size_str} * {element_size})"

            ptr_type = Pointer(dtype)
            self.builder.add_container(tmp_name, ptr_type, False)
            self.symbol_table[tmp_name] = ptr_type
            self.array_info[tmp_name] = {
                "ndim": result_ndim,
                "shapes": result_shapes,
                "shapes_runtime": result_shapes_runtime,
            }

            debug_info = DebugInfo()
            block_alloc = self.builder.add_block(debug_info)
            t_malloc = self.builder.add_malloc(block_alloc, total_size)
            t_ptr = self.builder.add_access(block_alloc, tmp_name, debug_info)
            self.builder.add_memlet(
                block_alloc, t_malloc, "_ret", t_ptr, "void", "", ptr_type, debug_info
            )

        loop_vars = []
        debug_info = DebugInfo()

        for dim_idx, (orig_dim, start_str, stop_str, step_str) in enumerate(slice_info):
            loop_var = self.builder.find_new_name(f"_slice_loop_{dim_idx}_")
            loop_vars.append((loop_var, orig_dim, start_str, step_str))

            if not self.builder.exists(loop_var):
                self.builder.add_container(loop_var, Scalar(PrimitiveType.Int64), False)
                self.symbol_table[loop_var] = Scalar(PrimitiveType.Int64)

            count_str = f"({stop_str} - {start_str})"
            self.builder.begin_for(loop_var, "0", count_str, "1", debug_info)

        src_indices = [""] * ndim
        dst_indices = []

        for orig_dim, index_str in index_info:
            src_indices[orig_dim] = index_str

        for loop_var, orig_dim, start_str, step_str in loop_vars:
            if step_str == "1":
                src_indices[orig_dim] = f"({start_str} + {loop_var})"
            else:
                src_indices[orig_dim] = f"({start_str} + {loop_var} * {step_str})"
            dst_indices.append(loop_var)

        src_linear = self._compute_linear_index(src_indices, shapes, value_str, ndim)
        if result_ndim > 0:
            dst_linear = self._compute_linear_index(
                dst_indices, result_shapes, tmp_name, result_ndim
            )
        else:
            dst_linear = "0"

        block = self.builder.add_block(debug_info)
        t_src = self.builder.add_access(block, value_str, debug_info)
        t_dst = self.builder.add_access(block, tmp_name, debug_info)
        t_task = self.builder.add_tasklet(
            block, TaskletCode.assign, ["_in"], ["_out"], debug_info
        )

        self.builder.add_memlet(
            block, t_src, "void", t_task, "_in", src_linear, None, debug_info
        )
        self.builder.add_memlet(
            block, t_task, "_out", t_dst, "void", dst_linear, None, debug_info
        )

        for _ in loop_vars:
            self.builder.end_for()

        return tmp_name

    def _compute_linear_index(self, indices, shapes, array_name, ndim):
        """Compute linear index from multi-dimensional indices."""
        if ndim == 0:
            return "0"

        linear_index = ""
        for i in range(ndim):
            term = str(indices[i])
            for j in range(i + 1, ndim):
                shape_val = shapes[j] if j < len(shapes) else f"_{array_name}_shape_{j}"
                term = f"(({term}) * {shape_val})"

            if i == 0:
                linear_index = term
            else:
                linear_index = f"({linear_index} + {term})"

        return linear_index

    def _is_array_index(self, node):
        """Check if a node represents an array that could be used as an index (gather)."""
        if isinstance(node, ast.Name):
            return node.id in self.array_info
        return False

    def _handle_gather(self, value_str, index_node, debug_info=None):
        """Handle gather operation: x[indices] where indices is an array."""
        if debug_info is None:
            debug_info = DebugInfo()

        if isinstance(index_node, ast.Name):
            idx_array_name = index_node.id
        else:
            idx_array_name = self.visit(index_node)

        if idx_array_name not in self.array_info:
            raise ValueError(f"Gather index must be an array, got {idx_array_name}")

        idx_shapes = self.array_info[idx_array_name].get("shapes", [])
        idx_ndim = self.array_info[idx_array_name]["ndim"]

        if idx_ndim != 1:
            raise NotImplementedError("Only 1D index arrays supported for gather")

        result_shape = idx_shapes[0] if idx_shapes else f"_{idx_array_name}_shape_0"

        dtype = Scalar(PrimitiveType.Double)
        if value_str in self.symbol_table:
            t = self.symbol_table[value_str]
            if isinstance(t, Pointer) and t.has_pointee_type():
                dtype = t.pointee_type

        idx_dtype = Scalar(PrimitiveType.Int64)
        if idx_array_name in self.symbol_table:
            t = self.symbol_table[idx_array_name]
            if isinstance(t, Pointer) and t.has_pointee_type():
                idx_dtype = t.pointee_type

        tmp_name = self.builder.find_new_name("_gather_")

        element_size = self.builder.get_sizeof(dtype)
        total_size = f"({result_shape} * {element_size})"

        ptr_type = Pointer(dtype)
        self.builder.add_container(tmp_name, ptr_type, False)
        self.symbol_table[tmp_name] = ptr_type
        self.array_info[tmp_name] = {"ndim": 1, "shapes": [result_shape]}

        block_alloc = self.builder.add_block(debug_info)
        t_malloc = self.builder.add_malloc(block_alloc, total_size)
        t_ptr = self.builder.add_access(block_alloc, tmp_name, debug_info)
        self.builder.add_memlet(
            block_alloc, t_malloc, "_ret", t_ptr, "void", "", ptr_type, debug_info
        )

        loop_var = self.builder.find_new_name("_gather_i_")
        self.builder.add_container(loop_var, Scalar(PrimitiveType.Int64), False)
        self.symbol_table[loop_var] = Scalar(PrimitiveType.Int64)

        idx_var = self.builder.find_new_name("_gather_idx_")
        self.builder.add_container(idx_var, idx_dtype, False)
        self.symbol_table[idx_var] = idx_dtype

        self.builder.begin_for(loop_var, "0", str(result_shape), "1", debug_info)

        block_load_idx = self.builder.add_block(debug_info)
        idx_arr_access = self.builder.add_access(
            block_load_idx, idx_array_name, debug_info
        )
        idx_var_access = self.builder.add_access(block_load_idx, idx_var, debug_info)
        tasklet_load = self.builder.add_tasklet(
            block_load_idx, TaskletCode.assign, ["_in"], ["_out"], debug_info
        )
        self.builder.add_memlet(
            block_load_idx,
            idx_arr_access,
            "void",
            tasklet_load,
            "_in",
            loop_var,
            None,
            debug_info,
        )
        self.builder.add_memlet(
            block_load_idx,
            tasklet_load,
            "_out",
            idx_var_access,
            "void",
            "",
            None,
            debug_info,
        )

        block_gather = self.builder.add_block(debug_info)
        src_access = self.builder.add_access(block_gather, value_str, debug_info)
        dst_access = self.builder.add_access(block_gather, tmp_name, debug_info)
        tasklet_gather = self.builder.add_tasklet(
            block_gather, TaskletCode.assign, ["_in"], ["_out"], debug_info
        )

        self.builder.add_memlet(
            block_gather,
            src_access,
            "void",
            tasklet_gather,
            "_in",
            idx_var,
            None,
            debug_info,
        )
        self.builder.add_memlet(
            block_gather,
            tasklet_gather,
            "_out",
            dst_access,
            "void",
            loop_var,
            None,
            debug_info,
        )

        self.builder.end_for()

        return tmp_name

    def _get_max_array_ndim_in_expr(self, node):
        """Get the maximum array dimensionality in an expression."""
        max_ndim = 0

        class NdimVisitor(ast.NodeVisitor):
            def __init__(self, array_info):
                self.array_info = array_info
                self.max_ndim = 0

            def visit_Name(self, node):
                if node.id in self.array_info:
                    ndim = self.array_info[node.id].get("ndim", 0)
                    self.max_ndim = max(self.max_ndim, ndim)
                return self.generic_visit(node)

        visitor = NdimVisitor(self.array_info)
        visitor.visit(node)
        return visitor.max_ndim

    def _handle_broadcast_slice_assignment(
        self, target, value, target_name, indices, target_ndim, value_ndim, debug_info
    ):
        """Handle slice assignment with broadcasting (e.g., 2D -= 1D)."""
        # Number of broadcast dimensions (outer loops)
        broadcast_dims = target_ndim - value_ndim

        shapes = self.array_info[target_name].get("shapes", [])

        # Create outer loops for broadcast dimensions
        outer_loop_vars = []
        for i in range(broadcast_dims):
            loop_var = self.builder.find_new_name(f"_bcast_iter_{i}_")
            outer_loop_vars.append(loop_var)

            if not self.builder.exists(loop_var):
                self.builder.add_container(loop_var, Scalar(PrimitiveType.Int64), False)
                self.symbol_table[loop_var] = Scalar(PrimitiveType.Int64)

            dim_size = shapes[i] if i < len(shapes) else f"_{target_name}_shape_{i}"
            self.builder.begin_for(loop_var, "0", dim_size, "1", debug_info)

        # Create a row view (reference) for the inner dimensions
        row_view_name = self.builder.find_new_name("_row_view_")

        # Get inner shape for the row view
        inner_shapes = shapes[broadcast_dims:] if len(shapes) > broadcast_dims else []

        # Determine element type from the target
        target_type = self.symbol_table.get(target_name)
        if isinstance(target_type, Pointer) and target_type.has_pointee_type():
            element_type = target_type.pointee_type
        else:
            element_type = Scalar(PrimitiveType.Double)

        # Create pointer type for row view
        row_type = Pointer(element_type)
        self.builder.add_container(row_view_name, row_type, False)
        self.symbol_table[row_view_name] = row_type

        # Register row view in array_info
        self.array_info[row_view_name] = {"ndim": value_ndim, "shapes": inner_shapes}

        # Create reference memlet: row_view = &target[i, 0, 0, ...]
        # The index is: outer_loop_vars joined, then zeros for inner dims
        ref_index_parts = outer_loop_vars[:]
        for _ in range(value_ndim):
            ref_index_parts.append("0")

        # Compute linearized index for reference
        # For target[i, j] with shape (n, m), linear index for row i is i * m
        linear_idx = outer_loop_vars[0] if outer_loop_vars else "0"
        for dim_idx in range(1, broadcast_dims):
            dim_size = (
                shapes[dim_idx]
                if dim_idx < len(shapes)
                else f"_{target_name}_shape_{dim_idx}"
            )
            linear_idx = f"({linear_idx}) * ({dim_size}) + {outer_loop_vars[dim_idx]}"

        # Multiply by inner dimension sizes to get the start of the row
        for dim_idx in range(broadcast_dims, target_ndim):
            dim_size = (
                shapes[dim_idx]
                if dim_idx < len(shapes)
                else f"_{target_name}_shape_{dim_idx}"
            )
            linear_idx = f"({linear_idx}) * ({dim_size})"

        # Create the reference memlet block
        block = self.builder.add_block(debug_info)
        t_src = self.builder.add_access(block, target_name, debug_info)
        t_dst = self.builder.add_access(block, row_view_name, debug_info)
        self.builder.add_reference_memlet(
            block, t_src, t_dst, linear_idx, row_type, debug_info
        )

        # Now handle the inner slice assignment with the row view
        # Create inner indices (all slices for the inner dimensions)
        inner_indices = [
            ast.Slice(lower=None, upper=None, step=None) for _ in range(value_ndim)
        ]

        # Create new target using row view
        new_target = ast.Subscript(
            value=ast.Name(id=row_view_name, ctx=ast.Load()),
            slice=(
                ast.Tuple(elts=inner_indices, ctx=ast.Load())
                if len(inner_indices) > 1
                else inner_indices[0]
            ),
            ctx=ast.Store(),
        )

        # Recursively handle the inner assignment (now same-dimension)
        self._handle_slice_assignment(
            new_target, value, row_view_name, inner_indices, debug_info
        )

        # Close outer loops
        for _ in outer_loop_vars:
            self.builder.end_for()

    def _handle_slice_assignment(
        self, target, value, target_name, indices, debug_info=None
    ):
        if debug_info is None:
            debug_info = DebugInfo()

        if target_name in self.array_info:
            ndim = self.array_info[target_name]["ndim"]
            if len(indices) < ndim:
                indices = list(indices)
                for _ in range(ndim - len(indices)):
                    indices.append(ast.Slice(lower=None, upper=None, step=None))

        # Check if the RHS contains a ufunc outer operation
        # If so, we handle it specially to avoid the loop transformation
        # which would destroy the slice shape information
        has_outer, ufunc_name, outer_node = contains_ufunc_outer(value)
        if has_outer:
            self._handle_ufunc_outer_slice_assignment(
                target, value, target_name, indices, debug_info
            )
            return

        # Count slice dimensions to determine effective target dimensionality
        # (slice indices produce array dimensions, point indices collapse them)
        target_slice_ndim = sum(1 for idx in indices if isinstance(idx, ast.Slice))
        value_max_ndim = self._get_max_array_ndim_in_expr(value)

        if (
            target_slice_ndim > 0
            and value_max_ndim > 0
            and target_slice_ndim > value_max_ndim
        ):
            # Broadcasting case: use row-by-row approach with reference memlets
            self._handle_broadcast_slice_assignment(
                target,
                value,
                target_name,
                indices,
                target_slice_ndim,
                value_max_ndim,
                debug_info,
            )
            return

        loop_vars = []
        new_target_indices = []

        for i, idx in enumerate(indices):
            if isinstance(idx, ast.Slice):
                loop_var = self.builder.find_new_name(f"_slice_iter_{len(loop_vars)}_")
                loop_vars.append(loop_var)

                if not self.builder.exists(loop_var):
                    self.builder.add_container(
                        loop_var, Scalar(PrimitiveType.Int64), False
                    )
                    self.symbol_table[loop_var] = Scalar(PrimitiveType.Int64)

                start_str = "0"
                if idx.lower:
                    start_str = self.visit(idx.lower)
                    if start_str.startswith("-"):
                        shapes = self.array_info[target_name].get("shapes", [])
                        dim_size = (
                            str(shapes[i])
                            if i < len(shapes)
                            else f"_{target_name}_shape_{i}"
                        )
                        start_str = f"({dim_size} {start_str})"

                stop_str = ""
                if idx.upper and not (
                    isinstance(idx.upper, ast.Constant) and idx.upper.value is None
                ):
                    stop_str = self.visit(idx.upper)
                    if stop_str.startswith("-") or stop_str.startswith("(-"):
                        shapes = self.array_info[target_name].get("shapes", [])
                        dim_size = (
                            str(shapes[i])
                            if i < len(shapes)
                            else f"_{target_name}_shape_{i}"
                        )
                        stop_str = f"({dim_size} {stop_str})"
                else:
                    shapes = self.array_info[target_name].get("shapes", [])
                    stop_str = (
                        str(shapes[i])
                        if i < len(shapes)
                        else f"_{target_name}_shape_{i}"
                    )

                step_str = "1"
                if idx.step:
                    step_str = self.visit(idx.step)

                count_str = f"({stop_str} - {start_str})"

                self.builder.begin_for(loop_var, "0", count_str, "1", debug_info)
                self.symbol_table[loop_var] = Scalar(PrimitiveType.Int64)

                new_target_indices.append(
                    ast.Name(
                        id=f"{start_str} + {loop_var} * {step_str}", ctx=ast.Load()
                    )
                )
            else:
                # Handle non-slice indices - need to normalize negative indices
                shapes = self.array_info[target_name].get("shapes", [])
                dim_size = shapes[i] if i < len(shapes) else f"_{target_name}_shape_{i}"
                normalized_idx = normalize_negative_index(idx, dim_size)
                new_target_indices.append(normalized_idx)

        rewriter = SliceRewriter(loop_vars, self.array_info, self)
        new_value = rewriter.visit(copy.deepcopy(value))

        new_target = copy.deepcopy(target)
        if len(new_target_indices) == 1:
            new_target.slice = new_target_indices[0]
        else:
            new_target.slice = ast.Tuple(elts=new_target_indices, ctx=ast.Load())

        target_str = self.visit(new_target)
        value_str = self.visit(new_value)
        self.builder.add_assignment(target_str, value_str, debug_info)

        for _ in loop_vars:
            self.builder.end_for()

    def _handle_ufunc_outer_slice_assignment(
        self, target, value, target_name, indices, debug_info=None
    ):
        """Handle slice assignment where RHS contains a ufunc outer operation.

        Example: path[:] = np.minimum(path[:], np.add.outer(path[:, k], path[k, :]))

        The strategy is:
        1. Evaluate the entire RHS expression, which will create a temporary array
           containing the result of the ufunc outer (potentially wrapped in other ops)
        2. Copy the temporary result to the target slice

        This avoids the loop transformation that would destroy slice shape info.
        """
        if debug_info is None:
            from docc.sdfg import DebugInfo

            debug_info = DebugInfo()

        # Evaluate the full RHS expression
        # This will:
        # - Create temp arrays for ufunc outer results
        # - Apply any wrapping operations (np.minimum, etc.)
        # - Return the name of the final result array
        result_name = self.visit(value)

        # Now we need to copy result to target slice
        # Count slice dimensions to determine if we need loops
        target_slice_ndim = sum(1 for idx in indices if isinstance(idx, ast.Slice))

        if target_slice_ndim == 0:
            # No slices on target - just simple assignment
            target_str = self.visit(target)
            block = self.builder.add_block(debug_info)
            t_src, src_sub = self._add_read(block, result_name, debug_info)
            t_dst = self.builder.add_access(block, target_str, debug_info)
            t_task = self.builder.add_tasklet(
                block, TaskletCode.assign, ["_in"], ["_out"], debug_info
            )
            self.builder.add_memlet(
                block, t_src, "void", t_task, "_in", src_sub, None, debug_info
            )
            self.builder.add_memlet(
                block, t_task, "_out", t_dst, "void", "", None, debug_info
            )
            return

        # We have slices on the target - need to create loops for copying
        # Get target array info
        target_info = self.array_info.get(target_name, {})
        target_shapes = target_info.get("shapes", [])

        loop_vars = []
        new_target_indices = []

        for i, idx in enumerate(indices):
            if isinstance(idx, ast.Slice):
                loop_var = self.builder.find_new_name(f"_copy_iter_{len(loop_vars)}_")
                loop_vars.append(loop_var)

                if not self.builder.exists(loop_var):
                    self.builder.add_container(
                        loop_var, Scalar(PrimitiveType.Int64), False
                    )
                    self.symbol_table[loop_var] = Scalar(PrimitiveType.Int64)

                start_str = "0"
                if idx.lower:
                    start_str = self.visit(idx.lower)

                stop_str = ""
                if idx.upper and not (
                    isinstance(idx.upper, ast.Constant) and idx.upper.value is None
                ):
                    stop_str = self.visit(idx.upper)
                else:
                    stop_str = (
                        target_shapes[i]
                        if i < len(target_shapes)
                        else f"_{target_name}_shape_{i}"
                    )

                step_str = "1"
                if idx.step:
                    step_str = self.visit(idx.step)

                count_str = f"({stop_str} - {start_str})"

                self.builder.begin_for(loop_var, "0", count_str, "1", debug_info)
                self.symbol_table[loop_var] = Scalar(PrimitiveType.Int64)

                new_target_indices.append(
                    ast.Name(
                        id=f"{start_str} + {loop_var} * {step_str}", ctx=ast.Load()
                    )
                )
            else:
                # Handle non-slice indices - need to normalize negative indices
                dim_size = (
                    target_shapes[i]
                    if i < len(target_shapes)
                    else f"_{target_name}_shape_{i}"
                )
                normalized_idx = normalize_negative_index(idx, dim_size)
                new_target_indices.append(normalized_idx)

        # Create assignment block: target[i,j,...] = result[i,j,...]
        block = self.builder.add_block(debug_info)

        # Access nodes
        t_src = self.builder.add_access(block, result_name, debug_info)
        t_dst = self.builder.add_access(block, target_name, debug_info)
        t_task = self.builder.add_tasklet(
            block, TaskletCode.assign, ["_in"], ["_out"], debug_info
        )

        # Source index - just use loop vars for flat array from ufunc outer
        # The ufunc outer result is a flat array of size M*N
        if len(loop_vars) == 2:
            # 2D case: result is indexed as i * N + j
            # Get the second dimension size from target shapes
            n_dim = (
                target_shapes[1]
                if len(target_shapes) > 1
                else f"_{target_name}_shape_1"
            )
            src_index = f"(({loop_vars[0]}) * ({n_dim}) + ({loop_vars[1]}))"
        elif len(loop_vars) == 1:
            src_index = loop_vars[0]
        else:
            # General case - compute linear index
            src_terms = []
            stride = "1"
            for i in range(len(loop_vars) - 1, -1, -1):
                if stride == "1":
                    src_terms.insert(0, loop_vars[i])
                else:
                    src_terms.insert(0, f"({loop_vars[i]} * {stride})")
                if i > 0:
                    dim_size = (
                        target_shapes[i]
                        if i < len(target_shapes)
                        else f"_{target_name}_shape_{i}"
                    )
                    stride = (
                        f"({stride} * {dim_size})" if stride != "1" else str(dim_size)
                    )
            src_index = " + ".join(src_terms) if src_terms else "0"

        # Target index - compute linear index (row-major order)
        # For 2D array with shape (M, N): linear_index = i * N + j
        target_index_parts = []
        for idx in new_target_indices:
            if isinstance(idx, ast.Name):
                target_index_parts.append(idx.id)
            else:
                target_index_parts.append(self.visit(idx))

        # Convert to linear index
        if len(target_index_parts) == 2:
            # 2D case
            n_dim = (
                target_shapes[1]
                if len(target_shapes) > 1
                else f"_{target_name}_shape_1"
            )
            target_index = (
                f"(({target_index_parts[0]}) * ({n_dim}) + ({target_index_parts[1]}))"
            )
        elif len(target_index_parts) == 1:
            target_index = target_index_parts[0]
        else:
            # General case - compute linear index with strides
            stride = "1"
            target_index = "0"
            for i in range(len(target_index_parts) - 1, -1, -1):
                idx_part = target_index_parts[i]
                if stride == "1":
                    term = idx_part
                else:
                    term = f"(({idx_part}) * ({stride}))"

                if target_index == "0":
                    target_index = term
                else:
                    target_index = f"({term} + {target_index})"

                if i > 0:
                    dim_size = (
                        target_shapes[i]
                        if i < len(target_shapes)
                        else f"_{target_name}_shape_{i}"
                    )
                    stride = (
                        f"({stride} * {dim_size})" if stride != "1" else str(dim_size)
                    )

        # Connect memlets
        self.builder.add_memlet(
            block, t_src, "void", t_task, "_in", src_index, None, debug_info
        )
        self.builder.add_memlet(
            block, t_task, "_out", t_dst, "void", target_index, None, debug_info
        )

        # End loops
        for _ in loop_vars:
            self.builder.end_for()

    def _is_indirect_access(self, node):
        """Check if a node represents an indirect array access (e.g., A[B[i]]).

        Returns True if the node is a subscript where the index itself is a subscript
        into an array (indirect access pattern).
        """
        if not isinstance(node, ast.Subscript):
            return False
        if isinstance(node.value, ast.Name):
            arr_name = node.value.id
            if arr_name in self.array_info:
                if isinstance(node.slice, ast.Subscript):
                    if isinstance(node.slice.value, ast.Name):
                        idx_arr_name = node.slice.value.id
                        if idx_arr_name in self.array_info:
                            return True
        return False

    def _contains_indirect_access(self, node):
        """Check if an AST node contains any indirect array access."""
        if isinstance(node, ast.Subscript):
            if isinstance(node.value, ast.Name):
                arr_name = node.value.id
                if arr_name in self.array_info:
                    return True
        elif isinstance(node, ast.BinOp):
            return self._contains_indirect_access(
                node.left
            ) or self._contains_indirect_access(node.right)
        elif isinstance(node, ast.UnaryOp):
            return self._contains_indirect_access(node.operand)
        return False

    def _materialize_indirect_access(
        self, node, debug_info=None, return_original_expr=False
    ):
        """Materialize an array access into a scalar variable using tasklet+memlets."""
        if not self.builder:
            expr = self.visit(node)
            return (expr, expr) if return_original_expr else expr

        if debug_info is None:
            debug_info = DebugInfo()

        if not isinstance(node, ast.Subscript):
            expr = self.visit(node)
            return (expr, expr) if return_original_expr else expr

        if not isinstance(node.value, ast.Name):
            expr = self.visit(node)
            return (expr, expr) if return_original_expr else expr

        arr_name = node.value.id
        if arr_name not in self.array_info:
            expr = self.visit(node)
            return (expr, expr) if return_original_expr else expr

        dtype = Scalar(PrimitiveType.Int64)
        if arr_name in self.symbol_table:
            t = self.symbol_table[arr_name]
            if isinstance(t, Pointer) and t.has_pointee_type():
                dtype = t.pointee_type

        tmp_name = self.builder.find_new_name("_idx_")
        self.builder.add_container(tmp_name, dtype, False)
        self.symbol_table[tmp_name] = dtype

        ndim = self.array_info[arr_name]["ndim"]
        shapes = self.array_info[arr_name].get("shapes", [])

        if isinstance(node.slice, ast.Tuple):
            indices = [self.visit(elt) for elt in node.slice.elts]
        else:
            indices = [self.visit(node.slice)]

        materialized_indices = []
        for idx_str in indices:
            if "(" in idx_str and idx_str.endswith(")"):
                materialized_indices.append(idx_str)
            else:
                materialized_indices.append(idx_str)

        linear_index = self._compute_linear_index(
            materialized_indices, shapes, arr_name, ndim
        )

        block = self.builder.add_block(debug_info)
        t_src = self.builder.add_access(block, arr_name, debug_info)
        t_dst = self.builder.add_access(block, tmp_name, debug_info)
        t_task = self.builder.add_tasklet(
            block, TaskletCode.assign, ["_in"], ["_out"], debug_info
        )

        self.builder.add_memlet(
            block, t_src, "void", t_task, "_in", linear_index, None, debug_info
        )
        self.builder.add_memlet(
            block, t_task, "_out", t_dst, "void", "", None, debug_info
        )

        if return_original_expr:
            original_expr = f"{arr_name}({linear_index})"
            return (tmp_name, original_expr)

        return tmp_name

    def _get_unique_id(self):
        self._unique_counter_ref[0] += 1
        return self._unique_counter_ref[0]

    def _element_type(self, name):
        if name in self.symbol_table:
            return element_type_from_sdfg_type(self.symbol_table[name])
        else:  # Constant
            if self._is_int(name):
                return Scalar(PrimitiveType.Int64)
            else:
                return Scalar(PrimitiveType.Double)

    def _is_int(self, operand):
        try:
            if operand.lstrip("-").isdigit():
                return True
        except ValueError:
            pass

        name = operand
        if "(" in operand and operand.endswith(")"):
            name = operand.split("(")[0]

        if name in self.symbol_table:
            t = self.symbol_table[name]

            def is_int_ptype(pt):
                return pt in [
                    PrimitiveType.Int64,
                    PrimitiveType.Int32,
                    PrimitiveType.Int8,
                    PrimitiveType.Int16,
                    PrimitiveType.UInt64,
                    PrimitiveType.UInt32,
                    PrimitiveType.UInt8,
                    PrimitiveType.UInt16,
                ]

            if isinstance(t, Scalar):
                return is_int_ptype(t.primitive_type)

            if type(t).__name__ == "Array" and hasattr(t, "element_type"):
                et = t.element_type
                if callable(et):
                    et = et()
                if isinstance(et, Scalar):
                    return is_int_ptype(et.primitive_type)

            if type(t).__name__ == "Pointer":
                if hasattr(t, "pointee_type"):
                    et = t.pointee_type
                    if callable(et):
                        et = et()
                    if isinstance(et, Scalar):
                        return is_int_ptype(et.primitive_type)
                if hasattr(t, "element_type"):
                    et = t.element_type
                    if callable(et):
                        et = et()
                    if isinstance(et, Scalar):
                        return is_int_ptype(et.primitive_type)

        return False

    def _add_read(self, block, expr_str, debug_info=None):
        try:
            if (block, expr_str) in self._access_cache:
                return self._access_cache[(block, expr_str)]
        except TypeError:
            pass

        if debug_info is None:
            debug_info = DebugInfo()

        if "(" in expr_str and expr_str.endswith(")"):
            name = expr_str.split("(")[0]
            subset = expr_str[expr_str.find("(") + 1 : -1]
            access = self.builder.add_access(block, name, debug_info)
            try:
                self._access_cache[(block, expr_str)] = (access, subset)
            except TypeError:
                pass
            return access, subset

        if self.builder.exists(expr_str):
            access = self.builder.add_access(block, expr_str, debug_info)
            subset = ""
            if expr_str in self.symbol_table:
                sym_type = self.symbol_table[expr_str]
                if isinstance(sym_type, Pointer):
                    if expr_str in self.array_info:
                        ndim = self.array_info[expr_str].get("ndim", 0)
                        if ndim == 0:
                            subset = "0"
                    else:
                        subset = "0"
            try:
                self._access_cache[(block, expr_str)] = (access, subset)
            except TypeError:
                pass
            return access, subset

        dtype = Scalar(PrimitiveType.Double)
        if self._is_int(expr_str):
            dtype = Scalar(PrimitiveType.Int64)
        elif expr_str == "true" or expr_str == "false":
            dtype = Scalar(PrimitiveType.Bool)

        const_node = self.builder.add_constant(block, expr_str, dtype, debug_info)
        try:
            self._access_cache[(block, expr_str)] = (const_node, "")
        except TypeError:
            pass
        return const_node, ""

daisytuner / docc / 22023884668

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous