diff --git a/corpus/baseline.json b/corpus/baseline.json index 48e6dfb..70891b3 100644 --- a/corpus/baseline.json +++ b/corpus/baseline.json @@ -1367,11 +1367,11 @@ "stopped_by": "static.rwset", "verdicts": { "static.rwset": { - "detail": "7/66 blocks disagree: cspl_interp/B002, cspleps/B003, csplint/B006, csplint/B012, csplint/B016 (+2 more)", + "detail": "5/66 blocks disagree: csplint/B006, csplint/B012, csplint/B016, csplint_square/B011, csplint_square/B015", "metrics": { "blocks_checked": 66, "blocks_deferred": 14, - "blocks_matched": 59, + "blocks_matched": 61, "blocks_waived": 0 }, "passed": false @@ -1418,11 +1418,11 @@ "stopped_by": "static.rwset", "verdicts": { "static.rwset": { - "detail": "18/124 blocks disagree: splrep_msg/B002, splprep/B014, splprep/B016, splprep/B019, splprep/B020 (+13 more)", + "detail": "1/124 blocks disagree: splrep_msg/B002", "metrics": { "blocks_checked": 124, "blocks_deferred": 8, - "blocks_matched": 106, + "blocks_matched": 123, "blocks_waived": 0 }, "passed": false @@ -1650,11 +1650,11 @@ "stopped_by": "static.rwset", "verdicts": { "static.rwset": { - "detail": "7/187 blocks disagree: genrand_encode/B002, genrand_encode/B003, genrand_decode/B002, genrand_load_state/B003, genrand_load_state/B004 (+2 more)", + "detail": "3/187 blocks disagree: genrand_encode/B002, genrand_encode/B003, genrand_decode/B002", "metrics": { "blocks_checked": 187, "blocks_deferred": 3, - "blocks_matched": 180, + "blocks_matched": 184, "blocks_waived": 0 }, "passed": false @@ -1904,23 +1904,42 @@ "oracle/f2py-golden": "ok", "store/fs-evidence": "ok", "transform/translate.numpy": "ok", - "verifier/differential.bitexact": "ok", - "verifier/static.rwset": "ok", - "verifier/symbolic.notary": "ok" + "verifier/differential.bitexact": "failed", + "verifier/static.rwset": "ok" }, - "stopped_by": null, + "stopped_by": "differential.bitexact", "verdicts": { "differential.bitexact": { - "detail": "80 points across 1 subprogram(s), all bit-exact", + "detail": "3 subprogram(s) could not be compared: searchsorted_dp: candidate raised: NameError: name 'SMALL' is not defined; searchsorted_idp: candidate raised: NameError:", "metrics": { - "bit_exact": 80, + "bit_exact": 90, "integer_mismatch": 0, - "integer_points": 0, + "integer_points": 10, "max_rel": 0.0, "max_ulp": 0, "nan_mismatch": 0, - "points": 80, + "points": 90, "subprograms": { + "searchsorted_dp": { + "error": "candidate raised: NameError: name 'SMALL' is not defined" + }, + "searchsorted_dpi": { + "error": "candidate raised: NameError: name 'SMALL' is not defined" + }, + "searchsorted_i": { + "bit_exact": 10, + "integer_mismatch": 0, + "integer_points": 10, + "max_rel": 0.0, + "max_ulp": 0, + "nan_mismatch": 0, + "points": 10, + "redrawn": 0, + "reshaped": 0 + }, + "searchsorted_idp": { + "error": "candidate raised: NameError: name 'SMALL' is not defined" + }, "sort": { "bit_exact": 80, "integer_mismatch": 0, @@ -1935,7 +1954,7 @@ }, "trials": 10 }, - "passed": true + "passed": false }, "static.rwset": { "detail": "62 blocks match", @@ -1946,13 +1965,6 @@ "blocks_waived": 0 }, "passed": true - }, - "symbolic.notary": { - "detail": "no rewrites to notarize; the translation is print-order faithful", - "metrics": { - "rewrites": 0 - }, - "passed": true } } }, @@ -2360,11 +2372,11 @@ "stopped_by": "static.rwset", "verdicts": { "static.rwset": { - "detail": "15/106 blocks disagree: slsqp/B019, slsqpb/B002, slsqpb/B003, slsqpb/B006, reset_bfgs_matrix/B002 (+10 more)", + "detail": "7/106 blocks disagree: slsqpb/B002, slsqpb/B003, slsqpb/B006, reset_bfgs_matrix/B002, lsq/B026 (+2 more)", "metrics": { "blocks_checked": 106, "blocks_deferred": 8, - "blocks_matched": 91, + "blocks_matched": 99, "blocks_waived": 0 }, "passed": false diff --git a/src/recast/fortran/expr.py b/src/recast/fortran/expr.py index 6403575..e93076b 100644 --- a/src/recast/fortran/expr.py +++ b/src/recast/fortran/expr.py @@ -156,7 +156,9 @@ def substitute(expr: Expr, name: str, replacement: Expr) -> Expr: For the one legal self-reference in an initializer, a kind inquiry on the constant being declared (``tol = max( 1.e-10_r8, epsilon(tol) )``): the reference carries the constant's kind and nothing else, and the fold - renders reals as 64-bit, so a 64-bit literal stands in for it. + renders reals as 64-bit, so a 64-bit literal stands in for it. Every + occurrence is replaced: a parameter cannot name itself anywhere else in + its own initializer, so there is no other occurrence to preserve. """ if expr.kind == "name" and expr.text == name: return replacement @@ -205,27 +207,28 @@ def render( REAL_CALLS = frozenset({"real", "dble", "sqrt"} | KIND_INQUIRIES) -def typed(expr: Expr) -> str | None: +def typed(expr: Expr, env: dict[str, str | None] | None = None) -> str | None: """``"real"``, ``"int"``, or ``None`` when a bare name leaves it open. Type inference the fold needs for exactly one decision: whether a ``/`` is Fortran's integer division. A real literal or a real-valued call anywhere in an operand makes the quotient real; ``int(...)`` and integer - literals make it integer; a name is whatever its initializer was, which - this tree does not carry. + literals make it integer; a name is what ``env`` says its declaration + was -- CLUBB's ``ep = Rd / Rv`` over two real parameters is a real + quotient, and a fold that guessed integer made it zero. """ if expr.kind == "real": return "real" if expr.kind == "int": return "int" if expr.kind == "name": - return None + return (env or {}).get(expr.text) if expr.kind == "call": if expr.text == "int": return "int" if expr.text in REAL_CALLS: return "real" - kinds = {typed(a) for a in expr.args} + kinds = {typed(a, env) for a in expr.args} if "real" in kinds: return "real" if kinds == {"int"}: @@ -233,24 +236,27 @@ def typed(expr: Expr) -> str | None: return None -def with_integer_division(expr: Expr, *, default_integer: bool | None = None) -> Expr: +def with_integer_division( + expr: Expr, *, default_integer: bool | None = None, env: dict[str, str | None] | None = None +) -> Expr: """The tree with every integer ``/`` spelled ``//``. Fortran divides two integers to an integer: ``nrk = runge_kutta_type / 10`` is 4, not 4.1. A quotient whose operands are both known integers is - marked; one with a name in it falls back to ``default_integer``, which - the caller sets from the whole initializer -- an expression with no real - literal and no real-valued call in it is integer arithmetic throughout, - because a name in it is an integer parameter or it would have had one. + marked; one with a name in it is typed by ``env`` (the declared types + of the constants resolved so far) and otherwise falls back to + ``default_integer``, which the caller sets from the whole initializer. """ if default_integer is None: - default_integer = typed(expr) != "real" + default_integer = typed(expr, env) != "real" if not expr.args: return expr - args = tuple(with_integer_division(a, default_integer=default_integer) for a in expr.args) + args = tuple( + with_integer_division(a, default_integer=default_integer, env=env) for a in expr.args + ) text = expr.text if expr.kind == "binary" and expr.text == "/": - kinds = {typed(a) for a in args} + kinds = {typed(a, env) for a in args} if kinds == {"int"} or ("real" not in kinds and default_integer): text = "//" return Expr(expr.kind, text, args) diff --git a/src/recast/fortran/flatten.py b/src/recast/fortran/flatten.py index e6527cc..100ef42 100644 --- a/src/recast/fortran/flatten.py +++ b/src/recast/fortran/flatten.py @@ -62,7 +62,8 @@ ] DERIVED = re.compile(r"UNKNOWN\(TYPE\((\w+)\)\)", re.IGNORECASE) -ALLOCATE_THIS = re.compile(r"allocate\s*\(\s*this\s*%\s*(\w+)\s*\(([^()]*)\)", re.I) +ALLOCATE_STMT = re.compile(r"\ballocate\s*\(", re.I) +COMPONENT_ALLOCATION = re.compile(r"(\w+)\s*%\s*(\w+)\s*\(([^()]*)\)") FORTRAN_TYPES = {"float64": "real(8)", "float32": "real(4)", "int32": "integer", "bool": "logical"} @@ -145,6 +146,20 @@ class FlatPlan: dim_constants: dict[str, int] = field(default_factory=dict) """Named extents of the original dummies that are tree constants (``a(nrk,nrk)``), so the flat signature can spell them as numbers.""" + extent_args: dict[str, list[Any]] = field(default_factory=dict) + """``{name: [object, component, axis]}``: an integer argument of the + adapter carrying an extent the plan could not spell -- an allocation + sized by the allocating routine's dummy when the planned subprogram + has none of that name (CLUBB's ``coef_wp4_implicit(1:ngrdcol, 1:nz)`` + under ``advance_wp2_wp3``, which takes ``nzm`` and ``nzt``). The + recorder writes it from ``size()``; the sampled gate sizes it like any + scalar an array's dims name.""" + + left_to_module: list[str] = field(default_factory=list) + """What the body reaches that the adapter does not carry: module state + its module keeps private (CLUBB's ``error_code % clubb_debug_level``), + a character component (``err_info % err_header``). Both sides run with + the default, and the plan says so rather than failing.""" patch_count: str = "np_" counter_prefix: str = "num_" @@ -152,6 +167,17 @@ class FlatPlan: def name(self) -> str: return f"{self.subprogram['name']}_flat" + def _component_names(self, text: str) -> str: + """``obj%comp`` in a bound -> the component's flat name, when the + plan carries it; left as written otherwise.""" + carried = {(o.name, c.name): c.flat for o in self.objects for c in o.components} + + def swap(match: re.Match[str]) -> str: + key = (match.group(1).lower(), match.group(2).lower()) + return carried.get(key, match.group(0)) + + return re.sub(r"(\w+)\s*%\s*(\w+)", swap, text) + @property def usable(self) -> bool: return not self.unsupported and bool(self.objects) @@ -183,6 +209,8 @@ def from_dict(cls, data: dict[str, Any]) -> FlatPlan: unsupported=list(data.get("unsupported", [])), states=states, dim_constants=dict(data.get("dim_constants", {})), + left_to_module=list(data.get("left_to_module", [])), + extent_args={k: list(v) for k, v in (data.get("extent_args") or {}).items()}, patch_count=data.get("patch_count", "np_"), counter_prefix=data.get("counter_prefix", "num_"), ) @@ -211,22 +239,33 @@ def flat_args(self) -> list[dict[str, Any]]: for dim in entry["dims"]: if dim.get("ub"): ub = str(dim["ub"]).strip().lower() - sized.append({**dim, "ub": str(self.dim_constants.get(ub, dim["ub"]))}) + # ``thvm(ngrdcol, gr%nzt)`` (CLUBB's calc_pressure): a + # dummy sized by a component of the object, spelled + # by that component's flat name, an argument here. + ub = self._component_names(ub) + sized.append({**dim, "ub": str(self.dim_constants.get(ub, ub))}) else: counter = f"{self.counter_prefix}{entry['name'].lower()}" extent = counter if counter in names else self.patch_count sized.append({"lb": "1", "ub": extent}) entry["dims"] = sized args.append(entry) - args.append( - { - "name": self.patch_count, - "dtype": "int32", - "intent": "IN", - "optional": False, - "dims": None, - } - ) + if not any(a["name"].lower() == self.patch_count for a in args): + # CLUBB passes ngrdcol explicitly: the driver's extent is a dummy + # already there, and declaring it twice does not compile. + args.append( + { + "name": self.patch_count, + "dtype": "int32", + "intent": "IN", + "optional": False, + "dims": None, + } + ) + for name in self.extent_args: + args.append( + {"name": name, "dtype": "int32", "intent": "IN", "optional": False, "dims": None} + ) for obj in self.objects: for comp in obj.components: args.append( @@ -313,14 +352,51 @@ def _state_declaration(name: str, root: Path) -> tuple[str, str] | None: continue # a dummy of that type, not the module's variable module = MODULE_DEFINITION.search(text) return match.group(1).lower(), (module.group(1).lower() if module else "") + # Declared in a list over several lines (CLUBB's four sponge settings + # under one ``type(sponge_damp_settings), public :: &``): the module + # record has read the declaration whole. + for path in sources(root): + record = _module_record(path, {}) + for entry in (record or {}).get("module_state", ()): + if str(entry.get("name", "")).lower() != name.lower(): + continue + derived = DERIVED.match(str(entry.get("dtype", ""))) + if derived: + return derived.group(1).lower(), str((record or {}).get("module", "")).lower() return None def _allocation_bounds(path: Path) -> dict[str, list[str]]: - """``component -> [axis bound text, ...]`` from ``allocate (this%c (…))``.""" + """``component -> [axis bound text, ...]`` from the module's ALLOCATE + statements: ``allocate (this%c (…))`` in the CLM family, and CLUBB's + ``allocate( gr%zm(ngrdcol,gr%nzm), gr%zt(ngrdcol,gr%nzt), … )`` -- one + statement over many components of an object named however the setup + routine names its dummy. Continuation lines and comments are folded + first; the first allocation of a component wins.""" out: dict[str, list[str]] = {} - for match in ALLOCATE_THIS.finditer(path.read_text(errors="replace")): - out.setdefault(match.group(1).lower(), [b.strip() for b in match.group(2).split(",")]) + text = path.read_text(errors="replace") + for start in ALLOCATE_STMT.finditer(text): + # The statement: from ``allocate(`` to its matching parenthesis, + # across ``&`` continuations, with trailing comments dropped. + depth, at = 0, start.end() - 1 + while at < len(text): + character = text[at] + if character == "!": + at = text.find("\n", at) + if at < 0: + break + continue + if character == "(": + depth += 1 + elif character == ")": + depth -= 1 + if depth == 0: + break + at += 1 + body = text[start.end() : at] + for item in COMPONENT_ALLOCATION.finditer(body.replace("&", " ")): + axes = [b.strip() for b in item.group(3).split(",")] + out.setdefault(item.group(2).lower(), axes) return out @@ -340,6 +416,11 @@ def spell(text: str) -> str | None: return text if text in constants: return str(constants[text]) + # A scalar component of the object itself -- CLUBB allocates + # ``gr%zm(ngrdcol, gr%nzm)`` -- spelled by that component's flat + # name once the plan's objects are known (``_bind_symbolic_extents``). + if re.fullmatch(r"\w+\s*%\s*\w+", text): + return text # An arithmetic bound over constants: ``-nlevsno+1``. expression = re.sub( r"[A-Za-z_]\w*", lambda m: str(constants.get(m.group(0).lower(), m.group(0))), text @@ -363,6 +444,8 @@ def spell(text: str) -> str | None: extent = patch if low == "1" else None elif re.fullmatch(r"-?\d+", low) and re.fullmatch(r"-?\d+", high): extent = str(int(high) - int(low) + 1) + elif low == "1": + extent = high else: extent = f"({high}) - ({low}) + 1" if extent is None: @@ -531,13 +614,39 @@ def take(stmt: Any) -> None: interesting = objects | set(aliases) # ``dummy = hybrid(...)`` parses as any of three node kinds depending # on what fparser could tell about the name. + # A call spelled by a generic (CLUBB's ``zt2zm_api`` over grid_class's + # specifics) is followed into every specific: which one Fortran + # picks depends on ranks this walk does not resolve, and the union + # of what they touch is what the adapter has to carry. + generics: dict[str, list[str]] = {} + for _path, module_record_, _sub in procedures.values(): + for generic, specifics in (module_record_.get("generics") or {}).items(): + generics.setdefault(generic.lower(), [n.lower() for n in specifics]) + expanded: list[tuple[Any, str]] = [] for call in [ *walk(node, f03.Call_Stmt), *walk(node, f03.Part_Ref), *walk(node, f03.Function_Reference), *walk(node, f03.Structure_Constructor), ]: - callee = str(call.children[0]).lower() + spelled_callee = str(call.children[0]).lower() + candidates = generics.get(spelled_callee) + if candidates is None: + expanded.append((call, spelled_callee)) + continue + # The specifics a call of this arity can reach: Fortran picks by + # rank too, which this walk does not resolve, so the union of + # what the arity-compatible ones touch is what the adapter carries. + given = len(call.children[1].children) if call.children[1] is not None else 0 + arity = { + specific: len(procedures[specific][2]["args"]) + for specific in candidates + if specific in procedures + } + exact = [n for n, count in arity.items() if count == given] + fitting = exact or [n for n, count in arity.items() if count >= given] + expanded.extend((call, callee) for callee in fitting) + for call, callee in expanded: if callee not in procedures: continue path, module_record, callee_record = procedures[callee] @@ -562,6 +671,20 @@ def take(stmt: Any) -> None: spelled = aliases[spelled].split("%", 1)[0] if spelled in objects or spelled in interesting: mapping[dummy] = spelled + elif ( + spelled not in declared_here + and spelled not in procedures + and any( + a["name"].lower() == dummy and DERIVED.match(str(a.get("dtype", ""))) + for a in callee_record["args"] + ) + ): + # A module-state object handed whole to a derived-type + # dummy (CLUBB passes sponge_layer_damping's profile to + # its sponge_damp_xm): followed under the callee's + # dummy, so the components the callee reads come back + # under the state's own name and the plan carries them. + mapping[dummy] = spelled # A procedure passed as an actual (``hybrid(..., func, ...)``) is # called back with the object; its own derived-type dummies are # mapped by name, which is how the model spells them. @@ -612,6 +735,7 @@ def take(stmt: Any) -> None: visited.add(key) callee_node = _subprogram_node(path, callee) if callee_node is None: + visited.discard(key) continue _, inner_reads, inner_writes = _accesses( callee_node, @@ -625,6 +749,9 @@ def take(stmt: Any) -> None: externals, companions, ) + # A guard against cycles, not a memo: the same callee followed + # from another site, under another mapping, is followed again. + visited.discard(key) callee_dummies = {a["name"].lower() for a in callee_record["args"]} for inner, target in ((inner_reads, reads), (inner_writes, writes)): for item in inner: @@ -776,6 +903,10 @@ def type_info(type_name: str) -> tuple[dict[str, Any], dict[str, list[str]]] | N touched.setdefault(obj, {}).setdefault(member, False) if name in writes: touched[obj][member] = True + # Every dummy's type first, so a bound in one object's allocation + # naming another object's component finds that type on record. + for type_name in dummies.values(): + type_info(type_name) for obj in sorted(touched): if obj in dummies: flat = FlatObject(name=obj, type_name=dummies[obj], kind="dummy") @@ -802,11 +933,37 @@ def type_info(type_name: str) -> tuple[dict[str, Any], dict[str, list[str]]] | N (type_files[flat.type_name],), kinds, ) + # A component named by a bound is carried as an input whether + # or not the body reads it: by a touched component's allocation + # (``gr%zm(ngrdcol, gr%nzm)``), by another object's + # (``damping_profile%tau_sponge_damp(gr%nzm)``), or by a dummy's + # declaration (``thvm(ngrdcol, gr%nzt)``). + named_in_bounds = " ".join( + " ".join(type_records[dummies[other]][1].get(member) or []) + for other, members in touched.items() + if other in dummies and dummies[other] in type_records + for member in members + ) + named_in_bounds += " " + " ".join( + str(d.get(k) or "") + for a in sub["args"] + for d in (a.get("dims") or ()) + for k in ("lb", "ub") + ) + for ref in re.findall(r"(\w+)\s*%\s*(\w+)", named_in_bounds): + if ref[0].lower() == obj and ref[1].lower() in comps: + touched[obj].setdefault(ref[1].lower(), False) for member, written in sorted(touched[obj].items()): spec = comps.get(member) if spec is None: plan.unsupported.append(f"{obj}%{member}: no such component") continue + if str(spec.get("dtype")) not in FORTRAN_TYPES: + # A character component (CLUBB's err_info%err_header): the + # adapter has no flat spelling for it and the physics reads + # it only to print. Left at the object's default, and said. + plan.left_to_module.append(f"{obj}%{member}: {spec.get('dtype')} not carried") + continue found_axes = bounds.get(member) if found_axes is None and spec.get("dims"): plan.unsupported.append(f"{obj}%{member}: no allocate statement found") @@ -915,6 +1072,10 @@ def _state_vars( if not record: continue module = str(record.get("module", "")).lower() + public = {str(n).lower() for n in record.get("public") or ()} + default_private = bool( + re.search(r"^\s*private\s*(?:!.*)?$", path.read_text(errors="replace"), re.I | re.M) + ) for entry in record.get("module_state", ()): name = str(entry["name"]).lower() # Every module *variable* the run may have set -- initialized in @@ -925,6 +1086,16 @@ def _state_vars( DERIVED.match(str(entry.get("dtype"))) or str(entry.get("dtype")) not in FORTRAN_TYPES ): + # A derived-type state variable is an *object* of the plan + # (carried by component, through the objects path); naming + # it here as left to the module misdescribed the sponge + # profiles CLUBB's solvers read through a companion. + continue + if default_private and name not in public: + # No ``use`` reaches it, so no adapter sets it: both sides + # run with the module's own default, and the plan says so. + seen.add(name) + plan.left_to_module.append(f"{module}%{name}") continue dims = entry.get("dims") or [] names = [ @@ -971,10 +1142,31 @@ def _bind_symbolic_extents(plan: FlatPlan) -> None: state name, which is a scalar argument of the adapter; a name no state answers for makes the component unsupported.""" by_name = {state.name: state.flat for state in plan.states if not state.extents} + # A dummy of the subprogram is an argument of the adapter, so an + # allocation sized by the allocating routine's dummy of the same name + # (``coef_wp4_implicit(1:ngrdcol,1:nz)``, CLUBB) is spelled as written. + dummies = { + str(a["name"]).lower() + for a in plan.subprogram["args"] + if not a.get("dims") and not DERIVED.match(str(a["dtype"])) + } + carried = {(obj.name, comp.name): comp.flat for obj in plan.objects for comp in obj.components} + flat_names = set(carried.values()) - def bind(text: str) -> str | None: + def bind(text: str, synthetic: str | None = None) -> str | None: + """``synthetic`` names the extent argument to stand in when the + text is one identifier nothing else answers for.""" missing: list[str] = [] + def component(match: re.Match[str]) -> str: + key = (match.group(1).lower(), match.group(2).lower()) + if key in carried: + return carried[key] + missing.append(match.group(0)) + return match.group(0) + + text = re.sub(r"(\w+)\s*%\s*(\w+)", component, text) + def swap(match: re.Match[str]) -> str: token = match.group(0) lowered = token.lower() @@ -982,17 +1174,30 @@ def swap(match: re.Match[str]) -> str: return token if lowered in by_name: return by_name[lowered] + if token in flat_names or lowered in dummies: + return token # a component bound just above, or a dummy missing.append(token) return token out = re.sub(r"[A-Za-z_]\w*", swap, text) + if missing and synthetic is not None: + bare = text.strip().strip("()").strip() + if re.fullmatch(r"[A-Za-z_]\w*", bare) and bare.lower() == missing[0].lower(): + return synthetic return None if missing else out for obj in plan.objects: kept = [] for comp in obj.components: - extents = [bind(e) for e in comp.extents] - bounds = [(bind(lo), bind(hi)) for lo, hi in comp.bounds] + names = [f"{comp.flat}_n{axis + 1}" for axis in range(len(comp.extents))] + extents = [bind(e, names[axis]) for axis, e in enumerate(comp.extents)] + bounds = [ + (bind(lo), bind(hi, names[axis]) if lo.strip() == "1" else bind(hi)) + for axis, (lo, hi) in enumerate(comp.bounds) + ] + for axis, extent in enumerate(extents): + if extent == names[axis]: + plan.extent_args[names[axis]] = [obj.name, comp.name, axis + 1] if any(e is None for e in extents) or any( lo is None or hi is None for lo, hi in bounds ): diff --git a/src/recast/fortran/frontend.py b/src/recast/fortran/frontend.py index c42652d..957ab41 100644 --- a/src/recast/fortran/frontend.py +++ b/src/recast/fortran/frontend.py @@ -394,7 +394,11 @@ def analyze(self, unit: Unit, root: Path) -> Facts: from recast.fortran._parse import STD, digest, f03 from recast.fortran._parse import parse as parse_file from recast.fortran.effects import side_channels - from recast.fortran.interface import _scope_of, companion_externals, subprogram_key + from recast.fortran.interface import ( + _scope_of, + companion_externals, + subprogram_key, + ) from recast.fortran.rwset import block_rwsets, scope_for path = self._source_of(unit, Path(root)) @@ -475,6 +479,18 @@ def analyze(self, unit: Unit, root: Path) -> Facts: table[local] = table[remote] for name, entry in table.items(): externals.setdefault(name, entry) + # A procedure use-imported from a stubbed module is a call the + # translation answers with a stub -- ``pass`` for CLUBB's stats_update + # -- so on this side too its actuals are neither read nor written. + # Procedures only: a constant the same module exports (constants_clubb + # is a stub and a table of constants) is a read on both sides. The + # stub module's own record says which names are procedures; a stub + # the tree does not carry is taken at its import list. + stub_procedures = self._stub_procedures(record, Path(root)) + for name in stub_procedures: + stubbed = {"kind": "subroutine", "out_positions": [], "buffer_positions": []} + externals.setdefault(name, {**stubbed, "stub": True}) + record = {**record, "stub_procedures": sorted(stub_procedures)} callgraph: dict[str, list[str]] = {} effects: dict[str, Any] = {} @@ -630,6 +646,36 @@ def _tree_kinds( pending.extend(record_of.get("use_statements", ())) return found + def _stub_procedures(self, record: dict[str, Any], root: Path) -> set[str]: + """The local names this unit imports from stubbed modules that are + procedures of theirs -- calls the translation stubs. A stub module the + tree does not carry contributes every name it is imported for.""" + from recast.fortran import interface as interface_mod + + index = self._module_index(root.resolve()) + names: set[str] = set() + for statement in record.get("use_statements", ()): + match = USE_STATEMENT.match(statement.strip()) + if not match or match.group("module").lower() not in self.stub_modules: + continue + module = match.group("module").lower() + imported = { + item.split("=>", 1)[0].strip().lower(): item.split("=>", 1)[-1].strip().lower() + for item in (match.group("only") or "").split(",") + if item.strip() + } + source = index.get(module) + procedures = None + if source is not None: + record_of = self._readable(source, interface_mod.extract, module) + if record_of is not None: + procedures = {str(sub["name"]).lower() for sub in record_of["subprograms"]} + procedures |= {g.lower() for g in record_of.get("generics") or {}} + for local, remote in imported.items(): + if procedures is None or remote in procedures: + names.add(local) + return names + def _companions( self, record: dict[str, Any], path: Path, root: Path ) -> tuple[list[dict[str, Any]], list[dict[str, str]]]: diff --git a/src/recast/fortran/interface.py b/src/recast/fortran/interface.py index 00c4c8c..ebdb9ae 100644 --- a/src/recast/fortran/interface.py +++ b/src/recast/fortran/interface.py @@ -1361,6 +1361,73 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]: for at, argument in enumerate(sub["args"]) if argument["intent"] in ("OUT", "INOUT") ], + # A buffer OUT is passed in and returned (the caller's storage), + # so the caller reads the actual as well as writing it -- the + # same rule the scope applies to its own subprograms (#38). + "buffer_positions": [ + at for at, argument in enumerate(sub["args"]) if argument.get("buffer") + ], + # What the caller reads: IN and INOUT actuals, and a buffer OUT. + # An INOUT actual is written *and* read; out_positions alone said + # only the first. + "read_positions": [ + at + for at, argument in enumerate(sub["args"]) + if argument["intent"] in ("IN", "INOUT", "UNKNOWN") or argument.get("buffer") + ], + # So a keyword actual (``rcond = rcond``, CLUBB's band_solve) lands + # on its own position and not on whichever comes next. + "arg_names": [str(argument["name"]).lower() for argument in sub["args"]], + # Handing the caller's own optional to one of these queries its + # presence -- a read of it, as ``present()`` is. + "optional_out_positions": [ + at + for at, argument in enumerate(sub["args"]) + if argument.get("optional") and argument["intent"] == "OUT" + ], + } + # The sibling's generics too: a call spells the generic (CLUBB's + # ``zt2zm_api`` over grid_class's specifics), and a name the scope does + # not know as a procedure it counts as a read of data. The entry's + # writes are the union over the specifics -- which agree, in every + # generic seen so far, on being functions with no OUT argument. + for generic, specifics in (record.get("generics") or {}).items(): + known = [(s, table[s]) for s in specifics if s in table] + if generic in table or not known: + continue + # Which specific a call reaches depends on its arity (and ranks the + # scope does not resolve): the entry carries every specific with its + # argument count, and the scope picks by the actuals it sees. A + # union over specifics of different arity marked the wrong + # positions -- CLUBB's tridiag_solve, zm2zt_api. + signatures = {s: next(x for x in record["subprograms"] if x["name"] == s) for s, _ in known} + arity = {s: len(sig["args"]) for s, sig in signatures.items()} + required = { + s: sum(1 for a in sig["args"] if not a.get("optional")) for s, sig in signatures.items() + } + table[generic] = { + "kind": known[0][1]["kind"], + "out_positions": sorted({at for _, entry in known for at in entry["out_positions"]}), + "buffer_positions": sorted( + {at for _, entry in known for at in entry.get("buffer_positions", [])} + ), + "specifics": [ + { + "name": s, + "args": arity[s], + "required": required[s], + "out_positions": entry["out_positions"], + "buffer_positions": entry.get("buffer_positions", []), + "read_positions": entry.get("read_positions", []), + "arg_names": entry.get("arg_names", []), + "optional_out_positions": entry.get("optional_out_positions", []), + # Two specifics of one arity (CLUBB's tridiag_solve + # with and without its optional rcond) are told apart + # by the ranks of their dummies. + "ranks": [len(a.get("dims") or []) for a in signatures[s]["args"]], + } + for s, entry in known + ], } return table diff --git a/src/recast/fortran/rwset.py b/src/recast/fortran/rwset.py index fa1f276..508715d 100644 --- a/src/recast/fortran/rwset.py +++ b/src/recast/fortran/rwset.py @@ -359,7 +359,15 @@ def _write_actual(actual: Any) -> None: """ if isinstance(actual, f03.Name): writes.add(str(actual).lower()) - elif isinstance(actual, (f03.Part_Ref, f03.Data_Ref)): + elif isinstance(actual, f03.Data_Ref): + # ``pdf_params%chi_1`` as an OUT actual (CLUBB's pdf_closure): + # the object is written; a component's *subscripts* are read, + # the component's name is not a variable of this scope. + writes.add(str(actual.children[0]).lower()) + for comp in actual.children[1:]: + if isinstance(comp, f03.Part_Ref) and comp.children[1] is not None: + reads.update(expr_reads(comp.children[1], scope)) + elif isinstance(actual, f03.Part_Ref): writes.add(str(actual.children[0]).lower()) for child in actual.children[1:]: reads.update(expr_reads(child, scope)) @@ -404,14 +412,76 @@ def call(stmt: Any) -> None: ) if name in dummies: reads.add(name) + optional_dummies = ( + {a["name"].lower() for a in scope.semantics.subprogram["args"] if a.get("optional")} + if (scope.semantics is not None) + else set() + ) + + def hands_on_presence(actual: Any) -> None: + """The caller's own optional passed to an optional OUT (CLUBB's + xm_wpxp_solve hands ``rcond = rcond`` to band_solve): the callee + asks whether it is present, which is a read of it on both sides + -- ``present(x)`` here, the ``want_x`` sentinel there.""" + if isinstance(actual, f03.Name) and str(actual).lower() in optional_dummies: + reads.add(str(actual).lower()) if callee is None: external = scope.externals.get(name) + if external and external.get("stub"): + return # a stubbed call: the translation drops it, actuals and all + if external and external.get("specifics"): + # A sibling's generic: the specific this arity reaches, or the + # union when none matches exactly. + # An optional dummy may be left off: a specific fits when the + # actuals fall between its required and its total count. + fitting = [ + x + for x in external["specifics"] + if x.get("required", x["args"]) <= len(actuals) <= x["args"] + ] + if len(fitting) > 1: + # Several fit the count (one's optional tail is the + # other's required one): the ranks of the bare-name + # actuals decide, where a declared rank is known. + def _agrees(x: dict[str, Any]) -> bool: + return all( + not isinstance(a, f03.Name) + or scope.ranks.get(str(a).lower()) is None + or scope.ranks[str(a).lower()] == r + for a, r in zip(actuals, x.get("ranks", []), strict=False) + ) + + fitting = [x for x in fitting if _agrees(x)] or fitting + if fitting: + external = fitting[0] out_positions = set(external.get("out_positions", [])) if external else set() + buffers = set(external.get("buffer_positions", [])) if external else set() + read_positions = ( + set(external["read_positions"]) + if external and "read_positions" in external + else None + ) + if external and external.get("arg_names"): + # Keyword actuals by name, positional ones in order, as for a + # callee of this module. + formals = [{"name": n} for n in external["arg_names"]] + actuals = _bind_actuals({"args": formals}, items) + optional_out = set(external.get("optional_out_positions", [])) if external else set() for j, actual in enumerate(actuals): + if actual is None: + continue if j in out_positions: write_target(actual) + if j in optional_out: + hands_on_presence(actual) + if read_positions is not None: + is_read = j in read_positions else: + # A buffer OUT of a sibling is read as well as written: + # the emitter passes the caller's storage in (#38). + is_read = j not in out_positions or j in buffers + if is_read: reads.update(expr_reads(actual, scope)) return @@ -427,6 +497,8 @@ def call(stmt: Any) -> None: reads.update(expr_reads(actual, scope)) if formal["intent"] in ("OUT", "INOUT"): _write_actual(actual) + if formal.get("optional") and formal["intent"] == "OUT": + hands_on_presence(actual) def visit(stmt: Any) -> None: if isinstance(stmt, f08.Block_Construct): diff --git a/src/recast/fortran/semantics.py b/src/recast/fortran/semantics.py index 4bdae54..f434f2c 100644 --- a/src/recast/fortran/semantics.py +++ b/src/recast/fortran/semantics.py @@ -442,8 +442,17 @@ def _reference_rank(self, node: Any) -> int: return 0 if name in self.procedures: return self._call_rank(self.procedures[name], items) - if name in self.companion_generics: - return 0 # the overload decides, and dispatch is a separate question + if name in self.companion_generics or name in self.generics: + # The overload decides: CLUBB's ``sat_mixrat_liq_api( ..., + # thlm2T_in_K_api( nzt, ngrdcol, thlm, exner, rcm ), ... )`` is + # an actual whose rank is the inner generic's 2-D specific's + # result rank, and the outer generic's dispatch needs it. + try: + specific = self.dispatch(name, items) + except AmbiguousDispatch: + return 0 + record = self.procedures.get(specific) + return self._call_rank(record, items) if record is not None else 0 if name in TRANSFORMATIONAL or name in STATE_QUERY - {"merge"}: return 0 if name in ELEMENTAL or name == "merge": @@ -455,10 +464,11 @@ def _reference_rank(self, node: Any) -> int: return sum(1 for s in items if isinstance(s, f03.Subscript_Triplet)) def _call_rank(self, record: dict[str, Any], items: list[Any]) -> int: - """An ELEMENTAL function broadcasts; anything else returns its result.""" + """An ELEMENTAL function broadcasts; anything else returns its result + -- an array where the result is declared with dimensions.""" if any("ELEMENTAL" in str(p).upper() for p in (record.get("prefixes") or [])): return self._broadcast_rank(items) - return 0 + return len(record.get("result_dims") or []) def _broadcast_rank(self, items: list[Any]) -> int: return max( diff --git a/src/recast/fortran/tree.py b/src/recast/fortran/tree.py index 1068cc9..c104922 100644 --- a/src/recast/fortran/tree.py +++ b/src/recast/fortran/tree.py @@ -159,18 +159,19 @@ def _evaluate( initialize it with something a parameter can be folded from.""" # Lazy, like ``render`` above: ``expr`` parses, and this module is imported # by paths that must stay importable without the ``fortran`` extra. - from recast.fortran.expr import python_call, with_integer_division + from recast.fortran.expr import python_call, typed, with_integer_division try: records = resolve([name], files) except unresolved: return None env: dict[str, Any] = {} + kinds: dict[str, str | None] = {} try: for entry in records: # Integer arithmetic where Fortran's ``/`` truncates. text = render( - with_integer_division(entry["expr"]), + with_integer_division(entry["expr"], env=kinds), real=lambda t: f"float('{t}')", integer=lambda t: t, name=lambda t: t.upper(), @@ -179,6 +180,7 @@ def _evaluate( scope = {"__builtins__": {}, "max": max, "min": min, "abs": abs, "int": int} scope.update({"float": float, "math": math, "sys": sys}) env[entry["name"].upper()] = eval(text, scope, dict(env)) # noqa: S307 + kinds[entry["name"]] = entry.get("dtype") or typed(entry["expr"], kinds) except Exception: # an initializer shape the renderer has no rule for return None return env.get(name.upper()) diff --git a/src/recast/fortran/use.py b/src/recast/fortran/use.py index 42f64c5..947a9d8 100644 --- a/src/recast/fortran/use.py +++ b/src/recast/fortran/use.py @@ -25,8 +25,9 @@ class UnresolvedConstant(RecastError): """A use-imported name whose initializer is in none of the given sources.""" -def harvest(path: Path) -> dict[str, tuple[Any, int | None]]: - """``name -> (initializer node, line)`` for module-level initialized entities. +def harvest(path: Path) -> dict[str, tuple[Any, int | None, str | None]]: + """``name -> (initializer node, line, declared base type)`` for module-level + initialized entities; the base type is ``real``, ``int`` or ``None``. Covers parameters and initialized ``save``/``protected`` variables alike: a constant that a physics module reads is a constant whether or not the @@ -34,7 +35,7 @@ def harvest(path: Path) -> dict[str, tuple[Any, int | None]]: of them are harvested. """ ast = parse(path) - out: dict[str, tuple[Any, int | None]] = {} + out: dict[str, tuple[Any, int | None, str | None]] = {} for mod in walk(ast, f03.Module): spec = next((c for c in mod.children if isinstance(c, f03.Specification_Part)), None) if spec is None: @@ -46,9 +47,14 @@ def harvest(path: Path) -> dict[str, tuple[Any, int | None]]: if item is not None and getattr(item, "span", None): line = item.span[0] break + # The declared base type, which is what says whether ``rd / rv`` + # is a real quotient: the fold cannot tell from two names. + base = str(decl.children[0]).split("(")[0].strip().upper() + declared = {"REAL": "real", "DOUBLE PRECISION": "real", "INTEGER": "int"}.get(base) for ent in walk(decl, f03.Entity_Decl): if ent.children[3] is not None: - out[str(ent.children[0]).lower()] = (ent.children[3].children[1], line) + initializer = ent.children[3].children[1] + out[str(ent.children[0]).lower()] = (initializer, line, declared) return out @@ -63,7 +69,7 @@ def resolve(symbols: list[str], sources: list[Path]) -> list[dict[str, Any]]: constant that silently becomes undefined downstream is far more expensive to diagnose than a failure here that names it. """ - table: dict[str, tuple[Any, int | None]] = {} + table: dict[str, tuple[Any, int | None, str | None]] = {} origin: dict[str, Path] = {} for path in sources: for name, rec in harvest(path).items(): @@ -80,7 +86,7 @@ def need(name: str) -> None: if name not in table: raise UnresolvedConstant(f"no initializer for {name!r} in {[str(s) for s in sources]}") seen.add(name) - node, line = table[name] + node, line, declared = table[name] # The one legal self-reference, a kind inquiry on the constant being # declared, stands for its kind alone; see ``expr.substitute``. expr: Expr = substitute(build(node), name, Expr("real", "1.0")) @@ -90,6 +96,7 @@ def need(name: str) -> None: { "name": name, "expr": expr, + "dtype": declared, "source": str(origin[name]), "line": line, "requested": name in requested, diff --git a/src/recast/oracle/dump_replay.py b/src/recast/oracle/dump_replay.py index 833c4c8..8ad33da 100644 --- a/src/recast/oracle/dump_replay.py +++ b/src/recast/oracle/dump_replay.py @@ -122,7 +122,18 @@ def parse_dump(text: str) -> tuple[dict[str, Any], dict[str, Any]]: all_integers = True # every value of the section so far is an integer literal def flush() -> None: - if not name or not values: + if not name: + return + if not values: + # A zero-extent record: a component the run never allocated + # (CLUBB's scalar-tracer arrays under sclr_dim = 0), written as + # ``name(1,88,0)``. An empty array of that shape is a value. + if dims_text and any(t.strip() == "0" for t in dims_text.split(",")): + empty: list[int] = [] + for t in dims_text.split(","): + sized = str(metadata.get(t.strip().lower(), t.strip())) + empty.append(int(sized) if sized.lstrip("-").isdigit() else 0) + (inputs if target == "INPUT" else outputs)[name] = np.zeros(tuple(empty), order="F") return # The recorder writes integer arrays with ``i0`` and reals with an # exponent, so a section whose every value is an integer literal is an @@ -179,6 +190,12 @@ def flush() -> None: metadata[key] = whole inputs[key] = np.int32(whole) except ValueError: + # A logical, written ``T`` / ``F`` by the recorder's l1 + # format: a value, not a diagnostic. + if text_value.upper() in ("T", "F", ".TRUE.", ".FALSE."): + metadata[key] = text_value.upper().startswith(("T", ".T")) + inputs[key] = np.bool_(metadata[key]) + continue # Not a number. Upstream swallows this with a bare # ``except``; narrowed to what can actually be raised # here, which changes no outcome and stops the clause diff --git a/src/recast/oracle/record.py b/src/recast/oracle/record.py index 808022e..dacad47 100644 --- a/src/recast/oracle/record.py +++ b/src/recast/oracle/record.py @@ -96,6 +96,25 @@ def _record_call( ) +def _record_component(lines: list[str], unit: str, tag: str, obj: Any, comp: Any) -> None: + """One component's record, guarded when it may not be there: a component + the model allocates only on some configurations (CLUBB's scalar-tracer + coefficients under ``sclr_dim = 0``) is written as a zero-extent record, + and ``reshape`` is never asked for storage that does not exist.""" + target = f"{obj.name}%{comp.name}" + rank = len(comp.extents) + if not comp.bounds: + _record_call(lines, unit, tag, comp.flat, target, rank, comp.dtype) + return + test = "associated" if comp.pointer else "allocated" + lines.append(f" if ({test}({target})) then") + _record_call(lines, unit, tag, comp.flat, target, rank, comp.dtype) + lines.append(" else") + zeros = "(" + ",".join("0" for _ in range(rank)) + ")" if rank else "" + lines.append(f" write ({unit}, '(a)') '# {tag}: {comp.flat}{zeros}'") + lines.append(" end if") + + def recorder_module( module: str, plans: list[FlatPlan], @@ -211,7 +230,12 @@ def recorder_module( spelled = FORTRAN_TYPES[str(a["dtype"])] dims = "(" + ",".join(":" for _ in a["dims"]) + ")" if a.get("dims") else "" lines.append(f" {spelled}, intent(in) :: {a['name']}{dims}") - lines.append(f" integer :: {patch}") + passes_patch = any(a["name"].lower() == patch.lower() for a in originals) + if not passes_patch: + # CLUBB passes ngrdcol as a dummy; then it is already declared and + # already the run's value, and a local of the same name would be + # a second declaration. + lines.append(f" integer :: {patch}") lines.append(" character(len=128) :: dims") # The patch count from the first component allocated over it. first = next( @@ -223,7 +247,9 @@ def recorder_module( ), None, ) - if first is None: + if passes_patch: + pass + elif first is None: lines.append(f" {patch} = 1") else: lines.append(f" {patch} = size({first[0].name}%{first[1].name}, 1)") @@ -244,6 +270,17 @@ def recorder_module( fmt = {"int32": "i0", "bool": "l1"}.get(str(a["dtype"]), "es25.17e3") lines.append(f" write ({u}, '(a,{fmt})') '# {a['name']} = ', {a['name']}") lines.append(f" write ({u}, '(a,i0)') '# {patch} = ', {patch}") + for extent, (owner, member, axis) in plan.extent_args.items(): + # An extent the plan could not spell: the run's own, from size(), + # zero when the component was never allocated. + component = next( + c for o in plan.objects if o.name == owner for c in o.components if c.name == member + ) + test = "associated" if component.pointer else "allocated" + lines.append( + f" write ({u}, '(a,i0)') '# {extent} = ', " + f"merge(size({owner}%{member}, {axis}), 0, {test}({owner}%{member}))" + ) for a in originals: if a.get("dims") and not DERIVED.match(str(a["dtype"])): @@ -252,15 +289,7 @@ def recorder_module( ) for obj in plan.objects: for comp in obj.components: - _record_call( - lines, - u, - "INPUT", - comp.flat, - f"{obj.name}%{comp.name}", - len(comp.extents), - comp.dtype, - ) + _record_component(lines, u, "INPUT", obj, comp) for state in plan.states: _record_call( lines, @@ -273,23 +302,18 @@ def recorder_module( _dims_text(state.extents), ) lines += [" else", f" if (n_{sname} > max_calls) return"] + # The plain OUT and INOUT dummies are outputs too (CLUBB's advance_* + # return wp2, wp3, ... beside what they write into the objects). + for a in originals: + if DERIVED.match(str(a["dtype"])) or a["intent"] not in ("OUT", "INOUT"): + continue + _record_call( + lines, u, "OUTPUT", a["name"], a["name"], len(a.get("dims") or ()), str(a["dtype"]) + ) for obj in plan.objects: for comp in obj.components: if comp.written: - _record_call( - lines, - u, - "OUTPUT", - comp.flat, - f"{obj.name}%{comp.name}", - len(comp.extents), - comp.dtype, - ) - for a in originals: - if a.get("dims") and not DERIVED.match(str(a["dtype"])) and a["intent"] != "IN": - _record_call( - lines, u, "OUTPUT", a["name"], a["name"], len(a["dims"]), str(a["dtype"]) - ) + _record_component(lines, u, "OUTPUT", obj, comp) for state in plan.states: if state.written: _record_call( @@ -329,6 +353,15 @@ def _continues(line: str) -> bool: return _strip_comment(line).rstrip().endswith("&") +def _last_code(span: list[str]) -> str: + """The last line of ``span`` that carries code: blank and comment-only + lines do not end a continued statement.""" + for line in reversed(span): + if _strip_comment(line).strip(): + return line + return span[-1] + + def _split_actuals(text: str) -> list[str]: """Split an actual-argument list on top-level commas only.""" parts: list[str] = [] @@ -374,10 +407,17 @@ def probe_tree( i = 0 while i < len(lines): # A call may continue over several lines: gather the statement. + # A blank or comment-only line between continuations (what cpp + # leaves of an ``#ifdef`` inside CLUBB's advance_clubb_core call) + # is part of the statement, not its end. span = [lines[i]] - while _continues(span[-1]) and i + len(span) < len(lines): + while _continues(_last_code(span)) and i + len(span) < len(lines): span.append(lines[i + len(span)]) - logical = " ".join(_strip_comment(ln).rstrip("&").strip().lstrip("&") for ln in span) + # ``a, & ! In`` leaves a space between the ampersand and the + # comment it trailed: strip blanks before the ampersands. + logical = " ".join( + _strip_comment(ln).strip().rstrip("&").strip().lstrip("&").strip() for ln in span + ) match = CALL.match(" " * (len(span[0]) - len(span[0].lstrip())) + logical.strip()) called = match.group("name").lower() if match else None if match and called and called in targets: diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py index 3232802..f4e77df 100644 --- a/src/recast/transform/jax/backend.py +++ b/src/recast/transform/jax/backend.py @@ -267,10 +267,27 @@ class CallRewrite(ast.NodeTransformer): def __init__(self, call_map, known_subs): self.map = call_map self.subs = known_subs + self.host_calls: list[str] = [] def visit_Call(self, node): self.generic_visit(node) f = node.func + if isinstance(f, ast.Name) and f.id == "_f_ecall" and node.args: + # ``_f_ecall(split, x[i - 1, :])``: the elemental broadcast of an + # emitted kernel is the runtime's jnp.vectorize over the kernel's + # implementation, its state closure appended; of a subprogram + # not being emitted, a host function no tracer can run. + callee = node.args[0] + if isinstance(callee, ast.Name) and callee.id in self.map: + info = self.map[callee.id] + node.args = [ + ast.Name(id=f"_{callee.id}_k_impl", ctx=ast.Load()), + *node.args[1:], + *[ast.Name(id=c, ctx=ast.Load()) for c in info["closure"]], + ] + elif isinstance(callee, ast.Name) and callee.id in self.subs: + raise JaxQueue(f"elemental call of non-emitted subprogram {callee.id}") + return node if isinstance(f, ast.Name): if f.id in self.map: info = self.map[f.id] @@ -284,13 +301,165 @@ def visit_Call(self, node): node.func = ast.Name(id=f"_{f.id}_k_impl", ctx=ast.Load()) node.args = pos + [ast.Name(id=c, ctx=ast.Load()) for c in info["closure"]] elif f.id in self.subs: - raise JaxQueue(f"calls non-emitted subprogram {f.id}") + # A subprogram of this module the port could not emit, + # called from one it can: the call stays the host's + # (``_host.``, the NumPy anchor's), under whatever + # guard the anchor put it -- a configuration switch, static + # at trace time, is the usual one -- and is never traced + # while the guard holds. If the guard ever lets it + # through, the trace fails on it by name, which the gate + # reports. Delegating the caller, and its callers up to + # the driver, for a path the run never takes was the + # alternative. The note names it. + self.host_calls.append(f.id) + node.func = ast.Attribute( + value=ast.Name(id="_host", ctx=ast.Load()), attr=f.id, ctx=ast.Load() + ) return node # ------------------------------------------------------- statement lowering +def _has_cycle(stmts) -> bool: + """Whether a ``continue`` of *this* loop sits in ``stmts`` -- at the top + level or under an ``if``; one inside a nested ``for`` is that loop's.""" + for s in stmts: + if isinstance(s, ast.Continue): + return True + if isinstance(s, ast.If) and (_has_cycle(s.body) or _has_cycle(s.orelse)): + return True + return False + + +def _cycle_to_else(stmts, rest): + """``stmts`` followed by ``rest``, with every ``continue`` folded away. + + Fortran's ``if ( ... ) then ... cycle end if`` followed by the rest of + the loop body (CLUBB's interpolators) is ``if c: A else: R`` -- the + continuation of the loop body moves into the branches that do not + cycle. Done on the Python AST before the fori_loop lowering, which has + no place for a ``continue``: a ``lax.cond`` branch is a function. + """ + out: list[ast.stmt] = [] + for at, s in enumerate(stmts): + if isinstance(s, ast.Continue): + return out # what follows is never reached on this path + if isinstance(s, ast.If) and (_has_cycle(s.body) or _has_cycle(s.orelse)): + tail = [*stmts[at + 1 :], *rest] + folded = ast.If( + test=s.test, + body=_cycle_to_else(s.body, copy.deepcopy(tail)) or [ast.Pass()], + orelse=_cycle_to_else(s.orelse, copy.deepcopy(tail)), + ) + return [*out, ast.copy_location(folded, s)] + out.append(s) + return [*out, *rest] + + +def _not_flag(flag: str) -> ast.expr: + """``jnp.logical_not(flag)``: a traced test. Python's ``not`` on a + tracer is the boolean conversion JAX refuses, and the expression + mapping that would spell it ran before the loop was lowered.""" + return ast.Call( + func=ast.Attribute( + value=ast.Name(id="jnp", ctx=ast.Load()), attr="logical_not", ctx=ast.Load() + ), + args=[ast.Name(id=flag, ctx=ast.Load())], + keywords=[], + ) + + +def _jnp_bool(value: bool) -> ast.expr: + """``jnp.bool_(value)``: a traced store, not a literal the backend + would treat as a trace-time constant and fail to carry.""" + return ast.Call( + func=ast.Attribute(value=ast.Name(id="jnp", ctx=ast.Load()), attr="bool_", ctx=ast.Load()), + args=[ast.Constant(value=value)], + keywords=[], + ) + + +def _breaks_at_level(stmts) -> bool: + """A ``break`` that belongs to this loop: not one inside a nested loop.""" + for st in stmts: + if isinstance(st, ast.Break): + return True + if isinstance(st, ast.For | ast.While): + continue + for field in ("body", "orelse"): + inner = getattr(st, field, None) + if isinstance(inner, list) and _breaks_at_level(inner): + return True + return False + + +def _guard_after_breaks(stmts, flag, kept, index): + """``break`` -> ``flag = True; kept = index``; what follows a statement + that may have broken, in the same block, runs under ``if not flag`` -- + recursively through branches, never into a nested loop (its breaks are + its own).""" + out: list[ast.stmt] = [] + for at, st in enumerate(stmts): + had_break = _breaks_at_level([st]) + if isinstance(st, ast.Break): + out.append( + ast.copy_location( + ast.Assign(targets=[ast.Name(id=flag, ctx=ast.Store())], value=_jnp_bool(True)), + st, + ) + ) + rewritten: ast.stmt = ast.copy_location( + ast.Assign( + targets=[ast.Name(id=kept, ctx=ast.Store())], + value=ast.Call( + func=ast.Attribute( + value=ast.Call( + func=ast.Attribute( + value=ast.Name(id="jnp", ctx=ast.Load()), + attr="asarray", + ctx=ast.Load(), + ), + args=[ast.Name(id=index, ctx=ast.Load())], + keywords=[], + ), + attr="astype", + ctx=ast.Load(), + ), + args=[ + ast.Attribute( + value=ast.Name(id=kept, ctx=ast.Load()), + attr="dtype", + ctx=ast.Load(), + ) + ], + keywords=[], + ), + ), + st, + ) + else: + rewritten = st + if not isinstance(st, ast.For | ast.While): + for field in ("body", "orelse"): + inner = getattr(st, field, None) + if isinstance(inner, list) and inner: + setattr(st, field, _guard_after_breaks(inner, flag, kept, index)) + out.append(rewritten) + if had_break: + rest = _guard_after_breaks(stmts[at + 1 :], flag, kept, index) + if rest: + out.append( + ast.If( + test=_not_flag(flag), + body=rest, + orelse=[], + ) + ) + return out + return out + + def _assigned_names(stmts): """Names stored by Assign statements, first-assignment order (nested fori_loop/cond results arrive as Tuple targets; static @@ -309,11 +478,16 @@ def add(n): # body before their use and never read after it; carried, they # would have to be initialized at the enclosing level, where # nothing assigns them. - if isinstance(t, ast.Name) and not re.fullmatch(r"_(?:hi_|cnt_|t)\d+", t.id): + # ``_out`` is the anchor's call-result tuple, assigned and + # unpacked on consecutive lines of one block: the block's own, + # never a carry (its shape changes from call to call). + if isinstance(t, ast.Name) and not re.fullmatch( + r"_(?:hi_|cnt_|lo_|st_|t)\d+|_out", t.id + ): add(t.id) elif isinstance(t, ast.Tuple): for e in t.elts: - if isinstance(e, ast.Name) and not re.fullmatch(r"_t\d+", e.id): + if isinstance(e, ast.Name) and not re.fullmatch(r"_t\d+|_out", e.id): add(e.id) elif isinstance(s, ast.If): for n in _assigned_names(s.body) + _assigned_names(s.orelse): @@ -321,10 +495,64 @@ def add(n): return out -def _static_test(test): +def _jnp_logic(node) -> str | None: + """``jnp.logical_not/and/or(...)``: which, or None.""" + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "jnp" + and node.func.attr in ("logical_not", "logical_and", "logical_or") + and not node.keywords + ): + return node.func.attr + return None + + +def _python_logic(node): + """The expression mapping's ``jnp.logical_*`` back to Python's ``not``, + ``and``, ``or`` -- for a test a Python ``if`` evaluates at trace time, + where a jnp op on a constant would be staged into a tracer.""" + which = _jnp_logic(node) + if which == "logical_not" and len(node.args) == 1: + return ast.UnaryOp(op=ast.Not(), operand=_python_logic(node.args[0])) + if which in ("logical_and", "logical_or") and len(node.args) == 2: + op = ast.And() if which == "logical_and" else ast.Or() + return ast.BoolOp(op=op, values=[_python_logic(a) for a in node.args]) + return node + + +def _logic_leaves(node) -> list[ast.expr]: + """The operands under the logical operators: what has to be concrete + for the Python form to be evaluable.""" + which = _jnp_logic(node) + if which is not None: + return [leaf for a in node.args for leaf in _logic_leaves(a)] + return [node] + + +def _module_constant(node) -> bool: + """``_mod.IIPDF_ADG1``: a use-associated module constant through its + module alias -- a Python value at trace time, like the bare upper-case + spelling of one the translation resolved.""" + return ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id.startswith("_") + and node.attr.isupper() + and len(node.attr) > 1 + ) + + +def _static_test(test, statics=frozenset()): """True for branch conditions decidable at trace time per the translate.py grammar: `x is [not] None` (Fortran PRESENT) and bare - `want_*` sentinels (optional-output flags, static under jit).""" + `want_*` sentinels (optional-output flags, static under jit), and a + comparison over module constants and the kernel's static scalar + arguments (``statics``: Python ints at trace time under jit). CLUBB's + ``if ( sclr_dim > 0 )`` guards stores into arrays the run never + allocated; lowered to lax.cond both arms are traced, and the store + into a (0, 0) array is an IndexError at trace time.""" if ( isinstance(test, ast.Compare) and len(test.ops) == 1 @@ -338,11 +566,24 @@ def constant(node): if isinstance(node, ast.Constant): return isinstance(node.value, (int, float)) and not isinstance(node.value, bool) if isinstance(node, ast.Name): - return node.id.isupper() and len(node.id) > 1 + return node.id in statics or (node.id.isupper() and len(node.id) > 1) + if _module_constant(node): + return True if isinstance(node, ast.BinOp): return constant(node.left) and constant(node.right) if isinstance(node, ast.UnaryOp): return constant(node.operand) + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and all(constant(a) for a in node.args) + and not node.keywords + ): + # ``_error_code.clubb_at_least_debug_level_api(2)``: a stand-in's + # function of constants, a Python value at trace time. (A kernel + # of constants is a concrete array; ``_f_concrete`` tells.) + return True return False # ``RUNGE_KUTTA_TYPE == I_10``: a comparison over module constants @@ -351,8 +592,14 @@ def constant(node): if isinstance(test, ast.Compare) and len(test.ops) == 1: return constant(test.left) and all(constant(c) for c in test.comparators) if isinstance(test, ast.BoolOp): - return all(_static_test(v) for v in test.values) - return False + return all(_static_test(v, statics) for v in test.values) + if isinstance(test, ast.UnaryOp) and isinstance(test.op, ast.Not): + return _static_test(test.operand, statics) + if _jnp_logic(test) is not None: + # The expression mapping ran first: ``.not. ( k == I )`` arrived as + # ``jnp.logical_not(k == I)``. Static when its operands are. + return all(_static_test(a, statics) for a in test.args) + return constant(test) def _names(ids, ctx): @@ -445,14 +692,33 @@ class KernelLowerer: ast.ImportFrom, ) - def __init__(self): + def __init__(self, bound=(), statics=()): self.n = 0 + self.statics = frozenset(statics) + # The names bound so far, in statement order: a loop carries only a + # name that exists before it. One first assigned inside the body + # (the anchor's ``_out = callee(...)`` result tuple, unpacked on the + # next lines) is the body's own -- carried, its initial value would + # be read before any assignment, and its shape may change between + # two calls in the same body. + self.bound = set(bound) + + def _bind(self, stmts): + for name in _assigned_names(stmts): + self.bound.add(name) def lower_block(self, stmts, depth): - out = [] + out: list[ast.stmt] = [] for s in stmts: + start = len(out) if isinstance(s, self.BANNED): raise JaxQueue(f"unsupported stmt {type(s).__name__}") + if isinstance(s, ast.Continue | ast.Break): + # A CYCLE the loop pass could not fold into a branch, or an + # EXIT: neither has a place in a fori_loop body. Delegated, + # not emitted -- a ``continue`` inside a ``lax.cond`` branch + # is a SyntaxError that takes the whole module down. + raise JaxQueue(f"{type(s).__name__.lower()} inside a lowered loop") if isinstance(s, ast.Return) and depth > 0: raise JaxQueue("return inside loop/branch body") if isinstance(s, ast.For): @@ -460,9 +726,14 @@ def lower_block(self, stmts, depth): elif isinstance(s, ast.If): out.extend(self.lower_if(s, depth)) elif isinstance(s, ast.Assign): - out.append(self.lower_assign(s)) + lowered = self.lower_assign(s) + out.extend(lowered if isinstance(lowered, list) else [lowered]) else: out.append(s) # Expr (docstring), Return at top level, Pass + # Bound by what was *emitted*: a lowered loop or branch binds its + # carried names and nothing else -- a body-local of one loop is + # not bound for the next. + self._bind(out[start:]) return out def lower_assign(self, s): @@ -473,7 +744,20 @@ def lower_assign(self, s): # multi-output intra-module call: a, b = _callee_k_impl(...) if all(isinstance(e, ast.Name) for e in t.elts): return s - raise JaxQueue("tuple target with non-name elements") + # ``lo[i - 1, :], hi[i - 1, :] = split(...)`` (CLUBB's column + # loops): the tuple through a temporary, each element a store + # of its own -- a subscript store is lowered on its own line. + self.n += 1 + tmp = f"_t{self.n}" + stores = [ast.Assign(targets=[ast.Name(id=tmp, ctx=ast.Store())], value=s.value)] + for index, element in enumerate(t.elts): + piece = ast.Subscript( + value=ast.Name(id=tmp, ctx=ast.Load()), + slice=ast.Constant(value=index), + ctx=ast.Load(), + ) + stores.append(self.lower_assign(ast.Assign(targets=[element], value=piece))) + return stores if isinstance(t, ast.Name): # strengthen scalar literal inits so fori_loop carries keep a # stable strong dtype (0.0 -> jnp.float64(0.0)) @@ -537,8 +821,6 @@ def lower_assign(self, s): raise JaxQueue(f"unsupported assign target {type(t).__name__}") def lower_for(self, s, depth): - if s.orelse: - raise JaxQueue("for-else") if not isinstance(s.target, ast.Name): raise JaxQueue("tuple loop target") it = s.iter @@ -549,37 +831,148 @@ def lower_for(self, s, depth): and not it.keywords ): raise JaxQueue("non-range for") - step = 1 + step: int | None = 1 # Annotated because the first branch would otherwise pin these to # Constant and the others assign general expressions to them. lo: ast.expr hi: ast.expr + stride: ast.expr | None = None if len(it.args) == 1: lo, hi = ast.Constant(value=0), it.args[0] elif len(it.args) == 2: lo, hi = it.args elif len(it.args) == 3: - step = _const_int(it.args[2]) lo, hi = it.args[0], it.args[1] + try: + step = _const_int(it.args[2]) + except JaxQueue: + step = None if step not in (1, -1): - raise JaxQueue("range step not +-1") + # ``do k = lb, ub, dir``: a stride that is a name or an + # expression (a grid's direction, +-1 at run time). Trips + # counted by the runtime, the index remapped from the + # trip: k = lo + t * stride. + step, stride = None, it.args[2] else: raise JaxQueue("malformed range") - body = self.lower_block(s.body, depth + 1) + body_stmts = _cycle_to_else(s.body, []) + flag: str | None = None + pre_flag: list[ast.stmt] = [] + if _breaks_at_level(body_stmts): + # EXIT: a flag the loop carries. Set where the break was, it + # guards the rest of that trip and the whole of every later + # one -- the loop runs its trips and does nothing after the + # exit, which is what the exit left it doing. A for-else + # (the DO ran to completion) runs after the loop under the + # same flag, inverted. + self.n += 1 + flag = f"_brk_{self.n}" + # The index at the exit is what the DO variable holds after the + # loop (CLUBB's window search reads it): captured with the + # flag, in the index's own dtype, and rebound after the loop. + kept = f"_kx_{self.n}" + body_stmts = [ + ast.If( + test=_not_flag(flag), + body=_guard_after_breaks(body_stmts, flag, kept, s.target.id), + orelse=[], + ) + ] + pre_flag = [ + ast.Assign(targets=[ast.Name(id=flag, ctx=ast.Store())], value=_jnp_bool(False)), + ast.Assign( + targets=[ast.Name(id=kept, ctx=ast.Store())], + value=ast.Call( + func=ast.Attribute( + value=ast.Name(id="jnp", ctx=ast.Load()), attr="asarray", ctx=ast.Load() + ), + args=[copy.deepcopy(lo)], + keywords=[ + ast.keyword( + arg="dtype", + value=ast.Attribute( + value=ast.Name(id="jnp", ctx=ast.Load()), + attr="int32", + ctx=ast.Load(), + ), + ) + ], + ), + ), + ] + self.bound.add(flag) + self.bound.add(kept) + elif s.orelse: + raise JaxQueue("for-else without a break") + bound_before = set(self.bound) + self.bound.add(s.target.id) + body = self.lower_block(body_stmts, depth + 1) + after: list[ast.stmt] = [] + if flag is not None: + after = [ + ast.Assign( + targets=[ast.Name(id=s.target.id, ctx=ast.Store())], + value=ast.Name(id=kept, ctx=ast.Load()), + ) + ] + if s.orelse and flag is not None: + # The for-else stores the DO variable's completion value, a + # Python int when the bounds are static; the carry it joins + # is the int32 index kept at the exit. + completion: list[ast.stmt] = [] + for st in s.orelse: + if ( + isinstance(st, ast.Assign) + and len(st.targets) == 1 + and isinstance(st.targets[0], ast.Name) + and st.targets[0].id == s.target.id + ): + st = ast.Assign( + targets=st.targets, + value=ast.Call( + func=ast.Attribute( + value=ast.Name(id="jnp", ctx=ast.Load()), + attr="asarray", + ctx=ast.Load(), + ), + args=[st.value], + keywords=[ + ast.keyword( + arg="dtype", + value=ast.Attribute( + value=ast.Name(id="jnp", ctx=ast.Load()), + attr="int32", + ctx=ast.Load(), + ), + ) + ], + ), + ) + completion.append(st) + after += self.lower_block( + [ast.If(test=_not_flag(flag), body=completion, orelse=[])], depth + ) + if not body or all(isinstance(b, ast.Pass) for b in body): + # A loop whose body lowered to nothing: statistics calls the + # stand-in dropped, a nested loop that went the same way. + # Fortran ran it for nothing; nothing is what it becomes. + return [] settled = _trace_constant_stores(body) - carried = [n for n in _assigned_names(body) if n != s.target.id and n not in settled] + carried = [ + n + for n in _assigned_names(body) + if n != s.target.id and n not in settled and n in bound_before + ] if not carried: raise JaxQueue("loop with no carried effects") self.n += 1 fname = f"_body_{self.n}" - carry_call = ast.Call( - func=ast.Attribute( - value=ast.Name(id="lax", ctx=ast.Load()), attr="fori_loop", ctx=ast.Load() - ), - args=[], - keywords=[], - ) + # Through the runtime's ``_f_fori``: a trip count that is static + # and empty (a loop over an array's zero-extent axis -- CLUBB's + # scalar tracers under sclr_dim = 0) is skipped rather than traced, + # because JAX refuses any index into a size-0 axis at trace time. + carry_call = ast.Call(func=ast.Name(id="_f_fori", ctx=ast.Load()), args=[], keywords=[]) result = ast.Assign( targets=[ast.Tuple(elts=_names(carried, ast.Store), ctx=ast.Store())], value=carry_call ) @@ -588,7 +981,46 @@ def lower_for(self, s, depth): if step == 1: fn = self._carry_fn(fname, [s.target.id, "_c"], carried, body) carry_call.args = [lo, hi, ast.Name(id=fname, ctx=ast.Load()), init] - return [fn, result] + return [*pre_flag, fn, result, *after] + + if stride is not None: + lo_name, st_name, cnt_name = f"_lo_{self.n}", f"_st_{self.n}", f"_cnt_{self.n}" + pre = [ + ast.Assign(targets=[ast.Name(id=lo_name, ctx=ast.Store())], value=lo), + ast.Assign(targets=[ast.Name(id=st_name, ctx=ast.Store())], value=stride), + ast.Assign( + targets=[ast.Name(id=cnt_name, ctx=ast.Store())], + value=ast.Call( + func=ast.Name(id="_f_trips", ctx=ast.Load()), + args=[ + ast.Name(id=lo_name, ctx=ast.Load()), + hi, + ast.Name(id=st_name, ctx=ast.Load()), + ], + keywords=[], + ), + ), + ] + remap = ast.Assign( + targets=[ast.Name(id=s.target.id, ctx=ast.Store())], + value=ast.BinOp( + left=ast.Name(id=lo_name, ctx=ast.Load()), + op=ast.Add(), + right=ast.BinOp( + left=ast.Name(id="_r", ctx=ast.Load()), + op=ast.Mult(), + right=ast.Name(id=st_name, ctx=ast.Load()), + ), + ), + ) + fn = self._carry_fn(fname, ["_r", "_c"], carried, [remap, *body]) + carry_call.args = [ + ast.Constant(value=0), + ast.Name(id=cnt_name, ctx=ast.Load()), + ast.Name(id=fname, ctx=ast.Load()), + init, + ] + return [*pre_flag, *pre, fn, result, *after] # step -1: Fortran DO k=hi,lo,-1 arrived as range(hi, stop, -1) # iterating hi..stop+1. Remap: t in [0, hi-stop), k = hi - t. @@ -616,7 +1048,7 @@ def lower_for(self, s, depth): ast.Name(id=fname, ctx=ast.Load()), init, ] - return [*pre, fn, result] + return [*pre_flag, *pre, fn, result, *after] def lower_if(self, s, depth): """if/elif/else -> lax.cond with the union of both branches' @@ -631,20 +1063,67 @@ def lower_if(self, s, depth): a None arg — is never traced).""" body = self.lower_block(s.body, depth + 1) orelse = self.lower_block(s.orelse, depth + 1) + if not body and not orelse: + # Both arms lowered to nothing (a debug print under a level + # check, statistics under their switch): no branch at all. + return [] + if not body: + body = [ast.Pass()] # a Python if needs a body; the else is the point if _static_test(s.test): - return [ast.If(test=s.test, body=body, orelse=orelse or [])] + return [ast.If(test=_python_logic(s.test), body=body, orelse=orelse or [])] + if _static_test(s.test, self.statics): + # Over the kernel's static scalar arguments: a Python if when + # the kernel is called through its jit wrapper (the arguments + # are Python ints), and the lax.cond when another kernel calls + # the implementation from a traced body (CLUBB's mono_cubic_interp + # takes its level indices as arguments, and its caller's loop + # index reaches it as a tracer). Decided at trace time by the + # runtime's ``_f_concrete``. + # The Python if is spelled with Python logic: ``jnp.logical_not`` + # of a constant is staged under jit like any other op, a tracer + # no ``if`` can convert. Whether the leaves (the comparisons, + # the names) are concrete is what decides the form. + python_form: list[ast.stmt] = [ + ast.If(test=_python_logic(s.test), body=body, orelse=orelse or []) + ] + try: + cond_form = self._cond_form(s, copy.deepcopy(body), copy.deepcopy(orelse)) + except JaxQueue: + return python_form # nothing a cond could carry: the branch is a guard alone + if not cond_form: + return python_form + return [ + ast.If( + test=ast.Call( + func=ast.Name(id="_f_concrete", ctx=ast.Load()), + args=[copy.deepcopy(leaf) for leaf in _logic_leaves(s.test)], + keywords=[], + ), + body=python_form, + orelse=cond_form, + ) + ] + return self._cond_form(s, body, orelse) + + def _cond_form(self, s, body, orelse): + """The lax.cond lowering of an if whose test is traced.""" carried = _assigned_names(body) for n in _assigned_names(orelse): if n not in carried: carried.append(n) if not carried: + # A guard around dropped logs/aborts, or around a call statement + # whose result nothing binds (a check the port left on the host, + # under a traced guard -- CLUBB's parameterization check under + # ``any(err_code == fatal)``): nothing to carry, and nothing a + # tracer could run. The tree's notes name the host calls. trivial = all( isinstance(st, ast.Pass) - or (isinstance(st, ast.Expr) and isinstance(st.value, ast.Constant)) + or (isinstance(st, ast.Expr) and isinstance(st.value, ast.Constant | ast.Call)) for st in [*body, *orelse] ) if trivial: - return [] # a guard around dropped logs/aborts: nothing to carry + return [] raise JaxQueue("IF with no carried effects") self.n += 1 t_name, f_name = f"_true_{self.n}", f"_false_{self.n}" @@ -705,7 +1184,16 @@ def split_params(fn_src): return params[:n_req], params[n_req:], fn_src.args.defaults -def emit_kernel(fn_src, sub, closure, call_map=None, known_subs=None, writes=()): +def emit_kernel( + fn_src, + sub, + closure, + call_map=None, + known_subs=None, + writes=(), + traced_scalars=frozenset(), + hosted=None, +): """Original numpy FunctionDef -> unparsed __k_impl source. Kernel signature: [required..., closure..., optional-with-defaults]. @@ -756,8 +1244,15 @@ def visit_FunctionDef(self, node): fn.body[at] = Extend().visit(stmt) _bind_writer_calls(fn, call_map or {}) ExprMap().visit(fn) - CallRewrite(call_map or {}, known_subs or set()).visit(fn) - fn.body = KernelLowerer().lower_block(fn.body, 0) + calls = CallRewrite(call_map or {}, known_subs or set()) + calls.visit(fn) + if hosted is not None: + hosted.extend(calls.host_calls) + nums, _ = static_spec(fn_src, sub, traced_scalars) + required, _, _ = split_params(fn_src) + statics = {required[pos] for pos in nums} + lowerer = KernelLowerer(bound={a.arg for a in fn.args.args}, statics=statics) + fn.body = lowerer.lower_block(fn.body, 0) ast.fix_missing_locations(fn) return ast.unparse(fn) @@ -836,7 +1331,12 @@ def static_spec(fn_src, sub, traced_scalars=frozenset()): # when it is module state (the ``__`` spelling: dtime_ml, dtstep # -- namelist configuration): an ordinary real dummy (xa, xb, # tol) must stay traced, or jvp/grad against it breaks. - if r["dtype"] in ("int32", "str") or ( + # A logical scalar dummy is a switch (a model's configuration + # flags): static, so its branches are Python ifs at trace time + # and a path the run never takes is never traced. When another + # kernel passes it a tracer, the runtime's ``_f_concrete`` + # takes the lax.cond form instead. + if r["dtype"] in ("int32", "str", "bool") or ( r["dtype"] in ("int64", "float64") and "__" in r["name"] ): nums.append(pos) @@ -846,16 +1346,17 @@ def static_spec(fn_src, sub, traced_scalars=frozenset()): def build_module( interface: dict[str, Any], tree: ast.Module, traced_scalars: frozenset[str] = frozenset() -) -> tuple[list[str], list[str], dict[str, str]]: +) -> tuple[list[str], list[str], dict[str, str], dict[str, list[str]]]: """Emit all kernels of one module to a fixpoint. - Returns (pieces, jitted, delegated) where pieces are source chunks + Returns (pieces, jitted, delegated, hosted) where pieces are source chunks (kernels + jit lines + wrappers + delegations + _JAX_KERNELS) and delegated maps name -> reason.""" fns = {n.name: n for n in tree.body if isinstance(n, ast.FunctionDef)} subs = {s["name"]: s for s in interface["subprograms"]} delegated = {} + hosted: dict[str, list[str]] = {} kernels = set() for name, rec in subs.items(): if name not in fns: @@ -891,12 +1392,25 @@ def build_module( } for name in sorted(emit_set): call_map = {k: v for k, v in call_map_all.items() if k != name} + calls_kept: list[str] = [] try: srcs[name] = emit_kernel( - fns[name], subs[name], closures[name], call_map, set(subs), wclosures[name] + fns[name], + subs[name], + closures[name], + call_map, + set(subs), + wclosures[name], + traced_scalars, + calls_kept, ) except JaxQueue as e: failed[name] = f"[emit] {e}" + continue + if calls_kept: + hosted[name] = sorted(set(calls_kept)) + else: + hosted.pop(name, None) if not failed: break for n, r in failed.items(): @@ -970,7 +1484,7 @@ def build_module( if rec["name"] in delegated and rec["name"] in fns: pieces.append(f"{rec['name']} = _host.{rec['name']}") pieces.append(f"_JAX_KERNELS = {sorted(jitted)!r}") - return pieces, jitted, delegated + return pieces, jitted, delegated, hosted HEADER = '''"""Machine-generated by recast.transform.jax -- JAX backend (EXPERIMENTAL). diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py index 05d3973..3c05951 100644 --- a/src/recast/transform/jax/runtime.py +++ b/src/recast/transform/jax/runtime.py @@ -28,6 +28,7 @@ """ import jax +import numpy as np jax.config.update("jax_enable_x64", True) @@ -37,8 +38,11 @@ # A star-import from the generated module must see the underscore names. __all__ = [ "_f_adjustl", + "_f_concrete", "_f_dim", + "_f_ecall", "_f_epsilon", + "_f_fori", "_f_huge", "_f_int_div", "_f_len_trim", @@ -48,10 +52,14 @@ "_f_modulo", "_f_nint", "_f_sign", + "_f_sqrt", "_f_tiny", "_f_trim", + "_f_trips", "_f_vceil", "_f_vdot", + "_f_verf", + "_f_verfc", "_f_vexp", "_f_vfloor", "_f_vlog", @@ -59,6 +67,7 @@ "_f_vmax", "_f_vmin", "_f_vpow", + "_f_vsum", "_fstr_eq", "jax", "jnp", @@ -147,6 +156,109 @@ def _f_vpow(a, b): return jnp.asarray(a) ** b +def _f_concrete(*values): + """Whether every value is known at trace time -- a Python or NumPy + scalar, or a concrete array -- rather than a tracer. A branch over a + kernel's static scalar argument is a Python ``if`` when the kernel runs + through its jit wrapper and a ``lax.cond`` when another kernel's traced + body calls its implementation; this decides which, over the leaves of + the test (its comparisons and names), since the test itself is spelled + with Python logic the Python ``if`` evaluates.""" + return not any(isinstance(x, jax.core.Tracer) for x in values) + + +def _f_fori(lo, hi, body, init): + """``lax.fori_loop`` unless the trip count is static and empty. + + A Fortran DO over an array's zero-extent axis (CLUBB's scalar tracers + under ``sclr_dim = 0``) runs no iteration; ``fori_loop`` would still + trace the body once, and JAX refuses any index into a size-0 axis at + trace time. A dynamic bound is left to ``fori_loop``. + """ + static = (int, np.integer) + if isinstance(lo, static) and isinstance(hi, static): + if int(hi) <= int(lo): + return init + # A static trip count: a scan over the indices (reverse- + # differentiable, as fori_loop's own scan form is), spelled int32 + # -- Fortran's default integer, the dtype every integer local and + # dummy carries, so a store of the index into one keeps its dtype + # across a lax.cond. fori_loop's scan form would count in the + # default int, int64 under x64. + indices = jnp.arange(int(lo), int(hi), dtype=jnp.int32) + + def step(carry, i): + return body(i, carry), None + + carry, _ = lax.scan(step, init, indices) + return carry + return lax.fori_loop(lo, hi, body, init) + + +def _f_trips(lo, hi, step): + """How many times ``range(lo, hi, step)`` runs -- Fortran's DO trip + count, ``max(0, (hi - lo + step - sign(step)) // step)`` -- as a + Python int when every bound is, else traced.""" + static = (int, np.integer) + if isinstance(lo, static) and isinstance(hi, static) and isinstance(step, static): + return len(range(int(lo), int(hi), int(step))) + return jnp.maximum(0, (hi - lo + step - jnp.sign(step)) // step) + + +def _f_ecall(fn, *args, **kw): + """ELEMENTAL procedure broadcast over array actuals: the scalar kernel + per element, in sequence, as the NumPy runtime's np.vectorize runs it. + + Not jnp.vectorize: under jit that is a vmap, which turns the kernel's + branches into selects and its loops into batched ones, and CLUBB's + hybrid PDF closure came out 1e7 ULP from the anchor that way (2 ULP + this way). lax.map keeps each element's own control flow.""" + arrays = [jnp.asarray(a) for a in args] + shape = jnp.broadcast_shapes(*[a.shape for a in arrays]) + flat = tuple(jnp.broadcast_to(a, shape).reshape(-1) for a in arrays) + outs = lax.map(lambda xs: fn(*xs, **kw), flat) + return jax.tree_util.tree_map(lambda o: o.reshape(shape), outs) + + +def _f_sqrt(x): + """Fortran SQRT: a NaN for a negative real, not an exception, and the + correctly rounded root otherwise -- what the NumPy shim does with + math.sqrt, and what jnp.sqrt does by itself.""" + return jnp.sqrt(x) + + +def _f_vsum(a, axis=None): + """Fortran SUM accumulates in element order; a sequential fori_loop + keeps the fold order the NumPy anchor uses (its ``_f_vsum``), so the + two sides differ by XLA's rounding alone and not by association.""" + arr = jnp.asarray(a) + if axis is None: + flat = jnp.ravel(arr, order="F") + + def body(i, s): + return s + flat[i] + + return lax.fori_loop(0, flat.shape[0], body, jnp.zeros((), dtype=arr.dtype)) + moved = jnp.moveaxis(arr, axis, 0) + + def body_axis(i, s): + return s + moved[i] + + return lax.fori_loop(0, moved.shape[0], body_axis, jnp.zeros(moved.shape[1:], dtype=arr.dtype)) + + +def _f_verf(x): + from jax.scipy.special import erf as _erf + + return _erf(x) + + +def _f_verfc(x): + from jax.scipy.special import erfc as _erfc + + return _erfc(x) + + def _f_vdot(a, b): """Fortran DOT_PRODUCT accumulates in order; sequential fori_loop keeps the fold order (XLA may still contract the FMA).""" diff --git a/src/recast/transform/jax/translate.py b/src/recast/transform/jax/translate.py index a70bbcb..f00d448 100644 --- a/src/recast/transform/jax/translate.py +++ b/src/recast/transform/jax/translate.py @@ -63,7 +63,7 @@ def apply(self, unit: Unit, facts: Facts, config: dict[str, Any]) -> Candidate: source = anchor.files[Path(f"{module}_numpy.py")].decode() tree = ast.parse(source) - pieces, jitted, delegated = build_module(facts.interface, tree) + pieces, jitted, delegated, hosted = build_module(facts.interface, tree) emitted = ( HEADER.format(module=module, constants=constants_stem, runtime=runtime_stem) + _signatures_of(tree) @@ -88,6 +88,9 @@ def apply(self, unit: Unit, facts: Facts, config: dict[str, Any]) -> Candidate: "anchor": f"{module}_numpy.py", "kernels": sorted(jitted), "delegated": dict(sorted(delegated.items())), + "host_calls": { + k: [f"{module}.{n}" for n in v] for k, v in sorted(hosted.items()) + }, "runtime": f"{runtime_stem}.py", }, }, diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py index 840a26f..fc294ff 100644 --- a/src/recast/transform/jax/tree.py +++ b/src/recast/transform/jax/tree.py @@ -193,7 +193,8 @@ def __init__( self.inits: dict[str, str] = {} # local -> "int32" | "float64", from its guard init self.loop_vars: set[str] = set() # loop counters: int64 under x64, cast when stored self.state_params: list[str] = [] # a helper's module state, taken as parameters - self.buffer_outs: dict[str, list[str]] = {} # anchor _out buffers, elided + self.buffer_outs: dict[str, list[str | None]] = {} # anchor _out buffers, elided + self.host_calls: list[str] = [] # companion procedures left on the host, by name self.masked: list[str] = [] # statements whose dynamic slices became masks self.static_loops: list[str] = [] # loops whose trip count became static @@ -354,6 +355,40 @@ def visit_Assign(self, node: ast.Assign) -> Any: return None if isinstance(node.value, ast.Call) and self._flat_callee(node.value) is not None: return self._rewrite_call(node.targets[0], node.value) + value = node.value + if ( + len(node.targets) == 1 + and isinstance(value, ast.Subscript) + and isinstance(value.value, ast.Name) + and value.value.id in self.buffer_outs + and isinstance(value.slice, ast.Constant) + and isinstance(value.slice.value, int) + ): + # ``stats = _out[0]`` after ``_out = callee(...)`` (the anchor + # unpacking an object or an array the callee handed back): the + # buffer was elided at the call and the flat outputs bound to + # the actuals directly, so an object's unpack is already true + # and an array's is the actual it names -- itself, usually. + actuals = self.buffer_outs[value.value.id] + at = value.slice.value + if at >= len(actuals): + raise NotFlat(f"{ast.unparse(node)}: the elided buffer has no slot {at}") + target = node.targets[0] + if isinstance(target, ast.Name) and target.id == "_": + # ``_ = _out[4]``: an output the anchor discards (an optional + # OUT the caller did not ask for, absent on both sides). + return None + if isinstance(target, ast.Name) and self.spelling.object_of(target) is not None: + return None # an object: its components bound at the call + actual = actuals[at] + if actual is None: + raise NotFlat(f"{ast.unparse(node)}: the elided buffer's slot {at} has no actual") + if ast.unparse(target) == actual: + return None + bound_value = ast.parse(actual, mode="eval").body + return self.visit( + ast.copy_location(ast.Assign(targets=[target], value=bound_value), node) + ) pending = self._companion_writes(node.value) if isinstance(node.value, ast.Call) else [] self.generic_visit(node) if pending: @@ -425,6 +460,10 @@ def visit_Expr(self, node: ast.Expr) -> Any: call = node.value if isinstance(call, ast.Call) and self._flat_callee(call) is not None: return self._rewrite_call(None, call) + if isinstance(call, ast.Call): + bound = self._bind_buffer_outputs(node, call) + if bound is not None: + return bound # ``_f_copy_out(dst, src)``: an in-place copy is an assignment here. if ( isinstance(call, ast.Call) @@ -625,7 +664,7 @@ def visit_If(self, node: ast.If) -> Any: # under tracing, and the recorded run it is gated on never did; the # check is dropped and named on the candidate, so the evidence says # the kernel no longer aborts where the model would. - if node.body and all(isinstance(s, (ast.Raise, ast.Expr, ast.Pass)) for s in node.body): + if node.body and all(isinstance(s, ast.Raise) or _inert(s) for s in node.body): if any(isinstance(s, ast.Raise) for s in node.body): self.aborts.append(ast.unparse(node.test)) if not node.orelse: @@ -638,10 +677,24 @@ def visit_If(self, node: ast.If) -> Any: replaced.extend(seen if isinstance(seen, list) else [seen]) return replaced or None self.generic_visit(node) - if all(isinstance(s, (ast.Pass, ast.Expr)) for s in node.body): + if all(_inert(s) for s in node.body): # ``if cond: write(iulog, ...)`` -- a log line the anchor already - # left as ``pass``; nothing to carry. - return [*node.orelse] if node.orelse else None + # left as ``pass``; nothing to carry. A call statement is not + # inert: a procedure the port left on the host, under its guard, + # must stay where it is. An else branch keeps its + # guard, inverted: the fold of an early return leaves exactly + # ``if returned: pass else: ``, and running the rest + # unconditionally was the returned path's outputs overwritten. + if not node.orelse: + return None + return ast.copy_location( + ast.If( + test=ast.UnaryOp(op=ast.Not(), operand=node.test), + body=list(node.orelse), + orelse=[], + ), + node, + ) return node def visit_Raise(self, node: ast.Raise) -> Any: @@ -673,6 +726,41 @@ def visit_Call(self, node: ast.Call) -> ast.AST: raise NotFlat( f"{ast.unparse(node.func)} takes the object and is called inside an expression" ) + if ( + isinstance(node.func, ast.Name) + and node.func.id == "_f_ecall" + and node.args + and isinstance(node.args[0], ast.Attribute) + and isinstance(node.args[0].value, ast.Name) + ): + # ``_f_ecall(_new.calculate_mixture_fraction, ...)``: the + # elemental broadcast of a companion's procedure. Its kernel's + # implementation goes under the vectorize; a procedure the + # companion's port delegated is a host function no tracer can + # run, so the caller is not flat either. + callee = node.args[0] + owner = callee.value + assert isinstance(owner, ast.Name) + module = self.spelling.modules.get(owner.id) + port = self.ports.get(module) if module is not None else None + if port is not None: + if callee.attr not in port["kernels"]: + raise NotFlat( + f"elemental call of {module}.{callee.attr}, which its port did not lower" + ) + if (port.get("closures") or {}).get(callee.attr): + raise NotFlat( + f"elemental call of {module}.{callee.attr}, which reads module state" + ) + self.companions.add(owner.id) + node.args[0] = ast.copy_location( + ast.Attribute( + value=ast.Name(id=f"{owner.id}_jax", ctx=ast.Load()), + attr=f"_{callee.attr}_k_impl", + ctx=ast.Load(), + ), + callee, + ) self.generic_visit(node) # ``int(x)`` and ``np.float64(x)`` on a traced value: the cast the # anchor spells with a Python or NumPy constructor is ``jnp``'s here. @@ -753,6 +841,55 @@ def _companion_writes(self, node: ast.Call) -> list[str]: flats.append(flat) return flats + def _callee_record(self, call: ast.Call) -> dict[str, Any] | None: + """The interface record of a kernel a call statement reaches: one of + this module's, or a companion port's.""" + func = call.func + if isinstance(func, ast.Name): + return (self.own.get("records") or {}).get(func.id) + if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): + module = self.spelling.modules.get(func.value.id) + port = self.ports.get(module) if module is not None else None + if port is None or func.attr not in port.get("kernels", ()): + return None + return (port.get("records") or {}).get(func.attr) + return None + + def _bind_buffer_outputs(self, node: ast.Expr, call: ast.Call) -> ast.stmt | None: + """``callee(..., up3)`` as a statement: the anchor's callee wrote the + OUT array in place and returned it; the caller ignored the return. + A kernel cannot write in place -- its return *is* the output -- so + the statement binds what comes back to the actuals it was passed + (CLUBB's advance_xp3 calling skx_module's xp3_lg_2005_ansatz: + without this every xp3 came back as the zeros it went in as).""" + record = self._callee_record(call) + if record is None or record.get("kind") == "function": + return None + outs = [ + (at, a) + for at, a in enumerate(record["args"]) + if a["intent"] in ("OUT", "INOUT") and not a.get("optional") + ] + if not outs: + return None + by_keyword = {k.arg: k.value for k in call.keywords if k.arg} + targets: list[ast.expr] = [] + for at, a in outs: + actual = call.args[at] if at < len(call.args) else by_keyword.get(_py(a["name"])) + if actual is None or not isinstance(actual, ast.Name | ast.Subscript): + return None + target = copy.deepcopy(actual) + target.ctx = ast.Store() + targets.append(target) + target_node: ast.expr = ( + targets[0] if len(targets) == 1 else ast.Tuple(elts=targets, ctx=ast.Store()) + ) + # As an assignment, through the assignment's own path: the callee's + # spelling, its state closure and its written state bind there. + assign = ast.Assign(targets=[target_node], value=call) + bound: ast.stmt | None = self.visit(ast.copy_location(assign, node)) + return bound + # -- calls into companions ------------------------------------------------ def _companion_call(self, node: ast.Call) -> ast.AST: @@ -771,7 +908,17 @@ def _companion_call(self, node: ast.Call) -> ast.AST: # answers, not physics, and are left as they are. return node if func.attr not in port["kernels"]: - raise NotFlat(f"calls {module}.{func.attr}, which its port did not lower") + # The companion's port left this one on the host (CLUBB's + # numerical_check.parameterization_check: character arguments, + # writes). The call stays the host's, under whatever guard the + # anchor put it -- ``if clubb_at_least_debug_level_api(2)``, a + # stand-in's answer at trace time, false with statistics off -- + # and is never traced while the guard holds. If the guard ever + # lets it through, the trace fails on it by name, which the + # gate reports; refusing the whole kernel for a check that + # never runs was the alternative. + self.host_calls.append(f"{module}.{func.attr}") + return node closure: list[ast.expr] = [] for state in port["closures"].get(func.attr, []): flat = f"{module}__{state}" @@ -818,6 +965,13 @@ def _flat_callee(self, call: ast.Call) -> FlatPlan | None: self.spelling.object_of(k.value) is not None for k in call.keywords ): return callee + optional = {a["name"].lower() for a in callee.subprogram["args"] if a.get("optional")} + dummies = [obj for obj in callee.objects if obj.kind == "dummy"] + if dummies and all(obj.name in optional for obj in dummies): + # Every object the callee takes is optional and this call + # leaves them all out (pdf_closure without its implicit + # coefficients): still the flat function, with them absent. + return callee return None def _specialize(self, name: str, call: ast.Call, source: dict[str, Any]) -> FlatPlan | None: @@ -897,6 +1051,7 @@ def _specialize(self, name: str, call: ast.Call, source: dict[str, Any]) -> Flat if not body or not isinstance(body[-1], ast.Return): body.append(ast.Return(value=_tuple(_outputs(plan)))) taken = [_py(a["name"]) for a in plan.flat_args if a["intent"] != "OUT"] + body = [*_absent_optionals(plan, taken), *body] flat = ast.FunctionDef( name=plan.name, args=ast.arguments( @@ -1050,6 +1205,28 @@ def visit_Name(self, node: ast.Name) -> ast.AST: def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any: callee = self._flat_callee(call) assert callee is not None and self.plan is not None + if isinstance(call.func, ast.Attribute) and isinstance(call.func.value, ast.Name): + module = self.spelling.modules.get(call.func.value.id) + port = self.ports.get(module) if module is not None else None + if port is not None and callee.name not in port["kernels"]: + # The companion's port left this flat function on the host + # (CLUBB's numerical_check.parameterization_check: character + # arguments). The anchor's statement stays as it is, under + # whatever guard the anchor put it -- never traced while + # the guard holds, failing by name if it ever does -- and + # the note names it; refusing the whole kernel for a check + # that never runs was the alternative. + self.host_calls.append(f"{module}.{callee.name}") + if isinstance(target, ast.Name): + # The anchor's result buffer is the host's now; the + # unpacks that follow read it, under the same guard. + self.buffer_outs.pop(target.id, None) + kept: ast.stmt = ( + ast.Expr(value=call) + if target is None + else ast.Assign(targets=[target], value=call) + ) + return ast.copy_location(kept, call) dummies = [a["name"] for a in callee.subprogram["args"]] actual_by_dummy: dict[str, ast.expr] = {} for at, given in enumerate(call.args): @@ -1060,10 +1237,21 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any: actual_by_dummy[keyword_.arg.lower().rstrip("_")] = keyword_.value # Which callee object is which caller object. objects: dict[str, str] = {} + optional_dummies = { + a["name"].lower() for a in callee.subprogram["args"] if a.get("optional") + } + # An optional object the caller leaves out: the callee's kernel + # still takes its components (the flat signature has them), and + # its ``present(obj)`` is false at trace time. ``None`` in every + # slot, and whatever the kernel hands back for them is dropped. + absent: set[str] = set() for obj in callee.objects: if obj.kind == "dummy": actual = actual_by_dummy.get(obj.name) passed: str | None = self.spelling.object_of(actual) if actual is not None else None + if passed is None and actual is None and obj.name in optional_dummies: + absent.add(obj.name) + continue if passed is None: raise NotFlat(f"{callee.subprogram['name']}: object {obj.name} not passed") objects[obj.name] = passed @@ -1079,6 +1267,44 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any: continue # returned, not taken -- the emitted convention if name == callee.patch_count: args.append(ast.Name(id=self.plan.patch_count, ctx=ast.Load())) + elif "__" in name and name.split("__", 1)[0] in absent: + args.append(ast.Constant(value=None)) + elif name in callee.extent_args: + # An extent the callee's plan could not spell (an allocatable + # component sized by an allocating routine's dummy): the + # caller's own extent argument for the same axis when it + # has one, else that axis of the caller's component, a + # static shape at trace time. + owner, comp, axis = callee.extent_args[name] + passed_obj = objects.get(owner) + own = next( + ( + n + for n, (o, c, ax) in self.plan.extent_args.items() + if o == passed_obj and c == comp and ax == axis + ), + None, + ) + if own is not None: + args.append(ast.Name(id=own, ctx=ast.Load())) + continue + sized = caller_components.get(passed_obj or "", {}).get(comp) + if sized is None: + raise NotFlat( + f"{callee.subprogram['name']}: extent {name} of {owner}%{comp}" + " not in the caller's plan" + ) + args.append( + ast.Subscript( + value=ast.Attribute( + value=ast.Name(id=sized, ctx=ast.Load()), + attr="shape", + ctx=ast.Load(), + ), + slice=ast.Constant(value=int(axis) - 1), # the plan's axis is 1-based + ctx=ast.Load(), + ) + ) elif "__" in name and (owner := name.split("__", 1)[0]) in objects: comp = name.split("__", 1)[1] flat_name: str | None = caller_components.get(objects[owner], {}).get(comp) @@ -1112,8 +1338,6 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any: elif isinstance(call.func, ast.Attribute) and isinstance(call.func.value, ast.Name): # A companion's flat function: its port's kernel, or nothing. module = self.spelling.modules[call.func.value.id] - if callee.name not in self.ports[module]["kernels"]: - raise NotFlat(f"calls {module}.{callee.name}, which its port did not lower") self.companions.add(call.func.value.id) func = ast.Attribute( value=ast.Name(id=f"{call.func.value.id}_jax", ctx=ast.Load()), @@ -1135,10 +1359,13 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any: anchor_targets = list(target.elts) elif target is not None: anchor_targets = [target] + # Every OUT/INOUT dummy, the optional ones included: the anchor's + # return tuple has a slot for each (an absent optional's is ``_``), + # and the unpacks that follow index it by position -- CLUBB's + # ``stats = _out[3]`` after calc_brunt_vaisala_freq_sqd, whose + # ``stats`` is an optional INOUT object. original_outs = [ - _py(a["name"]) - for a in callee.subprogram["args"] - if a["intent"] in ("OUT", "INOUT") and not a.get("optional") + _py(a["name"]) for a in callee.subprogram["args"] if a["intent"] in ("OUT", "INOUT") ] if callee.subprogram["kind"] == "function": original_outs = ["_result", *original_outs] @@ -1166,13 +1393,18 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any: for name in original_outs if name.lower() in actual_by_dummy } + # Positional, one entry per output the callee returns: the + # anchor unpacks ``_out[k]`` by that position, and an object's + # outputs (bound by component) leave their slot empty. self.buffer_outs[buffered.id] = [ - ast.unparse(slot[name]) for name in original_outs if name in slot + ast.unparse(slot[name]) if name in slot else None for name in original_outs ] targets: list[ast.expr] = [] follow: list[ast.stmt] = [] for name in _outputs(callee): - if "__" in name and (owner := name.split("__", 1)[0]) in objects: + if "__" in name and name.split("__", 1)[0] in absent: + targets.append(ast.Name(id="_", ctx=ast.Store())) + elif "__" in name and (owner := name.split("__", 1)[0]) in objects: comp = name.split("__", 1)[1] targets.append( ast.Name(id=caller_components[objects[owner]][comp], ctx=ast.Store()) @@ -1181,6 +1413,12 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any: targets.append(ast.Name(id=name, ctx=ast.Store())) else: where = slot.get(name) + passed_as = actual_by_dummy.get(name.lower()) + if passed_as is not None and _reshaped_name(passed_as) is not None: + # The actual is a reshaped view of a name (sequence + # association); the output comes back in the callee's + # shape, whatever the anchor's own target spelling. + where = passed_as if where is None: raise NotFlat(f"{callee.subprogram['name']}: output {name} has no target") bound: Any = self.visit(copy.deepcopy(where)) @@ -1190,9 +1428,38 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any: self.temps += 1 temp = f"_t{self.temps}" targets.append(ast.Name(id=temp, ctx=ast.Store())) - follow.append( - ast.Assign(targets=[bound], value=ast.Name(id=temp, ctx=ast.Load())) - ) + reshaped = _reshaped_name(bound) + if reshaped is not None: + # ``np.reshape(rhs, (n, m, 1), order='F')`` as the + # actual: sequence association, the anchor's view of + # the array in the callee's shape. What comes back + # is reshaped to the array's own shape, the same + # order, and rebinds it. + follow.append( + ast.Assign( + targets=[ast.Name(id=reshaped, ctx=ast.Store())], + value=ast.Call( + func=ast.Attribute( + value=ast.Name(id="jnp", ctx=ast.Load()), + attr="reshape", + ctx=ast.Load(), + ), + args=[ + ast.Name(id=temp, ctx=ast.Load()), + ast.Attribute( + value=ast.Name(id=reshaped, ctx=ast.Load()), + attr="shape", + ctx=ast.Load(), + ), + ], + keywords=[ast.keyword(arg="order", value=ast.Constant("F"))], + ), + ) + ) + else: + follow.append( + ast.Assign(targets=[bound], value=ast.Name(id=temp, ctx=ast.Load())) + ) if not targets: return ast.Expr(value=new_call) # One output comes back bare (the flat function returns a name, not @@ -1435,9 +1702,13 @@ def _rewritten_body(fn: ast.FunctionDef, rewrite: _Rewrite) -> list[ast.stmt]: rewrite.inits = _guard_inits(fn) rewrite.statics = frozenset(rewrite.statics) | _concrete_scalars(fn, rewrite) lowered: list[ast.stmt] = [] + body = [copy.deepcopy(s) for s in fn.body] + # Fixed-length windows with traced bounds first, so the mask rules never + # see them: ``x(k-w:k+w)`` is a gather at ``lo + arange(2w+1)``. + _window_slices(body, frozenset(rewrite.statics)) # Single exit first, on the anchor's own returns: the flat return that # replaces them is one statement at the end. - for statement in _single_exit([copy.deepcopy(s) for s in fn.body]): + for statement in _single_exit(body): result = rewrite.visit(statement) if result is None: continue @@ -1446,9 +1717,11 @@ def _rewritten_body(fn: ast.FunctionDef, rewrite: _Rewrite) -> list[ast.stmt]: n for n, kind in (rewrite.inits or {}).items() if kind == "int32" ) lowered = _WhileLoops(_integer_inits(fn), int_locals).visit_block(lowered) - # The goto-region flags are synthetic locals with no UB-guard init; an - # enclosing loop carries them, and the initial carry tuple needs a value - # before the loop. Every flag starts False at the top of the function. + # The goto-region and while-exit flags are synthetic locals with no + # UB-guard init; an enclosing loop or branch carries them (a while inside + # an if inside a do: CLUBB's grid interpolation), and the initial carry + # tuple needs a value before it. Every flag starts False at the top of + # the function; the one beside its loop resets it on every entry. flags = sorted( { n.id @@ -1456,7 +1729,7 @@ def _rewritten_body(fn: ast.FunctionDef, rewrite: _Rewrite) -> list[ast.stmt]: for n in ast.walk(stmt) if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Store) - and re.fullmatch(r"_(?:skip|restart)_\d+", n.id) + and re.fullmatch(r"_(?:skip|restart|done)_\d+", n.id) } ) lowered = [ @@ -1476,7 +1749,293 @@ def _rewritten_body(fn: ast.FunctionDef, rewrite: _Rewrite) -> list[ast.stmt]: and rewrite._is_dynamic(node.upper) ): raise NotFlat(f"a dynamic slice outside a store or a sum: {ast.unparse(node)}") - return lowered + return [*_rebuilt_objects(lowered, rewrite), *lowered] + + +def _rebuilt_objects(lowered: list[ast.stmt], rewrite: _Rewrite) -> list[ast.stmt]: + """A dummy object the body still hands whole to the host -- a stand-in's + query on it (``var_on_stats_list(stats, name)``), a check the port left + on the host under its guard -- is rebuilt once, at entry, from the + components the kernel takes, the way the NumPy flat wrapper rebuilds + it before its call. A Python record at trace time; the components are + what flow, and nothing carries the record.""" + if rewrite.plan is None: + return [] + # An optional dummy the plan leaves out is absent (``k = None`` at the + # top, by _absent_optionals): the body's ``present(k)`` must stay false. + optional = { + _py(a["name"]) for a in rewrite.plan.subprogram.get("args") or () if a.get("optional") + } + whole = { + node.id + for statement in lowered + for node in ast.walk(statement) + if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load) + } + rebuilt: list[ast.stmt] = [] + for obj in rewrite.plan.objects: + if obj.kind != "dummy" or obj.name not in whole or obj.name in optional: + continue + if not obj.components: + continue + rebuilt.append( + ast.fix_missing_locations( + ast.Assign( + targets=[ast.Name(id=obj.name, ctx=ast.Store())], + value=ast.Call( + func=ast.Attribute( + value=ast.Name(id="_host", ctx=ast.Load()), + attr="_Record", + ctx=ast.Load(), + ), + args=[], + keywords=[ + ast.keyword(arg=c.name, value=ast.Name(id=c.flat, ctx=ast.Load())) + for c in obj.components + ], + ), + ) + ) + ) + return rebuilt + + +Affine = tuple[dict[str, int], int] +"""A linear form: ``{name: coefficient}`` and a constant.""" + + +def _reshaped_name(node: ast.expr) -> str | None: + """``np.reshape(x, shape, order='F')`` / ``jnp.reshape(...)`` / + ``x.reshape(...)`` over a bare name: that name.""" + if not isinstance(node, ast.Call): + return None + func = node.func + if ( + isinstance(func, ast.Attribute) + and func.attr == "reshape" + and isinstance(func.value, ast.Name) + and func.value.id in ("np", "jnp") + and node.args + and isinstance(node.args[0], ast.Name) + ): + return node.args[0].id + if ( + isinstance(func, ast.Attribute) + and func.attr == "reshape" + and isinstance(func.value, ast.Name) + ): + return func.value.id + return None + + +def _static_name(name: str, statics: frozenset[str]) -> bool: + return name in statics or (name.isupper() and len(name) > 1) + + +def _static_atom(node: ast.expr, statics: frozenset[str]) -> bool: + """An expression over static names and constants only -- ``dir * n``, + a Python int at trace time -- taken whole as one term of a form.""" + names = [n for n in ast.walk(node) if isinstance(n, ast.Name)] + return bool(names) and all(_static_name(n.id, statics) for n in names) + + +def _affine( + node: ast.expr, env: dict[str, ast.expr | None], statics: frozenset[str], depth: int = 0 +) -> Affine | None: + """The linear form of an integer expression: ``{term: coefficient}`` + and a constant, a term being a traced name or a static atom (spelled + ``@``, over static names only). A local that was assigned an + affine expression in scope (``k_start = k - dir * n``) is + substituted; None where the expression is not linear.""" + if depth > 8: + return None + if isinstance(node, ast.Constant): + if isinstance(node.value, int) and not isinstance(node.value, bool): + return {}, int(node.value) + return None + if isinstance(node, ast.Name): + bound = env.get(node.id) + if bound is not None: + inner = _affine(bound, env, statics, depth + 1) + if inner is not None: + return inner + if _static_name(node.id, statics): + return {f"@{node.id}": 1}, 0 + return {node.id: 1}, 0 + if not isinstance(node, ast.Name | ast.Constant) and _static_atom(node, statics): + return {f"@{ast.unparse(node)}": 1}, 0 + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.UAdd | ast.USub): + inner = _affine(node.operand, env, statics, depth) + if inner is None: + return None + if isinstance(node.op, ast.UAdd): + return inner + return {k: -v for k, v in inner[0].items()}, -inner[1] + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add | ast.Sub | ast.Mult): + left = _affine(node.left, env, statics, depth) + right = _affine(node.right, env, statics, depth) + if left is None or right is None: + return None + if isinstance(node.op, ast.Mult): + if not left[0]: + scale, form = left[1], right + elif not right[0]: + scale, form = right[1], left + else: + return None + return {k: v * scale for k, v in form[0].items()}, form[1] * scale + sign = 1 if isinstance(node.op, ast.Add) else -1 + coeffs = dict(left[0]) + for k, v in right[0].items(): + coeffs[k] = coeffs.get(k, 0) + sign * v + return {k: v for k, v in coeffs.items() if v}, left[1] + sign * right[1] + if ( + isinstance(node, ast.Call) + and len(node.args) == 1 + and not node.keywords + and ( + (isinstance(node.func, ast.Name) and node.func.id == "int") + or ( + isinstance(node.func, ast.Attribute) + and node.func.attr in ("int32", "int64") + and isinstance(node.func.value, ast.Name) + and node.func.value.id in ("np", "jnp") + ) + ) + ): + return _affine(node.args[0], env, statics, depth) + return None + + +def _affine_ast(form: Affine) -> ast.expr: + """``2 * (dir * n) + 1`` back as an expression, over static atoms only.""" + coeffs, const = form + out: ast.expr = ast.Constant(value=const) + for key in sorted(coeffs): + atom: ast.expr = ast.parse(key[1:], mode="eval").body + term: ast.expr = ast.BinOp(left=ast.Constant(value=coeffs[key]), op=ast.Mult(), right=atom) + out = ast.BinOp(left=out, op=ast.Add(), right=term) + return out + + +class _Windows(ast.NodeTransformer): + """``x[i - 1, k_start - 1:k_end:dir]`` where ``k_start`` and ``k_end`` + are ``k -+ dir * n``: bounds that trace, a distance apart that does not. + The slice becomes the index ``lo + arange(trips) * step`` -- a gather + of static length at a traced offset, a scatter when stored -- with + the trip count the runtime's ``_f_trips(0, distance, step)`` over + static names, a Python int at trace time. Fortran's windows (CLUBB's + hole filler draws ``num_hf_draw_points`` levels either side of ``k``) + are this; a slice whose length depends on a traced bound is the mask + rules' business and is left to them.""" + + def __init__(self, env: dict[str, ast.expr | None], statics: frozenset[str]) -> None: + self.env = env + self.statics = statics + + def _dynamic(self, node: ast.expr) -> bool: + return any( + isinstance(n, ast.Name) and not _static_name(n.id, self.statics) for n in ast.walk(node) + ) + + def _window(self, element: ast.Slice) -> ast.expr | None: + if element.upper is None: + return None + lower: ast.expr = element.lower if element.lower is not None else ast.Constant(0) + if not (self._dynamic(lower) or self._dynamic(element.upper)): + return None + if element.step is not None and self._dynamic(element.step): + return None + hi = _affine(element.upper, self.env, self.statics) + lo = _affine(lower, self.env, self.statics) + if hi is None or lo is None: + return None + coeffs = dict(hi[0]) + for k, v in lo[0].items(): + coeffs[k] = coeffs.get(k, 0) - v + coeffs = {k: v for k, v in coeffs.items() if v} + if any(not k.startswith("@") for k in coeffs): + return None # a traced term survives: the length traces too + distance = _affine_ast((coeffs, hi[1] - lo[1])) + step: ast.expr = ( + copy.deepcopy(element.step) if element.step is not None else ast.Constant(1) + ) + trips = ast.Call( + func=ast.Name(id="_f_trips", ctx=ast.Load()), + args=[ast.Constant(0), distance, step], + keywords=[], + ) + offsets: ast.expr = _jnp("arange", [trips]) + if not (isinstance(step, ast.Constant) and step.value == 1): + offsets = ast.BinOp(left=offsets, op=ast.Mult(), right=copy.deepcopy(step)) + return ast.BinOp(left=copy.deepcopy(lower), op=ast.Add(), right=offsets) + + def visit_Subscript(self, node: ast.Subscript) -> ast.AST: + self.generic_visit(node) + elts = list(node.slice.elts) if isinstance(node.slice, ast.Tuple) else [node.slice] + changed = False + out: list[ast.expr] = [] + for element in elts: + window = self._window(element) if isinstance(element, ast.Slice) else None + if window is None: + out.append(element) + else: + out.append(window) + changed = True + if changed: + node.slice = ast.Tuple(elts=out, ctx=ast.Load()) if len(out) > 1 else out[0] + return node + + +def _window_slices(stmts: list[ast.stmt], statics: frozenset[str]) -> None: + """Rewrite the windows in a block in place, an affine local's + definition in scope where a bound names one.""" + + def block(body: list[ast.stmt], env: dict[str, ast.expr | None]) -> set[str]: + assigned: set[str] = set() + for st in body: + if isinstance(st, ast.Assign): + _Windows(env, statics).visit(st) + for target in st.targets: + if isinstance(target, ast.Name): + affine = _affine(st.value, env, statics) is not None + env[target.id] = st.value if affine else None + assigned.add(target.id) + elif isinstance(st, ast.AugAssign): + _Windows(env, statics).visit(st) + if isinstance(st.target, ast.Name): + env[st.target.id] = None + assigned.add(st.target.id) + elif isinstance(st, ast.If | ast.While): + _Windows(env, statics).visit(st.test) + inner = block(st.body, dict(env)) | block(st.orelse, dict(env)) + for name in inner: + env[name] = None + assigned |= inner + elif isinstance(st, ast.For): + _Windows(env, statics).visit(st.iter) + scope = dict(env) + if isinstance(st.target, ast.Name): + scope[st.target.id] = None + inner = block(st.body, scope) | block(st.orelse, dict(scope)) + if isinstance(st.target, ast.Name): + inner.add(st.target.id) + for name in inner: + env[name] = None + assigned |= inner + elif isinstance(st, ast.Try): + inner = block(st.body, dict(env)) + for handler in st.handlers: + inner |= block(handler.body, dict(env)) + inner |= block(st.orelse, dict(env)) | block(st.finalbody, dict(env)) + for name in inner: + env[name] = None + assigned |= inner + else: + _Windows(env, statics).visit(st) + return assigned + + block(stmts, {}) def _dim_sources(args: list[dict[str, Any]]) -> dict[str, tuple[str, int]]: @@ -1491,10 +2050,92 @@ def _dim_sources(args: list[dict[str, Any]]) -> dict[str, tuple[str, int]]: return sources +def _inert(statement: ast.stmt) -> bool: + """A statement with nothing to carry: ``pass``, or a bare constant (a + docstring, a log line the anchor left as a string).""" + return isinstance(statement, ast.Pass) or ( + isinstance(statement, ast.Expr) and isinstance(statement.value, ast.Constant) + ) + + def _has_return(stmts: list[ast.stmt]) -> bool: return any(isinstance(n, ast.Return) for s in stmts for n in ast.walk(s)) +def _fold_returns(stmts: list[ast.stmt], rest: list[ast.stmt]) -> list[ast.stmt] | None: + """``stmts`` followed by ``rest`` with every early ``return`` folded into + the branch structure: the continuation moves into the branches that do + not return, so the function's one return at the end is reached by all + paths with the outputs as the early return would have left them. None + when a return sits inside a loop, which this fold cannot express.""" + out: list[ast.stmt] = [] + for at, statement in enumerate(stmts): + if isinstance(statement, ast.Return): + return out # what follows is never reached on this path + if isinstance(statement, ast.For | ast.While) and _has_return([statement]): + return None + if isinstance(statement, ast.If) and _has_return([statement]): + # The continuation, itself folded: a return further down the + # same block would otherwise ride into the branch unfolded. + tail = _fold_returns(stmts[at + 1 :], rest) + if tail is None: + return None + body = _fold_returns(statement.body, copy.deepcopy(tail)) + orelse = _fold_returns(statement.orelse, copy.deepcopy(tail)) + if body is None or orelse is None: + return None + return [ + *out, + ast.copy_location( + ast.If(test=statement.test, body=body or [ast.Pass()], orelse=orelse), + statement, + ), + ] + out.append(statement) + return [*out, *rest] + + +def _guard_after_returns(stmts: list[ast.stmt]) -> list[ast.stmt]: + """``return`` -> ``_ret = True``; what follows a statement that may have + returned, in the same block, runs under one ``if not _ret`` -- the whole + remainder as one block, so a call's result buffer (``_out``, never a + carry) and its unpacks stay together -- recursively through branches + and loop bodies.""" + out: list[ast.stmt] = [] + for at, statement in enumerate(stmts): + had_return = _has_return([statement]) # before the rewrite takes the returns away + if isinstance(statement, ast.Return): + # ``jnp.bool_(True)``, not the literal: a literal store is a + # trace-time constant to the backend, which would not carry the + # flag out of the loop or the branch that set it. + rewritten: ast.stmt = ast.copy_location( + ast.Assign( + targets=[ast.Name(id="_ret", ctx=ast.Store())], + value=_jnp("bool_", [ast.Constant(True)]), + ), + statement, + ) + else: + rewritten = statement + for field in ("body", "orelse"): + inner = getattr(statement, field, None) + if isinstance(inner, list) and inner: + setattr(statement, field, _guard_after_returns(inner)) + out.append(rewritten) + if had_return: + rest = _guard_after_returns(stmts[at + 1 :]) + if rest: + out.append( + ast.If( + test=ast.UnaryOp(op=ast.Not(), operand=ast.Name(id="_ret", ctx=ast.Load())), + body=rest, + orelse=[], + ) + ) + return out + return out + + def _single_exit(body: list[ast.stmt]) -> list[ast.stmt]: """Early returns -- ``if f0 == 0: root = x0; return root`` -- become a flag and a value, every later statement runs under ``if not _ret``, and @@ -1504,13 +2145,45 @@ def _single_exit(body: list[ast.stmt]) -> list[ast.stmt]: if not early or not isinstance(body[-1], ast.Return) or body[-1].value is None: return body # The merge is one ``jnp.where`` over one value: a function with several - # outputs returns a tuple, which ``where`` cannot select. - for statement in body: - for node in ast.walk(statement): - if isinstance(node, ast.Return) and isinstance(node.value, ast.Tuple): - raise NotFlat( - f"an early return in a function with several outputs: {ast.unparse(node)}" - ) + # outputs returns a tuple, which ``where`` cannot select. When every + # early return hands back the same tuple the final one does (CLUBB's + # advance_clubb_core: ``if ( fatal ) return`` after each solver, the + # outputs as they stand), the returns fold into the branch structure + # instead: what follows a returning branch moves into the branches that + # do not return, and the one return at the end is reached by all. + tuples = [ + node + for statement in body + for node in ast.walk(statement) + if isinstance(node, ast.Return) and isinstance(node.value, ast.Tuple) + ] + if tuples: + final = body[-1] + final_value = final.value + same = final_value is not None and all( + node.value is not None and ast.dump(node.value) == ast.dump(final_value) + for node in tuples + ) + if not same: + raise NotFlat( + f"an early return in a function with several outputs: {ast.unparse(tuples[0])}" + ) + folded = _fold_returns(body[:-1], []) + if folded is not None: + return [*folded, final] + # A return inside a loop (advance_clubb_core checks the error code + # per column): a flag instead. The return sets it, every later + # statement at every level runs under ``if not _ret``, and the + # loop's remaining iterations do nothing; the one return at the + # end hands back the outputs as the early return left them. + return [ + ast.Assign( + targets=[ast.Name(id="_ret", ctx=ast.Store())], + value=_jnp("bool_", [ast.Constant(False)]), + ), + *_guard_after_returns(body[:-1]), + final, + ] class Returns(ast.NodeTransformer): def visit_Return(self, node: ast.Return) -> Any: @@ -1978,12 +2651,20 @@ def visit_While(self, node: ast.While) -> ast.AST: exited = True for statement in guarded: ast.fix_missing_locations(statement) - carried = [n for n in _assigned_names(guarded) if n != done] # type: ignore[no-untyped-call] - state = [*carried, done] try: lowered = KernelLowerer().lower_block(guarded, 1) # type: ignore[no-untyped-call] except JaxQueue as why: raise NotFlat(f"while loop body: {why}") from why + # The carries are what the *lowered* body stores: a subscript store + # (``idx(i) = k`` in a search loop) is a store to its base only once + # lowered. An inner loop's break flag and kept index are the body's + # own, initialized beside that loop, never a carry of this one. + carried = [ + n + for n in _assigned_names(lowered) # type: ignore[no-untyped-call] + if n != done and not re.fullmatch(r"_(?:brk|kx)_\d+", n) + ] + state = [*carried, done] unpack = ast.Assign( targets=[ ast.Tuple(elts=[ast.Name(id=n, ctx=ast.Store()) for n in state], ctx=ast.Store()) @@ -2116,6 +2797,14 @@ def _static_expression(node: ast.expr, statics: frozenset[str]) -> bool: return isinstance(node.value, (int, float)) and not isinstance(node.value, bool) if isinstance(node, ast.Name): return node.id.isupper() or node.id in statics + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id.startswith("_") + and node.attr.isupper() + and len(node.attr) > 1 + ): + return True # ``_mod.IIPDF_ADG1``: a module constant through its alias if isinstance(node, ast.BinOp) and isinstance( node.op, (ast.Add, ast.Sub, ast.Mult, ast.Div, ast.FloorDiv) ): @@ -2297,6 +2986,37 @@ def _static_names(args: list[dict[str, Any]]) -> frozenset[str]: ) +def _absent_optionals(plan: FlatPlan, taken: list[str]) -> list[ast.stmt]: + """Bindings for the optional dummies the flat signature leaves out. + + The plan drops an optional dummy: the adapter calls the original with it + absent, so its NumPy default (``None``) stands in. The kernel inlines + the original's body, which still names it -- CLUBB's grid interpolators + test ``present(zt_min)`` -- so the absence is spelled at the top: + ``zt_min = None``, and ``want_zt_min = False`` for an optional OUT, the + anchor's presence sentinel. A trace-time ``if x is not None`` then skips + the branch the Fortran skipped. + """ + out: list[ast.stmt] = [] + for a in plan.subprogram.get("args") or (): + if not a.get("optional"): + continue + name = _py(a["name"]) + if name in taken: + continue + out.append( + ast.Assign(targets=[ast.Name(id=name, ctx=ast.Store())], value=ast.Constant(value=None)) + ) + if a.get("intent") == "OUT": + out.append( + ast.Assign( + targets=[ast.Name(id=f"want_{name}", ctx=ast.Store())], + value=ast.Constant(value=False), + ) + ) + return out + + def flat_function( fn: ast.FunctionDef, plan: FlatPlan, @@ -2325,6 +3045,7 @@ def flat_function( if not body or not isinstance(body[-1], ast.Return): body.append(ast.Return(value=_tuple(_outputs(plan)))) taken = [_py(a["name"]) for a in plan.flat_args if a["intent"] != "OUT"] + body = [*_absent_optionals(plan, taken), *body] flat = ast.FunctionDef( name=plan.name, args=ast.arguments( @@ -2418,6 +3139,7 @@ def flattened_module( specialized: dict[str, tuple[ast.FunctionDef, FlatPlan, list[str]]] = {} aborts: dict[str, list[str]] = {} masked: dict[str, list[str]] = {} + host_calls: dict[str, list[str]] = {} static_loops: dict[str, list[str]] = {} companions: set[str] = set() planned = {p.subprogram["name"] for p in plans} @@ -2439,6 +3161,8 @@ def flattened_module( aborts[plan.name] = rewrite.aborts if rewrite.masked: masked[plan.name] = rewrite.masked + if rewrite.host_calls: + host_calls[plan.name] = sorted(set(rewrite.host_calls)) if rewrite.static_loops: static_loops[plan.name] = rewrite.static_loops companions |= rewrite.companions @@ -2461,6 +3185,8 @@ def flattened_module( aborts[name] = rewrite.aborts if rewrite.masked: masked[name] = rewrite.masked + if rewrite.host_calls: + host_calls[name] = sorted(set(rewrite.host_calls)) if rewrite.static_loops: static_loops[name] = rewrite.static_loops if rewrite.state_params: @@ -2486,6 +3212,7 @@ def flattened_module( "aborts_dropped": aborts, "masked": masked, "static_loops": static_loops, + "host_calls": host_calls, "companions": sorted(companions), } return module, {**interface, "subprograms": [*interface["subprograms"], *entries]}, notes @@ -2518,7 +3245,12 @@ def apply(self, unit: Unit, facts: Facts, config: dict[str, Any]) -> Candidate: flat_tree, interface, flat_notes = flattened_module( tree, facts.interface, plans, ported, bundled ) - pieces, jitted, delegated = build_module(interface, flat_tree, TRACED_SCALARS) + pieces, jitted, delegated, kept_on_host = build_module(interface, flat_tree, TRACED_SCALARS) + for name, kept in kept_on_host.items(): + # This module's own subprograms the backend left on the host, + # beside the companions' the flat rewrite did. + own = flat_notes.setdefault("host_calls", {}).setdefault(name, []) + flat_notes["host_calls"][name] = sorted(set(own) | {f"{module}.{n}" for n in kept}) # A flat function the backend delegated has a host to fall back on # only if the NumPy module carries its wrapper -- the gated ones. A # private subprogram's or a function's flat form has none, and a line @@ -2586,6 +3318,7 @@ def apply(self, unit: Unit, facts: Facts, config: dict[str, Any]) -> Candidate: "aborts_dropped": dict(sorted(flat_notes["aborts_dropped"].items())), "masked": dict(sorted(flat_notes["masked"].items())), "static_loops": dict(sorted(flat_notes["static_loops"].items())), + "host_calls": dict(sorted(flat_notes.get("host_calls", {}).items())), "companions": sorted(ported), "runtime": f"{runtime_stem}.py", "_ports": ported, diff --git a/src/recast/transform/numpy/constants.py b/src/recast/transform/numpy/constants.py index 7ffe6c5..524f49b 100644 --- a/src/recast/transform/numpy/constants.py +++ b/src/recast/transform/numpy/constants.py @@ -23,7 +23,7 @@ from pathlib import PurePath, PurePosixPath from typing import Any -from recast.fortran.expr import Expr, python_call, render, with_integer_division +from recast.fortran.expr import Expr, python_call, render, typed, with_integer_division __all__ = [ "constants_module", @@ -291,18 +291,21 @@ def use_constants_module(resolved: list[dict[str, Any]], module_name: str) -> st "import numpy as np", "", ] + env: dict[str, str | None] = {} for entry in resolved: - value = _python(entry["expr"]) + value = _python(entry["expr"], env) + env[entry["name"]] = entry.get("dtype") or typed(entry["expr"], env) where = f"{PurePath(entry['source']).name}:{entry['line']}" lines.append(f"{entry['name'].upper()} = {value} # {where}") return "\n".join(lines) + "\n" -def _python(expr: Expr) -> str: +def _python(expr: Expr, env: dict[str, str | None] | None = None) -> str: # Fortran divides two integers to an integer; ``with_integer_division`` - # spells those quotients ``//`` from the tree's own types. + # spells those quotients ``//`` from the tree's own types and the + # declared types of the constants before it. return render( - with_integer_division(expr), + with_integer_division(expr, env=env), real=lambda text: f"np.float64('{text}')", integer=lambda text: text, name=lambda text: text.upper(), diff --git a/src/recast/transform/numpy/expressions.py b/src/recast/transform/numpy/expressions.py index eb53977..41c57f7 100644 --- a/src/recast/transform/numpy/expressions.py +++ b/src/recast/transform/numpy/expressions.py @@ -59,7 +59,7 @@ DIM_KEYWORD = re.compile(r"dim\s*=\s*", re.I) -BOUND_TOKENS = re.compile(r"[A-Za-z_]\w*\s*%\s*[A-Za-z_]\w*|[A-Za-z_]\w*|\d+|[()+\-*/ ]") +BOUND_TOKENS = re.compile(r"[A-Za-z_]\w*\s*%\s*[A-Za-z_]\w*|[A-Za-z_]\w*|\d+|[()+\-*/, ]") """What a declared bound is allowed to be made of. Bound texts are simple by construction; anything richer refuses the statement that needed the bound.""" @@ -532,25 +532,46 @@ def extent(match: re.Match[str]) -> str: if substituted != text: text = substituted rendered, position = [], 0 + opens_intrinsic = False # the next "(" opens a max/min call + calls: list[bool] = [] # per open parenthesis: a max/min call? for match in BOUND_TOKENS.finditer(text): if match.start() != position: raise NoRule(f"dim expr {text!r}") position = match.end() - token = match.group(0) - if "%" in token: + piece = match.group(0) + if "%" in piece: # ``bounds%begp`` sizing a local: the component of a dummy, # which is an attribute of the same name on this side. - root, component = (t.strip() for t in token.split("%", 1)) + root, component = (t.strip() for t in piece.split("%", 1)) rendered.append(f"{self.names.symbol(root)}.{pysafe(component.lower())}") - elif re.match(r"[A-Za-z_]", token): - rendered.append(self.names.symbol(token)) - elif token.isdigit() and token not in ("0", "1", "2"): - hoisted = self.names.literals.get(token) + elif piece.lower() in ("max", "min") and text[match.end() :].lstrip().startswith("("): + # ``max(2, edsclr_dim)`` sizing a local (CLUBB's windm + # solver): Python spells the two intrinsics the same way, + # and a bound's operands are integers. The only calls a + # bound may carry; a comma is legal inside one of them alone. + rendered.append(piece.lower()) + opens_intrinsic = True + elif re.match(r"[A-Za-z_]", piece): + rendered.append(self.names.symbol(piece)) + elif piece.isdigit() and piece not in ("0", "1", "2"): + hoisted = self.names.literals.get(piece) if hoisted is None: - raise NoRule(f"declared dim literal {token}") + raise NoRule(f"declared dim literal {piece}") rendered.append(hoisted) + elif piece == "(": + calls.append(opens_intrinsic) + opens_intrinsic = False + rendered.append(piece) + elif piece == ")": + if calls: + calls.pop() + rendered.append(piece) + elif piece == ",": + if not (calls and calls[-1]): + raise NoRule(f"dim expr {text!r}") + rendered.append(piece) else: - rendered.append(token) + rendered.append(piece) if position != len(text): raise NoRule(f"dim expr {text!r}") return "".join(rendered) diff --git a/src/recast/transform/numpy/flat.py b/src/recast/transform/numpy/flat.py index 2404762..b82add5 100644 --- a/src/recast/transform/numpy/flat.py +++ b/src/recast/transform/numpy/flat.py @@ -77,6 +77,12 @@ def python_adapter(plans: list[FlatPlan]) -> str: lines.append(f" {', '.join(_py(n) + '_' for n in outs)}, = _out") for obj in plan.objects: if obj.kind == "dummy": + if obj.name in outs and obj.name not in {a["name"] for a in passed}: + # An intent(out) object is one the translation returns + # (sponge_layer_damping's profile, CLUBB): its written + # components are read off the returned object, not the + # record built from the inputs. + lines.append(f" {obj.name} = {_py(obj.name)}_") for comp in obj.components: if comp.written: lines.append(f" {comp.flat} = {obj.name}.{comp.name}") diff --git a/src/recast/transform/numpy/runtime.py b/src/recast/transform/numpy/runtime.py index 6762192..a0a2f44 100644 --- a/src/recast/transform/numpy/runtime.py +++ b/src/recast/transform/numpy/runtime.py @@ -203,6 +203,23 @@ def _f_vdot(a: Any, b: Any) -> Any: return np.dot(a, b) +def _f_vsum(a: Any, axis: Any = None) -> Any: + """Fortran SUM accumulates in element order; np.sum pairs terms and + rounds differently (CLUBB's vertical_integral: 12 ULP).""" + if _LIBM_STRICT: + arr = np.asarray(a) + if axis is None: + s = arr.dtype.type(0) if arr.dtype.kind in "fc" else 0 + for x in np.ravel(arr, order="F"): + s = s + x + return s + out = np.zeros(arr.shape[:axis] + arr.shape[axis + 1 :], dtype=arr.dtype) + for i in range(arr.shape[axis]): + out = out + np.take(arr, i, axis=axis) + return out + return np.sum(a, axis=axis) + + def _fstr_eq(a: str, b: str) -> bool: """Fortran character equality: pad shorter operand with blanks.""" return a.rstrip(" ") == b.rstrip(" ") diff --git a/src/recast/transform/numpy/statements.py b/src/recast/transform/numpy/statements.py index 6e01bf9..77c91bd 100644 --- a/src/recast/transform/numpy/statements.py +++ b/src/recast/transform/numpy/statements.py @@ -139,6 +139,51 @@ def derived_array(type_name: str, extents: list[str], known: dict[str, Any]) -> write using anything else is refused rather than silently list-directed.""" +def _loops_whose_index_is_read_after(subprogram: Any) -> set[int]: + """The DO constructs whose loop variable a later statement reads before + anything redefines it -- where Fortran's completion value (one step past + the end) and Python's (the last value) differ and are observed. + + A later DO over the same variable, or an assignment to it, redefines it + first and closes the question; ``i`` and ``k`` are reused by most loops + of a module, and an ``else`` on every one of them would be noise. The + search loops that read their index (CLUBB's ``lscale_width_vert_avg``) + are what this finds.""" + from recast.fortran.interface import names_in, node_span + + def index_of(statement: Any) -> str | None: + control = walk(statement, f03.Loop_Control) + if control and control[0].children[1] is not None: + return str(control[0].children[1][0]).lower() + return None + + marked: set[int] = set() + for loop in walk(subprogram, (f03.Block_Nonlabel_Do_Construct, f03.Block_Label_Do_Construct)): + do_statement = walk(loop, (f03.Nonlabel_Do_Stmt, f03.Label_Do_Stmt)) + variable = index_of(do_statement[0]) if do_statement else None + _, end_line = node_span(loop) + if variable is None or end_line is None: + continue + for statement in walk(subprogram): + item = getattr(statement, "item", None) + span = getattr(item, "span", None) if item is not None else None + if not span or span[0] <= end_line: + continue + if isinstance(statement, (f03.Nonlabel_Do_Stmt, f03.Label_Do_Stmt)): + if index_of(statement) == variable: + break # redefined by the next loop over it + # Another loop's header may still read it in its bounds + # (``do k_avg = k_avg_lower, k_avg_upper``): checked below. + if isinstance(statement, f03.Assignment_Stmt): + target = statement.children[0] + if isinstance(target, f03.Name) and str(target).lower() == variable: + break # redefined by assignment + if variable in names_in(statement): + marked.add(id(loop)) + break + return marked + + @dataclass class Statements: """Render Fortran statements for one subprogram. @@ -156,6 +201,9 @@ class Statements: externals: dict[str, dict[str, Any]] = field(default_factory=dict) """Procedures with an audited shim in the externals module.""" + stub_procedures: frozenset[str] = frozenset() + """Use-imported from a stubbed module (the frontend's list).""" + call_transforms: dict[str, Any] = field(default_factory=dict) """Callee -> a domain package's answer for it; see ``calls.CallSite``.""" @@ -185,6 +233,13 @@ class Statements: called_names: set[str] = field(default_factory=set) """Names this subprogram's body calls or subscripts. Filled by ``scan``.""" + index_read_after: set[int] = field(default_factory=set) + """DO constructs (by node id) whose index a later statement reads. + Filled by ``scan``. Fortran leaves a completed loop's index one step past + the end; Python's ``for`` leaves the last value, so these loops get an + ``else`` that sets the completion value (an EXIT, a ``break``, skips it, + as Fortran keeps the exit value).""" + exit_labels: dict[int, str] = field(default_factory=dict) """``id(do-construct)`` -> the label that means ``exit`` inside it.""" @@ -225,6 +280,7 @@ def scan(self, subprogram: Any) -> None: """ self.exit_labels = {} self.consumed_labels = set() + self.index_read_after = _loops_whose_index_is_read_after(subprogram) self.assigned_names = { str(a.children[0]).lower() for a in walk(subprogram, f03.Assignment_Stmt) @@ -1400,7 +1456,18 @@ def _do_construct_inner( f"{pad}for {name} in range({low}, " f"({high}) + (1 if ({step}) > 0 else -1), {step}):" ) - return [head, *self._loop_body(node, indent, cycle_name)] + lines = [head, *self._loop_body(node, indent, cycle_name)] + if id(node) in self.index_read_after: + # CLUBB's lscale_width_vert_avg searches with ``do k_avg_upper = + # k, ...; if (...) exit; end do`` and then integrates up to + # k_avg_upper: on completion Fortran's index is the first value + # past the end, m1 + n * m3, and ``for``'s is the last one. + # ``else`` runs exactly when no ``break`` did. + increment = step if step is not None else "1" + trips = f"max(0, (({high}) - ({low}) + ({increment})) // ({increment}))" + lines.append(f"{pad}else:") + lines.append(f"{pad} {name} = ({low}) + {trips} * ({increment})") + return lines def _caught_cycle(self, body: list[str], indent: int, cycle_name: str | None) -> list[str]: """A loop body, wrapped so a CYCLE naming *this* loop reaches its header.""" @@ -1682,6 +1749,15 @@ def _call(self, node: Any, indent: int) -> list[str]: # list says "the engine does not know this intrinsic" rather # than "this tree is missing a library". raise NoRule(f"intrinsic subroutine {name!r} has no rule") + if name in self.stub_procedures: + # A procedure of a stubbed module that no statement stub + # answers (CLUBB's lapack_band_solvex: the LAPACK path the + # run does not take). A raise on the statement keeps the + # branch around it -- deferring the block dropped the + # ``if ( method == lapack )`` with it, and the candidate + # raised on the path the run *does* take. + reason = f"{name}: procedure of a stubbed module, not ported" + return [f"{pad}raise NotImplementedError({reason!r})"] raise NoRule(f"call to external subroutine {name!r}") # Bind actuals to formals BY NAME for keyword arguments: Fortran @@ -1723,7 +1799,7 @@ def _call(self, node: Any, indent: int) -> list[str]: outputs.append("_") # the return tuple has fixed length continue if self.is_optional_output(formal): - inputs.append(f"want_{formal['name']}=True") + inputs.append(f"want_{formal['name']}={self._presence(actual)}") passes = formal["intent"] in ("IN", "INOUT", "UNKNOWN") or bool( formal.get("buffer") and self.buffer_out_arrays ) @@ -1872,6 +1948,21 @@ def _external_call(self, name: str, external: dict[str, Any], node: Any, pad: st return [f"{pad}{', '.join(outputs)} = {call}"] return [f"{pad}{call}"] + def _presence(self, actual: Any) -> str: + """``present()`` of what is passed to an optional OUT: true for a + value of the caller's own; the caller's own presence when the + actual is one of its optional dummies handed on (CLUBB's + xm_wpxp_solve passes ``rcond = rcond`` to band_solve, and takes the + LAPACK diagnostic path only when *its* caller asked for rcond).""" + if isinstance(actual, f03.Name): + name = str(actual).lower() + for declared in self.semantics.subprogram["args"]: + if declared["name"].lower() == name and declared["optional"]: + if self.is_optional_output(declared): + return f"want_{name}" + return f"({self.names.symbol(name)} is not None)" + return "True" + @staticmethod def is_optional_output(formal: dict[str, Any]) -> bool: """Optional OUT, the ``want_`` sentinel convention: a trailing diff --git a/src/recast/transform/numpy/subprograms.py b/src/recast/transform/numpy/subprograms.py index cf814e8..dd2a55f 100644 --- a/src/recast/transform/numpy/subprograms.py +++ b/src/recast/transform/numpy/subprograms.py @@ -84,15 +84,15 @@ def _token_pass_guessed(text: str, spelled: str) -> bool: ``_is_expression`` only asks whether Python can *parse* it. Three shapes parse and are still wrong: a call or array reference of an uppercased name; a ``//`` concatenation, which Python reads as floor division; and an - array constructor that was only part of the text (``reshape((/.../), - ...)``), where the search silently dropped everything around it. + array constructor the rendering dropped. """ if UPPERCASED_CALL.search(spelled): return True if "//" in text: return True - constructed = ARRAY_CONSTRUCTOR.search(text) - return constructed is not None and constructed.span() != (0, len(text)) + # A constructor inside a larger text is rendered where it stands now; + # the guess is a rendering that lost it. + return ARRAY_CONSTRUCTOR.search(text) is not None and "np.array(" not in spelled def _is_expression(text: str) -> str | bool: @@ -149,6 +149,9 @@ class Subprograms: companion_globals: dict[str, str] = field(default_factory=dict) externals: dict[str, dict[str, Any]] = field(default_factory=dict) remotes: dict[str, Remote] = field(default_factory=dict) + stub_procedures: frozenset[str] = frozenset() + """Procedures use-imported from a stubbed module; a call to one that no + statement stub answers is a raise, not a deferral of its block.""" function_stubs: dict[str, str] = field(default_factory=dict) statement_stubs: dict[str, str] = field(default_factory=dict) intrinsics: dict[str, Any] = field(default_factory=dict) @@ -333,6 +336,7 @@ def floors(self, name: str) -> Statements: names, expressions, externals=self.externals, + stub_procedures=self.stub_procedures, stubs=dict(self.statement_stubs), call_transforms=dict(self.call_transforms), poison_undefined=self.poison_undefined, @@ -1007,8 +1011,8 @@ def _token_parameter_value(text: str, local_parameters: frozenset[str]) -> str: # rather than anything this file would notice. exponent = text.replace(" ", "").split("_")[0] return "np.float64('" + exponent.replace("d", "e").replace("D", "e") + "')" - constructed = ARRAY_CONSTRUCTOR.search(text) - if constructed: + + def array_of(constructed: re.Match[str]) -> str: items = [strip_kind(item.strip()) for item in constructed.group(1).split(",")] if all(re.fullmatch(r"'[^']*'", item) for item in items): return f"np.array([{', '.join(items)}])" @@ -1026,6 +1030,25 @@ def case_of(match: re.Match[str]) -> str: token = match.group() return pysafe(token.lower()) if token.lower() in local_parameters else token.upper() + constructed = ARRAY_CONSTRUCTOR.search(text) + if constructed and constructed.span() == (0, len(text)): + return array_of(constructed) + if constructed: + # A constant expression over a constructor (CLUBB's saturation + # and pdf_closure: ``100._core_rknd * (/ 6.09868993_core_rknd, + # ... /)``): each constructor as an array where it stands, the + # rest through the token pass; NumPy is elementwise as Fortran + # is. A call around one (``reshape((/.../), ...)``) still goes + # to the parse, by the uppercased-call rule. + pieces: list[str] = [] + position = 0 + for found in ARRAY_CONSTRUCTOR.finditer(text): + pieces.append(IDENTIFIER.sub(case_of, strip_kind(text[position : found.start()]))) + pieces.append(array_of(found)) + position = found.end() + pieces.append(IDENTIFIER.sub(case_of, strip_kind(text[position:]))) + return "".join(pieces) + return IDENTIFIER.sub(case_of, strip_kind(text)) @staticmethod diff --git a/src/recast/transform/numpy/translate.py b/src/recast/transform/numpy/translate.py index 3891414..1842897 100644 --- a/src/recast/transform/numpy/translate.py +++ b/src/recast/transform/numpy/translate.py @@ -280,6 +280,7 @@ def apply(self, unit: Unit, facts: Facts, config: dict[str, Any]) -> Candidate: use_parameters=use_parameters, companion_globals=companion_globals, externals=facts.provenance.get("externals", {}), + stub_procedures=frozenset(facts.interface.get("stub_procedures") or ()), remotes=remotes, function_stubs=config.get("function_stubs", {}), statement_stubs=config.get("statement_stubs", {}), @@ -422,6 +423,9 @@ def _rwset_protocol( # The siblings' procedures too: `_wv.wv_sat_svp_water(t)` is a # call, and without these the alias rule would read it as data. | {remote.name for remote in assembler.remotes.values()} + # ... and the stubbed modules' (the frontend's list): a call to + # a stand-in function is a call, not a read of its name. + | {pysafe(name) for name in facts.interface.get("stub_procedures") or ()} ), "aliases": sorted( {remote.alias for remote in assembler.remotes.values()} diff --git a/src/recast/transform/numpy/vocabulary.py b/src/recast/transform/numpy/vocabulary.py index 708713e..cdbc26a 100644 --- a/src/recast/transform/numpy/vocabulary.py +++ b/src/recast/transform/numpy/vocabulary.py @@ -166,7 +166,7 @@ "minval": "np.min", "product": "np.prod", "size": "np.size", - "sum": "np.sum", + "sum": "_f_vsum", # With unit lower bounds -- which every translated array has -- the upper # bound and the extent are the same number. "ubound": "np.size", diff --git a/src/recast/verify/bitexact.py b/src/recast/verify/bitexact.py index 1d7df19..7b03cd1 100644 --- a/src/recast/verify/bitexact.py +++ b/src/recast/verify/bitexact.py @@ -164,6 +164,16 @@ def call(*values: Any) -> Any: return callback +def _returned(translated_out: Any) -> list[Any]: + """The values a candidate call handed back, as a list: a tuple's items, + one bare value, or none at all -- a subroutine with no OUT argument + returns ``None`` (CLUBB's finalize_tau_sponge_damp_api deallocates and + returns), and that is zero values, not one.""" + if translated_out is None: + return [] + return list(translated_out) if isinstance(translated_out, tuple) else [translated_out] + + def _extent(dim: dict[str, Any], dims: dict[str, int]) -> int: """An axis's extent: ``ub - lb + 1`` when a lower bound is declared (CLUBB's ``lhs(-2:2, ...)`` has five rows, not two), ``ub`` otherwise.""" @@ -405,6 +415,12 @@ def generable(name: str) -> bool: return False return True + # One the operator declared ungated is not compared: the declaration + # says the reference cannot be held -- on generated inputs (CLUBB's + # rcm_sat_adj iterates and error-stops on them) or on a recording + # (its sponge initializer leaves the levels below the layer + # undefined on both sides). The reason is reported beside the verdict. + declared_ungated = set(config.get("ungated") or {}) if recorded: # A recording names what it is a recording of, so the set to # compare is the set that was captured -- not every subprogram the @@ -416,13 +432,20 @@ def generable(name: str) -> bool: by_subprogram.setdefault(str(sample.get("subprogram", "")), []).append(sample) offered = sorted(by_subprogram) wanted = config.get("subprograms") or [ - name for name in offered if name in table and judged(name) + name + for name in offered + if name in table and judged(name) and name not in declared_ungated ] skipped = sorted(set(offered) - set(wanted)) else: by_subprogram = {} wanted = config.get("subprograms") or [ - name for name in wrappers if name in table and judged(name) and generable(name) + name + for name in wrappers + if name in table + and judged(name) + and generable(name) + and name not in declared_ungated ] skipped = sorted(set(wrappers) - set(wanted)) @@ -1135,6 +1158,10 @@ def _dominance( if dominant_at is None: return None magnitude = np.abs(reference) + if magnitude.size == 0: + # A zero-extent output (CLUBB's scalar tracers under + # sclr_dim = 0): nothing to weigh, and no maximum to take. + return [] if axis in ("all", None) or magnitude.ndim <= 1: scale = magnitude.max() else: @@ -1334,7 +1361,7 @@ def _paired_outputs( # therefore by exact name on both sides. Every required output was # preflighted before the candidate call; keep the same check here # as a fail-closed local invariant for direct callers. - mine = list(translated_out) if isinstance(translated_out, tuple) else [translated_out] + mine = _returned(translated_out) names = ( [sub.get("result") or "result"] if sub["kind"] == "function" @@ -1390,7 +1417,7 @@ def _paired_outputs( for a in outs_required ] - ours = list(translated_out) if isinstance(translated_out, tuple) else [translated_out] + ours = _returned(translated_out) if len(ours) != len(outs_all): return ( f"candidate returned {len(ours)} value(s) for " diff --git a/src/recast/verify/rwset.py b/src/recast/verify/rwset.py index 1b29a8e..8867389 100644 --- a/src/recast/verify/rwset.py +++ b/src/recast/verify/rwset.py @@ -362,6 +362,11 @@ def walk_stmt(node: ast.stmt) -> None: STUB_LINE = re.compile(r"^\s*(?:pass\s*#.*\(infra stub\)|#.*)$") +CONTROL_LINE = re.compile(r"^\s*(?:(?:if|elif|for|while)\b[^#]*:|else\s*:|pass)\s*(?:#.*)?$") +"""A line of control flow with nothing of its own: a condition, a loop +header, an ``else``, a ``pass``. Around stub markers it is still a stub.""" + + def stubbed_blocks(candidate: Candidate) -> dict[str, str]: """``"sub/Bnnn" -> reason`` for every block emitted as stub markers only. @@ -390,7 +395,10 @@ def stubbed_blocks(candidate: Candidate) -> dict[str, str]: continue body = [ln for ln in lines[span[0] - 1 : span[1]] if ln.strip()] stubs = [ln for ln in body if "(infra stub)" in ln] - if stubs and all(STUB_LINE.match(ln) for ln in body): + if stubs and all(STUB_LINE.match(ln) or CONTROL_LINE.match(ln) for ln in body): + # ``if (stats%l_sample) then / call stats_update(...) / end if`` + # (CLUBB) is a stub under a condition: the condition is read on + # both sides, and the only disagreement is the stub's actuals. calls = sorted({ln.split("#", 1)[1].split("(")[0].strip() for ln in stubs}) waived[f"{block['subprogram']}/{block['block']}"] = "framework stub: " + ", ".join( calls diff --git a/tests/test_dump_replay.py b/tests/test_dump_replay.py index 75df905..b5a6c87 100644 --- a/tests/test_dump_replay.py +++ b/tests/test_dump_replay.py @@ -361,3 +361,23 @@ def test_the_shipped_example_replays_bit_exact(tmp_path: Path) -> None: assert set(compared) == {"settle", "column_mass"} assert all(outcome["points"] > 0 for outcome in compared.values()) assert verdict.metrics["uncovered"] == [] + + +def test_a_logical_header_scalar_is_an_input() -> None: + """CLUBB's ``l_implemented``: the recorder writes a logical ``T``/``F``. + The parser took only numbers, so the replay had no value for it.""" + inputs, _ = parse_dump("# PROBE m.s: call=1\n# l_on = T\n# l_off = F\n# INPUT: x(1)\n1.0\n") + assert inputs["l_on"] is not None and bool(inputs["l_on"]) is True + assert bool(inputs["l_off"]) is False + assert inputs["l_on"].dtype == np.bool_ + + +def test_a_zero_extent_array_is_a_value() -> None: + """A component the run never allocated (CLUBB's scalar tracers under + ``sclr_dim = 0``) is written ``name(1,3,0)`` with nothing under it. + Dropped, the replay said the record carried no value for it.""" + text = "# PROBE m.s: call=1\n# INPUT: s(1,3,0)\n# OUTPUT: t(0)\n# OUTPUT: y(1)\n2.0\n" + inputs, outputs = parse_dump(text) + assert inputs["s"].shape == (1, 3, 0) + assert outputs["t"].shape == (0,) + assert outputs["y"].tolist() == [2.0] diff --git a/tests/test_f2py_oracle.py b/tests/test_f2py_oracle.py index 13892ef..96e2ecf 100644 --- a/tests/test_f2py_oracle.py +++ b/tests/test_f2py_oracle.py @@ -916,6 +916,179 @@ def two_outputs(): assert "partial output evidence is not a pass" in verdict.detail +def test_a_recorded_subroutine_with_no_output_returns_nothing(tmp_path: Path) -> None: + """CLUBB's finalize_tau_sponge_damp_api deallocates a component and + returns: no OUT argument, so its adapter returns ``None``. The gate + counted that as one value against zero out-intent arguments.""" + emitted = b"""\ +_SIGNATURES = { + "release": { + "kind": "subroutine", + "args": [ + {"name": "n", "dtype": "int32", "intent": "IN", "optional": False}, + ], + "result": None, + "result_dtype": None, + } +} + +def release(n): + return None +""" + candidate = Candidate( + unit="fortran:no_output", + transform="translate.numpy", + files={Path("no_output_numpy.py"): emitted}, + ) + ref = OracleRef( + unit=candidate.unit, + oracle="dump-replay", + key="k", + handle={ + "module": None, + "input_source": "recorded", + "return_convention": "recorded", + "samples": [ + { + "subprogram": "release", + "source": "release.txt", + "inputs": {"n": 3}, + "outputs": {}, + } + ], + }, + ) + verdict = BitexactVerifier().verify( + Unit(uid=candidate.unit, kind="module"), + candidate, + ref, + tmp_path / "work", + LocalExecutor(), + {}, + ) + # Nothing to compare is still not a pass -- but for the right reason. + assert "returned 1 value(s)" not in verdict.detail + assert "zero numerical points" in verdict.detail + + +def test_a_declared_ungated_subprogram_is_not_compared_on_a_recording(tmp_path: Path) -> None: + """CLUBB's sponge initializer leaves the levels below the layer undefined + on both sides; the operator's declaration says so, with the reason, and + the replay reported it -- then compared the heap against np.empty anyway.""" + emitted = b"""\ +_SIGNATURES = { + "fill": { + "kind": "subroutine", + "args": [ + {"name": "n", "dtype": "int32", "intent": "IN", "optional": False}, + {"name": "y", "dtype": "int32", "intent": "OUT", "optional": False}, + ], + "result": None, + "result_dtype": None, + } +} + +def fill(n): + return 2 +""" + candidate = Candidate( + unit="fortran:undefined_tail", + transform="translate.numpy", + files={Path("undefined_tail_numpy.py"): emitted}, + ) + ref = OracleRef( + unit=candidate.unit, + oracle="dump-replay", + key="k", + handle={ + "module": None, + "input_source": "recorded", + "return_convention": "recorded", + "samples": [ + { + "subprogram": "fill", + "source": "fill.txt", + "inputs": {"n": 3}, + "outputs": {"y": 1}, + } + ], + }, + ) + verdict = BitexactVerifier().verify( + Unit(uid=candidate.unit, kind="module"), + candidate, + ref, + tmp_path / "work", + LocalExecutor(), + {"ungated": {"fill": "the tail is undefined on both sides"}}, + ) + assert "differ" not in verdict.detail + assert "fill (the tail is undefined on both sides)" in verdict.detail + + +def test_a_zero_extent_output_has_no_dominant_value_to_weigh(tmp_path: Path) -> None: + """CLUBB's scalar tracers under ``sclr_dim = 0``: an output of shape + ``(1, 88, 0)`` on the recording. The tolerance gate's dominant-value + mask took the maximum of an empty array and the whole unit's verdict + was a plugin exception.""" + import numpy as np + + from recast.verify.tolerance import ToleranceVerifier + + emitted = b"""\ +import numpy as np +_SIGNATURES = { + "tracers": { + "kind": "subroutine", + "args": [ + {"name": "n", "dtype": "int32", "intent": "IN", "optional": False}, + {"name": "y", "dtype": "float64", "intent": "OUT", "optional": False, + "dims": [{"lb": "1", "ub": "n"}, {"lb": "1", "ub": "0"}]}, + {"name": "z", "dtype": "float64", "intent": "OUT", "optional": False}, + ], + "result": None, + "result_dtype": None, + } +} + +def tracers(n): + return np.zeros((n, 0)), 2.0 +""" + candidate = Candidate( + unit="fortran:tracers", + transform="translate.numpy", + files={Path("tracers_numpy.py"): emitted}, + ) + ref = OracleRef( + unit=candidate.unit, + oracle="dump-replay", + key="k", + handle={ + "module": None, + "input_source": "recorded", + "return_convention": "recorded", + "samples": [ + { + "subprogram": "tracers", + "source": "tracers.txt", + "inputs": {"n": 3}, + "outputs": {"y": np.zeros((3, 0)), "z": 2.0}, + } + ], + }, + ) + verdict = ToleranceVerifier().verify( + Unit(uid=candidate.unit, kind="module"), + candidate, + ref, + tmp_path / "work", + LocalExecutor(), + {"module_suffix": "_numpy.py", "dominant_axis": "all", "rel_scale": "array"}, + ) + assert "exception" not in verdict.detail and "zero-size" not in verdict.detail + assert verdict.confidence is not Confidence.FAILED, verdict.detail + + # --- the whole spine, against a real compiler -------------------------------- SOURCE = """\ diff --git a/tests/test_flatten.py b/tests/test_flatten.py index 35c75c0..bac823a 100644 --- a/tests/test_flatten.py +++ b/tests/test_flatten.py @@ -484,3 +484,336 @@ def test_the_adapter_declares_a_lower_bound_and_calls_through_a_generic() -> Non assert "use solve_mod, only: solve\n" in text # both specifics, one generic assert "call solve(n=n, x=x)" in text assert "solve_one(" not in text.replace("subroutine solve_one_flat(", "") + + +# --- a CLUBB-shaped object: many components in one ALLOCATE, sized by itself + +GRID = """\ +module grid_class + implicit none + private + public :: grid, setup_grid + integer, parameter :: t_above = 1, t_below = 2 + type grid + integer :: nzm, nzt + real(8), allocatable, dimension(:,:) :: zm, zt + real(8), allocatable, dimension(:,:,:) :: weights_zt2zm + real(8) :: grid_dir + end type grid +contains + subroutine setup_grid( ngrdcol, nzmax, gr ) + integer, intent(in) :: ngrdcol, nzmax + type(grid), intent(inout) :: gr + integer :: ierr + gr%nzm = nzmax + gr%nzt = nzmax - 1 + allocate( gr%zm(ngrdcol,gr%nzm), gr%zt(ngrdcol,gr%nzt), & ! two at once + gr%weights_zt2zm(ngrdcol,gr%nzm,t_above:t_below), & + stat=ierr ) + gr%grid_dir = 1.0d0 + end subroutine setup_grid +end module grid_class +""" + +COLUMN = """\ +module column_mod + use grid_class, only: grid + implicit none + private + public :: ddz +contains + subroutine ddz( nzm, ngrdcol, gr, x, dxdz ) + integer, intent(in) :: nzm, ngrdcol + type(grid), intent(in) :: gr + real(8), intent(in), dimension(ngrdcol, nzm) :: x + real(8), intent(out), dimension(ngrdcol, nzm) :: dxdz + integer :: i, k + do k = 1, nzm + do i = 1, ngrdcol + dxdz(i,k) = gr%grid_dir * x(i,k) * gr%zm(i,k) * gr%weights_zt2zm(i,k,1) + end do + end do + end subroutine ddz +end module column_mod +""" + +CLUBB_CONVENTIONS = FlatConventions(patch_count="ngrdcol", bounds_pattern=r"^ngrdcol$") + + +def test_an_object_allocated_many_at_once_and_sized_by_itself(tmp_path: Path) -> None: + """CLUBB's grid: one ALLOCATE over every component, the object named by + the setup routine's dummy rather than ``this``, an axis sized by another + component of the same object (``gr%nzm``) and one by the module's + private parameters (``t_above:t_below``). The plan carries ``nzm`` as an + input the body never reads, spells the extents by it, and declares the + driver's extent once, because ``ngrdcol`` is a dummy already.""" + (tmp_path / "grid_class.f90").write_text(GRID) + (tmp_path / "column_mod.f90").write_text(COLUMN) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:column_mod") + facts = frontend.analyze(unit, tmp_path) + (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS) + assert plan.usable, plan.unsupported + (gr,) = plan.objects + by_name = {c.name: c for c in gr.components} + assert set(by_name) == {"grid_dir", "zm", "weights_zt2zm", "nzm"} + assert by_name["nzm"].written is False + assert by_name["zm"].extents == ["ngrdcol", "gr__nzm"] + assert by_name["weights_zt2zm"].extents == ["ngrdcol", "gr__nzm", "2"] + assert by_name["weights_zt2zm"].bounds[2] == ("1", "2") + names = [a["name"] for a in plan.flat_args] + assert names.count("ngrdcol") == 1 + assert names.index("gr__nzm") < names.index("gr__zm") + text = fortran_adapter("column_mod", [plan], []) + assert "real(8), intent(in) :: gr__zm(ngrdcol, gr__nzm)" in text + assert "allocate(gr%zm(1:ngrdcol, 1:gr__nzm))" in text + + +COEFS = """\ +module coefs_mod + implicit none + private + public :: coefs_type, init_coefs + type coefs_type + real(8), allocatable, dimension(:,:) :: coef + end type coefs_type +contains + subroutine init_coefs( ngrdcol, nz, c ) + integer, intent(in) :: ngrdcol, nz + type(coefs_type), intent(out) :: c + allocate( c%coef(1:ngrdcol,1:nz) ) + c%coef = 0.0d0 + end subroutine init_coefs +end module coefs_mod +""" + +USES_COEFS = """\ +module solver_mod + use coefs_mod, only: coefs_type + implicit none + private + public :: apply +contains + subroutine apply( nzt, ngrdcol, c, x ) + integer, intent(in) :: nzt, ngrdcol + type(coefs_type), intent(in) :: c + real(8), intent(inout), dimension(ngrdcol, nzt) :: x + x = x * c%coef(:, 1:nzt) + end subroutine apply +end module solver_mod +""" + + +def test_an_extent_the_plan_cannot_spell_becomes_an_argument(tmp_path: Path) -> None: + """``coef`` is allocated ``(1:ngrdcol, 1:nz)`` by the initializer's dummy + ``nz``; the planned subprogram takes ``nzt``, not ``nz``. The extent is + the run's own: the plan makes it an integer argument, the recorder writes + it from ``size()``, and the adapters declare the component by it.""" + from recast.oracle.record import recorder_module + + (tmp_path / "coefs_mod.f90").write_text(COEFS) + (tmp_path / "solver_mod.f90").write_text(USES_COEFS) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod") + facts = frontend.analyze(unit, tmp_path) + (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS) + assert plan.usable, plan.unsupported + assert plan.extent_args == {"c__coef_n2": ["c", "coef", 2]} + (coef,) = plan.objects[0].components + assert coef.extents == ["ngrdcol", "c__coef_n2"] + names = [a["name"] for a in plan.flat_args] + assert names.index("c__coef_n2") < names.index("c__coef") + adapter = fortran_adapter("solver_mod", [plan], []) + assert "integer, intent(in) :: c__coef_n2" in adapter + assert "real(8), intent(in) :: c__coef(ngrdcol, c__coef_n2)" in adapter + recorder = recorder_module("solver_mod", [plan]) + assert "'# c__coef_n2 = ', merge(size(c%coef, 2), 0, allocated(c%coef))" in recorder + # ngrdcol is a dummy of the probe already: declared once, not assigned. + probe = recorder[recorder.index("subroutine rec_apply(") :] + assert probe.count("integer :: ngrdcol") == 0 + assert "ngrdcol = " not in probe.split("phase == 0")[0] + again = FlatPlan.from_dict(plan.to_dict()) + assert again.extent_args == plan.extent_args + + +def test_a_call_continued_with_trailing_comments_is_probed_whole(tmp_path: Path) -> None: + """CLUBB continues its calls as ``a, b, & ! In`` on every line: the + comment after the ampersand left a blank the joiner did not strip, and + the probe carried ``&`` into its argument list.""" + from recast.oracle.record import probe_tree + + (tmp_path / "coefs_mod.f90").write_text(COEFS) + (tmp_path / "solver_mod.f90").write_text(USES_COEFS) + (tmp_path / "driver_mod.f90").write_text( + """\ +module driver_mod + use coefs_mod, only: coefs_type + use solver_mod, only: apply + implicit none +contains + subroutine step( nzt, ngrdcol, c, x ) + integer, intent(in) :: nzt, ngrdcol + type(coefs_type), intent(in) :: c + real(8), intent(inout), dimension(ngrdcol, nzt) :: x + call apply( nzt, ngrdcol, & ! In + c, & ! In + x ) ! In/out + end subroutine step +end module driver_mod +""" + ) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod") + facts = frontend.analyze(unit, tmp_path) + plans = plans_for(facts, tmp_path, CLUBB_CONVENTIONS) + sites = probe_tree(tmp_path, tmp_path / "probed", {"solver_mod": plans}) + assert sites == {"apply": 1} + probed = (tmp_path / "probed" / "driver_mod.f90").read_text() + assert "call rec_apply(0, nzt, ngrdcol, c, x)" in probed + assert "&" not in probed.split("call rec_apply(0")[1].split("\n")[0] + + +def test_the_recorder_guards_a_component_the_run_may_not_allocate(tmp_path: Path) -> None: + """CLUBB allocates its scalar-tracer coefficients only when sclr_dim > 0; + reshape of an unallocated component faulted the recording run.""" + from recast.oracle.record import recorder_module + + (tmp_path / "coefs_mod.f90").write_text(COEFS) + (tmp_path / "solver_mod.f90").write_text(USES_COEFS) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod") + facts = frontend.analyze(unit, tmp_path) + (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS) + recorder = recorder_module("solver_mod", [plan]) + assert "if (allocated(c%coef)) then" in recorder + assert "'# INPUT: c__coef(0,0)'" in recorder + assert "merge(size(c%coef, 2), 0, allocated(c%coef))" in recorder + + +def test_the_recorder_writes_the_plain_out_dummies_too(tmp_path: Path) -> None: + """CLUBB's advance_* hand back ``wp2``, ``wp3``... as INOUT dummies beside + what they write into their objects. Recorded only through the objects, + the replay found no value for the required outputs (mono_flux_limiter's + ``low_lev_effect``).""" + from recast.oracle.record import recorder_module + + (tmp_path / "coefs_mod.f90").write_text(COEFS) + (tmp_path / "solver_mod.f90").write_text(USES_COEFS) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod") + facts = frontend.analyze(unit, tmp_path) + (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS) + recorder = recorder_module("solver_mod", [plan]) + assert "call rec_r1(u_apply, 'INPUT', 'x', trim(dims), reshape(x, (/size(x)/)))" in recorder + assert "call rec_r1(u_apply, 'OUTPUT', 'x', trim(dims), reshape(x, (/size(x)/)))" in recorder + + +SPONGE_STATE = """\ +module sponge_mod + implicit none + private + public :: profile_type, damp, sponge_profile, init_profile + type profile_type + real(8), allocatable :: tau(:) + integer :: n_sponge = 0 + end type profile_type + type(profile_type), public :: sponge_profile +contains + subroutine init_profile( nz, prof ) + integer, intent(in) :: nz + type(profile_type), intent(inout) :: prof + allocate( prof%tau(1:nz) ) + prof%tau = 1.0d0 + prof%n_sponge = nz + end subroutine init_profile + function damp( nzt, x, prof ) result( damped ) + integer, intent(in) :: nzt + real(8), intent(in) :: x(nzt) + type(profile_type), intent(in) :: prof + real(8) :: damped(nzt) + if ( allocated( prof%tau ) ) then + damped = x * prof%tau(1:nzt) + else + damped = x + end if + end function damp +end module sponge_mod +""" + +HANDS_STATE_TO_FUNCTION = """\ +module advance_mod + use coefs_mod, only: coefs_type + use sponge_mod, only: damp, sponge_profile + implicit none + private + public :: advance +contains + subroutine advance( nzt, ngrdcol, c, x ) + integer, intent(in) :: nzt, ngrdcol + type(coefs_type), intent(in) :: c + real(8), intent(inout), dimension(ngrdcol, nzt) :: x + integer :: i + do i = 1, ngrdcol + x(i, :) = damp( nzt, x(i, :), sponge_profile ) * c%coef(i, 1:nzt) + end do + end subroutine advance +end module advance_mod +""" + + +def test_module_state_handed_whole_to_a_function_is_carried(tmp_path: Path) -> None: + """CLUBB's advance_xm_wpxp passes sponge_layer_damping's profile object + to its sponge_damp_xm function, which reads ``profile%tau_sponge_damp``. + The walk followed only the caller's own dummies into callees, so the + profile was left to the module and the replay found it unallocated.""" + (tmp_path / "coefs_mod.f90").write_text(COEFS) + (tmp_path / "sponge_mod.f90").write_text(SPONGE_STATE) + (tmp_path / "advance_mod.f90").write_text(HANDS_STATE_TO_FUNCTION) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:advance_mod") + facts = frontend.analyze(unit, tmp_path) + (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS) + assert plan.usable, plan.unsupported + by_name = {o.name: o for o in plan.objects} + assert by_name["sponge_profile"].kind == "state" + assert by_name["sponge_profile"].module == "sponge_mod" + assert [c.name for c in by_name["sponge_profile"].components] == ["tau"] + assert "sponge_mod%sponge_profile" not in plan.left_to_module + + +def test_a_probe_spans_a_blank_line_inside_a_continued_call(tmp_path: Path) -> None: + """cpp leaves blank lines where an ``#ifdef`` stood inside CLUBB's + advance_clubb_core call; the probe took the blank line for the end of + the statement and found no call site to bracket.""" + from recast.oracle.record import probe_tree + + (tmp_path / "coefs_mod.f90").write_text(COEFS) + (tmp_path / "solver_mod.f90").write_text(USES_COEFS) + (tmp_path / "caller_mod.f90").write_text( + "module caller_mod\n" + " use coefs_mod, only: coefs_type\n" + " use solver_mod, only: apply\n" + " implicit none\n" + "contains\n" + " subroutine run( nzt, ngrdcol, c, x )\n" + " integer, intent(in) :: nzt, ngrdcol\n" + " type(coefs_type), intent(in) :: c\n" + " real(8), intent(inout), dimension(ngrdcol, nzt) :: x\n" + " call apply( nzt, ngrdcol, & ! in\n" + "\n" + " ! the object\n" + " c, &\n" + "\n" + " x )\n" + " end subroutine run\n" + "end module caller_mod\n" + ) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod") + facts = frontend.analyze(unit, tmp_path) + (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS) + sites = probe_tree(tmp_path, tmp_path / "probed", {"solver_mod": [plan]}) + assert sites == {"apply": 1} + probed = (tmp_path / "probed" / "caller_mod.f90").read_text() + assert "call rec_apply(0, nzt, ngrdcol, c, x)" in probed + assert "call rec_apply(1, nzt, ngrdcol, c, x)" in probed diff --git a/tests/test_fortran_analysis.py b/tests/test_fortran_analysis.py index 24e63e0..e63ebad 100644 --- a/tests/test_fortran_analysis.py +++ b/tests/test_fortran_analysis.py @@ -760,15 +760,16 @@ def test_a_call_splits_its_arguments_by_declared_intent(tmp_path: Path) -> None: """ -def test_a_component_name_is_read_on_the_out_argument_path_only(tmp_path: Path) -> None: +def test_a_component_name_is_not_a_read_on_the_out_argument_path(tmp_path: Path) -> None: """``b % q`` writes ``b``. On an assignment, ``q`` is an attribute and not a symbol; passed to an intent(out) dummy, the pipeline this came from - counts it as a read as well. + counted it as a read as well. - The two disagree, and the disagreement is preserved. Resolving it would be - a change to answers a bit-exact gate has been run against, and the two - sites in CAM where it shows are both in modules with no translation to - check the tidier answer against. + The disagreement was preserved until a translation showed the tidier + answer: CLUBB's pdf_closure passes ``pdf_params%chi_1`` and six more + components as OUT actuals, the candidate spells attributes and reads no + variable of those names, and the gate scored six blocks as disagreeing + over reads of variables the scope does not have. """ from recast.fortran import rwset @@ -781,10 +782,8 @@ def test_a_component_name_is_read_on_the_out_argument_path_only(tmp_path: Path) ) blocks = {b["id"]: b for b in rwset.block_rwsets(node, rwset.scope_for(record, "drive"))} # ``slot(:)`` is a caller-buffer OUT (#36), so ``b`` is read as well as - # written on the call (#38); ``q`` is the pipeline's component read. - assert blocks["B001"] == {"id": "B001", "reads": ["b", "n", "q"], "writes": ["b"]}, ( - "out-argument" - ) + # written on the call (#38); ``q`` is an attribute on both paths. + assert blocks["B001"] == {"id": "B001", "reads": ["b", "n"], "writes": ["b"]}, "out-argument" assert blocks["B002"] == {"id": "B002", "reads": ["n"], "writes": ["b"]}, "assignment" @@ -1012,7 +1011,8 @@ def test_companion_externals_derive_from_the_siblings_record(tmp_path: Path) -> """ record = interface.extract(_write(tmp_path, "sib.f90", sibling), kind_assumptions=KINDS) table = interface.companion_externals(record) - assert table["qsat_water"] == {"kind": "subroutine", "out_positions": [2, 3]} + assert table["qsat_water"]["kind"] == "subroutine" + assert table["qsat_water"]["out_positions"] == [2, 3] def _sub_node(tmp_path: Path, name: str): @@ -1881,6 +1881,316 @@ def test_the_specifics_of_a_public_generic_are_public_through_it(tmp_path: Path) assert record["generics"] == {"solve": ["solve_one", "solve_many"]} +def test_a_siblings_generic_is_a_procedure_to_the_read_write_scope(tmp_path: Path) -> None: + """A call into a companion spells the generic (CLUBB's ``zt2zm_api``); + without an entry for it the scope counted the name as a read of data.""" + record = interface.extract(_write(tmp_path, "solve.f90", PUBLIC_GENERIC)) + table = interface.companion_externals(record) + assert table["solve"]["kind"] == "subroutine" + assert table["solve"]["out_positions"] == [1, 2] # the union, for a call of no known arity + # ... and each specific with its arity, for the scope to pick by the actuals. + assert [(x["name"], x["args"], x["out_positions"]) for x in table["solve"]["specifics"]] == [ + ("solve_one", 2, [1]), + ("solve_many", 3, [2]), + ] + assert table["solve_one"]["out_positions"] == [1] + + +def test_a_quotient_of_real_parameters_is_a_real_quotient(tmp_path: Path) -> None: + """CLUBB's ``ep = Rd / Rv``: no literal in sight, two real parameters. + The old rule -- no real literal means integer arithmetic throughout -- + folded it to zero. The declared types decide.""" + import numpy as np + + from recast.transform.numpy.constants import use_constants_module + + src = """\ +module gas_mod + implicit none + integer, parameter :: core_rknd = 8 + real( kind = core_rknd ), parameter :: rd = 287.04_core_rknd, rv = 461.5_core_rknd + real( kind = core_rknd ), parameter :: ep = rd / rv + real( kind = core_rknd ), parameter :: ep2 = 1.0_core_rknd / ep + integer, parameter :: runge_kutta_type = 45 + integer, parameter :: nrk = runge_kutta_type / 10 +end module gas_mod +""" + _write(tmp_path, "gas.f90", src) + resolved = use.resolve(["ep2", "nrk"], [tmp_path / "gas.f90"]) + assert {r["name"]: r["dtype"] for r in resolved}["ep"] == "real" + scope: dict[str, object] = {} + exec(use_constants_module(resolved, "gas_mod"), scope) # generated text under test + assert scope["EP"] == np.float64(287.04) / np.float64(461.5) + assert scope["EP2"] == np.float64(1.0) / scope["EP"] + assert scope["NRK"] == 4 + + +STUBBED_CALLER = """\ +module budget_mod + use stats_mod, only: stats_type, stats_update + implicit none + private + public :: tend +contains + subroutine tend( n, x, stats ) + integer, intent(in) :: n + real, dimension(n), intent(inout) :: x + type(stats_type), intent(inout) :: stats + real, dimension(n) :: stats_tmp + x = 2.0 * x + if ( stats%l_sample ) then + stats_tmp = x / 2.0 + call stats_update( "x_budget", stats_tmp, stats ) + end if + end subroutine tend +end module budget_mod +""" + + +def test_a_call_into_a_stubbed_module_reads_and_writes_nothing(tmp_path: Path) -> None: + """CLUBB brackets its budgets with calls into stats_netcdf, a stub: the + translation emits ``pass`` for them, so the source side must not count + their actuals either -- ``stats_tmp`` was a read only the source saw.""" + from recast.fortran.frontend import FortranFrontend + + _write(tmp_path, "budget.f90", STUBBED_CALLER) + frontend = FortranFrontend(stub_modules=["stats_mod"]) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:budget_mod") + facts = frontend.analyze(unit, tmp_path) + assert facts.interface # analysed without the stub module in the tree + blocks = facts.effects["fortran:budget_mod/tend"]["blocks"] + reads = {name for block in blocks for name in block.get("reads", [])} + assert "stats_tmp" not in reads and "x_budget" not in reads + + +KEYWORD_CALLER = """\ +module caller_mod + use solve_mod, only: solve + implicit none + private + public :: run +contains + subroutine run( n, m, x, rc ) + integer, intent(in) :: n, m + real, intent(inout) :: x(n, m) + real, intent(out) :: rc + call solve( n, m, x, rcond = rc ) + end subroutine run +end module caller_mod +""" + +SOLVER_WITH_OPTIONAL = """\ +module solve_mod + implicit none + private + public :: solve + interface solve + module procedure solve_one, solve_many + end interface +contains + subroutine solve_one( n, x, rcond ) + integer, intent(in) :: n + real, intent(inout) :: x(n) + real, intent(out), optional :: rcond + x = 2.0 * x + if ( present( rcond ) ) rcond = 1.0 + end subroutine solve_one + subroutine solve_many( n, m, x, rcond ) + integer, intent(in) :: n, m + real, intent(inout) :: x(n, m) + real, intent(out), optional :: rcond + x = 2.0 * x + if ( present( rcond ) ) rcond = 1.0 + end subroutine solve_many +end module solve_mod +""" + + +def test_a_keyword_actual_into_a_siblings_generic_lands_on_its_own_position(tmp_path: Path) -> None: + """``call band_solve( ..., solut, rcond = rcond )`` (CLUBB): the keyword + names an optional OUT dummy at the end. Bound by position it fell on the + dummy before it and was read; bound by name it is written, not read -- + and the specific is picked with the optional counted.""" + from recast.fortran.frontend import FortranFrontend + + _write(tmp_path, "solve.f90", SOLVER_WITH_OPTIONAL) + _write(tmp_path, "caller.f90", KEYWORD_CALLER) + frontend = FortranFrontend() + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:caller_mod") + facts = frontend.analyze(unit, tmp_path) + (block,) = facts.effects["fortran:caller_mod/run"]["blocks"] + assert "rc" in block["writes"] and "rc" not in block["reads"] + assert "x" in block["writes"] and "x" in block["reads"] # INOUT: both + + +RANK_OVERLOADED_SOLVER = """ +module solve_mod + implicit none + private + public :: solve + interface solve + module procedure solve_one, solve_many + end interface solve +contains + subroutine solve_one( n, flag, a, x, rc ) + integer, intent(in) :: n + logical, intent(in) :: flag + real, intent(inout) :: a(n) + real, intent(out) :: x(n) + real, intent(out), optional :: rc + x = a + if ( present(rc) ) rc = 1.0 + end subroutine solve_one + subroutine solve_many( n, m, flag, a, x ) + integer, intent(in) :: n, m + logical, intent(in) :: flag + real, intent(inout) :: a(n, m) + real, intent(out) :: x(n, m) + x = a + end subroutine solve_many +end module solve_mod +""" + +RANK_CALLER = """ +module caller_mod + use solve_mod, only: solve + implicit none +contains + subroutine run( n, m, flag, a, x ) + integer, intent(in) :: n, m + logical, intent(in) :: flag + real, intent(inout) :: a(n, m) + real, intent(out) :: x(n, m) + call solve( n, m, flag, a, x ) + end subroutine run +end module caller_mod +""" + + +def test_specifics_of_one_arity_are_told_apart_by_rank(tmp_path: Path) -> None: + """CLUBB's ``tridiag_solve``: the single-rhs specific with its optional + ``rcond`` takes as many actuals as the multiple-rhs one without. Picked + by count alone the first won, and ``l_implemented`` -- an IN logical on + the position where the other specific's ``rhs`` sits -- was scored + written. The ranks of the actuals pick the specific whose dummies match.""" + from recast.fortran.frontend import FortranFrontend + + _write(tmp_path, "solve.f90", RANK_OVERLOADED_SOLVER) + _write(tmp_path, "caller.f90", RANK_CALLER) + frontend = FortranFrontend() + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:caller_mod") + facts = frontend.analyze(unit, tmp_path) + (block,) = facts.effects["fortran:caller_mod/run"]["blocks"] + assert "flag" not in block["writes"] and "flag" in block["reads"] + assert "x" in block["writes"] and "x" not in block["reads"] + assert "a" in block["writes"] and "a" in block["reads"] + + +COMPONENT_OUT_CALLER = """ +module pdf_mod + implicit none + type pdf_type + real, allocatable :: chi(:), eta(:) + end type pdf_type +contains + subroutine fill( n, chi, eta ) + integer, intent(in) :: n + real, intent(in) :: chi(n) + real, intent(out) :: eta(n) + eta = chi + end subroutine fill + subroutine run( n, p ) + integer, intent(in) :: n + type(pdf_type), intent(inout) :: p + call fill( n, p%chi, p%eta ) + end subroutine run +end module pdf_mod +""" + + +def test_a_component_out_actual_reads_no_variable_of_the_components_name(tmp_path: Path) -> None: + """``call fill( n, p%chi, p%eta )`` (CLUBB's pdf_closure passes + ``pdf_params%chi_1`` and friends as OUT actuals): the object is written. + The component's bare name was counted as a read of a variable ``eta`` + that the scope does not have; the translation, spelling the attribute, + read no such thing, and the block disagreed.""" + from recast.fortran.frontend import FortranFrontend + + _write(tmp_path, "pdf.f90", COMPONENT_OUT_CALLER) + frontend = FortranFrontend() + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:pdf_mod") + facts = frontend.analyze(unit, tmp_path) + (block,) = facts.effects["fortran:pdf_mod/run"]["blocks"] + assert "p" in block["writes"] and "p" in block["reads"] + assert "eta" not in block["reads"] and "chi" not in block["reads"] + + +HANDS_ON_SOLVER = """ +module solve_mod + implicit none +contains + subroutine solve( n, a, x, rc ) + integer, intent(in) :: n + real, intent(in) :: a(n) + real, intent(out) :: x(n) + real, intent(out), optional :: rc + x = a + if ( present(rc) ) rc = 1.0 + end subroutine solve +end module solve_mod +""" + +HANDS_ON_CALLER = """ +module relay_mod + use solve_mod, only: solve + implicit none +contains + subroutine outer( n, a, x, rc ) + integer, intent(in) :: n + real, intent(in) :: a(n) + real, intent(out) :: x(n) + real, intent(out), optional :: rc + call solve( n, a, x, rc = rc ) + end subroutine outer + subroutine own( n, a, x, rc ) + integer, intent(in) :: n + real, intent(in) :: a(n) + real, intent(out) :: x(n) + real, intent(out), optional :: rc + call inner( n, a, x, rc ) + end subroutine own + subroutine inner( n, a, x, rc ) + integer, intent(in) :: n + real, intent(in) :: a(n) + real, intent(out) :: x(n) + real, intent(out), optional :: rc + x = a + if ( present(rc) ) rc = 2.0 + end subroutine inner +end module relay_mod +""" + + +def test_an_optional_handed_on_to_an_optional_out_is_read_for_its_presence(tmp_path: Path) -> None: + """``call solve( ..., rc = rc )`` with the caller's own optional ``rc``: + the callee asks ``present(rc)``, which the translation spells as the + caller's ``want_rc`` sentinel -- a read of ``rc`` on the target side. + The source side scored only the write, and CLUBB's xm_wpxp_solve + disagreed on the one block that hands ``rcond`` to band_solve. Both a + sibling's procedure and one of this module count it.""" + from recast.fortran.frontend import FortranFrontend + + _write(tmp_path, "solve.f90", HANDS_ON_SOLVER) + _write(tmp_path, "relay.f90", HANDS_ON_CALLER) + frontend = FortranFrontend() + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:relay_mod") + facts = frontend.analyze(unit, tmp_path) + (block,) = facts.effects["fortran:relay_mod/outer"]["blocks"] + assert "rc" in block["reads"] and "rc" in block["writes"] + (block,) = facts.effects["fortran:relay_mod/own"]["blocks"] + assert "rc" in block["reads"] and "rc" in block["writes"] + + CALLBACK = """\ module callback_mod implicit none diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py index dae7354..c5af316 100644 --- a/tests/test_jax_transform.py +++ b/tests/test_jax_transform.py @@ -206,7 +206,7 @@ def test_a_module_state_write_threads_through_the_closure() -> None: }, ] } - pieces, jitted, _delegated = build_module(interface, tree) + pieces, jitted, _delegated, _hosted = build_module(interface, tree) assert sorted(jitted) == ["tick", "use_tick"] text = "\n\n".join(pieces) assert "_host.cache = _res[0]" in text # tick's wrapper stores the write back @@ -241,3 +241,1770 @@ def test_a_constant_table_is_read_through_jnp() -> None: text = _ast.unparse(node) assert "jnp.asarray(MDAYLEAP)[m - 1]" in text assert "x[i]" in text # a lowercase array is a traced value already + + +CYCLES = """\ +module cycle_demo + implicit none + integer, parameter :: r8 = selected_real_kind(12) +contains + subroutine clip_below(n, zlo, z, v, w) + integer, intent(in) :: n + real(r8), intent(in) :: zlo + real(r8), intent(in) :: z(n), v(n) + real(r8), intent(out) :: w(n) + integer :: k + do k = 1, n + if ( z(k) < zlo ) then + w(k) = 0.0_r8 + cycle + end if + w(k) = v(k) * 2.0_r8 + end do + end subroutine clip_below + subroutine first_above(n, zlo, z, kfound) + integer, intent(in) :: n + real(r8), intent(in) :: zlo + real(r8), intent(in) :: z(n) + integer, intent(out) :: kfound + integer :: k + kfound = 0 + do k = 1, n + if ( z(k) > zlo ) then + kfound = k + exit + end if + end do + end subroutine first_above + subroutine exit_index(n, zlo, z, kfound) + integer, intent(in) :: n + real(r8), intent(in) :: zlo + real(r8), intent(in) :: z(n) + integer, intent(out) :: kfound + integer :: k + do k = 1, n + if ( z(k) > zlo ) exit + end do + kfound = k + end subroutine exit_index +end module cycle_demo +""" + + +def test_a_cycle_folds_into_the_branch_and_an_exit_is_a_carried_flag(tmp_path: Path) -> None: + """CLUBB's interpolators: ``if ( ... ) then ... cycle end if`` in a DO + loop. The lowering passed the ``continue`` through into a ``lax.cond`` + branch -- a SyntaxError that took the whole emitted module down. Folded + into the branch structure it is a kernel. An EXIT is a flag the loop + carries: the trips after it do nothing, and the DO variable's value at + the exit is what the code after the loop reads (CLUBB's window search + and its sponge damping).""" + import importlib + import sys + + candidate = port(tmp_path, CYCLES, "cycle_demo") + assert candidate.notes["jax"]["kernels"] == ["clip_below", "exit_index", "first_above"] + assert candidate.notes["jax"]["delegated"] == {} + emitted = candidate.files[Path("cycle_demo_jax.py")].decode() + assert "continue" not in emitted.replace("continuation", "") + assert "break" not in emitted + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("cycle_demo_jax") + import numpy as np + + z = np.array([1.0, 2.0, 3.0, 4.0]) + v = np.array([1.0, 1.0, 1.0, 1.0]) + w = np.asarray(module.clip_below(4, 2.5, z, v)) + assert w.tolist() == [0.0, 0.0, 2.0, 2.0] + assert int(module.first_above(4, 2.5, z)) == 3 + assert int(module.first_above(4, 9.0, z)) == 0 + assert int(module.exit_index(4, 2.5, z)) == 3 + assert int(module.exit_index(4, 9.0, z)) == 5 # ran to completion: n + 1 + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"cycle_demo{suffix}", None) + + +SHIMS = """\ +module shim_demo + implicit none + integer, parameter :: r8 = selected_real_kind(12) +contains + subroutine norms(n, x, root, total, err) + integer, intent(in) :: n + real(r8), intent(in) :: x(n) + real(r8), intent(out) :: root(n), total, err(n) + root = sqrt( x ) + total = sum( x ) + err = erf( x ) + end subroutine norms +end module shim_demo +""" + + +def test_the_jax_runtime_carries_sqrt_sum_and_erf(tmp_path: Path) -> None: + """CLUBB's clipping and PDF closure: ``sqrt``, ``sum`` and ``erf`` reach + the kernels as ``_f_sqrt``, ``_f_vsum`` and ``_f_verf``, which the NumPy + runtime defines and the JAX one did not -- a NameError at the first + call, on every kernel of the unit.""" + import importlib + import sys + + candidate = port(tmp_path, SHIMS, "shim_demo") + assert candidate.notes["jax"]["kernels"] == ["norms"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("shim_demo_jax") + import numpy as np + + x = np.array([4.0, 9.0, -1.0]) + root, total, err = (np.asarray(v) for v in module.norms(3, x)) + assert root[:2].tolist() == [2.0, 3.0] and np.isnan(root[2]) + assert total == 12.0 + assert abs(err[0] - 0.9999999845827421) < 1e-12 + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"shim_demo{suffix}", None) + + +EMPTY_AXIS = """\ +module tracer_demo + implicit none + integer, parameter :: r8 = selected_real_kind(12) +contains + subroutine scale_tracers(n, m, f, x, y) + integer, intent(in) :: n, m + real(r8), intent(in) :: f + real(r8), intent(in) :: x(n, m) + real(r8), intent(out) :: y(n, m) + integer :: s + do s = 1, m + y(:, s) = f * x(:, s) + end do + end subroutine scale_tracers +end module tracer_demo +""" + + +def test_a_loop_over_a_zero_extent_axis_runs_no_iteration(tmp_path: Path) -> None: + """CLUBB's scalar tracers under ``sclr_dim = 0``: ``do sclr = 1, sclr_dim`` + over ``(ngrdcol, nzm, 0)`` arrays. ``fori_loop`` traced the body once + and JAX refused the index into the size-0 axis.""" + import importlib + import sys + + candidate = port(tmp_path, EMPTY_AXIS, "tracer_demo") + assert candidate.notes["jax"]["kernels"] == ["scale_tracers"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("tracer_demo_jax") + import numpy as np + + empty = np.zeros((3, 0), order="F") + assert np.asarray(module.scale_tracers(3, 0, 2.0, empty)).shape == (3, 0) + x = np.ones((3, 2), order="F") + assert np.asarray(module.scale_tracers(3, 2, 2.0, x)).tolist() == [[2.0, 2.0]] * 3 + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"tracer_demo{suffix}", None) + + +FLOORED = """\ +module floor_mod + implicit none + type coefs_type + real(8), allocatable :: coef(:, :) + end type coefs_type +contains + subroutine init_coefs( nz, ngrdcol, c ) + integer, intent(in) :: nz, ngrdcol + type(coefs_type), intent(inout) :: c + allocate( c%coef(1:ngrdcol, 1:nz) ) + c%coef = 2.0d0 + end subroutine init_coefs + subroutine apply( nzt, ngrdcol, c, x, floor ) + integer, intent(in) :: nzt, ngrdcol + type(coefs_type), intent(in) :: c + real(8), intent(inout), dimension(ngrdcol, nzt) :: x + real(8), intent(in), optional :: floor + x = x * c%coef(:, 1:nzt) + if ( present( floor ) ) then + x = max( x, floor ) + end if + end subroutine apply +end module floor_mod +""" + + +def test_a_flat_kernel_spells_the_optional_dummy_the_plan_leaves_out(tmp_path: Path) -> None: + """CLUBB's grid interpolators take an optional ``zt_min``; the plan drops + it (the adapter calls with it absent) and the kernel inlines a body that + tests ``zt_min is not None`` -- a NameError on every call.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "floor_mod.f90").write_text(FLOORED) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:floor_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + ported = candidate.files[Path("floor_mod_jax.py")].decode() + assert "floor = None" in ported + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("floor_mod_jax") + import numpy as np + + coef = np.full((2, 3), 2.0, order="F") + x = np.ones((2, 3), order="F") + assert "apply_flat" in candidate.notes["jax"]["kernels"] + result = np.asarray(module.apply_flat(3, 2, x, 3, coef)) + assert result.tolist() == [[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"floor_mod{suffix}", None) + + +PAIRS_IN_A_LOOP = """\ +module pair_demo + implicit none + integer, parameter :: r8 = selected_real_kind(12) +contains + elemental subroutine split(x, lo, hi) + real(r8), intent(in) :: x + real(r8), intent(out) :: lo, hi + lo = x - 1.0_r8 + hi = x + 1.0_r8 + end subroutine split + subroutine bracket(n, nz, x, lo, hi) + integer, intent(in) :: n, nz + real(r8), intent(in) :: x(n, nz) + real(r8), intent(out) :: lo(n, nz), hi(n, nz) + integer :: i + do i = 1, n + call split( x(i, :), lo(i, :), hi(i, :) ) + end do + end subroutine bracket +end module pair_demo +""" + + +def test_a_call_result_tuple_inside_a_loop_is_the_bodys_own(tmp_path: Path) -> None: + """The anchor spells a two-output call as ``_out = split(...)`` and + unpacks it on the next lines. Carried through the fori_loop, ``_out`` + was read for the initial carry before any assignment + (UnboundLocalError; CLUBB's new_hybrid_pdf_driver). A name not bound + before the loop is the body's own.""" + import importlib + import sys + + candidate = port(tmp_path, PAIRS_IN_A_LOOP, "pair_demo") + assert "bracket" in candidate.notes["jax"]["kernels"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("pair_demo_jax") + import numpy as np + + x = np.array([[1.0, 2.0], [3.0, 4.0]], order="F") + lo, hi = (np.asarray(v) for v in module.bracket(2, 2, x)) + assert lo.tolist() == [[0.0, 1.0], [2.0, 3.0]] and hi.tolist() == [[2.0, 3.0], [4.0, 5.0]] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"pair_demo{suffix}", None) + + +WRITER = """\ +module writer_mod + implicit none +contains + subroutine fill(n, x, y) + integer, intent(in) :: n + real(8), intent(in) :: x(n) + real(8) :: y(n) + y = x + 1.0d0 + end subroutine fill +end module writer_mod +""" + +CALLS_WRITER = """\ +module caller_mod + use writer_mod, only: fill + implicit none +contains + subroutine run(n, x, y) + integer, intent(in) :: n + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + call fill( n, x, y ) + y = y * 2.0d0 + end subroutine run +end module caller_mod +""" + + +def test_a_bare_call_binds_the_kernels_returned_buffer(tmp_path: Path) -> None: + """CLUBB declares ``xp3_lg_2005_ansatz``'s ``xp3`` without an intent and + the extension's frontend overrides it to OUT. The caller's anchor, seeing + no intent at the call, calls bare: the callee writes ``y`` in place and + returns it, and the caller ignores the return. A kernel cannot write in + place -- its return *is* the output -- so the statement binds what + comes back to the actual; without it advance_xp3 got back the zeros it + passed in for every xp3.""" + import importlib + import sys + + from recast.registry import REGISTRY + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + def overriding_frontend(**_config: object) -> FortranFrontend: + return FortranFrontend( + flatten=True, buffer_out_arrays="all", intent_overrides={"fill": {"y": "OUT"}} + ) + + REGISTRY.register("frontend", "fortran-y-out", overriding_frontend, replace=True) + (tmp_path / "writer_mod.f90").write_text(WRITER) + (tmp_path / "caller_mod.f90").write_text(CALLS_WRITER) + frontend = overriding_frontend() + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:caller_mod") + facts = frontend.analyze(unit, tmp_path) + conventions = TreeConventions(frontend="fortran-y-out") + candidate = TreeToJax(conventions).apply(unit, facts, {"root": str(tmp_path)}) + assert "run" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + ported = candidate.files[Path("caller_mod_jax.py")].decode() + assert "_writer_mod.fill(n, x, y)" in candidate.files[Path("caller_mod_numpy.py")].decode() + assert "y = _writer_mod_jax._fill_k_impl(n, x, y)" in ported + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("caller_mod_jax") + import numpy as np + + got = module.run(2, np.array([1.0, 2.0]), np.zeros(2)) + assert np.asarray(got).tolist() == [4.0, 6.0] + finally: + sys.path.remove(str(out)) + for name in list(sys.modules): + if name.startswith(("caller_mod", "writer_mod")): + sys.modules.pop(name, None) + + +ELEMENTAL_COMPANION = """\ +module elem_mod + implicit none + integer, parameter :: r8 = selected_real_kind(12) +contains + elemental subroutine split(x, lo, hi) + real(r8), intent(in) :: x + real(r8), intent(out) :: lo, hi + lo = x - 1.0_r8 + hi = x + 1.0_r8 + end subroutine split +end module elem_mod +""" + +CALLS_ELEMENTAL_COMPANION = """\ +module bracket_mod + use elem_mod, only: split + implicit none + integer, parameter :: r8 = selected_real_kind(12) +contains + subroutine bracket(n, nz, x, lo, hi) + integer, intent(in) :: n, nz + real(r8), intent(in) :: x(n, nz) + real(r8), intent(out) :: lo(n, nz), hi(n, nz) + integer :: i + do i = 1, n + call split( x(i, :), lo(i, :), hi(i, :) ) + end do + end subroutine bracket +end module bracket_mod +""" + + +def test_an_elemental_call_of_a_companion_broadcasts_its_kernel(tmp_path: Path) -> None: + """CLUBB's new_hybrid_pdf_driver calls new_hybrid_pdf's elemental + ``calculate_mixture_fraction`` over column slices: the anchor's + ``_f_ecall(_new.calculate_mixture_fraction, ...)``. Left as the host + attribute, the vectorize traced a NumPy function; the companion's kernel + implementation goes under it instead.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "elem_mod.f90").write_text(ELEMENTAL_COMPANION) + (tmp_path / "bracket_mod.f90").write_text(CALLS_ELEMENTAL_COMPANION) + frontend = FortranFrontend(flatten=True) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:bracket_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "bracket" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + ported = candidate.files[Path("bracket_mod_jax.py")].decode() + assert "_f_ecall(_elem_mod_jax._split_k_impl, " in ported + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("bracket_mod_jax") + import numpy as np + + x = np.array([[1.0, 2.0], [3.0, 4.0]], order="F") + lo, hi = (np.asarray(v) for v in module.bracket(2, 2, x)) + assert lo.tolist() == [[0.0, 1.0], [2.0, 3.0]] and hi.tolist() == [[2.0, 3.0], [4.0, 5.0]] + finally: + sys.path.remove(str(out)) + for name in list(sys.modules): + if name.startswith(("bracket_mod", "elem_mod")): + sys.modules.pop(name, None) + + +GUARDED_BY_A_STATIC = """\ +module guard_demo + implicit none +contains + subroutine first_tracer(n, m, x, y) + integer, intent(in) :: n, m + real(8), intent(in) :: x(n, m) + real(8), intent(out) :: y(n) + y = 0.0d0 + if ( m > 0 ) then + y(:) = x(:, 1) + end if + end subroutine first_tracer +end module guard_demo +""" + + +def test_a_branch_on_a_static_scalar_is_a_trace_time_if(tmp_path: Path) -> None: + """CLUBB guards its scalar-tracer stores with ``if ( sclr_dim > 0 )``. + Lowered to lax.cond both arms are traced, and the store into the + zero-extent (or never allocated) array is an IndexError at trace time. + ``sclr_dim`` is a static argument of the kernel -- a Python int under + jit -- so the branch is a Python if.""" + import importlib + import sys + + candidate = port(tmp_path, GUARDED_BY_A_STATIC, "guard_demo") + assert candidate.notes["jax"]["kernels"] == ["first_tracer"] + ported = candidate.files[Path("guard_demo_jax.py")].decode() + assert "if _f_concrete(m > 0):" in ported and "if m > 0:" in ported + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("guard_demo_jax") + import numpy as np + + assert np.asarray(module.first_tracer(2, 0, np.zeros((2, 0)))).tolist() == [0.0, 0.0] + x = np.array([[1.0, 2.0], [3.0, 4.0]], order="F") + assert np.asarray(module.first_tracer(2, 2, x)).tolist() == [1.0, 3.0] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"guard_demo{suffix}", None) + + +EARLY_RETURNS = """\ +module early_mod + implicit none + type coefs_type + real(8), allocatable :: coef(:, :) + end type coefs_type +contains + subroutine init_coefs( nz, ngrdcol, c ) + integer, intent(in) :: nz, ngrdcol + type(coefs_type), intent(inout) :: c + allocate( c%coef(1:ngrdcol, 1:nz) ) + c%coef = 2.0d0 + end subroutine init_coefs + subroutine apply( nzt, ngrdcol, c, x, y, bad, worse ) + integer, intent(in) :: nzt, ngrdcol + type(coefs_type), intent(in) :: c + real(8), intent(inout), dimension(ngrdcol, nzt) :: x + real(8), intent(out), dimension(ngrdcol, nzt) :: y + logical, intent(in) :: bad, worse + x = x * c%coef(:, 1:nzt) + y = x + if ( bad ) then + if ( worse ) then + return + end if + y = y + 1.0d0 + return + end if + y = y + 10.0d0 + end subroutine apply +end module early_mod +""" + + +def test_early_returns_of_the_same_tuple_fold_into_the_branches(tmp_path: Path) -> None: + """CLUBB's advance_clubb_core returns after each solver when the error + code says so, the outputs as they stand. The single-exit rewrite merges + one value with a where and refused a tuple; when every early return is + the final tuple, the continuation folds into the non-returning branches + and the kernel keeps one exit.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "early_mod.f90").write_text(EARLY_RETURNS) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:early_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "apply_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("early_mod_jax") + import numpy as np + + coef = np.full((1, 2), 2.0, order="F") + + def run(bad, worse): + x = np.ones((1, 2), order="F") + _x, y = module.apply_flat(2, 1, x, np.zeros((1, 2), order="F"), bad, worse, 2, coef) + return np.asarray(y).tolist() + + assert run(False, False) == [[12.0, 12.0]] + assert run(True, False) == [[3.0, 3.0]] + assert run(True, True) == [[2.0, 2.0]] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"early_mod{suffix}", None) + + +RETURNS_IN_A_LOOP = """\ +module loopret_mod + implicit none + type coefs_type + real(8), allocatable :: coef(:, :) + end type coefs_type +contains + subroutine init_coefs( nz, ngrdcol, c ) + integer, intent(in) :: nz, ngrdcol + type(coefs_type), intent(inout) :: c + allocate( c%coef(1:ngrdcol, 1:nz) ) + c%coef = 2.0d0 + end subroutine init_coefs + subroutine apply( nzt, ngrdcol, c, x, y, err ) + integer, intent(in) :: nzt, ngrdcol + type(coefs_type), intent(in) :: c + real(8), intent(inout), dimension(ngrdcol, nzt) :: x + real(8), intent(out), dimension(ngrdcol, nzt) :: y + integer, intent(in) :: err(ngrdcol) + integer :: i + real(8) :: lo, hi + x = x * c%coef(:, 1:nzt) + y = x + do i = 1, ngrdcol + if ( err(i) /= 0 ) then + return + end if + y(i, :) = y(i, :) + 1.0d0 + end do + y = y + 10.0d0 + call bracket( x(1, 1), lo, hi ) + y = y + lo + hi + end subroutine apply + elemental subroutine bracket( v, lo, hi ) + real(8), intent(in) :: v + real(8), intent(out) :: lo, hi + lo = v - 1.0d0 + hi = v + 1.0d0 + end subroutine bracket +end module loopret_mod +""" + + +def test_an_early_return_inside_a_loop_becomes_a_flag(tmp_path: Path) -> None: + """advance_clubb_core checks the error code per column inside a loop and + returns. No branch structure holds that; the return sets a flag, every + later statement runs under it, and the remaining iterations do nothing.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "loopret_mod.f90").write_text(RETURNS_IN_A_LOOP) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:loopret_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "apply_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("loopret_mod_jax") + import numpy as np + + coef = np.full((2, 1), 2.0, order="F") + + def run(err): + x = np.ones((2, 1), order="F") + _x, y = module.apply_flat( + 1, 2, x, np.zeros((2, 1), order="F"), np.array(err, dtype=np.int32), 1, coef + ) + return np.asarray(y).ravel().tolist() + + # x(1,1) is 2 after scaling: lo + hi = 4 on the path that reaches the + # bracket call (an elemental with two outputs, the anchor's _out). + assert run([0, 0]) == [17.0, 17.0] + assert run([0, 1]) == [3.0, 2.0] # the second column returns before its +1 and the rest + assert run([1, 0]) == [2.0, 2.0] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"loopret_mod{suffix}", None) + + +UNPACKS_AN_OBJECT = """\ +module unpack_mod + implicit none + type knobs_type + real(8) :: gain = 2.0d0 + end type knobs_type +contains + subroutine scale( nzt, ngrdcol, x, y, k ) + integer, intent(in) :: nzt, ngrdcol + real(8), intent(inout), dimension(ngrdcol, nzt) :: x + real(8), intent(inout), dimension(ngrdcol, nzt) :: y + type(knobs_type), intent(inout), optional :: k + if ( present( k ) ) k%gain = k%gain * 2.0d0 + x = x * 2.0d0 + y = x + 1.0d0 + end subroutine scale + subroutine step( nzt, ngrdcol, k, x, y ) + integer, intent(in) :: nzt, ngrdcol + type(knobs_type), intent(inout) :: k + real(8), intent(inout), dimension(ngrdcol, nzt) :: x + real(8), intent(inout), dimension(ngrdcol, nzt) :: y + call scale( nzt, ngrdcol, x, y, k ) + y = y + k%gain + end subroutine step +end module unpack_mod +""" + + +def test_unpacking_an_object_from_the_elided_call_buffer_is_already_true(tmp_path: Path) -> None: + """advance_clubb_core's anchor unpacks what a solver hands back -- + ``stats = _out[0]``, ``pdf_params = _out[4]`` -- from the buffer the + flat rewrite elides at the call: the flat outputs bound to the actuals, + so the object's unpack is already true and the statement goes; an + array's is the actual it names. The object here is an *optional* INOUT + (calc_brunt_vaisala_freq_sqd's ``stats``): its slot is in the anchor's + tuple all the same, and was not in the rewrite's list.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "unpack_mod.f90").write_text(UNPACKS_AN_OBJECT) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:unpack_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + assert "k = _out[2]" in candidate.files[Path("unpack_mod_numpy.py")].decode() + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("unpack_mod_jax") + import numpy as np + + x = np.ones((1, 2), order="F") + result = module.step_flat(2, 1, x, np.zeros((1, 2), order="F"), np.float64(2.0)) + got = [np.asarray(v).tolist() for v in result] + # The optional object is absent in the flat world (the plan leaves + # optional dummies out, as the adapter calls without them), so the + # callee's ``present(k)`` branch does not run: gain stays 2, y is + # x * 2 + 1 + 2. + assert got == [[[2.0, 2.0]], [[5.0, 5.0]], 2.0] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"unpack_mod{suffix}", None) + + +CHECKS = """\ +module checks_mod + implicit none +contains + subroutine complain( n, x, msg ) + integer, intent(in) :: n + real(8), intent(in) :: x(n) + character(len=*), intent(out) :: msg + msg = "fine" + if ( any( x < 0.0d0 ) ) msg = "negative" + end subroutine complain +end module checks_mod +""" + +GUARDED_CHECK = """\ +module guarded_mod + use checks_mod, only: complain + implicit none +contains + subroutine step( n, debug_level, x, y ) + integer, intent(in) :: n, debug_level + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + character(len=16) :: msg + y = 2.0d0 * x + if ( debug_level >= 2 ) then + call complain( n, y, msg ) + end if + end subroutine step +end module guarded_mod +""" + + +def test_a_companion_procedure_the_port_left_on_the_host_stays_under_its_guard( + tmp_path: Path, +) -> None: + """CLUBB's advance_clubb_core calls numerical_check's parameterization + check under ``clubb_at_least_debug_level_api(2)``, false with statistics + off; the check takes character arguments and its port leaves it on the + host. Refusing the whole step for a call that never runs was the + alternative: the call stays the host's under its trace-time guard, and + the note names it.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "checks_mod.f90").write_text(CHECKS) + (tmp_path / "guarded_mod.f90").write_text(GUARDED_CHECK) + frontend = FortranFrontend(flatten=True) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:guarded_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "step" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + assert candidate.notes["jax"]["host_calls"] == {"step": ["checks_mod.complain"]} + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("guarded_mod_jax") + import numpy as np + + assert np.asarray(module.step(2, 0, np.array([1.0, -1.0]))).tolist() == [2.0, -2.0] + finally: + sys.path.remove(str(out)) + for name in list(sys.modules): + if name.startswith(("guarded_mod", "checks_mod")): + sys.modules.pop(name, None) + + +OBJECT_CHECK = """\ +module ocheck_mod + implicit none + type knobs_type + real(8) :: gain = 2.0d0 + end type knobs_type +contains + subroutine inspect( n, k, x, msg ) + integer, intent(in) :: n + type(knobs_type), intent(inout) :: k + real(8), intent(in) :: x(n) + character(len=*), intent(out) :: msg + msg = "fine" + if ( any( x < k%gain ) ) msg = "small" + k%gain = k%gain + 1.0d0 + end subroutine inspect +end module ocheck_mod +""" + +GUARDED_OBJECT_CHECK = """\ +module oguarded_mod + use ocheck_mod, only: knobs_type, inspect + implicit none +contains + subroutine step( n, debug_level, k, x, y ) + integer, intent(in) :: n, debug_level + type(knobs_type), intent(inout) :: k + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + character(len=16) :: msg + y = k%gain * x + if ( debug_level >= 2 ) then + call inspect( n, k, y, msg ) + end if + y = y + k%gain + end subroutine step +end module oguarded_mod +""" + + +def test_a_flat_companion_the_port_left_on_the_host_stays_under_its_guard(tmp_path: Path) -> None: + """The same, through the flat-callee path: the check takes the object + (advance_clubb_core's parameterization check takes gr and err_info), so + it has a plan, and the rewrite of the call into the port's flat kernel + found none.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "ocheck_mod.f90").write_text(OBJECT_CHECK) + (tmp_path / "oguarded_mod.f90").write_text(GUARDED_OBJECT_CHECK) + frontend = FortranFrontend(flatten=True) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:oguarded_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + assert candidate.notes["jax"]["host_calls"] == {"step_flat": ["ocheck_mod.inspect"]} + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("oguarded_mod_jax") + import numpy as np + + got = module.step_flat(2, 0, np.array([1.0, -1.0]), np.zeros(2), 1, np.float64(2.0)) + y, gain = (np.asarray(v).tolist() for v in got) + assert y == [4.0, 0.0] and gain == 2.0 + finally: + sys.path.remove(str(out)) + for name in list(sys.modules): + if name.startswith(("oguarded_mod", "ocheck_mod")): + sys.modules.pop(name, None) + + +OBJECT_QUERY = """\ +module oquery_mod + use ocheck_mod, only: knobs_type + use stats_query_mod, only: var_on_list + implicit none +contains + subroutine step( n, k, x, y ) + integer, intent(in) :: n + type(knobs_type), intent(in) :: k + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + y = k%gain * x + if ( var_on_list( k, "gain" ) ) then + y = y + k%gain + end if + end subroutine step +end module oquery_mod +""" + +STATS_QUERY = """\ +module stats_query_mod + use ocheck_mod, only: knobs_type + implicit none +contains + logical function var_on_list( k, name ) + type(knobs_type), intent(in) :: k + character(len=*), intent(in) :: name + var_on_list = .true. + end function var_on_list +end module stats_query_mod +""" + + +def test_an_object_handed_whole_to_the_host_is_rebuilt_at_entry(tmp_path: Path) -> None: + """``if ( var_on_stats_list( stats, "rsat" ) )``: the query is a + framework stand-in's, and it takes the object whole -- which the kernel + took apart into components. Rebuilt once at entry from them, as the + NumPy flat wrapper does, instead of an UnboundLocalError on a name the + kernel never bound.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "ocheck_mod.f90").write_text(OBJECT_CHECK) + (tmp_path / "oquery_mod.f90").write_text(OBJECT_QUERY) + (tmp_path / "stats_query_mod.f90").write_text(STATS_QUERY) + frontend = FortranFrontend(flatten=True, stub_modules=["stats_query_mod"]) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:oquery_mod") + facts = frontend.analyze(unit, tmp_path) + conventions = TreeConventions( + stub_modules=frozenset({"stats_query_mod"}), + framework={"stats_query_mod": "def var_on_list(k, name):\n return name == 'gain'\n"}, + ) + candidate = TreeToJax(conventions).apply(unit, facts, {"root": str(tmp_path)}) + assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + emitted = candidate.files[Path("oquery_mod_jax.py")].decode() + assert "k = _host._Record(gain=k__gain)" in emitted + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("oquery_mod_jax") + import numpy as np + + got = module.step_flat(2, np.array([1.0, -1.0]), np.zeros(2), 2, np.float64(2.0)) + assert np.asarray(got).tolist() == [4.0, 0.0] + finally: + sys.path.remove(str(out)) + for name in list(sys.modules): + if name.startswith(("oquery_mod", "ocheck_mod", "stats_query_mod")): + sys.modules.pop(name, None) + + +EMPTY_LOOP = """\ +module emptyloop_mod + implicit none +contains + subroutine scale( n, x, y ) + integer, intent(in) :: n + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + integer :: i + y = 2.0d0 * x + do i = 1, n + continue + end do + end subroutine scale +end module emptyloop_mod +""" + + +def test_a_loop_whose_body_lowered_to_nothing_is_dropped(tmp_path: Path) -> None: + """CLUBB's clipping routines loop over the columns to sample statistics; + with the sampling calls dropped by their stand-in, the loop's body is + ``pass`` and there is nothing to carry. Fortran ran it for nothing.""" + import importlib + import sys + + candidate = port(tmp_path, EMPTY_LOOP, "emptyloop_mod") + assert candidate.notes["jax"]["kernels"] == ["scale"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("emptyloop_mod_jax") + import numpy as np + + got = module.scale(2, np.array([1.0, 3.0])) + assert np.asarray(got).tolist() == [2.0, 6.0] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"emptyloop_mod{suffix}", None) + + +STRIDED = """\ +module strided_mod + implicit none +contains + subroutine running( n, dir, x, y ) + integer, intent(in) :: n, dir + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + integer :: k, lb, ub + real(8) :: acc + if ( dir > 0 ) then + lb = 1 + ub = n + else + lb = n + ub = 1 + end if + acc = 0.0d0 + do k = lb, ub, dir + acc = acc + x(k) + y(k) = acc + end do + end subroutine running +end module strided_mod +""" + + +def test_a_loop_with_a_named_stride_runs_in_that_direction(tmp_path: Path) -> None: + """``do k = gr%k_lb_zt, gr%k_ub_zt, gr%grid_dir_indx``: CLUBB's grid + direction is a run-time +-1, a name in the stride slot. The trip count + is the runtime's, the index is remapped from the trip.""" + import importlib + import sys + + candidate = port(tmp_path, STRIDED, "strided_mod") + assert candidate.notes["jax"]["kernels"] == ["running"], candidate.notes["jax"] + assert "_f_trips(" in candidate.files[Path("strided_mod_jax.py")].decode() + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("strided_mod_jax") + import numpy as np + + x = np.array([1.0, 2.0, 4.0]) + up = np.asarray(module.running(3, 1, x)).tolist() + down = np.asarray(module.running(3, -1, x)).tolist() + assert up == [1.0, 3.0, 7.0] + assert down == [7.0, 6.0, 4.0] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"strided_mod{suffix}", None) + + +SWITCHED_CHECK = """\ +module switched_mod + use silent_mod, only: complain + implicit none +contains + subroutine step( n, l_check, x, y ) + integer, intent(in) :: n + logical, intent(in) :: l_check + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + y = 2.0d0 * x + if ( l_check ) then + call complain( n, y ) + end if + end subroutine step +end module switched_mod +""" + + +def test_a_logical_scalar_dummy_is_a_static_switch(tmp_path: Path) -> None: + """CLUBB's configuration flags are logical dummies (the model's + ``clubb_config_flags`` components). Static under jit: the branch is a + Python if at trace time, and the check the port left on the host -- + under a flag the run has off -- is never traced.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "silent_mod.f90").write_text(SILENT_CHECK) + (tmp_path / "switched_mod.f90").write_text(SWITCHED_CHECK) + frontend = FortranFrontend(flatten=True) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:switched_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "step" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + emitted = candidate.files[Path("switched_mod_jax.py")].decode() + assert "static_argnums=(0, 1)" in emitted + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("switched_mod_jax") + import numpy as np + + got = module.step(2, False, np.array([1.0, -1.0])) + assert np.asarray(got).tolist() == [2.0, -2.0] + finally: + sys.path.remove(str(out)) + for name in list(sys.modules): + if name.startswith(("switched_mod", "silent_mod")): + sys.modules.pop(name, None) + + +OMITS_AN_OBJECT = """\ +module omit_mod + implicit none + type knobs_type + real(8) :: gain = 2.0d0 + end type knobs_type +contains + subroutine scale( nzt, ngrdcol, x, y, k ) + integer, intent(in) :: nzt, ngrdcol + real(8), intent(inout), dimension(ngrdcol, nzt) :: x + real(8), intent(inout), dimension(ngrdcol, nzt) :: y + type(knobs_type), intent(inout), optional :: k + if ( present( k ) ) k%gain = k%gain * 2.0d0 + x = x * 2.0d0 + y = x + 1.0d0 + end subroutine scale + subroutine step( nzt, ngrdcol, k, x, y ) + integer, intent(in) :: nzt, ngrdcol + type(knobs_type), intent(inout) :: k + real(8), intent(inout), dimension(ngrdcol, nzt) :: x + real(8), intent(inout), dimension(ngrdcol, nzt) :: y + call scale( nzt, ngrdcol, x, y ) + y = y + k%gain + end subroutine step +end module omit_mod +""" + + +def test_an_optional_object_the_caller_leaves_out_is_absent_in_the_callee(tmp_path: Path) -> None: + """pdf_closure_driver_zm calls pdf_closure without its optional + ``pdf_implicit_coefs_terms``. The callee's kernel still takes the + object's components (its flat signature has them): ``None`` in each, + ``present()`` false at trace time, and what it hands back for them + dropped.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "omit_mod.f90").write_text(OMITS_AN_OBJECT) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:omit_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("omit_mod_jax") + import numpy as np + + x = np.ones((1, 2), order="F") + result = module.step_flat(2, 1, x, np.zeros((1, 2), order="F"), np.float64(2.0)) + got = [np.asarray(v).tolist() for v in result] + # gain is read, never written, so it is not returned; the callee's + # present(k) branch did not run: y is x * 2 + 1 + 2. + assert got == [[[2.0, 2.0]], [[5.0, 5.0]]] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"omit_mod{suffix}", None) + + +DISCARDS_AN_OUTPUT = """\ +module discard_mod + implicit none + type knobs_type + real(8) :: gain = 2.0d0 + end type knobs_type +contains + subroutine solve( nzt, ngrdcol, k, x, y, resid ) + integer, intent(in) :: nzt, ngrdcol + type(knobs_type), intent(in) :: k + real(8), intent(in), dimension(ngrdcol, nzt) :: x + real(8), intent(inout), dimension(ngrdcol, nzt) :: y + real(8), intent(out), dimension(ngrdcol, nzt), optional :: resid + y = k%gain * x + if ( present( resid ) ) resid = y - x + end subroutine solve + subroutine step( nzt, ngrdcol, k, x, y ) + integer, intent(in) :: nzt, ngrdcol + type(knobs_type), intent(in) :: k + real(8), intent(in), dimension(ngrdcol, nzt) :: x + real(8), intent(inout), dimension(ngrdcol, nzt) :: y + call solve( nzt, ngrdcol, k, x, y ) + y = y + 1.0d0 + end subroutine step +end module discard_mod +""" + + +def test_an_output_the_anchor_discards_is_dropped_from_the_elided_buffer(tmp_path: Path) -> None: + """CLUBB's solvers return an optional residual the caller does not ask + for: the anchor's ``_ = _out[4]`` names a slot with no actual, and the + unpack is dropped rather than refused.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "discard_mod.f90").write_text(DISCARDS_AN_OUTPUT) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:discard_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "_ = _out[" in candidate.files[Path("discard_mod_numpy.py")].decode() + assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("discard_mod_jax") + import numpy as np + + x = np.ones((1, 2), order="F") + got = module.step_flat(2, 1, x, np.zeros((1, 2), order="F"), np.float64(3.0)) + assert np.asarray(got).tolist() == [[4.0, 4.0]] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"discard_mod{suffix}", None) + + +SIZED_BY_AN_ALLOCATOR = """\ +module sized_mod + implicit none + type coefs_type + real(8), allocatable :: coef(:, :) + end type coefs_type +contains + subroutine init_coefs( ngrdcol, nz, p ) + integer, intent(in) :: ngrdcol, nz + type(coefs_type), intent(out) :: p + allocate( p%coef(1:ngrdcol, 1:nz) ) + p%coef = 0.0d0 + end subroutine init_coefs + subroutine apply_coefs( ngrdcol, p, x ) + integer, intent(in) :: ngrdcol + type(coefs_type), intent(in) :: p + real(8), intent(inout) :: x(ngrdcol) + integer :: i + do i = 1, ngrdcol + x(i) = x(i) + sum( p%coef(i, :) ) + end do + end subroutine apply_coefs + subroutine step( ngrdcol, nz, p, x ) + integer, intent(in) :: ngrdcol, nz + type(coefs_type), intent(inout) :: p + real(8), intent(inout) :: x(ngrdcol) + p%coef(:, 1) = p%coef(:, 1) + 1.0d0 + call apply_coefs( ngrdcol, p, x ) + end subroutine step +end module sized_mod +""" + + +def test_a_callee_extent_argument_is_the_callers_axis(tmp_path: Path) -> None: + """A component allocated by another routine's dummy (CLUBB's + ``coef_wp4_implicit(1:ngrdcol, 1:nz)``) reaches a callee that has no + dummy of that name as an extent argument of its plan. The caller + passes that axis of its own component.""" + import importlib + import sys + + from recast.fortran.flatten import plans_from_facts + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "sized_mod.f90").write_text(SIZED_BY_AN_ALLOCATOR) + frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"}) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:sized_mod") + facts = frontend.analyze(unit, tmp_path) + plans = {p.name: p for p in plans_from_facts(facts)} + assert plans["apply_coefs_flat"].extent_args, plans["apply_coefs_flat"] + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("sized_mod_jax") + import numpy as np + + coef = np.zeros((1, 2), order="F") + got = module.step_flat(1, 2, np.array([1.0]), coef) + assert [np.asarray(v).tolist() for v in got] == [[2.0], [[1.0, 0.0]]] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"sized_mod{suffix}", None) + + +OWN_CHECK_UNDER_A_SWITCH = """\ +module ownswitch_mod + implicit none +contains + subroutine complain( n, x, msg ) + integer, intent(in) :: n + real(8), intent(in) :: x(n) + character(len=*), intent(out) :: msg + msg = "fine" + if ( any( x < 0.0d0 ) ) msg = "negative" + end subroutine complain + subroutine step( n, l_check, x, y ) + integer, intent(in) :: n + logical, intent(in) :: l_check + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + character(len=16) :: msg + y = 2.0d0 * x + if ( l_check ) then + call complain( n, y, msg ) + end if + end subroutine step +end module ownswitch_mod +""" + + +def test_a_subprogram_of_the_module_the_port_could_not_emit_stays_on_the_host( + tmp_path: Path, +) -> None: + """pdf_closure_driver calls its zm variant under a switch the run has + off; fill_holes_vertical_api dispatches on a type. A same-module + callee the port could not emit no longer delegates its caller and the + callers above it: the call stays the host's (``_host.``) under + the anchor's guard, named in the notes, never traced while the guard + holds.""" + import importlib + import sys + + candidate = port(tmp_path, OWN_CHECK_UNDER_A_SWITCH, "ownswitch_mod") + assert candidate.notes["jax"]["kernels"] == ["step"], candidate.notes["jax"] + assert "complain" in candidate.notes["jax"]["delegated"] + assert candidate.notes["jax"]["host_calls"] == {"step": ["ownswitch_mod.complain"]} + emitted = candidate.files[Path("ownswitch_mod_jax.py")].decode() + assert "_host.complain(" in emitted + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("ownswitch_mod_jax") + import numpy as np + + got = module.step(2, False, np.array([1.0, -1.0])) + assert np.asarray(got).tolist() == [2.0, -2.0] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"ownswitch_mod{suffix}", None) + + +WINDOWED = """\ +module window_mod + implicit none +contains + subroutine fill( n, w, dir, x, y ) + integer, intent(in) :: n, w, dir + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + integer :: k, lo, hi + y = x + do k = 1 + w, n - w + lo = k - dir * w + hi = k + dir * w + if ( any( x(lo:hi:dir) < 0.0d0 ) ) then + y(k) = -1.0d0 + else + y(k) = maxval( x(lo:hi:dir) ) + end if + end do + end subroutine fill +end module window_mod +""" + + +def test_a_window_with_traced_bounds_a_static_distance_apart_is_a_gather(tmp_path: Path) -> None: + """CLUBB's sliding-window hole filler reads ``field(i, k_start:k_end:dir)`` + with ``k_start = k - dir * n`` and ``k_end = k + dir * n``: bounds that + trace (the loop index), a distance apart that does not. The slice is a + gather at ``lo + arange(trips) * step`` -- the mask rules, for a slice + whose length depends on a traced bound, never see it.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "window_mod.f90").write_text(WINDOWED) + frontend = FortranFrontend(flatten=True) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:window_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert candidate.notes["jax"]["kernels"] == ["fill"], candidate.notes["jax"] + emitted = candidate.files[Path("window_mod_jax.py")].decode() + assert "jnp.arange(_f_trips(0, " in emitted, emitted + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("window_mod_jax") + import numpy as np + + x = np.array([1.0, 5.0, 2.0, -3.0, 4.0, 6.0]) + got = np.asarray(module.fill(6, 1, 1, x)).tolist() + assert got == [1.0, 5.0, -1.0, -1.0, -1.0, 6.0] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"window_mod{suffix}", None) + + +RESHAPED_ACTUAL = """\ +module reshaped_mod + implicit none + type knobs_type + real(8) :: gain = 2.0d0 + end type knobs_type +contains + subroutine solve( n, k, rhs ) + integer, intent(in) :: n + type(knobs_type), intent(in) :: k + real(8), intent(inout) :: rhs(n, 1) + rhs(:, 1) = k%gain * rhs(:, 1) + end subroutine solve + subroutine step( n, k, x ) + integer, intent(in) :: n + type(knobs_type), intent(in) :: k + real(8), intent(inout) :: x(n) + call solve( n, k, x ) + x = x + 1.0d0 + end subroutine step +end module reshaped_mod +""" + + +def test_an_array_passed_through_a_reshape_comes_back_in_its_own_shape(tmp_path: Path) -> None: + """CLUBB hands a 2-d right-hand side to a solver declared over three + axes: sequence association, which the anchor spells as + ``np.reshape(rhs, (n, m, 1), order='F')`` for the actual. The kernel's + output is reshaped back to the array's own shape and rebinds it -- a + reshape is no store target.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "reshaped_mod.f90").write_text(RESHAPED_ACTUAL) + frontend = FortranFrontend(flatten=True) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:reshaped_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "reshape(x" in candidate.files[Path("reshaped_mod_numpy.py")].decode() + assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("reshaped_mod_jax") + import numpy as np + + got = module.step_flat(2, np.array([1.0, 3.0]), 2, np.float64(2.0)) + assert np.asarray(got).tolist() == [3.0, 7.0] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"reshaped_mod{suffix}", None) + + +EMPTY_GUARD = """\ +module emptyguard_mod + implicit none +contains + subroutine scale( n, level, x, y ) + integer, intent(in) :: n, level + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + y = 2.0d0 * x + if ( level > 0 ) then + print *, "scaled", n + end if + if ( level > 1 ) then + print *, "twice" + else + y = y + 1.0d0 + end if + end subroutine scale +end module emptyguard_mod +""" + + +def test_a_static_branch_whose_arms_lowered_to_nothing_is_no_branch(tmp_path: Path) -> None: + """``if ( clubb_at_least_debug_level( 0 ) ) then`` around a print: the + print is dropped, the static Python if was emitted with no body -- a + SyntaxError that took the whole emitted module down. Nothing in both + arms is no branch; nothing in one arm is ``pass``.""" + import importlib + import sys + + candidate = port(tmp_path, EMPTY_GUARD, "emptyguard_mod") + assert candidate.notes["jax"]["kernels"] == ["scale"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("emptyguard_mod_jax") + import numpy as np + + assert np.asarray(module.scale(2, 0, np.array([1.0, 3.0]))).tolist() == [3.0, 7.0] + assert np.asarray(module.scale(2, 2, np.array([1.0, 3.0]))).tolist() == [2.0, 6.0] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"emptyguard_mod{suffix}", None) + + +PDF_KINDS = """\ +module kinds_mod + implicit none + integer, parameter :: I_PLAIN = 1 + integer, parameter :: I_TWICE = 2 + logical, parameter :: L_QUINTIC = .false. +end module kinds_mod +""" + +DISPATCHES_ON_A_KIND = """\ +module dispatch_mod + use kinds_mod, only: I_PLAIN, I_TWICE + use silent_mod, only: complain + implicit none +contains + subroutine pick( n, kind, x, y ) + integer, intent(in) :: n, kind + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + if ( kind == I_TWICE ) then + y = 2.0d0 * x + else if ( kind == I_PLAIN ) then + y = x + else + y = 0.0d0 + call complain( n, y ) + end if + if ( .not. ( kind == I_PLAIN ) ) then + y = y + 1.0d0 + end if + end subroutine pick +end module dispatch_mod +""" + + +def test_a_dispatch_on_a_module_constant_through_its_alias_is_static(tmp_path: Path) -> None: + """``iipdf_type == _mod.IIPDF_ADG1``: pdf_closure picks its PDF by a + constant of model_flags, spelled through the module alias. Static, the + way the bare upper-case spelling is, so the arm the run never takes -- + with its host call -- is never traced.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "kinds_mod.f90").write_text(PDF_KINDS) + (tmp_path / "silent_mod.f90").write_text(SILENT_CHECK) + (tmp_path / "dispatch_mod.f90").write_text(DISPATCHES_ON_A_KIND) + frontend = FortranFrontend(flatten=True) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:dispatch_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "pick" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + emitted = candidate.files[Path("dispatch_mod_jax.py")].decode() + # The dual form: a Python if when the kind is concrete (through the jit + # wrapper it is), the lax.cond only for a traced caller. + assert "if _f_concrete(kind == _kinds_mod.I_TWICE):" in emitted, emitted + # ``.not. ( kind == I_PLAIN )``: the Python form's test is Python's + # ``not``, not jnp.logical_not, which jit stages into a tracer no Python + # if can convert (interpolation's ``.not. l_quintic_poly_interp``). + assert "if not kind == _kinds_mod.I_PLAIN:" in emitted, emitted + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("dispatch_mod_jax") + import numpy as np + + x = np.array([1.0, 3.0]) + assert np.asarray(module.pick(2, 2, x)).tolist() == [3.0, 7.0] + assert np.asarray(module.pick(2, 1, x)).tolist() == [1.0, 3.0] + finally: + sys.path.remove(str(out)) + for name in list(sys.modules): + if name.startswith(("dispatch_mod", "kinds_mod", "silent_mod")): + sys.modules.pop(name, None) + + +NESTED_WHILE = """\ +module nestedwhile_mod + implicit none +contains + subroutine locate( n, m, x, grid, idx ) + integer, intent(in) :: n, m + real(8), intent(in) :: x(n), grid(m) + integer, intent(out) :: idx(n) + integer :: i, k + logical :: calc_done + idx = 0 + do i = 1, n + if ( x(i) > 0.0d0 ) then + k = 1 + calc_done = .false. + do while ( .not. calc_done .and. k <= m ) + if ( grid(k) >= x(i) ) then + idx(i) = k + calc_done = .true. + end if + k = k + 1 + end do + end if + end do + end subroutine locate +end module nestedwhile_mod +""" + + +def test_a_while_inside_a_branch_inside_a_loop_has_its_flag_before_the_branch( + tmp_path: Path, +) -> None: + """interpolation's lin_interp_between_grids: a DO WHILE search under an + IF inside a DO. The while's exit flag is a carry of the enclosing cond, + so it needs a value before the branch -- at the top of the function, + like the goto-region flags -- not only beside its loop.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "nestedwhile_mod.f90").write_text(NESTED_WHILE) + frontend = FortranFrontend(flatten=True) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:nestedwhile_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert candidate.notes["jax"]["kernels"] == ["locate"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("nestedwhile_mod_jax") + import numpy as np + + x = np.array([0.5, -1.0, 2.5, 9.0]) + grid = np.array([1.0, 2.0, 3.0]) + got = np.asarray(module.locate(4, 3, x, grid)).tolist() + assert got == [1, 0, 3, 0] + finally: + sys.path.remove(str(out)) + for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"): + sys.modules.pop(f"nestedwhile_mod{suffix}", None) + + +SILENT_CHECK = """\ +module silent_mod + implicit none +contains + subroutine complain( n, x ) + integer, intent(in) :: n + real(8), intent(in) :: x(n) + if ( any( x < 0.0d0 ) ) print *, "negative", n + end subroutine complain +end module silent_mod +""" + +TRACED_GUARD_CHECK = """\ +module tguarded_mod + use silent_mod, only: complain + implicit none +contains + subroutine step( n, x, y ) + integer, intent(in) :: n + real(8), intent(in) :: x(n) + real(8), intent(out) :: y(n) + y = 2.0d0 * x + if ( y(1) < 0.0d0 ) then + call complain( n, y ) + end if + end subroutine step +end module tguarded_mod +""" + + +def test_a_host_only_call_under_a_traced_guard_carries_nothing(tmp_path: Path) -> None: + """The same check under a guard on the data (``any(err_code == fatal)``): + the branch binds nothing and nothing in it could run under a tracer, so + the lowering carries nothing rather than refusing the kernel.""" + import importlib + import sys + + from recast.transform.jax.tree import TreeToJax + from recast.transform.numpy.tree import TreeConventions + + (tmp_path / "silent_mod.f90").write_text(SILENT_CHECK) + (tmp_path / "tguarded_mod.f90").write_text(TRACED_GUARD_CHECK) + frontend = FortranFrontend(flatten=True) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:tguarded_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)}) + assert "step" in candidate.notes["jax"]["kernels"], candidate.notes["jax"] + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("tguarded_mod_jax") + import numpy as np + + assert np.asarray(module.step(2, np.array([-1.0, 1.0]))).tolist() == [-2.0, 2.0] + finally: + sys.path.remove(str(out)) + for name in list(sys.modules): + if name.startswith(("tguarded_mod", "silent_mod")): + sys.modules.pop(name, None) diff --git a/tests/test_jax_tree.py b/tests/test_jax_tree.py index d8bb8d7..01d7c19 100644 --- a/tests/test_jax_tree.py +++ b/tests/test_jax_tree.py @@ -76,7 +76,7 @@ def test_the_port_emits_the_flat_function_as_a_kernel(tree: Path) -> None: candidate = TreeToJax(CONVENTIONS).apply(unit, facts, {"root": str(tree)}) ported = candidate.files[Path("physics_mod_jax.py")].decode() assert "def _warm_flat_k_impl(" in ported - assert "lax.fori_loop" in ported and ".at[p - 1, ic - 1].set(" in ported + assert "_f_fori(" in ported and ".at[p - 1, ic - 1].set(" in ported assert "_JAX_KERNELS = ['reset_flat', 'warm_flat']" in ported # The originals take the object and stay host-delegated; the flat # signatures reach the ported module's table for the gate. diff --git a/tests/test_loud_refusals.py b/tests/test_loud_refusals.py index 6cc3621..83fc9bf 100644 --- a/tests/test_loud_refusals.py +++ b/tests/test_loud_refusals.py @@ -44,10 +44,10 @@ subroutine prologue_refusals(n, out1, w) integer, intent(in) :: n - real(r8), intent(out) :: out1(max(n, 2)) + real(r8), intent(out) :: out1(mod(n, 3) + 2) type(wide_t), intent(out) :: w integer, parameter :: grid(2, 2) = reshape((/ 1, 2, 3, 4 /), (/ 2, 2 /)) - real(r8) :: scr(max(n, 2)) + real(r8) :: scr(mod(n, 3) + 2) out1 = 0.0_r8 scr = 0.0_r8 w%rows = 0.0_r8 @@ -137,10 +137,10 @@ def test_every_prologue_refusal_is_recorded_and_raises(source: Path, renderer: M prologue = [entry for entry in deferred if entry["block"].startswith("P")] assert [entry["block"] for entry in prologue] == ["P001", "P002", "P003", "P004"] reasons = "\n".join(entry["reason"] for entry in prologue) - assert "out-arg out1: allocation refused (dim expr 'MAX(n, 2)')" in reasons + assert "out-arg out1: allocation refused (dim expr 'MOD(n, 3) + 2')" in reasons assert "out-arg w: INTENT(OUT) derived-type dummy not materialized" in reasons assert "local parameter grid" in reasons - assert "local array scr: extent not resolvable (dim expr 'MAX(n, 2)')" in reasons + assert "local array scr: extent not resolvable (dim expr 'MOD(n, 3) + 2')" in reasons # The old wording is gone, and every refusal raises. assert "allocation skipped" not in body assert "prologue skipped" not in body diff --git a/tests/test_numpy_runtime.py b/tests/test_numpy_runtime.py index 21d045f..2746b4e 100644 --- a/tests/test_numpy_runtime.py +++ b/tests/test_numpy_runtime.py @@ -358,3 +358,24 @@ def test_copy_out_writes_the_overlap_and_leaves_the_rest() -> None: runtime._f_copy_out(same, 7.0) assert list(same) == [7.0, 7.0] runtime._f_copy_out(None, np.ones(2)) # nothing to write into, no error + + +def test_sum_accumulates_in_fortran_element_order() -> None: + """gfortran's inlined SUM is a loop in element order; np.sum pairs its + terms and rounds differently -- CLUBB's vertical_integral drifted 12 ULP. + The sequential helper matches the loop exactly, whole or along an axis.""" + import numpy as np + + from recast.transform.numpy import runtime + + rng = np.random.default_rng(7) + a = np.asfortranarray(rng.uniform(-1e6, 1e6, size=(37, 23))) + loop = np.float64(0) + for x in np.ravel(a, order="F"): + loop = loop + x + assert runtime._f_vsum(a) == loop + along = np.zeros(23) + for i in range(37): + along = along + a[i, :] + assert np.array_equal(runtime._f_vsum(a, axis=0), along) + assert runtime._f_vsum(np.array([1, 2, 3], dtype=np.int32)) == 6 diff --git a/tests/test_numpy_translate.py b/tests/test_numpy_translate.py index 4796ccd..c39757b 100644 --- a/tests/test_numpy_translate.py +++ b/tests/test_numpy_translate.py @@ -291,6 +291,260 @@ def test_logical_operators_on_arrays_are_elementwise(tmp_path: Path) -> None: assert "not l_ok" not in module +SEARCH_LOOP = """\ +module search_mod + implicit none + private + public :: first_above +contains + subroutine first_above( n, z, zmax, k_found, k_scan ) + integer, intent(in) :: n + real, dimension(n), intent(in) :: z + real, intent(in) :: zmax + integer, intent(out) :: k_found, k_scan + integer :: k, kk + do k = 1, n + if ( z(k) > zmax ) exit + end do + k_found = k + do kk = 1, n, 2 + k_scan = kk + end do + ! The bounds of a loop over another variable read kk: still a read. + do k = 1, kk + k_scan = k_scan + 0 + end do + k_scan = kk + end subroutine first_above +end module search_mod +""" + + +def test_a_loop_index_read_after_the_loop_has_the_completion_value(tmp_path: Path) -> None: + """CLUBB's lscale_width_vert_avg searches with ``do k = ...; if (...) + exit; end do`` and integrates up to ``k`` afterwards. On completion + Fortran leaves the index one step past the end -- ``n + 1`` for a unit + step, the first odd value past ``n`` for a step of two -- and after an + EXIT it keeps the exit value. Python's ``for`` leaves the last value.""" + import importlib + import sys + + (tmp_path / "search_mod.f90").write_text(SEARCH_LOOP) + frontend = FortranFrontend() + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:search_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = NumpyTranslation().apply(unit, facts, {"root": tmp_path}) + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + text = (out / "search_mod_numpy.py").read_text() + # The first two loops' indices are read after them (one in a later + # loop's bounds); the last loop's index k is not. + assert text.count("max(0, ") == 2 + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("search_mod_numpy") + import numpy as np + + z = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32) + k_found, k_scan = module.first_above(5, z, np.float32(3.5)) + assert (k_found, k_scan) == (4, 7) # exit at z(4); 1,3,5 then one step past + k_found, k_scan = module.first_above(5, z, np.float32(9.0)) + assert (k_found, k_scan) == (6, 7) # completed: one past n + finally: + sys.path.remove(str(out)) + sys.modules.pop("search_mod_numpy", None) + + +STUBBED_PATH = """ +module lapack_wrap + implicit none +contains + subroutine band_solvex( n, a, x ) + integer, intent(in) :: n + real, intent(inout) :: a(n) + real, intent(out) :: x(n) + x = a + end subroutine band_solvex +end module lapack_wrap +""" + +CHOOSES_A_SOLVER = """ +module solver_mod + use lapack_wrap, only: band_solvex + implicit none +contains + subroutine solve( method, n, m, a, x ) + integer, intent(in) :: method, n, m + real, intent(inout) :: a(n) + real, intent(out) :: x(n) + real :: work(n, max(2, m)) + work = 0.0 + if ( method == 1 ) then + call band_solvex( n, a, x ) + else + x = a + work(:, 1) + end if + end subroutine solve +end module solver_mod +""" + + +def test_a_call_into_a_stubbed_module_raises_on_its_own_line(tmp_path: Path) -> None: + """CLUBB's matrix_solver_wrapper chooses LAPACK or its own LU solver by + a run-time flag; ``lapack_wrap`` is stubbed. With no rule for the call + the whole IF was deferred -- condition and LU branch included -- and the + candidate raised on the path the run takes. The raise belongs to the + statement; the branch around it stays. And ``work(n, max(2, m))`` + (windm's ``rhs``) is a bound Python can spell.""" + import importlib + import sys + + (tmp_path / "lapack_wrap.f90").write_text(STUBBED_PATH) + (tmp_path / "solver_mod.f90").write_text(CHOOSES_A_SOLVER) + frontend = FortranFrontend(stub_modules=["lapack_wrap"]) + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod") + facts = frontend.analyze(unit, tmp_path) + assert facts.interface["stub_procedures"] == ["band_solvex"] + candidate = NumpyTranslation().apply(unit, facts, {"root": tmp_path}) + assert not candidate.deferred + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + text = (out / "solver_mod_numpy.py").read_text() + assert "max(2, m)" in text + raised = "raise NotImplementedError('band_solvex: procedure of a stubbed module, not ported')" + assert raised in text + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("solver_mod_numpy") + import numpy as np + + a = np.array([1.0, 2.0], dtype=np.float32) + _a, x = module.solve(2, 2, 1, a) # INOUT a and OUT x come back + assert x.tolist() == [1.0, 2.0] + with pytest.raises(NotImplementedError): + module.solve(1, 2, 1, a) + finally: + sys.path.remove(str(out)) + sys.modules.pop("solver_mod_numpy", None) + + +HANDS_ON_AN_OPTIONAL = """ +module relay_mod + implicit none +contains + subroutine inner( x, y, rc ) + real, intent(in) :: x + real, intent(out) :: y + real, intent(out), optional :: rc + y = 2.0 * x + if ( present(rc) ) rc = 1.0 / x + end subroutine inner + subroutine outer( x, y, rc, scale ) + real, intent(in) :: x + real, intent(out) :: y + real, intent(out), optional :: rc + real, intent(in), optional :: scale + call inner( x, y, rc = rc ) + if ( present(scale) ) y = y * scale + end subroutine outer +end module relay_mod +""" + + +def test_an_optional_handed_on_carries_its_own_presence(tmp_path: Path) -> None: + """``call inner( x, y, rc = rc )`` where ``rc`` is the caller's own + optional OUT: present in the callee exactly when present in the caller. + Rendered ``want_rc=True`` it was always present, and CLUBB's + xm_wpxp_solve took the LAPACK diagnostic path on every call.""" + import importlib + import sys + + (tmp_path / "relay_mod.f90").write_text(HANDS_ON_AN_OPTIONAL) + frontend = FortranFrontend() + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:relay_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = NumpyTranslation().apply(unit, facts, {"root": tmp_path}) + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + text = (out / "relay_mod_numpy.py").read_text() + assert "want_rc=want_rc" in text + assert "want_rc=True" not in text + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("relay_mod_numpy") + import numpy as np + + y, rc = module.outer(np.float32(4.0)) + assert y == 8.0 and rc != 0.25 # not asked for, not computed + y, rc = module.outer(np.float32(4.0), want_rc=True) + assert y == 8.0 and rc == 0.25 + finally: + sys.path.remove(str(out)) + sys.modules.pop("relay_mod_numpy", None) + + +SCALED_CONSTRUCTOR = """ +module fit_mod + implicit none + integer, parameter :: r8 = selected_real_kind(15) +contains + subroutine polynomial( n, x, y ) + integer, intent(in) :: n + real(r8), intent(in) :: x(n) + real(r8), intent(out) :: y(n) + real(r8), dimension(3), parameter :: & + a = 100._r8 * (/ 6.09868993_r8, 0.499320233_r8, 0.184672631E-01_r8 /) + real(r8), dimension(3), parameter :: b = (/ 1._r8, 2._r8, 3._r8 /) / 2._r8 + integer :: i + do i = 1, n + y(i) = a(1) + a(2) * x(i) + a(3) * x(i)**2 + b(3) + end do + end subroutine polynomial +end module fit_mod +""" + + +def test_a_constant_expression_over_an_array_constructor_is_a_value(tmp_path: Path) -> None: + """CLUBB's saturation and pdf_closure (#26): ``100._core_rknd * (/ ... /)`` + as a local parameter. The token pass rendered a bare constructor and + handed anything around one to the parser, whose literals were never + hoisted; the whole subprogram was a NotImplementedError.""" + import importlib + import sys + + (tmp_path / "fit_mod.f90").write_text(SCALED_CONSTRUCTOR) + frontend = FortranFrontend() + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:fit_mod") + facts = frontend.analyze(unit, tmp_path) + candidate = NumpyTranslation().apply(unit, facts, {"root": tmp_path}) + assert not candidate.deferred + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + text = (out / "fit_mod_numpy.py").read_text() + assert "a = 100. * np.array([6.09868993, 0.499320233, 0.184672631E-01])" in text + assert "b = np.array([1., 2., 3.]) / 2." in text + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("fit_mod_numpy") + import numpy as np + + y = module.polynomial(2, np.array([0.0, 1.0])) + a = 100.0 * np.array([6.09868993, 0.499320233, 0.184672631e-01]) + # To rounding: the point is the constructor, not the summation order. + assert np.allclose(y, [a[0] + 1.5, a[0] + a[1] + a[2] + 1.5], rtol=0, atol=1e-9) + finally: + sys.path.remove(str(out)) + sys.modules.pop("fit_mod_numpy", None) + + CALLBACK = """\ module callback_mod implicit none @@ -363,3 +617,91 @@ def test_the_callback_translation_passes_the_dataflow_gate(callback_candidate) - verdict = ReadWriteSetVerifier().check(unit, candidate, Path("."), LocalExecutor(), {}) assert verdict.confidence is Confidence.SAMPLED, verdict.detail assert verdict.metrics["blocks_matched"] == verdict.metrics["blocks_checked"] > 0 + + +NESTED_GENERICS = """ +module sat_mod + implicit none + private + public :: t_api, rsat_api + interface t_api + module procedure t_k, t_2d + end interface t_api + interface rsat_api + module procedure rsat_k, rsat_2d + end interface rsat_api +contains + function t_k( thl ) result( t ) + real, intent(in) :: thl + real :: t + t = thl + 1.0 + end function t_k + function t_2d( n, thl ) result( t ) + integer, intent(in) :: n + real, intent(in) :: thl(n, 2) + real :: t(n, 2) + t = thl + 1.0 + end function t_2d + function rsat_k( p, t ) result( r ) + real, intent(in) :: p, t + real :: r + r = t / p + end function rsat_k + function rsat_2d( n, p, t ) result( r ) + integer, intent(in) :: n + real, intent(in) :: p(n, 2), t(n, 2) + real :: r(n, 2) + r = t / p + end function rsat_2d +end module sat_mod +""" + +CALLS_NESTED_GENERICS = """ +module core_mod + use sat_mod, only: t_api, rsat_api + implicit none +contains + subroutine step( n, p, thl, rsat ) + integer, intent(in) :: n + real, intent(in) :: p(n, 2), thl(n, 2) + real, intent(out) :: rsat(n, 2) + rsat = rsat_api( n, p, t_api( n, thl ) ) + end subroutine step +end module core_mod +""" + + +def test_a_generic_whose_actual_is_another_generics_result_is_dispatched(tmp_path: Path) -> None: + """CLUBB's advance_clubb_core: ``sat_mixrat_liq_api( ..., thlm2T_in_K_api( + ... ), ... )``. The inner generic's result was ranked scalar without + looking, so the outer one matched no specific and its block was deferred + -- and a deferred block takes the whole subprogram out of the gate.""" + import importlib + import sys + + (tmp_path / "sat_mod.f90").write_text(NESTED_GENERICS) + (tmp_path / "core_mod.f90").write_text(CALLS_NESTED_GENERICS) + frontend = FortranFrontend() + unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:core_mod") + facts = frontend.analyze(unit, tmp_path) + from recast.transform.numpy.tree import TreeTranslation + + candidate = TreeTranslation().apply(unit, facts, {"root": str(tmp_path)}) + assert not candidate.deferred, candidate.deferred + out = tmp_path / "emitted" + out.mkdir() + for path, content in candidate.files.items(): + (out / path.name).write_bytes(content) + text = (out / "core_mod_numpy.py").read_text() + assert "rsat_2d(" in text and "t_2d(" in text + sys.path.insert(0, str(out)) + try: + module = importlib.import_module("core_mod_numpy") + import numpy as np + + p = np.full((2, 2), 2.0, dtype=np.float32, order="F") + thl = np.ones((2, 2), dtype=np.float32, order="F") + assert module.step(2, p, thl).tolist() == [[1.0, 1.0], [1.0, 1.0]] + finally: + sys.path.remove(str(out)) + sys.modules.pop("core_mod_numpy", None)