Dump Metal codegen result to a temporary source file

[k-ye]

Concisely describe the proposed feature

I'd like to output the Metal codegen result to an actual source file. Currently, we write the common Metal helper and runtime code as string literals in C++. This approach is becoming very hard to maintain or iterate now, in order to support #593 .

Questions

I saw a sandbox folder is being created when running in dev mode. But that seems to be used to only hold of the Taichi runtime .ll and .bc, not the compiled user's Taichi kernels. So, are we outputting any tmp source file for any backend? I know that OpenGL doesn't, and the legacy Taichi did. But what about the LLVM backends?

On the other hand, when evolving Taichi to the LLVM backend, if there was a particular reason not to generate such files, then I'm fine keeping the helpers just as string literals..

[yuanming-hu]

So, are we outputting any tmp source file for any backend? I know that OpenGL doesn't, and the legacy Taichi did. But what about the LLVM backends?

Currently, no. But for debugging proposed, we may end up emitting intermediate LLVM IR to the sandbox folder. So feel free to use it :-)

One worry about the sandbox folder: I think we only use it for development mode now. A little work will have to be done to enable it for release mode. In that case, we need to carefully ensure the sandbox folders are deleted in time so that users' disk space won't get occupied. For a normal run it's fine, but special treatments are needed when users' process crashes halfway (and leaving the sandboxes there...)

[k-ye]

One worry about the sandbox folder: I think we only use it for development mode now. A little work will have to be done to enable it for release mode.

Ah i see. I think I will defer this approach due to the additional works needed...

[yuanming-hu]

In the worst case, since each source file would be just < 100 KB, maybe it's not too bad to leave them there...

[k-ye]

In the worst case, since each source file would be just < 100 KB, maybe it's not too bad to leave them there...

Hmm, this could probably be a bit surprising to the users... But I guess that by creating tmp files in /tmp + registering signal handlers at exit, it would already cover lots of the cases.. I did a quick search but didn't seem to find a particular module that sets up a sandbox dir and cleans up automatically on shutdown. tempfile may be the closest...

[archibate]

+1. dumping codegen result can be really helpful. In fact, OpenGL backend is already doing this for debug purpose:

taichi/taichi/codegen/codegen_opengl.cpp

Lines 53 to 57 in d3559d0

	#ifdef _GLSL_DEBUG
	TI_INFO("source of kernel [{}] * {}:\n{}", kernel_name, num_groups, kernel_source_code);
	std::ofstream(fmt::format("/tmp/{}.comp", kernel_name))
	.write(kernel_source_code.c_str(), kernel_source_code.size());
	#endif

[archibate]

In the worst case, since each source file would be just < 100 KB, maybe it's not too bad to leave them there...

Hmm, this could probably be a bit surprising to the users... But I guess that by creating tmp files in /tmp + registering signal handlers at exit, it would already cover lots of the cases.. I did a quick search but didn't seem to find a particular module that sets up a sandbox dir and cleans up automatically on shutdown. tempfile may be the closest...

Registering atexit callback doesn't remove the sandbox if taichi crashed accidentally.
Maybe we can remove it the next time start up?
I mean, we can name sandbox as /tmp/taichi-$PID, then on each start up, detect&remove that dir if we found the process PID is no longer exist.

[k-ye]

+1. dumping codegen result can be really helpful. In fact, OpenGL backend is already doing this for debug purpose:

Yep, I am also printing the source code to stdout. But this issue is not only about debugging. We have more and more Metal kernels that are part of the Taichi runtime, and should be shared by all user Taichi kernels. It will be much easier to improve these runtime kernels if they are in their native source format.

I mean, we can name sandbox as /tmp/taichi-$PID, then on each start up, detect&remove that dir if we found the process PID is no longer exist.

SG, actually i'm thinking even less elegant than that. Just leave the files in /tmp and let OS clean them up...

[archibate]

Yep, I am also printing the source code to stdout. But this issue is not only about debugging. We have more and more Metal kernels that are part of the Taichi runtime, and should be shared by all user Taichi kernels. It will be much easier to improve these runtime kernels if they are in their native source format.

Maybe we want something like

1. Write a kernel in python.
2. Export kernel source file: ti.export(kernel, 'kernel.comp')
3. (now edit kernel.comp to improve it)
4. Load kernel from source: improved_kernel = ti.import('kernel.comp')
5. call improved_kernel().

ti.export here may be related to #439 #394. (respectively: kernel.so, kernel.js).

[k-ye]

+1, I think what you described is more relevant to #439 .

The problem in this issue is more around the runtime kernels and the backend-specific helpers that are part of Taichi. For example, it would be easier for me to iterate the development if these helpers

taichi/taichi/platform/metal/helpers.metal.h

Lines 35 to 62 in d3559d0

	T union_cast(G g) {
	// For some reason, if I emit taichi/common.h's union_cast(), Metal failed
	// to compile. More strangely, if I copy the generated code to XCode as a
	// Metal kernel, it compiled successfully...
	static_assert(sizeof(T) == sizeof(G), "Size mismatch");
	return *reinterpret_cast<thread const T *>(&g);
	}
	inline int ifloordiv(int lhs, int rhs) {
	const int intm = (lhs / rhs);
	return (((lhs * rhs < 0) && (rhs * intm != lhs)) ? (intm - 1) : intm);
	}
	float fatomic_fetch_add(device float *dest, const float operand) {
	// A huge hack! Metal does not support atomic floating point numbers
	// natively.
	bool ok = false;
	float old_val = 0.0f;
	while (!ok) {
	old_val = *dest;
	float new_val = (old_val + operand);
	ok = atomic_compare_exchange_weak_explicit(
	(device atomic_int *)dest, (thread int *)(&old_val),
	*((thread int *)(&new_val)), metal::memory_order_relaxed,
	metal::memory_order_relaxed);
	}
	return old_val;
	})

are in native Metal, instead of being a string literal. (These helpers may be trivial, but I'm gonna add more runtime Metal code, and it's becoming a bit messy now.)

[k-ye]

Currently the Metal shaders are organized inside backends/metal/shaders. It can be edited as normal CPP files, but can still be emitted as string literals inside codegen. This reduces the necessity of this issue, so closing it now