[Feature Request] codegen should be able to produce scalar output

[jjsjann123]

Opening an issue to track this. Want to get an idea on how easy it is for us to plumb support for codegen to output scalar values.

Currently the cases we are looking at are where the input is just a scalar. So this can already be done with ExpressionEvaluation, the plumbing might be a little bit ugly on the runtime system, since we need to glue the results from the kernel execution; But I expect support on codegen would be equally tedious/ugly?!

I'm quite neutral on how this should be supported.

cc'ing @kevinstephano @naoyam

[naoyam]

I don't think it would be a lot of work for the codegen itself. Writes to an output scalar would be through a pointer, so the final code string would need to be updated, but pretty much that's it?

I'd expect it would mostly about the executor.

[jjsjann123]

Writes to an output scalar would be through a pointer.

good point... but that's a device pointer?! Can we return a host scalar that can be queried without mem copy?!?!

[naoyam]

Are you saying doing some computations entirely on host without generating a kernel? That's a very different story.

[jjsjann123]

Not necessarily a complete host thing, but we might want to just output some scalar values inside the kernel?
i.e.

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = makeSymbolicTensor(2);
  fusion.addInput(tv1);
  auto reshape_shape0 = IrBuilder::create<Int>();
  fusion.addInput(reshape_shape0);
  auto reshape_shape1 = IrBuilder::create<Int>();
  fusion.addInput(reshape_shape1);
  auto tv2 = reshape(tv0, {reshape_shape0, reshape_shape1});
  auto tv3 = add(tv1, tv2);
  fusion.addOutput(tv3);
  
  auto numel = add(reshape_shape0, reshape_shape1);  # scalar operation;
  fusion.addOutput(numel);  # output scalar value;

I start to remember why we don't have this in the first place. We would want the output numel to be on host if it's viewed as a scalar by the client.
I'm leaning towards that we need to plumb this in our integration and not having fusion/codegen to handle this, unless we want to compile cpu code, which might not be a bad idea neither. 😜

[naoyam]

So, I think we could relatively easily support returning a scalar from a kernel. It'd be computed on device, not host.

Looks like that's not what we want but we want to compute it on host? Like you mentioned, we could extend ExpressionEvaluator, but it currently only supports some limited types of ops and data types.

Why does it need to be computed on host?

[jjsjann123]

Sorry for the confusion. Not necessarily computed on host, but we need to return a host value. I thought scalars usually are considered host value and it's a fair game to peek at what scalar values are.

correct me if I'm wrong, but I don't think we can pass by reference on a __global__ function, so it doesn't look clear to me how we can return a host scalar from the kernel.

[naoyam]

I think we could return a scalar by some workaround, for example storing it in a zero-dimensional tensor. We would definitely need to do something specific for it on the host side as well, but I think it's doable.

So, a kernel would have a parameter like Tensor<float, 0> T0, which would be used like T0[0] = some_scalar_value; in the kernel, so after the kernel completion, some_scalar_value could be retrieved from the T0 parameter.

[jjsjann123]

I still don't see how this would work.

Tensor<float, 0> T0, which would be used like T0[0] = some_scalar_value; ... some_scalar_value could be retrieved from the T0 parameter.

T0 needs to be a pointer on GPU, in order for kernel to write values via T0[0] = .... Afterwards, if we want to retrieve the value to use on host code, it's a device sync.

[jjsjann123]

The closest thing we have is CpuScalarTensor support. But that's basically a hack leveraging the memcpy on function arguments: [reference https://docs.nvidia.com/cuda/cuda-c-programming-guide/#global-function-argument-processing]

Here's our code:

Fuser/csrc/executor_kernel_arg.h

Lines 308 to 312 in e1a420e

	template <size_t size>
	struct CpuScalarTensorArg : public ArgAbstract {
	std::array<std::byte, size> instance_;
	DEF_HELPEE_FUNC(CpuScalarTensor, instance_)
	};

Fuser/runtime/tensor.cu

Lines 30 to 38 in e1a420e

	// Specialization for 0-dim case that's easy to pass in a CPU based tensor.
	template <typename T>
	struct CpuScalarTensor {
	__device__ T& operator[](int i) {
	return data;
	};
	T data;
	};

But I'm not aware of anything like that for us to return anything from device.

[naoyam]

T0 needs to be a pointer on GPU, in order for kernel to write values via T0[0] = .... Afterwards, if we want to retrieve the value to use on host code, it's a device sync.

Yes, as long as it's computed on the device, there must be a device sync.

[kevinstephano]

I thought, with Pytorch, that moving a Scalar from the GPU to CPU was done explicitly and Scalars in Pytorch actually aren't used to describe GPU calculations. It is always 1-element tensors with the usage of item() to move the item from the GPU->CPU, therefore, the usage of a synchronization is explicit.

[jjsjann123]

I feel that we have been confusing scalar tensor with scalar here in the whole thread.

A not-so-good example here: #315 (comment)
Or rather, if we just want to output a scalar value in nvfuser.

  s0 = fd.define_scalar()
  s1 = fd.define_scalar()
  s2 = fd.ops.add(s0, s1)
  fd.add_output(s2)  # we can use a better error message here today. More importantly, we might want to actually support this.

[naoyam]

WIP PR #4162