[DoNotMerge] Implement withBackend.

[dan-zheng]

The design of withBackend is explained here.

Non-nested invocations withBackend should work as expected. (Test matrix-matrix-dot-with-backend added in this PR is verified to work on GPU.) Nested invocations of withBackend should also work, but have not yet been tested.

Please do not merge until these type erasure FIXMEs is fixed:

def transfer[T: Manifest](from: Backend, to: Backend)(data: T): T = {
  ...
  data match {
    case t: Tensor => transferTensor(t).asInstanceOf[T]
    // FIXME: Type erasure makes matching `Seq[Tensor]` ineffective.
    case tensors: Seq[Tensor] => tensors.foreach(transferTensor).asInstanceOf[T]
    // FIXME: "abstract type pattern Rep[Unit] is unchecked since it is eliminated by erasure"
    // Critical to fix this because it is a catch-all. Fix using `TypeTag`?
    // This case is exercised when `withBackend` is invoked with a function that returns Unit.
    case _: Rep[Unit] | Unit => data /* no-op */
    case _ =>
      System.err.println(s"'data' has unknown type: ${data.getClass.toString}")
      ???
  }
}

I tried to use TypeTag but it didn't work:

def transfer[T: TypeTag](from: Backend, to: Backend)(data: T): T = {
  ...
  typeOf[T] match {
    case t if t =:= typeOf[Rep[Unit]] => /* error: No TypeTag available for Rep[Unit] */
}

I encounter the same problem if I try to use Manifest:

def transfer[T: TypeTag](from: Backend, to: Backend)(data: T): T = {
  System.out.println(s"Is Rep[Unit]? ${manifest[T] == manifest[Rep[Unit]]}")
  // error: No Manifest available for Rep[Unit].
}

One idea is to define extra overloads for withBackend, but this doesn't work because it leads to ambiguity errors:

// This overload transfers the result.
def withBackend[T: Manifest, U: Manifest](b: Backend)(input: T)(f: T => U): U = { ... }
// This overload does not transfers the result.
def withBackend[T: Manifest](b: Backend)(input: T)(f: T => Rep[Unit]): Rep[Unit] = { ... }

def withCPU[T: Manifest, U: Manifest](input: T)(f: T => U): U = withBackend[T, U](BackendCPU())(input)(f)
def withCPU[T: Manifest](input: T)(f: T => Rep[Unit]): Rep[Unit] = withBackend[T](BackendCPU())(input)(f)

// error: ambiguous reference to overloaded definition
// Need to explicitly specify generic parameters, which is highly unusable.
withCPU(Tensor.ones(2, 2)) { ... }

Does anyone have ideas for working around type erasure? @GSAir

[feiwang3311]

is this branch still DoNotMerge?

[dan-zheng]

Yes, I haven't solved the type erasure bug yet.

The problem is I can't check whether a generic type parameter T is equal to Rep[Unit]. I tried manifest[T] == manifest[Rep[Unit]] and typeOf[T] =:= typeOf[Rep[Unit]] but they result in the error No (Manifest|TypeTag) available for Rep[Unit].

[GSAir]

I am not sure I understand the design of the transfer function. In which situation can we have both Tensors and Unit? (Or Rep[Unit])

To solve that problem don't use Rep[Unit] but Unit instead, this is the same thing.

[dan-zheng]

withBackend takes a closure argument of type T => U. I want to pattern match on T and U to verify whether they're valid and to determine how to perform transfer. U can be Tensor (or a product of Tensor), or Unit if the closure just does sth like print.

I originally tried comparing with Unit (e.g. manifest[U] == manifest[Unit]) but it didn't evaluate to true because the actual result type was lifted to Rep[Unit]. I'll try again.

[TiarkRompf]

Let's not spend cycles on this right now. Plenty of priority things to do.

[feiwang3311]

Let's drop this for now. We should try to get small CNN case running on GPU ASAP