Fix concretization and execution with padded expanded empty inputs by jacobhinkle · Pull Request #876 · NVIDIA/Fuser

jacobhinkle · 2023-09-13T14:15:33Z

This PR:

Changes the analyzeResizes pass in concretization to inspect expanded extents
Changes root->rfactor propagation in concretization to return early when IterType is already concretized. This avoids propagating Broadcast in this test when we have already marked the resized ID as Iteration.
Changes PrecomputedValues::bindInputs to bind not only metadata but also the actual TensorView arguments.

I noticed that the ExpressionEvaluator used during compilation contained more bound scalars than the one used at execution where we fail to evaluate the extent. We had i0 and i2 bound at execution, but we did not have T0 bound, so we could not compute getMetaData(T0). At compilation, T0 was bound so there was no problem until execution.

Note that at compilation, we use auto expr_eval = executor_utils::bindInputs(args, kernel); whereas at compilation we use evaluatorPrecomputedValues()->bindInputs(args);. The difference is that PrecomputedValues::bindInputs will call bindTensorMetaData instead of binding the actual tensor. This PR also binds the actual tensor in addition to its metadata in that method.

Fixes #870

jacobhinkle · 2023-09-13T16:50:18Z

!build

jacobhinkle · 2023-09-15T11:45:45Z

!build

zasdfgbnm · 2023-09-17T17:17:39Z

Could you check the benchmark result for the cpp and python benchmarks in #749

jjsjann123 · 2023-09-18T18:07:04Z

Could you check the benchmark result for the cpp and python benchmarks in #749

Out of curiosity, are we expecting trivial changes in this PR to have big impact on overhead?!

zasdfgbnm · 2023-09-18T18:45:49Z

Could you check the benchmark result for the cpp and python benchmarks in #749

Out of curiosity, are we expecting trivial changes in this PR to have big impact on overhead?!

No, I don't expect so. Just want to confirm.

jacobhinkle · 2023-09-18T20:02:52Z

Could you check the benchmark result for the cpp and python benchmarks in #749

The shape inference benchmarks look like this on main:

-------------------------------------------------------------------------------------------
Benchmark                                                 Time             CPU   Iterations
-------------------------------------------------------------------------------------------
LayerNormBackward_ShapeInference                       7898 us         7898 us           84
LayerNormForward_ShapeInference                        1994 us         1994 us          353
LayerNormBackward_NoShapeInferenceCachedBaseline       1339 us         1339 us          521
LayerNormForward_NoShapeInferenceCachedBaseline         467 us          467 us         1499

and like this on this branch:

-------------------------------------------------------------------------------------------
Benchmark                                                 Time             CPU   Iterations
-------------------------------------------------------------------------------------------
LayerNormBackward_ShapeInference                       8194 us         8194 us           88
LayerNormForward_ShapeInference                        1948 us         1948 us          343
LayerNormBackward_NoShapeInferenceCachedBaseline       1263 us         1263 us          547
LayerNormForward_NoShapeInferenceCachedBaseline         461 us          461 us         1475

Where can I find the python benchmarks?

zasdfgbnm · 2023-09-18T21:53:13Z

Where can I find the python benchmarks?

The python one is in #749 (comment)

zasdfgbnm · 2023-09-18T21:55:43Z

csrc/evaluator_common.cpp

-    } else {
-      bindValue(input->evaluatorIndex(), *args[i]);
    }
+    bindValue(input->evaluatorIndex(), *args[i]);


Oops, I think I remembered why I felt worried about this. At some point, I tried to do the same thing, but eventually gave up. Because doing so will make PrecomputedValues hold ownership of tensors, and as a result, memory usage will go up a lot.

Ah that makes sense. So we don't want PrecomputedValues or ExpressionEvaluator to own any tensor buffers, but it might be nice to still hold a reference them so that we can evaluate GetMetaData.

We could bind a Pointer instead but I think we might still wind up owning the tensor in the PolymorphicValue that points to. We could also introduce a TensorMeta class that doesn't own its data but holds sizes, strides, dtype, device, and data pointer. I am not sure whether that would slow down evaluation of CPU aten ops though, since we'd need to construct an at::Tensor when evaluating those ops then.

OK I'm trying to understand the existing TensorMetaData class now. I think maybe we could bind these values as the the outputs of IrBuilder::metadataExpr(tv) instead of binding to tv directly?

In 26a1f44 I changed this from binding TVs to binding extents and TensorMetaData. That way we can evaluate the extents directly and also the metadata expressions.

