search

comp-imaging / ProxImaL

A domain-specific language for image optimization.
MIT License
112 stars 29 forks source link

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

Open soufianekhiat opened 3 years ago

soufianekhiat commented 3 years ago

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 https://github.com/comp-imaging/ProxImaL/issues/57#issuecomment-811969044

SteveDiamond commented 3 years ago

@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 commented 2 years ago

@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 commented 2 years ago

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?

https://github.com/comp-imaging/ProxImaL/blob/b7c752c3b798634c17c11501f0e0312db3fc3038/proximal/utils/cuda_codegen.py#L275-L281

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 commented 1 year ago

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().

https://github.com/comp-imaging/ProxImaL/blob/e3d08d73483f7173f8c3e321c723e8e9cb8c725c/proximal/templates/problem-definition.h.j2#L35-L44