Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 24 additions & 7 deletions src/tir/transforms/lower_tvm_builtin.cc
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,17 @@
namespace tvm {
namespace tir {

namespace {
Call MakeMemCopyHelper(const CallNode* op, std::string packed_func_name) {
PrimExpr dst = op->args[0];
PrimExpr src = op->args[1];
PrimExpr size = op->args[2];

return Call(DataType::Int(32), builtin::tvm_call_packed(),
{StringImm(packed_func_name), dst, src, size});
}
} // namespace

class StackSizeChecker : public StmtExprVisitor {
public:
struct StackSizes {
Expand Down Expand Up @@ -73,10 +84,19 @@ class StackSizeChecker : public StmtExprVisitor {
return MakeShape(op);
} else if (op->op.same_as(builtin::tvm_stack_make_array())) {
return MakeArray(op);
} else if (op->op.same_as(builtin::mem_copy())) {
return MakeMemCopy(op);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can't you just call MakeCallPacked for this?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not quite, because the additional string for the name of the PackedFunc to execute takes up another spot on the stack. But since I needed to test it in order to convince myself of that, I've added a comment to explain the reasoning.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I also figured that by having the MakeMemCopyHelper be shared between StackSizeChecker and BuiltInLower makes it harder for them to get out of sync in the future.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That makes sense. Does that mean that the name "nonexistent_function" is actually significant in the sense that its length equals or exceeds the length of the actually called function name? That would be worth a comment (if that's the case).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question to ask, and no, the function name isn't significant. The stack size here is in number of arguments, where a string only counts as a single argument.

} else {
return StmtExprVisitor::VisitExpr_(op);
}
}

void MakeMemCopy(const CallNode* op) {
Call call_packed = MakeMemCopyHelper(op, "nonexistent_function");

return VisitExpr(call_packed);
}

// call shape
void MakeShape(const CallNode* op) {
// if args.size() == 0, it is still valid and represents a scalar
Expand Down Expand Up @@ -346,15 +366,12 @@ class BuiltinLower : public StmtExprMutator {
}

PrimExpr MakeMemCopy(const CallNode* op) {
PrimExpr dst = op->args[0];
PrimExpr src = op->args[1];
PrimExpr size = op->args[2];
std::stringstream packed_func_name;
packed_func_name << "device_api." << runtime::DeviceName(device_type_.as<IntImmNode>()->value)
<< ".mem_copy";

std::string fdevapi_prefix =
"device_api." + std::string(runtime::DeviceName(device_type_.as<IntImmNode>()->value));
Call call_packed = MakeMemCopyHelper(op, packed_func_name.str());

Call call_packed = Call(DataType::Int(32), builtin::tvm_call_packed(),
{StringImm(fdevapi_prefix + ".mem_copy"), dst, src, size});
return VisitExpr(call_packed);
}

Expand Down
3 changes: 2 additions & 1 deletion tests/python/contrib/test_hexagon/test_cache_read_write.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ def intrin_func(ins, outs):


@requires_hexagon_toolchain
def test_cache_read_write(android_serial_number, tvm_tracker_host, tvm_tracker_port):
def test_cache_read_write(android_serial_number, tvm_tracker_host, tvm_tracker_port, adb_server_socket):
size = 128
outer_shape = (size,)
factor = 16
Expand Down Expand Up @@ -115,6 +115,7 @@ def test_cache_read_write(android_serial_number, tvm_tracker_host, tvm_tracker_p
"rpc_tracker_host": tvm_tracker_host,
"rpc_tracker_port": tvm_tracker_port,
"rpc_server_port": 7070,
"adb_server_socket": adb_server_socket,
}
launcher = HexagonLauncher(serial_number=android_serial_number, rpc_info=rpc_info)
launcher.upload(dso_binary_path, dso_binary)
Expand Down