Start cuBLAS backend support.

[dan-zheng]

A first step towards #8.

• Add cuBLAS/cuDNN code generators and drivers.
• Implement dot for BackendCublas, refactor Backend trait.
• Add cuBLAS code generation test.
• Implement withBackend device placement function.

[dan-zheng]

Relevant info for code review:

• There's a lot of code duplication in the code generators and DSL drivers. Let's tackle that in follow-up commits.
• The cuBLAS tests fail on my GPU machine for some reason. If someone else has a GPU, please try to run the tests. I'm not sure why this occurs, cudaMallocManaged is used to allocate tensors.

$ ./snippet asdf
[1]    1914 bus error (core dumped)  ./snippet asdf
$ gdb snippet
...
Thread 1 "snippet" received signal SIGBUS, Bus error.
Snippet (x0=<optimized out>) at snippet.cpp:144
144	float x32 = x30 - x31;

• It would be nice to set up GPU CI, but I don't think there are free services providing it. One solution adopted by other codebases is Jenkins CI + a personal GPU machine.

[TiarkRompf]

I wanted to try it on my iMac but it seems like there is no CUDA for Mojave yet ...

[TiarkRompf]

This looks like a great start! Couple of comments:

• Let's separate out the withBackend functionality into a separate PR. I think we need to get a basic GPU-only version working first and I expect getting all the transfer logic right with multiple backends will take several rounds of debugging, and perhaps we want to consider alternative designs as well.
• Instead of globally rewiring the codegen for ArrayNew, I'd propose to introduce a malloc like operation in the backend trait, and some of way of freeing memory in a coarse grained way (freeAll or something like this). This makes memory management for Tensors a responsibility of the CPU/GPU backend.
• The code duplication between the codegen back-ends in indeed unfortunate. It seems like the GPU-enabled codegen templates only need to add things, so I'd propose to refactor this with proper extension points (i.e. something like `includes += "<cuda_runtime.h>"

[TiarkRompf]

Rewiring NewArray globally is problematic for various reasons. I think a better way would be to have explicit allocation operations for GPU. I'm also concerned about performance of CUDA's Unified Memory layer.

[dan-zheng]

This makes complete sense. Adding a mallocTensor operation in Backend seems like a good approach. mallocTensor returns Rep[Array[Float]] and should replace invocations of NewArray[Float].

I'll try to flesh this out.

[TiarkRompf]

These don't have to be in the codegen template, but could be explicitly invoked by the CUBLAS backend trait

[dan-zheng]

Could you please clarify the design you have in mind?

[TiarkRompf]

I was thinking explicit init and cleanup methods in the backend interface.

[dan-zheng]

I wonder how to hook up init/cleanup Backend-defined methods with the codegen template. AFAICT, the DslDriverXXX trait (which contains DslGenXXX) and TensorExp trait (which contains Backend) are decoupled, so I'm not sure how to invoke Backend-defined methods in DslGenXXX, where codegen templates are defined.

Perhaps I'm missing something, do you have some idea on how to hook up things?

[TiarkRompf]

Ah, I didn't mean that the code generator would invoke these methods, but rather that they should be called from the frontend somewhere. Does that help clear things up?

[dan-zheng]

I'm sorry, I'm still a bit confused. Do you mean something like this?

val foo = new DslDriverCublas[String, Unit] with TensorExp {
  backend = new BackendCublas

  @virtualize
  def snippet(x: Rep[String]): Rep[Unit] = {
    backend.init()
    // ...
    Tensor.assertEqual(...)
    backend.cleanup()
  }
}

If so, I don't quite believe this is the best approach. I would argue backend setup/cleanup are implementation details that should be hidden from DSL users if possible. Also, if backend.init() initializes cublasHandle_t handle, then it really should be performed before any computation in the snippet. Burdening DSL users with that responsibility is suboptimal.

[TiarkRompf]

I wasn't necessarily thinking users should call these, but for example init/reset could be ideally hidden in a withGPU { ... } construct or even as part of the DSLDriver class before and after invoking snippet.

[TiarkRompf]

In general, my point here is that codegen templates should be as minimal as possible, and we should define the application logic as far as possible using metaprogramming.

[dan-zheng]

I see, that makes sense - thanks for the clarification.

Along those lines, a straightforward idea is to create a function wrapping snippet and compiling that function during codegen:

// `wrapper` could use a better name.
def wrapper(x: Rep[String]): Rep[Unit] = {
  backend.init()
  val result = snippet(x)
  backend.cleanup()
  return result
}

// During code generation:
lazy val code: String = {
  val source = new java.io.StringWriter()
  // Compile wrapper function.
  codegen.emitSource(wrapper, "Snippet", new java.io.PrintWriter(source))
  source.toString
}

If this seems reasonable (and it works), I'll implement it.

Note: this wrapper does add complexity to the simple mental model of "the snippet function is compiled, WYSIWYG". It would be good to make setup and cleanup idempotent if possible, in case DSL users call the functions themselves within snippet.

[TiarkRompf]

That looks good, yes. Let's not worry about idempotency at this point.

[dan-zheng]

My next steps:

• Add explicit tensor allocation (and freeing) operations.
• Remove backend setup/cleanup from codegen template.