self.assertEqual(f(x), scripted(x))
self.assertAllFused(scripted.graph_for(x))
+ @unittest.skipIf(IS_WINDOWS, "NYI: fuser support for Windows")
+ @unittest.skipIf(not RUN_CUDA, "fuser requires CUDA")
+ @skipIfRocm
+ def test_scalar_arg_cuda(self):
+ def fn_test_scalar_arg(x, p):
+ # type: (Tensor, float) -> Tensor
+ return p * (x * x + x)
+
+ x = torch.randn(4, 4, dtype=torch.float, device='cuda')
+ p = 3
+ scripted = torch.jit.script(fn_test_scalar_arg, (x, p))
+ self.assertEqual(fn_test_scalar_arg(x, p), scripted(x, p))
+ self.assertAllFused(scripted.graph_for(x, p))
+ x.requires_grad_(True)
+ out = scripted(x, p)
+ self.assertAllFused(scripted.graph_for(x, p), except_for=("aten::size", "prim::BroadcastSizes"))
+
@unittest.skipIf(IS_WINDOWS or IS_SANDCASTLE, "NYI: fuser support for Windows or Sandcastle")
@enable_cpu_fuser
def test_fuser_deduplication(self):
std::string generateKernel(
const std::string& name,
const Graph& graph,
- const std::vector<std::pair<const Value*, const TensorDesc>>& inputs,
+ const std::vector<std::pair<const Value*, const c10::optional<TensorDesc>>>& inputs,
const std::vector<std::pair<const Value*, const TensorDesc>>& outputs,
const bool use_cuda) {
TemplateEnv env;
// Lambda for writing arguments
auto emitFormal = [&](const Value* n, const TensorDesc& desc) {
- std::string tensor =
- "t" +
+ env.d(
+ "formal_index",
+ formals.size() +
+ 1); // + 1 because the first argument is the linearIndex
+ std::string tensor =
+ "t" +
+ std::to_string(
+ formals.size()); // can't be unique() because Param may be an output
+ const auto nDim = desc.nDim();
+ emitIndexingFor(tensorOffsets, tensor, nDim, desc.lastIsContiguous());
+ env.s("tensor", tensor);
+ env.d("nDim", nDim);
+ env.s("scalar_type", scalarTypeName(desc.scalar_type));
+ formals.push_back(
+ format("TensorInfo<${scalar_type},${nDim}> ${tensor}", env));
+ argument_loads.push_back(format(
+ "*static_cast<TensorInfo<${scalar_type},${nDim}>*>(args[${formal_index}])",
+ env));
+ };
+
+ auto emitScalarFormal = [&](const Value* n){
+ env.d(
+ "formal_index",
+ formals.size() +
+ 1); // + 1 because the first argument is the linearIndex
+ std::string scalar =
+ "s" +
std::to_string(
formals.size()); // can't be unique() because Param may be an output
- const auto nDim = desc.nDim();
- emitIndexingFor(tensorOffsets, tensor, nDim, desc.lastIsContiguous());
- env.s("tensor", tensor);
env.d(
"formal_index",
formals.size() +
1); // + 1 because the first argument is the linearIndex
- env.d("nDim", nDim);
- env.s("scalar_type", scalarTypeName(desc.scalar_type));
- formals.push_back(
- format("TensorInfo<${scalar_type},${nDim}> ${tensor}", env));
+ env.s("scalar", scalar);
+ env.s("scalar_type", variableType(n->type()));
+ formals.push_back(format("${scalar_type} ${scalar}", env));
argument_loads.push_back(format(
- "*static_cast<TensorInfo<${scalar_type},${nDim}>*>(args[${formal_index}])",
- env));
+ "*static_cast<${scalar_type}*>(args[${formal_index}])", env));
};
+
// Writes input parameters
for (const auto& input : inputs) {
- emitFormal(input.first, input.second);
+ if (input.second.has_value()){
+ emitFormal(input.first, *input.second);
+ } else {
+ emitScalarFormal(input.first);
+ }
}
// Writes output parameters
// Note: conversion from half is only supported for CUDA kernels.
// The conversion immediately converts fp16 inputs to float.
// Access for other types is common to CUDA and CPU kernels.
- const auto is_half = (input.second.scalar_type == at::ScalarType::Half);
- if (is_half) {
- AT_ASSERT(use_cuda);
- env.s(
- "access",
- format("__half2float(t${formal}.data[t${formal}_offset])", env));
- has_half_tensor = true;
+ if (input.second.has_value()) {
+ const auto is_half = input.second.has_value() && ((*input.second).scalar_type == at::ScalarType::Half);
+ if (is_half) {
+ AT_ASSERT(use_cuda);
+ env.s(
+ "access",
+ format("__half2float(t${formal}.data[t${formal}_offset])", env));
+ has_half_tensor = true;
+ } else {
+ env.s("access", format("t${formal}.data[t${formal}_offset]", env));
+ }
+ env.s("lhs_type", calcScalarTypeName(input.second.value().scalar_type));
} else {
- env.s("access", format("t${formal}.data[t${formal}_offset]", env));
+ env.s("access", format("s${formal}", env));
+ env.s("lhs_type", variableType(input.first->type()));
}
- env.s("lhs_type", calcScalarTypeName(input.second.scalar_type));
-
body << format("${lhs_type} ${node} = ${access};\n", env);
}
TORCH_API std::string generateKernel(
const std::string& name,
const Graph& graph,
- const std::vector<std::pair<const Value*, const TensorDesc>>& inputs,
+ const std::vector<std::pair<const Value*, const c10::optional<TensorDesc>>>& inputs,
const std::vector<std::pair<const Value*, const TensorDesc>>& outputs,
const bool use_cuda);
// Creates chunk and flattened input descriptions
std::vector<PartitionDesc> chunk_desc;
- std::vector<std::pair<const Value*, const TensorDesc>> flat_inputs;
+ std::vector<std::pair<const Value*, const c10::optional<TensorDesc>>> flat_inputs;
{
size_t input_index = 0;
for (const auto& p : graph->inputs()) {
+ if (p->type()->isSubtypeOf(FloatType::get())) {
+ flat_inputs.emplace_back(p, c10::nullopt);
+ }
if (!p->type()->isSubtypeOf(TensorType::get())) {
continue;
}
static bool shouldExpandArgs(
const KernelSpec& spec,
std::vector<at::Tensor>& args,
- std::vector<int64_t>& map_size) {
+ std::vector<int64_t>& map_size) {
return expandArgs(spec, args, map_size, /*dry_run=*/true);
}
const FusedKernel& fusion,
const at::Device device,
const at::ArrayRef<at::Tensor>& inputs,
+ const at::ArrayRef<IValue>& all_inputs,
std::vector<at::Tensor>& outputs) {
// Fails if fusion and given inputs disagree
AT_ASSERT(inputs.size() == fusion.inputDesc().size());
numel = computeNumel(map_size);
}
+ // compute number of scalar inputs and convert them to float
+ std::vector<float> scalar_inputs;
+ scalar_inputs.reserve(all_inputs.size());
+ for (auto const &input: all_inputs){
+ if (input.isDouble()) scalar_inputs.push_back(input.to<float>());
+ }
+
// Computes the storage needed to store TensorInfo structs for inputs and
// outputs.
size_t uncompressedDim = fusion.inputDesc().at(0).contiguity.size();
// A vector of arguments to the kernel (numel, *input_desc_s, *output_desc_s)
std::vector<void*> arguments;
- arguments.reserve(3 + flat_inputs_size + flat_outputs_size);
+ arguments.reserve(3 + scalar_inputs.size() + flat_inputs_size + flat_outputs_size);
arguments.push_back(&numel);
auto addTensorInfoRaw = [&](const TensorDesc& desc,
}
}
}
+ // Adds scalar arguments
+ for (float &s: scalar_inputs){
+ arguments.push_back(&s);
+ }
// Adds (flattened) output arguments
outputs.reserve(fusion.outputDesc().size());
// Launches fusion
std::vector<at::Tensor> raw_outputs;
- launchFusion(*(*maybe_kernel), device, inputs, raw_outputs);
+ launchFusion(*(*maybe_kernel), device, inputs, all_inputs, raw_outputs);
auto outputs = fmap(spec.outputMapAndSizes(), [&](const OutputMapAndSize& omap) {
if (omap.needsSumToSize()) {
if (!simple_mappable.find(node)) {
return false;
}
- // Check that all non-tensor inputs are constant
for (Value* input : node->inputs()) {
- if (input->type()->isSubtypeOf(TensorType::get())) {
+ if (input->type()->isSubtypeOf(TensorType::get()) || input->type()->isSubtypeOf(FloatType::get())) {
continue;
}
if (input->node()->kind() != prim::Constant) {
group->insertInput(tensor_insert_idx, input);
tensor_insert_idx++;
} else if (
- n->kind() == aten::_grad_sum_to_size &&
- input->type()->isSubtypeOf(ListType::ofInts())) {
+ (input->type()->isSubtypeOf(FloatType::get()) && input->node()->kind() != prim::Constant) ||
+ (n->kind() == aten::_grad_sum_to_size &&
+ input->type()->isSubtypeOf(ListType::ofInts()))) {
auto in_group = subgraph.addInput();
in_group->setType(input->type());
inputs_map[input] = in_group;
projects / platform / upstream / pytorch.git / commitdiff