• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

daisytuner / docc / 22851518700

09 Mar 2026 11:34AM UTC coverage: 64.448% (+0.03%) from 64.415%
22851518700

Pull #556

github

web-flow
Merge 523df0a24 into 64dd02640
Pull Request #556: [MLIR] add support for linalg broadcast

24493 of 38004 relevant lines covered (64.45%)

386.53 hits per line

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

73.43
/python/docc/python/python_program.py
1
import inspect
4✔
2
import shutil
4✔
3
import textwrap
4✔
4
import ast
4✔
5
import os
4✔
6
import getpass
4✔
7
import hashlib
4✔
8
import ml_dtypes
4✔
9
import numpy as np
4✔
10
from typing import Annotated, get_origin, get_args, Any, Optional
4✔
11

12
from docc.sdfg import (
4✔
13
    Scalar,
14
    PrimitiveType,
15
    Pointer,
16
    Structure,
17
    Array,
18
    Type,
19
    Tensor,
20
    StructuredSDFG,
21
    StructuredSDFGBuilder,
22
)
23
from docc.compiler.docc_program import DoccProgram
4✔
24
from docc.compiler.compiled_sdfg import CompiledSDFG
4✔
25
from docc.python.ast_parser import ASTParser
4✔
26
from docc.python.types import element_type_from_sdfg_type
4✔
27
from docc.python.target_registry import get_target_schedule_fn, get_target_compile_fn
4✔
28

29

30
def _compile_wrapper(self, output_folder=None):
4✔
31
    """Wrapper to allow StructuredSDFG.compile() to return a CompiledSDFG."""
32
    lib_path = self._compile(output_folder)
×
33
    return CompiledSDFG(lib_path, self)
×
34

35

36
# Monkey-patch StructuredSDFG to add compile method
37
StructuredSDFG.compile = _compile_wrapper
4✔
38

39

40
def _map_python_type(dtype):
4✔
41
    """Map Python/numpy types to SDFG types."""
42
    # If it is already a sdfg Type, return it
43
    if isinstance(dtype, Type):
4✔
44
        return dtype
×
45

46
    # Handle Annotated for Arrays
47
    if get_origin(dtype) is Annotated:
4✔
48
        args = get_args(dtype)
4✔
49
        base_type = args[0]
4✔
50
        metadata = args[1:]
4✔
51

52
        if base_type is np.ndarray:
4✔
53
            # Convention: Annotated[np.ndarray, shape, dtype]
54
            shape = metadata[0]
4✔
55
            elem_type = Scalar(PrimitiveType.Double)  # Default
4✔
56

57
            if len(metadata) > 1:
4✔
58
                possible_dtype = metadata[1]
4✔
59
                elem_type = _map_python_type(possible_dtype)
4✔
60

61
            return Pointer(elem_type)
4✔
62

63
    # Handle numpy.ndarray[Shape, DType]
64
    if get_origin(dtype) is np.ndarray:
4✔
65
        args = get_args(dtype)
×
66
        # args[0] is shape, args[1] is dtype
67
        if len(args) >= 2:
×
68
            elem_type = _map_python_type(args[1])
×
69
            return Pointer(elem_type)
×
70

71
    # Simple mapping for python types
72
    if dtype is float or dtype is np.float64:
4✔
73
        return Scalar(PrimitiveType.Double)
4✔
74
    elif dtype is np.float32:
4✔
75
        return Scalar(PrimitiveType.Float)
×
76
    elif dtype is bool or dtype is np.bool_:
4✔
77
        return Scalar(PrimitiveType.Bool)
4✔
78
    elif dtype is int or dtype is np.int64:
4✔
79
        return Scalar(PrimitiveType.Int64)
4✔
80
    elif dtype is np.int32:
4✔
81
        return Scalar(PrimitiveType.Int32)
4✔
82
    elif dtype is np.int16:
×
83
        return Scalar(PrimitiveType.Int16)
×
84
    elif dtype is np.int8:
×
85
        return Scalar(PrimitiveType.Int8)
×
86
    elif dtype is np.uint64:
×
87
        return Scalar(PrimitiveType.UInt64)
×
88
    elif dtype is np.uint32:
×
89
        return Scalar(PrimitiveType.UInt32)