I think it is OK for ExpressionEvaluator to bind tensors, actually, we are evaluating tensor expressions. Because we do not carry instances of ExpressionEvaluator for long time.

But if we bind tens to tv in ExpressionEvaluator ee then later we bind a PrecomputedValues pv to ee, then we evaluate an expression with tv as an input, won't that populate pv with tens?

No, I don't think so. PrecomputedValues -> ExpressionEvaluator is one-directional.

Fuser/csrc/expr_evaluator.cpp

Lines 162 to 166 in 5ff26b9

if (precomputed_values_ && precomputed_values_->ready()) {

if (precomputed_values_->getMaybeValueFor(value).hasValue()) {

return precomputed_values_->getMaybeValueFor(value);

}

}

Got it. I'll revert the last PR then and add a note that we bind to ExpressionEvaluator but should avoid binding to PrecomputedValues.

jacobhinkle · 2023-09-19T13:47:29Z

The python one is in #749 (comment)

Sorry I missed that.

[---------------------  ---------------------]
           |  Before #876 PR  |  After #876 PR
1 threads: -----------------------------------
      2    |       17.7       |       18.7    
      4    |       22.6       |       23.8    
      8    |       33.0       |       34.5    
      16   |       52.8       |       53.4    
      32   |       98.6       |       98.4    
      64   |      192.4       |      194.2    
      128  |      411.1       |      406.6    

Times are in microseconds (us).

jacobhinkle · 2023-09-19T15:35:21Z

!build

This reverts commit 26a1f44.

jacobhinkle · 2023-09-21T14:24:57Z

csrc/polymorphic_value.h

+  //! This is a shallow comparison operator that just checks whether we point to
+  //! the same exact Struct
+  bool operator==(const StructHandle& other) const {
+    return struct_ptr_ == other.struct_ptr_;
+  }
+


After binding a TensorMetaData handle in PolymorphicValues, without this I was hitting the following error:

what(): ret.has_value() INTERNAL ASSERT FAILED at "/opt/pytorch/nvfuser/lib/dynamic_type/src/dynamic_type/dynamic_type.h":741, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues. Cannot compute N7nvfuser12StructHandleE == N7nvfuser12StructHandleE : incompatible type
Exception raised from operator== at /opt/pytorch/nvfuser/lib/dynamic_type/src/dynamic_type/dynamic_type.h:741 (most recent call first):

Adding this shallow pointer comparison fixes it, but @zasdfgbnm please let me know if this is safe to add.

Alternatively, I could add another special case in PolymorphicValue_functions::isSame?

I think this makes sense

jacobhinkle · 2023-09-21T14:39:05Z

csrc/scheduler/normalization_inner.h

 };

