Add some features to operations interface

src/xtc/cli/explore.py

@@ -99,6 +99,31 @@ def xtc_conv2d_graph(
     return gb.graph
 
 
+def xtc_conv2d_nhwc_frsc_graph(
+    n: int,
+    h: int,
+    w: int,
+    f: int,
+    r: int,
+    s: int,
+    c: int,
+    SH: int,
+    SW: int,
+    ftype: str,
+    name: str = "matmul",
+) -> Graph:
+    import xtc.graphs.xtc.op as O
+
+    dtype = DTYPES_MAP[ftype]
+    a = O.tensor((n, h * SH + r - 1, w * SW + s - 1, c), dtype, name="A")
+    b = O.tensor((f, r, s, c), dtype, name="B")
+    with O.graph(name=name) as gb:
+        tp = O.transpose(b, axes=(1, 2, 3, 0), name="T")
+        O.conv2d(a, tp, stride=(SH, SW), name="O")
+    graph = gb.graph
+    return graph
+
+
 def tvm_impl(graph: Graph) -> tuple[Backend, str]:
     from xtc.backends.tvm import Backend
 
@@ -210,12 +235,12 @@ def compile_one(
             )
         )
     if args.save_temps:
-        compile_args.update(
-            dict(
-                save_temps=True,
-                save_temps_dir=f"{args.save_temps_dir}/{ident}",
-            )
-        )
+        tmp_path = Path(args.save_temps_dir) / ident
+        tmp_path.mkdir(parents=True, exist_ok=True)
+        compile_args.update(dict(save_temps=True, save_temps_dir=str(tmp_path)))
+        with open(tmp_path / f"{ident}_graph.txt", "w") as outf:
+            print(graph, file=outf)
+
     assert args.eval == "eval"
     compiler = impl.get_compiler(**compile_args)
     module = compiler.compile(schedule=schedule)
@@ -252,7 +277,7 @@ def load_and_evaluate_sample(
     results, code, error_msg = evaluate()
     if code == 0:
         time = min(results)
-        logger.debug("  Evaluated: %s: %s: time: %.2f msecs", ident, in_x, time * 1000)
+        logger.debug("  Evaluated: %s: %s: time: %.3f msecs", ident, in_x, time * 1000)
     else:
         time = 0
         logger.error("Error evaluating: %s: %s: %d: %s", ident, in_x, code, error_msg)
@@ -524,7 +549,6 @@ def peak_time(args: NS) -> float:
     assert flops != 0, f"unable to evaluate machine flops for type {dtype}"
     dims_names = OPERATORS[args.operator]["dims"]
     dims_map = {k: v for k, v in zip(dims_names, args.dims)}
-
     flop = mulall([d for k, d in dims_map.items() if k.lower() == k])
     time = flop / flops / args.threads
     return time
@@ -634,22 +658,23 @@ def optimize(args: NS):
         }
         evaluate_sample(strategy, schedule, graph, args, callbacks=callbacks)
         for row in result_callback._rows:
-            in_x, time, peak, backend = row[-4:]
+            in_x, time, peak, backend = row[-4]
             tqdm.write(
-                f"Schedule: {backend}: {in_x}: time: {time * 1000:.2f} msecs, peak perf: {peak * 100:.2f}%"
+                f"Schedule: {backend}: {in_x}: time: {time * 1000:.3f} msecs, peak perf: {peak * 100:.2f}%"
             )
     else:
         search_some(strategy, graph, args)
 
 
-def get_strategy(graph: Graph, args: NS) -> Strategy:
-    def strat_name(name: str) -> str:
-        alias = STRATEGIES_ALIASES.get(name)
-        if alias is None:
-            return name
-        return strat_name(alias)
+def strategy_name(name: str) -> str:
+    alias = STRATEGIES_ALIASES.get(name)
+    if alias is None:
+        return name
+    return strategy_name(alias)
 
-    name = strat_name(args.strategy)
+
+def get_strategy(graph: Graph, args: NS) -> Strategy:
+    name = strategy_name(args.strategy)
     options = dict(
         threads=args.threads,
         max_unroll=args.max_unroll,
@@ -716,6 +741,27 @@ def strat_name(name: str) -> str:
         },
         "default_strategy": "tile_oo",
     },