×
90
    elif dtype is np.uint16:
×
91
        return Scalar(PrimitiveType.UInt16)
×
92
    elif dtype is np.uint8:
×
93
        return Scalar(PrimitiveType.UInt8)
×
94

95
    # Handle Python classes - map to Structure type
96
    if inspect.isclass(dtype):
×
97
        # Use the class name as the structure name
98
        return Pointer(Structure(dtype.__name__))
×
99

100
    return dtype
×
101

102

103
class PythonProgram(DoccProgram):
4✔
104

105
    def __init__(
4✔
106
        self,
107
        func,
108
        target: str = "none",
109
        category: str = "server",
110
        instrumentation_mode: Optional[str] = None,
111
        capture_args: Optional[bool] = None,
112
        remote_tuning: bool = False,
113
    ):
114
        super().__init__(
4✔
115
            name=func.__name__,
116
            target=target,
117
            category=category,
118
            instrumentation_mode=instrumentation_mode,
119
            capture_args=capture_args,
120
            remote_tuning=remote_tuning,
121
        )
122
        self.func = func
4✔
123
        self._last_structure_member_info = {}
4✔
124

125
    def __call__(self, *args: Any) -> Any:
4✔
126
        # JIT compile and run
127
        compiled = self.compile(*args)
4✔
128
        res = compiled(*args)
4✔
129

130
        # Handle return value conversion based on annotation
131
        sig = inspect.signature(self.func)
4✔
132
        ret_annotation = sig.return_annotation
4✔
133

134
        if ret_annotation is not inspect.Signature.empty:
4✔
135
            if get_origin(ret_annotation) is Annotated:
4✔
136
                type_args = get_args(ret_annotation)
4✔
137
                if len(type_args) >= 1 and type_args[0] is np.ndarray:
4✔
138
                    shape = None
4✔
139
                    if len(type_args) >= 2:
4✔
140
                        shape = type_args[1]
4✔
141

142
                    if shape is not None:
4✔
143
                        try:
4✔
144
                            return np.ctypeslib.as_array(res, shape=shape)
4✔
145
                        except Exception:
×
146
                            pass
×
147

148
        # Try to infer return shape from metadata
149
        if hasattr(compiled, "get_return_shape"):
4✔
150
            shape = compiled.get_return_shape(*args)
4✔
151
            if shape is not None:
4✔
152
                try:
×
153
                    return np.ctypeslib.as_array(res, shape=shape)
×
154
                except Exception:
×
155
                    pass
×
156

157
        return res
4✔
158

159
    def compile(
4✔
160
        self,
161
        *args: Any,
162
        output_folder: Optional[str] = None,
163
        instrumentation_mode: Optional[str] = None,
164
        capture_args: Optional[bool] = None,
165
    ) -> CompiledSDFG:
166
        original_output_folder = output_folder
4✔
167

168
        # Resolve options
169
        if instrumentation_mode is None:
4✔
170
            instrumentation_mode = self.instrumentation_mode
4✔
171
        if capture_args is None:
4✔
172
            capture_args = self.capture_args
4✔
173

174
        # Check environment variable DOCC_CI
175
        docc_ci = os.environ.get("DOCC_CI", "")
4✔
176
        if docc_ci:
4✔
177
            if docc_ci == "regions":
×
178
                if instrumentation_mode is None:
×
179
                    instrumentation_mode = "ols"
×
180
            elif docc_ci == "arg-capture":
×
181
                if capture_args is None:
×
182
                    capture_args = True
×
183
            else:
184
                # Full mode (or unknown value treated as full)
185
                if instrumentation_mode is None:
×
186
                    instrumentation_mode = "ols"
×
187
                if capture_args is None:
×
188
                    capture_args = True
×
189

190
        # Defaults
191
        if instrumentation_mode is None:
4✔
192
            instrumentation_mode = ""
4✔
193
        if capture_args is None:
4✔
194
            capture_args = False
4✔
195

196
        # 1. Analyze arguments and shapes
197
        arg_types = []
4✔
198
        shape_values = []  # List of unique shape values found
4✔
199
        shape_sources = []  # List of (arg_idx, dim_idx) for each unique shape value
4✔
200