-TORCH_CUDA_CU_API std::shared_ptr<ReductionParams> getInnerPersistentHeuristics(
+std::shared_ptr<ReductionParams> getInnerPersistentHeuristics(


Lintrunner was failing...

jacobhinkle · 2023-09-21T14:39:23Z

!build

jacobhinkle · 2023-09-21T17:41:58Z

csrc/evaluator_common.cpp

+  metadata->logical_size = tensor.sizes();
+  metadata->logical_stride = tensor.strides();
+  metadata->alloc_size = tensor.sizes();
+  metadata->alloc_stride = tensor.strides();


How should we handle these lines? Currently this is failing in the test AllocationDomainTest.NHWC4d_To_NHWC4d_CUDA. In that test we have a memcpy fusion with channels-last inputs and outputs and we schedule vectorization of size 4. I get the error

terminate called after throwing an instance of 'nvfuser::nvfError'
what(): is_contiguous || size == 1 || is_expanded_broadcasting || (still_rightmost && stride == 1) || (!still_rightmost && stride % word_size == 0) INTERNAL ASSERT FAILED at "/opt/pytorch/nvfuser/csrc/executor_utils.cpp":545, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues. Vectorization of T0_g[ iS15{( (( (( getMetaData(T0) )).logical_size ))[0] )}, iS16{( (( (( getMetaData(T0) )).logical_size ))[1] )}, iS17{( (( (( getMetaData(T0) )).logical_size ))[2] )}, iS18{(
(( (( getMetaData(T0) )).logical_size ))[3] )} ] with word size 4 not possible due to invalid stride. Domain: iS18{( (( (( getMetaData(T0) )).logical_size ))[3] )}, stride: 21, cur_contig_stride: 1, non contig due to slice: 0

Inspecting this a bit I see that logical_stride and alloc_stride are the same because of the line above, but they should differ for this channels-last example.

Sorry, missed this message. Is this still the case?

Oh this comment is outdated now since I realize we should just use ExpressionEvaluator to get the metadata object here instead of trying (and failing) to replicate it as I did previously.

jacobhinkle · 2023-09-21T19:55:44Z

!build

jacobhinkle · 2023-09-22T12:40:10Z

I think this fixes the original problem. However, in doing so it exposes #596 in the pytest_ops.py slice correctness tests. Those tests include a fusion like this:

Fusion before concretization:
Inputs:  T0_g[ iS0{i0}, iS1{i1}, iS2{i2} ], float
Outputs:
  T1_g[ ?S4{1}rf, ?S6{6}rf, ?S8{5}rf ], float
%kernel_math {
T1_g[ ?S4{1}rf, ?S6{6}rf, ?S8{5}rf ]
   = slice( T0_g[ iS0{i0}, iS1{i1}, iS2{i2} ], { {1, 2, 1} {0, 6, 1} {3, 8, 1} } )
}

The output IterDomains are currently symbolic, even though the extents are constant. These extents should actually depend on the inputs since if they are zero, even ?S4 might have zero extent. The problem at hand is that before this PR (e.g. on main commit b96d569) this is concretized as

Concretized Fusion:
Inputs:
  T0_g[ iS0{i0}, iS1{i1}, iS2{i2} ], float
Outputs:
  T1_g[ iS12{1}rf, iS13{6}rf, iS14{5}rf ], float

%kernel_math {T1_g[ iS12{1}rf, iS13{6}rf, iS14{5}rf ]
   = slice( T0_g[ iS0{i0}, iS1{i1}, iS2{i2} ], { {1, 2, 1} {0, 6, 1} {3, 8, 1} } )}

That is, the output iterdomain is concretized as Iteration since the input is marked as Iteration with size 1. This is somewhat related to #900 since in that proposal we would concretize ?S0 as broadcast as well. Anyway, this is the reverse problem from what is fixed in this PR, but the erroneous concretization was helping avoid triggering #596 previously. I would like to wait until a solution is found for that issue before merging this. If instead the current fix is needed urgently, I can disable the slice tests temporarily; i think this is the only failure in our test suite.

zasdfgbnm · 2023-09-22T16:01:52Z

csrc/evaluator_common.cpp

+  // we do not want them to own large objects.
+  // To do this we create a temporary ExpressionEvaluator so that we can compute
+  // the metadata once, then save it
+  ExpressionEvaluator ee;


There's some sharp edges here. For example, if between rfactor domain and root domain, there's a split whose factor is a fusion input, then this will fail, because we do not have it binded. But for now, I think it's fine because we do not have such case. Let's leave this as is. In the future we might need to refactor if this becomes an issue.

Great point. We should at minimum bind this here so that any precomputed values that are already bound will be available to ee. I'll also add a comment about this pitfall.

jjsjann123 · 2023-09-22T16:22:19Z

...
That is, the output iterdomain is concretized as Iteration since the input is marked as Iteration with size 1. This is somewhat related to #900 since in that proposal we would concretize ?S0 as broadcast as well. Anyway, this is the reverse problem from what is fixed in this PR, but the erroneous concretization was helping avoid triggering #596 previously. I would like to wait until a solution is found for that issue before merging this. If instead the current fix is needed urgently, I can disable the slice tests temporarily; i think this is the only failure in our test suite.

>Fusion before concretization:
Inputs:  T0_g[ iS0{i0}, iS1{i1}, iS2{i2} ], float
Outputs:
  T1_g[ ?S4{1}rf, ?S6{6}rf, ?S8{5}rf ], float
%kernel_math {
T1_g[ ?S4{1}rf, ?S6{6}rf, ?S8{5}rf ]
   = slice( T0_g[ iS0{i0}, iS1{i1}, iS2{i2} ], { {1, 2, 1} {0, 6, 1} {3, 8, 1} } )
}

These extents should actually depend on the inputs since if they are zero, even ?S4 might have zero extent.

IIUC, all extends from a slice is dependent on the actual input sizes. i.e. if we slice out of range, the behavior is that it'll just crop at the boundary so output range is indeed dynamic.

So this is backing up the idea in #900 😄

jacobhinkle · 2023-09-22T16:38:07Z

all extends from a slice is dependent on the actual input sizes

This is true. It's incorrect now as we assume 0 <= start <= stop <= extent in slice, but #460 would replace these with the proper expressions.

jacobhinkle · 2023-09-27T14:36:08Z

!build

jacobhinkle · 2023-10-04T17:09:04Z

!build

jacobhinkle added 2 commits September 13, 2023 10:14

Add test

ea50620

Use expanded extent and mutated out_id

ed8f5b1

jacobhinkle changed the title ~~Fix 870~~ [WIP] Fix 870 Sep 13, 2023

Bind tensorviews in PrecomputedValues::bindInputs

2ae5bf9

jacobhinkle changed the title ~~[WIP] Fix 870~~ Fix concretization and evaluation with padded expanded empty inputs Sep 13, 2023

jacobhinkle requested review from jjsjann123 and zasdfgbnm September 13, 2023 16:51

jacobhinkle marked this pull request as ready for review September 13, 2023 16:51

Merge branch 'main' into fix_870

d615688

jacobhinkle changed the title ~~Fix concretization and evaluation with padded expanded empty inputs~~ Fix concretization and execution with padded expanded empty inputs Sep 15, 2023

Merge branch 'main' into fix_870

9f379eb

zasdfgbnm reviewed Sep 18, 2023

View reviewed changes

Bind extents and TensorMetaData instead of tensor itself

26a1f44

jacobhinkle added 2 commits September 19, 2023 14:44

Revert "Bind extents and TensorMetaData instead of tensor itself"

a21ae8c

This reverts commit 26a1f44.

Add comment about binding tensors to ExpressionEvaluator

ef85ea1

zasdfgbnm mentioned this pull request Sep 19, 2023

Enable evaluation of all ternary ops in ExpressionEvaluator #895

Merged

jacobhinkle added 3 commits September 21, 2023 08:10

Only bind tv metadata in PrecomputedValues

3f7e19f

Add C++ test

41d6b2d

Add StructHandle::operator==

84ed3b4

jacobhinkle commented Sep 21, 2023

View reviewed changes

Merge remote-tracking branch 'origin/main' into fix_870

887e093

Silence lint error about TORCH_CUDA_CU_API

c9f174a

jacobhinkle commented Sep 21, 2023

View reviewed changes

jacobhinkle and others added 2 commits September 21, 2023 15:55

Use ExpressionEvaluator to build metadata for PV

70c06ba

Merge branch 'main' into fix_870

9624ca3

zasdfgbnm approved these changes Sep 21, 2023

View reviewed changes

Merge remote-tracking branch 'origin/main' into fix_870

b66d638

jacobhinkle mentioned this pull request Sep 22, 2023

Resolve broadcasts resulting from a PadOp #610

Merged

zasdfgbnm reviewed Sep 22, 2023

View reviewed changes

jacobhinkle added 3 commits September 22, 2023 12:57

Bind this and add note about dangers

139580f

Merge remote-tracking branch 'origin/main' into fix_870

12feadd

Merge remote-tracking branch 'origin/main' into fix_870

205c8b2

Merge branch 'main' into fix_870

30807f2

jacobhinkle merged commit d70e19e into main Oct 4, 2023

jacobhinkle deleted the fix_870 branch October 4, 2023 19:14

	if (precomputed_values_ && precomputed_values_->ready()) {
	if (precomputed_values_->getMaybeValueFor(value).hasValue()) {
	return precomputed_values_->getMaybeValueFor(value);
	}
	}

Conversation

jacobhinkle commented Sep 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jacobhinkle commented Sep 13, 2023

Uh oh!

jacobhinkle commented Sep 15, 2023

Uh oh!

zasdfgbnm commented Sep 17, 2023

Uh oh!

jjsjann123 commented Sep 18, 2023

Uh oh!

zasdfgbnm commented Sep 18, 2023

Uh oh!

jacobhinkle commented Sep 18, 2023

Uh oh!

zasdfgbnm commented Sep 18, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jacobhinkle Sep 19, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jacobhinkle Sep 19, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jacobhinkle commented Sep 19, 2023

Uh oh!

jacobhinkle commented Sep 19, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jacobhinkle commented Sep 21, 2023

Uh oh!

jacobhinkle Sep 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jacobhinkle commented Sep 21, 2023

Uh oh!

jacobhinkle commented Sep 22, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jacobhinkle Sep 22, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jjsjann123 commented Sep 22, 2023

Uh oh!

jacobhinkle commented Sep 22, 2023

Uh oh!

jacobhinkle commented Sep 27, 2023

Uh oh!

jacobhinkle commented Oct 4, 2023

Uh oh!

Reviewers

Assignees

jacobhinkle commented Sep 13, 2023 •

edited

Loading

jacobhinkle Sep 19, 2023 •

edited

Loading

jacobhinkle Sep 19, 2023 •

edited

Loading

jacobhinkle Sep 21, 2023 •

edited

Loading

jacobhinkle Sep 22, 2023 •

edited

Loading