Generating halide-optimized routine from K.fwd and K.adj

[soufianekhiat]

Fork from other discussion:
Component level specification

We could achieve the same pattern with ProxImaL. In ideal world that would be like:

1. Describe the Problem [ProxImaL]
2. Compile to have the solver as Halide Kernel (C++ or Python). [ProxImaL ->Halide]
3. Use Halide to build a scheduler or auto_schedule for {CUDA, OpenCL, x86, ...} [Halide]

I suggest that we fork out a separate Github issue to discuss this. @jrk and @SteveDiamond can chime in as well.

My first thought is that ProxImaL, just like FlexISP, generates a cyclic compute graph. The lack of compute cycles in Halide is a fundamental strength to customize the compute schedule (or weakness in your use-case?).

As the authors suggested in the original article, I think the ProxImaL project can still grow by

1. fusing more proximal.linops.* nodes in the proximal.CompGraph into a single Halide pipeline; and
2. fusing the individual proximal update steps in ADMM into Halide pipeline.

But eventually the Halide-generate code will have to hand back control to the ADMM while-loop to complete the "cycle". Whether the while-loop has to be coded in C++ or to remain in Python, I do not know.

Originally posted by @antonysigma in #57 (comment)

[SteveDiamond]

@soufianekhiat this proposal makes a lot of sense as a future development goal. When we created proximal it was envisioned more as an exploratory toolbox, and so the dependency on Python is there to make it easy to try out different approaches. But a more "production" use case would be very interesting as well.

[antonysigma]

@SteveDiamond revisiting the proposal 1 year later, may I suggest changing the issue heading from "Component Level Specification" to "Generating halide-optimized routine from K.fwd and K.adj" ?

Here, K means the Proximal generated mapping z = K x. Feel free to paraphrase the title so it makes sense.

[antonysigma]

The first impression is that the mapping K, being a CompGraph, contains subgraphs that cannot be converted into halide code. So, the CompGraph will need to have graph query and manipulation algorithms to fuse individual halide-friendly sub-graphs into the corresponding (single) node.

I see that these graph algorithms, in proximal.utils.cuda_codegen.CudaSubgraph, are already implemented and bundled with pycuda_codegen. Is there a possibility to turn the code into a base class? So that both pycuda_codegen and halide_codegen can derive from the base class?

ProxImaL/proximal/utils/cuda_codegen.py

Lines 275 to 281 in b7c752c

	class CudaSubGraph:
	"""
	This class splits a linear operation graph into multiple computable subgraphs.
	The constructor is intended to be called with the end node of a comp graph.
	The comp graph is split in pieces with single kernels, and the pieces
	are recursively stored in the "dependent_subgraphs" member.
	"""

Otherwise, it sound too much code bloat to me.

Another idea: call external library Networkx to create, query, and manipulate our CompGraph. How much effort is required?

[antonysigma]

Post PR #65 update:

This issue is regarding how to modify proximal to implement the following code, namely K.halide_forward_graph.gen_code() and K.halide_adjoint_graph.gen_code().

ProxImaL/proximal/templates/problem-definition.h.j2

Lines 35 to 44 in e3d08d7

	FuncTuple<N> forward(const Func& z) {
	/* Begin code-generation */
	{{ K.halide_forward_graph.gen_code() }}
	/* End code-generation */
	}
	Func adjoint(const FuncTuple<N>& u) {
	/* Begin code-generation */
	{{ K.halide_adjoint_graph.gen_code() }}
	/* End code-generation */