201
        # Mapping from (arg_idx, dim_idx) -> unique_shape_idx
202
        arg_shape_mapping = {}
4✔
203

204
        # First pass: collect scalar integer arguments and their values
205
        sig = inspect.signature(self.func)
4✔
206
        params = list(sig.parameters.items())
4✔
207
        scalar_int_params = {}  # Maps value -> parameter name (first one wins)
4✔
208
        for i, ((name, param), arg) in enumerate(zip(params, args)):
4✔
209
            if isinstance(arg, (int, np.integer)) and not isinstance(
4✔
210
                arg, (bool, np.bool_)
211
            ):
212
                val = int(arg)
4✔
213
                if val not in scalar_int_params:
4✔
214
                    scalar_int_params[val] = name
4✔
215

216
        for i, arg in enumerate(args):
4✔
217
            t = self._infer_type(arg)
4✔
218
            arg_types.append(t)
4✔
219

220
            if isinstance(arg, np.ndarray):
4✔
221
                for dim_idx, dim_val in enumerate(arg.shape):
4✔
222
                    # Check if we've seen this value
223
                    if dim_val in shape_values:
4✔
224
                        # Reuse
225
                        u_idx = shape_values.index(dim_val)
4✔
226
                    else:
227
                        # New
228
                        u_idx = len(shape_values)
4✔
229
                        shape_values.append(dim_val)
4✔
230
                        shape_sources.append((i, dim_idx))
4✔
231

232
                    arg_shape_mapping[(i, dim_idx)] = u_idx
4✔
233

234
        # Detect scalar-shape equivalences: which shape indices have a matching scalar param
235
        # Maps unique_shape_idx -> scalar parameter name
236
        shape_to_scalar = {}
4✔
237
        for s_idx, s_val in enumerate(shape_values):
4✔
238
            if s_val in scalar_int_params:
4✔
239
                shape_to_scalar[s_idx] = scalar_int_params[s_val]
4✔
240

241
        # 2. Signature - include scalar-shape equivalences for correct caching
242
        mapping_sig = sorted(arg_shape_mapping.items())
4✔
243
        equiv_sig = sorted(shape_to_scalar.items())
4✔
244
        type_sig = ", ".join(self._type_to_str(t) for t in arg_types)
4✔
245
        signature = f"{type_sig}|{mapping_sig}|{equiv_sig}"
4✔
246

247
        if output_folder is None:
4✔
248
            filename = inspect.getsourcefile(self.func)
4✔
249
            hash_input = f"{filename}|{self.name}|{self.target}|{self.category}|{self.capture_args}|{self.instrumentation_mode}|{self.remote_tuning}|{signature}".encode(
4✔
250
                "utf-8"
251
            )
252
            stable_id = hashlib.sha256(hash_input).hexdigest()[:16]
4✔
253

254
            docc_tmp = os.environ.get("DOCC_TMP")
4✔
255
            if docc_tmp:
4✔
256
                output_folder = f"{docc_tmp}/{self.name}-{stable_id}"
×
257
            else:
258
                user = os.getenv("USER")
4✔
259
                if not user:
4✔
260
                    user = getpass.getuser()
4✔
261
                output_folder = f"/tmp/{user}/DOCC/{self.name}-{stable_id}"
4✔
262

263
        if original_output_folder is None and signature in self.cache:
4✔
264
            return self.cache[signature]
4✔
265

266
        # 3. Build SDFG
267
        if os.path.exists(output_folder):
4✔
268
            # Multiple python processes running the same code?
269
            shutil.rmtree(output_folder)
4✔
270
        sdfg, out_args, out_shapes, out_strides = self._build_sdfg(
4✔
271
            arg_types, args, arg_shape_mapping, shape_values, shape_to_scalar
272
        )
273
        sdfg.validate()
4✔
274

275
        # Tensor targets keep tensor nodes
276
        if self.target != "onnx":
4✔
277
            sdfg.expand()
4✔
278

279
        # Simplify pipelines
280
        sdfg.simplify()
4✔
281

282
        # Normalization for scheduling
283
        if self.target != "none":
4✔
284
            sdfg.normalize()
4✔
285

286
        sdfg.dump(output_folder)
4✔
287

288
        # Schedule if target is specified
