Move impl to C++

Author: peterbell10

python/src/specialize.cc

@@ -6,6 +6,7 @@
 #include <functional>
 #include <pybind11/pybind11.h>
 #include <string>
+#include <string_view>
 #include <unordered_map>
 #include <utility>
 
@@ -605,9 +606,151 @@ PyObject *specialize_impl(PyObject *self, PyObject *const *args,
   return PyTuple_Pack(2, type.ptr(), key.ptr());
 }
 
+bool visit_make_tensordesc_args(PyObject *arg, PyObject *sig,
+                                PyObject *relevant_paths,
+                                PyObject *tensordesc_meta,
+                                bool has_tensordesc_meta, PyObject *base_args,
+                                PyObject *make_tensordesc_arg,
+                                Py_ssize_t *tensordesc_idx, PyObject *result) {
+  assert(PyTuple_Check(sig));
+  auto arg_fast =
+      from_new_ref(PySequence_Fast(arg, "Expected iterable args node"));
+  if (!arg_fast)
+    return false;
+
+  Py_ssize_t arg_len = PySequence_Fast_GET_SIZE(arg_fast.ptr());
+  Py_ssize_t sig_len = PyTuple_GET_SIZE(sig);
+  assert(sig_len == arg_len || !"Invalid signature");
+  Py_ssize_t len = arg_len;
+
+  for (Py_ssize_t i = 0; i < len; ++i) {
+    PyObject *a = PySequence_Fast_GET_ITEM(arg_fast.ptr(), i);
+    PyObject *s = PyTuple_GET_ITEM(sig, i);
+
+    if (PyUnicode_CheckExact(s)) {
+      Py_ssize_t size;
+      const char *type_cstr = PyUnicode_AsUTF8AndSize(s, &size);
+      if (!type_cstr)
+        return false;
+
+      // if not s.startswith("tensordesc")
+      std::string_view tensordesc = "tensordesc";
+      std::string_view type_str(type_cstr, size);
+      if (type_str.substr(0, tensordesc.length()) != tensordesc) {
+        if (PyList_Append(result, a) < 0)
+          return false;
+        continue;
+      }
+
+      PyObject *meta = Py_None;
+      if (has_tensordesc_meta) {
+        // Borrowed reference
+        meta = PyList_GetItem(tensordesc_meta, *tensordesc_idx);
+        if (!meta)
+          return false;
+      }
+
+      PyObject *vector_args[] = {a, meta, base_args};
+      auto desc_args = from_new_ref(
+          PyObject_Vectorcall(make_tensordesc_arg, vector_args, 3, nullptr));
+      if (!desc_args)
+        return false;
+      // list.extend(desc_args)
+      if (PyList_SetSlice(result, PY_SSIZE_T_MAX, PY_SSIZE_T_MAX,
+                          desc_args.ptr()) < 0)
+        return false;
+
+      *tensordesc_idx += 1;
+      continue;
+    }
+
+    auto key = from_new_ref(PyLong_FromSsize_t(i));
+    if (!key)
+      return false;
+
+    // Borrowed ref
+    PyObject *inner_relevant_paths =
+        PyDict_GetItemWithError(relevant_paths, key.ptr());
+    if (PyErr_Occurred())
+      return false;
+
+    if (!inner_relevant_paths) {
+      // Short-circuit if tuple doesn't contain any tensordesc args
+      if (PyList_Append(result, a) < 0)
+        return false;
+      continue;
+    }
+
+    // Recurse into tuple
+    auto inner_res = from_new_ref(PyList_New(0));
+    if (!inner_res)
+      return false;
+    if (!visit_make_tensordesc_args(
+            a, s, inner_relevant_paths, tensordesc_meta, has_tensordesc_meta,
+            base_args, make_tensordesc_arg, tensordesc_idx, inner_res.ptr()))
+      return false;
+
+    auto inner_tuple = from_new_ref(PyList_AsTuple(inner_res.ptr()));
+    if (!inner_tuple)
+      return false;
+    if (PyList_Append(result, inner_tuple.ptr()) < 0)
+      return false;
+  }
+  return true;
+}
+
+PyObject *make_tensordesc_args(PyObject *self, PyObject *const *args,
+                               Py_ssize_t nargs) {
+  if (nargs != 6) {
+    PyErr_SetString(PyExc_TypeError,
+                    "make_tensordesc_args expected 6 arguments");
+    return nullptr;
+  }
+
+  PyObject *kernel_args = args[0];
+  PyObject *signature = args[1];
+  PyObject *relevant_paths = args[2];
+  PyObject *tensordesc_meta = args[3];
+  PyObject *base_args = args[4];
+  PyObject *make_tensordesc_arg = args[5];
+
+  bool has_tensordesc_meta = tensordesc_meta != Py_None;
+  if (has_tensordesc_meta && !PyList_CheckExact(tensordesc_meta)) {
+    PyErr_SetString(PyExc_TypeError, "Expected tensordesc_meta to be a list");
+    return nullptr;
+  }
+
+  auto result = from_new_ref(PyList_New(0));
+  if (!result)
+    return nullptr;
+
+  Py_ssize_t tensordesc_idx = 0;
+  if (!visit_make_tensordesc_args(kernel_args, signature, relevant_paths,
+                                  tensordesc_meta, has_tensordesc_meta,
+                                  base_args, make_tensordesc_arg,
+                                  &tensordesc_idx, result.ptr()))
+    return nullptr;
+
+  if (has_tensordesc_meta) {
+    Py_ssize_t meta_len = PySequence_Size(tensordesc_meta);
+    if (meta_len < 0)
+      return nullptr;
+
+    if (tensordesc_idx != meta_len) {
+      PyErr_SetString(PyExc_ValueError,
+                      "make_tensordesc_args: tensordesc_idx != meta_len");
+      return nullptr;
+    }
+  }
+
+  return result.release().ptr();
+}
+
 static PyMethodDef module_methods[] = {
     {"native_specialize_impl", (PyCFunction)specialize_impl, METH_FASTCALL,
      nullptr},
+    {"make_tensordesc_args", (PyCFunction)make_tensordesc_args, METH_FASTCALL,
+     "Helper to translate tensordesc args"},
     {nullptr, nullptr, 0, nullptr} // sentinel
 };
 

