mul: convert inputs to result type.

[ysiraichi]

Fix: #7084

This PR fixes a data-type related problem for mul operation. It does so by creating a structure OpConfig that behaves similarly to DoBinaryOp. The difference is that it takes care of pre/post-processing of inputs and outputs, casting them to the correct data-type.

Problem

t = torch.rand(10, dtype=torch.half).to(xm.xla_device())
s = torch.tensor(10, dtype=torch.double).to(xm.xla_device())
out = torch.mul_(t, s)

• Tensor.mul_ is dispatched to its CompositeExplicitAutograd kernel
  • It wraps the scalar into a tensor, and calls torch.mul (functional version)
• DoBinaryOp is called
  • Computes at::result_type (let's call it common_dtype) and passes it on to bin_op
  • Note that UnwrapNumber does nothing, since s is a tensor with is_wrapped_number_ unset
• The computed common_dtype is passed on to tensor_methods::mul
  • Creates an IR node with data-type common_dtype
  • Does nothing with its inputs
• Later, when BuildMul is called, we have 2 XlaOp with different data-types: f16 and f64
  • BuildMul promotes f16 to f64
• The output is common_dtype (torch.float16), but the actal XlaOp is f64

Solution

Following PyTorch behavior [1, 2, 3], I created OpConfig: a structure that let us specify common pre/post-processing on inputs and outputs.

Affected Models

• timm_nfnet (training+nondynamo)

cc @miladm @JackCaoG @lezcano

[JackCaoG]

hmm I am surprise that bf16 and f64's at::result_type is f64..

[ysiraichi]

hmm I am surprise that bf16 and f64's at::result_type is f64.

This is a bit confusing, so let me try and clarify the cases, given op(bf16, f64):

bf16	f64	at::result_type	PromoteType
tensor	scalar	bf16	f64
scalar	tensor	f64	f64
tensor	tensor	f64	f64

[lezcano]

The tl;dr is: Scalars don't promote unless they are of a different kind.
Here are the exact rules: https://pytorch.org/docs/stable/tensor_attributes.html

[bhavya01]

@ysiraichi This PR is failing the mul operation on TPUs.

It fails this check https://github.com/pytorch/xla/blob/master/torch_xla/csrc/aten_xla_type.cpp#L161

>>> import torch
>>> import torch_xla
>>> x = torch.tensor([1,2,3]).to('xla')
>>> y = torch.tensor([2,4,5]).to('xla')
>>> x*y
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: torch_xla/csrc/aten_xla_type.cpp:161 : Check failed: it != inputs_.end() 
*** Begin stack trace ***
	tsl::CurrentStackTrace[abi:cxx11]()
	torch_xla::XLANativeFunctions::mul(at::Tensor const&, at::Tensor const&)
	
	at::_ops::mul_Tensor::redispatch(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&)
	
	
	at::_ops::mul_Tensor::call(at::Tensor const&, at::Tensor const&)
	
	PyNumber_Multiply
	_PyEval_EvalFrameDefault
	
	PyEval_EvalCode
	
	_PyRun_InteractiveLoopObject
	
	PyRun_AnyFileExFlags
	
	Py_BytesMain
	
	__libc_start_main
	
*** End stack trace ***

[ysiraichi]

Apparently, you are not the only one: #7266

[vanbasten23]

This PR is also impacting DDP:

[JackCaoG]

Let me take a look this afternoon

[JackCaoG]

I can't repo this issue but I do see that half of our internal TPU test crashed because of this. Let me revert this pr for now while figuring out what happened.