+    "conv2d_nhwc_frsc": {
+        "dims": ["n", "h", "w", "f", "r", "s", "c", "SH", "SW"],
+        "default_dims": [1, 112, 112, 64, 7, 7, 3, 2, 2],
+        "default_type": "f32",
+        "inputs": [
+            ["n", "h * SH + r - 1", "w * SW + s - 1", "c"],
+            ["f", "r", "s", "c"],
+        ],
+        "outputs": [["n", "h", "w", "f"]],
+        "reference_impl": None,  # defaults to graph evaluation
+        "operation": xtc_conv2d_nhwc_frsc_graph,
+        "backends": {
+            "mlir": {
+                "implementer": mlir_impl,
+            },
+            "tvm": {
+                "implementer": tvm_impl,
+            },
+        },
+        "default_strategy": "tile_oo",
+    },
     "relu": {
         "dims": ["i"],
         "default_dims": [512 * 1024],

src/xtc/csrcs/runtimes/host/evaluate_perf.c

@@ -30,10 +30,10 @@ typedef int (*packed_func_t)(PackedArg *, int *, int, PackedArg *, int *);
 
 #define NUMBER_FACTOR 2
 
-#define define_evaluateN(FUNC, ...)                                     \
+#define define_evaluateN(FUNC, ...)
 {                                                                       \
   assert(repeat > 0);                                                   \
-  assert(number > 0);                                                   \
+  assert(number >= 0);                                                  \
   assert(min_repeat_ms >= 0);                                           \
                                                                         \
   int fd = -1;                                                          \
@@ -50,9 +50,13 @@ typedef int (*packed_func_t)(PackedArg *, int *, int, PackedArg *, int *);
   }                                                                     \
   open_perf_events(events_num, events, perf_fds);                       \
                                                                         \
-  mem_barrier();                                                        \
-  (void)func(__VA_ARGS__);                                              \
-  mem_barrier();                                                        \
+  if (number > 0) {                                                     \
+    mem_barrier();                                                      \
+    (void)func(__VA_ARGS__);                                            \
+    mem_barrier();                                                      \
+  } else {                                                              \
+    number = 1;                                                         \
+  }                                                                     \
                                                                         \
   for (int r = 0; r < repeat; r++) {                                    \
     double elapsed;                                                     \

src/xtc/graphs/xtc/builder.py

@@ -6,6 +6,7 @@
 
 from .graph import XTCGraph
 from .context import XTCGraphContext
+from .expr import XTCExpr
 
 
 class graph_builder:
@@ -25,3 +26,9 @@ def __exit__(self, *_: Any) -> None:
     def graph(self) -> XTCGraph:
         assert self._graph is not None, "can't get graph inside builder context"
         return self._graph
+
+    def set_outputs(self, *outs: XTCExpr) -> None:
+        return XTCGraphContext.current.set_outputs(*outs)
+
+    def set_inputs(self, *inps: XTCExpr) -> None:
+        return XTCGraphContext.current.set_inputs(*inps)

src/xtc/graphs/xtc/context.py

@@ -35,14 +35,23 @@ def add_outputs(self, *outs: XTCExpr) -> None:
     def add_inputs(self, *inps: XTCExpr) -> None:
         self._inputs.extend(inps)
 
+    def set_outputs(self, *outs: XTCExpr) -> None:
+        self._outputs = list(outs)
+
+    def set_inputs(self, *inps: XTCExpr) -> None:
+        self._inputs = list(inps)
+
     def _infer_inputs(self, inps_seed: list[XTCExpr]) -> list[XTCExpr]:
         defs = set(self._exprs)
         uses = []
         for expr in self._exprs:
             if isinstance(expr, XTCOpExpr):
                 for arg in expr.args:
                     uses.append(arg)
-        inputs = inps_seed + [use for use in uses if use not in defs]
+        inputs = list({use._idx: use for use in uses if use not in defs}.values())
+        # when no order specified, sort by expr id
+        inputs = sorted(inputs, key=lambda x: x._idx)
+        inputs = inps_seed + inputs
         inputs = list({expr._idx: expr for expr in inputs}.values())
         return inputs
 
@@ -113,11 +122,5 @@ def append(self, expr: XTCExpr, name: str | None = None) -> XTCExpr:
             scope.add_expr(expr, name)
         return expr
 
-    def outputs(self, *outs: XTCExpr) -> None:
-        return self.current.add_outputs(*outs)
-
-    def inputs(self, *inps: XTCExpr) -> None:
-        return self.current.add_inputs(*inps)
-
 
 XTCGraphContext = XTCGraphScopes()

src/xtc/graphs/xtc/expr.py

@@ -14,6 +14,7 @@
 from .operators import (
     XTCOperator,
     XTCOperTensor,
+    XTCOperCompute,
     XTCOperMatmul,
     XTCOperRelu,
     XTCOperConv2D,
@@ -76,6 +77,10 @@ def forward_types(
     @abstractmethod
     def forward(self, inputs: Sequence[Tensor]) -> Sequence[XTCTensor]: ...
 
+    @property
+    @abstractmethod
+    def op(self) -> XTCOperator: ...
+
     @property
     @abstractmethod
     def op_name(self) -> str: ...
@@ -150,6 +155,11 @@ def type(self) -> XTCTensorType:
         assert isinstance(self._value, XTCTensor)
         return self._value.type
 
+    @property
+    @override
+    def op(self) -> XTCOperTensor:
+        return self._op
+
     @property
     @override
     def op_name(self) -> str:
@@ -180,11 +190,16 @@ def __str__(self) -> str:
 
 
 class XTCOpExpr(XTCExpr):
-    def __init__(self, op: XTCOperator, args: ArgumentsType) -> None:
+    def __init__(self, op: XTCOperCompute, args: ArgumentsType) -> None:
         super().__init__()
         self._op = op
         self._args = args
 
+    @property
+    @override
+    def op(self) -> XTCOperCompute:
+        return self._op
+
     @property
     @override
     def op_name(self) -> str:

src/xtc/graphs/xtc/node.py

@@ -136,9 +136,7 @@ def _operation(self) -> XTCOperation:
         assert self._outputs_types is not None
         inputs_types = self._inputs_types
         outputs_types = self._outputs_types
-        assert isinstance(self._expr, XTCTensorExpr) or isinstance(
-            self._expr, XTCOpExpr
-        )
+        assert isinstance(self._expr, XTCOpExpr)
         return self._expr._op.get_operation(
             inps_types=inputs_types,
             outs_types=outputs_types,

src/xtc/graphs/xtc/operation.py

@@ -5,19 +5,21 @@
 from typing_extensions import override
 from collections.abc import Mapping, Sequence
 from typing import Any
+import json
 
+from xtc.graphs.xtc.data import XTCTensorType
 from xtc.itf.graph import Operation
 from xtc.itf.graph.operation import AccessesMaps
-from xtc.itf.data import TensorType
+from xtc.utils.math import mulall
 
 
 class XTCOperation(Operation):
     def __init__(
         self,
         name: str,
         attrs: Mapping[str, Any],
-        inputs_types: Sequence[TensorType],
-        outputs_types: Sequence[TensorType],
+        inputs_types: Sequence[XTCTensorType],
+        outputs_types: Sequence[XTCTensorType],
         dims: Mapping[str, int | str],
         kinds: Sequence[str],
         inps_maps: Sequence[Sequence[str]],
@@ -48,12 +50,12 @@ def attrs(self) -> Mapping[str, Any]:
 
     @property
     @override
-    def inputs_types(self) -> Sequence[TensorType]:
+    def inputs_types(self) -> Sequence[XTCTensorType]:
         return self._inputs_types
 
     @property
     @override
-    def outputs_types(self) -> Sequence[TensorType]:
+    def outputs_types(self) -> Sequence[XTCTensorType]:
         return self._outputs_types
 
     @property
@@ -69,3 +71,30 @@ def dims_kind(self, kind: str) -> Sequence[str]:
     @override
     def accesses_maps(self) -> AccessesMaps:
         return self._maps
+
+    @property
+    @override
+    def ops_count(self) -> int:
+        # Assume single output, hence estimate
+        # ops as the product of all dimensions
+        # in the iteration space
+        shape = self._outputs_types[0].constant_shape
+        ops_count = mulall(list(shape))
+        return ops_count
+
+    @property
+    @override
+    def ops_dtype(self) -> str:
+        # Assume single output, hence estimate
+        # dtype as the first output dtype
+        return self._outputs_types[0].constant_dtype
+
+    @property
+    @override
+    def signature(self) -> list[Any]:
+        # Normalize json
+        return json.loads(
+            json.dumps(
+                [self.name, list(self.dims.values()), self.ops_dtype, dict(self.attrs)]
+            )
+        )