Allow strings as constant kernel arguments

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

cext/py.h

@@ -256,6 +256,12 @@ static inline PyPtr getattr(const PyPtr& obj, const char* attrname) {
     return getattr(obj.get(), attrname);
 }
 
+static inline void pyunicode_intern_in_place(PyPtr* s) {
+    PyObject* raw = s->release();
+    PyUnicode_InternInPlace(&raw);
+    *s = steal(raw);
+}
+
 struct ErrorGuard {
     SavedException exc;
 

cext/tile_kernel.cpp

@@ -573,6 +573,7 @@ struct AggregateArgType {
     X(Int) \
     X(Float) \
     X(None_) \
+    X(String) \
     X(Enum) \
     X(NativeDType) \
     X(ForeignDType)
@@ -651,6 +652,7 @@ enum class PythonArgKind : uint8_t {
     ConstantInt,
     ConstantFloat,
     ConstantNone,
+    ConstantString,
     IdentityConstant,
     ForeignDTypeConstant,
     // A torch.Tensor that we can access via torch._C._to_dlpack
@@ -675,6 +677,7 @@ static inline PythonArgKind constant_kind_as_arg_kind(ConstantKind kind) {
     case ConstantKind::Int: return PythonArgKind::ConstantInt;
     case ConstantKind::Float: return PythonArgKind::ConstantFloat;
     case ConstantKind::None_: return PythonArgKind::ConstantNone;
+    case ConstantKind::String: return PythonArgKind::ConstantString;
     case ConstantKind::Enum: return PythonArgKind::IdentityConstant;
     case ConstantKind::NativeDType: return PythonArgKind::IdentityConstant;
     case ConstantKind::ForeignDType: return PythonArgKind::ForeignDTypeConstant;
@@ -689,6 +692,7 @@ static ParameterKind::Category param_category_from_pyarg_kind(PythonArgKind k) {
     case PythonArgKind::ConstantInt: return ParameterKind::ConstantInt;
     case PythonArgKind::ConstantFloat: return ParameterKind::ConstantFloat;
     case PythonArgKind::ConstantNone: return ParameterKind::ConstantNone;
+    case PythonArgKind::ConstantString: return ParameterKind::IdentityConstant;
     case PythonArgKind::IdentityConstant: return ParameterKind::IdentityConstant;
     case PythonArgKind::ForeignDTypeConstant: return ParameterKind::IdentityConstant;
     case PythonArgKind::TorchTensorDlpack: return ParameterKind::Array;
@@ -915,6 +919,9 @@ static std::optional<ConstantKind> classify_constant(PyObject* obj, bool kernel_
     if (obj == Py_None)
         return ConstantKind::None_;
 
+    if (PyUnicode_CheckExact(obj))
+        return ConstantKind::String;
+
     if (PyObject_TypeCheck(obj, reinterpret_cast<PyTypeObject*>(g_enum_Enum_type)))
         return ConstantKind::Enum;
 
@@ -2097,6 +2104,7 @@ static void extract_identity_constant(PyObject* object, Vec<int64_t>* constants,
     identity_constants->push_back(object);
 }
 
+
 static PyPtr parse_identity_constant_constraint(ConstantCursor& cursor,
                                                 const Vec<PyObject*>& identity_constants) {
     int64_t address = cursor.next();
@@ -2107,6 +2115,21 @@ static PyPtr parse_identity_constant_constraint(ConstantCursor& cursor,
     CHECK_UNREACHABLE;
 }
 
+static Status extract_string_constant(PyObject* pyobj, Vec<int64_t>* constants,
+                                      Vec<PyObject*>* identity_constants,
+                                      Vec<PyPtr>* pyarg_refs) {
+    if (!PyUnicode_CHECK_INTERNED(pyobj)) {
+        PyPtr ref = newref(pyobj);
+        pyunicode_intern_in_place(&ref);
+        if (!PyUnicode_CHECK_INTERNED(ref.get()))
+            return raise(PyExc_RuntimeError, "Failed to intern a string kernel argument");
+        pyobj = ref.get();
+        pyarg_refs->push_back(std::move(ref));
+    }
+    extract_identity_constant(pyobj, constants, identity_constants);
+    return OK;
+}
+
 static Status extract_foreign_dtype_constant(PyObject* object, Vec<int64_t>* constants,
                                              Vec<PyObject*>* identity_constants) {
     HashMap<PyPtr, ForeignDTypeInfo>::Item* item = get_foreign_dtype_registry()->find(object);
@@ -2408,6 +2431,9 @@ static Status extract_arg(const DriverApi* driver, PyObject* obj, PythonArgKind
         return OK;
     case PythonArgKind::ConstantNone:
         return OK;
+    case PythonArgKind::ConstantString:
+        return extract_string_constant(obj, &helper.constants, &helper.identity_constants,
+                                       &helper.pyarg_refs);
     case PythonArgKind::IdentityConstant:
         extract_identity_constant(obj, &helper.constants, &helper.identity_constants);
         return OK;

src/cuda/tile/_ir/typing_support.py

@@ -137,6 +137,8 @@ def type_of_constant_python_value(val, typing_hooks: TypingHooks) -> Type:
             return typing_hooks.get_tensor_like_type(dtype_of_constant_scalar(val), ())
         case ConstantKind.None_:
             return NONE
+        case ConstantKind.String:
+            return StringTy(val)
         case ConstantKind.Enum:
             return EnumTy(val)
         case ConstantKind.NativeDType:
@@ -145,8 +147,6 @@ def type_of_constant_python_value(val, typing_hooks: TypingHooks) -> Type:
             return _get_dtype_spec(foreign_dtype_object_to_native(val))
         case _: assert False
 
-    if isinstance(val, str):
-        return StringTy(val)
     if val is Ellipsis:
         return ELLIPSIS
     if isinstance(val, slice):

src/cuda/tile/compilation/_name_mangling.py

@@ -212,6 +212,8 @@ def _mangle_constraint(p: ParameterConstraint, alias_group_map: dict[str, int],
                 return "F" + f"{i:016x}"
             case ConstantKind.None_:
                 return "Cn_"
+            case ConstantKind.String:
+                return "Cs_" + _mangle_string(p.value)
             case ConstantKind.Enum:
                 assert cconv_v3_enabled()
                 return "Ce_" + _mangle_enum_constant(p.value, collected_globals)
@@ -249,6 +251,8 @@ def _demangle_constraint(cursor: _Cursor,
         return ConstantConstraint(f)
     elif c == "Cn_" and cconv_v3_enabled():
         return ConstantConstraint(None)
+    elif c == "Cs_" and cconv_v3_enabled():
+        return ConstantConstraint(_demangle_string(cursor))
     elif c == "Ce_" and cconv_v3_enabled():
         return ConstantConstraint(_demangle_enum_constant(cursor, allowed_globals))
     elif c == "Cd_" and cconv_v3_enabled():

test/test_constant.py

@@ -92,3 +92,22 @@ def kern(c: ct.Constant, out):
     out = torch.full((), -1, dtype=torch.int32, device="cuda")
     ct.launch(torch.cuda.current_stream(), (1,), kern, (123, out))
     assert out.tolist() == 0
+
+
+@pytest.mark.skipif(not cconv_v3_enabled(), reason="Requires cconv3 enabled")
+def test_str_kernel_argument():
+    @ct.kernel
+    def kern(c: ct.Constant, out):
+        ct.scatter(out, 0, c == "hello")
+        ct.scatter(out, 1, c == "test string for test_str_kernel_argument!")
+
+    # Repeat a few times to make sure we exercise different branches that depend on string interning
+    for i in range(3):
+        out = torch.full((2,), -1, dtype=torch.int32, device="cuda")
+        ct.launch(torch.cuda.current_stream(), (1,), kern, ("hello", out))
+        assert out.tolist() == [1, 0]
+
+        out = torch.full((2,), -1, dtype=torch.int32, device="cuda")
+        ct.launch(torch.cuda.current_stream(), (1,), kern,
+                  ("test string for test_str_kernel_argument!", out))
+        assert out.tolist() == [0, 1]

test/test_name_mangling.py

@@ -458,6 +458,13 @@ class MyEnum(Enum):
         "_Cn__I123_Cn_",
         id="none_constant",
     ),
+
+    # String constant
+    pytest.param(
+        ["Hello", "world!"],
+        "_Cs_Hello_z_Cs_world_21_z",
+        id="string_constant",
+    ),
 ] if cconv_v3_enabled() else [])
 @pytest.mark.skipif(not cconv_v3_enabled(), reason="Requires cconv3 enabled")
 def test_name_mangling_cutile_python_v3(parameters, expected_suffix):