289
        if self.target != "none":
4✔
290
            # Check for custom registered target first
291
            custom_schedule_fn = get_target_schedule_fn(self.target)
4✔
292
            if custom_schedule_fn is not None:
4✔
293
                custom_schedule_fn(
4✔
294
                    sdfg, self.category, {"remote_tuning": self.remote_tuning}
295
                )
296
            else:
297
                sdfg.schedule(self.target, self.category, self.remote_tuning)
4✔
298

299
        self.last_sdfg = sdfg
4✔
300

301
        sdfg.dump(output_folder, "post_sched")
4✔
302

303
        custom_compile_fn = get_target_compile_fn(self.target)
4✔
304
        if custom_compile_fn is not None:
4✔
305
            lib_path = custom_compile_fn(
4✔
306
                sdfg, output_folder, instrumentation_mode, capture_args, {}
307
            )
308
        else:
309
            lib_path = sdfg._compile(
4✔
310
                output_folder=output_folder,
311
                target=self.target,
312
                instrumentation_mode=instrumentation_mode,
313
                capture_args=capture_args,
314
            )
315

316
        # Build ONNX model from JSON if target is onnx (after _compile creates the JSON)
317
        if self.target == "onnx":
4✔
318
            from docc.python.targets.onnx_model_builder import convert_json_to_onnx
×
319

320
            onnx_model_path = convert_json_to_onnx(output_folder)
×
321
            if onnx_model_path:
×
322
                print(f"Generated ONNX models: {onnx_model_path}")
×
323

324
        # 5. Create CompiledSDFG
325
        compiled = CompiledSDFG(
4✔
326
            lib_path,
327
            sdfg,
328
            shape_sources,
329
            self._last_structure_member_info,
330
            out_args,
331
            out_shapes,
332
            out_strides,
333
        )
334

335
        # Cache if using default output folder
336
        if original_output_folder is None:
4✔
337
            self.cache[signature] = compiled
4✔
338

339
        return compiled
4✔
340

341
    def to_sdfg(self, *args: Any) -> StructuredSDFG:
4✔
342
        arg_types = [self._infer_type(arg) for arg in args]
×
343

344
        # Build shape mapping
345
        shape_values = []
×
346
        shape_sources = []
×
347
        arg_shape_mapping = {}
×
348

349
        sig = inspect.signature(self.func)
×
350
        params = list(sig.parameters.items())
×
351
        scalar_int_params = {}
×
352
        for i, ((name, param), arg) in enumerate(zip(params, args)):
×
353
            if isinstance(arg, (int, np.integer)) and not isinstance(
×
354
                arg, (bool, np.bool_)
355
            ):
356
                val = int(arg)
×
357
                if val not in scalar_int_params:
×
358
                    scalar_int_params[val] = name
×
359

360
        for i, arg in enumerate(args):
×
361
            if isinstance(arg, np.ndarray):
×
362
                for dim_idx, dim_val in enumerate(arg.shape):
×
363
                    if dim_val in shape_values:
×
364
                        u_idx = shape_values.index(dim_val)
×
365
                    else:
366
                        u_idx = len(shape_values)
×
367
                        shape_values.append(dim_val)
×
368
                        shape_sources.append((i, dim_idx))
×
369
                    arg_shape_mapping[(i, dim_idx)] = u_idx
×
370

371
        shape_to_scalar = {}
×
372
        for s_idx, s_val in enumerate(shape_values):
×
373
            if s_val in scalar_int_params:
×
374
                shape_to_scalar[s_idx] = scalar_int_params[s_val]
×
375

376
        sdfg, _, _, _ = self._build_sdfg(
×
377
            arg_types, args, arg_shape_mapping, shape_values, shape_to_scalar
378
        )
379
        return sdfg
×
380

381
    def _convert_inputs(self, args: tuple) -> tuple:
4✔
382
        return args
×
383

384
    def _convert_outputs(self, result: Any, original_args: tuple) -> Any:
4✔
385
        return result
×
386

387
    def _get_signature(self, arg_types):
4✔
388
        return ", ".join(self._type_to_str(t) for t in arg_types)
×
389

390
    def _type_to_str(self, t):
4✔
391
        if isinstance(t, Scalar):
