Apply the magic-zero protection to each indexed domain individually for predicate indexing #1846

naoyam · 2022-07-19T06:38:47Z

In non-predicate indexing, the protection is done per index expression
as the expressions for indexed domains are summed together. Only the
loop index corresponding to the inner-most loop is protected by the
magic zero.

Unlike non-predicate indexing, each indexed domain has its own
expression in predicate indexing, so each indexed domain must be
protected individually.

This means that we can't just modify the initial loop index map as is
done for non-predicate indexing as the map is shared by all indexed
domains, where only some of them may protect certain loop indices.

Example: FusionInsertMagicZero1

Before:

    #pragma unroll
    for(nvfuser_index_t i29 = 0; i29 < 32; ++i29) {
      #pragma unroll
      for(nvfuser_index_t i30 = 0; i30 < 2; ++i30) {
        int64_t i75;
        i75 = ((i28 % (ceilDiv(T0.size[1], 2))) * 2) + (i30 + nvfuser_zero);
        if ((((((i28 / (ceilDiv(T0.size[1], 2))) * 32) + i29) < T0.size[0]) && (i75 < T0.size[1]))) {
          T2[((((i28 / (ceilDiv(T0.size[1], 2))) * 32) + i29) * T0.size[1]) + i75]
             = T1[(i29 * 2) + i30];
        }

After:

    #pragma unroll
    for(nvfuser_index_t i29 = 0; i29 < 32; ++i29) {
      #pragma unroll
      for(nvfuser_index_t i30 = 0; i30 < 2; ++i30) {
        int64_t i75;
        i75 = ((i28 % (ceilDiv(T0.size[1], 2))) * 2) + (i30 + nvfuser_zero);
        if ((((((i28 / (ceilDiv(T0.size[1], 2))) * 32) + (i29 + nvfuser_zero)) < T0.size[0]) && (i75 < T0.size[1]))) {
          T2[((((i28 / (ceilDiv(T0.size[1], 2))) * 32) + i29) * T0.size[1]) + i75]
             = T1[(i29 * 2) + i30];
        }

Notice that the two conditional expressions of the predicate of T2 are both protected by magic zero.

predicate indexing. In non-predicate indexing, the protection is done per index expression as the expressions for indexed domains are summed together. Only the loop index corresponding to the inner-most loop is protected by the magic zero. Unlike non-predicate indexing, each indexed domain has its own expression in predicate indexing, so each indexed domain must be protected individually. This means that we can't just modify the initial loop index map as is done for non-predicate indexing as the map is shared by all indexed domains, where only some of them may protect certain loop indices.

shmsong

Overall looks good.

Thanks for extending the insertion logic and cleanup some of the existing lines too.

Just wanted to discuss on some minor points before stamping.

shmsong · 2022-07-19T23:29:33Z

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

+  }
+
+  // Clone expression after recurisvely replacing inputs
+  void handle(UnaryOp* uop) override {


Just curious if ir_utils::replaceValInExpr could work here. I guess the main difference is if we could just substitute the uses of replaced vals instead of replicating the whole dependent graph.

Slightly nervous about having multiple dispatch pass for value replacement but would be happy to help maintaining if it is necessary.

The new one effectively clones a given val while applying replacements, whereas replaceValInExpr just applies replacements. Maybe the new one should be renamed to indicate a given val is cloned. Cloning is required as we don't want to affect other predicate expressions.

Maybe leave a comment or issue as it would be nice to try and unify them (have ir_utils::replaceValInExpr optionally clone the DAG you're replacing).

shmsong · 2022-07-19T23:30:45Z

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

+  return loop->isUnrolled() && (!ref_dom_simple || !ind_simple);
+}
+
+void protextNonPredicateIndexWithMagicZero(


Just double checking if I'm understanding correctly, is this code motion from lower_index.h?

I agree this is a good place for this function if that's the case.

That's right.

Should the name be protect instead of protext?

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

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

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

csarofeen

Minor comments.

csarofeen · 2022-07-20T01:31:37Z

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

+
+  bool ref_dom_simple =
+      (reference_domain == nullptr ? true
+                                   : reference_domain->definition() != nullptr);


Am I reading this wrong?
Reference domain is simple if the domain is a null pointer, else if the reference domain has a definition?
Instead of a ternary operation here couldn't it just be
reference_domain == nullptr || reference_domain->definition() != nullptr

Haha, I just copied your original code. I'll apply the cleanup.

csarofeen · 2022-07-20T01:32:21Z

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

+                                   : reference_domain->definition() != nullptr);
+  bool ind_simple =
+      (ind == nullptr ? true
+                      : ind->definition() != nullptr && !ind->isZeroInt());


Same here, uncertain why we need a ternary.

csarofeen · 2022-07-20T01:33:14Z

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

+  return loop->isUnrolled() && (!ref_dom_simple || !ind_simple);
+}
+
+void protextNonPredicateIndexWithMagicZero(


Should the name be protect instead of protext?

csarofeen · 2022-07-20T01:36:40Z

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

+//!
+//! This should be only used for non-predicate indexing.
+//!
+//! No protection is done if none of loop is determined to require


Nit: none of the loops

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

csarofeen · 2022-07-20T01:56:55Z

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

+ public:
+  //! Apply replacements to index as specified in
+  //! replacement_map. index is assumed to consist only from Int and
+  //! NamedScalar


Since it's pretty targeted should we make this OptIn instead of OptOut?

csarofeen · 2022-07-20T02:01:37Z

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

+  }
+
+  // Clone expression after recurisvely replacing inputs
+  void handle(UnaryOp* uop) override {


Maybe leave a comment or issue as it would be nice to try and unify them (have ir_utils::replaceValInExpr optionally clone the DAG you're replacing).

shmsong

LGTM as well after the remaining minor fixes are addressed.

Syncing nvfuser devel branch to upstream master. https://github.com/csarofeen/pytorch/ Code changes includes: - codegen improvements: 1. removes un-necessary sync from redundant thread compute analysis 2. symmetric API for BestEffortReplay 3. support merge on trivial reductions 4. Ampere async copy improvements - bug fixes: 1. vectorization bug fixes 2. type inference patch : fixes upstream pytorch#81725 3. segmenter bug fix with deterministic iteration ordering - parser update 1. added leaky_relu - scheduler 1. normalization scheduler clean up. 2. simplifies matmul scheduling with new transform propagator 3. merge all dimensions in PW scheduler 4. various gemm related improvements - debuggability 1. nsight compute support 2. debug dump for InlinePropagator 3. Add `UnaryOpType::Print` Squashed commits to WAR github API Commits that's actually in this PR from the devel branch: ``` dfe02f3 Merge remote-tracking branch 'csarofeen/devel' into HEAD 1617373 Add `TensorViewBuilder::shape(std::vector<Val*> shape)` (#1884) 7cfb779 Merge pull request #1887 from csarofeen/upstream_merge_0803 3399f6d Merge remote-tracking branch 'origin/viable/strict' into HEAD 01208f5 Add `UnaryOpType::Print` which can be helpful for debugging (#1878) 0646522 Remove redundant TORCH_INTERNAL_ASSERT in lower_magic_zero.cpp (#1881) 7bc76aa Fix most inlined propagator for mismatched dims (#1875) 501f4aa Nonaffine swizzle formulation ep.2: Loop swizzle variant. (#1826) d863d69 Ampere async copy ep.2: circular buffering extension to support pipelined matmul operand load (#1827) e0ae11a Larger sized mma instructions to support full vectorization (#1824) 9bb4cf7 fragment iteration to support fully unrolled mma ops (#1823) a48270a Merge all dims in pointwise scheduler (#1872) 172fb36 Make MostInlined and BestEffort inline propagation no longer assert replayed (#1868) a64462a Allow trivial reduction to be merged (#1871) 440102b Symmetric API for BestEffortReplay (#1870) d1caf33 Some misc cleanups/refactor split out from #1854 (#1867) 1013eda Remove some welford specific logic. (#1864) 51589d3 Some cleanups on tests and heuristics params (#1866) a6b3e70 Segmenter bug fix, and deterministic iteration ordering. (#1865) 1b665b9 Add nullptr checks to IrBuilder (#1861) 1cd9451 Simplify matmul scheduling with the new transform propagator. (#1817) bbc1fb9 Add leaky_relu operation (#1852) e842a9b Minor cleanup in pointwise scheduler (#1858) 9ee850c Fix stringstream usage (#1857) 20a36c1 Improve nsight compute support (#1855) 4059103 Remove debugging `true ||` from getPointwiseHeuristics (#1822) 01117bf Misc cleanup (#1853) 5cc6494 Apply the magic-zero protection to each indexed domain individually for predicate indexing (#1846) 92e6f02 Cleanup normalization scheduler (#1845) db89c65 Type inference patch (#1848) 102fe93 Add debug dump for InlinePropagator (#1847) b7a4d93 Redundant thread compute analysis to avoid un-necessary sync insertion (#1687) 942be5b Upstream ci build fixes (#1842) 0b83645 Fix vectorization bug introduced in #1831 (#1840) 63630f1 Move MaxProducerPosUpdater into InlinePropagator::tearDown (#1825) 9135a96 Fix transpose benchmark dtype (#1839) 2c9a6c0 Add extra configurability to `parallelizeAllLike` (#1831) ``` RUN_TORCHBENCH: nvfuser Differential Revision: [D38543000](https://our.internmc.facebook.com/intern/diff/D38543000) Pull Request resolved: pytorch#83067 Approved by: https://github.com/davidberard98

naoyam force-pushed the magic_zero_predicates branch from cc1fe6b to 8036578 Compare July 19, 2022 06:47

naoyam added 7 commits July 19, 2022 00:17

cleanup

75f751f

cleanup

28e40cf

cleanup

62bac70

cleanup

c94cbe0

cleanup

0180e9e

fix

1137883

Test cleanup and validation

e8d5208

naoyam changed the title ~~[WIP] Apply the magic-zero protection to each indexed domain individually for predicate indexing~~ Apply the magic-zero protection to each indexed domain individually for predicate indexing Jul 19, 2022

naoyam requested a review from csarofeen July 19, 2022 15:29

update comment

62c7252

csarofeen requested a review from shmsong July 19, 2022 20:28

shmsong reviewed Jul 20, 2022

View reviewed changes

pr feedback

97281d4

csarofeen approved these changes Jul 20, 2022

View reviewed changes

shmsong approved these changes Jul 20, 2022

View reviewed changes

naoyam added 5 commits July 20, 2022 08:39

PR feedback

d773828

Add a comment to replaceValInIndexVal

52fc1e8

cleanup

2485ac5

Merge branch 'devel' into magic_zero_predicates

dd870ee

just testing

bf5b41b

shmsong mentioned this pull request Jul 20, 2022

Track dangling node change in indexing math generation. #1850

Open

final cleanup

b8fe441

naoyam merged commit 5cc6494 into devel Jul 21, 2022

naoyam deleted the magic_zero_predicates branch July 21, 2022 01:55

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Apply the magic-zero protection to each indexed domain individually for predicate indexing #1846

Apply the magic-zero protection to each indexed domain individually for predicate indexing #1846

naoyam commented Jul 19, 2022 •

edited

Loading

shmsong left a comment

shmsong Jul 19, 2022

naoyam Jul 20, 2022

csarofeen Jul 20, 2022

shmsong Jul 19, 2022

naoyam Jul 20, 2022

csarofeen Jul 20, 2022

csarofeen left a comment

csarofeen Jul 20, 2022

naoyam Jul 20, 2022

csarofeen Jul 20, 2022

csarofeen Jul 20, 2022

csarofeen Jul 20, 2022

csarofeen Jul 20, 2022

csarofeen Jul 20, 2022

shmsong left a comment

Apply the magic-zero protection to each indexed domain individually for predicate indexing #1846

Apply the magic-zero protection to each indexed domain individually for predicate indexing #1846

Conversation

naoyam commented Jul 19, 2022 • edited Loading

shmsong left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

csarofeen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shmsong left a comment

Choose a reason for hiding this comment

naoyam commented Jul 19, 2022 •

edited

Loading