Output allocate patch (#1790)

Caching strides along with sizes. This is to support current expand, which introduces non-contiguous output tensor

Author: jjsjann123

torch/csrc/jit/codegen/cuda/executor.cpp

@@ -762,8 +762,9 @@ std::vector<at::Tensor> FusionExecutor::runFusion(
       if (outputs.empty()) {
         FUSER_PERF_SCOPE("ExecutorRunFusion::OutputAlloc");
         for (const auto i : c10::irange(executor_entry->output_sizes.size())) {
-          allocated_outputs.push_back(at::native::empty_cuda(
+          allocated_outputs.push_back(at::native::empty_strided_cuda(
               executor_entry->output_sizes[i],
+              executor_entry->output_strides[i],
               executor_entry->output_types[i],
               c10::nullopt,
               options_.device,
@@ -934,6 +935,7 @@ std::vector<at::Tensor> FusionExecutor::runFusion(
       executor_entry->io_alias_indices = alias_indices;
       for (const auto& output : allocated_outputs) {
         executor_entry->output_sizes.push_back(output.sizes().vec());
+        executor_entry->output_strides.push_back(output.strides().vec());
         executor_entry->output_types.push_back(output.scalar_type());
       }
 

torch/csrc/jit/codegen/cuda/executor.h

@@ -74,6 +74,7 @@ class TORCH_CUDA_CU_API FusionExecutor : public NonCopyable {
     LaunchParams launch_params;
     std::vector<std::pair<int, int>> io_alias_indices;
     std::vector<std::vector<int64_t>> output_sizes;
+    std::vector<std::vector<int64_t>> output_strides;
     std::vector<at::ScalarType> output_types;
     std::vector<std::vector<int64_t>> buffer_sizes;
     std::vector<at::ScalarType> buffer_types;

torch/csrc/jit/codegen/cuda/test/test_gpu_view.cpp

@@ -972,6 +972,40 @@ TEST_F(NVFuserTest, FusionComputeAtRootDomainMapWithView_CUDA) {
       tv1->axis(1)->toString());
 }
 
+TEST_F(NVFuserTest, FusionExpandRepro_CUDA) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  const std::vector<int64_t> input_shape1{4, 1, 1};
+  const std::vector<int64_t> input_shape2{4, 3, 2};
+
+  auto tv0 = makeConcreteTensor({-1, 1, 1});
+  fusion.addInput(tv0);
+  auto tv1 = makeSymbolicTensor(3);
+  fusion.addInput(tv1);
+
+  auto tv2 = expand_as(tv0, tv1);
+  fusion.addOutput(tv2);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  at::Tensor at_x = at::randn(input_shape1, options);
+  at::Tensor at_y = at::randn(input_shape2, options);
+  std::vector<IValue> aten_inputs = {at_x, at_y};
+
+  FusionExecutor fe;
+  fe.compileFusion(&fusion);
+  LaunchParams l_params;
+  auto outputs = fe.runFusion(aten_inputs, {}, l_params, 0);
+
+  auto out = at_x.expand_as(at_y);
+
+  testValidate(&fusion, outputs, aten_inputs, {out}, __LINE__, __FILE__);
+
+  // second run to verify cached output allocation
+  outputs = fe.runFusion(aten_inputs, {}, l_params, 0);
+  testValidate(&fusion, outputs, aten_inputs, {out}, __LINE__, __FILE__);
+}
+
 } // namespace jit
 } // namespace torch
 #endif // #if defined(USE_CUDA)