4✔
392
            return f"Scalar({t.primitive_type})"
4✔
393
        elif isinstance(t, Array):
4✔
394
            return f"Array({self._type_to_str(t.element_type)}, {t.num_elements})"
×
395
        elif isinstance(t, Pointer):
4✔
396
            return f"Pointer({self._type_to_str(t.pointee_type)})"
4✔
397
        elif isinstance(t, Structure):
4✔
398
            return f"Structure({t.name})"
4✔
399
        return str(t)
×
400

401
    def _infer_type(self, arg):
4✔
402
        if isinstance(arg, (float, np.float64)):
4✔
403
            return Scalar(PrimitiveType.Double)
4✔
404
        elif isinstance(arg, np.float32):
4✔
405
            return Scalar(PrimitiveType.Float)
4✔
406
        elif isinstance(arg, (bool, np.bool_)):
4✔
407
            return Scalar(PrimitiveType.Bool)
4✔
408
        elif isinstance(arg, (int, np.int64)):
4✔
409
            return Scalar(PrimitiveType.Int64)
4✔
410
        elif isinstance(arg, np.int32):
4✔
411
            return Scalar(PrimitiveType.Int32)
4✔
412
        elif isinstance(arg, np.int16):
4✔
413
            return Scalar(PrimitiveType.Int16)
×
414
        elif isinstance(arg, np.int8):
4✔
415
            return Scalar(PrimitiveType.Int8)
×
416
        elif isinstance(arg, np.uint64):
4✔
417
            return Scalar(PrimitiveType.UInt64)
×
418
        elif isinstance(arg, np.uint32):
4✔
419
            return Scalar(PrimitiveType.UInt32)
×
420
        elif isinstance(arg, np.uint16):
4✔
421
            return Scalar(PrimitiveType.UInt16)
×
422
        elif isinstance(arg, np.uint8):
4✔
423
            return Scalar(PrimitiveType.UInt8)
×
424
        elif isinstance(arg, np.ndarray):
4✔
425
            # Map dtype
426
            if arg.dtype == np.float64:
4✔
427
                elem_type = Scalar(PrimitiveType.Double)
4✔
428
            elif arg.dtype == np.float32:
4✔
429
                elem_type = Scalar(PrimitiveType.Float)
4✔
430
            elif arg.dtype == np.float16:
4✔
431
                elem_type = Scalar(PrimitiveType.Half)
×
432
            elif arg.dtype == ml_dtypes.bfloat16:
4✔
433
                elem_type = Scalar(PrimitiveType.BFloat)
4✔
434
            elif arg.dtype == np.bool_:
4✔
435
                elem_type = Scalar(PrimitiveType.Bool)
4✔
436
            elif arg.dtype == np.int64:
4✔
437
                elem_type = Scalar(PrimitiveType.Int64)
4✔
438
            elif arg.dtype == np.int32:
4✔
439
                elem_type = Scalar(PrimitiveType.Int32)
4✔
440
            elif arg.dtype == np.int16:
×
441
                elem_type = Scalar(PrimitiveType.Int16)
×
442
            elif arg.dtype == np.int8:
×
443
                elem_type = Scalar(PrimitiveType.Int8)
×
444
            elif arg.dtype == np.uint64:
×
445
                elem_type = Scalar(PrimitiveType.UInt64)
×
446
            elif arg.dtype == np.uint32:
×
447
                elem_type = Scalar(PrimitiveType.UInt32)
×
448
            elif arg.dtype == np.uint16:
×
449
                elem_type = Scalar(PrimitiveType.UInt16)
×
450
            elif arg.dtype == np.uint8:
×
451
                elem_type = Scalar(PrimitiveType.UInt8)
×
452
            else:
453
                raise ValueError(f"Unsupported numpy dtype: {arg.dtype}")
×
454

455
            return Pointer(elem_type)
4✔
456
        elif isinstance(arg, str):
4✔
457
            # Explicitly reject strings - they are not supported
458
            raise ValueError(f"Unsupported argument type: {type(arg)}")
4✔
459
        else:
460
            # Check if it's a class instance
461
            if hasattr(arg, "__class__") and not isinstance(arg, type):
4✔
462
                # It's an instance of a class, return pointer to Structure
