This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 019da20b5 [FIX][Codegen] Correct C emission and Target host
serialization (#172)
019da20b5 is described below
commit 019da20b566d9b3e1118819f2d58530204289b2f
Author: Lei Wang <[email protected]>
AuthorDate: Wed Nov 4 02:48:47 2026 +0800
[FIX][Codegen] Correct C emission and Target host serialization (#172)
This PR fixes three correctness issues in shared C code generation and
Target serialization.
### Preserve string literal bytes
`StringImm` currently writes its contents directly between quotes.
Embedded quotes, backslashes, or newlines can therefore produce invalid
C source or change the string's contents. Route it through the existing
`PrintEscapedCString` helper:
| Input value, in Python notation | Previous C literal | Fixed C literal
|
| --- | --- | --- |
| `'a "quoted" message'` | `"a "quoted" message"` | `"a \"quoted\"
message"` |
| `'path\\to\\file\\'` | `"path\to\file\"` | `"path\\to\\file\\"` |
| `'first\nsecond'` | An actual newline inside the quoted literal |
`"first\nsecond"` |
The existing helper, also used for assertion messages, has a separate
issue with numeric escape boundaries. C hexadecimal escapes consume all
following hexadecimal digits, and octal escapes consume up to three
octal digits. The byte sequences below omit the terminating NUL:
| Intended bytes | Previous helper output | Fixed output |
| --- | --- | --- |
| `00 37` — NUL followed by `7` | `"\07"` | `"\0007"` |
| `01 46` — byte `0x01` followed by `F` | `"\x01F"` | `"\001F"` |
| `C3 A9 46` — UTF-8 for `éF` | `"\xc3\xa9F"` | `"\303\251F"` |
For example, `"\x01F"` represents one byte, `0x1F`, instead of the
intended two bytes, `0x01` and `0x46`. Emitting exactly three octal
digits per numeric byte escape makes the boundary explicit, even when
the next character is another digit.
There is also a C11 trigraph case: `"??=abc"` becomes `"#abc"` before
string literals are parsed. Emitting `"\?\?=abc"` preserves the original
`??=abc` bytes. The helper now escapes question marks as well.
### Emit initializer setup before tensor declarations
`DispatchDeclTensor` starts writing the pointer declaration before
printing its initializer. However, printing an expression such as
`stack_alloca` can itself emit setup statements. For example:
```python
local = T.decl_tensor(
T.stack_alloca("shape", 2),
ty_args=[T.Tensor((2,), "int64")],
)
```
Previously, the stack declaration could be emitted in the middle of the
pointer initializer:
```c
int64_t* local = ((int64_t*) TVMFFIAny stack[1];
((int64_t*)stack));
```
Generate the initializer expression text first, allowing its setup
statements to be emitted before starting the complete pointer
declaration:
```c
TVMFFIAny stack[1];
int64_t* local = ((int64_t*)stack);
```
### Invalidate Target string caches when changing the host
Once a target's string representation has been cached, cloning it and
changing its host also copies the stale `str_repr_`. The resulting
object can have the correct host while its serialized string still
describes the previous one.
For example, adding a host after populating the cache:
```python
from tvm_tirx.target import Target
original = Target({"kind": "cuda", "arch": "sm_80"})
str(original) # Populate the cache before cloning.
updated = Target(original, host="c")
restored = Target(str(updated))
assert updated.host.kind.name == "c"
# Before: restored.host is None.
# After: restored.host.kind.name == "c".
```
The same issue affects replacing and removing a host. In each row below,
`str(original)` has already populated the cache:
| Original host | Operation | Host recovered from the string before this
fix | After this fix |
| --- | --- | --- | --- |
| None | `Target(original, host="c")` | None | `c` |
| `c` | `original.with_host("llvm")` | `c` | `llvm` |
| `c` | `original.with_host(None)` | `c` | None |
Clear `str_repr_` on the new object in `WithHost`, and route the
constructor that takes a host and `WithoutHost()` through that same
path. This keeps serialization consistent when adding, replacing, or
removing a host.
---
include/tvm/target/target.h | 11 ++------
src/codegen/c/codegen_c.cc | 38 +++++++++++++++------------
src/target/target.cc | 1 +
tests/cpp/target/target_test.cc | 24 +++++++++++++++++
tests/python/codegen/c/test_basic.py | 50 ++++++++++++++++++++++++++++++++++++
5 files changed, 99 insertions(+), 25 deletions(-)
diff --git a/include/tvm/target/target.h b/include/tvm/target/target.h
index 4e12f08ed..7843098fb 100644
--- a/include/tvm/target/target.h
+++ b/include/tvm/target/target.h
@@ -183,12 +183,7 @@ class Target : public ffi::ObjectRef {
* \param target The Target typed object with host field undefined for target
* \param host The Target typed object for target host
*/
- explicit Target(Target target, ffi::Optional<Target> host) {
- ffi::ObjectPtr<TargetObj> n = ffi::make_object<TargetObj>(*target.get());
- n->host =
- host.has_value() ?
ffi::Optional<ffi::ObjectRef>(std::move(host).value()) : std::nullopt;
- data_ = std::move(n);
- }
+ explicit Target(Target target, ffi::Optional<Target> host) :
Target(WithHost(target, host)) {}
explicit Target(ffi::UnsafeInit tag, ffi::ObjectPtr<TargetObj> n) :
ffi::ObjectRef(tag) {
data_ = std::move(n);
}
@@ -200,9 +195,7 @@ class Target : public ffi::ObjectRef {
/*! \return The target with the host stripped out */
Target WithoutHost() const {
if ((*this)->GetHost()) {
- auto output = ffi::make_object<TargetObj>(*get());
- output->host = std::nullopt;
- return Target(ffi::UnsafeInit{}, std::move(output));
+ return WithHost(*this, std::nullopt);
} else {
return *this;
}
diff --git a/src/codegen/c/codegen_c.cc b/src/codegen/c/codegen_c.cc
index 668a8dbd6..bd259d528 100644
--- a/src/codegen/c/codegen_c.cc
+++ b/src/codegen/c/codegen_c.cc
@@ -582,7 +582,7 @@ void CodeGenC::Dispatch_(const FloatImmObj* op,
std::ostream& os) { // NOLINT(*
PrintConst(op, os, this);
}
void CodeGenC::Dispatch_(const StringImmObj* op, std::ostream& os) { //
NOLINT(*)
- os << "\"" << op->value << "\"";
+ PrintEscapedCString(op->value, os);
}
template <typename T>
@@ -980,22 +980,23 @@ void CodeGenC::DispatchDeclTensor(const BindObj* op,
const CallObj* buffer_call)
}
alloc_storage_scope_[buffer.get()] = scope;
+ PtrType pointer_type(PrimType(dtype), scope);
+ // Packed argument decoding may already have established this typed view.
+ const auto* cast = data.as<CallObj>();
+ Expr initializer = data;
+ if (!(cast && cast->op.same_as(tirx::reinterpret_op()) &&
+ ffi::StructuralEqual()(cast->ty, pointer_type))) {
+ initializer = Call(pointer_type, tirx::reinterpret_op(), {data});
+ }
+ // An initializer such as stack_alloca emits preparation statements. Emit
+ // those before starting the declaration so they cannot split its
initializer.
+ std::string value = PrintExpr(initializer);
this->PrintIndent();
if (IsScopePartOfType()) {
PrintStorageScope(scope, stream);
}
- PtrType pointer_type(PrimType(dtype), scope);
PrintType(pointer_type, stream);
- stream << ' ' << AllocVarID(buffer.get()) << " = ";
- // Packed argument decoding may already have established this typed view.
- const auto* cast = data.as<CallObj>();
- if (cast && cast->op.same_as(tirx::reinterpret_op()) &&
- ffi::StructuralEqual()(cast->ty, pointer_type)) {
- PrintExpr(data, stream);
- } else {
- PrintExpr(Call(pointer_type, tirx::reinterpret_op(), {data}), stream);
- }
- stream << ";\n";
+ stream << ' ' << AllocVarID(buffer.get()) << " = " << value << ";\n";
RegisterHandleType(buffer.get(), PrimType(dtype));
}
@@ -1357,6 +1358,10 @@ void CodeGenC::PrintEscapedCString(const std::string&
str, std::ostream& os) {
case '\\':
os << "\\\\";
break;
+ case '?':
+ // Prevent C trigraph replacement before string literals are parsed.
+ os << "\\?";
+ break;
case '\n':
os << "\\n";
break;
@@ -1379,13 +1384,14 @@ void CodeGenC::PrintEscapedCString(const std::string&
str, std::ostream& os) {
os << "\\v";
break;
case '\0':
- os << "\\0";
+ os << "\\000";
break;
default:
if (c < 0x20 || c >= 0x7f) {
- // Non-printable: emit as hex escape
- os << "\\x" << std::setfill('0') << std::setw(2) << std::hex <<
static_cast<int>(c)
- << std::dec;
+ // Three octal digits encode exactly one byte. Hex escapes can
+ // consume following hex digits, including ordinary ASCII text.
+ os << '\\' << static_cast<char>('0' + ((c >> 6) & 7))
+ << static_cast<char>('0' + ((c >> 3) & 7)) <<
static_cast<char>('0' + (c & 7));
} else {
os << static_cast<char>(c);
}
diff --git a/src/target/target.cc b/src/target/target.cc
index 7c82dbe62..faa73ff46 100644
--- a/src/target/target.cc
+++ b/src/target/target.cc
@@ -54,6 +54,7 @@ class TargetInternal {
ffi::ObjectPtr<TargetObj> n = ffi::make_object<TargetObj>(*target.get());
n->host =
target_host.has_value() ?
ffi::Optional<ffi::ObjectRef>(target_host.value()) : std::nullopt;
+ n->str_repr_ = ffi::String();
return Target(ffi::UnsafeInit{}, std::move(n));
}
diff --git a/tests/cpp/target/target_test.cc b/tests/cpp/target/target_test.cc
index d657a2fb2..857b81789 100644
--- a/tests/cpp/target/target_test.cc
+++ b/tests/cpp/target/target_test.cc
@@ -141,6 +141,30 @@ TEST(TargetCreation,
RoundTripCanonicalizerFeaturesNestedHost) {
ASSERT_TRUE(reconstructed_host.value()->GetAttr<ffi::String>("mcpu").has_value());
}
+TEST(TargetCreation, HostConstructorInvalidatesStringCache) {
+ Target original(ffi::Map<ffi::String, ffi::Any>{{"kind",
ffi::String("cuda")},
+ {"arch",
ffi::String("sm_80")}});
+ ffi::String original_string = original->str();
+ Target combined(original, Target("c"));
+ EXPECT_EQ(Target(combined->str())->GetHost().value()->kind->name, "c");
+ EXPECT_FALSE(original->GetHost().has_value());
+ EXPECT_EQ(original->str(), original_string);
+}
+
+TEST(TargetCreation, WithoutHostInvalidatesStringCache) {
+ Target original(ffi::Map<ffi::String, ffi::Any>{
+ {"kind", ffi::String("cuda")},
+ {"arch", ffi::String("sm_80")},
+ {"host", ffi::Map<ffi::String, ffi::Any>{{"kind", ffi::String("c")}}},
+ });
+ ffi::String original_string = original->str();
+ Target without_host = original.WithoutHost();
+ EXPECT_FALSE(Target(without_host->str())->GetHost().has_value());
+ EXPECT_TRUE(original->GetHost().has_value());
+ EXPECT_EQ(original->str(), original_string);
+ EXPECT_TRUE(without_host.WithoutHost().same_as(without_host));
+}
+
TEST(TargetCreationFail, UnknownNonFeatureKeyStillFails) {
// Verify that unknown non-feature.* keys still fail schema validation
ffi::Map<ffi::String, ffi::Any> config = {
diff --git a/tests/python/codegen/c/test_basic.py
b/tests/python/codegen/c/test_basic.py
index bfd8c5bbc..1f9391b6b 100644
--- a/tests/python/codegen/c/test_basic.py
+++ b/tests/python/codegen/c/test_basic.py
@@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.
+import shutil
+
import numpy as np
import pytest
@@ -399,5 +401,53 @@ def test_unsupported_vector_diagnostic():
assert lines[index + 1] == " " * start + "^" * len(expression)
[email protected](
+ "message,index",
+ [
+ ('"\\\n', 0),
+ ("\0" + "7", 0),
+ ("éF", 1),
+ ("??=", 0),
+ ],
+)
+def test_string_literal_escapes(message, index, tmp_path):
+ compiler = shutil.which("cc") or shutil.which("gcc") or
shutil.which("clang")
+ if compiler is None:
+ pytest.skip("A C11 compiler is required")
+ expected = message.encode("utf-8")
+ size = len(expected)
+
+ @T.function
+ def read_literal(Out: T.Tensor((1,), "uint8")):
+ data = T.decl_tensor(
+ T.reinterpret(message, ty=T.Ptr("uint8")),
ty_args=[T.Tensor((size,), "uint8")]
+ )
+ # The chosen index stays in bounds even if a broken escape shortens
the literal.
+ Out[0] = data[index]
+
+ compiled = tvm_tirx.compile(read_literal, target="c")
+ path = str(tmp_path / "literal.so")
+ compiled.export_library(path, cc=compiler, options=["-std=c11"])
+ loaded = tvm_tirx.runtime.load_module(path)
+ out = tvm_tirx.runtime.tensor(np.zeros(1, dtype="uint8"))
+ loaded["read_literal"](out)
+ assert out.numpy()[0] == expected[index]
+
+
+def test_decl_tensor_stack_initializer():
+ @T.function
+ def stack_initializer(Out: T.Tensor((2,), "int64")):
+ local = T.decl_tensor(T.stack_alloca("shape", 2),
ty_args=[T.Tensor((2,), "int64")])
+ local[0] = T.int64(17)
+ local[1] = T.int64(23)
+ Out[0] = local[0]
+ Out[1] = local[1]
+
+ executable = tvm_tirx.compile(stack_initializer, target="c")
+ out = tvm_tirx.runtime.tensor(np.zeros(2, dtype="int64"))
+ executable(out)
+ np.testing.assert_array_equal(out.numpy(), [17, 23])
+
+
if __name__ == "__main__":
tvm_tirx.testing.main()