python/test/unit/runtime/test_driver.py

@@ -4,7 +4,7 @@
 
 import triton
 import triton.language as tl
-from triton.backends.driver import expand_signature, wrap_descriptors
+from triton.backends.driver import expand_signature, wrap_handle_tensordesc_impl
 
 
 def test_is_lazy():
@@ -69,7 +69,7 @@ def test_expand_signature_with_aggregate_tensordesc():
     assert expanded[2:] == ["nvTmaDesc", "i32", "i32", "i32", "i32", "i64", "i64", "i64", "i64"]
 
 
-def test_wrap_descriptors_handles_aggregate_arguments():
+def test_wrap_tensordesc_handles_aggregate_arguments():
     signature = {0: ("tensordesc<fp16[16,16]>", "i32"), 1: "i64", 2: "tensordesc<fp16[16,16]>"}
     outer_meta = {"tag": "outer"}
     launcher_calls = []
@@ -81,24 +81,24 @@ def launcher(*args):
     def make_descriptor(arg, meta, base_args):
         return [("desc", arg, meta, base_args[0]), ("shape", arg)]
 
-    wrapped = wrap_descriptors(launcher, signature, [None, outer_meta], make_descriptor)
+    wrapped = wrap_handle_tensordesc_impl(launcher, signature, [None, outer_meta], make_descriptor)
     assert wrapped("meta0", "meta1", (("A", 7), 9, "B")) == "ok"
 
     assert len(launcher_calls) == 1
     assert launcher_calls[0][0] == "meta0"
     assert launcher_calls[0][1] == "meta1"
-    assert launcher_calls[0][2] == (
+    assert launcher_calls[0][2] == [
         (("desc", "A", None, "meta0"), ("shape", "A"), 7),
         9,
         ("desc", "B", outer_meta, "meta0"),
         ("shape", "B"),
-    )
+    ]
 
 
-def test_wrap_descriptors_is_noop_without_tensordesc():
+def test_wrap_tensordesc_is_noop_without_tensordesc():
 
     def launcher(*args):
         return args
 
-    wrapped = wrap_descriptors(launcher, {0: "i32", 1: ("i64", "constexpr")}, None, lambda *_: [])
+    wrapped = wrap_handle_tensordesc_impl(launcher, {0: "i32", 1: ("i64", "constexpr")}, None, lambda *_: [])
     assert wrapped is launcher

python/triton/backends/driver.py

@@ -1,9 +1,9 @@
 from abc import ABCMeta, abstractmethod
 import re
 from typing import Callable, List, Protocol, Sequence
-from collections import defaultdict
 
 from triton._utils import find_paths_if
+from triton._C.libtriton import make_tensordesc_args
 
 
 def decompose_descriptor(arg):
@@ -16,48 +16,25 @@ def _is_descriptor(arg):
     return isinstance(arg, str) and arg.startswith("tensordesc")
 
 
-def _make_tensordesc_args(args, signature, relevant_paths, tensordesc_meta, base_args, make_tensordesc_arg):
-
-    def visit(arg, sig, relevant_paths, tensordesc_idx, result):
-        for i, (a, s) in enumerate(zip(arg, sig)):
-            rel_paths = relevant_paths.get(i, None)
-            if rel_paths is None:
-                result.append(a)
-            elif len(rel_paths) == 0:
-                meta = tensordesc_meta[tensordesc_idx] if tensordesc_meta else None
-                result.extend(make_tensordesc_arg(a, meta, base_args))
-                tensordesc_idx += 1
-            else:
-                inner_res = []
-                tensordesc_idx = visit(a, s, rel_paths, tensordesc_idx, inner_res)
-                result.append(tuple(inner_res))
-        return tensordesc_idx
-
-    result = []
-    tensordesc_idx = visit(args, signature, relevant_paths, 0, result)
-    assert not tensordesc_meta or tensordesc_idx == len(tensordesc_meta)
-    return result
-
-
 def wrap_handle_tensordesc_impl(launcher, signature, tensordesc_meta, make_tensordesc_arg):
-    signature = tuple(signature.values()) if hasattr(signature, "values") else tuple(signature)
+    signature = tuple(signature.values())
     tensordesc_paths = find_paths_if(signature, lambda _, x: _is_descriptor(x))
     if len(tensordesc_paths) == 0:
         return launcher
 
-    # Build a tree to speed up checking, it will look like:
-    # signature = ['tensordesc', 'i32', ('i32', 'tensordesc')]
+    # Build a tree to speed up tensordesc type checking, e.g.
+    # signature = ('tensordesc', 'i32', ('i32', 'tensordesc'))
     # relevant_paths = {0: {}, 2: {1: {}}}
-    relevant_paths = defaultdict(defaultdict)
+    relevant_paths = {}
     for path in tensordesc_paths:
         cur = relevant_paths
         for step in path:
-            cur = cur[step]
+            cur = cur.setdefault(step, {})
 
     def inner(*args):
         base_args = args[:-1]
         kernel_args = args[-1]
-        wrapped = _make_tensordesc_args(
+        wrapped = make_tensordesc_args(
             kernel_args,
             signature,
             relevant_paths,