463
                return Pointer(Structure(arg.__class__.__name__))
4✔
464
            raise ValueError(f"Unsupported argument type: {type(arg)}")
×
465

466
    def _build_sdfg(
4✔
467
        self,
468
        arg_types,
469
        args,
470
        arg_shape_mapping,
471
        shape_values,
472
        shape_to_scalar=None,
473
    ):
474
        if shape_to_scalar is None:
4✔
475
            shape_to_scalar = {}
×
476
        sig = inspect.signature(self.func)
4✔
477

478
        # Handle return type - always void for SDFG, output args used for returns
479
        return_type = Scalar(PrimitiveType.Void)
4✔
480
        infer_return_type = True
4✔
481

482
        # Parse return annotation to determine output arguments if possible
483
        explicit_returns = []
4✔
484
        if sig.return_annotation is not inspect.Signature.empty:
4✔
485
            infer_return_type = False
4✔
486

487
            # Helper to normalize annotation to list of types
488
            def normalize_annotation(ann):
4✔
489
                # Handle Tuple[type, ...]
490
                origin = get_origin(ann)
4✔
491
                if origin is tuple:
4✔
492
                    type_args = get_args(ann)
×
493
                    # Tuple[()] or Tuple w/o args
494
                    if not type_args:
×
495
                        return []
×
496
                    # Tuple[int, float]
497
                    if len(type_args) > 0 and type_args[-1] is not Ellipsis:
×
498
                        return [_map_python_type(t) for t in type_args]
×
499
                    # Tuple[int, ...] - not supported for fixed number of returns yet?
500
                    # For now assume fixed tuple
501
                    return [_map_python_type(t) for t in type_args]
×
502
                else:
503
                    return [_map_python_type(ann)]
4✔
504

505
            explicit_returns = normalize_annotation(sig.return_annotation)
4✔
506
            for rt in explicit_returns:
4✔
507
                if not isinstance(rt, Type):
4✔
508
                    # Fallback if map failed (e.g. invalid annotation)
509
                    infer_return_type = True
×
510
                    explicit_returns = []
×
511
                    break
×
512

513
        builder = StructuredSDFGBuilder(f"{self.name}_sdfg", return_type)
4✔
514

515
        # Add pre-defined return arguments if we know them
516
        if not infer_return_type:
4✔
517
            for i, dtype in enumerate(explicit_returns):
4✔
518
                # Scalar -> Pointer(Scalar)
519
                # Array -> Already Pointer(Scalar). Keep it.
520
                arg_type = dtype
4✔
521
                if isinstance(dtype, Scalar):
4✔
522
                    arg_type = Pointer(dtype)
4✔
523

524
                builder.add_container(f"_docc_ret_{i}", arg_type, is_argument=True)
4✔
525

526
        # Register structure types for any class arguments
527
        # Also track member name to index mapping for each structure
528
        structures_to_register = {}
4✔
529
        structure_member_info = {}  # Maps struct_name -> {member_name: (index, type)}
4✔
530
        for i, (arg, dtype) in enumerate(zip(args, arg_types)):
4✔
531
            if isinstance(dtype, Pointer) and dtype.has_pointee_type():
4✔
532
                pointee = dtype.pointee_type
4✔
533
                if isinstance(pointee, Structure):
4✔
534
                    struct_name = pointee.name
4✔
535
                    if struct_name not in structures_to_register:
4✔
536
                        # Get class from arg to introspect members
537
                        if hasattr(arg, "__dict__"):
4✔
538
                            # Use __dict__ to get only instance attributes
539
                            # Sort by name to ensure consistent ordering
540
                            # Note: This alphabetical ordering is used to define the
541
                            # structure layout and must match the order expected by
542
                            # the backend code generation
543
                            member_types = []
4✔
544
                            member_names = []
4✔
545
                            for attr_name, attr_value in sorted(arg.__dict__.items()):
4✔
546
                                if not attr_name.startswith("_"):
4✔
547
                                    # Infer member type from instance attribute
548
                                    # Check bool before int since bool is subclass of int
549
                                    member_type = None
4✔
550
                                    if isinstance(attr_value, bool):
4✔
551
                                        member_type = Scalar(PrimitiveType.Bool)
×
552
                                    elif isinstance(attr_value, (int, np.int64)):
4✔
553
                                        member_type = Scalar(PrimitiveType.Int64)
×
554
                                    elif isinstance(attr_value, (float, np.float64)):
4✔
555
                                        member_type = Scalar(PrimitiveType.Double)
4✔
556
                                    elif isinstance(attr_value, np.int32):
×
557
                                        member_type = Scalar(PrimitiveType.Int32)
×
558
                                    elif isinstance(attr_value, np.float32):
×
559
                                        member_type = Scalar(PrimitiveType.Float)
×
560
                                    # TODO: Consider using np.integer and np.floating abstract types
561
                                    # for more comprehensive numpy type coverage
562
                                    # TODO: Add support for nested structures and arrays
563

564
                                    if member_type is not None:
4✔
565
                                        member_types.append(member_type)
4✔
566
                                        member_names.append(attr_name)
4✔
567

568
                            if member_types:
4✔
569
                                structures_to_register[struct_name] = member_types
4✔
570
                                # Build member name to (index, type) mapping
571
                                structure_member_info[struct_name] = {
4✔
572
                                    name: (idx, mtype)
573
                                    for idx, (name, mtype) in enumerate(
574
                                        zip(member_names, member_types)
575
                                    )
576
                                }
577

578
        # Store structure_member_info for later use in CompiledSDFG
579
        self._last_structure_member_info = structure_member_info
4✔
580

581
        # Register all discovered structures with the builder
582
        for struct_name, member_types in structures_to_register.items():
4✔
583
            builder.add_structure(struct_name, member_types)
4✔
584

585
        # Register arguments
586
        params = list(sig.parameters.items())
4✔
587
        if len(params) != len(arg_types):
4✔
588
            raise ValueError(
×
589
                f"Argument count mismatch: expected {len(params)}, got {len(arg_types)}"
590
            )
591

592
        # Add regular arguments
593
        tensor_table = {}
4✔
594
        for i, ((name, param), dtype, arg) in enumerate(zip(params, arg_types, args)):
4✔
595
            builder.add_container(name, dtype, is_argument=True)
4✔
596

597
            # Store layout information for arrays
598
            if isinstance(arg, np.ndarray):
4✔
599
                element_type = element_type_from_sdfg_type(dtype)
4✔
600

601
                shapes = []
4✔
602
                for dim_idx in range(arg.ndim):
4✔
603
                    dim_val = arg.shape[dim_idx]
4✔
604
                    if dim_val == 1:
4✔
605
                        # Always use literal "1" for size-1 dimensions to enable
606
                        # proper broadcasting detection
607
                        shapes.append("1")
4✔
608
                    else:
609
                        u_idx = arg_shape_mapping[(i, dim_idx)]
4✔
610
                        if u_idx in shape_to_scalar:
4✔
611
                            shapes.append(shape_to_scalar[u_idx])
4✔
612
                        else:
613
                            shapes.append(f"_s{u_idx}")
4✔
614

615
                strides = []
4✔
616
                if arg.flags["C_CONTIGUOUS"]:
4✔
617
                    # Row-major: stride[i] = product of shapes[i+1:]
618
                    for dim_idx in range(arg.ndim):
4✔
619
                        if dim_idx == arg.ndim - 1:
4✔
620
                            strides.append("1")
4✔
621
                        else:
622
                            suffix_shapes = shapes[dim_idx + 1 :]
4✔
623
                            if len(suffix_shapes) == 1:
4✔
624
                                strides.append(suffix_shapes[0])
4✔
625
                            else:
626
                                strides.append("(" + " * ".join(suffix_shapes) + ")")
4✔
627
                elif arg.flags["F_CONTIGUOUS"]:
4✔
628
                    # Column-major: stride[i] = product of shapes[:i]
629
                    for dim_idx in range(arg.ndim):
4✔
630
                        if dim_idx == 0:
4✔
631
                            strides.append("1")
4✔
632
                        else:
633
                            prefix_shapes = shapes[:dim_idx]
4✔
634
                            if len(prefix_shapes) == 1:
4✔
635
                                strides.append(prefix_shapes[0])
4✔
636
                            else:
637
                                strides.append("(" + " * ".join(prefix_shapes) + ")")
4✔
638
                else:
639
                    # Non-contiguous: use actual stride values
640
                    for dim_idx in range(arg.ndim):
4✔
641
                        stride_val = arg.strides[dim_idx] // arg.itemsize
4✔
642
                        strides.append(f"{stride_val}")
4✔
643

644
                offset = "0"
4✔
645
                tensor_table[name] = Tensor(element_type, shapes, strides, offset)
4✔
646

647
        # Add unified shape arguments only for shapes without scalar equivalents
648
        # and skip size-1 dimensions (they use literal "1" instead)
649
        for i in range(len(shape_values)):
4✔
650
            if i not in shape_to_scalar and shape_values[i] != 1:
4✔
651
                builder.add_container(
4✔
652
                    f"_s{i}", Scalar(PrimitiveType.Int64), is_argument=True
653
                )
654

655
        # Create symbol table for parser
656
        container_table = {}
4✔
657
        for i, ((name, param), dtype, arg) in enumerate(zip(params, arg_types, args)):
4✔
658
            container_table[name] = dtype
4✔
659

660
        for i in range(len(shape_values)):
4✔
661
            if i not in shape_to_scalar and shape_values[i] != 1:
4✔
662
                container_table[f"_s{i}"] = Scalar(PrimitiveType.Int64)
4✔
663

664
        # Parse AST
665
        source_lines, start_line = inspect.getsourcelines(self.func)
4✔
666
        source = textwrap.dedent("".join(source_lines))
4✔
667
        tree = ast.parse(source)
4✔
668
        ast.increment_lineno(tree, start_line - 1)
4✔
669
        func_def = tree.body[0]
4✔
670

671
        filename = inspect.getsourcefile(self.func)
4✔
672
        function_name = self.func.__name__
4✔
673

674
        # Combine globals with closure variables (closure takes precedence)
675
        combined_globals = dict(self.func.__globals__)
4✔
676
        if self.func.__closure__ is not None and self.func.__code__.co_freevars:
4✔
677
            for name, cell in zip(
4✔
678
                self.func.__code__.co_freevars, self.func.__closure__
679
            ):
680
                combined_globals[name] = cell.cell_contents
4✔
681

682
        parser = ASTParser(
4✔
683
            builder,
684
            tensor_table,
685
            container_table,
686
            filename,
687
            function_name,
688
            infer_return_type=infer_return_type,
689
            globals_dict=combined_globals,
690
            structure_member_info=structure_member_info,
691
        )
692
        for node in func_def.body:
4✔
693
            parser.visit(node)
4✔
694

695
        # Emit hoisted allocations at function entry
696
        parser.memory_handler.emit_allocations()
4✔
697

698
        sdfg = builder.move()
4✔
699
        # Mark return arguments metadata
700
        out_args = []
4✔
701
        for name in sdfg.arguments:
4✔
702
            if name.startswith("_docc_ret_"):
4✔
703
                out_args.append(name)
4✔
704

705
        return (
4✔
706
            sdfg,
707
            out_args,
708
            parser.captured_return_shapes,
709
            parser.captured_return_strides,
710
        )
711

712

713
def native(
4✔
714
    func=None,
715
    *,
716
    target="none",
717
    category="server",
718
    instrumentation_mode=None,
719
    capture_args=None,
720
    remote_tuning=False,
721
):
722
    """Decorator to create a PythonProgram from a Python function.
723

724
    Example:
725
        @native
726
        def my_function(x: np.ndarray) -> np.ndarray:
727
            return x * 2
728

729
        result = my_function(np.array([1.0, 2.0, 3.0]))
730
    """
731
    if func is None:
4✔
732
        return lambda f: PythonProgram(
4✔
733
            f,
734
            target=target,
735
            category=category,
736
            instrumentation_mode=instrumentation_mode,
737
            capture_args=capture_args,
738
            remote_tuning=remote_tuning,
739
        )
740
    return PythonProgram(
4✔
741
        func,
742
        target=target,
743
        category=category,
744
        instrumentation_mode=instrumentation_mode,
745
        capture_args=capture_args,
746
        remote_tuning=remote_tuning,
747
    )
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc