From 33115db78d76295e7170b4fa9fb55f93683e0631 Mon Sep 17 00:00:00 2001 From: Ashutosh0x Date: Wed, 3 Jun 2026 14:36:16 +0530 Subject: [PATCH 01/48] fix: add upper bound validation for num_segments in UnsortedSegmentReductionOp Add upper bound validation for num_segments in UnsortedSegmentReductionOp::Compute to prevent integer overflow in output tensor size calculations. Changes: - Reject num_segments > 2^31 with InvalidArgument error before allocation - Document the 2G bound rationale (safe across all architectures and dtypes) - Add Python regression test verifying num_segments > 2G raises InvalidArgumentError Fixes #117549 --- .../core/kernels/segment_reduction_ops_impl.h | 18 +++++++++--- .../math_ops/segment_reduction_ops_test.py | 29 ++++++++++++++++++- 2 files changed, 42 insertions(+), 5 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index ad7d9f8e4916b4..82d610ba683136 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -518,12 +518,22 @@ class UnsortedSegmentReductionOp : public OpKernel { internal::ValidateUnsortedSegmentReduction( this, context, data, segment_ids, num_segments)); const auto segment_flat = segment_ids.flat(); - const Index output_rows = internal::SubtleMustCopy(static_cast( + const int64_t num_segments_val = num_segments.dtype() == DT_INT32 ? num_segments.scalar()() - : num_segments.scalar()())); - OP_REQUIRES(context, output_rows >= 0, - errors::InvalidArgument("Input num_segments == ", output_rows, + : num_segments.scalar()(); + OP_REQUIRES(context, num_segments_val >= 0, + errors::InvalidArgument("Input num_segments == ", + num_segments_val, " must not be negative.")); + // Guard against excessively large num_segments that would cause integer + // overflow in output tensor size calculations, leading to illegal memory + // access or process abort (see #117549). + const int64_t kMaxSegments = static_cast(1) << 31; // 2G + OP_REQUIRES(context, num_segments_val <= kMaxSegments, + errors::InvalidArgument( + "Input num_segments == ", num_segments_val, + " is too large. Must be at most ", kMaxSegments)); + const Index output_rows = internal::SubtleMustCopy(static_cast(num_segments_val)); TensorShape output_shape; OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(output_rows)); for (int i = segment_ids.dims(); i < data.dims(); i++) { diff --git a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py index 24cfb3eec109ec..91b0fdcd909302 100644 --- a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py @@ -660,7 +660,7 @@ def testBadNumSegments(self): unsorted = math_ops.unsorted_segment_sum( np.ones((3)), segment_ids=898042203, num_segments=num_segments) with self.assertRaisesOpError( - "Encountered overflow when multiplying | must not be negative" + "Encountered overflow when multiplying | must not be negative | is too large" ): self.evaluate(unsorted) @@ -724,6 +724,33 @@ def testNaNPropagation(self): self.assertTrue(np.isnan(evaluated_max[0][0])) self.assertAllClose(evaluated_max[0][1], 3.0) + def testNumSegmentsOverflow(self): + """Regression test: num_segments > 2G must raise InvalidArgumentError. + + Excessively large num_segments values (near INT64_MAX) cause integer + overflow in internal size calculations, leading to illegal memory access + or process abort. The upper bound of 2^31 prevents this. See #117549. + """ + data = constant_op.constant([1, 2, 3, 4], dtype=dtypes_lib.int32) + segment_ids = constant_op.constant([0, 1, 0, 1], dtype=dtypes_lib.int64) + # Just above the 2G limit (2^31 + 1). + num_segments = constant_op.constant((1 << 31) + 1, dtype=dtypes_lib.int64) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + "is too large"): + self.evaluate(math_ops.unsorted_segment_sum( + data=data, segment_ids=segment_ids, num_segments=num_segments)) + + def testNumSegmentsOverflow_Int32Index(self): + """Regression test: num_segments > 2G must raise InvalidArgumentError even with int32 index. + """ + data = constant_op.constant([1, 2, 3, 4], dtype=dtypes_lib.int32) + segment_ids = constant_op.constant([0, 1, 0, 1], dtype=dtypes_lib.int32) + # Just above the 2G limit (2^31 + 1). + num_segments = constant_op.constant((1 << 31) + 1, dtype=dtypes_lib.int64) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + "is too large"): + self.evaluate(math_ops.unsorted_segment_sum( + data=data, segment_ids=segment_ids, num_segments=num_segments)) class SparseSegmentReductionHelper(SegmentReductionHelper): From 82cb5155dcb6e2a965e6875959f7056dcd6f3944 Mon Sep 17 00:00:00 2001 From: Ashutosh Kumar Singh <161562995+Ashutosh0x@users.noreply.github.com> Date: Sat, 25 Jul 2026 20:41:26 +0530 Subject: [PATCH 02/48] Cap num_segments bound by Index max to prevent int32 overflow Addresses review feedback on #120176: the 2G guard now caps kMaxSegments at std::numeric_limits::max(), so num_segments == 2^31 is rejected for signed int32 indices instead of overflowing the static_cast. Adds an int32-index boundary regression test at exactly 2^31 and wraps an over-long assert regex to satisfy PyLint. --- .../core/kernels/segment_reduction_ops_impl.h | 9 +++++++-- .../math_ops/segment_reduction_ops_test.py | 14 ++++++++++---- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index 82d610ba683136..a6b098e283c1c4 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -527,8 +527,13 @@ class UnsortedSegmentReductionOp : public OpKernel { " must not be negative.")); // Guard against excessively large num_segments that would cause integer // overflow in output tensor size calculations, leading to illegal memory - // access or process abort (see #117549). - const int64_t kMaxSegments = static_cast(1) << 31; // 2G + // access or process abort (see #117549). For 32-bit Index types the + // effective limit is the largest value representable by Index (kint32max), + // since a larger value would overflow the static_cast below; cap + // kMaxSegments accordingly. + const int64_t kMaxSegments = std::min( + static_cast(1) << 31, // 2G + static_cast(std::numeric_limits::max())); OP_REQUIRES(context, num_segments_val <= kMaxSegments, errors::InvalidArgument( "Input num_segments == ", num_segments_val, diff --git a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py index 91b0fdcd909302..7c337cd7662a6b 100644 --- a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py @@ -660,7 +660,8 @@ def testBadNumSegments(self): unsorted = math_ops.unsorted_segment_sum( np.ones((3)), segment_ids=898042203, num_segments=num_segments) with self.assertRaisesOpError( - "Encountered overflow when multiplying | must not be negative | is too large" + "Encountered overflow when multiplying | must not be negative | " + "is too large" ): self.evaluate(unsorted) @@ -741,12 +742,17 @@ def testNumSegmentsOverflow(self): data=data, segment_ids=segment_ids, num_segments=num_segments)) def testNumSegmentsOverflow_Int32Index(self): - """Regression test: num_segments > 2G must raise InvalidArgumentError even with int32 index. + """num_segments == 2^31 must be rejected when the index type is int32. + + With a signed int32 index the largest representable segment count is + kint32max (2^31 - 1), so exactly 2^31 must be flagged as too large + rather than silently overflowing the cast to Index. See #117549. """ data = constant_op.constant([1, 2, 3, 4], dtype=dtypes_lib.int32) segment_ids = constant_op.constant([0, 1, 0, 1], dtype=dtypes_lib.int32) - # Just above the 2G limit (2^31 + 1). - num_segments = constant_op.constant((1 << 31) + 1, dtype=dtypes_lib.int64) + # Exactly the 2G boundary (2^31): valid for an int64 index but out of + # range for a signed int32 index. + num_segments = constant_op.constant(1 << 31, dtype=dtypes_lib.int64) with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "is too large"): self.evaluate(math_ops.unsorted_segment_sum( From 12b2fed532063d81ed6d73e071967f3900f2b74e Mon Sep 17 00:00:00 2001 From: Ashutosh0x Date: Sun, 26 Jul 2026 23:35:31 +0530 Subject: [PATCH 03/48] Copy num_segments before validation and apply clang-format Addresses review feedback on #120176: - Wrap the initial read of num_segments in internal::SubtleMustCopy, as the original code did, so the value is copied out of the (possibly shared) input buffer before it is validated. Previously the copy happened after the range checks, leaving a TOCTOU window. - Reformat the changed region with clang-format (Google style) so no line exceeds 80 columns. - Restore the two blank lines before SparseSegmentReductionHelper in segment_reduction_ops_test.py. --- .../core/kernels/segment_reduction_ops_impl.h | 21 ++++++++++++------- .../math_ops/segment_reduction_ops_test.py | 1 + 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index a6b098e283c1c4..7f4c8b4be04914 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -518,13 +518,17 @@ class UnsortedSegmentReductionOp : public OpKernel { internal::ValidateUnsortedSegmentReduction( this, context, data, segment_ids, num_segments)); const auto segment_flat = segment_ids.flat(); - const int64_t num_segments_val = - num_segments.dtype() == DT_INT32 ? num_segments.scalar()() - : num_segments.scalar()(); - OP_REQUIRES(context, num_segments_val >= 0, - errors::InvalidArgument("Input num_segments == ", - num_segments_val, - " must not be negative.")); + // Copy the value out of the (possibly shared) input buffer before any + // validation, so that the value that is checked is the same one that is + // subsequently used (see SubtleMustCopy). + const int64_t num_segments_val = internal::SubtleMustCopy( + num_segments.dtype() == DT_INT32 + ? static_cast(num_segments.scalar()()) + : num_segments.scalar()()); + OP_REQUIRES( + context, num_segments_val >= 0, + errors::InvalidArgument("Input num_segments == ", num_segments_val, + " must not be negative.")); // Guard against excessively large num_segments that would cause integer // overflow in output tensor size calculations, leading to illegal memory // access or process abort (see #117549). For 32-bit Index types the @@ -538,7 +542,8 @@ class UnsortedSegmentReductionOp : public OpKernel { errors::InvalidArgument( "Input num_segments == ", num_segments_val, " is too large. Must be at most ", kMaxSegments)); - const Index output_rows = internal::SubtleMustCopy(static_cast(num_segments_val)); + const Index output_rows = + internal::SubtleMustCopy(static_cast(num_segments_val)); TensorShape output_shape; OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(output_rows)); for (int i = segment_ids.dims(); i < data.dims(); i++) { diff --git a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py index 7c337cd7662a6b..5e832ca30106ea 100644 --- a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py @@ -758,6 +758,7 @@ def testNumSegmentsOverflow_Int32Index(self): self.evaluate(math_ops.unsorted_segment_sum( data=data, segment_ids=segment_ids, num_segments=num_segments)) + class SparseSegmentReductionHelper(SegmentReductionHelper): def _sparse_input(self, input_shape, num_indices, dtype=dtypes_lib.int32): From 9916bf6e5559ac533c9e3f8ef43343fdee9ab21d Mon Sep 17 00:00:00 2001 From: Ashutosh0x Date: Wed, 29 Jul 2026 23:01:52 +0530 Subject: [PATCH 04/48] Use absl::InvalidArgumentError and absl::StrCat for the new errors Matches the surrounding code, which already uses absl::InvalidArgumentError throughout this file, and the style dmiltr3 had to apply by hand to get the change through the internal presubmits. The message text is unchanged, so the two regression tests still assert the same strings. --- tensorflow/core/kernels/segment_reduction_ops_impl.h | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index 7f4c8b4be04914..24c6d729126738 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -525,10 +525,10 @@ class UnsortedSegmentReductionOp : public OpKernel { num_segments.dtype() == DT_INT32 ? static_cast(num_segments.scalar()()) : num_segments.scalar()()); - OP_REQUIRES( - context, num_segments_val >= 0, - errors::InvalidArgument("Input num_segments == ", num_segments_val, - " must not be negative.")); + OP_REQUIRES(context, num_segments_val >= 0, + absl::InvalidArgumentError( + absl::StrCat("Input num_segments == ", num_segments_val, + " must not be negative."))); // Guard against excessively large num_segments that would cause integer // overflow in output tensor size calculations, leading to illegal memory // access or process abort (see #117549). For 32-bit Index types the @@ -539,9 +539,9 @@ class UnsortedSegmentReductionOp : public OpKernel { static_cast(1) << 31, // 2G static_cast(std::numeric_limits::max())); OP_REQUIRES(context, num_segments_val <= kMaxSegments, - errors::InvalidArgument( + absl::InvalidArgumentError(absl::StrCat( "Input num_segments == ", num_segments_val, - " is too large. Must be at most ", kMaxSegments)); + " is too large. Must be at most ", kMaxSegments))); const Index output_rows = internal::SubtleMustCopy(static_cast(num_segments_val)); TensorShape output_shape; From 218b09a815d4d13bc1ce8d110f1a3820860120e9 Mon Sep 17 00:00:00 2001 From: Ashutosh0x Date: Thu, 30 Jul 2026 07:12:33 +0530 Subject: [PATCH 05/48] Apply SubtleMustCopy to the tensor reads, not the derived values Read each num_segments scalar through internal::SubtleMustCopy inside the branch, before the widening cast, so the copy actually guards the read from the shared input buffer. Drop the redundant SubtleMustCopy on output_rows, which was applied to a local const int64_t that cannot be mutated asynchronously. --- tensorflow/core/kernels/segment_reduction_ops_impl.h | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index 24c6d729126738..e517bbbf4c6608 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -521,10 +521,10 @@ class UnsortedSegmentReductionOp : public OpKernel { // Copy the value out of the (possibly shared) input buffer before any // validation, so that the value that is checked is the same one that is // subsequently used (see SubtleMustCopy). - const int64_t num_segments_val = internal::SubtleMustCopy( + const int64_t num_segments_val = num_segments.dtype() == DT_INT32 - ? static_cast(num_segments.scalar()()) - : num_segments.scalar()()); + ? internal::SubtleMustCopy(num_segments.scalar()()) + : internal::SubtleMustCopy(num_segments.scalar()()); OP_REQUIRES(context, num_segments_val >= 0, absl::InvalidArgumentError( absl::StrCat("Input num_segments == ", num_segments_val, @@ -542,8 +542,7 @@ class UnsortedSegmentReductionOp : public OpKernel { absl::InvalidArgumentError(absl::StrCat( "Input num_segments == ", num_segments_val, " is too large. Must be at most ", kMaxSegments))); - const Index output_rows = - internal::SubtleMustCopy(static_cast(num_segments_val)); + const Index output_rows = static_cast(num_segments_val); TensorShape output_shape; OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(output_rows)); for (int i = segment_ids.dims(); i < data.dims(); i++) { From 12ad00140e43207310c332f6512ac05e349f3ff4 Mon Sep 17 00:00:00 2001 From: Ashutosh0x Date: Thu, 30 Jul 2026 20:25:19 +0530 Subject: [PATCH 06/48] Bound num_segments by kint32max for every index type The previous bound was min(2^31, Index max), which resolved to kint32max for a 32-bit index but to 2^31 for a 64-bit one. Those differ by one, and the extra value is the unsafe one: 2^31 is not representable in a signed 32-bit integer, so it casts to -2147483648 and would wrap in any downstream code that indexes with int32. Using kint32max for all index types removes both the wrap and the asymmetry where a 64-bit index admitted one more segment than a 32-bit one, for no reason a caller could observe. Nothing practical is given up: the CPU functor allocates a std::vector row_counter(num_segments), so kint32max segments is already tens of gigabytes of bookkeeping before any output element is written. Adds testNumSegmentsOverflow_Int64Index, which pins 2^31 as rejected on the 64-bit path -- the case the old bound let through and which no test covered. Existing expectations are unaffected: testNumSegmentsOverflow (2^31 + 1), testNumSegmentsOverflow_Int32Index (2^31) and testBadNumSegments all still reject, and the error strings are unchanged. --- .../core/kernels/segment_reduction_ops_impl.h | 22 +++++++++++++------ .../math_ops/segment_reduction_ops_test.py | 20 +++++++++++++++-- 2 files changed, 33 insertions(+), 9 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index e517bbbf4c6608..3debc1f655a8c8 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -531,13 +531,21 @@ class UnsortedSegmentReductionOp : public OpKernel { " must not be negative."))); // Guard against excessively large num_segments that would cause integer // overflow in output tensor size calculations, leading to illegal memory - // access or process abort (see #117549). For 32-bit Index types the - // effective limit is the largest value representable by Index (kint32max), - // since a larger value would overflow the static_cast below; cap - // kMaxSegments accordingly. - const int64_t kMaxSegments = std::min( - static_cast(1) << 31, // 2G - static_cast(std::numeric_limits::max())); + // access or process abort (see #117549). + // + // The bound is kint32max for every Index type, not 2^31. Those differ by + // one, and the smaller value is the safe one: 2^31 is not representable in + // a signed 32-bit integer, so any downstream code that indexes with int32 + // would wrap it to a negative value. Using kint32max also removes the + // asymmetry where a 64-bit index admitted one more segment than a 32-bit + // one for no reason a caller could observe or rely on. + // + // Nothing practical is given up. A single segment costs at least an entry + // in the CPU functor's std::vector row_counter, so kint32max + // segments is already tens of gigabytes of bookkeeping before any output + // element is written. + const int64_t kMaxSegments = + static_cast(std::numeric_limits::max()); OP_REQUIRES(context, num_segments_val <= kMaxSegments, absl::InvalidArgumentError(absl::StrCat( "Input num_segments == ", num_segments_val, diff --git a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py index 5e832ca30106ea..19fc209f4a7e23 100644 --- a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py @@ -750,8 +750,24 @@ def testNumSegmentsOverflow_Int32Index(self): """ data = constant_op.constant([1, 2, 3, 4], dtype=dtypes_lib.int32) segment_ids = constant_op.constant([0, 1, 0, 1], dtype=dtypes_lib.int32) - # Exactly the 2G boundary (2^31): valid for an int64 index but out of - # range for a signed int32 index. + # Exactly the 2^31 boundary: one past what a signed int32 can represent. + num_segments = constant_op.constant(1 << 31, dtype=dtypes_lib.int64) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + "is too large"): + self.evaluate(math_ops.unsorted_segment_sum( + data=data, segment_ids=segment_ids, num_segments=num_segments)) + + def testNumSegmentsOverflow_Int64Index(self): + """num_segments == 2^31 must be rejected with an int64 index too. + + The limit is kint32max for every index type, not 2^31. The two differ by + one, and an earlier revision of this check allowed 2^31 whenever Index was + wide enough to hold it -- which left a value that is not representable in a + signed 32-bit integer reaching any downstream code that indexes with int32. + This pins the 64-bit path to the same bound. See #117549. + """ + data = constant_op.constant([1, 2, 3, 4], dtype=dtypes_lib.int32) + segment_ids = constant_op.constant([0, 1, 0, 1], dtype=dtypes_lib.int64) num_segments = constant_op.constant(1 << 31, dtype=dtypes_lib.int64) with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "is too large"): From 81285cf9835d43b908ffc68a17f53514db7f640d Mon Sep 17 00:00:00 2001 From: Ashutosh0x Date: Thu, 30 Jul 2026 20:28:27 +0530 Subject: [PATCH 07/48] Tighten the comment on the num_segments bound Comment only; the bound, the error strings and the tests are unchanged. Says why kint32max is used for every index type without restating it three times, and keeps the note that the CPU functor's row_counter allocation makes the limit academic in practice. --- .../core/kernels/segment_reduction_ops_impl.h | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index 3debc1f655a8c8..f47f4fa96969b8 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -533,17 +533,11 @@ class UnsortedSegmentReductionOp : public OpKernel { // overflow in output tensor size calculations, leading to illegal memory // access or process abort (see #117549). // - // The bound is kint32max for every Index type, not 2^31. Those differ by - // one, and the smaller value is the safe one: 2^31 is not representable in - // a signed 32-bit integer, so any downstream code that indexes with int32 - // would wrap it to a negative value. Using kint32max also removes the - // asymmetry where a 64-bit index admitted one more segment than a 32-bit - // one for no reason a caller could observe or rely on. - // - // Nothing practical is given up. A single segment costs at least an entry - // in the CPU functor's std::vector row_counter, so kint32max - // segments is already tens of gigabytes of bookkeeping before any output - // element is written. + // Use kint32max for every index type. Although an int64 index could + // represent 2^31, that value is not representable as a signed int32 and + // would wrap if any downstream code performs 32-bit indexing; a uniform + // bound avoids that edge case. Nothing practical is lost, since the CPU + // functor already allocates std::vector row_counter(num_segments). const int64_t kMaxSegments = static_cast(std::numeric_limits::max()); OP_REQUIRES(context, num_segments_val <= kMaxSegments, From 4f880468ad11ee81bcf355ea7562298d25485839 Mon Sep 17 00:00:00 2001 From: Ashutosh0x Date: Wed, 2 Sep 2026 08:33:11 +0530 Subject: [PATCH 08/48] Make both ternary branches explicitly int64_t The int32 branch now casts to int64_t rather than relying on the implicit conversion the ternary's common type performs, so both branches have the same type at the point of selection. The cast sits outside SubtleMustCopy rather than inside it. That ordering is what keeps the earlier review fix intact: SubtleMustCopy takes a const T& and reads through a volatile pointer to the address it binds, so it has to bind to the tensor buffer to be worth anything. Casting first would hand it a temporary and quietly restore the problem. Casting its result is free -- the volatile read has already happened on the buffer, and widening a copied int32 cannot lose information. No behavioural change: the ternary already yielded int64_t by promotion, so this only makes the conversion explicit and its position deliberate. Error strings are untouched and the three regression tests still assert the same messages. Also corrects a docstring that still described the bound as 2^31; it has been kint32max since 12ad001. --- tensorflow/core/kernels/segment_reduction_ops_impl.h | 3 ++- .../kernel_tests/math_ops/segment_reduction_ops_test.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index f47f4fa96969b8..c92156f26151fd 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -523,7 +523,8 @@ class UnsortedSegmentReductionOp : public OpKernel { // subsequently used (see SubtleMustCopy). const int64_t num_segments_val = num_segments.dtype() == DT_INT32 - ? internal::SubtleMustCopy(num_segments.scalar()()) + ? static_cast( + internal::SubtleMustCopy(num_segments.scalar()())) : internal::SubtleMustCopy(num_segments.scalar()()); OP_REQUIRES(context, num_segments_val >= 0, absl::InvalidArgumentError( diff --git a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py index 19fc209f4a7e23..ddd9e5b39da545 100644 --- a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py @@ -726,11 +726,11 @@ def testNaNPropagation(self): self.assertAllClose(evaluated_max[0][1], 3.0) def testNumSegmentsOverflow(self): - """Regression test: num_segments > 2G must raise InvalidArgumentError. + """Regression test: an oversized num_segments raises InvalidArgumentError. Excessively large num_segments values (near INT64_MAX) cause integer overflow in internal size calculations, leading to illegal memory access - or process abort. The upper bound of 2^31 prevents this. See #117549. + or process abort. The kint32max upper bound prevents this. See #117549. """ data = constant_op.constant([1, 2, 3, 4], dtype=dtypes_lib.int32) segment_ids = constant_op.constant([0, 1, 0, 1], dtype=dtypes_lib.int64) From bf50ce5520aa77c122258c141b7dc670aeb014d2 Mon Sep 17 00:00:00 2001 From: abhirupone Date: Tue, 25 Aug 2026 20:36:57 +0530 Subject: [PATCH 09/48] test(process_util): add negative thread test for NumInterOpThreadsFromSessionOptions Adds NumThreadsNegative in process_util_test.cc to verify that NumInterOpThreadsFromSessionOptions returns a positive thread count when inter_op_parallelism_threads is negative. Follow-up to #124954 and commit 1d31888b3e0. --- tensorflow/core/common_runtime/process_util_test.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorflow/core/common_runtime/process_util_test.cc b/tensorflow/core/common_runtime/process_util_test.cc index d1a53e8f411854..20a190c839967b 100644 --- a/tensorflow/core/common_runtime/process_util_test.cc +++ b/tensorflow/core/common_runtime/process_util_test.cc @@ -28,6 +28,12 @@ TEST(ProcessUtilTest, NumThreads) { EXPECT_EQ(10, NumInterOpThreadsFromSessionOptions(opts)); } +TEST(ProcessUtilTest, NumThreadsNegative) { + SessionOptions opts; + opts.config.set_inter_op_parallelism_threads(-1); + EXPECT_GT(NumInterOpThreadsFromSessionOptions(opts), 0); +} + TEST(ProcessUtilTest, ThreadPool) { SessionOptions opts; opts.config.set_inter_op_parallelism_threads(10); From 081a82a2304f3cc7077f5adebac72f57e56e387e Mon Sep 17 00:00:00 2001 From: Rishikesh R <145715786+dynamo-pentester@users.noreply.github.com> Date: Mon, 7 Sep 2026 00:49:02 +0530 Subject: [PATCH 10/48] Validate non-negative values in fractional max pool Add checks to ensure height and width sequence tensor values are non-negative. --- tensorflow/core/kernels/fractional_max_pool_op.cc | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tensorflow/core/kernels/fractional_max_pool_op.cc b/tensorflow/core/kernels/fractional_max_pool_op.cc index 5ae827dc25032a..8735fc0e0691fe 100644 --- a/tensorflow/core/kernels/fractional_max_pool_op.cc +++ b/tensorflow/core/kernels/fractional_max_pool_op.cc @@ -344,6 +344,10 @@ class FractionalMaxPoolGradOp : public OpKernel { for (int64_t hs = 0; hs < height_seq_tensor.dim_size(0) - 1; ++hs) { // height start and end. const int64_t height_start = height_seq_tensor_flat(hs); + OP_REQUIRES(context, height_start >= 0, + absl::InvalidArgumentError(absl::StrCat( + "Row sequence tensor values must not be negative, got ", + height_seq_tensor.DebugString()))); int64_t height_end = overlapping_ ? height_seq_tensor_flat(hs + 1) : height_seq_tensor_flat(hs + 1) - 1; height_end = std::min(height_end, height_max); @@ -354,6 +358,10 @@ class FractionalMaxPoolGradOp : public OpKernel { (b * output_size[1] + hs) * output_size[2] + ws; // width start and end. const int64_t width_start = width_seq_tensor_flat(ws); + OP_REQUIRES(context, width_start >= 0, + absl::InvalidArgumentError(absl::StrCat( + "Col sequence tensor values must not be negative, got ", + width_seq_tensor.DebugString()))); int64_t width_end = overlapping_ ? width_seq_tensor_flat(ws + 1) : width_seq_tensor_flat(ws + 1) - 1; width_end = std::min(width_end, width_max); From c3bf3b3b933724e768b7c607dc6b1e3e3bb0c80e Mon Sep 17 00:00:00 2001 From: adi-IL Date: Sat, 22 Aug 2026 04:46:23 +0530 Subject: [PATCH 11/48] fix: validate output signatures in LookupTableExportOp to prevent crash on dtype mismatch --- tensorflow/core/kernels/lookup_table_op.cc | 5 +++++ .../data_structures/lookup_ops_test.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/tensorflow/core/kernels/lookup_table_op.cc b/tensorflow/core/kernels/lookup_table_op.cc index f882223f760e59..e4b0ae6dddf91c 100644 --- a/tensorflow/core/kernels/lookup_table_op.cc +++ b/tensorflow/core/kernels/lookup_table_op.cc @@ -1042,6 +1042,11 @@ class LookupTableExportOp : public LookupTableOpKernel { OP_REQUIRES_OK(ctx, GetTable(ctx, &table)); core::ScopedUnref unref_me(table); + DataTypeVector expected_inputs = {expected_input_0_}; + DataTypeVector expected_outputs = {table->key_dtype(), + table->value_dtype()}; + OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, expected_outputs)); + OP_REQUIRES_OK(ctx, table->ExportValues(ctx)); } }; diff --git a/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py b/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py index 9aba606f69b6ff..6c3cceb4f6bbe6 100644 --- a/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py +++ b/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py @@ -116,6 +116,23 @@ def testStaticHashTable(self, is_anonymous): ) self.assertCountEqual([0, 1, 2], self.evaluate(exported_values_tensor)) + def testExportSignatureMismatch(self, is_anonymous): + if is_anonymous and not tf2.enabled(): + self.skipTest(SKIP_ANONYMOUS_IN_TF1_REASON) + table = self.getHashTable()( + lookup_ops.KeyValueTensorInitializer( + constant_op.constant([0], dtype=dtypes.int64), + constant_op.constant([1], dtype=dtypes.int64)), + -1, + experimental_is_anonymous=is_anonymous) + self.initialize_table(table) + with self.assertRaises((errors_impl.InvalidArgumentError, ValueError)): + self.evaluate( + gen_lookup_ops.lookup_table_export_v2( + table.resource_handle, + Tkeys=dtypes.string, + Tvalues=dtypes.string)) + def testStaticHashTableFindHighRank(self, is_anonymous): if is_anonymous and not tf2.enabled(): self.skipTest(SKIP_ANONYMOUS_IN_TF1_REASON) From 6d8039efd262633ea8a039fd74d854857d0d2f97 Mon Sep 17 00:00:00 2001 From: gauravarun21 Date: Wed, 16 Sep 2026 09:59:59 +0200 Subject: [PATCH 12/48] Fix out-of-bounds access in GPU SparseSegment*Grad kernels SparseSegmentGradFunctor (used by SparseSegmentSumGrad, SparseSegmentMeanGrad, and SparseSegmentSqrtNGrad on GPU) never validated its `indices` and `segment_ids` inputs before using them to read/write device memory. The CPU functor bounds-checks every value against the output-row count and the incoming gradient's row count before using it, returning InvalidArgumentError on failure; the GPU path had no equivalent check, so a malformed or out-of-range index caused an out-of-bounds device read/write, typically surfacing as a fatal `cudaErrorIllegalAddress` abort that also corrupts the CUDA context for the rest of the process. Add the missing validation to the GPU functor by copying the (typically small) `indices`/`segment_ids` arrays to the host and replicating the same checks and error messages the CPU functor already uses. Several existing tests for this were already restricted to `use_gpu=False` specifically to avoid the crash; this removes that restriction so they now exercise the fixed GPU path directly. Fixes #113146. --- .../kernels/segment_reduction_ops_gpu.cu.h | 69 +++++++++++++++++++ .../math_ops/segment_reduction_ops_test.py | 10 +-- 2 files changed, 74 insertions(+), 5 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h index 2221347395abaa..9cd518036efdd5 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h +++ b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h @@ -22,6 +22,7 @@ limitations under the License. #include "xla/tsl/platform/statusor.h" #include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/kernels/gpu_prim.h" #include "tensorflow/core/kernels/gpu_prim_helpers.h" @@ -1003,6 +1004,74 @@ struct SparseSegmentGradFunctor { const Index nouter = indices_vec.dimension(0); const Index noutput = output_flat.dimension(0); + // The CPU functor bounds-checks every `indices` value against `noutput` + // and every `segment_ids` value against `nsegments` before using them to + // read/write `input`/`output`, returning an InvalidArgumentError on + // failure. This GPU path performed no such validation, so an + // out-of-range indices or segment_ids value caused an out-of-bounds + // read or write on the device instead of a normal, recoverable error. + // Validate here by copying both (typically small) arrays to the host and + // replicating the same checks. + se::Stream* stream = context->op_device_context()->stream(); + OP_REQUIRES(context, stream != nullptr, + absl::InternalError("No GPU stream available.")); + + AllocatorAttributes host_alloc_attr; + host_alloc_attr.set_on_host(true); + host_alloc_attr.set_gpu_compatible(true); + Tensor indices_host, segment_host; + OP_REQUIRES_OK(context, + context->allocate_temp(DataTypeToEnum::value, + TensorShape({nouter}), &indices_host, + host_alloc_attr)); + OP_REQUIRES_OK(context, + context->allocate_temp(DataTypeToEnum::value, + TensorShape({nouter}), &segment_host, + host_alloc_attr)); + OP_REQUIRES_OK( + context, + stream->Memcpy(indices_host.flat().data(), + stream_executor::DeviceAddressBase( + const_cast(indices_vec.data()), + nouter * sizeof(Index)), + nouter * sizeof(Index))); + OP_REQUIRES_OK( + context, + stream->Memcpy(segment_host.flat().data(), + stream_executor::DeviceAddressBase( + const_cast(segment_vec.data()), + nouter * sizeof(SegmentId)), + nouter * sizeof(SegmentId))); + OP_REQUIRES_OK(context, stream->BlockHostUntilDone()); + { + const auto indices_host_vec = indices_host.vec(); + const auto segment_host_vec = segment_host.vec(); + + // Fast check, matching the CPU functor: `segment_vec` is assumed + // sorted in non-decreasing order, so its last element is the maximum + // segment id. This produces the same "Invalid number of segments" + // message as the CPU functor for the common (sorted) case where the + // incoming gradient simply doesn't have enough rows. + const SegmentId last_segment_id = + internal::SubtleMustCopy(segment_host_vec(nouter - 1)); + OP_REQUIRES(context, last_segment_id + 1 <= nsegments, + absl::InvalidArgumentError("Invalid number of segments")); + + for (Index i = 0; i < nouter; ++i) { + const Index output_idx = + internal::SubtleMustCopy(indices_host_vec(i)); + OP_REQUIRES(context, FastBoundsCheck(output_idx, noutput), + absl::InvalidArgumentError(absl::StrCat( + "Index ", output_idx, " out of range [0, ", noutput, + ")."))); + const SegmentId idx = internal::SubtleMustCopy(segment_host_vec(i)); + OP_REQUIRES( + context, FastBoundsCheck(idx, nsegments), + absl::InvalidArgumentError(absl::StrCat( + "Segment id ", idx, " out of range [0, ", nsegments, ")."))); + } + } + // Allocate and compute segment weights (for Mean/SqrtN operations only). Tensor weights; using Tweights = typename RealTypeIfComplex::type; diff --git a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py index 2f7adecacfd6bf..3f9a8b888c893f 100644 --- a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py @@ -1281,7 +1281,7 @@ def testGradientIndicesInvalid1(self): ] segment_indices = [0, 1, 2, 2] tf_indices = [8, 3, 0, 10] - with self.session(use_gpu=False): + with self.session(): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) with self.assertRaisesOpError(r"Index 10 out of range \[0, 10\)"): @@ -1296,7 +1296,7 @@ def testGradientIndicesInvalid2(self): ] segment_indices = [0, 1, 2, 2] tf_indices = [8, 3, -1, 9] - with self.session(use_gpu=False): + with self.session(): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) with self.assertRaisesOpError(r"Index -1 out of range \[0, 10\)"): @@ -1312,7 +1312,7 @@ def testGradientSegmentsInvalid1(self): ] segment_indices = [0, 1, 1, 4] # 5 segments tf_indices = [8, 3, 0, 9] - with self.session(use_gpu=False): + with self.session(): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) with self.assertRaisesOpError("Invalid number of segments"): @@ -1327,7 +1327,7 @@ def testGradientSegmentsInvalid2(self): ] segment_indices = [0, 1, 2, 0] tf_indices = [8, 3, 0, 9] - with self.session(use_gpu=False): + with self.session(): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) with self.assertRaisesOpError(r"Segment id 1 out of range \[0, 1\)"): @@ -1342,7 +1342,7 @@ def testGradientSegmentsInvalid3(self): ] segment_indices = [-1, 0, 1, 1] tf_indices = [8, 3, 0, 9] - with self.session(use_gpu=False): + with self.session(): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) with self.assertRaisesOpError(r"Segment id -1 out of range \[0, 2\)"): From 4b118850acecd8cf20956f67a0f49ea15797a173 Mon Sep 17 00:00:00 2001 From: gauravarun21 Date: Wed, 16 Sep 2026 15:29:45 +0200 Subject: [PATCH 13/48] Address review: guard nouter==0 and avoid full-array D2H copies - Wrap the GPU validation logic in `if (nouter > 0)` so it no longer reads segment_vec[-1] when indices_vec is empty. - Replace the synchronous full-array copy of `indices`/`segment_ids` to the host with a device-side bounds-check kernel that scans both arrays in parallel and records only the position of a violation (if any) via GpuAtomicMin into a 2-element int32 buffer. Only that tiny buffer is transferred back to the host in the common (valid-input) case; the actual offending value is fetched with one additional single-element copy only when a violation was actually found, to build the error message. Addresses review feedback from @dmiltr3. --- .../kernels/segment_reduction_ops_gpu.cu.h | 179 ++++++++++++------ 1 file changed, 126 insertions(+), 53 deletions(-) diff --git a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h index 9cd518036efdd5..cbd40d8cf2ee7f 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h +++ b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h @@ -988,6 +988,46 @@ absl::Status SparseSegmentReductionFunctor::operator()( /*output=*/output.data()); } +// Finds the position of an out-of-range `indices` value (against `noutput`) +// and/or `segment_ids` value (against `nsegments`), if any. `bad_index_pos` +// and `bad_segment_pos` must already be initialized to +// std::numeric_limits::max(); a value is left unchanged if no +// violation of that kind is found, or set to the position of one such +// violation (not necessarily the first in `i` order) otherwise. +template +__global__ void SparseSegmentGradBoundsCheckKernel( + Index nouter, Index noutput, SegmentId nsegments, + const Index* __restrict__ indices_vec, // [nouter] + const SegmentId* __restrict__ segment_vec, // [nouter] + int32_t* __restrict__ bad_index_pos, // [1] + int32_t* __restrict__ bad_segment_pos) { // [1] + for (Index i : GpuGridRangeX(nouter)) { + if (!FastBoundsCheck(indices_vec[i], noutput)) { + GpuAtomicMin(bad_index_pos, static_cast(i)); + } + if (!FastBoundsCheck(segment_vec[i], nsegments)) { + GpuAtomicMin(bad_segment_pos, static_cast(i)); + } + } +} + +template +absl::Status LaunchSparseSegmentGradBoundsCheckKernel( + const GPUDevice& d, Index nouter, Index noutput, SegmentId nsegments, + const Index* indices_vec, const SegmentId* segment_vec, + int32_t* bad_index_pos, int32_t* bad_segment_pos) { + TF_ASSIGN_OR_RETURN( + GpuLaunchConfig64 config, + GetGpuLaunchConfig64(nouter, d, + &SparseSegmentGradBoundsCheckKernel, + /*dynamic_shared_memory_size=*/0, + /*block_size_limit=*/0)); + return GpuLaunchKernel(SparseSegmentGradBoundsCheckKernel, + config.block_count, config.thread_per_block, 0, + d.stream(), nouter, noutput, nsegments, indices_vec, + segment_vec, bad_index_pos, bad_segment_pos); +} + template struct SparseSegmentGradFunctor { void operator()(OpKernelContext* context, @@ -1010,65 +1050,98 @@ struct SparseSegmentGradFunctor { // failure. This GPU path performed no such validation, so an // out-of-range indices or segment_ids value caused an out-of-bounds // read or write on the device instead of a normal, recoverable error. - // Validate here by copying both (typically small) arrays to the host and - // replicating the same checks. - se::Stream* stream = context->op_device_context()->stream(); - OP_REQUIRES(context, stream != nullptr, - absl::InternalError("No GPU stream available.")); - - AllocatorAttributes host_alloc_attr; - host_alloc_attr.set_on_host(true); - host_alloc_attr.set_gpu_compatible(true); - Tensor indices_host, segment_host; - OP_REQUIRES_OK(context, - context->allocate_temp(DataTypeToEnum::value, - TensorShape({nouter}), &indices_host, - host_alloc_attr)); - OP_REQUIRES_OK(context, - context->allocate_temp(DataTypeToEnum::value, - TensorShape({nouter}), &segment_host, - host_alloc_attr)); - OP_REQUIRES_OK( - context, - stream->Memcpy(indices_host.flat().data(), - stream_executor::DeviceAddressBase( - const_cast(indices_vec.data()), - nouter * sizeof(Index)), - nouter * sizeof(Index))); - OP_REQUIRES_OK( - context, - stream->Memcpy(segment_host.flat().data(), - stream_executor::DeviceAddressBase( - const_cast(segment_vec.data()), - nouter * sizeof(SegmentId)), - nouter * sizeof(SegmentId))); - OP_REQUIRES_OK(context, stream->BlockHostUntilDone()); - { - const auto indices_host_vec = indices_host.vec(); - const auto segment_host_vec = segment_host.vec(); + // Validate here, without copying the full `indices`/`segment_ids` + // arrays to the host: a device-side kernel does the O(nouter) scan, and + // only a couple of small flags come back over the (synchronous) D2H + // transfer in the common, valid-input case. + if (nouter > 0) { + se::Stream* stream = context->op_device_context()->stream(); + OP_REQUIRES(context, stream != nullptr, + absl::InternalError("No GPU stream available.")); // Fast check, matching the CPU functor: `segment_vec` is assumed // sorted in non-decreasing order, so its last element is the maximum - // segment id. This produces the same "Invalid number of segments" - // message as the CPU functor for the common (sorted) case where the - // incoming gradient simply doesn't have enough rows. - const SegmentId last_segment_id = - internal::SubtleMustCopy(segment_host_vec(nouter - 1)); - OP_REQUIRES(context, last_segment_id + 1 <= nsegments, + // segment id. A single element is enough to validate `nsegments` + // against the well-formed/sorted case and produce the same + // "Invalid number of segments" message the CPU functor uses for it. + ScratchSpace last_segment_id_host(context, 1, + /*on_host=*/true); + OP_REQUIRES_OK( + context, + stream->Memcpy( + last_segment_id_host.mutable_data(), + stream_executor::DeviceAddressBase( + const_cast(segment_vec.data()) + (nouter - 1), + sizeof(SegmentId)), + sizeof(SegmentId))); + OP_REQUIRES_OK(context, stream->BlockHostUntilDone()); + OP_REQUIRES(context, *last_segment_id_host.data() + 1 <= nsegments, absl::InvalidArgumentError("Invalid number of segments")); - for (Index i = 0; i < nouter; ++i) { - const Index output_idx = - internal::SubtleMustCopy(indices_host_vec(i)); - OP_REQUIRES(context, FastBoundsCheck(output_idx, noutput), + // Full per-element validation, matching the CPU functor: catches + // malformed input that isn't actually sorted (so the fast check above + // can't tell it's invalid from the last element alone). Runs as a + // device-side scan; `bad_positions` only ever carries 2 int32s back to + // the host, regardless of `nouter`. + Tensor bad_positions; + OP_REQUIRES_OK(context, context->allocate_temp( + DT_INT32, TensorShape({2}), &bad_positions)); + int32_t* const bad_positions_ptr = bad_positions.flat().data(); + int32_t* const bad_index_pos_ptr = bad_positions_ptr; + int32_t* const bad_segment_pos_ptr = bad_positions_ptr + 1; + stream_executor::DeviceAddressBase bad_positions_device( + bad_positions_ptr, 2 * sizeof(int32_t)); + OP_REQUIRES_OK(context, + stream->Memset32(&bad_positions_device, + std::numeric_limits::max(), + 2 * sizeof(int32_t))); + OP_REQUIRES_OK( + context, LaunchSparseSegmentGradBoundsCheckKernel( + device, nouter, noutput, nsegments, indices_vec.data(), + segment_vec.data(), bad_index_pos_ptr, + bad_segment_pos_ptr)); + ScratchSpace bad_positions_host(context, 2, /*on_host=*/true); + OP_REQUIRES_OK(context, + stream->Memcpy(bad_positions_host.mutable_data(), + bad_positions_device, + 2 * sizeof(int32_t))); + OP_REQUIRES_OK(context, stream->BlockHostUntilDone()); + const int32_t bad_index_pos = bad_positions_host.data()[0]; + const int32_t bad_segment_pos = bad_positions_host.data()[1]; + + if (bad_index_pos != std::numeric_limits::max()) { + ScratchSpace bad_index_host(context, 1, /*on_host=*/true); + OP_REQUIRES_OK( + context, + stream->Memcpy( + bad_index_host.mutable_data(), + stream_executor::DeviceAddressBase( + const_cast(indices_vec.data()) + bad_index_pos, + sizeof(Index)), + sizeof(Index))); + OP_REQUIRES_OK(context, stream->BlockHostUntilDone()); + OP_REQUIRES(context, false, + absl::InvalidArgumentError(absl::StrCat( + "Index ", *bad_index_host.data(), " out of range [0, ", + noutput, ")."))); + } + if (bad_segment_pos != std::numeric_limits::max()) { + ScratchSpace bad_segment_host(context, 1, + /*on_host=*/true); + OP_REQUIRES_OK( + context, + stream->Memcpy( + bad_segment_host.mutable_data(), + stream_executor::DeviceAddressBase( + const_cast(segment_vec.data()) + + bad_segment_pos, + sizeof(SegmentId)), + sizeof(SegmentId))); + OP_REQUIRES_OK(context, stream->BlockHostUntilDone()); + OP_REQUIRES(context, false, absl::InvalidArgumentError(absl::StrCat( - "Index ", output_idx, " out of range [0, ", noutput, - ")."))); - const SegmentId idx = internal::SubtleMustCopy(segment_host_vec(i)); - OP_REQUIRES( - context, FastBoundsCheck(idx, nsegments), - absl::InvalidArgumentError(absl::StrCat( - "Segment id ", idx, " out of range [0, ", nsegments, ")."))); + "Segment id ", *bad_segment_host.data(), + " out of range [0, ", nsegments, ")."))); } } From da78f1a32290f5caafcf82f5fee1ad7a7c2065d9 Mon Sep 17 00:00:00 2001 From: gauravarun21 Date: Thu, 17 Sep 2026 18:01:01 +0200 Subject: [PATCH 14/48] Reject nouter that doesn't fit in int32 before bounds-check kernel The bounds-check kernel packs each violation's position into an int32 buffer via GpuAtomicMin. If Index is int64_t and nouter exceeded INT32_MAX, a position beyond INT32_MAX would truncate to a negative value when narrowed, which would then be used as a negative pointer offset when fetching the offending value for the error message. Guard against this the same way SparseSegmentGradV2Functor already does for its own nouter/ninner (this file, above): reject the input with a clear error instead of proceeding with a value that doesn't fit. --- .../core/kernels/segment_reduction_ops_gpu.cu.h | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h index cbd40d8cf2ee7f..84a45d7bfa5182 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h +++ b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h @@ -1055,6 +1055,19 @@ struct SparseSegmentGradFunctor { // only a couple of small flags come back over the (synchronous) D2H // transfer in the common, valid-input case. if (nouter > 0) { + // The bounds-check kernel below packs each violation's position into + // an int32 (matching SparseSegmentGradV2Functor's existing + // `nouter`/`ninner` handling above, which notes that neither is + // expected to be huge), so reject inputs where a position wouldn't + // fit rather than silently truncating it. + OP_REQUIRES( + context, + static_cast(nouter) <= + std::numeric_limits::max(), + absl::InvalidArgumentError(absl::StrCat( + "Indices vector of length ", nouter, + " is too large to fit in int32."))); + se::Stream* stream = context->op_device_context()->stream(); OP_REQUIRES(context, stream != nullptr, absl::InternalError("No GPU stream available.")); From f448a3eb394667a8ce9a9c95d635c4dc0c011caf Mon Sep 17 00:00:00 2001 From: Levon Ter-Grigoryan Date: Thu, 24 Sep 2026 02:35:19 -0700 Subject: [PATCH 15/48] [XLA:GPU] Add collective memory space to layout and use it as collective memory space color. PiperOrigin-RevId: 987355696 --- .../backends/gpu/runtime/collective_thunk.cc | 7 +- .../gpu/tests/collective_ops_e2e_test.cc | 12 +- .../gpu/tests/collective_ops_ffi_test.cc | 16 +-- third_party/xla/xla/layout.h | 1 + .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc | 2 +- .../pjrt/gpu/se_gpu_topology_description.cc | 1 + .../pjrt/se/pjrt_stream_executor_client.cc | 1 + .../xla/pjrt/se/stream_executor_executable.cc | 1 + .../gpu/execution_stream_assignment_test.cc | 4 +- .../xla/xla/service/gpu/gpu_compiler_test.cc | 134 +++++++++--------- .../gpu/gpu_memory_space_assignment.cc | 18 ++- .../service/gpu/gpu_memory_space_assignment.h | 5 +- .../gpu/gpu_memory_space_assignment_test.cc | 34 ++--- third_party/xla/xla/xla.proto | 2 +- 14 files changed, 130 insertions(+), 108 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc index f4b30ced48c070..0ad331a60d8761 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc @@ -47,6 +47,7 @@ limitations under the License. #include "xla/hlo/ir/collective_op_group_mode.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/layout.h" #include "xla/primitive_util.h" #include "xla/runtime/buffer_use.h" #include "xla/runtime/device_id.h" @@ -71,8 +72,6 @@ limitations under the License. namespace xla::gpu { namespace { -static constexpr int64_t kCollectiveMemorySpaceColor = 1; - bool IsTypeSupportedBy(PrimitiveType element_type, Thunk::Kind reduction_op) { switch (element_type) { case S8: @@ -404,13 +403,13 @@ absl::Status CollectiveThunk::Prepare(const PrepareParams& params) { if (CanUseSymmetricBuffer() && config().use_symmetric_buffer) { for (const Buffer& buffer : buffers_) { - if (buffer.source_memory_space == kCollectiveMemorySpaceColor) { + if (buffer.source_memory_space == Layout::kCollectiveMemorySpace) { ABSL_RETURN_IF_ERROR( params.collective_memory_requests->RequestSymmetricAllocation( clique_key, buffer.source_buffer.slice.index())); } - if (buffer.destination_memory_space == kCollectiveMemorySpaceColor) { + if (buffer.destination_memory_space == Layout::kCollectiveMemorySpace) { ABSL_RETURN_IF_ERROR( params.collective_memory_requests->RequestSymmetricAllocation( clique_key, buffer.destination_buffer.slice.index())); diff --git a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc index 0f5721059b85ef..15fec5e22fed50 100644 --- a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc @@ -48,6 +48,7 @@ limitations under the License. #include "xla/hlo/testlib/pattern_matcher_gmock.h" #include "xla/hlo/testlib/verified_hlo_module.h" #include "xla/hlo/utils/hlo_matchers.h" +#include "xla/layout.h" #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/service/gpu/backend_configs.pb.h" @@ -3406,9 +3407,11 @@ ENTRY main { test_runner().HloModuleFromWrapped(executable.get())); const HloInstruction* ag_start = FindCollectiveStarts(executable_module, HloOpcode::kAllGather).at(0); - // Both ag and its producer should have collective memory space 1 - EXPECT_EQ(ag_start->shape().tuple_shapes()[1].layout().memory_space(), 1); - EXPECT_EQ(ag_start->operand(0)->shape().layout().memory_space(), 1); + // Both ag and its producer should have collective memory space. + EXPECT_EQ(ag_start->shape().tuple_shapes()[1].layout().memory_space(), + Layout::kCollectiveMemorySpace); + EXPECT_EQ(ag_start->operand(0)->shape().layout().memory_space(), + Layout::kCollectiveMemorySpace); } TEST_F(CollectiveOpsTestE2E, @@ -3457,7 +3460,8 @@ ROOT tuple = (bf16[1024,1024]{1,0}, bf16[]) tuple(all-reduce-done, all-reduce-do // space. for (auto ar : all_ar) { EXPECT_EQ(ar->operand(0)->opcode(), HloOpcode::kCopy); - EXPECT_EQ(ar->operand(0)->shape().layout().memory_space(), 1); + EXPECT_EQ(ar->operand(0)->shape().layout().memory_space(), + Layout::kCollectiveMemorySpace); } } diff --git a/third_party/xla/xla/backends/gpu/tests/collective_ops_ffi_test.cc b/third_party/xla/xla/backends/gpu/tests/collective_ops_ffi_test.cc index 7b02691885318b..8cce5f8a4b1236 100644 --- a/third_party/xla/xla/backends/gpu/tests/collective_ops_ffi_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/collective_ops_ffi_test.cc @@ -1227,8 +1227,8 @@ TEST_P(AllReduceTest, DeviceAllReduce) { ENTRY test_computation { id = u32[] replica-id() - in = u32[]{:S(1)} copy(id) - all-reduce = u32[]{:S(1)} custom-call(in), + in = u32[]{:S(7)} copy(id) + all-reduce = u32[]{:S(7)} custom-call(in), custom_call_target="__xla_test_$0_device_all_reduce", api_version=API_VERSION_TYPED_FFI ROOT out = u32[] copy(all-reduce) @@ -1316,7 +1316,7 @@ TEST_P(AllReduceTest, MulticastAllReduce) { ENTRY test_computation { c0 = u32[] constant(1) - in = u32[]{:S(1)} copy(c0) + in = u32[]{:S(7)} copy(c0) all-reduce = u32[] custom-call(in), custom_call_target="__xla_test_$0_multimem_all_reduce", api_version=API_VERSION_TYPED_FFI @@ -1360,7 +1360,7 @@ TEST_P(AllReduceTest, SymMulticastAllReduce) { ENTRY test_computation { c0 = u32[] constant(1) - in = u32[]{:S(1)} copy(c0) + in = u32[]{:S(7)} copy(c0) all-reduce = u32[] custom-call(in), custom_call_target="__xla_test_$0_sym_multimem_all_reduce", api_version=API_VERSION_TYPED_FFI @@ -1405,7 +1405,7 @@ TEST_P(AllReduceTest, SymPeerAllReduce) { ENTRY test_computation { id = u32[] replica-id() - in = u32[]{:S(1)} copy(id) + in = u32[]{:S(7)} copy(id) all-reduce = u32[] custom-call(in), custom_call_target="__xla_test_$0_sym_peer_all_reduce", api_version=API_VERSION_TYPED_FFI @@ -1442,7 +1442,7 @@ INSTANTIATE_TEST_SUITE_P( }); // Same as DeviceAllReduce, but uses frontend_attributes to specify memory -// spaces instead of hardcoded S(1). +// spaces instead of hardcoded S(7). TEST_F(CollectiveOpsTestFFI, DeviceAllReduceWithFrontendAttributes) { if (device_count() < kNumReplicas) { GTEST_SKIP() << "Test requires at least " << kNumReplicas << " devices (" @@ -1462,8 +1462,8 @@ TEST_F(CollectiveOpsTestFFI, DeviceAllReduceWithFrontendAttributes) { custom_call_target="__xla_test_blocked_device_all_reduce", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } ROOT out = u32[] copy(all-reduce) } diff --git a/third_party/xla/xla/layout.h b/third_party/xla/xla/layout.h index 276243c2971f61..c25424c463abc2 100644 --- a/third_party/xla/xla/layout.h +++ b/third_party/xla/xla/layout.h @@ -404,6 +404,7 @@ class Layout { static constexpr int64_t kDefaultMemorySpace = 0; static constexpr int64_t kGenericFastMemorySpace = 1; + static constexpr int64_t kCollectiveMemorySpace = 7; static constexpr int64_t kHostMemorySpace = 5; static constexpr int64_t kUnpinnedHostMemorySpace = 6; int64_t memory_space() const { return memory_space_; } diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc index 2feae674faa7b3..efcd52db06227a 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc @@ -3215,7 +3215,7 @@ constexpr char kAttrPlacedModule[] = R"( s = f32[512,1024] custom-call(t), custom_call_target="RecordBufferAddress", api_version=API_VERSION_TYPED_FFI, output_to_operand_aliasing={{}: (0, {})}, - frontend_attributes={results_memory_spaces="{0:1}"} + frontend_attributes={results_memory_spaces="{0:7}"} ROOT r = f32[512,1024] add(s, s) })"; diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc index a4cb5769b865ea..b25d2b2f6a63f9 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc @@ -425,6 +425,7 @@ StreamExecutorGpuTopologyDescription::GetMemorySpaceKindForShape( return GetMemorySpaceKindIds()[1]; break; case Layout::kGenericFastMemorySpace: + case Layout::kCollectiveMemorySpace: case Layout::kDefaultMemorySpace: break; default: diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc index 2d5ad0a0168708..362c40dfb1a82c 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc @@ -1761,6 +1761,7 @@ absl::StatusOr MemoryKindFromSimpleShape( case Layout::kHostMemorySpace: return PinnedHostMemorySpace::kKind; case Layout::kGenericFastMemorySpace: + case Layout::kCollectiveMemorySpace: case Layout::kDefaultMemorySpace: return default_memory_kind; default: diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable.cc b/third_party/xla/xla/pjrt/se/stream_executor_executable.cc index ace503c89b6b52..9c80957bd98a10 100644 --- a/third_party/xla/xla/pjrt/se/stream_executor_executable.cc +++ b/third_party/xla/xla/pjrt/se/stream_executor_executable.cc @@ -262,6 +262,7 @@ absl::StatusOr MemoryKindFromLayout( return PinnedHostMemorySpace::kKind; case Layout::kGenericFastMemorySpace: case Layout::kDefaultMemorySpace: + case Layout::kCollectiveMemorySpace: return default_memory_kind; default: return InvalidArgument("Unexpected memory space %d in output layout", diff --git a/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc b/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc index d50dd053f74420..76826c35af37c4 100644 --- a/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc +++ b/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc @@ -109,8 +109,8 @@ TEST_F(ExecutionStreamAssignmentTest, CopyStartStreamIdTest) { HloModule Module ENTRY CopyStartAndCopyDone { - p0 = f32[2,3]{1,0:S(1)} parameter(0) - copy-start = (f32[2,3]{1,0:S(2)}, f32[2,3]{1,0:S(1)}, u32[]) copy-start(p0) + p0 = f32[2,3]{1,0:S(7)} parameter(0) + copy-start = (f32[2,3]{1,0:S(2)}, f32[2,3]{1,0:S(7)}, u32[]) copy-start(p0) ROOT copy-done = f32[2,3]{1,0:S(2)} copy-done(copy-start) } )"; diff --git a/third_party/xla/xla/service/gpu/gpu_compiler_test.cc b/third_party/xla/xla/service/gpu/gpu_compiler_test.cc index 0283257a50c84e..d11dc58f45c9cf 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler_test.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler_test.cc @@ -1909,7 +1909,7 @@ TEST_F(GpuCompilerTest, MosaicMultimemRequiresSymmetricMemoryCopies) { p_multimem = s32[1] parameter(0) p_non_coll = s32[1] parameter(1) - cc_multimem = (s32[1]{0}) custom-call(p_multimem), custom_call_target="mosaic_gpu_v2", frontend_attributes={operands_memory_spaces="{0:1}",results_memory_spaces="{0:1}"}, api_version=API_VERSION_TYPED_FFI + cc_multimem = (s32[1]{0}) custom-call(p_multimem), custom_call_target="mosaic_gpu_v2", frontend_attributes={operands_memory_spaces="{0:7}",results_memory_spaces="{0:7}"}, api_version=API_VERSION_TYPED_FFI res_multimem = s32[1] get-tuple-element(cc_multimem), index=0 cc_non_coll = (s32[1]{0}) custom-call(p_non_coll), custom_call_target="mosaic_gpu_v2", api_version=API_VERSION_TYPED_FFI @@ -1931,16 +1931,16 @@ TEST_F(GpuCompilerTest, MosaicMultimemRequiresSymmetricMemoryCopies) { // CHECK-DAG: [[P_NON:%[^ ]+]] = s32[1]{0} parameter(1) // Multimem input parameters are copied to S1 - // CHECK-DAG: [[COPY_MULTI_IN:%copy[^ ]*]] = s32[1]{0:S(1)} copy([[P_MULTI]]) + // CHECK-DAG: [[COPY_MULTI_IN:%copy[^ ]*]] = s32[1]{0:S(7)} copy([[P_MULTI]]) - // CHECK-DAG: [[CC_MULTI:%[^ ]+]] = (s32[1]{0:S(1)}) custom-call([[COPY_MULTI_IN]]){{.*}}results_memory_spaces={0:1} + // CHECK-DAG: [[CC_MULTI:%[^ ]+]] = (s32[1]{0:S(7)}) custom-call([[COPY_MULTI_IN]]){{.*}}results_memory_spaces={0:7} // CHECK-DAG: [[CC_NON:%[^ ]+]] = (s32[1]{0}) custom-call([[P_NON]]) // Extracting from the 1-element tuples returned by custom calls (all index=0) - // CHECK-DAG: [[GTE_MULTI:%[^ ]+]] = s32[1]{0:S(1)} get-tuple-element([[CC_MULTI]]), index=0 + // CHECK-DAG: [[GTE_MULTI:%[^ ]+]] = s32[1]{0:S(7)} get-tuple-element([[CC_MULTI]]), index=0 // CHECK-DAG: [[GTE_NON:%[^ ]+]] = s32[1]{0} get-tuple-element([[CC_NON]]), index=0 - // The S1 -> S0 isolation copies + // The S7 -> S0 isolation copies // CHECK-DAG: [[COPY_OUT_MULTI:%copy[^ ]*]] = s32[1]{0} copy([[GTE_MULTI]]) // The final ROOT is pure S0 @@ -1972,7 +1972,7 @@ TEST_F(GpuCompilerTest, MosaicCollectiveMetadataRequiresSymmetricMemoryCopies) { HloModule test ENTRY main { p = s32[1] parameter(0) - mosaic = (s32[1]{0}) custom-call(p), custom_call_target="mosaic_gpu_v2", api_version=API_VERSION_TYPED_FFI, backend_config={uses_xla_collective_metadata=true}, frontend_attributes={operands_memory_spaces="{0:1}", results_memory_spaces="{0:1}"} + mosaic = (s32[1]{0}) custom-call(p), custom_call_target="mosaic_gpu_v2", api_version=API_VERSION_TYPED_FFI, backend_config={uses_xla_collective_metadata=true}, frontend_attributes={operands_memory_spaces="{0:7}", results_memory_spaces="{0:7}"} ROOT res = s32[1]{0} get-tuple-element(mosaic), index=0 } )"; @@ -1993,9 +1993,9 @@ TEST_F(GpuCompilerTest, MosaicCollectiveMetadataRequiresSymmetricMemoryCopies) { constexpr absl::string_view kExpected = R"( // CHECK: [[P:%[^ ]+]] = s32[1]{0} parameter(0) - // CHECK: [[COPY_IN:%copy[^ ]*]] = s32[1]{0:S(1)} copy([[P]]) - // CHECK: [[MOSAIC:%[^ ]+]] = (s32[1]{0:S(1)}) custom-call([[COPY_IN]]){{.*}}backend_config={uses_xla_collective_metadata=true} - // CHECK: [[GTE:%[^ ]+]] = s32[1]{0:S(1)} get-tuple-element([[MOSAIC]]), index=0 + // CHECK: [[COPY_IN:%copy[^ ]*]] = s32[1]{0:S(7)} copy([[P]]) + // CHECK: [[MOSAIC:%[^ ]+]] = (s32[1]{0:S(7)}) custom-call([[COPY_IN]]){{.*}}backend_config={uses_xla_collective_metadata=true} + // CHECK: [[GTE:%[^ ]+]] = s32[1]{0:S(7)} get-tuple-element([[MOSAIC]]), index=0 // CHECK: ROOT [[COPY_OUT:%copy[^ ]*]] = s32[1]{0} copy([[GTE]]) )"; EXPECT_THAT( @@ -2064,8 +2064,8 @@ TEST_P(OneShotRaggedAllToAllMemSpaceTest, DirectUsage) { constexpr absl::string_view kS1TwoCopies = R"( // CHECK: %output = f32[16]{0} parameter(1) - // CHECK: [[COPY1:%copy[0-9.]*]] = f32[16]{0:S(1)} copy(%output) - // CHECK: [[RA2A:%[^ ]+]] = f32[16]{0:S(1)} ragged-all-to-all(%input, [[COPY1]], + // CHECK: [[COPY1:%copy[0-9.]*]] = f32[16]{0:S(7)} copy(%output) + // CHECK: [[RA2A:%[^ ]+]] = f32[16]{0:S(7)} ragged-all-to-all(%input, [[COPY1]], // CHECK: ROOT %copy.{{[0-9]+}} = f32[16]{0} copy([[RA2A]]) )"; @@ -2140,10 +2140,10 @@ TEST_P(OneShotRaggedAllToAllMemSpaceTest, LoopUsage) { constexpr absl::string_view kS1TwoCopies = R"( // CHECK: [[OUTPUT1:%output[0-9.]*]] = f32[16]{0} parameter(1) - // CHECK: [[COPY1:%copy[0-9.]*]] = f32[16]{0:S(1)} copy([[OUTPUT1]]) - // CHECK: %tuple = (s32[], f32[16]{0}, f32[16]{0:S(1)}) tuple(%copy{{.*}}, %input{{.*}}, [[COPY1]]) - // CHECK: %while = (s32[], f32[16]{0}, f32[16]{0:S(1)}) while(%tuple) - // CHECK: [[RESULT:%result[0-9.]*]] = f32[16]{0:S(1)} get-tuple-element(%while), index=2 + // CHECK: [[COPY1:%copy[0-9.]*]] = f32[16]{0:S(7)} copy([[OUTPUT1]]) + // CHECK: %tuple = (s32[], f32[16]{0}, f32[16]{0:S(7)}) tuple(%copy{{.*}}, %input{{.*}}, [[COPY1]]) + // CHECK: %while = (s32[], f32[16]{0}, f32[16]{0:S(7)}) while(%tuple) + // CHECK: [[RESULT:%result[0-9.]*]] = f32[16]{0:S(7)} get-tuple-element(%while), index=2 // CHECK: ROOT %copy{{.*}} = f32[16]{0} copy([[RESULT]]) )"; @@ -2212,8 +2212,8 @@ ENTRY test_computation { )"; constexpr absl::string_view kS1TwoCopies = R"( - // CHECK: [[COPY0:%copy[0-9.]*]] = u32[2]{0:S(1)} copy(%p) - // CHECK: [[PERMUTE:%[^ ]+]] = u32[2]{0:S(1)} collective-permute([[COPY0]]) + // CHECK: [[COPY0:%copy[0-9.]*]] = u32[2]{0:S(7)} copy(%p) + // CHECK: [[PERMUTE:%[^ ]+]] = u32[2]{0:S(7)} collective-permute([[COPY0]]) // CHECK: ROOT %copy{{.*}} = u32[2]{0} copy([[PERMUTE]]) )"; @@ -2314,10 +2314,10 @@ TEST_P(RequiresCollectiveSymmetricMemorySpaceTest, LoopUsage) { constexpr absl::string_view kS1TwoCopies = R"( // CHECK:ENTRY %entry_computation (input: u32[2]) -> u32[2] { // CHECK: %input = u32[2]{0} parameter(0) - // CHECK: [[COPY1:%copy[0-9.]*]] = u32[2]{0:S(1)} copy(%input) - // CHECK: %tuple = (s32[], u32[2]{0:S(1)}) tuple(%copy{{.*}}, [[COPY1]]) - // CHECK: %while = (s32[], u32[2]{0:S(1)}) while(%tuple) - // CHECK: [[RESULT:%result[0-9.]*]] = u32[2]{0:S(1)} get-tuple-element(%while) + // CHECK: [[COPY1:%copy[0-9.]*]] = u32[2]{0:S(7)} copy(%input) + // CHECK: %tuple = (s32[], u32[2]{0:S(7)}) tuple(%copy{{.*}}, [[COPY1]]) + // CHECK: %while = (s32[], u32[2]{0:S(7)}) while(%tuple) + // CHECK: [[RESULT:%result[0-9.]*]] = u32[2]{0:S(7)} get-tuple-element(%while) // CHECK: ROOT %copy{{.*}} = u32[2]{0} copy([[RESULT]]) )"; @@ -2428,8 +2428,8 @@ TEST_P(GpuCompilerParametersCopyCollectiveMemoryTest, DirectUsage) { )"; constexpr absl::string_view kS1TwoCopies = R"( - // CHECK: [[COPY_IN:%copy[0-9.]*]] = s32[1]{0:S(1)} copy(%parameter_used_by_collective) - // CHECK: [[ALL_REDUCE:%[^ ]+]] = s32[1]{0:S(1)} all-reduce([[COPY_IN]]) + // CHECK: [[COPY_IN:%copy[0-9.]*]] = s32[1]{0:S(7)} copy(%parameter_used_by_collective) + // CHECK: [[ALL_REDUCE:%[^ ]+]] = s32[1]{0:S(7)} all-reduce([[COPY_IN]]) // CHECK: ROOT %copy.{{[0-9]+}} = s32[1]{0} copy([[ALL_REDUCE]]) )"; @@ -2538,10 +2538,10 @@ TEST_P(GpuCompilerParametersCopyCollectiveMemoryTest, LoopUsage) { constexpr absl::string_view kS1TwoCopies = R"( // CHECK: %input = s32[1]{0} parameter(0) - // CHECK: %copy.{{[0-9]+}} = s32[1]{0:S(1)} copy(%input) - // CHECK: %tuple = (s32[], s32[1]{0:S(1)}) tuple({{.*}}, %copy.{{[0-9]+}}) - // CHECK: %while = (s32[], s32[1]{0:S(1)}) while(%tuple) - // CHECK: %result.2.0 = s32[1]{0:S(1)} get-tuple-element(%while), index=1 + // CHECK: %copy.{{[0-9]+}} = s32[1]{0:S(7)} copy(%input) + // CHECK: %tuple = (s32[], s32[1]{0:S(7)}) tuple({{.*}}, %copy.{{[0-9]+}}) + // CHECK: %while = (s32[], s32[1]{0:S(7)}) while(%tuple) + // CHECK: %result.2.0 = s32[1]{0:S(7)} get-tuple-element(%while), index=1 // CHECK: ROOT %copy.{{[0-9]+}} = s32[1]{0} copy(%result.2.0) )"; @@ -2754,8 +2754,8 @@ TEST_P(FrontendAttributesMemorySpaceTest, DirectUsage) { custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -2779,8 +2779,8 @@ TEST_P(FrontendAttributesMemorySpaceTest, DirectUsage) { // Regardless of the input_output_alias it should be two copies constexpr absl::string_view expected_check = R"( // CHECK: %p = f32[16]{0} parameter(0) - // CHECK: [[COPY0:%copy[0-9.]*]] = f32[16]{0:S(1)} copy(%p) - // CHECK: %cc = f32[16]{0:S(1)} custom-call([[COPY0]]) + // CHECK: [[COPY0:%copy[0-9.]*]] = f32[16]{0:S(7)} copy(%p) + // CHECK: %cc = f32[16]{0:S(7)} custom-call([[COPY0]]) // CHECK: ROOT %copy{{.*}} = f32[16]{0} copy(%cc) )"; @@ -2811,8 +2811,8 @@ TEST_P(FrontendAttributesMemorySpaceTest, LoopUsage) { custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } new_loop_counter = s32[] add(loop_counter, s32[] constant(1)) ROOT result = (s32[], f32[16]) tuple(new_loop_counter, cc) @@ -2846,10 +2846,10 @@ TEST_P(FrontendAttributesMemorySpaceTest, LoopUsage) { // Regardless of the input_output_alias it should be two copies constexpr absl::string_view expected_check = R"( // CHECK: %input = f32[16]{0} parameter(0) - // CHECK: [[COPY0:%copy[0-9.]*]] = f32[16]{0:S(1)} copy(%input) - // CHECK: %tuple = (s32[], f32[16]{0:S(1)}) tuple(%copy{{.*}}, [[COPY0]]) - // CHECK: %while = (s32[], f32[16]{0:S(1)}) while(%tuple) - // CHECK: [[RESULT:%result[0-9.]*]] = f32[16]{0:S(1)} get-tuple-element(%while), index=1 + // CHECK: [[COPY0:%copy[0-9.]*]] = f32[16]{0:S(7)} copy(%input) + // CHECK: %tuple = (s32[], f32[16]{0:S(7)}) tuple(%copy{{.*}}, [[COPY0]]) + // CHECK: %while = (s32[], f32[16]{0:S(7)}) while(%tuple) + // CHECK: [[RESULT:%result[0-9.]*]] = f32[16]{0:S(7)} get-tuple-element(%while), index=1 // CHECK: ROOT %copy{{.*}} = f32[16]{0} copy([[RESULT]]) )"; @@ -2867,13 +2867,13 @@ TEST_F(GpuCompilerTest, HloModule test, input_output_alias={ {}: (0, {}) } ENTRY test_computation { - p = f32[16]{0:S(1)} parameter(0) - ROOT cc = f32[16]{0:S(1)} custom-call(p), + p = f32[16]{0:S(7)} parameter(0) + ROOT cc = f32[16]{0:S(7)} custom-call(p), custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -2923,13 +2923,13 @@ TEST_F(GpuCompilerTest, HloModule test, input_output_alias={ {}: (0, {}) } ENTRY test_computation { - p = f32[16]{0:S(1)} parameter(0) - ROOT cc = f32[16]{0:S(1)} custom-call(p), + p = f32[16]{0:S(7)} parameter(0) + ROOT cc = f32[16]{0:S(7)} custom-call(p), custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -2963,13 +2963,13 @@ TEST_F(GpuCompilerTest, HloModule test ENTRY test_computation { - p = f32[16]{0:S(1)} parameter(0) - ROOT cc = f32[16]{0:S(1)} custom-call(p), + p = f32[16]{0:S(7)} parameter(0) + ROOT cc = f32[16]{0:S(7)} custom-call(p), custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -3004,8 +3004,8 @@ TEST_F(GpuCompilerTest, custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -3040,8 +3040,8 @@ TEST_F(GpuCompilerTest, custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -3071,16 +3071,16 @@ TEST_F(GpuCompilerTest, HloModule test, input_output_alias={ {0}: (0, {}) } ENTRY test_computation { - p0 = f32[16]{0:S(1)} parameter(0) - cc = (f32[16]{0:S(1)}, f32[16]{0:S(1)}) custom-call(p0), + p0 = f32[16]{0:S(7)} parameter(0) + cc = (f32[16]{0:S(7)}, f32[16]{0:S(7)}) custom-call(p0), custom_call_target="mock_call", frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1,1:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7,1:7}" } - elem0 = f32[16]{0:S(1)} get-tuple-element(cc), index=0 - elem1 = f32[16]{0:S(1)} get-tuple-element(cc), index=1 - ROOT out = (f32[16]{0:S(1)}, f32[16]{0:S(1)}) tuple(elem1, elem0) + elem0 = f32[16]{0:S(7)} get-tuple-element(cc), index=0 + elem1 = f32[16]{0:S(7)} get-tuple-element(cc), index=1 + ROOT out = (f32[16]{0:S(7)}, f32[16]{0:S(7)}) tuple(elem1, elem0) } )"; @@ -3243,7 +3243,7 @@ TEST_F(GpuCompilerTest, SymmetricBuffersFilter) { // - channel 3 does NOT have S(1) (f32, 8192 bytes - filtered out by size) constexpr absl::string_view expected_check = R"( - // CHECK-DAG: f32[1024]{0:S(1)}{{.*}}all-reduce{{.*}}channel_id=1 + // CHECK-DAG: f32[1024]{0:S(7)}{{.*}}all-reduce{{.*}}channel_id=1 // CHECK-DAG: s32[1024]{0}{{.*}}all-reduce{{.*}}channel_id=2 // CHECK-DAG: f32[2048]{0}{{.*}}all-reduce{{.*}}channel_id=3 )"; @@ -3300,8 +3300,8 @@ TEST_F(GpuCompilerTest, SymmetricBuffersMultipleCollectives) { const HloModule* optimized_module = optimized_module_and_executable.first; constexpr absl::string_view expected_check = R"( - // CHECK-DAG: f32[1024]{0:S(1)}{{.*}}all-reduce{{.*}}channel_id=1 - // CHECK-DAG: f32[1024]{0:S(1)}{{.*}}all-gather{{.*}}channel_id=2 + // CHECK-DAG: f32[1024]{0:S(7)}{{.*}}all-reduce{{.*}}channel_id=1 + // CHECK-DAG: f32[1024]{0:S(7)}{{.*}}all-gather{{.*}}channel_id=2 // CHECK-DAG: f32[2048]{0}{{.*}}all-reduce{{.*}}channel_id=3 )"; @@ -3363,8 +3363,8 @@ TEST_F(GpuCompilerTest, SymmetricBuffersSeveralFilters) { const HloModule* optimized_module = optimized_module_and_executable.first; constexpr absl::string_view expected_check = R"( - // CHECK-DAG: f32[1024]{0:S(1)}{{.*}}all-reduce{{.*}}channel_id=1 - // CHECK-DAG: s32[2048]{0:S(1)}{{.*}}all-gather{{.*}}channel_id=2 + // CHECK-DAG: f32[1024]{0:S(7)}{{.*}}all-reduce{{.*}}channel_id=1 + // CHECK-DAG: s32[2048]{0:S(7)}{{.*}}all-gather{{.*}}channel_id=2 )"; EXPECT_THAT(RunFileCheck( @@ -3425,8 +3425,8 @@ TEST_F(GpuCompilerTest, SymmetricBuffersOverlappingFilters) { const HloModule* optimized_module = optimized_module_and_executable.first; constexpr absl::string_view expected_check = R"( - // CHECK-DAG: f32[1024]{0:S(1)}{{.*}}all-reduce{{.*}}channel_id=1 - // CHECK-DAG: f32[2048]{0:S(1)}{{.*}}all-reduce{{.*}}channel_id=2 + // CHECK-DAG: f32[1024]{0:S(7)}{{.*}}all-reduce{{.*}}channel_id=1 + // CHECK-DAG: f32[2048]{0:S(7)}{{.*}}all-reduce{{.*}}channel_id=2 )"; EXPECT_THAT(RunFileCheck( diff --git a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc index c9eda097160eea..d598c075274480 100644 --- a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc +++ b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc @@ -69,6 +69,10 @@ const absl::NoDestructor> absl::StatusOr AsMemorySpaceColor(int64_t memory_space) { switch (memory_space) { + case 1: + // Legacy value for collective memory space before + // Layout::kCollectiveMemorySpace (7) was introduced + return MemorySpaceColor::kCollective; case static_cast(MemorySpaceColor::kDefault): case static_cast(MemorySpaceColor::kCollective): case static_cast(MemorySpaceColor::kTempBuffer): @@ -76,8 +80,9 @@ absl::StatusOr AsMemorySpaceColor(int64_t memory_space) { default: return InvalidArgument( "Invalid memory space %d. " - "Valid values are 0 (default), 1 (collective), 2 (temp).", - memory_space); + "Valid values are %d (default), %d (collective), %d (temp).", + memory_space, MemorySpaceColor::kDefault, + MemorySpaceColor::kCollective, MemorySpaceColor::kTempBuffer); } } @@ -265,8 +270,15 @@ absl::StatusOr DetermineBufferColor( // space from the layout. const HloPosition& defining_position = value->defining_position(); if (defining_position.shape().has_layout()) { - const BufferValue::Color memory_space = + BufferValue::Color memory_space = defining_position.shape().layout().memory_space(); + if (memory_space == 1) { + // Legacy value for collective memory space before + // Layout::kCollectiveMemorySpace (7) was introduced. + memory_space = + static_cast(MemorySpaceColor::kCollective); + } + if (memory_space != 0) { candidates.push_back(memory_space); } diff --git a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h index 6fbda3c1cdd28a..8433f555458bca 100644 --- a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h +++ b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h @@ -23,6 +23,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/layout.h" #include "xla/service/buffer_assignment.h" #include "xla/xla.pb.h" @@ -39,11 +40,11 @@ inline constexpr absl::string_view kResultsMemorySpacesAttr = enum class MemorySpaceColor { // Corresponds to stream_executor::MemoryTypes::kDefault or kUnified. // This memory can be allocated with any device allocation API. - kDefault = 0, + kDefault = Layout::kDefaultMemorySpace, // Corresponds to stream_executor::MemoryTypes::kCollective. This memory // should be compatible with symmetric memory requirements. - kCollective = 1, + kCollective = Layout::kCollectiveMemorySpace, // Temp buffers can be allocated within separate memory space (if // xla_gpu_temp_buffer_use_separate_color is set). This improves cuda-graphs diff --git a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc index 814f38efdbd121..e2605cbe591816 100644 --- a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc +++ b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc @@ -441,7 +441,7 @@ TEST_F(GpuMemorySpaceAssignmentTest, HloModule m ENTRY main { - ROOT %custom-call.9 = (f16[8], f16[8]) custom-call(), custom_call_target="mosaic_gpu_v2", frontend_attributes={results_memory_spaces="{0:1}"} + ROOT %custom-call.9 = (f16[8], f16[8]) custom-call(), custom_call_target="mosaic_gpu_v2", frontend_attributes={results_memory_spaces="{0:7}"} } )"; @@ -474,14 +474,14 @@ TEST_F(GpuMemorySpaceAssignmentTest, } TEST(ParseIndexMemorySpacePairsTest, SinglePair) { - ASSERT_OK_AND_ASSIGN(auto pairs, ParseIndexMemorySpacePairs("{0:1}")); + ASSERT_OK_AND_ASSIGN(auto pairs, ParseIndexMemorySpacePairs("{0:7}")); ASSERT_EQ(pairs.size(), 1); EXPECT_EQ(pairs[0].first, 0); EXPECT_EQ(pairs[0].second, MemorySpaceColor::kCollective); } TEST(ParseIndexMemorySpacePairsTest, MultiplePairs) { - ASSERT_OK_AND_ASSIGN(auto pairs, ParseIndexMemorySpacePairs("{0:1,2:2}")); + ASSERT_OK_AND_ASSIGN(auto pairs, ParseIndexMemorySpacePairs("{0:7,2:2}")); ASSERT_EQ(pairs.size(), 2); EXPECT_EQ(pairs[0].first, 0); EXPECT_EQ(pairs[0].second, MemorySpaceColor::kCollective); @@ -496,7 +496,7 @@ TEST(ParseIndexMemorySpacePairsTest, EmptyBraces) { TEST(ParseIndexMemorySpacePairsTest, WhitespaceHandling) { ASSERT_OK_AND_ASSIGN(auto pairs, - ParseIndexMemorySpacePairs("{ 0 : 1 , 2 : 0 }")); + ParseIndexMemorySpacePairs("{ 0 : 7 , 2 : 0 }")); ASSERT_EQ(pairs.size(), 2); EXPECT_EQ(pairs[0].first, 0); EXPECT_EQ(pairs[0].second, MemorySpaceColor::kCollective); @@ -505,7 +505,7 @@ TEST(ParseIndexMemorySpacePairsTest, WhitespaceHandling) { } TEST(ParseIndexMemorySpacePairsTest, MissingBraces) { - EXPECT_FALSE(ParseIndexMemorySpacePairs("0:1").ok()); + EXPECT_FALSE(ParseIndexMemorySpacePairs("0:7").ok()); } TEST(ParseIndexMemorySpacePairsTest, InvalidPairFormat) { @@ -542,7 +542,7 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallOperandMemorySpace) { p1 = f32[1024]{0} parameter(1) ROOT custom-call = f32[1024]{0} custom-call(p0, p1), custom_call_target="my_custom_call", - frontend_attributes={operands_memory_spaces="{0:1}"} + frontend_attributes={operands_memory_spaces="{0:7}"} } )"; @@ -557,8 +557,9 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallOperandMemorySpace) { DependencyHloOrdering ordering(module.get()); EXPECT_OK(colorer(alias_analysis.get(), ordering)); - // Operand 0 (p0) should be colored with memory space 1. - EXPECT_EQ(FindColorByName(*alias_analysis, "p0"), 1); + // Operand 0 (p0) should be colored with memory space kCollective. + EXPECT_EQ(FindColorByName(*alias_analysis, "p0"), + (int)MemorySpaceColor::kCollective); // Operand 1 (p1) should remain in default memory space. EXPECT_EQ(FindColorByName(*alias_analysis, "p1"), (int)MemorySpaceColor::kDefault); @@ -572,7 +573,7 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallResultMemorySpace) { p0 = f32[1024]{0} parameter(0) ROOT custom-call = f32[1024]{0} custom-call(p0), custom_call_target="my_custom_call", - frontend_attributes={results_memory_spaces="{0:1}"} + frontend_attributes={results_memory_spaces="{0:7}"} } )"; @@ -587,8 +588,9 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallResultMemorySpace) { DependencyHloOrdering ordering(module.get()); EXPECT_OK(colorer(alias_analysis.get(), ordering)); - // The custom call result should be colored with memory space 1. - EXPECT_EQ(FindColorByName(*alias_analysis, "custom-call"), 1); + // The custom call result should be colored with memory space kCollective. + EXPECT_EQ(FindColorByName(*alias_analysis, "custom-call"), + (int)MemorySpaceColor::kCollective); } TEST_F(GpuMemorySpaceAssignmentTest, CustomCallTupleResultMemorySpace) { @@ -599,7 +601,7 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallTupleResultMemorySpace) { p0 = f32[1024]{0} parameter(0) custom-call = (f32[1024]{0}, f32[512]{0}) custom-call(p0), custom_call_target="my_custom_call", - frontend_attributes={results_memory_spaces="{0:1,1:1}"} + frontend_attributes={results_memory_spaces="{0:7,1:7}"} ROOT gte = f32[1024]{0} get-tuple-element(custom-call), index=0 } )"; @@ -621,8 +623,8 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallTupleResultMemorySpace) { if (value->instruction()->name() != "custom-call") continue; const ShapeIndex& idx = value->defining_index(); if (idx.size() == 1 && (idx[0] == 0 || idx[0] == 1)) { - EXPECT_EQ(value->color(), 1) - << "Tuple element " << idx[0] << " should have color 1"; + EXPECT_EQ(value->color(), (int)MemorySpaceColor::kCollective) + << "Tuple element " << idx[0] << " should have color kCollective"; } } } @@ -642,8 +644,8 @@ TEST_P(GpuMosaicCollectiveMemorySpaceAssignmentTest, HloModule m ENTRY main { ROOT %custom-call.9 = (f16[8])", - IsMosaicWithCollectiveMetadata() ? "{0:S(1)}" : "", R"(, f16[8])", - IsMosaicWithCollectiveMetadata() ? "{0:S(1)}" : "", + IsMosaicWithCollectiveMetadata() ? "{0:S(7)}" : "", R"(, f16[8])", + IsMosaicWithCollectiveMetadata() ? "{0:S(7)}" : "", R"() custom-call(), custom_call_target="mosaic_gpu_v2" } )"); diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto index 36f544733c1151..44ef782f522cc1 100644 --- a/third_party/xla/xla/xla.proto +++ b/third_party/xla/xla/xla.proto @@ -1867,7 +1867,7 @@ message DebugOptions { // for all instructions. Accepts textproto of autotuner.CandidateConfigs. optional string xla_candidate_configs_file = 544; - // If true, allows skipping defensive copy insertion for S(1) collective + // If true, allows skipping defensive copy insertion for collective // memory parameters that have input-output aliasing and execute on all // available devices in the topology. optional bool xla_gpu_enable_persistent_symmetric_memory = 545; From 3445433c0b89c978bad184087669abfee25c7ccc Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 02:55:49 -0700 Subject: [PATCH 16/48] PR #49314: Bump astral-sh/setup-uv from 10.0.1 to 10.1.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/49314 Bumps [astral-sh/setup-uv](https://github.com/astral-sh/setup-uv) from 10.0.1 to 10.1.0.
Release notes

Sourced from astral-sh/setup-uv's releases.

v10.1.0 🌈 New output python-runtime-idand respect NO_PROXY

Changes

This release adds more bheind the scene security improvements and also 2 small improvements.

NO_PROXY

This action now respects no_proxy/NO_PROXY environment variables which were previously ignored.

New output python-runtime-id

The new output python-runtime-id can be used to know which python version exactly was installed if you use activate-environment. See pyca/cryptography#15572 for details on why this can be useful.

🐛 Bug fixes

🚀 Enhancements

🧰 Maintenance

📚 Documentation

⬆️ Dependency updates

Commits

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=astral-sh/setup-uv&package-manager=github_actions&previous-version=10.0.1&new-version=10.1.0)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Copybara import of the project: -- 454adc369ab7b2a599e818b54fdb1ced4f464ea6 by dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>: Bump astral-sh/setup-uv from 10.0.1 to 10.1.0 Bumps [astral-sh/setup-uv](https://github.com/astral-sh/setup-uv) from 10.0.1 to 10.1.0. - [Release notes](https://github.com/astral-sh/setup-uv/releases) - [Commits](https://github.com/astral-sh/setup-uv/compare/20cfd1bf945f4377ade1205e4dbc17946fc9a30d...bec219d24cd3e171d82865faccec33120bb574f4) --- updated-dependencies: - dependency-name: astral-sh/setup-uv dependency-version: 10.1.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Merging this change closes #49314 PiperOrigin-RevId: 987366858 --- third_party/xla/.github/workflows/clang_format.yml | 2 +- third_party/xla/.github/workflows/zizmor.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/xla/.github/workflows/clang_format.yml b/third_party/xla/.github/workflows/clang_format.yml index 359d5b78d2a2d7..e267d670b023d8 100644 --- a/third_party/xla/.github/workflows/clang_format.yml +++ b/third_party/xla/.github/workflows/clang_format.yml @@ -41,7 +41,7 @@ jobs: - name: "Fetch HEAD of main branch" run: git fetch origin main --depth=1 - name: Install uv dependencies - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 - name: "Run clang-format check" env: TARGET_REF: "origin/${{ github.base_ref || 'main' }}" diff --git a/third_party/xla/.github/workflows/zizmor.yml b/third_party/xla/.github/workflows/zizmor.yml index 3bcae86a2acdb9..ee405f39080dfb 100644 --- a/third_party/xla/.github/workflows/zizmor.yml +++ b/third_party/xla/.github/workflows/zizmor.yml @@ -33,7 +33,7 @@ jobs: persist-credentials: false - name: Install uv dependencies - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 - name: Run zizmor run: uvx zizmor --format=github .github xla From 2dd5ba03ea49be8d64c4f27d2df1742f5a28df55 Mon Sep 17 00:00:00 2001 From: linchen1-robot Date: Thu, 24 Sep 2026 03:08:52 -0700 Subject: [PATCH 17/48] PR #49372: [ROCm]Enable vectorization.hlo check on ROCm mi200 and mi350. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/49372 📝 Summary of Changes - Enable `//xla/backends/gpu/tests:vectorization.hlo.test` on mi200 and mi350. Same `gpu/` lit as CUDA. Removes those two specs from the disabled list. - FileCheck LLVM `store <4 x i8>` at `--stage=llvm-before-optimizations`, the same stage and width as the H100/A100 lines. - Gate `--stage=ptx` with `%if !IS_ROCM`. `AMDGPUCompiler::CompileTargetBinary` returns an empty asm text, and `hlo-opt` has no GCN stage. NVIDIA PTX checks stay. 🎯 Justification Explain why this change is important and which workload benefits from this change. 🚀 Kind of Contribution 🧪 Tests Copybara import of the project: -- 9b122c87b8c8923be80148b641034d1a5330d3ad by linchen1 : Enable vectorization.hlo check on ROCm mi200 and mi350. Signed-off-by: linchen1 Merging this change closes #49372 PiperOrigin-RevId: 987373576 --- third_party/xla/xla/backends/gpu/tests/BUILD | 2 -- third_party/xla/xla/backends/gpu/tests/vectorization.hlo | 4 +++- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 7bc17e886098fe..f41f0d00fde282 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -766,7 +766,6 @@ lit_test_suite_for_gpus( "single_instruction.hlo", "reduce_unnested.hlo", "reduction_vectorization_sm_all.hlo", - "vectorization.hlo", "pdl.hlo", ], "mi300": [ @@ -779,7 +778,6 @@ lit_test_suite_for_gpus( "mi350": [ "dot_bf16.hlo", "single_instruction.hlo", - "vectorization.hlo", "pdl.hlo", ], "mi450": [ diff --git a/third_party/xla/xla/backends/gpu/tests/vectorization.hlo b/third_party/xla/xla/backends/gpu/tests/vectorization.hlo index 56a4747a3195e8..d2cbe6e7058d1c 100644 --- a/third_party/xla/xla/backends/gpu/tests/vectorization.hlo +++ b/third_party/xla/xla/backends/gpu/tests/vectorization.hlo @@ -1,5 +1,5 @@ // RUN: hlo-opt %s --platform=gpu --stage=llvm-before-optimizations --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb | FileCheck --check-prefixes=CHECK-LLVM-%{GPU} %s -// RUN: hlo-opt %s --platform=gpu --stage=ptx --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb | FileCheck --check-prefixes=CHECK-PTX-%{GPU} %s +// RUN: %if !IS_ROCM %{ hlo-opt %s --platform=gpu --stage=ptx --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb | FileCheck --check-prefixes=CHECK-PTX-%{GPU} %s %} // CHECK-LLVM-b200: store <16 x i8> // CHECK-PTX-b200: st.global.v4.b32 @@ -12,6 +12,8 @@ // CHECK-PTX-a100_pcie_80: st.global.b32 // CHECK-LLVM-a6000: store <4 x i8> // CHECK-PTX-a6000: st.global.b32 +// CHECK-LLVM-mi200: store <4 x i8> +// CHECK-LLVM-mi350: store <4 x i8> HloModule m8, is_scheduled=true From f64587a59116560752ba450e9ad6791719a31a90 Mon Sep 17 00:00:00 2001 From: Christian Sigg Date: Thu, 24 Sep 2026 03:09:29 -0700 Subject: [PATCH 18/48] Update TensorIR to v0.1.4. PiperOrigin-RevId: 987373870 --- .../tensor_ir/patches/pow_op.patch | 36 +++ .../xla/third_party/tensor_ir/patches/series | 1 + .../xla/third_party/tensor_ir/tensor_ir.BUILD | 290 ++++++++++-------- .../xla/third_party/tensor_ir/workspace.bzl | 4 +- .../xla/xla/backends/gpu/codegen/tools/BUILD | 2 +- 5 files changed, 210 insertions(+), 123 deletions(-) create mode 100644 third_party/xla/third_party/tensor_ir/patches/pow_op.patch diff --git a/third_party/xla/third_party/tensor_ir/patches/pow_op.patch b/third_party/xla/third_party/tensor_ir/patches/pow_op.patch new file mode 100644 index 00000000000000..d202fe1aa92cc3 --- /dev/null +++ b/third_party/xla/third_party/tensor_ir/patches/pow_op.patch @@ -0,0 +1,36 @@ +diff --git a/lib/Conversion/TensorToCudaTile/OutlinedComputeOps.cpp b/lib/Conversion/TensorToCudaTile/OutlinedComputeOps.cpp +--- a/lib/Conversion/TensorToCudaTile/OutlinedComputeOps.cpp ++++ b/lib/Conversion/TensorToCudaTile/OutlinedComputeOps.cpp +@@ -293,7 +293,7 @@ + getSignedness(exponentType), + cuda_tile::RoundingMode::NEAREST_EVEN); + } +- rewriter.replaceOpWithNewOp(op, operands[0], rhs); ++ rewriter.replaceOpWithNewOp(op, operands[0], rhs); + return success(); + }) + .Case([&](auto) { +diff --git a/test/Conversion/TensorToCudaTile/iterspace_matmul.mlir b/test/Conversion/TensorToCudaTile/iterspace_matmul.mlir +--- a/test/Conversion/TensorToCudaTile/iterspace_matmul.mlir ++++ b/test/Conversion/TensorToCudaTile/iterspace_matmul.mlir +@@ -267,7 +267,7 @@ + // CHECK: %[[RHS2:.*]] = reshape %[[RHS]] : [[TILE_RHS]] -> tile<2x128x16xf32> + // CHECK: %[[MMA:.*]] = mmaf %[[LHS2]], %[[RHS2]], %[[ACCUM]] + // CHECK: %[[RESHAPE:.*]] = reshape %[[MMA]] : tile<2x64x16xf32> -> [[TILE_RS:tile<1x1x2x8x8x4x4xf32>]] +-// CHECK: %[[TRANS:.*]] = permute %[[RESHAPE]] [2, 3, 0, 4, 5, 1, 6] : [[TILE_RS]] -> [[TILE_TR:tile<2x8x1x8x4x1x4xf32>]] ++// CHECK: %[[TRANS:.*]] = permute %[[RESHAPE]] [2, 3, {{[01]}}, 4, 5, {{[01]}}, 6] : [[TILE_RS]] -> [[TILE_TR:tile<2x8x1x8x4x1x4xf32>]] + // CHECK: %[[RESULT:.*]] = broadcast %[[TRANS]] : [[TILE_TR]] -> tile<2x8x1x8x4x2x4xf32> + // CHECK: %[[OUT_VIEW:.*]] = make_tensor_view %[[OUT]], shape = [8, 8, 8, 8, 4, 8, 4], strides = [65536, 8192, 1024, 128, 32, 4, 1] + // CHECK: %[[OUT_PVIEW:.*]] = make_partition_view %[[OUT_VIEW]] +diff --git a/test/Conversion/TensorToCudaTile/reduction_padding.mlir b/test/Conversion/TensorToCudaTile/reduction_padding.mlir +--- a/test/Conversion/TensorToCudaTile/reduction_padding.mlir ++++ b/test/Conversion/TensorToCudaTile/reduction_padding.mlir +@@ -238,7 +238,7 @@ + // CHECK: %[[BLOCK:.*]], %{{.*}}, %{{.*}} = get_tile_block_id : tile + // CHECK: %[[LOOP1:.*]] = for %[[IVAR1:.*]] in (%[[ZERO]] to %[[C2]], step %[[ONE]]) + // CHECK-SAME: iter_values(%[[IARG1:.*]] = %[[ACCUM]]) -> ([[TILE]]) +-// CHECK: %[[A_MUL:.*]] = muli %[[IVAR]], %[[C16]] : tile ++// CHECK: %[[A_MUL:.*]] = muli %[[IVAR1]], %[[C16]] : tile + // CHECK: %[[A_MUL1:.*]] = reshape %[[A_MUL]] : tile -> tile<1xi32> + // CHECK: %[[A_MUL2:.*]] = broadcast %[[A_MUL1]] : tile<1xi32> -> tile<16xi32> + // CHECK: %[[A_IOTA:.*]] = iota : tile<16xi32> diff --git a/third_party/xla/third_party/tensor_ir/patches/series b/third_party/xla/third_party/tensor_ir/patches/series index 9b81d4c8e0dfda..2e39ebeb534e1d 100644 --- a/third_party/xla/third_party/tensor_ir/patches/series +++ b/third_party/xla/third_party/tensor_ir/patches/series @@ -1,2 +1,3 @@ unused_variable.patch symbol_op_interface.patch +pow_op.patch diff --git a/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD b/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD index 3d591efd26df72..1a97f4faac5f5d 100644 --- a/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD +++ b/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD @@ -1,6 +1,22 @@ load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library") +load("@rules_cc//cc:cc_binary.bzl", "cc_binary") load("@rules_cc//cc:cc_library.bzl", "cc_library") +cc_binary( + name = "tensor_ir-tblgen", + srcs = glob([ + "tools/tensor_ir-tblgen/*.cpp", + "tools/tensor_ir-tblgen/*.h", + ]), + deps = [ + "@llvm-project//llvm:Support", + "@llvm-project//llvm:TableGen", + "@llvm-project//mlir:MlirTableGenMain", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TableGen", + ], +) + exports_files(["LICENSE"]) td_library( @@ -14,10 +30,105 @@ td_library( "@llvm-project//mlir:GPUOpsTdFiles", "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", "@llvm-project//mlir:OpBaseTdFiles", + "@llvm-project//mlir:PtrTdFiles", "@llvm-project//mlir:SideEffectInterfacesTdFiles", ], ) +gentbl_cc_library( + name = "TensorIRTypesIncGen", + tbl_outs = [ + ( + [ + "-gen-typedef-decls", + "-typedefs-dialect=nv_tensor_ir", + ], + "include/tensor_ir/Dialect/TensorTypes.h.inc", + ), + ( + [ + "-gen-typedef-defs", + "-typedefs-dialect=nv_tensor_ir", + ], + "include/tensor_ir/Dialect/TensorTypes.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Dialect/TensorTypes.td", + deps = [":TensorIRTdFiles"], +) + +td_library( + name = "TensorIROptionsTdFiles", + srcs = [ + "include/tensor_ir/Options/Enums.td", + "include/tensor_ir/Options/Options.td", + "include/tensor_ir/Options/OptionsBase.td", + ], + includes = ["include"], + deps = [ + ":TensorIRTdFiles", + "@llvm-project//mlir:OpBaseTdFiles", + ], +) + +gentbl_cc_library( + name = "TensorIROptionsEnumsIncGen", + tbl_outs = [ + ( + ["-gen-enum-decls"], + "include/tensor_ir/Options/OptionsEnums.h.inc", + ), + ( + ["-gen-enum-defs"], + "include/tensor_ir/Options/OptionsEnums.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Options/Enums.td", + deps = [":TensorIROptionsTdFiles"], +) + +gentbl_cc_library( + name = "TensorIRCompilerOptionsIncGen", + tbl_outs = [ + ( + ["-gen-component-options"], + "include/tensor_ir/Options/CompilerOptions.h.inc", + ), + ], + tblgen = ":tensor_ir-tblgen", + td_file = "include/tensor_ir/Options/Options.td", + deps = [":TensorIROptionsTdFiles"], +) + +cc_library( + name = "NVTensorIRCompilerOptions", + srcs = [ + "lib/Options/BytecodeVersionOptions.cpp", + "lib/Options/CompilerInvocation.cpp", + "lib/Options/Options.cpp", + ], + hdrs = [ + "include/tensor_ir/Options/BytecodeVersionOptions.h", + "include/tensor_ir/Options/Options.h", + "include/tensor_ir/Options/OptionsDetail.h", + "include/tensor_ir/Options/OptionsEnums.h", + ], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRDialect", + ":NVTensorIRSupport", + ":NVTensorIRUtils", + ":TensorIRCompilerOptionsIncGen", + ":TensorIROptionsEnumsIncGen", + "@cuda_tile//:CudaTileBytecode", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Support", + ], +) + gentbl_cc_library( name = "TensorIRDialectIncGen", tbl_outs = [ @@ -204,19 +315,6 @@ cc_library( ], ) -cc_library( - name = "NVTensorIRCudaApi", - srcs = ["lib/Support/CudaApi.cpp"], - hdrs = ["include/tensor_ir/Support/CudaApi.h"], - includes = ["include"], - visibility = ["//visibility:public"], - deps = [ - ":NVTensorIRSupport", - "@llvm-project//llvm:Support", - "@local_config_cuda//cuda:cuda_headers", - ], -) - cc_library( name = "NVTensorIRDialect", srcs = [ @@ -224,6 +322,7 @@ cc_library( "lib/Dialect/TensorAttrs.cpp", "lib/Dialect/TensorDialect.cpp", "lib/Dialect/TensorOps.cpp", + "lib/Dialect/TensorTypes.cpp", ], hdrs = [ "include/tensor_ir/Dialect/TensorIR.h", @@ -240,14 +339,18 @@ cc_library( ":TensorIROpInterfacesIncGen", ":TensorIROpsCanonicalizationIncGen", ":TensorIROpsIncGen", + ":TensorIRTypesIncGen", "@llvm-project//llvm:Support", "@llvm-project//mlir:BytecodeOpInterface", "@llvm-project//mlir:ControlFlowInterfaces", + "@llvm-project//mlir:DialectUtils", "@llvm-project//mlir:FunctionInterfaces", "@llvm-project//mlir:GPUDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:InferTypeOpInterface", + "@llvm-project//mlir:PtrDialect", "@llvm-project//mlir:SideEffectInterfaces", + "@llvm-project//mlir:ViewLikeInterface", ], ) @@ -289,11 +392,14 @@ cc_library( includes = ["include"], visibility = ["//visibility:public"], deps = [ + ":NVTensorIRCompilerOptions", + ":NVTensorIRConversionOptions", ":NVTensorIRDialect", - ":NVTensorIRRuntime", + ":NVTensorIRKernelArgLayout", ":NVTensorIRUtils", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", ], ) @@ -305,176 +411,120 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":NVTensorIRAnalysis", + ":NVTensorIRCompilerOptions", ":NVTensorIRDialect", ":NVTensorIRSupport", ":NVTensorIRUtils", ":TensorIRTransformPassesIncGen", "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:DialectUtils", "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:GPUDialect", "@llvm-project//mlir:IR", + "@llvm-project//mlir:MemRefDialect", "@llvm-project//mlir:Pass", + "@llvm-project//mlir:PtrDialect", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:SideEffectInterfaces", "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", ], ) cc_library( - name = "NVTensorIRToCudaTileConversion", - srcs = glob(["lib/Conversion/TensorToCudaTile/*.cpp"]), - hdrs = [ - "include/tensor_ir/Conversion/TensorToCudaTile/Options.h", - "include/tensor_ir/Conversion/TensorToCudaTile/TensorToCudaTile.h", - "include/tensor_ir/Conversion/TensorToCudaTile/TensorToCudaTileInternal.h", - ], + name = "NVTensorIRConversionOptions", + hdrs = ["include/tensor_ir/Conversion/TensorToCudaTile/Options.h"], includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRAnalysis", - ":NVTensorIRDialect", - ":NVTensorIRSupport", - ":NVTensorIRUtils", - ":TensorToCudaTileConversionPassIncGen", - "@cuda_tile//:CudaTileDialect", + ":NVTensorIRCompilerOptions", "@llvm-project//llvm:Support", - "@llvm-project//mlir:DialectUtils", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Pass", - "@llvm-project//mlir:SideEffectInterfaces", - "@llvm-project//mlir:TransformUtils", ], ) cc_library( - name = "NVTensorIRCudaTilePipelines", - srcs = ["lib/Compiler/CudaTile/Pipelines.cpp"], - hdrs = ["include/tensor_ir/Compiler/CudaTile/Pipelines.h"], + name = "NVTensorIRKernelArgLayout", + hdrs = ["include/tensor_ir/Runtime/CudaTile/KernelArgLayout.h"], includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRDialect", - ":NVTensorIRToCudaTileConversion", - ":NVTensorIRTransform", - "@llvm-project//mlir:Pass", - "@llvm-project//mlir:Transforms", + ":NVTensorIRConversionOptions", + "@llvm-project//llvm:Support", ], ) cc_library( - name = "NVTensorIRRuntime", - srcs = ["lib/Runtime/CudaTileRuntimeKernel.cpp"], - hdrs = [ - "include/tensor_ir/Runtime/CudaTile/CudaTileRuntimeKernel.h", - "include/tensor_ir/Runtime/CudaTile/KernelArgLayout.h", - "include/tensor_ir/Runtime/CudaTile/KernelLaunchHelpers.h", - "include/tensor_ir/Runtime/CudaTile/RuntimeOperandAccessor.h", - "include/tensor_ir/Runtime/IRuntimeKernel.h", - "include/tensor_ir/Runtime/Types.h", - ], + name = "NVTensorIRRuntimeTypes", + hdrs = ["include/tensor_ir/Runtime/Types.h"], includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRCudaApi", ":NVTensorIRSupport", - "@cuda_tile//:CudaTileBytecode", - "@cuda_tile//:CudaTileDialect", - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Support", - "@local_config_cuda//cuda:cuda_headers", ], ) cc_library( - name = "NVTensorIRReference", - srcs = glob(["lib/Reference/*.cpp"]), + name = "NVTensorIRToCudaTileConversion", + srcs = glob(["lib/Conversion/TensorToCudaTile/*.cpp"]), hdrs = [ - "include/tensor_ir/Reference/reference_graph.h", - "include/tensor_ir/Reference/reference_node.h", - "include/tensor_ir/Reference/simplified_tensor.h", - "include/tensor_ir/Reference/tensor_memory.h", - "lib/Reference/constant_utils.h", - ], - includes = [ - "include", - "lib/Reference", + "include/tensor_ir/Conversion/TensorToCudaTile/TensorToCudaTile.h", + "include/tensor_ir/Conversion/TensorToCudaTile/TensorToCudaTileInternal.h", ], + includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRCudaApi", + ":NVTensorIRAnalysis", + ":NVTensorIRCompilerOptions", + ":NVTensorIRConversionOptions", ":NVTensorIRDialect", - ":NVTensorIRRuntime", ":NVTensorIRSupport", ":NVTensorIRUtils", + ":TensorToCudaTileConversionPassIncGen", + "@cuda_tile//:CudaTileDialect", "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Support", - "@local_config_cuda//cuda:cuda_headers", - ], -) - -cc_library( - name = "NVTensorIRRegistration", - srcs = ["lib/Registration/Registration.cpp"], - hdrs = ["include/tensor_ir/Registration/Registration.h"], - includes = ["include"], - visibility = ["//visibility:public"], - deps = [ - ":NVTensorIRDialect", "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:DialectUtils", "@llvm-project//mlir:FuncDialect", - "@llvm-project//mlir:FuncExtensions", + "@llvm-project//mlir:GPUDialect", "@llvm-project//mlir:IR", + "@llvm-project//mlir:MemRefDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:PtrDialect", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:SideEffectInterfaces", + "@llvm-project//mlir:TransformUtils", ], ) cc_library( - name = "NVTensorIRCompiler", - srcs = [ - "lib/Compiler/Compiler.cpp", - "lib/Compiler/CudaTile/CudaTileCompiler.cpp", - "lib/Compiler/CudaTile/CudaTileFrontend.cpp", - "lib/Compiler/CudaTile/TileIRAssembly.cpp", - ], - hdrs = [ - "include/tensor_ir/Compiler/CompileOptions.h", - "include/tensor_ir/Compiler/Compiler.h", - "include/tensor_ir/Compiler/CudaTile/CudaTileCompiler.h", - "include/tensor_ir/Compiler/CudaTile/CudaTileFrontend.h", - "include/tensor_ir/Compiler/CudaTile/TileIRAssembly.h", - ], + name = "NVTensorIRCudaTilePipelines", + srcs = ["lib/Compiler/CudaTile/Pipelines.cpp"], + hdrs = ["include/tensor_ir/Compiler/CudaTile/Pipelines.h"], includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRAnalysis", - ":NVTensorIRCudaTilePipelines", + ":NVTensorIRConversionOptions", ":NVTensorIRDialect", - ":NVTensorIRRegistration", - ":NVTensorIRRuntime", - ":NVTensorIRSupport", ":NVTensorIRToCudaTileConversion", - ":NVTensorIRUtils", - "@cuda_tile//:CudaTileBytecode", - "@cuda_tile//:CudaTileDialect", - "@llvm-project//llvm:Support", - "@llvm-project//mlir:ArithDialect", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Parser", + ":NVTensorIRTransform", + "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:Pass", - "@llvm-project//mlir:Support", + "@llvm-project//mlir:Transforms", ], ) cc_library( - name = "NVTensorIRCAPI", - srcs = ["lib/CAPI/TensorIR.cpp"], - hdrs = ["include/tensor_ir-c/TensorIR.h"], + name = "NVTensorIRTileIRAssembly", + srcs = ["lib/Compiler/CudaTile/TileIRAssembly.cpp"], + hdrs = ["include/tensor_ir/Compiler/CudaTile/TileIRAssembly.h"], includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRCompiler", - ":NVTensorIRDialect", - ":NVTensorIRRegistration", - ":NVTensorIRRuntime", + ":NVTensorIRSupport", + ":NVTensorIRUtils", + "@cuda_tile//:CudaTileBytecode", "@llvm-project//llvm:Support", - "@llvm-project//mlir:CAPIIRHeaders", ], ) diff --git a/third_party/xla/third_party/tensor_ir/workspace.bzl b/third_party/xla/third_party/tensor_ir/workspace.bzl index 4ed33e9d8047f6..f5bc5c190d1b80 100644 --- a/third_party/xla/third_party/tensor_ir/workspace.bzl +++ b/third_party/xla/third_party/tensor_ir/workspace.bzl @@ -5,8 +5,8 @@ load("//third_party:repo.bzl", "tf_mirror_urls") def repo(): """Imports Tensor IR.""" - TENSOR_IR_COMMIT = "63692d79629e6f32a1d8757695590a59e0adbafd" - TENSOR_IR_SHA256 = "b80794d7c2bfb1bc1ca432d892977becba8d35ec6c18c586acbd648ccc8074dd" + TENSOR_IR_COMMIT = "664d385996371db2a8de67b76f82fc89e2470109" + TENSOR_IR_SHA256 = "846a7bb37ef6a67b0224bd9cb25583ad0d84370a9d2565e3c5a8223c0dcf91f0" quilt_http_archive( name = "tensor_ir", diff --git a/third_party/xla/xla/backends/gpu/codegen/tools/BUILD b/third_party/xla/xla/backends/gpu/codegen/tools/BUILD index b14db5fafbf8f4..51e8db082d5405 100644 --- a/third_party/xla/xla/backends/gpu/codegen/tools/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/tools/BUILD @@ -108,9 +108,9 @@ xla_cc_binary( "@llvm-project//mlir:Pass", "@llvm-project//mlir:Support", "@llvm-project//mlir:TranslateLib", + "@tensor_ir//:NVTensorIRConversionOptions", "@tensor_ir//:NVTensorIRCudaTilePipelines", "@tensor_ir//:NVTensorIRDialect", - "@tensor_ir//:NVTensorIRToCudaTileConversion", ], ) From d2a201ac0e35720e9624d0c6bebdb16f8af16e97 Mon Sep 17 00:00:00 2001 From: Levon Ter-Grigoryan Date: Thu, 24 Sep 2026 03:18:06 -0700 Subject: [PATCH 19/48] [XLA:GPU] Include device failure details in PlatformUtil::GetStreamExecutors when no supported devices are found. This simplifies debugging for example when VMM API is not available. See: https://github.com/openxla/xla/issues/49252 PiperOrigin-RevId: 987377773 --- third_party/xla/xla/service/BUILD | 1 + third_party/xla/xla/service/platform_util.cc | 8 ++++-- .../xla/xla/service/platform_util_test.cc | 25 +++++++++++++++++++ 3 files changed, 32 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index bfb0cec85ccb07..11bed11424bc76 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -1050,6 +1050,7 @@ xla_cc_test( srcs = ["platform_util_test.cc"], deps = [ ":platform_util", + "//xla/stream_executor:mock_platform", "//xla/stream_executor/cuda:cuda_platform_id", "//xla/stream_executor/host:host_platform_id", "//xla/stream_executor/rocm:rocm_platform_id", diff --git a/third_party/xla/xla/service/platform_util.cc b/third_party/xla/xla/service/platform_util.cc index 19bed4a689a89b..a1b1866ae08d87 100644 --- a/third_party/xla/xla/service/platform_util.cc +++ b/third_party/xla/xla/service/platform_util.cc @@ -29,6 +29,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/ascii.h" #include "absl/strings/match.h" +#include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/strings/string_view.h" #include "xla/debug_options_flags.h" @@ -271,6 +272,7 @@ PlatformUtil::GetStreamExecutors( } std::vector out; + std::vector device_errors; for (int i = 0; i < executors.size(); ++i) { absl::StatusOr& se = executors[i]; if (se.ok()) { @@ -279,11 +281,13 @@ PlatformUtil::GetStreamExecutors( LOG(ERROR) << "Failed to create stream executor for device " << platform->Name() << ":" << device_ordinals[i] << ": " << se.status().message(); + device_errors.push_back(absl::StrCat("Device ", device_ordinals[i], ": ", + se.status().message())); } } if (out.empty()) { - return Internal("no supported devices found for platform %s", - platform->Name()); + return Internal("no supported devices found for platform %s: %s", + platform->Name(), absl::StrJoin(device_errors, ", ")); } return out; } diff --git a/third_party/xla/xla/service/platform_util_test.cc b/third_party/xla/xla/service/platform_util_test.cc index 78bb47c211afc6..b0ef13178f18d0 100644 --- a/third_party/xla/xla/service/platform_util_test.cc +++ b/third_party/xla/xla/service/platform_util_test.cc @@ -15,12 +15,15 @@ limitations under the License. #include "xla/service/platform_util.h" +#include + #include #include #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "xla/stream_executor/cuda/cuda_platform_id.h" #include "xla/stream_executor/host/host_platform_id.h" +#include "xla/stream_executor/mock_platform.h" #include "xla/stream_executor/rocm/rocm_platform_id.h" #include "xla/stream_executor/sycl/sycl_platform_id.h" @@ -29,6 +32,10 @@ namespace { using ::absl_testing::IsOkAndHolds; using ::absl_testing::StatusIs; +using ::testing::HasSubstr; +using ::testing::NiceMock; +using ::testing::Return; +using ::testing::ReturnRefOfCopy; TEST(PlatformUtilTest, GetPlatformIdFromCanonicalName) { EXPECT_THAT(PlatformUtil::GetPlatformIdFromCanonicalName("host"), @@ -55,5 +62,23 @@ TEST(PlatformUtilTest, GetPlatformIdFromCanonicalName) { StatusIs(absl::StatusCode::kInvalidArgument)); } +TEST(PlatformUtilTest, GetStreamExecutorsReportsAllDeviceErrors) { + NiceMock platform; + ON_CALL(platform, id()) + .WillByDefault(Return(stream_executor::cuda::kCudaPlatformId)); + ON_CALL(platform, Name()).WillByDefault(ReturnRefOfCopy(std::string("CUDA"))); + ON_CALL(platform, VisibleDeviceCount()).WillByDefault(Return(2)); + ON_CALL(platform, ExecutorForDevice(0)) + .WillByDefault(Return(absl::InternalError("out of memory"))); + ON_CALL(platform, ExecutorForDevice(1)) + .WillByDefault(Return(absl::UnavailableError("device busy"))); + + EXPECT_THAT( + PlatformUtil::GetStreamExecutors(&platform), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("no supported devices found for platform CUDA: " + "Device 0: out of memory, Device 1: device busy"))); +} + } // namespace } // namespace xla From 6dc9dc2a5f53151c9e093f8dab99ab59b409d2ea Mon Sep 17 00:00:00 2001 From: linchen1-robot Date: Thu, 24 Sep 2026 03:18:50 -0700 Subject: [PATCH 20/48] PR #49033: [ROCm]Add ROCm HSACO compile-then-pack coverage and fix clang-offload-bundler target IDs. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/49033 📝 Summary of Changes - Add rocm-only //xla/stream_executor/gpu:gpu_hsaco_bundle_test: compile inlined AMDGPU LLVM IR with amdgpu::CompileToHsaco, then pack the HSACO with BundleGpuAsm(HsacoImage). - Fix ROCm 7 clang-offload-bundler --targets= IDs in gpu/asm_compiler.cc (4-field host and HIP triples). - Does not tag-drop //xla/stream_executor/cuda:subprocess_compilation_no_fakes_test (cuda-only). This is extra coverage, not a name-matched twin. 🎯 Justification Add more unit test coverage on ROCm platform. 🚀 Kind of Contribution 🧪 Tests Copybara import of the project: -- 981e93ac52fbeef919ccda3fdaf17cd7a2c6dd8f by Lin Chen1 : Add ROCm HSACO compile-then-pack coverage and fix clang-offload-bundler target IDs. Signed-off-by: Lin Chen1 Merging this change closes #49033 PiperOrigin-RevId: 987378077 --- third_party/xla/xla/stream_executor/gpu/BUILD | 25 +++++ .../xla/stream_executor/gpu/asm_compiler.cc | 10 +- .../gpu/gpu_hsaco_bundle_test.cc | 106 ++++++++++++++++++ 3 files changed, 137 insertions(+), 4 deletions(-) create mode 100644 third_party/xla/xla/stream_executor/gpu/gpu_hsaco_bundle_test.cc diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD index 463ae18dbb36be..6a4ee1dfccabe7 100644 --- a/third_party/xla/xla/stream_executor/gpu/BUILD +++ b/third_party/xla/xla/stream_executor/gpu/BUILD @@ -350,6 +350,31 @@ cc_library( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", + "@tsl//tsl/platform:path", + ], +) + +xla_test( + name = "gpu_hsaco_bundle_test", + srcs = ["gpu_hsaco_bundle_test.cc"], + backends = ["gpu"], + tags = ["rocm-only"], + deps = [ + ":asm_compiler", + "//xla:debug_options_flags", + "//xla/service/gpu/llvm_gpu_backend:amdgpu_backend", + "//xla/stream_executor:platform", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor/rocm:rocm_platform", # buildcleaner: keep + "//xla/tsl/platform:env", + "//xla/tsl/platform:rocm_rocdl_path", + "//xla/tsl/platform:statusor", + "@com_google_googletest//:gtest_main", # buildcleaner: keep + "@llvm-project//llvm:AsmParser", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:ir_headers", "@tsl//tsl/platform:path", ], ) diff --git a/third_party/xla/xla/stream_executor/gpu/asm_compiler.cc b/third_party/xla/xla/stream_executor/gpu/asm_compiler.cc index 70b955d9bbdd01..79a0cecf616723 100644 --- a/third_party/xla/xla/stream_executor/gpu/asm_compiler.cc +++ b/third_party/xla/xla/stream_executor/gpu/asm_compiler.cc @@ -58,13 +58,14 @@ absl::StatusOr> BundleGpuAsm( findRocmExecutable("llvm/bin/clang-offload-bundler", rocm_root_dir); // Initialise the "--inputs" / "--targets" arguments for the - // clang-offload-bundler with a dummy file / host target triple... - // clang-offload-bundler requires 1 and only 1 host target triple + // clang-offload-bundler with a dummy file / host target ID. + // clang-offload-bundler requires exactly one host target ID. + // Each ID is '----[-]'. std::ostringstream inputs_list; std::ostringstream targets_list; inputs_list << "/dev/null"; - targets_list << "host-x86_64-unknown-linux"; + targets_list << "host-x86_64-unknown-linux-gnu"; // Write images to temporary files. std::vector image_paths; @@ -79,7 +80,8 @@ absl::StatusOr> BundleGpuAsm( env, img_path, std::string(img.bytes.begin(), img.bytes.end()))); VLOG(2) << "image written to " << img_path; inputs_list << "," << img_path; - targets_list << ",hip-amdgcn-amd-amdhsa-" << img.gfx_arch; + // Empty is the extra '-' before gfx (hip-amdgcn-amd-amdhsa--gfx...). + targets_list << ",hip-amdgcn-amd-amdhsa--" << img.gfx_arch; image_paths.push_back(std::move(img_path)); } absl::Cleanup image_files_cleaner = [&image_paths] { diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_hsaco_bundle_test.cc b/third_party/xla/xla/stream_executor/gpu/gpu_hsaco_bundle_test.cc new file mode 100644 index 00000000000000..cb5804eda665eb --- /dev/null +++ b/third_party/xla/xla/stream_executor/gpu/gpu_hsaco_bundle_test.cc @@ -0,0 +1,106 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include + +#include +#include +#include "llvm/AsmParser/Parser.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/Support/SourceMgr.h" +#include "xla/debug_options_flags.h" +#include "xla/service/gpu/llvm_gpu_backend/amdgpu_backend.h" +#include "xla/stream_executor/gpu/asm_compiler.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" +#include "xla/stream_executor/stream_executor.h" +#include "xla/tsl/platform/env.h" +#include "xla/tsl/platform/rocm_rocdl_path.h" +#include "xla/tsl/platform/statusor.h" +#include "tsl/platform/path.h" + +namespace stream_executor { +namespace { + +TEST(GpuHsacoBundleTest, CompileToHsacoThenBundleGpuAsmWorks) { + const std::string rocm_root = tsl::RocmRoot(); + if (rocm_root.empty()) { + GTEST_SKIP() << "ROCm root is empty; clang-offload-bundler is unavailable"; + } + const std::string bundler = + tsl::io::JoinPath(rocm_root, "llvm/bin/clang-offload-bundler"); + if (!tsl::Env::Default()->FileExists(bundler).ok()) { + GTEST_SKIP() << "clang-offload-bundler not found at " << bundler; + } + + ASSERT_OK_AND_ASSIGN(Platform * platform, + PlatformManager::PlatformWithName("ROCM")); + ASSERT_GT(platform->VisibleDeviceCount(), 0); + ASSERT_OK_AND_ASSIGN(StreamExecutor * executor, + platform->ExecutorForDevice(0)); + const std::string gfx_arch = + executor->GetDeviceDescription().rocm_compute_capability().gfx_version(); + ASSERT_FALSE(gfx_arch.empty()); + + // CompileToHsaco takes an LLVM module; this kernel has no host or ROCDL deps. + constexpr char kIr[] = R"( +target triple = "amdgcn-amd-amdhsa" +define amdgpu_kernel void @simple_add(ptr addrspace(1) %a, ptr addrspace(1) %b, + ptr addrspace(1) %c) { +entry: + %tid = call i32 @llvm.amdgcn.workitem.id.x() + %tidx = zext i32 %tid to i64 + %a_ptr = getelementptr float, ptr addrspace(1) %a, i64 %tidx + %b_ptr = getelementptr float, ptr addrspace(1) %b, i64 %tidx + %c_ptr = getelementptr float, ptr addrspace(1) %c, i64 %tidx + %a_val = load float, ptr addrspace(1) %a_ptr, align 4 + %b_val = load float, ptr addrspace(1) %b_ptr, align 4 + %sum = fadd float %a_val, %b_val + store float %sum, ptr addrspace(1) %c_ptr, align 4 + ret void +} +declare i32 @llvm.amdgcn.workitem.id.x() +)"; + + llvm::LLVMContext llvm_context; + llvm::SMDiagnostic parse_err; + std::unique_ptr module = + llvm::parseAssemblyString(kIr, parse_err, llvm_context); + ASSERT_NE(module, nullptr) << parse_err.getMessage().str(); + // EmitModuleToHsaco joins the module id into a temp path. + module->setModuleIdentifier("gpu_hsaco_bundle_test"); + + ASSERT_OK_AND_ASSIGN( + xla::gpu::amdgpu::HsacoResult hsaco_result, + xla::gpu::amdgpu::CompileToHsaco( + module.get(), + executor->GetDeviceDescription().gpu_compute_capability(), + xla::GetDebugOptionsFromFlags(), "gpu_hsaco_bundle_test")); + ASSERT_FALSE(hsaco_result.hsaco.empty()); + + HsacoImage image; + image.gfx_arch = gfx_arch; + image.bytes = std::move(hsaco_result.hsaco); + ASSERT_OK_AND_ASSIGN(std::vector bundle, + BundleGpuAsm({image}, rocm_root)); + ASSERT_FALSE(bundle.empty()); +} + +} // namespace +} // namespace stream_executor From 435c4e858e22a96d4f476c629a48f7e17d80e651 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 24 Sep 2026 03:24:39 -0700 Subject: [PATCH 21/48] Automated Code Change PiperOrigin-RevId: 987380727 --- tensorflow/lite/examples/label_image/label_image_test.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorflow/lite/examples/label_image/label_image_test.cc b/tensorflow/lite/examples/label_image/label_image_test.cc index b42769dc8eaefe..65a2e3a40da7d0 100644 --- a/tensorflow/lite/examples/label_image/label_image_test.cc +++ b/tensorflow/lite/examples/label_image/label_image_test.cc @@ -16,8 +16,10 @@ limitations under the License. #include "tensorflow/lite/examples/label_image/label_image.h" #include +#include #include #include +#include #include #include #include From 57dbb888acf77611c4af628fa8dee2c89d90a90e Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Thu, 24 Sep 2026 04:19:32 -0700 Subject: [PATCH 22/48] [XLA] Use root `native_custom_call_handler_allowlist` in non-GPU XLA target visibility Replace `"//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist"` in the `visibility` lists of core non-GPU XLA targets (`hlo`, `compiler`, `buffer_assignment`, `kernel_arguments`, `launch_dim`, `kernel_spec`) with `"//xla:native_custom_call_handler_allowlist"` and include `native_custom_call_thunks/...` in `native_custom_call_handler_allowlist` so non-GPU XLA packages do not reference a `package_group` inside `xla/backends/gpu/`. PiperOrigin-RevId: 987404559 --- third_party/xla/xla/codegen/emitters/BUILD | 2 +- third_party/xla/xla/hlo/ir/BUILD | 2 +- third_party/xla/xla/service/BUILD | 4 ++-- third_party/xla/xla/stream_executor/BUILD | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/third_party/xla/xla/codegen/emitters/BUILD b/third_party/xla/xla/codegen/emitters/BUILD index b3e3a248ebc56e..0e987b5d9a85fc 100644 --- a/third_party/xla/xla/codegen/emitters/BUILD +++ b/third_party/xla/xla/codegen/emitters/BUILD @@ -224,7 +224,7 @@ cc_library( hdrs = ["kernel_arguments.h"], visibility = [ ":friends", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", ], deps = [ "//xla:frontend_attributes", diff --git a/third_party/xla/xla/hlo/ir/BUILD b/third_party/xla/xla/hlo/ir/BUILD index 9e44d36b13c01c..a25fbd8777f887 100644 --- a/third_party/xla/xla/hlo/ir/BUILD +++ b/third_party/xla/xla/hlo/ir/BUILD @@ -63,7 +63,7 @@ cc_library( ], visibility = internal_visibility([ ":friends", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", ]), deps = [ ":backend_config", diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 11bed11424bc76..395738a19efe68 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -1632,7 +1632,7 @@ cc_library( visibility = internal_visibility([ "//xla/internal:hwi_internal", "//xla/backends/gpu/ffi:__pkg__", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", "//xla/backends/gpu/runtime:__subpackages__", "//xla/backends/gpu/tests:__pkg__", "//xla/backends/gpu/transforms:__subpackages__", @@ -1891,7 +1891,7 @@ cc_library( ], visibility = internal_visibility([ ":friends", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", ]), deps = [ ":buffer_assignment_proto_cc", diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD index fbfaaba6f928b0..2d863d1a66fc41 100644 --- a/third_party/xla/xla/stream_executor/BUILD +++ b/third_party/xla/xla/stream_executor/BUILD @@ -485,7 +485,7 @@ cc_library( hdrs = ["launch_dim.h"], visibility = internal_visibility([ ":friends", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", ]), deps = [ ":launch_dim_proto_cc", @@ -874,7 +874,7 @@ cc_library( hdrs = ["kernel_spec.h"], visibility = internal_visibility([ ":friends", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", ]), deps = [ ":kernel", From 8c9cfc9f5af15383395941a882f7f0a745ffd76f Mon Sep 17 00:00:00 2001 From: Alexander Belyaev Date: Thu, 24 Sep 2026 04:58:17 -0700 Subject: [PATCH 23/48] [XLA:CPU] Legalize narrow float vector memory operations in CPU xtile. CPU targets without specific ISA extensions (such as AVX512-BF16) lack legal vector register classes for narrow floats (e.g., bf16, fp16). Using these types causes LLVM to scalarize them and generate libcalls. PiperOrigin-RevId: 987420616 --- .../backends/cpu/codegen/fusion_compiler.cc | 3 + .../tiled/tests/elementwise/add_bf16_mask.hlo | 72 ++++++ .../tests/elementwise/convert_f32_bf16.hlo | 2 +- .../tiled/tests/elementwise/is_finite.hlo | 6 +- .../transpose_convert_bf16_f32.hlo | 10 +- .../transpose_convert_f32_bf16.hlo | 10 +- .../transpose/transpose_broadcast_f16.hlo | 40 +++ .../transpose/transpose_convert_bf16.hlo | 50 ++++ .../cpu/codegen/tiled/transforms/BUILD | 3 + .../legalize_narrow_float_storage.cc | 228 ++++++++++++++++++ .../tiled/transforms/lower_memref_bitcast.cc | 78 ++++++ .../cpu/codegen/tiled/transforms/passes.td | 46 ++++ .../tests/legalize_narrow_float_storage.mlir | 153 ++++++++++++ .../tests/lower_memref_bitcast.mlir | 43 ++++ .../xla/xla/codegen/xtile/ir/tests/ops.mlir | 33 +++ .../xla/xla/codegen/xtile/ir/xtile_ops.cc | 44 ++++ .../xla/xla/codegen/xtile/ir/xtile_ops.td | 25 ++ 17 files changed, 835 insertions(+), 11 deletions(-) create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/add_bf16_mask.hlo create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_broadcast_f16.hlo create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_convert_bf16.hlo create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/transforms/legalize_narrow_float_storage.cc create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_memref_bitcast.cc create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/legalize_narrow_float_storage.mlir create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_memref_bitcast.mlir diff --git a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc index b7c1aa55f3304f..fee83871c42993 100644 --- a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc +++ b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc @@ -407,6 +407,7 @@ void AddNewXtileToVectorPasses(mlir::OpPassManager& pm) { pm.addPass(xtile::createExpandXtileComplexOpsPass()); pm.addPass(xtile::createStablehloLowerToArithPass()); pm.addPass(xtile::createLegalizeUnsignedIntegersAsSignlessPass()); + pm.addPass(cpu::createLegalizeNarrowFloatStoragePass()); pm.addPass(mlir::createCanonicalizerPass()); pm.addPass(cpu::createVectorizeXTilePass()); pm.addPass(cpu::createLowerXTileEntryPass()); @@ -484,6 +485,8 @@ void AddNewVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max, options.vectorTransposeLowering = mlir::vector::VectorTransposeLowering::Shuffle16x16; pm.addPass(mlir::createConvertVectorToLLVMPass(options)); + pm.addPass(cpu::createLowerMemRefBitcastPass()); + pm.addPass(cpu::createLowerToLLVMPass()); pm.addPass(mlir::memref::createExpandStridedMetadataPass()); pm.addPass(emitters::createSafeIntegerArithmeticPass()); diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/add_bf16_mask.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/add_bf16_mask.hlo new file mode 100644 index 00000000000000..412198446a0d55 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/add_bf16_mask.hlo @@ -0,0 +1,72 @@ +// RUN: fusion_to_xtile %s | fusion_compiler_opt --xtile-cpu-xtile-to-vector \ +// RUN: --xtile-cpu-vector-to-llvm | FileCheck %s --check-prefix=CHECK-LLVM +// RUN: test_correctness %s + +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector --xtile-cpu-new-vector-to-llvm \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR-LLVM +// RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true + +// The tile does not divide the minor dimension, so the reads are guarded by an +// `scf.if` that picks between an in-bounds read and a masked read. None of the +// data movement may happen on bf16 vectors: x86 has no legal vector register +// class for bf16 without AVX512-BF16, so SelectionDAG would scalarize the merge +// and the masked accesses and emit one `__truncsfbf2` libcall per element. +// `xtile-cpu-legalize-narrow-float-storage` views the buffers as i16, so bf16 +// only shows up right at `arith.extf`/`arith.truncf`. + +fusion { + p0 = bf16[100,100] parameter(0) + p1 = bf16[100,100] parameter(1) + c0 = f32[100,100] convert(p0) + c1 = f32[100,100] convert(p1) + a = f32[100,100] add(c0, c1) + ROOT r = bf16[100,100] convert(a) +} + +// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl +// NEW-VECTOR-SAME: %[[ARG0:.*]]: memref<100x100xbf16>, %[[ARG1:.*]]: memref<100x100xbf16>, %[[ARG2:.*]]: memref<100x100xbf16> + +// NEW-VECTOR-DAG: %[[PAD:.*]] = arith.constant 0 : i16 +// NEW-VECTOR-DAG: %[[OUT:.*]] = xtile.memref_bitcast %[[ARG2]] : memref<100x100xbf16> -> memref<100x100xi16> +// NEW-VECTOR-DAG: %[[IN0:.*]] = xtile.memref_bitcast %[[ARG0]] : memref<100x100xbf16> -> memref<100x100xi16> +// NEW-VECTOR: scf.for +// NEW-VECTOR: %[[COND:.*]] = arith.cmpi sge +// NEW-VECTOR: %[[R0:.*]] = scf.if %[[COND]] -> (vector<1x16xi16>) { +// NEW-VECTOR: vector.transfer_read %[[IN0]]{{.*}}, %[[PAD]] {in_bounds = [true, true]} : memref<100x100xi16>, vector<1x16xi16> +// NEW-VECTOR: } else { +// NEW-VECTOR: %[[MASK:.*]] = vector.create_mask {{.*}} : vector<1x16xi1> +// NEW-VECTOR: vector.transfer_read %[[IN0]]{{.*}}, %[[PAD]], %[[MASK]] : memref<100x100xi16>, vector<1x16xi16> +// NEW-VECTOR: } +// NEW-VECTOR: %[[R0_BF16:.*]] = arith.bitcast %[[R0]] : vector<1x16xi16> to vector<1x16xbf16> +// NEW-VECTOR: arith.extf %[[R0_BF16]] : vector<1x16xbf16> to vector<1x16xf32> +// NEW-VECTOR: %[[TRUNC:.*]] = arith.truncf %{{.*}} : vector<1x16xf32> to vector<1x16xbf16> +// NEW-VECTOR: %[[TRUNC_I16:.*]] = arith.bitcast %[[TRUNC]] : vector<1x16xbf16> to vector<1x16xi16> +// NEW-VECTOR: vector.transfer_write %[[TRUNC_I16]], %[[OUT]]{{.*}} : vector<1x16xi16>, memref<100x100xi16> + +// NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl +// NEW-VECTOR-LLVM-NOT: xtile.memref_bitcast +// NEW-VECTOR-LLVM: ^bb{{[0-9]+}}(%{{.*}}: !llvm.array<1 x vector<16xi16>>) +// NEW-VECTOR-LLVM: ^bb{{[0-9]+}}(%{{.*}}: !llvm.array<1 x vector<16xi16>>) +// NEW-VECTOR-LLVM-NOT: llvm.intr.masked.{{.*}}vector<16xbf16> +// NEW-VECTOR-LLVM: llvm.intr.masked.store{{.*}}vector<16xi16> +// NEW-VECTOR-LLVM-NOT: llvm.intr.masked.{{.*}}vector<16xbf16> + +// CHECK-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl + +ENTRY main { + p0 = bf16[100,100] parameter(0) + p1 = bf16[100,100] parameter(1) + ROOT fusion = bf16[100,100] fusion(p0, p1), kind=kLoop, calls=fusion, + backend_config={ + "fusion_backend_config":{ + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["1", "16"]}], + "num_ctas":"1" + } + } + } +} diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo index 85cd8460d6f1e6..0556729980ea28 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo @@ -16,7 +16,7 @@ fusion { // NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl // NEW-VECTOR-LLVM: %[[LOAD:.*]] = llvm.load %{{.*}} : !llvm.ptr -> vector<16xf32> // NEW-VECTOR-LLVM-COUNT-2: llvm.call @xla.fptrunc.v8f32.to.v8bf16(%{{.*}}) : (vector<8xf32>) -> vector<8xbf16> -// NEW-VECTOR-LLVM: llvm.store %{{.*}}, %{{.*}} : vector<16xbf16>, !llvm.ptr +// NEW-VECTOR-LLVM: llvm.store %{{.*}}, %{{.*}} : vector<16xi16>, !llvm.ptr ENTRY main { p0 = f32[1, 16] parameter(0) diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo index cdb8e7604f65f0..05965d7a9060e5 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo @@ -42,9 +42,11 @@ // NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl // NEW-VECTOR-SAME: %[[ARG:[^:]+]]: memref<1x9xf16>, %[[OUT:[^:]+]]: memref<1x9xi8>, %[[PID:[^:]+]]: index +// NEW-VECTOR: %[[ARG_I16:.*]] = xtile.memref_bitcast %[[ARG]] : memref<1x9xf16> -> memref<1x9xi16> // NEW-VECTOR: %[[MASK:.*]] = vector.constant_mask [1, 9] : vector<1x16xi1> -// NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %[[ARG]]{{.*}}, %[[MASK]]{{.*}} : memref<1x9xf16>, vector<1x16xf16> -// NEW-VECTOR: %[[FINITE:.*]] = math.isfinite %[[READ]] : vector<1x16xf16> +// NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %[[ARG_I16]]{{.*}}, %[[MASK]]{{.*}} : memref<1x9xi16>, vector<1x16xi16> +// NEW-VECTOR: %[[READ_F16:.*]] = arith.bitcast %[[READ]] : vector<1x16xi16> to vector<1x16xf16> +// NEW-VECTOR: %[[FINITE:.*]] = math.isfinite %[[READ_F16]] : vector<1x16xf16> // NEW-VECTOR: %[[CAST:.*]] = arith.extui %[[FINITE]] : vector<1x16xi1> to vector<1x16xi8> // NEW-VECTOR: vector.transfer_write %[[CAST]], %[[OUT]]{{.*}}, %[[MASK]]{{.*}} : vector<1x16xi8>, memref<1x9xi8> diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo index 2ac596a5d7e59e..82c69b890bd617 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo @@ -61,16 +61,18 @@ fusion { // CHECK-LLVM: llvm.alloca // NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl -// NEW-VECTOR: %[[PAD:.*]] = arith.constant 0.000000e+00 : bf16 +// NEW-VECTOR: %[[PAD:.*]] = arith.constant 0 : i16 +// NEW-VECTOR: %[[IN:.*]] = xtile.memref_bitcast %{{.*}} : memref<1x256x56x56xbf16> -> memref<1x256x56x56xi16> // NEW-VECTOR: scf.for -// NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %{{.*}}[%{{.*}}], %[[PAD]] {in_bounds = [true, true, true, true]} : memref<1x256x56x56xbf16>, vector<1x8x1x8xbf16> -// NEW-VECTOR: %[[EXT:.*]] = arith.extf %[[READ]] : vector<1x8x1x8xbf16> to vector<1x8x1x8xf32> +// NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %[[IN]][%{{.*}}], %[[PAD]] {in_bounds = [true, true, true, true]} : memref<1x256x56x56xi16>, vector<1x8x1x8xi16> +// NEW-VECTOR: %[[CAST:.*]] = arith.bitcast %[[READ]] : vector<1x8x1x8xi16> to vector<1x8x1x8xbf16> +// NEW-VECTOR: %[[EXT:.*]] = arith.extf %[[CAST]] : vector<1x8x1x8xbf16> to vector<1x8x1x8xf32> // NEW-VECTOR: %[[TRANS:.*]] = vector.transpose %[[EXT]], [0, 2, 3, 1] : vector<1x8x1x8xf32> to vector<1x1x8x8xf32> // NEW-VECTOR: vector.transfer_write %[[TRANS]], %{{.*}}[%{{.*}}] {in_bounds = [true, true, true, true]} : vector<1x1x8x8xf32>, memref<1x56x56x256xf32> // NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl // NEW-VECTOR-LLVM-NOT: llvm.alloca -// NEW-VECTOR-LLVM-COUNT-8: llvm.load %{{.*}} : !llvm.ptr -> vector<8xbf16> +// NEW-VECTOR-LLVM-COUNT-8: llvm.load %{{.*}} : !llvm.ptr -> vector<8xi16> // NEW-VECTOR-LLVM: llvm.shufflevector // NEW-VECTOR-LLVM-COUNT-8: llvm.store %{{.*}}, %{{.*}} : vector<8xf32>, !llvm.ptr diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo index c28222377502bf..cca9185f4e72bc 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo @@ -65,15 +65,17 @@ fusion { // NEW-VECTOR: scf.for // NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %{{.*}}[%{{.*}}], %[[PAD]] {in_bounds = [true, true, true, true]} : memref<1x56x56x256xf32>, vector<1x1x8x8xf32> // NEW-VECTOR: %[[TRUNC:.*]] = arith.truncf %[[READ]] : vector<1x1x8x8xf32> to vector<1x1x8x8xbf16> -// NEW-VECTOR: %[[TRANS:.*]] = vector.transpose %[[TRUNC]], [0, 3, 1, 2] : vector<1x1x8x8xbf16> to vector<1x8x1x8xbf16> -// NEW-VECTOR: vector.transfer_write %[[TRANS]], %{{.*}}[%{{.*}}] {in_bounds = [true, true, true, true]} : vector<1x8x1x8xbf16>, memref<1x256x56x56xbf16> +// NEW-VECTOR: %[[CAST:.*]] = arith.bitcast %[[TRUNC]] : vector<1x1x8x8xbf16> to vector<1x1x8x8xi16> +// NEW-VECTOR: %[[TRANS:.*]] = vector.transpose %[[CAST]], [0, 3, 1, 2] : vector<1x1x8x8xi16> to vector<1x8x1x8xi16> +// NEW-VECTOR: vector.transfer_write %[[TRANS]], %{{.*}}[%{{.*}}] {in_bounds = [true, true, true, true]} : vector<1x8x1x8xi16>, memref<1x256x56x56xi16> // NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl // NEW-VECTOR-LLVM-NOT: llvm.alloca // NEW-VECTOR-LLVM-COUNT-8: llvm.load %{{.*}} : !llvm.ptr -> vector<8xf32> // NEW-VECTOR-LLVM-COUNT-8: llvm.call @xla.fptrunc.v8f32.to.v8bf16 -// NEW-VECTOR-LLVM: llvm.shufflevector -// NEW-VECTOR-LLVM-COUNT-8: llvm.store %{{.*}}, %{{.*}} : vector<8xbf16>, !llvm.ptr +// NEW-VECTOR-LLVM-COUNT-8: llvm.bitcast %{{.*}} : vector<8xbf16> to vector<8xi16> +// NEW-VECTOR-LLVM: llvm.shufflevector %{{.*}}, %{{.*}} [{{.*}}] : vector<8xi16> +// NEW-VECTOR-LLVM-COUNT-8: llvm.store %{{.*}}, %{{.*}} : vector<8xi16>, !llvm.ptr ENTRY %main { %input = f32[1,56,56,256]{3,2,1,0} parameter(0) diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_broadcast_f16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_broadcast_f16.hlo new file mode 100644 index 00000000000000..e32862353cdc8d --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_broadcast_f16.hlo @@ -0,0 +1,40 @@ +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR +// RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true + +// f16 has the same problem as bf16 on x86 without AVX512-FP16. The broadcast +// and the transpose only move data, so they run on i16; f16 is only visible +// at the compute. + +fusion { + p0 = f16[64,48] parameter(0) + p1 = f16[48] parameter(1) + t = f16[48,64] transpose(p0), dimensions={1,0} + b = f16[48,64] broadcast(p1), dimensions={0} + ROOT a = f16[48,64] add(t, b) +} + +// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl +// NEW-VECTOR-DAG: xtile.memref_bitcast %{{.*}} : memref<64x48xf16> -> memref<64x48xi16> +// NEW-VECTOR-DAG: xtile.memref_bitcast %{{.*}} : memref<48xf16> -> memref<48xi16> +// NEW-VECTOR-DAG: xtile.memref_bitcast %{{.*}} : memref<48x64xf16> -> memref<48x64xi16> +// NEW-VECTOR: vector.transpose %{{.*}} : vector<8x8xi16> to vector<8x8xi16> +// NEW-VECTOR: vector.broadcast %{{.*}} : vector<8x1xi16> to vector<8x8xi16> +// NEW-VECTOR-NOT: vector.transpose {{.*}}f16> +// NEW-VECTOR-NOT: vector.broadcast {{.*}}f16> +// NEW-VECTOR-NOT: vector.transfer_{{.*}}f16> + +ENTRY main { + p0 = f16[64,48] parameter(0) + p1 = f16[48] parameter(1) + ROOT fusion = f16[48,64] fusion(p0, p1), kind=kLoop, calls=fusion, + backend_config={ + "fusion_backend_config":{ + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["8", "8"]}], + "num_ctas":"1" + } + } + } +} diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_convert_bf16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_convert_bf16.hlo new file mode 100644 index 00000000000000..f69f4d8820ff57 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_convert_bf16.hlo @@ -0,0 +1,50 @@ +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector --xtile-cpu-new-vector-to-llvm \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR-LLVM +// RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true + +// The output side of an NHWC -> NCHW layout change in bf16, as emitted around a +// depthwise convolution. The convert happens before the transpose, so without +// help the 16x16 transpose shuffles bf16 vectors, which x86 scalarizes into one +// `__truncsfbf2` libcall per element. The transpose only moves data, so it must +// run on i16 instead. + +fusion { + p0 = f32[1,56,56,256]{3,2,1,0} parameter(0) + c = bf16[1,56,56,256]{3,2,1,0} convert(p0) + t = bf16[1,256,56,56]{1,3,2,0} transpose(c), dimensions={0,3,1,2} + ROOT copy = bf16[1,256,56,56]{3,2,1,0} copy(t) +} + +// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl +// NEW-VECTOR: %[[OUT:.*]] = xtile.memref_bitcast %{{.*}} : memref<1x256x56x56xbf16> -> memref<1x256x56x56xi16> +// NEW-VECTOR: %[[TRUNC:.*]] = arith.truncf %{{.*}} : vector<1x1x16x16xf32> to vector<1x1x16x16xbf16> +// NEW-VECTOR: %[[CAST:.*]] = arith.bitcast %[[TRUNC]] : vector<1x1x16x16xbf16> to vector<1x1x16x16xi16> +// NEW-VECTOR: %[[T:.*]] = vector.transpose %[[CAST]], [0, 3, 1, 2] : vector<1x1x16x16xi16> to vector<1x16x1x16xi16> +// NEW-VECTOR: vector.transfer_write %[[T]], %[[OUT]] + +// The transpose shuffles and all memory accesses are on i16. The only bf16 +// vectors left join the two halves produced by the 8-wide truncation intrinsic +// right before the bitcast, which LLVM keeps in registers. +// NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl +// NEW-VECTOR-LLVM-NOT: llvm.shufflevector {{.*}}[0, 16, 1, 17, {{.*}}bf16> +// NEW-VECTOR-LLVM: llvm.shufflevector {{.*}}[0, 16, 1, 17, {{.*}} : vector<16xi16> +// NEW-VECTOR-LLVM-NOT: llvm.shufflevector {{.*}}[0, 16, 1, 17, {{.*}}bf16> +// NEW-VECTOR-LLVM-NOT: llvm.store {{.*}}bf16> +// NEW-VECTOR-LLVM-NOT: llvm.intr.masked.store {{.*}}bf16> + +ENTRY main { + p0 = f32[1,56,56,256]{3,2,1,0} parameter(0) + ROOT fusion = bf16[1,256,56,56]{3,2,1,0} fusion(p0), kind=kLoop, calls=fusion, + backend_config={ + "fusion_backend_config":{ + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["1", "16", "1", "16"]}], + "num_ctas":"12" + } + } + } +} diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD index 6dc1b55ac96e50..ce7c4c57144e07 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD @@ -68,7 +68,9 @@ cc_library( "fuse_elementwise_pass.cc", "hoist_alloca.cc", "initialize_allocs.cc", + "legalize_narrow_float_storage.cc", "linalg_elementwise_to_vector_pass.cc", + "lower_memref_bitcast.cc", "lower_xtile_entry.cc", "memref_copy_to_loops.cc", "shlo_to_vector.cc", @@ -108,6 +110,7 @@ cc_library( "@llvm-project//mlir:FuncTransforms", "@llvm-project//mlir:FuncUtils", "@llvm-project//mlir:IR", + "@llvm-project//mlir:LLVMCommonConversion", "@llvm-project//mlir:LLVMDialect", "@llvm-project//mlir:LinalgDialect", "@llvm-project//mlir:LinalgTransforms", diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/legalize_narrow_float_storage.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/legalize_narrow_float_storage.cc new file mode 100644 index 00000000000000..12a682cdcc04f8 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/legalize_narrow_float_storage.cc @@ -0,0 +1,228 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/Support/LogicalResult.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" // IWYU pragma: keep +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/IR/Types.h" +#include "mlir/IR/Value.h" +#include "mlir/Pass/Pass.h" // IWYU pragma: keep +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "stablehlo/dialect/StablehloOps.h" +#include "xla/backends/cpu/codegen/tiled/transforms/passes.h" // IWYU pragma: keep +#include "xla/codegen/xtile/ir/xtile_ops.h" + +namespace xla::cpu { + +#define GEN_PASS_DEF_LEGALIZENARROWFLOATSTORAGEPASS +#include "xla/backends/cpu/codegen/tiled/transforms/passes.h.inc" + +namespace { + +namespace ma = ::mlir::arith; +namespace shlo = ::mlir::stablehlo; +namespace tensor = ::mlir::tensor; + +using ::mlir::IntegerType; +using ::mlir::Location; +using ::mlir::LogicalResult; +using ::mlir::MemRefType; +using ::mlir::OpBuilder; +using ::mlir::PatternRewriter; +using ::mlir::RankedTensorType; +using ::mlir::Type; +using ::mlir::Value; + +// Returns the integer type used to carry values of `type` through data +// movement, or nullptr if `type` is not a narrow float. +IntegerType GetStorageType(Type type) { + auto float_type = mlir::dyn_cast(type); + if (float_type == nullptr) { + return {}; + } + unsigned width = float_type.getWidth(); + if (width != 8 && width != 16) { + return {}; + } + return IntegerType::get(type.getContext(), width); +} + +// Same shape as `type`, with the element type replaced by its storage type, or +// nullptr if the element type is not a narrow float. +RankedTensorType GetStorageTensorType(Type type) { + auto tensor_type = mlir::dyn_cast(type); + if (tensor_type == nullptr) { + return {}; + } + IntegerType storage = GetStorageType(tensor_type.getElementType()); + if (storage == nullptr) { + return {}; + } + return tensor_type.clone(storage); +} + +MemRefType GetStorageMemRefType(MemRefType type) { + IntegerType storage = GetStorageType(type.getElementType()); + if (!storage) { + return nullptr; + } + return MemRefType::Builder(type).setElementType(storage); +} + +// Bitcasts a tensor or scalar `value` to `type`. +Value Bitcast(OpBuilder& builder, Location loc, Type type, Value value) { + if (value.getType() == type) { + return value; + } + if (mlir::isa(type)) { + return tensor::BitcastOp::create(builder, loc, type, value); + } + return ma::BitcastOp::create(builder, loc, type, value); +} + +// Maps `value` from its narrow float type to the storage type. +Value ToStorage(OpBuilder& builder, Location loc, Value value) { + Type type = value.getType(); + if (RankedTensorType tensor_type = GetStorageTensorType(type)) { + return Bitcast(builder, loc, tensor_type, value); + } + if (IntegerType scalar_type = GetStorageType(type)) { + return Bitcast(builder, loc, scalar_type, value); + } + return value; +} + +// op(x : fN) -> bitcast(op(bitcast(x) : iN) : iN -> fN) +template +LogicalResult RetypeDataMovement(OpTy op, PatternRewriter& rewriter) { + Type type = op.getType(); + RankedTensorType storage_type = GetStorageTensorType(type); + if (!storage_type) { + return rewriter.notifyMatchFailure(op, "not a narrow float tensor."); + } + llvm::SmallVector operands = llvm::map_to_vector( + op->getOperands(), + [&](Value operand) { return ToStorage(rewriter, op.getLoc(), operand); }); + auto storage_op = OpTy::create(rewriter, op.getLoc(), storage_type, operands, + op.getProperties(), + op->getDiscardableAttrDictionary().getValue()); + rewriter.replaceOpWithNewOp(op, type, storage_op); + return mlir::success(); +} + +// bitcast(bitcast(x)) -> x, if the types round-trip. +LogicalResult FoldBitcastRoundTrip(tensor::BitcastOp bitcast, + PatternRewriter& rewriter) { + auto producer = bitcast.getSource().getDefiningOp(); + if (!producer || producer.getSource().getType() != bitcast.getType()) { + return rewriter.notifyMatchFailure(bitcast, "not a round trip."); + } + rewriter.replaceOp(bitcast, producer.getSource()); + return mlir::success(); +} + +// Returns a view of narrow float `buffer` as its storage type. The view is +// created right after the definition of `buffer`, so it dominates every use of +// `buffer` and can be shared by all its tile accesses. +Value GetStorageBuffer(PatternRewriter& rewriter, + mlir::TypedValue buffer) { + MemRefType storage_type = GetStorageMemRefType(buffer.getType()); + for (mlir::Operation* user : buffer.getUsers()) { + auto view = mlir::dyn_cast(user); + if (view && view.getType() == storage_type) { + return view; + } + } + OpBuilder::InsertionGuard guard(rewriter); + if (auto arg = mlir::dyn_cast(buffer)) { + rewriter.setInsertionPointToStart(arg.getOwner()); + } else { + rewriter.setInsertionPointAfterValue(buffer); + } + return xtile::MemRefBitcastOp::create(rewriter, buffer.getLoc(), storage_type, + buffer); +} + +// extract(buffer : fN) -> bitcast(extract(view(buffer) : iN) : iN -> fN) +LogicalResult RetypeExtractTile(xtile::ExtractTileOp op, + PatternRewriter& rewriter) { + if (!GetStorageMemRefType(op.getSource().getType())) { + return rewriter.notifyMatchFailure(op, "not a narrow float buffer."); + } + Value buffer = GetStorageBuffer(rewriter, op.getSource()); + auto storage_tile = xtile::ExtractTileOp::create( + rewriter, op.getLoc(), GetStorageTensorType(op.getType()), buffer, + op.getOffsets(), op.getFullTileShapeAttr(), op.getStridesAttr()); + storage_tile->setDiscardableAttrs(op->getDiscardableAttrDictionary()); + rewriter.replaceOpWithNewOp(op, op.getType(), + storage_tile); + return mlir::success(); +} + +// insert(tile : fN, buffer : fN) -> insert(bitcast(tile) : iN, view(buffer)) +LogicalResult RetypeInsertTile(xtile::InsertTileOp op, + PatternRewriter& rewriter) { + if (!GetStorageMemRefType(op.getDestination().getType())) { + return rewriter.notifyMatchFailure(op, "not a narrow float buffer."); + } + Value buffer = GetStorageBuffer(rewriter, op.getDestination()); + Value storage_tile = ToStorage(rewriter, op.getLoc(), op.getSource()); + rewriter.modifyOpInPlace(op, [&] { + op.getSourceMutable().assign(storage_tile); + op.getDestinationMutable().assign(buffer); + }); + return mlir::success(); +} + +class LegalizeNarrowFloatStoragePass + : public impl::LegalizeNarrowFloatStoragePassBase< + LegalizeNarrowFloatStoragePass> { + public: + using LegalizeNarrowFloatStoragePassBase::LegalizeNarrowFloatStoragePassBase; + + void runOnOperation() override { + mlir::MLIRContext* context = &getContext(); + mlir::RewritePatternSet patterns(context); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(FoldBitcastRoundTrip); + patterns.add(RetypeExtractTile); + patterns.add(RetypeInsertTile); + if (mlir::failed( + mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) { + signalPassFailure(); + return; + } + } +}; + +} // namespace +} // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_memref_bitcast.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_memref_bitcast.cc new file mode 100644 index 00000000000000..76863e2a7506e8 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_memref_bitcast.cc @@ -0,0 +1,78 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include "mlir/Conversion/LLVMCommon/LoweringOptions.h" +#include "mlir/Conversion/LLVMCommon/Pattern.h" +#include "mlir/Conversion/LLVMCommon/TypeConverter.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" // IWYU pragma: keep +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Interfaces/DataLayoutInterfaces.h" +#include "mlir/Pass/Pass.h" // IWYU pragma: keep +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/DialectConversion.h" +#include "xla/backends/cpu/codegen/tiled/transforms/passes.h" // IWYU pragma: keep +#include "xla/codegen/xtile/ir/xtile_ops.h" + +namespace xla::cpu { + +#define GEN_PASS_DEF_LOWERMEMREFBITCASTPASS +#include "xla/backends/cpu/codegen/tiled/transforms/passes.h.inc" + +namespace { + +struct LowerMemRefBitcast + : public mlir::ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + + mlir::LogicalResult matchAndRewrite( + xtile::MemRefBitcastOp op, OpAdaptor adaptor, + mlir::ConversionPatternRewriter& rewriter) const override { + rewriter.replaceOp(op, adaptor.getSource()); + return mlir::success(); + } +}; + +class LowerMemRefBitcastPass + : public impl::LowerMemRefBitcastPassBase { + public: + using LowerMemRefBitcastPassBase::LowerMemRefBitcastPassBase; + + void runOnOperation() override { + mlir::ModuleOp module = getOperation(); + mlir::MLIRContext* context = &getContext(); + mlir::LowerToLLVMOptions options(context, mlir::DataLayout(module)); + mlir::LLVMTypeConverter type_converter(context, options); + + mlir::ConversionTarget target(*context); + target.addIllegalOp(); + target.markUnknownOpDynamicallyLegal([](mlir::Operation*) { return true; }); + + mlir::RewritePatternSet patterns(context); + patterns.add(type_converter); + if (mlir::failed(mlir::applyPartialConversion(module, target, + std::move(patterns)))) { + signalPassFailure(); + return; + } + } +}; + +} // namespace +} // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td index 151d1e3f851c9b..1545a80a5a6007 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td @@ -115,6 +115,52 @@ def UnrollVectorsPass : Pass<"xtile-cpu-unroll-vectors"> { } +def LegalizeNarrowFloatStoragePass + : Pass<"xtile-cpu-legalize-narrow-float-storage", "mlir::ModuleOp"> { + let summary = "Moves data movement of narrow floats onto same-width integers."; + + let description = [{ + x86 has no legal vector register class for float types narrower than 32 + bits unless the matching ISA extension is present (AVX512-BF16, + AVX512-FP16, ...). SelectionDAG therefore scalarizes any such vector that + is not immediately extended or truncated, promoting every element to f32 + and paying a libcall (e.g. `__truncsfbf2`) per element on the way back. + This happens for phis, shuffles, masked loads and stores and selects of + narrow float vectors, i.e. for exactly the ops that only move data around. + + This pass treats 8- and 16-bit floats as storage-only types: + + * Tile accesses to narrow float buffers go through an + `xtile.memref_bitcast` view of the buffer as the same-width integer type, + followed by a `tensor.bitcast` back to the float type. + * These bitcasts are then moved through data movement ops (transpose, + reshape, broadcast, slice, concatenate, mask, select) towards the ops that + compute on the float values, e.g. `arith.extf`/`arith.truncf`, where the + remaining pairs cancel. + + Ops the pass does not know simply stop the bitcasts, so it never changes + semantics; at worst the float type survives on that op, as before. + }]; + + let dependentDialects = [ + "::mlir::arith::ArithDialect", + "::mlir::tensor::TensorDialect", + "::xla::xtile::XTileDialect", + ]; +} + +def LowerMemRefBitcastPass + : Pass<"xtile-cpu-lower-memref-bitcast", "mlir::ModuleOp"> { + let summary = "Lowers xtile.memref_bitcast to LLVM."; + + let description = [{ + The memref descriptor holds an opaque pointer, so its LLVM type does not + depend on the element type and the op lowers to nothing. + }]; + + let dependentDialects = ["::mlir::LLVM::LLVMDialect"]; +} + def UnpackSubByteVectorWritePass : Pass<"xtile-cpu-unpack-sub-byte-vector-write"> { let summary = "Unpacks sub-byte vector writes."; diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/legalize_narrow_float_storage.mlir b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/legalize_narrow_float_storage.mlir new file mode 100644 index 00000000000000..093fcfacba9e86 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/legalize_narrow_float_storage.mlir @@ -0,0 +1,153 @@ +// Copyright 2026 The OpenXLA Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================== + +// RUN: fusion_compiler_opt %s -xtile-cpu-legalize-narrow-float-storage \ +// RUN: -split-input-file | FileCheck %s + +func.func @extract_extf(%arg0: memref<16xbf16>, %arg1: index) -> tensor<8xf32> { + %0 = xtile.extract %arg0[%arg1] [8] [1] : memref<16xbf16> -> tensor<8xbf16> + %1 = arith.extf %0 : tensor<8xbf16> to tensor<8xf32> + return %1 : tensor<8xf32> +} +// CHECK-LABEL: func.func @extract_extf( +// CHECK-SAME: %[[ARG0:.*]]: memref<16xbf16>, %[[ARG1:.*]]: index) +// CHECK: %[[VIEW:.*]] = xtile.memref_bitcast %[[ARG0]] : memref<16xbf16> -> memref<16xi16> +// CHECK: %[[TILE:.*]] = xtile.extract %[[VIEW]][%[[ARG1]]] [8] [1] : memref<16xi16> -> tensor<8xi16> +// CHECK: %[[CAST:.*]] = tensor.bitcast %[[TILE]] : tensor<8xi16> to tensor<8xbf16> +// CHECK: arith.extf %[[CAST]] : tensor<8xbf16> to tensor<8xf32> + +// ----- + +func.func @copy(%arg0: memref<16xbf16>, %arg1: memref<16xbf16>, %arg2: index) { + %0 = xtile.extract %arg0[%arg2] [8] [1] : memref<16xbf16> -> tensor<8xbf16> + xtile.insert %0 into %arg1[%arg2] [8] [1] : tensor<8xbf16> -> memref<16xbf16> + return +} +// CHECK-LABEL: func.func @copy( +// CHECK-SAME: %[[ARG0:.*]]: memref<16xbf16>, %[[ARG1:.*]]: memref<16xbf16>, %[[ARG2:.*]]: index) +// CHECK-DAG: %[[SRC:.*]] = xtile.memref_bitcast %[[ARG0]] : memref<16xbf16> -> memref<16xi16> +// CHECK-DAG: %[[DST:.*]] = xtile.memref_bitcast %[[ARG1]] : memref<16xbf16> -> memref<16xi16> +// CHECK: %[[TILE:.*]] = xtile.extract %[[SRC]][%[[ARG2]]] [8] [1] : memref<16xi16> -> tensor<8xi16> +// CHECK-NOT: tensor.bitcast +// CHECK: xtile.insert %[[TILE]] into %[[DST]][%[[ARG2]]] [8] [1] : tensor<8xi16> -> memref<16xi16> + +// ----- + +// All tile accesses to a buffer share one view of it. +func.func @shared_view(%arg0: memref<16xf16>, %arg1: index, %arg2: index) -> tensor<8xf32> { + %0 = xtile.extract %arg0[%arg1] [8] [1] : memref<16xf16> -> tensor<8xf16> + %1 = xtile.extract %arg0[%arg2] [8] [1] : memref<16xf16> -> tensor<8xf16> + %2 = arith.extf %0 : tensor<8xf16> to tensor<8xf32> + %3 = arith.extf %1 : tensor<8xf16> to tensor<8xf32> + %4 = arith.addf %2, %3 : tensor<8xf32> + return %4 : tensor<8xf32> +} +// CHECK-LABEL: func.func @shared_view( +// CHECK: %[[VIEW:.*]] = xtile.memref_bitcast %{{.*}} : memref<16xf16> -> memref<16xi16> +// CHECK-NOT: xtile.memref_bitcast +// CHECK: xtile.extract %[[VIEW]] +// CHECK: xtile.extract %[[VIEW]] + +// ----- + +// The transpose of the truncated tile runs on the storage type. +func.func @truncf_transpose(%arg0: memref<8x8xf32>, %arg1: memref<8x8xbf16>, %arg2: index) { + %0 = xtile.extract %arg0[%arg2, %arg2] [8, 8] [1, 1] : memref<8x8xf32> -> tensor<8x8xf32> + %1 = arith.truncf %0 : tensor<8x8xf32> to tensor<8x8xbf16> + %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<8x8xbf16>) -> tensor<8x8xbf16> + xtile.insert %2 into %arg1[%arg2, %arg2] [8, 8] [1, 1] : tensor<8x8xbf16> -> memref<8x8xbf16> + return +} +// CHECK-LABEL: func.func @truncf_transpose( +// CHECK-SAME: %[[ARG0:.*]]: memref<8x8xf32>, %[[ARG1:.*]]: memref<8x8xbf16>, %[[ARG2:.*]]: index) +// CHECK: %[[DST:.*]] = xtile.memref_bitcast %[[ARG1]] : memref<8x8xbf16> -> memref<8x8xi16> +// CHECK: %[[TILE:.*]] = xtile.extract %[[ARG0]] +// CHECK: %[[TRUNC:.*]] = arith.truncf %[[TILE]] : tensor<8x8xf32> to tensor<8x8xbf16> +// CHECK: %[[CAST:.*]] = tensor.bitcast %[[TRUNC]] : tensor<8x8xbf16> to tensor<8x8xi16> +// CHECK: %[[TRANS:.*]] = stablehlo.transpose %[[CAST]], dims = [1, 0] : (tensor<8x8xi16>) -> tensor<8x8xi16> +// CHECK: xtile.insert %[[TRANS]] into %[[DST]] + +// ----- + +// A chain of data movement ops between tile accesses runs entirely on the +// storage type, without any bitcasts in between. +func.func @data_movement_chain(%arg0: memref<4x8xf8E4M3FN>, %arg1: memref<8x8xf8E4M3FN>, %arg2: index) { + %0 = xtile.extract %arg0[%arg2, %arg2] [4, 8] [1, 1] : memref<4x8xf8E4M3FN> -> tensor<4x8xf8E4M3FN> + %1 = stablehlo.concatenate %0, %0, dim = 0 : (tensor<4x8xf8E4M3FN>, tensor<4x8xf8E4M3FN>) -> tensor<8x8xf8E4M3FN> + %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<8x8xf8E4M3FN>) -> tensor<8x8xf8E4M3FN> + %3 = stablehlo.reshape %2 : (tensor<8x8xf8E4M3FN>) -> tensor<64xf8E4M3FN> + %4 = stablehlo.slice %3 [0:32] : (tensor<64xf8E4M3FN>) -> tensor<32xf8E4M3FN> + %5 = stablehlo.broadcast_in_dim %4, dims = [1] : (tensor<32xf8E4M3FN>) -> tensor<2x32xf8E4M3FN> + %6 = stablehlo.reshape %5 : (tensor<2x32xf8E4M3FN>) -> tensor<8x8xf8E4M3FN> + xtile.insert %6 into %arg1[%arg2, %arg2] [8, 8] [1, 1] : tensor<8x8xf8E4M3FN> -> memref<8x8xf8E4M3FN> + return +} +// CHECK-LABEL: func.func @data_movement_chain( +// CHECK-NOT: tensor.bitcast +// CHECK: %[[TILE:.*]] = xtile.extract {{.*}} : memref<4x8xi8> -> tensor<4x8xi8> +// CHECK: %[[CONCAT:.*]] = stablehlo.concatenate %[[TILE]], %[[TILE]], dim = 0 : (tensor<4x8xi8>, tensor<4x8xi8>) -> tensor<8x8xi8> +// CHECK: %[[TRANS:.*]] = stablehlo.transpose %[[CONCAT]], dims = [1, 0] : (tensor<8x8xi8>) -> tensor<8x8xi8> +// CHECK: %[[RESHAPE:.*]] = stablehlo.reshape %[[TRANS]] : (tensor<8x8xi8>) -> tensor<64xi8> +// CHECK: %[[SLICE:.*]] = stablehlo.slice %[[RESHAPE]] [0:32] : (tensor<64xi8>) -> tensor<32xi8> +// CHECK: %[[BCAST:.*]] = stablehlo.broadcast_in_dim %[[SLICE]], dims = [1] : (tensor<32xi8>) -> tensor<2x32xi8> +// CHECK: %[[RESHAPE2:.*]] = stablehlo.reshape %[[BCAST]] : (tensor<2x32xi8>) -> tensor<8x8xi8> +// CHECK: xtile.insert %[[RESHAPE2]] into {{.*}} : tensor<8x8xi8> -> memref<8x8xi8> +// CHECK-NOT: tensor.bitcast + +// ----- + +func.func @select(%arg0: memref<16xbf16>, %arg1: memref<16xbf16>, %arg2: tensor<8xi1>, %arg3: index) { + %0 = xtile.extract %arg0[%arg3] [8] [1] : memref<16xbf16> -> tensor<8xbf16> + %1 = xtile.extract %arg1[%arg3] [8] [1] : memref<16xbf16> -> tensor<8xbf16> + %2 = arith.select %arg2, %0, %1 : tensor<8xi1>, tensor<8xbf16> + xtile.insert %2 into %arg0[%arg3] [8] [1] : tensor<8xbf16> -> memref<16xbf16> + return +} +// CHECK-LABEL: func.func @select( +// CHECK-SAME: %[[ARG0:.*]]: memref<16xbf16>, %[[ARG1:.*]]: memref<16xbf16>, %[[COND:.*]]: tensor<8xi1>, %[[ARG3:.*]]: index) +// CHECK-DAG: %[[VIEW0:.*]] = xtile.memref_bitcast %[[ARG0]] +// CHECK-DAG: %[[VIEW1:.*]] = xtile.memref_bitcast %[[ARG1]] +// CHECK: %[[LHS:.*]] = xtile.extract %[[VIEW0]] +// CHECK: %[[RHS:.*]] = xtile.extract %[[VIEW1]] +// CHECK: %[[SEL:.*]] = arith.select %[[COND]], %[[LHS]], %[[RHS]] : tensor<8xi1>, tensor<8xi16> +// CHECK: xtile.insert %[[SEL]] into %[[VIEW0]] + +// ----- + +func.func @mask(%arg0: memref<16xbf16>, %arg1: memref<16xbf16>, %arg2: index) { + %cst = arith.constant 0.0 : bf16 + %0 = xtile.extract %arg0[%arg2] [8] [1] : memref<16xbf16> -> tensor<8xbf16> + %1 = xtile.mask %0 bounds [5], %cst : tensor<8xbf16> + xtile.insert %1 into %arg1[%arg2] [8] [1] : tensor<8xbf16> -> memref<16xbf16> + return +} +// CHECK-LABEL: func.func @mask( +// CHECK: %[[ZERO:.*]] = arith.constant 0 : i16 +// CHECK: %[[TILE:.*]] = xtile.extract {{.*}} : memref<16xi16> -> tensor<8xi16> +// CHECK: %[[MASKED:.*]] = xtile.mask %[[TILE]] bounds [5], %[[ZERO]] : tensor<8xi16> +// CHECK: xtile.insert %[[MASKED]] into {{.*}} : tensor<8xi16> -> memref<16xi16> + +// ----- + +// Types that are not 8- or 16-bit floats are left alone. +func.func @not_narrow(%arg0: memref<16xf32>, %arg1: memref<16xf4E2M1FN>, %arg2: index) -> (tensor<8xf32>, tensor<8xf4E2M1FN>) { + %0 = xtile.extract %arg0[%arg2] [8] [1] : memref<16xf32> -> tensor<8xf32> + %1 = xtile.extract %arg1[%arg2] [8] [1] : memref<16xf4E2M1FN> -> tensor<8xf4E2M1FN> + return %0, %1 : tensor<8xf32>, tensor<8xf4E2M1FN> +} +// CHECK-LABEL: func.func @not_narrow( +// CHECK-NOT: xtile.memref_bitcast +// CHECK: xtile.extract %{{.*}} : memref<16xf32> -> tensor<8xf32> +// CHECK: xtile.extract %{{.*}} : memref<16xf4E2M1FN> -> tensor<8xf4E2M1FN> diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_memref_bitcast.mlir b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_memref_bitcast.mlir new file mode 100644 index 00000000000000..7dbd101c3cf36c --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_memref_bitcast.mlir @@ -0,0 +1,43 @@ +// Copyright 2026 The OpenXLA Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================== + +// RUN: fusion_compiler_opt %s -xtile-cpu-lower-memref-bitcast \ +// RUN: -split-input-file | FileCheck %s + +func.func @memref_bitcast(%arg0: memref<16xbf16>, %arg1: index) -> vector<8xi16> { + %c0 = arith.constant 0 : i16 + %0 = xtile.memref_bitcast %arg0 : memref<16xbf16> -> memref<16xi16> + %1 = vector.transfer_read %0[%arg1], %c0 : memref<16xi16>, vector<8xi16> + return %1 : vector<8xi16> +} +// CHECK-LABEL: func.func @memref_bitcast( +// CHECK-SAME: %[[ARG0:.*]]: memref<16xbf16>, %[[ARG1:.*]]: index) +// CHECK-NOT: xtile.memref_bitcast +// CHECK: %[[DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : memref<16xbf16> to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> +// CHECK: %[[VIEW:.*]] = builtin.unrealized_conversion_cast %[[DESC]] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<16xi16> +// CHECK: vector.transfer_read %[[VIEW]][%[[ARG1]]] + +// ----- + +func.func @memref_bitcast_2d(%arg0: memref<4x8xf8E4M3FN>) -> memref<4x8xi8> { + %0 = xtile.memref_bitcast %arg0 : memref<4x8xf8E4M3FN> -> memref<4x8xi8> + return %0 : memref<4x8xi8> +} +// CHECK-LABEL: func.func @memref_bitcast_2d( +// CHECK-SAME: %[[ARG0:.*]]: memref<4x8xf8E4M3FN>) +// CHECK-NOT: xtile.memref_bitcast +// CHECK: %[[DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : memref<4x8xf8E4M3FN> to !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> +// CHECK: %[[VIEW:.*]] = builtin.unrealized_conversion_cast %[[DESC]] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> to memref<4x8xi8> +// CHECK: return %[[VIEW]] diff --git a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir index 4a8457949af424..d05c7e37e3741d 100644 --- a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir +++ b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir @@ -282,3 +282,36 @@ func.func @memref_layout_shape_size_mismatch(%arg0: memref<1024xf32, #xtile.layo func.func @memref_layout_is_not_a_permutation(%arg0: memref<1024xf32, #xtile.layout<[1]>>) { return } + +// ----- + +func.func @memref_bitcast(%arg0: memref<4x16xbf16, #xtile.layout<[0, 1]>>) + -> memref<4x16xi16, #xtile.layout<[0, 1]>> { + %0 = xtile.memref_bitcast %arg0 : memref<4x16xbf16, #xtile.layout<[0, 1]>> -> memref<4x16xi16, #xtile.layout<[0, 1]>> + return %0 : memref<4x16xi16, #xtile.layout<[0, 1]>> +} + +// ----- + +func.func @memref_bitcast_shape_mismatch(%arg0: memref<4x16xbf16>) -> memref<16x4xi16> { + // expected-error@+1 {{requires the same shape for all operands and results}} + %0 = xtile.memref_bitcast %arg0 : memref<4x16xbf16> -> memref<16x4xi16> + return %0 : memref<16x4xi16> +} + +// ----- + +func.func @memref_bitcast_layout_mismatch(%arg0: memref<4x16xbf16, #xtile.layout<[0, 1]>>) + -> memref<4x16xi16, #xtile.layout<[1, 0]>> { + // expected-error@+1 {{does not match result layout}} + %0 = xtile.memref_bitcast %arg0 : memref<4x16xbf16, #xtile.layout<[0, 1]>> -> memref<4x16xi16, #xtile.layout<[1, 0]>> + return %0 : memref<4x16xi16, #xtile.layout<[1, 0]>> +} + +// ----- + +func.func @memref_bitcast_width_mismatch(%arg0: memref<16xbf16>) -> memref<16xi32> { + // expected-error@+1 {{element bit widths differ: 'bf16' vs 'i32'}} + %0 = xtile.memref_bitcast %arg0 : memref<16xbf16> -> memref<16xi32> + return %0 : memref<16xi32> +} diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc index aada99430a4ce4..3e6a2c34a280af 100644 --- a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc +++ b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc @@ -177,6 +177,50 @@ mlir::LogicalResult EntryFuncOp::verify() { return mlir::success(); } +mlir::LogicalResult MemRefBitcastOp::verify() { + mlir::MemRefType source_type = getSource().getType(); + mlir::MemRefType result_type = getResult().getType(); + + if (source_type.getLayout() != result_type.getLayout()) { + return emitOpError() << "source layout: " << source_type.getLayout() + << " does not match result layout: " + << result_type.getLayout(); + } + if (source_type.getMemorySpace() != result_type.getMemorySpace()) { + return emitOpError() << "source and result memory spaces differ"; + } + + mlir::Type source_element_type = source_type.getElementType(); + mlir::Type result_element_type = result_type.getElementType(); + if (!source_element_type.isIntOrFloat() || + !result_element_type.isIntOrFloat()) { + return emitOpError() << "element types must be integer or float, got: " + << source_element_type << " and " + << result_element_type; + } + if (source_element_type.getIntOrFloatBitWidth() != + result_element_type.getIntOrFloatBitWidth()) { + return emitOpError() << "element bit widths differ: " << source_element_type + << " vs " << result_element_type; + } + return mlir::success(); +} + +mlir::OpFoldResult MemRefBitcastOp::fold(FoldAdaptor) { + if (getSource().getType() == getType()) { + return getSource(); + } + // bitcast(bitcast(x)) -> x, or a single bitcast from x otherwise. + if (auto producer = getSource().getDefiningOp()) { + if (producer.getSource().getType() == getType()) { + return producer.getSource(); + } + getSourceMutable().assign(producer.getSource()); + return getResult(); + } + return {}; +} + mlir::TypedValue ExtractTileOp::getBuffer() { return getSource(); } diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td index aa03b558989d4a..1cbaab2442e73b 100644 --- a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td +++ b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td @@ -187,6 +187,31 @@ def SelectBufferOp : XTile_Op<"select_buffer", [Pure]> { }]; } +def MemRefBitcastOp : XTile_Op<"memref_bitcast", + [Pure, SameOperandsAndResultShape]> { + let summary = "Reinterprets the element type of a memref."; + let description = [{ + Shape, layout and memory space are unchanged, so no data is moved: this is + purely a change of how the bits of each element are interpreted. + + Example: + + ```mlir + %0 = xtile.memref_bitcast %arg0 : memref<4x16xbf16> -> memref<4x16xi16> + ``` + }]; + + let arguments = (ins AnyMemRef:$source); + let results = (outs AnyMemRef:$result); + + let assemblyFormat = [{ + $source `:` type($source) `->` type($result) attr-dict + }]; + + let hasVerifier = 1; + let hasFolder = 1; +} + def ExtractTileOp : XTile_Op<"extract", [TiledBufferInterface]> { let summary = "Extract a tile from a memref."; let description = [{ From d9ecfe9451299efe87b26b5f06efd7389c6026eb Mon Sep 17 00:00:00 2001 From: Enver Kayaaslan Date: Thu, 24 Sep 2026 05:13:47 -0700 Subject: [PATCH 24/48] Reverts 84df1268af5c065ada6aea0704152d5de6b2f47d PiperOrigin-RevId: 987427673 --- .../xla/xla/hlo/transforms/simplifiers/BUILD | 4 +- .../simplifiers/hlo_constant_folding.cc | 100 +++- .../simplifiers/hlo_constant_folding.h | 49 +- .../simplifiers/hlo_constant_folding_test.cc | 454 +++++++++++++++++- 4 files changed, 593 insertions(+), 14 deletions(-) diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD index 1f3f0aa3423a74..dae4b21fdff4b6 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD +++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD @@ -1217,7 +1217,6 @@ cc_library( "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/time", - "@tsl//tsl/platform:errors", ], ) @@ -1225,6 +1224,7 @@ xla_cc_test( name = "hlo_constant_folding_test", srcs = ["hlo_constant_folding_test.cc"], deps = [ + ":flatten_call_graph", ":hlo_constant_folding", "//xla:literal", "//xla:literal_util", @@ -1233,12 +1233,14 @@ xla_cc_test( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/parser:hlo_parser", + "//xla/hlo/testlib:filecheck", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/testlib:pattern_matcher_gmock", "//xla/hlo/testlib:test", "//xla/hlo/utils:hlo_matchers", "//xla/service:pattern_matcher", "//xla/tsl/platform:statusor", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc index 4f9f4e286bf1af..6cefdf8def86ab 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include #include +#include #include #include @@ -47,7 +48,6 @@ limitations under the License. #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" -#include "tsl/platform/errors.h" namespace xla { @@ -447,6 +447,84 @@ absl::StatusOr PropagateIdenticalConstantArguments( return changed; } +// TODO(enver): Share CloneComputation with other passes, such as +// FlattenCallGraph, by moving to HloModule or HloComputation. +HloComputation* CloneComputation(HloModule* module, + HloComputation* computation) { + if (!module->has_schedule() || + !module->schedule().is_computation_scheduled(computation)) { + return module->AddEmbeddedComputation(computation->Clone()); + } + auto [clone, clone_sequence] = computation->CloneWithSchedule(); + HloComputation* clone_ptr = module->AddEmbeddedComputation(std::move(clone)); + module->schedule().set_sequence(clone_ptr, clone_sequence); + return clone_ptr; +} + +using SpecializationKey = HloConstantFolding::SpecializationKey; + +absl::StatusOr SpecializeCalls( + HloModule* module, HloComputation* computation, + absl::flat_hash_map& + specialization_cache, + std::vector& computation_versions) { + auto caller_instructions = computation->caller_instructions(); + if (caller_instructions.empty()) { + return false; + } + if (!absl::c_all_of(caller_instructions, [](const HloInstruction* instr) { + return instr->opcode() == HloOpcode::kCall; + })) { + return false; + } + + if (caller_instructions.size() > 1) { + // Sort the caller instructions by their unique id to make the compilation + // deterministic. + absl::c_sort(caller_instructions, + [](const HloInstruction* a, const HloInstruction* b) { + return a->unique_id() < b->unique_id(); + }); + } + + bool changed = false; + bool original_computation_used = false; + for (HloInstruction* caller : caller_instructions) { + std::vector> arguments; + arguments.reserve(caller->operand_count()); + for (const HloInstruction* operand : caller->operands()) { + if (operand->opcode() == HloOpcode::kConstant) { + arguments.push_back(LiteralSlice(operand->literal())); + } else { + arguments.push_back(std::nullopt); + } + } + SpecializationKey key{computation, std::move(arguments)}; + auto it = specialization_cache.find(key); + HloComputation* target_comp = nullptr; + if (it != specialization_cache.end()) { + target_comp = it->second; + } else { + if (!original_computation_used) { + // The first argument combination for this computation uses the original + // computation. + target_comp = computation; + original_computation_used = true; + } else { + target_comp = CloneComputation(module, computation); + computation_versions.push_back(target_comp); + } + specialization_cache.emplace(key, target_comp); + } + + if (caller->to_apply() != target_comp) { + caller->set_to_apply(target_comp); + changed = true; + } + } + return changed; +} + } // namespace absl::StatusOr HloConstantFolding::RunOnComputation( @@ -606,6 +684,7 @@ absl::StatusOr HloConstantFolding::RunOnComputation( absl::StatusOr HloConstantFolding::RunImpl( HloModule* module, const absl::flat_hash_set& execution_threads) { + specialization_cache_.clear(); // Limit the constant folding to 0 iterations to skip folding loops in the // default case. This retains the behavior from before while loop support in // HloEvaluator and may be revised. @@ -627,10 +706,21 @@ absl::StatusOr HloConstantFolding::RunImpl( // arguments from callers to callees. for (auto it = computations.rbegin(); it != computations.rend(); ++it) { HloComputation* computation = *it; - ABSL_ASSIGN_OR_RETURN(bool computation_changed, - RunOnComputation(computation, evaluator.get(), - is_foldable_computation)); - changed |= computation_changed; + // TODO(b/260601110): Early exit the computation if all callers are already + // constant folded. + std::vector computation_versions; + computation_versions.push_back(computation); + ABSL_ASSIGN_OR_RETURN(bool did_specialize, + SpecializeCalls(module, computation, specialization_cache_, + computation_versions)); + changed |= did_specialize; + + for (HloComputation* computation_version : computation_versions) { + ABSL_ASSIGN_OR_RETURN(bool version_changed, + RunOnComputation(computation_version, evaluator.get(), + is_foldable_computation)); + changed |= version_changed; + } } return changed; } diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.h b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.h index 2c4bb8a45049e7..00fdd843bdf336 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.h +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.h @@ -19,13 +19,16 @@ limitations under the License. #include #include #include +#include +#include #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" -#include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/pass/hlo_pass_interface.h" +#include "xla/literal.h" #include "xla/shape.h" namespace xla { @@ -82,6 +85,49 @@ class HloConstantFolding : public HloModulePass { explicit HloConstantFolding(const Options& options) : options_(options) {} absl::string_view name() const override { return "constant_folding"; } + struct SpecializationKey { + const HloComputation* original_computation = nullptr; + std::vector> arguments; + + bool operator==(const SpecializationKey& other) const { + if (original_computation != other.original_computation || + arguments.size() != other.arguments.size()) { + return false; + } + for (size_t i = 0; i < arguments.size(); ++i) { + if (arguments[i].has_value() != other.arguments[i].has_value()) { + return false; + } + if (arguments[i].has_value()) { + if (!arguments[i]->Equal(*other.arguments[i], + /*layout_sensitive=*/true)) { + return false; + } + } + } + return true; + } + + template + friend H AbslHashValue(H h, const SpecializationKey& key) { + h = H::combine(std::move(h), key.original_computation, + key.arguments.size()); + for (const auto& arg : key.arguments) { + if (arg.has_value()) { + h = H::combine(std::move(h), true, Literal::AbslHashable(*arg)); + } else { + h = H::combine(std::move(h), false); + } + } + return h; + } + }; + + const absl::flat_hash_map& + specialization_cache() const { + return specialization_cache_; + } + protected: // Run constant folding operations on the given module. Returns whether the // module was changed (constant expressions folded). @@ -99,6 +145,7 @@ class HloConstantFolding : public HloModulePass { static std::atomic slow_op_counter_; Options options_; + absl::flat_hash_map specialization_cache_; }; } // namespace xla diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc index 12f90085cdc2c1..22506ad1b04e88 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc @@ -21,15 +21,18 @@ limitations under the License. #include #include +#include "absl/status/status_matchers.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/parser/hlo_parser.h" +#include "xla/hlo/testlib/filecheck.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/pattern_matcher_gmock.h" #include "xla/hlo/testlib/test.h" +#include "xla/hlo/transforms/simplifiers/flatten_call_graph.h" #include "xla/hlo/utils/hlo_matchers.h" #include "xla/layout_util.h" #include "xla/literal.h" @@ -46,6 +49,8 @@ limitations under the License. namespace xla { namespace { +using ::absl_testing::IsOkAndHolds; + namespace op = xla::testing::opcode_matchers; namespace m = xla::match; using HloConstantFoldingTest = HloHardwareIndependentTestBase; @@ -759,7 +764,7 @@ TEST_F(HloConstantFoldingTest, HloConstantFolding constant_folding; TF_ASSERT_OK_AND_ASSIGN(bool result, RunHloPass(&constant_folding, module.get())); - EXPECT_FALSE(result); + EXPECT_TRUE(result); } TEST_F(HloConstantFoldingTest, InterproceduralMultipleCallsitesSomeConstants) { @@ -792,11 +797,10 @@ TEST_F(HloConstantFoldingTest, InterproceduralMultipleCallsitesSomeConstants) { RunHloPass(&constant_folding, module.get())); EXPECT_TRUE(result); HloComputation* fn = module->GetComputationWithName("Fn"); - EXPECT_THAT(fn->root_instruction(), - GmockMatch(m::Add(m::Subtract(m::Multiply(m::ConstantScalar(1), - m::Parameter(1)), - m::Parameter(2)), - m::ConstantScalar(2)))); + EXPECT_THAT( + fn->root_instruction(), + GmockMatch(m::Add(m::Subtract(m::ConstantScalar(2), m::Parameter(2)), + m::ConstantScalar(2)))); } TEST_F(HloConstantFoldingTest, @@ -832,7 +836,7 @@ TEST_F(HloConstantFoldingTest, EXPECT_TRUE(result); HloComputation* fn = module->GetComputationWithName("Fn"); EXPECT_THAT(fn->root_instruction(), - GmockMatch(m::Subtract(m::Add(m::Add(), m::Parameter(1)), + GmockMatch(m::Subtract(m::Add(m::Constant(), m::Parameter(1)), m::Constant()))); } @@ -1400,5 +1404,441 @@ TEST_F(HloConstantFoldingTest, LateOptionsDontFoldGteWithControlDependency) { EXPECT_FALSE(result); } +TEST_F(HloConstantFoldingTest, ComputationCalledTwiceOnDifferentConstants) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({15, 15}) + + // CHECK-LABEL: %subcomp.clone ( + // CHECK: constant({25, 25}) + subcomp { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(p1, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %call1 = s32[2]{0} call(%c10, %p0.1), to_apply=%subcomp + // CHECK: %call2 = s32[2]{0} call(%c20, %p1.1), to_apply=%subcomp.clone + ENTRY entry { + p0.1 = s32[2] parameter(0) + p1.1 = s32[2] parameter(1) + c10 = s32[2] constant({10, 10}) + c20 = s32[2] constant({20, 20}) + call1 = s32[2] call(c10, p0.1), to_apply=subcomp + call2 = s32[2] call(c20, p1.1), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + } + + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + FlattenCallGraph flatten; + EXPECT_THAT(flatten.Run(module.get()), IsOkAndHolds(true)); + + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), + IsOkAndHolds(true)); + } +} + +TEST_F(HloConstantFoldingTest, ComputationCalledThriceOnTwoDifferentConstants) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({15, 15}) + + // CHECK-LABEL: %subcomp.clone ( + // CHECK: constant({25, 25}) + + // CHECK-LABEL: %subcomp.clone.1 ( + // CHECK: constant({25, 25}) + subcomp { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(p1, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %call1 = s32[2]{0} call(%c10, %p0.1), to_apply=%subcomp + // CHECK: %call2 = s32[2]{0} call(%c20, %p0.1), to_apply=%subcomp.clone + // CHECK: %call3 = s32[2]{0} call(%c20, %p0.1), to_apply=%subcomp.clone.1 + ENTRY entry { + p0.1 = s32[2] parameter(0) + c10 = s32[2] constant({10, 10}) + c20 = s32[2] constant({20, 20}) + call1 = s32[2] call(c10, p0.1), to_apply=subcomp + call2 = s32[2] call(c20, p0.1), to_apply=subcomp + call3 = s32[2] call(c20, p0.1), to_apply=subcomp + add1 = s32[2] add(call1, call2) + ROOT out = s32[2] add(add1, call3) + })"; + + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + } + + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + FlattenCallGraph flatten; + EXPECT_THAT(flatten.Run(module.get()), IsOkAndHolds(true)); + + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), + IsOkAndHolds(true)); + } +} + +TEST_F(HloConstantFoldingTest, + ComputationCalledTwiceOnDifferentConstantsSingleArgument) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({5, 5}) + subcomp { + p0 = s32[2] parameter(0) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(add0, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: ROOT {{.*}} = s32[2]{0} constant({80, 80}) + ENTRY entry { + c10 = s32[2] constant({10, 10}) + c20 = s32[2] constant({20, 20}) + call1 = s32[2] call(c10), to_apply=subcomp + call2 = s32[2] call(c20), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, + ComputationCalledTwiceOnConstantAndOnNonConstant) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({15, 15}) + + // CHECK-LABEL: %subcomp.clone ( + // CHECK: constant({5, 5}) + subcomp { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(p1, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %call1 = s32[2]{0} call(%c10, %p0.1), to_apply=%subcomp + // CHECK: %call2 = s32[2]{0} call(%p0.1, %p1.1), to_apply=%subcomp.clone + // CHECK: ROOT %out = s32[2]{0} add(%call1, %call2) + ENTRY entry { + p0.1 = s32[2] parameter(0) + p1.1 = s32[2] parameter(1) + c10 = s32[2] constant({10, 10}) + call1 = s32[2] call(c10, p0.1), to_apply=subcomp + call2 = s32[2] call(p0.1, p1.1), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + } + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + FlattenCallGraph flatten; + EXPECT_THAT(flatten.Run(module.get()), IsOkAndHolds(true)); + + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), + IsOkAndHolds(true)); + } +} + +TEST_F(HloConstantFoldingTest, + ComputationCalledTwiceOnConstantAndOnNonConstantSingleArgument) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({5, 5}) + subcomp { + p0 = s32[2] parameter(0) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(add0, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %constant = s32[2]{0} constant({30, 30}) + // CHECK: %call2 = s32[2]{0} call(%p0.1), to_apply=%subcomp + // CHECK: ROOT %out = s32[2]{0} add(%constant, %call2) + ENTRY entry { + p0.1 = s32[2] parameter(0) + c10 = s32[2] constant({10, 10}) + call1 = s32[2] call(c10), to_apply=subcomp + call2 = s32[2] call(p0.1), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, ComputationCalledTwiceOnSameConstants) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({15, 15}) + subcomp { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(p1, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %call1 = s32[2]{0} call(%c10, %p0.1), to_apply=%subcomp + // CHECK: %call2 = s32[2]{0} call(%c10, %p1.1), to_apply=%subcomp + // CHECK: ROOT %out = s32[2]{0} add(%call1, %call2) + ENTRY entry { + p0.1 = s32[2] parameter(0) + p1.1 = s32[2] parameter(1) + c10 = s32[2] constant({10, 10}) + call1 = s32[2] call(c10, p0.1), to_apply=subcomp + call2 = s32[2] call(c10, p1.1), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, + ComputationCalledTwiceOnSameConstantsSingleArgument) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({5, 5}) + subcomp { + p0 = s32[2] parameter(0) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(add0, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: ROOT {{.*}} = s32[2]{0} constant({60, 60}) + ENTRY entry { + p0.1 = s32[2] parameter(0) + c10 = s32[2] constant({10, 10}) + call1 = s32[2] call(c10), to_apply=subcomp + call2 = s32[2] call(c10), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, ComputationCalledOnceOnConstants) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({15, 15}) + subcomp { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(p1, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %call1 = s32[2]{0} call(%c10, %p0.1), to_apply=%subcomp + // CHECK: ROOT %out = s32[2]{0} negate(%call1) + ENTRY entry { + p0.1 = s32[2] parameter(0) + p1.1 = s32[2] parameter(1) + c10 = s32[2] constant({10, 10}) + call1 = s32[2] call(c10, p0.1), to_apply=subcomp + ROOT out = s32[2] negate(call1) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, NestedCallOnConstant) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %bar ( + // CHECK-NEXT: %p0.bar = s32[2]{0} parameter(0) + // CHECK-NEXT: %constant.1 = s32[2]{0} constant({10, 10}) + // CHECK-NEXT: %p1.bar = s32[2]{0} parameter(1) + // CHECK-NEXT: ROOT %out.bar = s32[2]{0} add(%constant.1, %p1.bar) + bar { + p0.bar = s32[2] parameter(0) + p1.bar = s32[2] parameter(1) + negate.bar = s32[2] negate(p0.bar) + ROOT out.bar = s32[2] add(negate.bar, p1.bar) + } + + // CHECK-LABEL: %foo ( + // CHECK-NEXT: %p0.foo = s32[2]{0} parameter(0) + // CHECK-NEXT: %constant = s32[2]{0} constant({-10, -10}) + // CHECK-NEXT: %p1.foo = s32[2]{0} parameter(1) + // CHECK-NEXT: %call.bar = s32[2]{0} call(%constant, %p1.foo), to_apply=%bar + // CHECK-NEXT: ROOT %out.foo = s32[2]{0} negate(%call.bar) + foo { + p0.foo = s32[2] parameter(0) + p1.foo = s32[2] parameter(1) + negate.foo = s32[2] negate(p0.foo) + call.bar = s32[2] call(negate.foo, p1.foo), to_apply=bar + ROOT out.foo = s32[2] negate(call.bar) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK-NEXT: %c10 = s32[2]{0} constant({10, 10}) + // CHECK-NEXT: %p0 = s32[2]{0} parameter(0) + // CHECK-NEXT: %call.foo = s32[2]{0} call(%c10, %p0), to_apply=%foo + // CHECK-NEXT: ROOT %out = s32[2]{0} negate(%call.foo) + ENTRY entry { + p0 = s32[2] parameter(0) + c10 = s32[2] constant({10, 10}) + call.foo = s32[2] call(c10, p0), to_apply=foo + ROOT out = s32[2] negate(call.foo) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, NestedCallAndFooCalledTwice) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %bar ( + // CHECK-NEXT: %p0.bar = s32[2]{0} parameter(0) + // CHECK-NEXT: %negate.bar = s32[2]{0} negate(%p0.bar) + // CHECK-NEXT: ROOT %out.bar = s32[2]{0} add(%negate.bar, %negate.bar) + bar { + p0.bar = s32[2] parameter(0) + negate.bar = s32[2] negate(p0.bar) + ROOT out.bar = s32[2] add(negate.bar, negate.bar) + } + + // CHECK-LABEL: %foo ( + // CHECK-NEXT: %p0.foo = s32[2]{0} parameter(0) + // CHECK-NEXT: %constant.2 = s32[2]{0} constant({20, 20}) + // CHECK-NEXT: %p1.foo = s32[2]{0} parameter(1) + // CHECK-NEXT: ROOT %out.foo = s32[2]{0} add(%constant.2, %p1.foo) + foo { + p0.foo = s32[2] parameter(0) + p1.foo = s32[2] parameter(1) + negate.foo = s32[2] negate(p0.foo) + call.bar = s32[2] call(negate.foo), to_apply=bar + ROOT out.foo = s32[2] add(call.bar, p1.foo) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK-NEXT: %c10 = s32[2]{0} constant({10, 10}) + // CHECK-NEXT: %p0 = s32[2]{0} parameter(0) + // CHECK-NEXT: %call.0 = s32[2]{0} call(%c10, %p0), to_apply=%foo + // CHECK-NEXT: %constant = s32[2]{0} constant({30, 30}) + // CHECK-NEXT: ROOT %out = s32[2]{0} add(%call.0, %constant) + ENTRY entry { + p0 = s32[2] parameter(0) + c10 = s32[2] constant({10, 10}) + call.0 = s32[2] call(c10, p0), to_apply=foo + call.1 = s32[2] call(c10, c10), to_apply=foo + ROOT out = s32[2] add(call.0, call.1) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, NestedCallWhereOriginalCallerIsFolded) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %bar ( + // CHECK: %constant.1 = s32[2]{0} constant({3, 3}) + bar { + p0.bar = s32[2] parameter(0) + c1 = s32[2] constant({1, 1}) + c2 = s32[2] constant({2, 2}) + c3 = s32[2] add(c1, c2) + ROOT out.bar = s32[2] add(p0.bar, c3) + } + + foo { + p0.foo = s32[2] parameter(0) + p1.foo = s32[2] parameter(1) + call.bar = s32[2] call(p0.foo), to_apply=bar + ROOT out.foo = s32[2] add(call.bar, p1.foo) + } + + ENTRY entry { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c10 = s32[2] constant({10, 10}) + call.0 = s32[2] call(c10, p0), to_apply=foo + call.1 = s32[2] call(p1, p0), to_apply=foo + ROOT out = s32[2] add(call.0, call.1) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + } // namespace } // namespace xla From de517cb79aaa40921362f24463f74022b63faad8 Mon Sep 17 00:00:00 2001 From: Christian Sigg Date: Thu, 24 Sep 2026 05:31:25 -0700 Subject: [PATCH 25/48] Add defensive checks for P2P symmetric memory support and null windows. Add pre-checks for P2P symmetric memory support via cudaDeviceCanAccessPeer, post-checks for null windows returned from ncclCommWindowRegister, and defensive null window checks in multimem_addr, peer_addr, and the destructor. PiperOrigin-RevId: 987435182 --- .../gpu/collectives/nccl_symmetric_memory.cc | 46 +++++++++++++++---- 1 file changed, 36 insertions(+), 10 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/collectives/nccl_symmetric_memory.cc b/third_party/xla/xla/backends/gpu/collectives/nccl_symmetric_memory.cc index 49c3433a4f76b9..6531960fb4d149 100644 --- a/third_party/xla/xla/backends/gpu/collectives/nccl_symmetric_memory.cc +++ b/third_party/xla/xla/backends/gpu/collectives/nccl_symmetric_memory.cc @@ -35,7 +35,7 @@ limitations under the License. // Include NCCL after XLA headers. #include "third_party/nccl/nccl.h" -#include "third_party/nccl/nccl_device.h" +#include "third_party/nccl/nccl_device.h" // IWYU pragma: keep #include "xla/tsl/concurrency/executor.h" namespace xla::gpu { @@ -79,16 +79,23 @@ NcclSymmetricMemory::Create(std::shared_ptr comm_state, stream_executor::DeviceAddressBase addr, const std::shared_ptr executor, stream_executor::StreamExecutor* stream_executor) { - ncclWindow_t win; + if (comm_state == nullptr) { + return absl::InvalidArgumentError("NCCL comm_state is null."); + } + + ncclWindow_t win = nullptr; ncclResult_t nccl_status; { - VLOG(3) << absl::StrFormat( - "Create NCCL symmetric memory on comm=%p from: ptr=%p; size=%ld", - comm_state->comm, addr.opaque(), addr.size()); absl::MutexLock lock(comm_state->mutex); - nccl_status = - ncclCommWindowRegister(comm_state->comm, addr.opaque(), addr.size(), - &win, NCCL_WIN_COLL_SYMMETRIC); + ncclComm_t comm = comm_state->comm; + if (comm == nullptr) { + return absl::InvalidArgumentError("NCCL comm is null."); + } + VLOG(3) << absl::StrFormat( + "Create NCCL symmetric memory on comm=%p from: ptr=%p; size=%ld", comm, + addr.opaque(), addr.size()); + nccl_status = ncclCommWindowRegister(comm, addr.opaque(), addr.size(), &win, + NCCL_WIN_COLL_SYMMETRIC); } if (nccl_status != ncclSuccess) { @@ -96,11 +103,22 @@ NcclSymmetricMemory::Create(std::shared_ptr comm_state, return XLA_NCCL_STATUS(nccl_status); } + if (win == nullptr) { + LogInterconnectStatus(stream_executor); + return absl::FailedPreconditionError( + "ncclCommWindowRegister returned a null window because NCCL symmetric " + "memory is not supported on this communicator (e.g. P2P/NVLink or " + "CUMEM is unavailable on fractional/virtualized GPU slices)."); + } + return absl::WrapUnique( new NcclSymmetricMemory(comm_state, win, addr, executor)); } NcclSymmetricMemory::~NcclSymmetricMemory() { + if (win_ == nullptr) { + return; + } absl::Status status = Execute( [&] { @@ -126,9 +144,13 @@ stream_executor::DeviceAddressBase NcclSymmetricMemory::addr() const { absl::StatusOr NcclSymmetricMemory::multimem_addr() const { #if (NCCL_VERSION_CODE >= 22900) || defined(USE_NCCL_HOST_API) + if (win_ == nullptr) { + return absl::FailedPreconditionError( + "Cannot get multimem address from null NCCL window."); + } void* multimem = nullptr; XLA_NCCL_RETURN_IF_ERROR(ncclGetLsaMultimemDevicePointer(win_, 0, &multimem)); - if (multimem) { + if (multimem != nullptr) { return stream_executor::DeviceAddressBase(multimem, addr_.size()); } #endif @@ -139,10 +161,14 @@ NcclSymmetricMemory::multimem_addr() const { absl::StatusOr NcclSymmetricMemory::peer_addr(RankId peer) const { #if (NCCL_VERSION_CODE >= 22902) || defined(USE_NCCL_HOST_API) + if (win_ == nullptr) { + return absl::FailedPreconditionError( + "Cannot get peer address from null NCCL window."); + } void* peer_addr = nullptr; XLA_NCCL_RETURN_IF_ERROR( ncclGetPeerDevicePointer(win_, 0, peer.value(), &peer_addr)); - if (peer_addr) { + if (peer_addr != nullptr) { return stream_executor::DeviceAddressBase(peer_addr, addr_.size()); } return absl::FailedPreconditionError(absl::StrFormat( From 6d4756c6b2920dd0edd9ec23d2f9947f316731dd Mon Sep 17 00:00:00 2001 From: Ilya Tikhonovskiy Date: Thu, 24 Sep 2026 05:58:26 -0700 Subject: [PATCH 26/48] [XLA:GPU] Add cuDNN scaled dot device test suite. Adds cudnn_scaled_dot_test covering block-scaled dot support in cuDNN on Blackwell. Parameterized across MXFP8 and NVFP4 type combinations and transposition layouts, reporting lowering stages (composite, optimized HLO, cuDNN graph representation) and execution parity against reference decomposition. PiperOrigin-RevId: 987447766 --- .../xla/xla/backends/gpu/codegen/BUILD | 44 ++ .../gpu/codegen/cudnn_scaled_dot_test.cc | 570 ++++++++++++++++++ 2 files changed, 614 insertions(+) create mode 100644 third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc diff --git a/third_party/xla/xla/backends/gpu/codegen/BUILD b/third_party/xla/xla/backends/gpu/codegen/BUILD index 0877e5e0921bf2..52bfcd607cb8e3 100644 --- a/third_party/xla/xla/backends/gpu/codegen/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/BUILD @@ -112,6 +112,50 @@ xla_test( ], ) +xla_test( + name = "cudnn_scaled_dot_test", + srcs = ["cudnn_scaled_dot_test.cc"], + backends = ["b200"], + tags = ["cuda-only"], + deps = [ + "//xla:error_spec", + "//xla:literal", + "//xla:shape_util", + "//xla:types", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", + "//xla/backends/autotuner:backends_proto_cc", + "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base", + "//xla/backends/gpu/transforms:composite_rewriter", + "//xla/backends/gpu/transforms:cudnn_custom_call_compiler", + "//xla/backends/gpu/transforms:cudnn_fusion_compiler", + "//xla/backends/gpu/transforms:scaled_dot_rewriter", + "//xla/hlo/ir:hlo", + "//xla/hlo/parser:hlo_parser", + "//xla/hlo/testlib:verified_hlo_module", + "//xla/service/gpu:backend_configs_cc", + "//xla/service/gpu:ir_emission_utils", + "//xla/stream_executor:device_description", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor:semantic_version", + "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor/cuda:cuda_compute_capability", + "//xla/tests:hlo_interpreter_reference_mixin", + "//xla/tests:test_utils", + "//xla/tsl/platform:test", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/random:distributions", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "custom", srcs = ["custom.cc"], diff --git a/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc b/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc new file mode 100644 index 00000000000000..8a9a98e060af43 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc @@ -0,0 +1,570 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include "absl/algorithm/container.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/random/uniform_int_distribution.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_replace.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "xla/backends/autotuner/backends.pb.h" +#include "xla/backends/gpu/tests/hlo_pjrt_gpu_test_base.h" +#include "xla/backends/gpu/transforms/composite_rewriter.h" +#include "xla/backends/gpu/transforms/cudnn_custom_call_compiler.h" +#include "xla/backends/gpu/transforms/cudnn_fusion_compiler.h" +#include "xla/backends/gpu/transforms/scaled_dot_rewriter.h" +#include "xla/error_spec.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/parser/hlo_parser.h" +#include "xla/hlo/testlib/verified_hlo_module.h" +#include "xla/literal.h" +#include "xla/primitive_util.h" +#include "xla/service/gpu/backend_configs.pb.h" +#include "xla/service/gpu/ir_emission_utils.h" +#include "xla/stream_executor/cuda/cuda_compute_capability.h" +#include "xla/stream_executor/device_description.h" +#include "xla/stream_executor/platform_manager.h" +#include "xla/stream_executor/semantic_version.h" +#include "xla/stream_executor/stream_executor.h" +#include "xla/tests/hlo_interpreter_reference_mixin.h" +#include "xla/tests/test_utils.h" +#include "xla/tsl/platform/test.h" +#include "xla/types.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" + +namespace xla { +namespace gpu { +namespace { + +class CudnnScaledDotTestBase + : public HloInterpreterReferenceMixin { + public: + se::StreamExecutor* stream_executor() const { + auto platform = + se::PlatformManager::PlatformWithId(stream_executor_platform_id()); + CHECK_OK(platform); + auto executor = (*platform)->ExecutorForDevice(0); + CHECK_OK(executor); + return *executor; + } + + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = HloPjRtGpuTestBase::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_autotune_level(0); + debug_options.set_xla_gpu_cudnn_gemm_fusion_level(2); + debug_options.clear_xla_gpu_experimental_autotune_backends(); + debug_options.add_xla_gpu_experimental_autotune_backends( + autotuner::Backend::CUDNN); + debug_options.add_xla_gpu_experimental_autotune_backends( + autotuner::Backend::CUBLASLT); + debug_options.set_xla_gpu_experimental_scaled_dot_with_triton(false); + return debug_options; + } + + se::CudaComputeCapability get_cuda_cc() const { + return device_description().cuda_compute_capability(); + } + + bool IsAtLeastBlackwellWithCuDnn9() const { + return get_cuda_cc().IsAtLeastBlackwell() && + gpu_target_config() + .device_description.dnn_version() + .major_version() >= 9; + } + + absl::Status PopulateScale(Literal* scale, std::minstd_rand0* engine) { + switch (scale->shape().element_type()) { + case F8E8M0FNU: { + absl::uniform_int_distribution exponent_distribution(-4, 1); + return scale->Populate( + [&](absl::Span /*indices*/) { + return float8_e8m0fnu( + std::ldexp(1.0f, exponent_distribution(*engine))); + }); + } + case F8E4M3FN: + for (float8_e4m3fn& value : scale->data()) { + value = float8_e4m3fn(std::abs(static_cast(value))); + } + return absl::OkStatus(); + case F8E5M2: + for (float8_e5m2& value : scale->data()) { + value = float8_e5m2(std::abs(static_cast(value))); + } + return absl::OkStatus(); + default: + return absl::InvalidArgumentError( + absl::StrCat("Unsupported scale type: ", + PrimitiveType_Name(scale->shape().element_type()))); + } + } + + absl::StatusOr> MakeScaledDotArguments( + const HloModule* module) { + std::minstd_rand0 engine; + FakeArgumentsOptions options; + options.engine = &engine; + ABSL_ASSIGN_OR_RETURN(std::vector arguments, + MakeFakeArguments(module, options)); + if (arguments.size() < 4) { + return absl::InternalError(absl::StrCat( + "Expected at least 4 scaled-dot arguments, got ", arguments.size())); + } + ABSL_RETURN_IF_ERROR(PopulateScale(&arguments[2], &engine)); + ABSL_RETURN_IF_ERROR(PopulateScale(&arguments[3], &engine)); + return arguments; + } +}; + +enum class CuDnnSupportResult { + kSupported, // Lowered to native cuDNN block-scaled graph, execution succeeds + kNotSupported, // Not supported by cuDNN block scaling, falls back to + // decomposed dot + kFailed, // Lowered to native cuDNN block-scaled graph, but execution + // exhibits discrepancy +}; + +absl::string_view CuDnnSupportResultToString(CuDnnSupportResult result) { + switch (result) { + case CuDnnSupportResult::kSupported: + return "supported"; + case CuDnnSupportResult::kNotSupported: + return "not_supported"; + case CuDnnSupportResult::kFailed: + return "failed"; + } +} + +std::ostream& operator<<(std::ostream& os, CuDnnSupportResult result) { + return os << CuDnnSupportResultToString(result); +} + +struct ScaledDotCoverageTestCase { + PrimitiveType lhs_type; + PrimitiveType rhs_type; + PrimitiveType scale_type; + int block_size; + bool lhs_k_minor; + bool rhs_k_minor; + CuDnnSupportResult expected_result; +}; + +std::vector GetCoverageTestCases() { + return { + // 1. MXFP8 E4M3FN x E4M3FN (scale E8M0FNU, block 32) + {F8E4M3FN, F8E4M3FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E4M3FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E4M3FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E4M3FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kSupported}, + + // 2. MXFP8 E4M3FN x E5M2 (scale E8M0FNU, block 32) + {F8E4M3FN, F8E5M2, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E5M2, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E5M2, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E5M2, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kSupported}, + + // 3. E4M3FN x F4E2M1FN (scale E8M0FNU, block 32) + {F8E4M3FN, F4E2M1FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F4E2M1FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F4E2M1FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F4E2M1FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 4. MXFP8 E5M2 x E4M3FN (scale E8M0FNU, block 32) + {F8E5M2, F8E4M3FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E4M3FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E4M3FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E4M3FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kSupported}, + + // 5. MXFP8 E5M2 x E5M2 (scale E8M0FNU, block 32) - unsupported by cuDNN + {F8E5M2, F8E5M2, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E5M2, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E5M2, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E5M2, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 6. E5M2 x F4E2M1FN (scale E8M0FNU, block 32) + {F8E5M2, F4E2M1FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F4E2M1FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F4E2M1FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F4E2M1FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 7. F4E2M1FN x E4M3FN (scale E8M0FNU, block 32) + {F4E2M1FN, F8E4M3FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E4M3FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E4M3FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E4M3FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 8. F4E2M1FN x E5M2 (scale E8M0FNU, block 32) + {F4E2M1FN, F8E5M2, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E5M2, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E5M2, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E5M2, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 9. F4E2M1FN x F4E2M1FN (scale E8M0FNU, block 32) + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 10. F4E2M1FN x F4E2M1FN (scale E8M0FNU, block 16) + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16, false, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16, false, true, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16, true, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16, true, true, + CuDnnSupportResult::kNotSupported}, + + // 11. NVFP4 F4E2M1FN x F4E2M1FN (scale E4M3FN, block 16) + {F4E2M1FN, F4E2M1FN, F8E4M3FN, 16, false, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E4M3FN, 16, false, true, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E4M3FN, 16, true, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E4M3FN, 16, true, true, + CuDnnSupportResult::kFailed}, + }; +} + +std::string ScaledDotCoverageTestParamToString( + const ::testing::TestParamInfo& info) { + const auto& tc = info.param; + return absl::StrCat( + PrimitiveType_Name(tc.lhs_type), "_", PrimitiveType_Name(tc.rhs_type), + "_", PrimitiveType_Name(tc.scale_type), "_Block", tc.block_size, "_Lhs", + tc.lhs_k_minor, "_Rhs", tc.rhs_k_minor); +} + +class CudnnScaledDotCoverageTest + : public CudnnScaledDotTestBase, + public ::testing::WithParamInterface {}; + +TEST_P(CudnnScaledDotCoverageTest, Executes) { + const auto& param = GetParam(); + const bool lhs_k_minor = param.lhs_k_minor; + const bool rhs_k_minor = param.rhs_k_minor; + + auto cc = get_cuda_cc(); + std::string device_name = "Unknown"; + if (cc.IsAtLeastBlackwell()) { + device_name = "B200"; + } else if (cc.IsAtLeastHopper()) { + device_name = "H100"; + } else { + device_name = "Pre-Hopper"; + } + + std::string lhs_name = + primitive_util::LowercasePrimitiveTypeName(param.lhs_type); + std::string rhs_name = + primitive_util::LowercasePrimitiveTypeName(param.rhs_type); + std::string scale_name = + primitive_util::LowercasePrimitiveTypeName(param.scale_type); + std::string lhs_display = + lhs_k_minor ? lhs_name : absl::StrCat(lhs_name, ".T"); + std::string rhs_display = + rhs_k_minor ? absl::StrCat(rhs_name, ".T") : rhs_name; + + LOG(ERROR) << "Report Device: " << device_name; + LOG(ERROR) << "Report LHS: " << lhs_display; + LOG(ERROR) << "Report RHS: " << rhs_display; + LOG(ERROR) << "Report Scale: " << scale_name; + LOG(ERROR) << "Report BlockSize: " << param.block_size; + + constexpr int64_t m = 128; + constexpr int64_t n = 128; + constexpr int64_t k = 256; + const int64_t scale_k = k / param.block_size; + + const std::string lhs_shape = + lhs_k_minor ? absl::StrCat(m, ",", k) : absl::StrCat(k, ",", m); + const std::string rhs_shape = + rhs_k_minor ? absl::StrCat(n, ",", k) : absl::StrCat(k, ",", n); + const std::string lhs_scale_shape = lhs_k_minor + ? absl::StrCat(m, ",", scale_k) + : absl::StrCat(scale_k, ",", m); + const std::string rhs_scale_shape = rhs_k_minor + ? absl::StrCat(n, ",", scale_k) + : absl::StrCat(scale_k, ",", n); + + // 1. Stage: After Composite + std::string after_composite = "N/A"; + constexpr absl::string_view kCompositeHloTemplate = R"hlo( +HloModule test_composite + +%xla.scaled_dot.1 { + %p0 = $lhs_type[$lhs_shape]{1,0} parameter(0) + %p1 = $rhs_type[$rhs_shape]{1,0} parameter(1) + %p2 = $scale_type[$lhs_scale_shape]{1,0} parameter(2) + %p3 = $scale_type[$rhs_scale_shape]{1,0} parameter(3) + %zero = bf16[] constant(0) + ROOT %dummy = bf16[$output_shape]{1,0} broadcast(%zero), dimensions={} +} + +ENTRY %main { + %lhs = $lhs_type[$lhs_shape]{1,0} parameter(0) + %rhs = $rhs_type[$rhs_shape]{1,0} parameter(1) + %lhs_scale = $scale_type[$lhs_scale_shape]{1,0} parameter(2) + %rhs_scale = $scale_type[$rhs_scale_shape]{1,0} parameter(3) + ROOT %call = bf16[$output_shape]{1,0} call(%lhs, %rhs, %lhs_scale, %rhs_scale), + to_apply=%xla.scaled_dot.1, + is_composite=true, + frontend_attributes={ + composite.attributes="{dimension_numbers=[[[$lhs_contracting_dim],[$rhs_contracting_dim]],[[],[]]]}", + composite.name="xla.scaled_dot", + composite.version="1" + } +} +)hlo"; + + std::string composite_hlo = + absl::StrReplaceAll(kCompositeHloTemplate, + {{"$lhs_type", lhs_name}, + {"$rhs_type", rhs_name}, + {"$scale_type", scale_name}, + {"$lhs_shape", lhs_shape}, + {"$rhs_shape", rhs_shape}, + {"$lhs_scale_shape", lhs_scale_shape}, + {"$rhs_scale_shape", rhs_scale_shape}, + {"$output_shape", absl::StrCat(m, ",", n)}, + {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, + {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); + + if (auto comp_module = ParseAndReturnVerifiedModule(composite_hlo); + comp_module.ok()) { + CompositeRewriter rewriter; + if (auto rewrite_status = rewriter.Run(comp_module->get()); + rewrite_status.ok()) { + const auto* root = + (*comp_module)->entry_computation()->root_instruction(); + after_composite = HloOpcodeString(root->opcode()); + } + } + LOG(ERROR) << "Report AfterComposite: " << after_composite; + EXPECT_EQ(after_composite, "scaled-dot"); + + // 2. Stage: Optimized HLO (compiled directly from composite_hlo) + std::string optimized_hlo = "N/A"; + HloInstruction* custom_call_or_fusion = nullptr; + auto optimized_module_or = GetOptimizedModule(composite_hlo); + if (!optimized_module_or.ok()) { + LOG(ERROR) << "GetOptimizedModule error: " << optimized_module_or.status(); + } + if (optimized_module_or.ok()) { + HloModule& optimized_module = **optimized_module_or; + for (HloComputation* comp : optimized_module.computations()) { + for (HloInstruction* instr : comp->instructions()) { + if (instr->opcode() == HloOpcode::kCustomCall && + instr->custom_call_target() == "__cudnn$blockScaledDot") { + optimized_hlo = "__cudnn$blockScaledDot"; + custom_call_or_fusion = instr; + break; + } + if (instr->opcode() == HloOpcode::kFusion) { + if (auto gpu_cfg = instr->backend_config(); + gpu_cfg.ok() && + gpu_cfg->fusion_backend_config().kind() == "__cudnn$fusion") { + const HloComputation* fused_comp = + instr->fused_instructions_computation(); + bool has_scaled_dot = false; + for (const HloInstruction* fused_instr : + fused_comp->instructions()) { + if (fused_instr->opcode() == HloOpcode::kScaledDot) { + has_scaled_dot = true; + break; + } + } + optimized_hlo = has_scaled_dot ? "__cudnn$fusion" : "dot"; + custom_call_or_fusion = instr; + break; + } + } + if (instr->opcode() == HloOpcode::kScaledDot) { + optimized_hlo = "scaled-dot"; + } else if (instr->opcode() == HloOpcode::kDot && + optimized_hlo == "N/A") { + optimized_hlo = "dot"; + } + } + if (custom_call_or_fusion != nullptr) { + break; + } + } + } + LOG(ERROR) << "Report OptimizedHLO: " << optimized_hlo; + + // 3. Stage: cuDNN Graph + std::string cudnn_graph_pattern = "N/A"; + if (custom_call_or_fusion != nullptr && IsAtLeastBlackwellWithCuDnn9() && + optimized_hlo != "dot" && optimized_module_or.ok()) { + HloModule& optimized_module = **optimized_module_or; + if (custom_call_or_fusion->opcode() == HloOpcode::kFusion) { + BinaryMap compilation_results; + CuDnnFusionCompiler cudnn_compiler(stream_executor()->AsDnn(), + device_description(), + compilation_results); + auto cloned = optimized_module.Clone(); + auto res = cudnn_compiler.Run(cloned.get()); + if (res.ok()) { + cudnn_graph_pattern = "BLOCK_SCALE_DEQUANTIZE + MATMUL"; + } else { + cudnn_graph_pattern = absl::StrCat("FAILED: ", res.status().message()); + } + } else if (custom_call_or_fusion->opcode() == HloOpcode::kCustomCall) { + BinaryMap compilation_results; + CuDnnCustomCallCompiler custom_call_compiler(stream_executor()->AsDnn(), + device_description(), + compilation_results); + auto cloned = optimized_module.Clone(); + auto res = custom_call_compiler.Run(cloned.get()); + if (res.ok() && !compilation_results.empty()) { + cudnn_graph_pattern = "BLOCK_SCALE_DEQUANTIZE + MATMUL"; + } else if (!res.ok()) { + cudnn_graph_pattern = absl::StrCat("FAILED: ", res.status().message()); + } + } + } + + // 4. Stage: Execution Status + std::string status = "FAILURE"; + if (!IsAtLeastBlackwellWithCuDnn9()) { + GTEST_SKIP() + << "Block scaled dot cuDNN execution requires Blackwell and cuDNN 9+."; + } + if (optimized_module_or.ok()) { + auto ref_module_or = ParseAndReturnVerifiedModule(composite_hlo); + if (ref_module_or.ok()) { + CompositeRewriter composite_pass; + (void)composite_pass.Run(ref_module_or->get()); + ScaledDotRewriter decompose_pass; + if (auto decomposed = decompose_pass.Run(ref_module_or->get()); + decomposed.ok()) { + auto ref_optimized_or = + GetOptimizedModule((*ref_module_or)->ToString()); + if (ref_optimized_or.ok()) { + auto run_result = RunAndCompareTwoModules( + std::move(*optimized_module_or), std::move(*ref_optimized_or), + ErrorSpec{/*aabs=*/0.02, /*arel=*/0.05}, + /*run_hlo_passes=*/false); + status = run_result ? "SUCCESS" : "FAILURE"; + if (!run_result) { + LOG(ERROR) << "RunAndCompare error: " << run_result.message(); + } + } else { + LOG(ERROR) << "ref_optimized error: " << ref_optimized_or.status(); + } + } + } + } + + CuDnnSupportResult actual = CuDnnSupportResult::kNotSupported; + if (optimized_hlo == "__cudnn$fusion" || + optimized_hlo == "__cudnn$blockScaledDot") { + actual = (cudnn_graph_pattern == "BLOCK_SCALE_DEQUANTIZE + MATMUL" && + status == "SUCCESS") + ? CuDnnSupportResult::kSupported + : CuDnnSupportResult::kFailed; + } + + LOG(ERROR) << "Report CuDnnGraph: " << CuDnnSupportResultToString(actual); + LOG(ERROR) << "Report CuDnnGraphPattern: " << cudnn_graph_pattern; + LOG(ERROR) << "Report Status: " << status; + + CuDnnSupportResult expected = param.expected_result; + EXPECT_EQ(actual, expected); + + switch (expected) { + case CuDnnSupportResult::kSupported: + EXPECT_EQ(optimized_hlo, "__cudnn$fusion"); + EXPECT_EQ(cudnn_graph_pattern, "BLOCK_SCALE_DEQUANTIZE + MATMUL"); + EXPECT_EQ(status, "SUCCESS"); + break; + case CuDnnSupportResult::kFailed: + EXPECT_EQ(optimized_hlo, "__cudnn$fusion"); + EXPECT_EQ(cudnn_graph_pattern, "BLOCK_SCALE_DEQUANTIZE + MATMUL"); + EXPECT_EQ(status, "FAILURE"); + break; + case CuDnnSupportResult::kNotSupported: + EXPECT_EQ(optimized_hlo, "dot"); + EXPECT_EQ(cudnn_graph_pattern, "N/A"); + EXPECT_EQ(status, "SUCCESS"); + break; + } +} + +INSTANTIATE_TEST_SUITE_P(CudnnScaledDotCoverageTestSuite, + CudnnScaledDotCoverageTest, + ::testing::ValuesIn(GetCoverageTestCases()), + ScaledDotCoverageTestParamToString); + +} // namespace +} // namespace gpu +} // namespace xla From e7a7327618b81c1c8fe48a36ecaa1baa9ca8f26f Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Thu, 24 Sep 2026 06:31:54 -0700 Subject: [PATCH 27/48] [XLA:GPU] Require bitcasts to keep bitwidth in CopyFusion AlgebraicSimplifier can rewrite bitcast-convert instructions into bitcast even when the element bitwidth changes (when the minor-most dimension is contiguous). While PriorityFusion and GpuFusible guard against fusing width-changing bitcasts using hlo_instruction_utils::KeepsBitwidth, CopyFusion previously skipped through any single-user bitcast without checking its bitwidth. Fusing a width-changing bitcast and its consuming copy into a producer fusion breaks the MLIR fusion emitters because GetBitcastMap requires matching element counts between roots and elemental lowering cannot emit a scalar mhlo.bitcast_convert across different bitwidths. Guard kBitcast checks in CopyFusion with hlo_instruction_utils::KeepsBitwidth. PiperOrigin-RevId: 987463670 --- .../xla/xla/backends/gpu/transforms/BUILD | 1 + .../backends/gpu/transforms/copy_fusion.cc | 10 +++-- .../gpu/transforms/copy_fusion_test.cc | 44 +++++++++++++++++++ 3 files changed, 52 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index fcb31cc5e44121..2013ccd9e9dc89 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -763,6 +763,7 @@ cc_library( deps = [ "//xla/codegen:ir_emission_utils", "//xla/hlo/ir:hlo", + "//xla/hlo/ir:hlo_instruction_utils", "//xla/hlo/pass:hlo_pass", "//xla/hlo/utils:hlo_traversal", "//xla/service:call_graph", diff --git a/third_party/xla/xla/backends/gpu/transforms/copy_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/copy_fusion.cc index e3d9e9c2943cb5..072800c0eb340a 100644 --- a/third_party/xla/xla/backends/gpu/transforms/copy_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/copy_fusion.cc @@ -29,6 +29,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/codegen/ir_emission_utils.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instruction_utils.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_original_value.h" #include "xla/hlo/utils/hlo_traversal.h" @@ -53,7 +54,9 @@ bool OnlyElementwiseOpsReachableFromParams(HloComputation* fused_computation) { for (auto user : hlo->users()) { if ((!user->IsElementwiseOnOperand(user->operand_index(hlo)) || HloPredicateIsOp(user)) && - HloPredicateIsNotOp(user)) { + !(HloPredicateIsOp(user) && + hlo_instruction_utils::KeepsBitwidth(*user)) && + HloPredicateIsNotOp(user)) { return false; } if (visited.insert(user).second) { @@ -101,9 +104,10 @@ absl::StatusOr CopyFusion::DoCopyFusion( } copy_user = copy_user->users()[0]; } - // Skip bitcast ops. + // Skip bitcast ops that keep the bitwidth. if (HloPredicateIsOp(copy_user) && - copy_user->user_count() == 1) { + copy_user->user_count() == 1 && + hlo_instruction_utils::KeepsBitwidth(*copy_user)) { copy_user = copy_user->users()[0]; } if (HloPredicateIsOp(copy_user) && diff --git a/third_party/xla/xla/backends/gpu/transforms/copy_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/copy_fusion_test.cc index ddc4950643d6e8..7780daebd6efc9 100644 --- a/third_party/xla/xla/backends/gpu/transforms/copy_fusion_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/copy_fusion_test.cc @@ -676,5 +676,49 @@ TEST_F(CopyFusionTest, PropagateOriginalValue) { EXPECT_EQ(fusion->original_value()->ToString(), R"(({"transpose"}, {}, {}))"); } +TEST_F(CopyFusionTest, CopyFusionBitcastConvertSameBitwidthIsFused) { + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnVerifiedModule(absl::StrCat(kModulePrefix, R"( + fused_computation { + p0 = f32[16,32]{1,0} parameter(0) + ROOT neg = f32[16,32]{1,0} negate(p0) + } + + ENTRY main { + p0 = f32[16,32]{1,0} parameter(0) + fusion = f32[16,32]{1,0} fusion(p0), kind=kLoop, calls=fused_computation + bitcast = s32[16,32]{1,0} bitcast(fusion) + copy = s32[16,32]{1,0} copy(bitcast) + ROOT t = (f32[16,32]{1,0}, s32[16,32]{1,0}) tuple(fusion, copy) + })"))); + ASSERT_OK_AND_ASSIGN(bool changed, cf_.Run(module.get())); + EXPECT_TRUE(changed); + const HloInstruction* root = module->entry_computation()->root_instruction(); + const HloInstruction* fusion = nullptr; + ASSERT_THAT(root, GmockMatch(m::Tuple(m::GetTupleElement(m::Fusion(&fusion)), + m::GetTupleElement()))); + EXPECT_THAT(fusion->fused_expression_root(), + GmockMatch(m::Tuple(m::Negate(), m::Copy(m::Bitcast())))); +} + +TEST_F(CopyFusionTest, CopyFusionBitcastConvertDifferentBitwidthIsNotFused) { + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnVerifiedModule(absl::StrCat(kModulePrefix, R"( + fused_computation { + p0 = f32[16,32]{1,0} parameter(0) + ROOT neg = f32[16,32]{1,0} negate(p0) + } + + ENTRY main { + p0 = f32[16,32]{1,0} parameter(0) + fusion = f32[16,32]{1,0} fusion(p0), kind=kLoop, calls=fused_computation + bitcast = s8[16,32,4]{2,1,0} bitcast(fusion) + copy = s8[16,32,4]{2,1,0} copy(bitcast) + ROOT t = (f32[16,32]{1,0}, s8[16,32,4]{2,1,0}) tuple(fusion, copy) + })"))); + ASSERT_OK_AND_ASSIGN(bool changed, cf_.Run(module.get())); + EXPECT_FALSE(changed); +} + } // namespace gpu } // namespace xla From 9b12b73b2ea8bbd2948614766e9f93b25b3dc734 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= Date: Thu, 24 Sep 2026 06:38:07 -0700 Subject: [PATCH 28/48] Document XLA:GPU compiler/runtime compatibility requirements. Now that we have this compatibility window, we'll often need to ask OSS developers to change their PRs to avoid breaking it (e.g. https://github.com/openxla/xla/pull/46865). So we should have some documentation that we can point to. PiperOrigin-RevId: 987466715 --- third_party/xla/docs/contributing.md | 35 ++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/third_party/xla/docs/contributing.md b/third_party/xla/docs/contributing.md index fa08305d0b2a5c..66a713ab9480bd 100644 --- a/third_party/xla/docs/contributing.md +++ b/third_party/xla/docs/contributing.md @@ -107,6 +107,41 @@ the XLA runner could load and execute the program. That means that there should be no pointers to `HloInstruction` or to other parts of the compiler or the `StreamExecutor`. +##### GPU Compatibility Window {#gpu-compatibility-window} + +XLA:GPU guarantees the following compatibility windows between the compiler +that produces a serialized program and the runtime that executes it: + +* **6 months backward compatibility**: a runtime must be able to execute + programs compiled by a version of XLA up to 6 months older. +* **2 weeks forward compatibility**: a program compiled by a newer version of + XLA must run on a runtime that is up to 2 weeks older. + +**Backward compatibility** means the runtime must keep executing everything +older compilers may have emitted, so existing thunks must not be removed from +the runtime. To retire a thunk, first stop emitting it from the compiler, and +only remove its runtime support 6 months later. + +**Forward compatibility** means the compiler must not emit anything the runtime +cannot yet deserialize or execute. Changes that affect the serialized thunk +sequence, such as adding a new thunk kind or a field the runtime must +understand, must therefore land in two stages: + +1. **Runtime first.** Add the thunk, its proto definitions, and + deserialization support, without emitting it from the compiler. +2. **Compiler after 2 weeks.** Once the runtime change has been submitted for + at least 2 weeks, enable emission in the compiler (e.g. thunk emitter, HLO + passes that produce the new op). + +There are two ways to do this: + +* Split the change into two PRs and send the compiler PR once the runtime + change is 2 weeks old. +* Land both in a single PR, but guard the compiler side behind a + [feature flag](#code-standards) that is off by default, and flip it in a + follow-up after 2 weeks. This also lets you test the full path end to end + in the meantime. + ### Code standards * *Coding style*: We follow [Google's code style guide](https://google.github.io/styleguide/). From 277107e9a3df5a4cefc621db5685cf7f168e5e15 Mon Sep 17 00:00:00 2001 From: Toli Yevtushenko Date: Thu, 24 Sep 2026 07:22:28 -0700 Subject: [PATCH 29/48] [XLA] Update TODO comments on closed duplicate bugs to active canonical bugs PiperOrigin-RevId: 987488947 --- third_party/xla/xla/backends/gpu/runtime/BUILD | 4 ++-- third_party/xla/xla/backends/gpu/tests/BUILD | 4 ++-- third_party/xla/xla/stream_executor/cuda/BUILD | 8 ++++---- third_party/xla/xla/tsl/platform/default/env.cc | 4 ++-- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index d96b6dfca36538..c1fdc1ecc51f15 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -4283,8 +4283,8 @@ xla_test( "xla_amdgpu_any", "no_oss", ], - "gb200": ["broken"], # TODO(b/505963304): re-enable when fixed. - "gb300": ["broken"], # TODO(b/505963304): re-enable when fixed. + "gb200": ["broken"], # TODO(b/505963049): re-enable when fixed. + "gb300": ["broken"], # TODO(b/505963049): re-enable when fixed. }, backends = ["gpu"], disabled_backends = [], diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index f41f0d00fde282..9775a8ccd7691f 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -1421,8 +1421,8 @@ xla_test( "gpu": [ "multi_gpu", ], - "gb200": ["broken"], # TODO(b/505963304): re-enable when fixed. - "gb300": ["broken"], # TODO(b/505963304): re-enable when fixed. + "gb200": ["broken"], # TODO(b/505963049): re-enable when fixed. + "gb300": ["broken"], # TODO(b/505963049): re-enable when fixed. }, backends = ["gpu"], deps = [ diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD index 409f089c94351a..363a2b440b0e2b 100644 --- a/third_party/xla/xla/stream_executor/cuda/BUILD +++ b/third_party/xla/xla/stream_executor/cuda/BUILD @@ -1279,8 +1279,8 @@ xla_test( "gpu": [ "multi_gpu", ], - "gb200": ["broken"], # TODO(b/505963304): re-enable when fixed. - "gb300": ["broken"], # TODO(b/505963304): re-enable when fixed. + "gb200": ["broken"], # TODO(b/505963049): re-enable when fixed. + "gb300": ["broken"], # TODO(b/505963049): re-enable when fixed. }, backends = ["gpu"], tags = ["cuda-only"], @@ -1721,8 +1721,8 @@ xla_test( "gpu": [ "multi_gpu", ], - "gb200": ["broken"], # TODO(b/505963304): re-enable when fixed. - "gb300": ["broken"], # TODO(b/505963304): re-enable when fixed. + "gb200": ["broken"], # TODO(b/505963049): re-enable when fixed. + "gb300": ["broken"], # TODO(b/505963049): re-enable when fixed. }, backends = ["gpu"], tags = ["cuda-only"], diff --git a/third_party/xla/xla/tsl/platform/default/env.cc b/third_party/xla/xla/tsl/platform/default/env.cc index 6ed6443235f56d..4b52100e7d6c06 100644 --- a/third_party/xla/xla/tsl/platform/default/env.cc +++ b/third_party/xla/xla/tsl/platform/default/env.cc @@ -206,7 +206,7 @@ class PosixEnv : public Env { } void SchedClosure(absl::AnyInvocable closure) override { - // TODO(b/27290852): Spawning a new thread here is wasteful, but + // TODO(b/119634930): Spawning a new thread here is wasteful, but // needed to deal with the fact that many `closure` functions are // blocking in the current codebase. std::thread closure_thread(std::move(closure)); @@ -215,7 +215,7 @@ class PosixEnv : public Env { void SchedClosureAfter(int64_t micros, absl::AnyInvocable closure) override { - // TODO(b/27290852): Consuming a thread here is wasteful, but this + // TODO(b/119634930): Consuming a thread here is wasteful, but this // code is (currently) only used in the case where a step fails // (AbortStep). This could be replaced by a timer thread SchedClosure([this, micros, closure = std::move(closure)]() mutable { From 8a022e7d11df218ec1f07c02afe1fc6c006ea493 Mon Sep 17 00:00:00 2001 From: Mikhail Goncharov Date: Thu, 24 Sep 2026 07:28:55 -0700 Subject: [PATCH 30/48] [XLA:GPU] simplify multiple expressions in tiling space That will be used later to add constraints to one indexing map. Should also marginally imporove compilation time. Also call simplify in tiling propagation only if we are not in "symbolic" case where we will not use them simplified anyway. PiperOrigin-RevId: 987492063 --- .../xla/xla/codegen/tiling/experimental/BUILD | 1 + .../xla/codegen/tiling/experimental/tile.cc | 49 +++++++++-- .../xla/codegen/tiling/experimental/tile.h | 11 +++ .../tiling/experimental/tile_propagation.cc | 13 ++- .../tiling/experimental/tiling_space.cc | 16 ++-- .../tiling/experimental/tiling_space.h | 5 +- .../tiling/experimental/tiling_space_test.cc | 83 +++++++++++++------ .../xla/xla/hlo/analysis/symbolic_map.h | 1 + 8 files changed, 136 insertions(+), 43 deletions(-) diff --git a/third_party/xla/xla/codegen/tiling/experimental/BUILD b/third_party/xla/xla/codegen/tiling/experimental/BUILD index 1c3c1cefb0ab2a..fd9be74a3cf56d 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/BUILD +++ b/third_party/xla/xla/codegen/tiling/experimental/BUILD @@ -198,6 +198,7 @@ xla_cc_test( "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", ], ) diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile.cc b/third_party/xla/xla/codegen/tiling/experimental/tile.cc index 20947f5d663f00..15623bf136dbd3 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tile.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tile.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" @@ -225,17 +226,51 @@ void Tile::Replace(const llvm::DenseMap& map) { } } +void SimplifyDimTiles( + llvm::ArrayRef> dim_tile_groups, + const TilingSpace& space) { + int64_t total_dim_tiles = 0; + for (const auto& group : dim_tile_groups) { + total_dim_tiles += group.size(); + } + if (total_dim_tiles == 0) { + return; + } + llvm::SmallVector expressions; + expressions.reserve(total_dim_tiles * 4); + for (auto& group : dim_tile_groups) { + for (DimTile& dim_tile : group) { + expressions.push_back(dim_tile.offset); + expressions.push_back(dim_tile.size); + expressions.push_back(dim_tile.stride); + expressions.push_back(dim_tile.upper_bound); + } + } + expressions = space.SimplifyExpressions(expressions); + CHECK_EQ(expressions.size(), total_dim_tiles * 4); + int idx = 0; + for (auto& group : dim_tile_groups) { + for (DimTile& dim_tile : group) { + dim_tile.offset = expressions[idx++]; + dim_tile.size = expressions[idx++]; + dim_tile.stride = expressions[idx++]; + dim_tile.upper_bound = expressions[idx++]; + } + } +} + +void SimplifyDimTiles(llvm::MutableArrayRef dim_tiles, + const TilingSpace& space) { + SimplifyDimTiles( + llvm::ArrayRef>(&dim_tiles, 1), space); +} + void DimTile::Simplify(const TilingSpace& space) { - offset = space.SimplifyExpression(offset); - size = space.SimplifyExpression(size); - stride = space.SimplifyExpression(stride); - upper_bound = space.SimplifyExpression(upper_bound); + SimplifyDimTiles(*this, space); } void Tile::Simplify() { - for (DimTile& dim_tile : llvm::concat(dim_tiles_, replica_ids_)) { - dim_tile.Simplify(*tiling_space_); - } + SimplifyDimTiles({dim_tiles_, replica_ids_}, *tiling_space_); } Tile Tile::CloneWithNewDims(llvm::SmallVector new_dim_tiles) const { diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile.h b/third_party/xla/xla/codegen/tiling/experimental/tile.h index dbe519ae89df6e..3c958a17d30489 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tile.h +++ b/third_party/xla/xla/codegen/tiling/experimental/tile.h @@ -143,6 +143,17 @@ H AbslHashValue(H h, const DimTile& dim_tile) { return H::combine(std::move(h), static_cast(dim_tile_hash)); } +// Simplifies groups of DimTiles using the dimension and symbol bounds of the +// given tiling space. +void SimplifyDimTiles( + llvm::ArrayRef> dim_tile_groups, + const TilingSpace& space); + +// Simplifies a list of DimTiles using the dimension and symbol bounds of the +// given tiling space. +void SimplifyDimTiles(llvm::MutableArrayRef dim_tiles, + const TilingSpace& space); + // Tile is a collection of tilings for every dimension of output tensor // of an HLO instruction. TiledHloInstruction associates a Tile // with an HLO instruction. diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc index edd62cb0501e6d..8b2abc3b0c8219 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc @@ -1340,14 +1340,13 @@ absl::Status PropagateTileThroughMinimalReshape( mlir_context) + 1; - target_dim_tiles[target_id].Simplify(tiling_space); - if (tiling_space.IsSymbolic()) { VLOG(2) << "Skipping reshape contiguity check as tile sizes are not " "assigned yet."; return absl::OkStatus(); } + target_dim_tiles[target_id].Simplify(tiling_space); return VerifyReshapeContiguity( source_info.dims, source_info.tiles, target_info.dims, absl::MakeSpan(&target_dim_tiles[target_id], 1)); @@ -1369,8 +1368,6 @@ absl::Status PropagateTileThroughMinimalReshape( target_dim_tiles[target_id].size = upper_bounds_inclusive[i] - offsets[i] + 1; target_dim_tiles[target_id].upper_bound = upper_bounds_inclusive[i] + 1; - target_dim_tiles[target_id].Simplify(tiling_space); - target_info.tiles[i] = target_dim_tiles[target_id]; } if (tiling_space.IsSymbolic()) { @@ -1379,6 +1376,14 @@ absl::Status PropagateTileThroughMinimalReshape( return absl::OkStatus(); } + for (auto [i, target_id] : llvm::enumerate(target_info.ids)) { + target_info.tiles[i] = target_dim_tiles[target_id]; + } + SimplifyDimTiles(target_info.tiles, tiling_space); + for (auto [i, target_id] : llvm::enumerate(target_info.ids)) { + target_dim_tiles[target_id] = target_info.tiles[i]; + } + return VerifyReshapeContiguity(source_info.dims, source_info.tiles, target_info.dims, target_info.tiles); } diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc index 12f46e69b2e2e7..f010e9db706c48 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc @@ -618,18 +618,24 @@ void TilingSpace::InitSimplificationIndexing() { } } -SymbolicExpr TilingSpace::SimplifyExpression(const SymbolicExpr& expr) const { +llvm::SmallVector TilingSpace::SimplifyExpressions( + const llvm::SmallVector& expressions) const { if (is_symbolic_) { - return expr.Canonicalize(); + llvm::SmallVector simplified_expressions; + simplified_expressions.reserve(expressions.size()); + for (const auto& expr : expressions) { + simplified_expressions.push_back(expr.Canonicalize()); + } + return simplified_expressions; } - + // TODO(b/565301234): add constraints from tiling space? SymbolicMap map = SymbolicMap::Get(mlir_context(), dimensions_.size(), - rt_vars_.size(), {expr}); + rt_vars_.size(), expressions); IndexingMap indexing_map(map, dim_vars_indexing_, range_vars_indexing_, rt_vars_indexing_); indexing_map.Simplify(IndexingMap::SimplifyPointDimensions::kPreserve); - return indexing_map.GetSymbolicMap().GetResults()[0]; + return std::move(indexing_map).GetSymbolicMap().GetResults(); } absl::StatusOr>> diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h index fb7f920c8757cf..4c1c5a24b50234 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h +++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h @@ -194,9 +194,10 @@ class TilingSpace { bool IsSymbolic() const { return is_symbolic_; } - // Simplifies an expression using actual dimension and symbol bounds + // Simplifies expressions using actual dimension and symbol bounds // based on the assigned tile sizes and runtime variable bounds. - SymbolicExpr SimplifyExpression(const SymbolicExpr& expr) const; + llvm::SmallVector SimplifyExpressions( + const llvm::SmallVector& expressions) const; // Returns the list of valid tilings for the tiling space. absl::StatusOr>> GetValidTilings(); diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc index 5ff15bdce99216..8bb35fcda03251 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc @@ -26,7 +26,9 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/strings/string_view.h" +#include "llvm/ADT/SmallVector.h" #include "mlir/IR/MLIRContext.h" +#include "xla/codegen/tiling/experimental/tile.h" #include "xla/hlo/analysis/indexing_test_utils.h" #include "xla/hlo/analysis/symbolic_expr.h" #include "xla/hlo/analysis/symbolic_map.h" @@ -42,6 +44,7 @@ namespace { using ::absl_testing::StatusIs; using ::mlir::MLIRContext; +using ::testing::ElementsAre; using ::testing::HasSubstr; MATCHER_P(MatchString, tiling_space_string, "") { @@ -441,7 +444,7 @@ TEST_F(TilingSpaceSameShapeMultiOutputTest, TwoOutputsEqualShapesParallelDims) { )")); } -class TilingSpaceSimplifyExpressionTest : public TilingSpaceTest { +class TilingSpaceSimplifyExpressionsTest : public TilingSpaceTest { public: void SetUp() override { TilingSpaceTest::SetUp(); @@ -463,41 +466,71 @@ class TilingSpaceSimplifyExpressionTest : public TilingSpaceTest { CHECK_OK(tiling_space_->AssignTileSizes({16, 2})); } + SymbolicExpr ParseExpr(absl::string_view expr_str) { + return ParseSymbolicExpr(expr_str, &mlir_context_, /*num_dims=*/2); + } + std::unique_ptr tiling_space_; }; -TEST_F(TilingSpaceSimplifyExpressionTest, ModRemovedIfLessThanDivisor) { - SymbolicExpr tid_0 = CreateDimExpr(0, &mlir_context_); - EXPECT_EQ(tiling_space_->SimplifyExpression((tid_0 * 8) % 96), tid_0 * 8); +TEST_F(TilingSpaceSimplifyExpressionsTest, ModRemovedIfLessThanDivisor) { + EXPECT_THAT( + tiling_space_->SimplifyExpressions({ParseExpr("(d0 * 8) mod 96")}), + ElementsAre(ParseExpr("d0 * 8"))); } -TEST_F(TilingSpaceSimplifyExpressionTest, FloorDivFactorsDivisor) { - SymbolicExpr tid_1 = CreateDimExpr(1, &mlir_context_); - EXPECT_EQ(tiling_space_->SimplifyExpression((tid_1 * 2).floorDiv(10)), - tid_1.floorDiv(5)); +TEST_F(TilingSpaceSimplifyExpressionsTest, MultipleExpressionsSimplified) { + EXPECT_THAT(tiling_space_->SimplifyExpressions({ParseExpr("(d0 * 8) mod 96"), + ParseExpr("(d1 * 2) / 10"), + ParseExpr("d0 * 16 + 500")}), + ElementsAre(ParseExpr("d0 * 8"), ParseExpr("d1 / 5"), + ParseExpr("d0 * 16 + 500"))); } -TEST_F(TilingSpaceSimplifyExpressionTest, - ExpressionUnchangedIfNotAlgebraicallyFolds) { - SymbolicExpr tid_0 = CreateDimExpr(0, &mlir_context_); - EXPECT_EQ(tiling_space_->SimplifyExpression(tid_0 * 16 + 500), - tid_0 * 16 + 500); +TEST_F(TilingSpaceSimplifyExpressionsTest, DimTileSimplify) { + DimTile dt{ParseExpr("(d0 * 8) mod 96"), ParseExpr("(d1 * 2) / 10"), + ParseExpr("d0 * 16 + 500"), + ParseExpr("(d0 * 16 + d1 * 2) mod 200")}; + dt.Simplify(*tiling_space_); + EXPECT_EQ(dt.offset, ParseExpr("d0 * 8")); + EXPECT_EQ(dt.size, ParseExpr("d1 / 5")); + EXPECT_EQ(dt.stride, ParseExpr("d0 * 16 + 500")); + EXPECT_EQ(dt.upper_bound, ParseExpr("d0 * 16 + d1 * 2")); } -TEST_F(TilingSpaceSimplifyExpressionTest, NestedFloorDivFactorsDivisor) { - auto expr = ParseSymbolicExpr("(d0 * 16 + d1 * 2) / 200", &mlir_context_, - /*num_dims=*/2); - EXPECT_EQ(tiling_space_->SimplifyExpression(expr), - ParseSymbolicExpr("(d0 * 8 + d1) / 100", &mlir_context_, - /*num_dims=*/2)); +TEST_F(TilingSpaceSimplifyExpressionsTest, SimplifyDimTiles) { + llvm::SmallVector dim_tiles = { + {ParseExpr("(d0 * 8) mod 96"), ParseExpr("(d1 * 2) / 10"), + ParseExpr("d0 * 16 + 500"), ParseExpr("(d0 * 16 + d1 * 2) mod 200")}, + {ParseExpr("(d1 * 2) / 10"), ParseExpr("(d0 * 8) mod 96"), + ParseExpr("d0 * 16 + 500"), ParseExpr("(d0 * 16 + d1 * 2) mod 200")}}; + SimplifyDimTiles(dim_tiles, *tiling_space_); + EXPECT_EQ(dim_tiles[0].offset, ParseExpr("d0 * 8")); + EXPECT_EQ(dim_tiles[0].size, ParseExpr("d1 / 5")); + EXPECT_EQ(dim_tiles[0].stride, ParseExpr("d0 * 16 + 500")); + EXPECT_EQ(dim_tiles[0].upper_bound, ParseExpr("d0 * 16 + d1 * 2")); + EXPECT_EQ(dim_tiles[1].offset, ParseExpr("d1 / 5")); + EXPECT_EQ(dim_tiles[1].size, ParseExpr("d0 * 8")); + EXPECT_EQ(dim_tiles[1].stride, ParseExpr("d0 * 16 + 500")); + EXPECT_EQ(dim_tiles[1].upper_bound, ParseExpr("d0 * 16 + d1 * 2")); + + llvm::SmallVector empty_dim_tiles; + SimplifyDimTiles(empty_dim_tiles, *tiling_space_); + EXPECT_TRUE(empty_dim_tiles.empty()); } -TEST_F(TilingSpaceSimplifyExpressionTest, NestedModRemovedIfLessThanDivisor) { - auto expr = ParseSymbolicExpr("(d0 * 16 + d1 * 2) mod 200", &mlir_context_, - /*num_dims=*/2); - EXPECT_EQ( - tiling_space_->SimplifyExpression(expr), - ParseSymbolicExpr("d0 * 16 + d1 * 2", &mlir_context_, /*num_dims=*/2)); +TEST_F(TilingSpaceSimplifyExpressionsTest, SimplifyDimTilesGroups) { + llvm::SmallVector group1 = { + {ParseExpr("(d0 * 8) mod 96"), ParseExpr("(d1 * 2) / 10"), + ParseExpr("d0 * 16 + 500"), ParseExpr("(d0 * 16 + d1 * 2) mod 200")}}; + llvm::SmallVector group2 = { + {ParseExpr("(d1 * 2) / 10"), ParseExpr("(d0 * 8) mod 96"), + ParseExpr("d0 * 16 + 500"), ParseExpr("(d0 * 16 + d1 * 2) mod 200")}}; + SimplifyDimTiles({group1, group2}, *tiling_space_); + EXPECT_EQ(group1[0].offset, ParseExpr("d0 * 8")); + EXPECT_EQ(group1[0].size, ParseExpr("d1 / 5")); + EXPECT_EQ(group2[0].offset, ParseExpr("d1 / 5")); + EXPECT_EQ(group2[0].size, ParseExpr("d0 * 8")); } TEST_F(TilingSpaceTest, ClonePerformsDeepCopies) { diff --git a/third_party/xla/xla/hlo/analysis/symbolic_map.h b/third_party/xla/xla/hlo/analysis/symbolic_map.h index 95883e7b20f123..08e66cb75c4684 100644 --- a/third_party/xla/xla/hlo/analysis/symbolic_map.h +++ b/third_party/xla/xla/hlo/analysis/symbolic_map.h @@ -20,6 +20,7 @@ limitations under the License. #include #include #include +#include #include "absl/log/check.h" #include "absl/types/span.h" From 1d4cc1bac908b8884d3b353c0957a6a42a89c8a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= Date: Thu, 24 Sep 2026 07:38:15 -0700 Subject: [PATCH 31/48] Parse split-proto merge records from a riegeli::Chain without flattening. `HandleProtoMergeRecord` read each record into a flat `absl::string_view` and then called `MergeFromString`. For records larger than a riegeli buffer (e.g. an HLO module carrying hundreds of MB of constants) the decoded record spans several blocks, and producing a flat view forces riegeli to copy it into a contiguous scratch buffer first. Read the record as a `riegeli::Chain` instead, which shares the decoded blocks, and parse it with `riegeli::ParseMessage(..., set_merge(true))`, which consumes the chain directly. PiperOrigin-RevId: 987496597 --- third_party/xla/xla/util/split_proto/BUILD | 2 ++ .../xla/util/split_proto/split_proto_reader.cc | 17 ++++++++++++----- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/third_party/xla/xla/util/split_proto/BUILD b/third_party/xla/xla/util/split_proto/BUILD index 3f8d0dc579cff4..8a5122362547d3 100644 --- a/third_party/xla/xla/util/split_proto/BUILD +++ b/third_party/xla/xla/util/split_proto/BUILD @@ -74,7 +74,9 @@ cc_library( "@com_google_absl//absl/strings:string_view", "@com_google_protobuf//:protobuf", "@riegeli//riegeli/base:any", + "@riegeli//riegeli/base:chain", "@riegeli//riegeli/bytes:reader", + "@riegeli//riegeli/messages:parse_message", "@riegeli//riegeli/records:record_position", "@riegeli//riegeli/records:record_reader", ], diff --git a/third_party/xla/xla/util/split_proto/split_proto_reader.cc b/third_party/xla/xla/util/split_proto/split_proto_reader.cc index 47fc1419aa08ab..3e3a18284ccf3f 100644 --- a/third_party/xla/xla/util/split_proto/split_proto_reader.cc +++ b/third_party/xla/xla/util/split_proto/split_proto_reader.cc @@ -27,7 +27,9 @@ limitations under the License. #include "google/protobuf/message.h" #include "google/protobuf/reflection.h" #include "riegeli/base/any.h" +#include "riegeli/base/chain.h" #include "riegeli/bytes/reader.h" +#include "riegeli/messages/parse_message.h" #include "riegeli/records/record_position.h" #include "riegeli/records/record_reader.h" #include "xla/tsl/platform/errors.h" @@ -52,15 +54,20 @@ absl::Status ReadRecord(riegeli::RecordReader& record_reader, T& record) { template absl::Status HandleProtoMergeRecord(riegeli::RecordReader& record_reader, google::protobuf::Message& proto) { - absl::string_view record_data; + // Read the record as a `Chain` rather than a flat `string_view`: large + // records (e.g. HLO modules with big constants) span multiple blocks of the + // decoded chunk, and reading them as a flat view would force an extra copy + // to make them contiguous. The `Chain` shares the decoded blocks instead and + // the proto parser can consume it directly. + riegeli::Chain record_data; TF_RETURN_WITH_CONTEXT_IF_ERROR(ReadRecord(record_reader, record_data), "failed to read proto merge record data for ", proto.GetTypeName()); - if (!proto.MergeFromString(record_data)) { - return absl::InternalError(absl::StrFormat( - "Failed to parse proto merge record for %s", proto.GetTypeName())); - } + TF_RETURN_WITH_CONTEXT_IF_ERROR( + riegeli::ParseMessage(record_data, proto, + riegeli::ParseMessageOptions().set_merge(true)), + "Failed to parse proto merge record for ", proto.GetTypeName()); return absl::OkStatus(); } From b08a6e3e6c9e9028b886b1feedd8fbbeec8997b2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= Date: Thu, 24 Sep 2026 07:40:00 -0700 Subject: [PATCH 32/48] Write `ExecutableAndOptionsProto` records uncompressed. `serialized_executable` is already a snappy-compressed split `GpuExecutableProto` (see `WriteSplitGpuExecutable`). Compressing the outer container again saves no space (dropping it adds 0.06% on a 330 MB executable), but it costs an extra compression pass on write and a full decompression pass on load. Riegeli stores the compression type per chunk, so readers don't change and existing artifacts still load. Alternative considered: leave the inner proto uncompressed and compress the outer one instead. That would mean deciding per use case when to skip compression, which adds logic for little gain, since the outer proto is small apart from `serialized_executable`. PiperOrigin-RevId: 987497408 --- .../split_proto/split_executable_and_options_writer.cc | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/util/split_proto/split_executable_and_options_writer.cc b/third_party/xla/xla/util/split_proto/split_executable_and_options_writer.cc index 8ef9c2ba1779c3..828373418e79ef 100644 --- a/third_party/xla/xla/util/split_proto/split_executable_and_options_writer.cc +++ b/third_party/xla/xla/util/split_proto/split_executable_and_options_writer.cc @@ -78,8 +78,13 @@ ExecutableAndOptionsProto GetProtoWithoutSerializedExecutable( absl::Status WriteSplitExecutableAndOptions( const ExecutableAndOptionsProto& executable_and_options, std::unique_ptr writer) { - riegeli::RecordWriter record_writer(std::move(writer), - GetGpuSplitProtoOptions()); + riegeli::RecordWriter record_writer( + std::move(writer), + // The bulk of this proto is `serialized_executable`, which is already a + // compressed split proto. Compressing it again does not reduce the size + // but adds a full (de)compression pass over the whole executable when + // serializing and when loading it. + GetSplitProtoRiegeliOptions(SplitProtoCompression::kNone)); SplitProtoManifest manifest = BuildManifest(); TF_RETURN_WITH_CONTEXT_IF_ERROR( WriteRecord(record_writer, manifest), From 9c7723c4db682b4738098256007f413934adf009 Mon Sep 17 00:00:00 2001 From: Mohammed Anany Date: Thu, 24 Sep 2026 07:40:42 -0700 Subject: [PATCH 33/48] Return a ready empty vector from typed JoinFutures on empty input `Future>({})` cannot select the templated converting constructors of `Future`, because a braced-init-list cannot deduce `U`. Overload resolution therefore falls back to the move constructor from a value-initialized `Future`, and both typed `JoinFutures` overloads returned an invalid future for an empty span instead of a ready empty vector. Calling `Map`, `OnReady` or `Await` on it dereferences a null `AsyncValue`. PiperOrigin-RevId: 987497834 --- third_party/xla/xla/tsl/concurrency/future.h | 4 ++-- .../xla/xla/tsl/concurrency/future_test.cc | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/tsl/concurrency/future.h b/third_party/xla/xla/tsl/concurrency/future.h index aca4ac089c050e..487057effb65f5 100644 --- a/third_party/xla/xla/tsl/concurrency/future.h +++ b/third_party/xla/xla/tsl/concurrency/future.h @@ -1719,7 +1719,7 @@ Future> JoinFutures(absl::Span> futures) { VLOG(2) << "tsl::JoinFutures: " << futures.size() << " futures"; if (futures.empty()) { - return Future>({}); + return Future>(std::vector{}); } auto [promise, future] = MakePromise>(); @@ -1743,7 +1743,7 @@ Future> JoinFutures(absl::Span> futures) { VLOG(2) << "tsl::JoinFutures: " << futures.size() << " futures"; if (futures.empty()) { - return Future>({}); + return Future>(std::vector{}); } auto [promise, future] = MakePromise>(); diff --git a/third_party/xla/xla/tsl/concurrency/future_test.cc b/third_party/xla/xla/tsl/concurrency/future_test.cc index 43a71c6e6b268d..d60575ffc83cd7 100644 --- a/third_party/xla/xla/tsl/concurrency/future_test.cc +++ b/third_party/xla/xla/tsl/concurrency/future_test.cc @@ -1087,6 +1087,25 @@ TEST(FutureTest, JoinCopyableFuturesError) { EXPECT_EQ(join_two.Await().status(), absl::InternalError("error0")); } +TEST(FutureTest, JoinEmptyCopyableFutures) { + std::vector> futures; + Future> join = JoinFutures(futures); + ASSERT_TRUE(join.IsValid()); + EXPECT_TRUE(join.IsReady()); + ASSERT_OK_AND_ASSIGN(std::vector v, join.Await()); + EXPECT_TRUE(v.empty()); +} + +TEST(FutureTest, JoinEmptyMoveOnlyFutures) { + std::vector>> futures; + Future>> join = + JoinFutures>(absl::MakeSpan(futures)); + ASSERT_TRUE(join.IsValid()); + EXPECT_TRUE(join.IsReady()); + ASSERT_OK_AND_ASSIGN(auto v, std::move(join).Await()); + EXPECT_TRUE(v.empty()); +} + TEST(FutureTest, JoinMoveOnlyFuture) { auto [promise0, future0] = MakePromise>(); auto [promise1, future1] = MakePromise>(); From bf988457bbfa1e5624d6c9bf83453e964b23410e Mon Sep 17 00:00:00 2001 From: Aliia Khasanova Date: Thu, 24 Sep 2026 08:20:32 -0700 Subject: [PATCH 34/48] Add update_goldens.py to regenerate AOT compatibility goldens Provides a script to automate regenerating and promoting AOT compatibility golden snapshots from Bazel undeclared test outputs into `/executables//v/`. Also updates `AOTInterceptionPjrtClient::VerifyAgainstGolden` and `test_lib.cc` (`GetUndeclaredOutputsDir`) to reference `update_goldens.py` in golden comparison failure and precondition error messages. The golden comparison failure message prints the failing test's full label (from `TEST_TARGET`), which is what the script takes as its argument. PiperOrigin-RevId: 987522513 --- .../test_lib.cc | 7 +++-- .../xla/tests/aot_interception_pjrt_client.cc | 27 ++++++++++++++----- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc index 98d1c661a30001..065a3691bfb7ef 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc @@ -71,14 +71,17 @@ bool IsEnvFlagEnabled(const char* name) { } // The undeclared-outputs directory exists only when the golden-update helper +// (`third_party/tensorflow/compiler/xla/tests/aot_compatibility_experimental/google/update_goldens.py`) // drives the test. Shared so both callers enforce it with the same message. absl::StatusOr GetUndeclaredOutputsDir() { const char* out_dir = std::getenv("TEST_UNDECLARED_OUTPUTS_DIR"); if (out_dir == nullptr) { return absl::FailedPreconditionError( "TEST_UNDECLARED_OUTPUTS_DIR is unset. Golden updates must be driven " - "by the golden-update helper for this package, which sets it; running " - "the test target directly will not work."); + "by the golden-update helper for this package " + "(Google-internal: third_party/tensorflow/compiler/xla/tests/" + "aot_compatibility_experimental/google/update_goldens.py), which sets " + "it; running the test target directly will not work."); } return std::string(out_dir); } diff --git a/third_party/xla/xla/tests/aot_interception_pjrt_client.cc b/third_party/xla/xla/tests/aot_interception_pjrt_client.cc index b308f43752c725..5b1fab26b8c1da 100644 --- a/third_party/xla/xla/tests/aot_interception_pjrt_client.cc +++ b/third_party/xla/xla/tests/aot_interception_pjrt_client.cc @@ -16,6 +16,7 @@ limitations under the License. #include "xla/tests/aot_interception_pjrt_client.h" #include +#include #include #include #include @@ -119,6 +120,23 @@ absl::Status CompareStructurally( return absl::OkStatus(); } +std::string RegenerationHint(absl::string_view target_name) { + // Bazel exports the running test's full label as TEST_TARGET, which is what + // update_goldens.py takes. Fall back to a placeholder label if it is unset. + const char* test_target = std::getenv("TEST_TARGET"); + std::string label = test_target != nullptr + ? std::string(test_target) + : absl::StrCat("//:", target_name); + return absl::StrCat( + "\n\nRegenerate the goldens for target `", label, + "` and add the newly generated v directory to the change after " + "reviewing the differences.\n" + "(Google-internal: run, from the workspace root, " + "third_party/tensorflow/compiler/xla/tests/" + "aot_compatibility_experimental/google/update_goldens.py ", + label, ")"); +} + } // namespace absl::StatusOr @@ -271,12 +289,9 @@ absl::Status AOTInterceptionPjrtClient::VerifyAgainstGolden( tsl::io::Basename(tsl::io::Dirname(tsl::io::Dirname(artifact_path_))); return absl::Status( diff_status.code(), - absl::StrCat( - "Golden artifact verification failed for ", artifact_path_, ".\n", - diff_status.message(), "\n\nRegenerate the goldens for target `", - target_name, - "` and add the newly generated v directory to the change " - "after reviewing the differences.")); + absl::StrCat("Golden artifact verification failed for ", artifact_path_, + ".\n", diff_status.message(), + RegenerationHint(target_name))); } return absl::OkStatus(); From 615d5ee70e34da548fd4fe149544b5e0c62d548b Mon Sep 17 00:00:00 2001 From: Aliia Khasanova Date: Thu, 24 Sep 2026 09:19:00 -0700 Subject: [PATCH 35/48] Split collective_ops AOT goldens into per-arch directories Goldens move to `executables/collective_ops_aot_test_2gpu_/`, one directory per backend. gb300 stays excluded through the existing `disabled_backends` entry (b/491194726); no gb300 target is generated, so `DetectGpuArchToken()` has no gb300 branch and a gb300 device would hit its `LOG(FATAL)`. PiperOrigin-RevId: 987552652 --- .../aot_compatibility_experimental/BUILD | 5 + .../aot_compatibility_experimental/gpu/BUILD | 2 +- .../gpu/collective_ops_aot_test.cc | 6 +- .../v1/AllGatherMixedTypes.pbtxt | 2180 ++++++++++++++++ .../v1/ReduceScatter.pbtxt | 1615 ++++++++++++ .../v2/AllGatherMixedTypes.pbtxt | 2190 +++++++++++++++++ .../v2/ReduceScatter.pbtxt | 1624 ++++++++++++ .../v1/AllGatherMixedTypes.pbtxt | 2180 ++++++++++++++++ .../v1/ReduceScatter.pbtxt | 1615 ++++++++++++ .../v2/AllGatherMixedTypes.pbtxt | 2190 +++++++++++++++++ .../v2/ReduceScatter.pbtxt | 1624 ++++++++++++ .../v1}/AllGatherMixedTypes.pbtxt | 664 ++--- .../v1}/ReduceScatter.pbtxt | 480 ++-- .../v2}/AllGatherMixedTypes.pbtxt | 680 ++--- .../v2}/ReduceScatter.pbtxt | 497 ++-- .../test_lib.cc | 36 + .../aot_compatibility_experimental/test_lib.h | 7 + 17 files changed, 16053 insertions(+), 1542 deletions(-) create mode 100644 third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/AllGatherMixedTypes.pbtxt create mode 100644 third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/ReduceScatter.pbtxt create mode 100644 third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/AllGatherMixedTypes.pbtxt create mode 100644 third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/ReduceScatter.pbtxt create mode 100644 third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/AllGatherMixedTypes.pbtxt create mode 100644 third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/ReduceScatter.pbtxt create mode 100644 third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/AllGatherMixedTypes.pbtxt create mode 100644 third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/ReduceScatter.pbtxt rename third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/{collective_ops_aot_test_2gpu/v2 => collective_ops_aot_test_2gpu_h100/v1}/AllGatherMixedTypes.pbtxt (60%) rename third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/{collective_ops_aot_test_2gpu/v2 => collective_ops_aot_test_2gpu_h100/v1}/ReduceScatter.pbtxt (67%) rename third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/{collective_ops_aot_test_2gpu/v1 => collective_ops_aot_test_2gpu_h100/v2}/AllGatherMixedTypes.pbtxt (60%) rename third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/{collective_ops_aot_test_2gpu/v1 => collective_ops_aot_test_2gpu_h100/v2}/ReduceScatter.pbtxt (66%) diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD b/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD index fa8e84382ea258..b35d17d3e5fcbf 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD @@ -15,12 +15,17 @@ cc_library( deps = [ "//xla:xla_proto_cc", "//xla/pjrt:pjrt_client", + "//xla/stream_executor:platform", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor:stream_executor_h", "//xla/tests:aot_interception_pjrt_client", "//xla/tests:hlo_test_base", "//xla/tests:pjrt_client_registry", "//xla/tsl/platform:env", "//xla/tsl/platform:logging", "//xla/tsl/platform:test", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD index 944a9784e441ba..6b59ed77da4a0f 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD @@ -24,7 +24,7 @@ xla_test( name = "collective_ops_aot_test", srcs = ["collective_ops_aot_test.cc"], backends = _GPU_TARGETS, - data = glob(["executables/collective_ops_aot_test_2gpu/**/*.pbtxt"]), + data = glob(["executables/collective_ops_aot_test_2gpu_*/**/*.pbtxt"]), disabled_backends = [ # TODO(b/491194726): Serialized executable compute capability 6.0 does not match sm_10.3a. "gb300", diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc index 6dcdf613e6804e..95f19a45dbee22 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc @@ -142,7 +142,8 @@ std::vector GetTestParamsOrDie( INSTANTIATE_TEST_SUITE_P( BackwardsCompatibility, CollectiveOpsAotTest, ValuesIn(GetTestParamsOrDie(GetAotTestParamsForBackwardsCompatibility( - "collective_ops_aot_test_2gpu", AOTTestPlatform::kGpu))), + absl::StrCat("collective_ops_aot_test_2gpu_", DetectGpuArchToken()), + AOTTestPlatform::kGpu))), [](const TestParamInfo& info) { return absl::StrCat("v", info.param.version); }); @@ -150,7 +151,8 @@ INSTANTIATE_TEST_SUITE_P( INSTANTIATE_TEST_SUITE_P( GoldenFileVerification, CollectiveOpsAotTest, ValuesIn(GetTestParamsOrDie(GetAotTestParamsForGoldenFileVerification( - "collective_ops_aot_test_2gpu", AOTTestPlatform::kGpu))), + absl::StrCat("collective_ops_aot_test_2gpu_", DetectGpuArchToken()), + AOTTestPlatform::kGpu))), [](const TestParamInfo& info) { return absl::StrCat("v", info.param.version); }); diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/AllGatherMixedTypes.pbtxt new file mode 100644 index 00000000000000..4e7a607fc3ae22 --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/AllGatherMixedTypes.pbtxt @@ -0,0 +1,2180 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "test_computation" + computations { + name: "fused_transpose" + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.6.1" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.6.1" + } + id: 4294967297 + operand_ids: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.1.1" + opcode: "transpose" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.1.1" + } + dimensions: 1 + dimensions: 0 + id: 4294967296 + operand_ids: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + } + id: 1 + root_id: 4294967296 + is_fusion_computation: true + } + computations { + name: "fused_transpose.1" + instructions { + name: "param_0.3" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.3" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.4.1" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.4.1" + } + id: 8589934593 + operand_ids: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.3" + opcode: "transpose" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.3" + } + dimensions: 1 + dimensions: 0 + id: 8589934592 + operand_ids: 8589934593 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.3" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "fused_broadcast" + instructions { + name: "param_0.5" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.5" + } + id: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p1.1.1" + opcode: "convert" + shape { + element_type: F32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p1.1.1" + } + id: 12884901889 + operand_ids: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.13.1" + opcode: "broadcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.13.1" + } + id: 12884901888 + operand_ids: 12884901889 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.5" + } + id: 3 + root_id: 12884901888 + is_fusion_computation: true + } + computations { + name: "wrapped_broadcast_computation" + instructions { + name: "param_0.7" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.7" + } + id: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.12.1" + opcode: "broadcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.12.1" + } + id: 21474836480 + operand_ids: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.7" + } + id: 5 + root_id: 21474836480 + is_fusion_computation: true + } + computations { + name: "test_computation" + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "wrapped_broadcast" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "wrapped_broadcast" + } + fusion_kind: "kLoop" + id: 9 + operand_ids: 0 + called_computation_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_broadcast_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_broadcast_fusion" + } + fusion_kind: "kLoop" + id: 8 + operand_ids: 0 + called_computation_ids: 3 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "allgather.3" + opcode: "all-gather" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "all-gather-start" + } + dimensions: 1 + id: 10 + operand_ids: 9 + operand_ids: 8 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + instructions { + name: "ag0.1" + opcode: "get-tuple-element" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag0.1" + } + id: 1 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion.1" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion.1" + } + fusion_kind: "kLoop" + id: 7 + operand_ids: 1 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.24" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.24" + } + id: 4 + operand_ids: 7 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ag1.1" + opcode: "get-tuple-element" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag1.1" + } + tuple_index: 1 + id: 2 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion" + } + fusion_kind: "kLoop" + id: 6 + operand_ids: 2 + called_computation_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.1.0" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.1.0" + } + id: 5 + operand_ids: 6 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "out.1" + opcode: "tuple" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "out.1" + } + id: 3 + operand_ids: 4 + operand_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + root_id: 3 + } + host_program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + schedule { + sequences { + key: 0 + value { + instruction_ids: 0 + instruction_ids: 8 + instruction_ids: 9 + instruction_ids: 10 + instruction_ids: 1 + instruction_ids: 7 + instruction_ids: 2 + instruction_ids: 6 + instruction_ids: 4 + instruction_ids: 5 + instruction_ids: 3 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "7b2e4f79c5524b5cbc566da178aa7fda" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + id: 9 + size: 4 + defined_at { + } + } + logical_buffers { + id: 10 + size: 8 + defined_at { + instruction_id: 9 + } + } + logical_buffers { + id: 11 + size: 8 + defined_at { + instruction_id: 8 + } + } + logical_buffers { + id: 12 + size: 16 + defined_at { + instruction_id: 10 + } + } + logical_buffers { + id: 13 + size: 16 + defined_at { + shape_index: 0 + instruction_id: 10 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + shape_index: 1 + instruction_id: 10 + } + } + logical_buffers { + id: 15 + size: 16 + defined_at { + instruction_id: 7 + } + } + logical_buffers { + id: 16 + size: 16 + defined_at { + instruction_id: 6 + } + } + logical_buffers { + id: 17 + size: 16 + defined_at { + instruction_id: 3 + } + } + buffer_allocations { + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 10 + size: 8 + element_type: U32 + } + assigned { + logical_buffer_id: 15 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 1 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 9 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 13 + size: 16 + element_type: U32 + } + assigned { + logical_buffer_id: 16 + size: 16 + element_type: F32 + } + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 17 + size: 16 + element_type: TUPLE + } + is_tuple: true + } + buffer_allocations { + index: 3 + size: 520 + assigned { + logical_buffer_id: 11 + offset: 512 + size: 8 + element_type: F32 + } + assigned { + logical_buffer_id: 12 + size: 16 + element_type: TUPLE + } + assigned { + logical_buffer_id: 14 + offset: 256 + size: 16 + element_type: F32 + } + is_tuple: true + } + heap_simulator_traces { + events { + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + events { + kind: FREE + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + kind: FREE + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 15 + } + peak_buffers { + index: 1 + logical_buffer_ids: 16 + } + peak_buffers { + index: 2 + logical_buffer_ids: 17 + } + peak_buffers { + index: 3 + logical_buffer_ids: 12 + logical_buffer_ids: 11 + logical_buffer_ids: 14 + } + } + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + buffer_allocations { + values { + size: 16 + maybe_live_out: true + } + values { + index: 1 + size: 16 + maybe_live_out: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + is_tuple: true + } + values { + index: 3 + size: 520 + is_tuple: true + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + output_info_map { + shape_index { + indexes: 0 + } + output_info { + } + } + output_info_map { + shape_index { + indexes: 1 + } + output_info { + allocation_index: 1 + } + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_broadcast_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_broadcast_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000.nv.capmerc.text.loop_broadcast_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_broadcast_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\014\000\000@\300\005\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\001\000\000\000\000\000\000\000\000\000\000\000\320\006\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004\020\000\000\000\004\360\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_broadcast_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "wrapped_broadcast" + thunk_id: 3 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "wrapped_broadcast" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000.nv.capmerc.text.wrapped_broadcast\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.wrapped_broadcast\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000Z\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\013\000\000@\300\002\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\007\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\004\020\000\000\000\004\340\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\340\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\342\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\002\000\000\000\000\000\000z\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\302\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000i\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\267\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\340\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000w\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\232\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\264\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "wrapped_broadcast" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "allgather.3" + thunk_id: 4 + } + all_gather_thunk { + collective_config { + operand_element_type: U32 + operand_element_type: F32 + } + buffers { + element_count: 2 + source_buffer { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + buffers { + element_count: 2 + source_buffer { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion.1" + thunk_id: 5 + } + custom_kernel_thunk { + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion_1" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000.nv.capmerc.text.loop_transpose_fusion_1\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion_1\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000r\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\006\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\002\000\000\000\000\000\000\230\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\004\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\324\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\311\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\362\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\204\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\225\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\330\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion_1" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion" + thunk_id: 6 + } + custom_kernel_thunk { + args { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000.nv.capmerc.text.loop_transpose_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/ReduceScatter.pbtxt new file mode 100644 index 00000000000000..0f7f9346f06368 --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/ReduceScatter.pbtxt @@ -0,0 +1,1615 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "main" + computations { + name: "add" + instructions { + name: "lhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "lhs" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "rhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "rhs" + } + parameter_number: 1 + id: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "add.1" + opcode: "add" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "add.1" + } + id: 2 + operand_ids: 0 + operand_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameter_names: "lhs" + parameter_names: "rhs" + } + root_id: 2 + } + computations { + name: "fused_select" + instructions { + name: "param_2.2" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_2.2" + } + parameter_number: 2 + id: 8589934597 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "zero_2" + opcode: "constant" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + } + literal { + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + u32s: 0 + } + id: 8589934598 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p.2" + opcode: "compare" + shape { + element_type: PRED + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p.2" + } + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 + comparison_direction: "EQ" + frontend_attributes { + } + comparison_type: "UNSIGNED" + statistics_viz { + } + comparison_order: "TOTAL" + } + instructions { + name: "pb.2" + opcode: "broadcast" + shape { + element_type: PRED + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "pb.2" + } + id: 8589934593 + operand_ids: 8589934596 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_1.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_1.1" + } + parameter_number: 1 + id: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "data.2" + opcode: "select" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "data.2" + } + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + parameter_names: "param_1.1" + parameter_names: "param_2.2" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "main" + instructions { + name: "c0_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 1 + u32s: 2 + u32s: 3 + u32s: 4 + u32s: 5 + u32s: 6 + u32s: 7 + u32s: 8 + } + id: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "c1_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 10 + u32s: 11 + u32s: 12 + u32s: 13 + u32s: 14 + u32s: 15 + u32s: 16 + u32s: 17 + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + id: 4294967296 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_select_fusion" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_select_fusion" + } + fusion_kind: "kLoop" + id: 4294967299 + operand_ids: 4294967297 + operand_ids: 4294967298 + operand_ids: 4294967296 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ars.3" + opcode: "reduce-scatter" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "reduce-scatter-start" + } + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + id: 1 + root_id: 4294967300 + } + host_program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + entry_computation_id: 1 + schedule { + sequences { + key: 0 + value { + instruction_ids: 1 + instruction_ids: 0 + instruction_ids: 2 + } + } + sequences { + key: 1 + value { + instruction_ids: 4294967296 + instruction_ids: 4294967298 + instruction_ids: 4294967297 + instruction_ids: 4294967299 + instruction_ids: 4294967300 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "c36d8870be58d238935bebefbbc5b827" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + size: 4 + defined_at { + } + } + logical_buffers { + id: 1 + size: 4 + defined_at { + instruction_id: 1 + } + } + logical_buffers { + id: 2 + size: 4 + defined_at { + instruction_id: 2 + } + } + logical_buffers { + id: 10 + size: 32 + defined_at { + instruction_id: 4294967297 + } + } + logical_buffers { + id: 11 + size: 32 + defined_at { + instruction_id: 4294967298 + } + } + logical_buffers { + id: 12 + size: 4 + defined_at { + instruction_id: 4294967296 + } + } + logical_buffers { + id: 13 + size: 32 + defined_at { + instruction_id: 4294967299 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + instruction_id: 4294967300 + } + } + buffer_allocations { + size: 32 + assigned { + logical_buffer_id: 10 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 1 + size: 32 + assigned { + logical_buffer_id: 11 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 12 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 14 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 3 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 2 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 4 + size: 4 + is_thread_local: true + assigned { + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 5 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 1 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 6 + size: 32 + assigned { + logical_buffer_id: 13 + size: 32 + element_type: U32 + } + } + heap_simulator_traces { + events { + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + events { + kind: FREE + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 10 + } + peak_buffers { + index: 1 + logical_buffer_ids: 11 + } + peak_buffers { + index: 2 + logical_buffer_ids: 14 + } + peak_buffers { + index: 3 + logical_buffer_ids: 2 + } + peak_buffers { + index: 4 + logical_buffer_ids: 0 + } + peak_buffers { + index: 5 + logical_buffer_ids: 1 + } + peak_buffers { + index: 6 + logical_buffer_ids: 13 + } + } + binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000p\013\000\000\000\000\000\000\360\006\000\000\000\000\000\000\nd\000\006@\0008\000\004\000@\000\022\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000.nv.merc.debug_frame\000.nv.merc.rela.nv.constant.pic\000.nv.merc.nv.constant.pic\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant.pic\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000.nv.constant4\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\013\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\021\000\014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\021\000\014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\002\000\001\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\254\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\035 \014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\035 \014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\\\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\001\000\000\000\000\000\000J\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\007\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\203\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\260\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\004\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000r\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000g\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\367\000\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\0000\000\000\000\000\000\000\000\021\000\000\000\017\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\001\000\000}\000\000p\002\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000 \005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\000\000\000\000\000\000\000\020\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + buffer_allocations { + values { + size: 32 + is_constant: true + } + values { + index: 1 + size: 32 + is_constant: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + } + values { + index: 3 + size: 4 + is_thread_local: true + } + values { + index: 4 + size: 4 + is_thread_local: true + } + values { + index: 5 + size: 4 + is_thread_local: true + } + values { + index: 6 + size: 32 + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + constants { + symbol_name: "buffer_for_c1_1" + allocation_index: 1 + } + constants { + symbol_name: "buffer_for_c0_1" + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_select_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 32 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 32 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + written: false + written: false + written: false + written: true + custom_kernel { + name: "loop_select_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\030\025\000\000\000\000\000\000\230\017\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000.nv.capmerc.text.loop_select_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_select_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000^\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\200\003 \000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\344\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000.\016\000\202{\004\377\000\340\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\346\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\342\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\343\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\021\000\000@\200\271\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\004\n\370\000\004\000\000\000A\002\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\203\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\030\000\000\000\002\"\016\006\000\000B\000\000\000\001\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\002\"\016\006\000\000B\000\000\000A\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\0028\0162\370\000@\021\000\000\000\000\202\001\n\000\000\002\001@\001\000\000\000\000\000\000\000\000\000\000\000P\010\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\004\020\000\000\000\004\220\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\001\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\002\000\000\000\000\000\000\177\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\305\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\370\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\254\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\005\000\000\000\000\000\000\204\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\272\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\365\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\343\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\021\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\240\014\000\000\000\000\000\000\016\001\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000V\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\256\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000k\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \016\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000|\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000D\016\000\000\000\000\000\000\224\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\272\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000\240\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 4 + kernel_name: "loop_select_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 2 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 3 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 8 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "ars.3" + thunk_id: 3 + } + reduce_scatter_thunk { + collective_config { + operand_element_type: U32 + } + buffers { + element_count: 8 + source_buffer { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + reduction_kind: REDUCTION_KIND_SUM + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/AllGatherMixedTypes.pbtxt new file mode 100644 index 00000000000000..49c9f9b9b06505 --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/AllGatherMixedTypes.pbtxt @@ -0,0 +1,2190 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "test_computation" + computations { + name: "fused_transpose" + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.6.1" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.6.1" + } + id: 4294967297 + operand_ids: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.1.1" + opcode: "transpose" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.1.1" + } + dimensions: 1 + dimensions: 0 + id: 4294967296 + operand_ids: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + } + id: 1 + root_id: 4294967296 + is_fusion_computation: true + } + computations { + name: "fused_transpose.1" + instructions { + name: "param_0.3" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.3" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.4.1" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.4.1" + } + id: 8589934593 + operand_ids: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.3" + opcode: "transpose" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.3" + } + dimensions: 1 + dimensions: 0 + id: 8589934592 + operand_ids: 8589934593 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.3" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "fused_broadcast" + instructions { + name: "param_0.5" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.5" + } + id: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p1.1.1" + opcode: "convert" + shape { + element_type: F32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p1.1.1" + } + id: 12884901889 + operand_ids: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.13.1" + opcode: "broadcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.13.1" + } + id: 12884901888 + operand_ids: 12884901889 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.5" + } + id: 3 + root_id: 12884901888 + is_fusion_computation: true + } + computations { + name: "wrapped_broadcast_computation" + instructions { + name: "param_0.7" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.7" + } + id: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.12.1" + opcode: "broadcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.12.1" + } + id: 21474836480 + operand_ids: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.7" + } + id: 5 + root_id: 21474836480 + is_fusion_computation: true + } + computations { + name: "test_computation" + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "wrapped_broadcast" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "wrapped_broadcast" + } + fusion_kind: "kLoop" + id: 9 + operand_ids: 0 + called_computation_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_broadcast_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_broadcast_fusion" + } + fusion_kind: "kLoop" + id: 8 + operand_ids: 0 + called_computation_ids: 3 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "allgather.3" + opcode: "all-gather" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "all-gather-start" + } + dimensions: 1 + id: 10 + operand_ids: 9 + operand_ids: 8 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + instructions { + name: "ag0.1" + opcode: "get-tuple-element" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag0.1" + } + id: 1 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion.1" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion.1" + } + fusion_kind: "kLoop" + id: 7 + operand_ids: 1 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.24" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.24" + } + id: 4 + operand_ids: 7 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ag1.1" + opcode: "get-tuple-element" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag1.1" + } + tuple_index: 1 + id: 2 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion" + } + fusion_kind: "kLoop" + id: 6 + operand_ids: 2 + called_computation_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.1.0" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.1.0" + } + id: 5 + operand_ids: 6 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "out.1" + opcode: "tuple" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "out.1" + } + id: 3 + operand_ids: 4 + operand_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + root_id: 3 + } + host_program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + schedule { + sequences { + key: 0 + value { + instruction_ids: 0 + instruction_ids: 8 + instruction_ids: 9 + instruction_ids: 10 + instruction_ids: 1 + instruction_ids: 7 + instruction_ids: 2 + instruction_ids: 6 + instruction_ids: 4 + instruction_ids: 5 + instruction_ids: 3 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "7b2e4f79c5524b5cbc566da178aa7fda" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + id: 9 + size: 4 + defined_at { + } + } + logical_buffers { + id: 10 + size: 8 + defined_at { + instruction_id: 9 + } + } + logical_buffers { + id: 11 + size: 8 + defined_at { + instruction_id: 8 + } + } + logical_buffers { + id: 12 + size: 16 + defined_at { + instruction_id: 10 + } + } + logical_buffers { + id: 13 + size: 16 + defined_at { + shape_index: 0 + instruction_id: 10 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + shape_index: 1 + instruction_id: 10 + } + } + logical_buffers { + id: 15 + size: 16 + defined_at { + instruction_id: 7 + } + } + logical_buffers { + id: 16 + size: 16 + defined_at { + instruction_id: 6 + } + } + logical_buffers { + id: 17 + size: 16 + defined_at { + instruction_id: 3 + } + } + buffer_allocations { + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 10 + size: 8 + element_type: U32 + } + assigned { + logical_buffer_id: 15 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 1 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 9 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 13 + size: 16 + element_type: U32 + } + assigned { + logical_buffer_id: 16 + size: 16 + element_type: F32 + } + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 17 + size: 16 + element_type: TUPLE + } + is_tuple: true + } + buffer_allocations { + index: 3 + size: 520 + assigned { + logical_buffer_id: 11 + offset: 512 + size: 8 + element_type: F32 + } + assigned { + logical_buffer_id: 12 + size: 16 + element_type: TUPLE + } + assigned { + logical_buffer_id: 14 + offset: 256 + size: 16 + element_type: F32 + } + is_tuple: true + } + heap_simulator_traces { + events { + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + events { + kind: FREE + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + kind: FREE + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 15 + } + peak_buffers { + index: 1 + logical_buffer_ids: 16 + } + peak_buffers { + index: 2 + logical_buffer_ids: 17 + } + peak_buffers { + index: 3 + logical_buffer_ids: 12 + logical_buffer_ids: 11 + logical_buffer_ids: 14 + } + } + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + buffer_allocations { + values { + size: 16 + maybe_live_out: true + } + values { + index: 1 + size: 16 + maybe_live_out: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + is_tuple: true + } + values { + index: 3 + size: 520 + is_tuple: true + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + output_info_map { + shape_index { + indexes: 0 + } + output_info { + } + } + output_info_map { + shape_index { + indexes: 1 + } + output_info { + allocation_index: 1 + } + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_broadcast_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_broadcast_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000.nv.capmerc.text.loop_broadcast_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_broadcast_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\014\000\000@\300\005\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\001\000\000\000\000\000\000\000\000\000\000\000\320\006\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004\020\000\000\000\004\360\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_broadcast_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "wrapped_broadcast" + thunk_id: 3 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "wrapped_broadcast" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000.nv.capmerc.text.wrapped_broadcast\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.wrapped_broadcast\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000Z\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\013\000\000@\300\002\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\007\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\004\020\000\000\000\004\340\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\340\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\342\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\002\000\000\000\000\000\000z\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\302\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000i\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\267\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\340\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000w\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\232\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\264\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "wrapped_broadcast" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "allgather.3" + thunk_id: 4 + } + all_gather_thunk { + collective_config { + operand_element_type: U32 + operand_element_type: F32 + } + buffers { + element_count: 2 + source_buffer { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + buffers { + element_count: 2 + source_buffer { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion.1" + thunk_id: 5 + } + custom_kernel_thunk { + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion_1" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000.nv.capmerc.text.loop_transpose_fusion_1\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion_1\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000r\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\006\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\002\000\000\000\000\000\000\230\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\004\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\324\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\311\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\362\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\204\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\225\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\330\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion_1" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion" + thunk_id: 6 + } + custom_kernel_thunk { + args { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000.nv.capmerc.text.loop_transpose_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/ReduceScatter.pbtxt new file mode 100644 index 00000000000000..07c084a61d3c00 --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/ReduceScatter.pbtxt @@ -0,0 +1,1624 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "main" + computations { + name: "add" + instructions { + name: "lhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "lhs" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "rhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "rhs" + } + parameter_number: 1 + id: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "add.1" + opcode: "add" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "add.1" + } + id: 2 + operand_ids: 0 + operand_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameter_names: "lhs" + parameter_names: "rhs" + } + root_id: 2 + } + computations { + name: "fused_select" + instructions { + name: "param_2.2" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_2.2" + } + parameter_number: 2 + id: 8589934597 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "zero_2" + opcode: "constant" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + } + literal { + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + u32s: 0 + } + id: 8589934598 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p.2" + opcode: "compare" + shape { + element_type: PRED + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p.2" + } + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 + comparison_direction: "EQ" + frontend_attributes { + } + statistics_viz { + } + comparison_order: "TOTAL" + } + instructions { + name: "pb.2" + opcode: "broadcast" + shape { + element_type: PRED + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "pb.2" + } + id: 8589934593 + operand_ids: 8589934596 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_1.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_1.1" + } + parameter_number: 1 + id: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "data.2" + opcode: "select" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "data.2" + } + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + parameter_names: "param_1.1" + parameter_names: "param_2.2" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "main" + instructions { + name: "c0_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 1 + u32s: 2 + u32s: 3 + u32s: 4 + u32s: 5 + u32s: 6 + u32s: 7 + u32s: 8 + } + id: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "c1_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 10 + u32s: 11 + u32s: 12 + u32s: 13 + u32s: 14 + u32s: 15 + u32s: 16 + u32s: 17 + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + id: 4294967296 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_select_fusion" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_select_fusion" + } + fusion_kind: "kLoop" + id: 4294967299 + operand_ids: 4294967297 + operand_ids: 4294967298 + operand_ids: 4294967296 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ars.3" + opcode: "reduce-scatter" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "reduce-scatter-start" + } + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + id: 1 + root_id: 4294967300 + } + host_program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + entry_computation_id: 1 + schedule { + sequences { + key: 0 + value { + instruction_ids: 1 + instruction_ids: 0 + instruction_ids: 2 + } + } + sequences { + key: 1 + value { + instruction_ids: 4294967296 + instruction_ids: 4294967298 + instruction_ids: 4294967297 + instruction_ids: 4294967299 + instruction_ids: 4294967300 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "c36d8870be58d238935bebefbbc5b827" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + size: 4 + defined_at { + } + } + logical_buffers { + id: 1 + size: 4 + defined_at { + instruction_id: 1 + } + } + logical_buffers { + id: 2 + size: 4 + defined_at { + instruction_id: 2 + } + } + logical_buffers { + id: 10 + size: 32 + defined_at { + instruction_id: 4294967297 + } + } + logical_buffers { + id: 11 + size: 32 + defined_at { + instruction_id: 4294967298 + } + } + logical_buffers { + id: 12 + size: 4 + defined_at { + instruction_id: 4294967296 + } + } + logical_buffers { + id: 13 + size: 32 + defined_at { + instruction_id: 4294967299 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + instruction_id: 4294967300 + } + } + buffer_allocations { + size: 32 + assigned { + logical_buffer_id: 10 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 1 + size: 32 + assigned { + logical_buffer_id: 11 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 12 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 14 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 3 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 2 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 4 + size: 4 + is_thread_local: true + assigned { + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 5 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 1 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 6 + size: 32 + assigned { + logical_buffer_id: 13 + size: 32 + element_type: U32 + } + } + heap_simulator_traces { + events { + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + events { + kind: FREE + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 10 + } + peak_buffers { + index: 1 + logical_buffer_ids: 11 + } + peak_buffers { + index: 2 + logical_buffer_ids: 14 + } + peak_buffers { + index: 3 + logical_buffer_ids: 2 + } + peak_buffers { + index: 4 + logical_buffer_ids: 0 + } + peak_buffers { + index: 5 + logical_buffer_ids: 1 + } + peak_buffers { + index: 6 + logical_buffer_ids: 13 + } + } + binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000p\013\000\000\000\000\000\000\360\006\000\000\000\000\000\000\nd\000\006@\0008\000\004\000@\000\022\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000.nv.merc.debug_frame\000.nv.merc.rela.nv.constant.pic\000.nv.merc.nv.constant.pic\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant.pic\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000.nv.constant4\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\013\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\021\000\014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\021\000\014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\002\000\001\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\254\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\035 \014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\035 \014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\\\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\001\000\000\000\000\000\000J\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\007\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\203\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\260\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\004\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000r\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000g\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\367\000\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\0000\000\000\000\000\000\000\000\021\000\000\000\017\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\001\000\000}\000\000p\002\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000 \005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\000\000\000\000\000\000\000\020\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + buffer_allocations { + values { + size: 32 + is_constant: true + } + values { + index: 1 + size: 32 + is_constant: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + } + values { + index: 3 + size: 4 + is_thread_local: true + } + values { + index: 4 + size: 4 + is_thread_local: true + } + values { + index: 5 + size: 4 + is_thread_local: true + } + values { + index: 6 + size: 32 + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + constants { + symbol_name: "buffer_for_c1_1" + allocation_index: 1 + } + constants { + symbol_name: "buffer_for_c0_1" + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_select_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 32 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 32 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + written: false + written: false + written: false + written: true + custom_kernel { + name: "loop_select_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\030\025\000\000\000\000\000\000\230\017\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000.nv.capmerc.text.loop_select_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_select_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000^\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\200\003 \000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\344\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000.\016\000\202{\004\377\000\340\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\346\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\342\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\343\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\021\000\000@\200\271\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\004\n\370\000\004\000\000\000A\002\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\203\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\030\000\000\000\002\"\016\006\000\000B\000\000\000\001\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\002\"\016\006\000\000B\000\000\000A\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\0028\0162\370\000@\021\000\000\000\000\202\001\n\000\000\002\001@\001\000\000\000\000\000\000\000\000\000\000\000P\010\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\004\020\000\000\000\004\220\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\001\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\002\000\000\000\000\000\000\177\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\305\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\370\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\254\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\005\000\000\000\000\000\000\204\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\272\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\365\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\343\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\021\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\240\014\000\000\000\000\000\000\016\001\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000V\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\256\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000k\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \016\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000|\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000D\016\000\000\000\000\000\000\224\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\272\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000\240\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 4 + kernel_name: "loop_select_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 2 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 3 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 8 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "ars.3" + thunk_id: 3 + } + reduce_scatter_thunk { + collective_config { + operand_element_type: U32 + } + buffers { + element_count: 8 + source_buffer { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + reduction_kind: REDUCTION_KIND_SUM + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/AllGatherMixedTypes.pbtxt new file mode 100644 index 00000000000000..78bf6596da5d5a --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/AllGatherMixedTypes.pbtxt @@ -0,0 +1,2180 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "test_computation" + computations { + name: "fused_transpose" + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.6.1" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.6.1" + } + id: 4294967297 + operand_ids: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.1.1" + opcode: "transpose" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.1.1" + } + dimensions: 1 + dimensions: 0 + id: 4294967296 + operand_ids: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + } + id: 1 + root_id: 4294967296 + is_fusion_computation: true + } + computations { + name: "fused_transpose.1" + instructions { + name: "param_0.3" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.3" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.4.1" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.4.1" + } + id: 8589934593 + operand_ids: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.3" + opcode: "transpose" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.3" + } + dimensions: 1 + dimensions: 0 + id: 8589934592 + operand_ids: 8589934593 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.3" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "fused_broadcast" + instructions { + name: "param_0.5" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.5" + } + id: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p1.1.1" + opcode: "convert" + shape { + element_type: F32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p1.1.1" + } + id: 12884901889 + operand_ids: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.13.1" + opcode: "broadcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.13.1" + } + id: 12884901888 + operand_ids: 12884901889 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.5" + } + id: 3 + root_id: 12884901888 + is_fusion_computation: true + } + computations { + name: "wrapped_broadcast_computation" + instructions { + name: "param_0.7" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.7" + } + id: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.12.1" + opcode: "broadcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.12.1" + } + id: 21474836480 + operand_ids: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.7" + } + id: 5 + root_id: 21474836480 + is_fusion_computation: true + } + computations { + name: "test_computation" + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "wrapped_broadcast" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "wrapped_broadcast" + } + fusion_kind: "kLoop" + id: 9 + operand_ids: 0 + called_computation_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_broadcast_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_broadcast_fusion" + } + fusion_kind: "kLoop" + id: 8 + operand_ids: 0 + called_computation_ids: 3 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "allgather.3" + opcode: "all-gather" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "all-gather-start" + } + dimensions: 1 + id: 10 + operand_ids: 9 + operand_ids: 8 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + instructions { + name: "ag0.1" + opcode: "get-tuple-element" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag0.1" + } + id: 1 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion.1" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion.1" + } + fusion_kind: "kLoop" + id: 7 + operand_ids: 1 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.24" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.24" + } + id: 4 + operand_ids: 7 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ag1.1" + opcode: "get-tuple-element" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag1.1" + } + tuple_index: 1 + id: 2 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion" + } + fusion_kind: "kLoop" + id: 6 + operand_ids: 2 + called_computation_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.1.0" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.1.0" + } + id: 5 + operand_ids: 6 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "out.1" + opcode: "tuple" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "out.1" + } + id: 3 + operand_ids: 4 + operand_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + root_id: 3 + } + host_program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + schedule { + sequences { + key: 0 + value { + instruction_ids: 0 + instruction_ids: 8 + instruction_ids: 9 + instruction_ids: 10 + instruction_ids: 1 + instruction_ids: 7 + instruction_ids: 2 + instruction_ids: 6 + instruction_ids: 4 + instruction_ids: 5 + instruction_ids: 3 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "7b2e4f79c5524b5cbc566da178aa7fda" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + id: 9 + size: 4 + defined_at { + } + } + logical_buffers { + id: 10 + size: 8 + defined_at { + instruction_id: 9 + } + } + logical_buffers { + id: 11 + size: 8 + defined_at { + instruction_id: 8 + } + } + logical_buffers { + id: 12 + size: 16 + defined_at { + instruction_id: 10 + } + } + logical_buffers { + id: 13 + size: 16 + defined_at { + shape_index: 0 + instruction_id: 10 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + shape_index: 1 + instruction_id: 10 + } + } + logical_buffers { + id: 15 + size: 16 + defined_at { + instruction_id: 7 + } + } + logical_buffers { + id: 16 + size: 16 + defined_at { + instruction_id: 6 + } + } + logical_buffers { + id: 17 + size: 16 + defined_at { + instruction_id: 3 + } + } + buffer_allocations { + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 10 + size: 8 + element_type: U32 + } + assigned { + logical_buffer_id: 15 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 1 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 9 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 13 + size: 16 + element_type: U32 + } + assigned { + logical_buffer_id: 16 + size: 16 + element_type: F32 + } + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 17 + size: 16 + element_type: TUPLE + } + is_tuple: true + } + buffer_allocations { + index: 3 + size: 520 + assigned { + logical_buffer_id: 11 + offset: 512 + size: 8 + element_type: F32 + } + assigned { + logical_buffer_id: 12 + size: 16 + element_type: TUPLE + } + assigned { + logical_buffer_id: 14 + offset: 256 + size: 16 + element_type: F32 + } + is_tuple: true + } + heap_simulator_traces { + events { + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + events { + kind: FREE + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + kind: FREE + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 15 + } + peak_buffers { + index: 1 + logical_buffer_ids: 16 + } + peak_buffers { + index: 2 + logical_buffer_ids: 17 + } + peak_buffers { + index: 3 + logical_buffer_ids: 12 + logical_buffer_ids: 11 + logical_buffer_ids: 14 + } + } + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + buffer_allocations { + values { + size: 16 + maybe_live_out: true + } + values { + index: 1 + size: 16 + maybe_live_out: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + is_tuple: true + } + values { + index: 3 + size: 520 + is_tuple: true + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + output_info_map { + shape_index { + indexes: 0 + } + output_info { + } + } + output_info_map { + shape_index { + indexes: 1 + } + output_info { + allocation_index: 1 + } + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_broadcast_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_broadcast_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000.nv.capmerc.text.loop_broadcast_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_broadcast_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\014\000\000@\300\005\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\001\000\000\000\000\000\000\000\000\000\000\000\320\006\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004\020\000\000\000\004\360\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_broadcast_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "wrapped_broadcast" + thunk_id: 3 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "wrapped_broadcast" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000.nv.capmerc.text.wrapped_broadcast\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.wrapped_broadcast\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000Z\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\013\000\000@\300\002\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\007\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\004\020\000\000\000\004\340\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\340\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\342\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\002\000\000\000\000\000\000z\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\302\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000i\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\267\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\340\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000w\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\232\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\264\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "wrapped_broadcast" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "allgather.3" + thunk_id: 4 + } + all_gather_thunk { + collective_config { + operand_element_type: U32 + operand_element_type: F32 + } + buffers { + element_count: 2 + source_buffer { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + buffers { + element_count: 2 + source_buffer { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion.1" + thunk_id: 5 + } + custom_kernel_thunk { + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion_1" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000.nv.capmerc.text.loop_transpose_fusion_1\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion_1\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000r\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\006\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\002\000\000\000\000\000\000\230\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\004\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\324\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\311\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\362\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\204\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\225\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\330\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion_1" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion" + thunk_id: 6 + } + custom_kernel_thunk { + args { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000.nv.capmerc.text.loop_transpose_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/ReduceScatter.pbtxt new file mode 100644 index 00000000000000..69429709332001 --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/ReduceScatter.pbtxt @@ -0,0 +1,1615 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "main" + computations { + name: "add" + instructions { + name: "lhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "lhs" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "rhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "rhs" + } + parameter_number: 1 + id: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "add.1" + opcode: "add" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "add.1" + } + id: 2 + operand_ids: 0 + operand_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameter_names: "lhs" + parameter_names: "rhs" + } + root_id: 2 + } + computations { + name: "fused_select" + instructions { + name: "param_2.2" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_2.2" + } + parameter_number: 2 + id: 8589934597 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "zero_2" + opcode: "constant" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + } + literal { + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + u32s: 0 + } + id: 8589934598 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p.2" + opcode: "compare" + shape { + element_type: PRED + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p.2" + } + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 + comparison_direction: "EQ" + frontend_attributes { + } + comparison_type: "UNSIGNED" + statistics_viz { + } + comparison_order: "TOTAL" + } + instructions { + name: "pb.2" + opcode: "broadcast" + shape { + element_type: PRED + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "pb.2" + } + id: 8589934593 + operand_ids: 8589934596 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_1.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_1.1" + } + parameter_number: 1 + id: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "data.2" + opcode: "select" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "data.2" + } + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + parameter_names: "param_1.1" + parameter_names: "param_2.2" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "main" + instructions { + name: "c0_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 1 + u32s: 2 + u32s: 3 + u32s: 4 + u32s: 5 + u32s: 6 + u32s: 7 + u32s: 8 + } + id: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "c1_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 10 + u32s: 11 + u32s: 12 + u32s: 13 + u32s: 14 + u32s: 15 + u32s: 16 + u32s: 17 + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + id: 4294967296 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_select_fusion" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_select_fusion" + } + fusion_kind: "kLoop" + id: 4294967299 + operand_ids: 4294967297 + operand_ids: 4294967298 + operand_ids: 4294967296 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ars.3" + opcode: "reduce-scatter" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "reduce-scatter-start" + } + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + id: 1 + root_id: 4294967300 + } + host_program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + entry_computation_id: 1 + schedule { + sequences { + key: 0 + value { + instruction_ids: 1 + instruction_ids: 0 + instruction_ids: 2 + } + } + sequences { + key: 1 + value { + instruction_ids: 4294967296 + instruction_ids: 4294967298 + instruction_ids: 4294967297 + instruction_ids: 4294967299 + instruction_ids: 4294967300 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "c36d8870be58d238935bebefbbc5b827" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + size: 4 + defined_at { + } + } + logical_buffers { + id: 1 + size: 4 + defined_at { + instruction_id: 1 + } + } + logical_buffers { + id: 2 + size: 4 + defined_at { + instruction_id: 2 + } + } + logical_buffers { + id: 10 + size: 32 + defined_at { + instruction_id: 4294967297 + } + } + logical_buffers { + id: 11 + size: 32 + defined_at { + instruction_id: 4294967298 + } + } + logical_buffers { + id: 12 + size: 4 + defined_at { + instruction_id: 4294967296 + } + } + logical_buffers { + id: 13 + size: 32 + defined_at { + instruction_id: 4294967299 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + instruction_id: 4294967300 + } + } + buffer_allocations { + size: 32 + assigned { + logical_buffer_id: 10 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 1 + size: 32 + assigned { + logical_buffer_id: 11 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 12 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 14 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 3 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 2 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 4 + size: 4 + is_thread_local: true + assigned { + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 5 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 1 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 6 + size: 32 + assigned { + logical_buffer_id: 13 + size: 32 + element_type: U32 + } + } + heap_simulator_traces { + events { + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + events { + kind: FREE + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 10 + } + peak_buffers { + index: 1 + logical_buffer_ids: 11 + } + peak_buffers { + index: 2 + logical_buffer_ids: 14 + } + peak_buffers { + index: 3 + logical_buffer_ids: 2 + } + peak_buffers { + index: 4 + logical_buffer_ids: 0 + } + peak_buffers { + index: 5 + logical_buffer_ids: 1 + } + peak_buffers { + index: 6 + logical_buffer_ids: 13 + } + } + binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000p\013\000\000\000\000\000\000\360\006\000\000\000\000\000\000\nd\000\006@\0008\000\004\000@\000\022\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000.nv.merc.debug_frame\000.nv.merc.rela.nv.constant.pic\000.nv.merc.nv.constant.pic\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant.pic\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000.nv.constant4\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\013\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\021\000\014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\021\000\014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\002\000\001\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\254\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\035 \014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\035 \014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\\\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\001\000\000\000\000\000\000J\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\007\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\203\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\260\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\004\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000r\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000g\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\367\000\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\0000\000\000\000\000\000\000\000\021\000\000\000\017\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\001\000\000}\000\000p\002\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000 \005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\000\000\000\000\000\000\000\020\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + buffer_allocations { + values { + size: 32 + is_constant: true + } + values { + index: 1 + size: 32 + is_constant: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + } + values { + index: 3 + size: 4 + is_thread_local: true + } + values { + index: 4 + size: 4 + is_thread_local: true + } + values { + index: 5 + size: 4 + is_thread_local: true + } + values { + index: 6 + size: 32 + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + constants { + symbol_name: "buffer_for_c1_1" + allocation_index: 1 + } + constants { + symbol_name: "buffer_for_c0_1" + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_select_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 32 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 32 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + written: false + written: false + written: false + written: true + custom_kernel { + name: "loop_select_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\030\025\000\000\000\000\000\000\230\017\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000.nv.capmerc.text.loop_select_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_select_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000^\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\200\003 \000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\344\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000.\016\000\202{\004\377\000\340\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\346\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\342\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\343\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\021\000\000@\200\271\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\004\n\370\000\004\000\000\000A\002\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\203\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\030\000\000\000\002\"\016\006\000\000B\000\000\000\001\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\002\"\016\006\000\000B\000\000\000A\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\0028\0162\370\000@\021\000\000\000\000\202\001\n\000\000\002\001@\001\000\000\000\000\000\000\000\000\000\000\000P\010\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\004\020\000\000\000\004\220\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\001\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\002\000\000\000\000\000\000\177\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\305\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\370\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\254\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\005\000\000\000\000\000\000\204\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\272\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\365\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\343\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\021\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\240\014\000\000\000\000\000\000\016\001\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000V\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\256\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000k\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \016\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000|\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000D\016\000\000\000\000\000\000\224\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\272\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000\240\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 4 + kernel_name: "loop_select_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 2 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 3 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 8 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "ars.3" + thunk_id: 3 + } + reduce_scatter_thunk { + collective_config { + operand_element_type: U32 + } + buffers { + element_count: 8 + source_buffer { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + reduction_kind: REDUCTION_KIND_SUM + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/AllGatherMixedTypes.pbtxt new file mode 100644 index 00000000000000..4ea9837395c53d --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/AllGatherMixedTypes.pbtxt @@ -0,0 +1,2190 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "test_computation" + computations { + name: "fused_transpose" + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.6.1" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.6.1" + } + id: 4294967297 + operand_ids: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.1.1" + opcode: "transpose" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.1.1" + } + dimensions: 1 + dimensions: 0 + id: 4294967296 + operand_ids: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + } + id: 1 + root_id: 4294967296 + is_fusion_computation: true + } + computations { + name: "fused_transpose.1" + instructions { + name: "param_0.3" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.3" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.4.1" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.4.1" + } + id: 8589934593 + operand_ids: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.3" + opcode: "transpose" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.3" + } + dimensions: 1 + dimensions: 0 + id: 8589934592 + operand_ids: 8589934593 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.3" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "fused_broadcast" + instructions { + name: "param_0.5" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.5" + } + id: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p1.1.1" + opcode: "convert" + shape { + element_type: F32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p1.1.1" + } + id: 12884901889 + operand_ids: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.13.1" + opcode: "broadcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.13.1" + } + id: 12884901888 + operand_ids: 12884901889 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.5" + } + id: 3 + root_id: 12884901888 + is_fusion_computation: true + } + computations { + name: "wrapped_broadcast_computation" + instructions { + name: "param_0.7" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.7" + } + id: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.12.1" + opcode: "broadcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.12.1" + } + id: 21474836480 + operand_ids: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.7" + } + id: 5 + root_id: 21474836480 + is_fusion_computation: true + } + computations { + name: "test_computation" + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "wrapped_broadcast" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "wrapped_broadcast" + } + fusion_kind: "kLoop" + id: 9 + operand_ids: 0 + called_computation_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_broadcast_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_broadcast_fusion" + } + fusion_kind: "kLoop" + id: 8 + operand_ids: 0 + called_computation_ids: 3 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "allgather.3" + opcode: "all-gather" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "all-gather-start" + } + dimensions: 1 + id: 10 + operand_ids: 9 + operand_ids: 8 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + instructions { + name: "ag0.1" + opcode: "get-tuple-element" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag0.1" + } + id: 1 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion.1" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion.1" + } + fusion_kind: "kLoop" + id: 7 + operand_ids: 1 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.24" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.24" + } + id: 4 + operand_ids: 7 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ag1.1" + opcode: "get-tuple-element" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag1.1" + } + tuple_index: 1 + id: 2 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion" + } + fusion_kind: "kLoop" + id: 6 + operand_ids: 2 + called_computation_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.1.0" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.1.0" + } + id: 5 + operand_ids: 6 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "out.1" + opcode: "tuple" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "out.1" + } + id: 3 + operand_ids: 4 + operand_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + root_id: 3 + } + host_program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + schedule { + sequences { + key: 0 + value { + instruction_ids: 0 + instruction_ids: 8 + instruction_ids: 9 + instruction_ids: 10 + instruction_ids: 1 + instruction_ids: 7 + instruction_ids: 2 + instruction_ids: 6 + instruction_ids: 4 + instruction_ids: 5 + instruction_ids: 3 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "7b2e4f79c5524b5cbc566da178aa7fda" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + id: 9 + size: 4 + defined_at { + } + } + logical_buffers { + id: 10 + size: 8 + defined_at { + instruction_id: 9 + } + } + logical_buffers { + id: 11 + size: 8 + defined_at { + instruction_id: 8 + } + } + logical_buffers { + id: 12 + size: 16 + defined_at { + instruction_id: 10 + } + } + logical_buffers { + id: 13 + size: 16 + defined_at { + shape_index: 0 + instruction_id: 10 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + shape_index: 1 + instruction_id: 10 + } + } + logical_buffers { + id: 15 + size: 16 + defined_at { + instruction_id: 7 + } + } + logical_buffers { + id: 16 + size: 16 + defined_at { + instruction_id: 6 + } + } + logical_buffers { + id: 17 + size: 16 + defined_at { + instruction_id: 3 + } + } + buffer_allocations { + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 10 + size: 8 + element_type: U32 + } + assigned { + logical_buffer_id: 15 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 1 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 9 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 13 + size: 16 + element_type: U32 + } + assigned { + logical_buffer_id: 16 + size: 16 + element_type: F32 + } + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 17 + size: 16 + element_type: TUPLE + } + is_tuple: true + } + buffer_allocations { + index: 3 + size: 520 + assigned { + logical_buffer_id: 11 + offset: 512 + size: 8 + element_type: F32 + } + assigned { + logical_buffer_id: 12 + size: 16 + element_type: TUPLE + } + assigned { + logical_buffer_id: 14 + offset: 256 + size: 16 + element_type: F32 + } + is_tuple: true + } + heap_simulator_traces { + events { + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + events { + kind: FREE + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + kind: FREE + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 15 + } + peak_buffers { + index: 1 + logical_buffer_ids: 16 + } + peak_buffers { + index: 2 + logical_buffer_ids: 17 + } + peak_buffers { + index: 3 + logical_buffer_ids: 12 + logical_buffer_ids: 11 + logical_buffer_ids: 14 + } + } + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + buffer_allocations { + values { + size: 16 + maybe_live_out: true + } + values { + index: 1 + size: 16 + maybe_live_out: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + is_tuple: true + } + values { + index: 3 + size: 520 + is_tuple: true + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + output_info_map { + shape_index { + indexes: 0 + } + output_info { + } + } + output_info_map { + shape_index { + indexes: 1 + } + output_info { + allocation_index: 1 + } + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_broadcast_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_broadcast_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000.nv.capmerc.text.loop_broadcast_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_broadcast_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\014\000\000@\300\005\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\001\000\000\000\000\000\000\000\000\000\000\000\320\006\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004\020\000\000\000\004\360\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_broadcast_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "wrapped_broadcast" + thunk_id: 3 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "wrapped_broadcast" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000.nv.capmerc.text.wrapped_broadcast\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.wrapped_broadcast\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000Z\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\013\000\000@\300\002\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\007\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\004\020\000\000\000\004\340\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\340\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\342\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\002\000\000\000\000\000\000z\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\302\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000i\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\267\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\340\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000w\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\232\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\264\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "wrapped_broadcast" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "allgather.3" + thunk_id: 4 + } + all_gather_thunk { + collective_config { + operand_element_type: U32 + operand_element_type: F32 + } + buffers { + element_count: 2 + source_buffer { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + buffers { + element_count: 2 + source_buffer { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion.1" + thunk_id: 5 + } + custom_kernel_thunk { + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion_1" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000.nv.capmerc.text.loop_transpose_fusion_1\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion_1\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000r\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\006\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\002\000\000\000\000\000\000\230\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\004\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\324\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\311\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\362\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\204\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\225\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\330\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion_1" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion" + thunk_id: 6 + } + custom_kernel_thunk { + args { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000.nv.capmerc.text.loop_transpose_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/ReduceScatter.pbtxt new file mode 100644 index 00000000000000..886fabba15a64e --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/ReduceScatter.pbtxt @@ -0,0 +1,1624 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "main" + computations { + name: "add" + instructions { + name: "lhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "lhs" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "rhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "rhs" + } + parameter_number: 1 + id: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "add.1" + opcode: "add" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "add.1" + } + id: 2 + operand_ids: 0 + operand_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameter_names: "lhs" + parameter_names: "rhs" + } + root_id: 2 + } + computations { + name: "fused_select" + instructions { + name: "param_2.2" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_2.2" + } + parameter_number: 2 + id: 8589934597 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "zero_2" + opcode: "constant" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + } + literal { + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + u32s: 0 + } + id: 8589934598 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p.2" + opcode: "compare" + shape { + element_type: PRED + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p.2" + } + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 + comparison_direction: "EQ" + frontend_attributes { + } + statistics_viz { + } + comparison_order: "TOTAL" + } + instructions { + name: "pb.2" + opcode: "broadcast" + shape { + element_type: PRED + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "pb.2" + } + id: 8589934593 + operand_ids: 8589934596 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_1.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_1.1" + } + parameter_number: 1 + id: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "data.2" + opcode: "select" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "data.2" + } + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + parameter_names: "param_1.1" + parameter_names: "param_2.2" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "main" + instructions { + name: "c0_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 1 + u32s: 2 + u32s: 3 + u32s: 4 + u32s: 5 + u32s: 6 + u32s: 7 + u32s: 8 + } + id: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "c1_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 10 + u32s: 11 + u32s: 12 + u32s: 13 + u32s: 14 + u32s: 15 + u32s: 16 + u32s: 17 + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + id: 4294967296 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_select_fusion" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_select_fusion" + } + fusion_kind: "kLoop" + id: 4294967299 + operand_ids: 4294967297 + operand_ids: 4294967298 + operand_ids: 4294967296 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ars.3" + opcode: "reduce-scatter" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "reduce-scatter-start" + } + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + id: 1 + root_id: 4294967300 + } + host_program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + entry_computation_id: 1 + schedule { + sequences { + key: 0 + value { + instruction_ids: 1 + instruction_ids: 0 + instruction_ids: 2 + } + } + sequences { + key: 1 + value { + instruction_ids: 4294967296 + instruction_ids: 4294967298 + instruction_ids: 4294967297 + instruction_ids: 4294967299 + instruction_ids: 4294967300 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "c36d8870be58d238935bebefbbc5b827" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + size: 4 + defined_at { + } + } + logical_buffers { + id: 1 + size: 4 + defined_at { + instruction_id: 1 + } + } + logical_buffers { + id: 2 + size: 4 + defined_at { + instruction_id: 2 + } + } + logical_buffers { + id: 10 + size: 32 + defined_at { + instruction_id: 4294967297 + } + } + logical_buffers { + id: 11 + size: 32 + defined_at { + instruction_id: 4294967298 + } + } + logical_buffers { + id: 12 + size: 4 + defined_at { + instruction_id: 4294967296 + } + } + logical_buffers { + id: 13 + size: 32 + defined_at { + instruction_id: 4294967299 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + instruction_id: 4294967300 + } + } + buffer_allocations { + size: 32 + assigned { + logical_buffer_id: 10 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 1 + size: 32 + assigned { + logical_buffer_id: 11 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 12 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 14 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 3 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 2 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 4 + size: 4 + is_thread_local: true + assigned { + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 5 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 1 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 6 + size: 32 + assigned { + logical_buffer_id: 13 + size: 32 + element_type: U32 + } + } + heap_simulator_traces { + events { + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + events { + kind: FREE + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 10 + } + peak_buffers { + index: 1 + logical_buffer_ids: 11 + } + peak_buffers { + index: 2 + logical_buffer_ids: 14 + } + peak_buffers { + index: 3 + logical_buffer_ids: 2 + } + peak_buffers { + index: 4 + logical_buffer_ids: 0 + } + peak_buffers { + index: 5 + logical_buffer_ids: 1 + } + peak_buffers { + index: 6 + logical_buffer_ids: 13 + } + } + binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000p\013\000\000\000\000\000\000\360\006\000\000\000\000\000\000\nd\000\006@\0008\000\004\000@\000\022\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000.nv.merc.debug_frame\000.nv.merc.rela.nv.constant.pic\000.nv.merc.nv.constant.pic\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant.pic\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000.nv.constant4\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\013\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\021\000\014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\021\000\014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\002\000\001\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\254\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\035 \014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\035 \014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\\\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\001\000\000\000\000\000\000J\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\007\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\203\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\260\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\004\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000r\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000g\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\367\000\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\0000\000\000\000\000\000\000\000\021\000\000\000\017\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\001\000\000}\000\000p\002\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000 \005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\000\000\000\000\000\000\000\020\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + buffer_allocations { + values { + size: 32 + is_constant: true + } + values { + index: 1 + size: 32 + is_constant: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + } + values { + index: 3 + size: 4 + is_thread_local: true + } + values { + index: 4 + size: 4 + is_thread_local: true + } + values { + index: 5 + size: 4 + is_thread_local: true + } + values { + index: 6 + size: 32 + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + constants { + symbol_name: "buffer_for_c1_1" + allocation_index: 1 + } + constants { + symbol_name: "buffer_for_c0_1" + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_select_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 32 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 32 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + written: false + written: false + written: false + written: true + custom_kernel { + name: "loop_select_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\030\025\000\000\000\000\000\000\230\017\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000.nv.capmerc.text.loop_select_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_select_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000^\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\200\003 \000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\344\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000.\016\000\202{\004\377\000\340\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\346\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\342\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\343\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\021\000\000@\200\271\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\004\n\370\000\004\000\000\000A\002\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\203\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\030\000\000\000\002\"\016\006\000\000B\000\000\000\001\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\002\"\016\006\000\000B\000\000\000A\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\0028\0162\370\000@\021\000\000\000\000\202\001\n\000\000\002\001@\001\000\000\000\000\000\000\000\000\000\000\000P\010\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\004\020\000\000\000\004\220\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\001\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\002\000\000\000\000\000\000\177\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\305\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\370\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\254\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\005\000\000\000\000\000\000\204\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\272\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\365\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\343\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\021\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\240\014\000\000\000\000\000\000\016\001\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000V\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\256\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000k\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \016\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000|\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000D\016\000\000\000\000\000\000\224\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\272\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000\240\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 4 + kernel_name: "loop_select_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 2 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 3 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 8 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "ars.3" + thunk_id: 3 + } + reduce_scatter_thunk { + collective_config { + operand_element_type: U32 + } + buffers { + element_count: 8 + source_buffer { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + reduction_kind: REDUCTION_KIND_SUM + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/AllGatherMixedTypes.pbtxt similarity index 60% rename from third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/AllGatherMixedTypes.pbtxt rename to third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/AllGatherMixedTypes.pbtxt index de78ec91592c00..cb7e12a41013aa 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/AllGatherMixedTypes.pbtxt +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/AllGatherMixedTypes.pbtxt @@ -1,9 +1,9 @@ -# proto-file: third_party/tensorflow/compiler/xla/tools/split_proto/split_proto_cli.proto -# proto-message: xla.split_proto_cli.DeserializedSplitExecutableAndOptions - +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable executable_and_options { compile_options { executable_build_options { + device_ordinal: -1 result_layout { element_type: TUPLE tuple_shapes { @@ -39,7 +39,7 @@ executable_and_options { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -65,7 +65,6 @@ executable_and_options { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -78,14 +77,10 @@ executable_and_options { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -129,6 +124,7 @@ executable_and_options { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -137,7 +133,7 @@ executable_and_options { xla_cmd_buffer_trace_cache_size: 16 xla_gpu_temp_buffer_use_separate_color: false xla_syntax_sugar_async_ops: false - xla_gpu_autotune_gemm_rtol: 0.1 + xla_gpu_autotune_gemm_rtol: 0.01 xla_enable_command_buffers_during_profiling: false xla_gpu_cudnn_gemm_max_plans: 5 xla_cpu_parallel_codegen_split_count: 32 @@ -199,7 +195,6 @@ executable_and_options { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -208,7 +203,6 @@ executable_and_options { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -219,14 +213,13 @@ executable_and_options { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -236,21 +229,32 @@ executable_and_options { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" xla_enable_hlo_sharding_v3: false xla_gpu_log_minmax: false @@ -261,6 +265,7 @@ executable_and_options { xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -280,12 +285,12 @@ executable_and_options { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -294,6 +299,35 @@ executable_and_options { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" } num_replicas: 2 num_partitions: 1 @@ -311,9 +345,9 @@ executable_and_options { allow_spmd_sharding_propagation_to_parameters: false process_count: 1 memory_fitting_level: EFFORT_O2 - slice_size: 2 } } + pjrt_client_name: "PjRtStreamExecutorClient" } gpu_executable { hlo_module_with_config { @@ -338,6 +372,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.1" } id: 4294967298 frontend_attributes { @@ -361,6 +396,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.6.1" } id: 4294967297 operand_ids: 4294967298 @@ -385,6 +421,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "transpose.1.1" } dimensions: 1 dimensions: 0 @@ -444,6 +481,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.3" } id: 8589934594 frontend_attributes { @@ -467,6 +505,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.4.1" } id: 8589934593 operand_ids: 8589934594 @@ -491,6 +530,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "transpose.3" } dimensions: 1 dimensions: 0 @@ -544,6 +584,7 @@ gpu_executable { } } metadata { + scheduling_name: "param_0.5" } id: 12884901890 frontend_attributes { @@ -561,6 +602,7 @@ gpu_executable { } } metadata { + scheduling_name: "p1.1.1" } id: 12884901889 operand_ids: 12884901890 @@ -585,6 +627,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "broadcast.13.1" } id: 12884901888 operand_ids: 12884901889 @@ -618,156 +661,6 @@ gpu_executable { root_id: 12884901888 is_fusion_computation: true } - computations { - name: "async_computation" - instructions { - name: "param_0.6" - opcode: "parameter" - shape { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - metadata { - } - id: 17179869184 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "param_1" - opcode: "parameter" - shape { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - metadata { - } - parameter_number: 1 - id: 17179869185 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "allgather.2" - opcode: "all-gather" - shape { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - metadata { - } - dimensions: 1 - id: 17179869186 - operand_ids: 17179869184 - operand_ids: 17179869185 - frontend_attributes { - } - statistics_viz { - } - collective_device_list { - } - } - program_shape { - parameters { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - parameters { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - result { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - parameter_names: "param_0.6" - parameter_names: "param_1" - } - id: 4 - root_id: 17179869186 - } computations { name: "wrapped_broadcast_computation" instructions { @@ -780,6 +673,7 @@ gpu_executable { } } metadata { + scheduling_name: "param_0.7" } id: 21474836481 frontend_attributes { @@ -803,6 +697,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "broadcast.12.1" } id: 21474836480 operand_ids: 21474836481 @@ -848,6 +743,7 @@ gpu_executable { } } metadata { + scheduling_name: "id.1" } frontend_attributes { } @@ -870,18 +766,16 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "wrapped_broadcast" } fusion_kind: "kLoop" - id: 11 + id: 9 operand_ids: 0 called_computation_ids: 5 frontend_attributes { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "loop_broadcast_fusion" @@ -899,6 +793,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_broadcast_fusion" } fusion_kind: "kLoop" id: 8 @@ -908,87 +803,10 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } - } - instructions { - name: "all-gather-start" - opcode: "async-start" - shape { - element_type: TUPLE - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - } - metadata { - } - id: 9 - operand_ids: 11 - operand_ids: 8 - called_computation_ids: 4 - frontend_attributes { - } - statistics_viz { - } - backend_config_payload { - id: 1 - } } instructions { - name: "all-gather-done" - opcode: "async-done" + name: "allgather.3" + opcode: "all-gather" shape { element_type: TUPLE tuple_shapes { @@ -1017,13 +835,21 @@ gpu_executable { } } metadata { + scheduling_name: "all-gather-start" } + dimensions: 1 id: 10 operand_ids: 9 + operand_ids: 8 frontend_attributes { } statistics_viz { } + collective_device_list { + } + backend_config_payload { + id: 0 + } } instructions { name: "ag0.1" @@ -1041,6 +867,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "ag0.1" } id: 1 operand_ids: 10 @@ -1065,6 +892,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_transpose_fusion.1" } fusion_kind: "kLoop" id: 7 @@ -1074,9 +902,6 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "bitcast.24" @@ -1091,6 +916,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.24" } id: 4 operand_ids: 7 @@ -1115,6 +941,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "ag1.1" } tuple_index: 1 id: 2 @@ -1140,6 +967,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_transpose_fusion" } fusion_kind: "kLoop" id: 6 @@ -1149,9 +977,6 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "bitcast.1.0" @@ -1166,6 +991,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.1.0" } id: 5 operand_ids: 6 @@ -1199,6 +1025,7 @@ gpu_executable { } } metadata { + scheduling_name: "out.1" } id: 3 operand_ids: 4 @@ -1262,26 +1089,17 @@ gpu_executable { value { instruction_ids: 0 instruction_ids: 8 - instruction_ids: 11 instruction_ids: 9 instruction_ids: 10 - instruction_ids: 2 instruction_ids: 1 instruction_ids: 7 - instruction_ids: 4 + instruction_ids: 2 instruction_ids: 6 + instruction_ids: 4 instruction_ids: 5 instruction_ids: 3 } } - sequences { - key: 4 - value { - instruction_ids: 17179869185 - instruction_ids: 17179869184 - instruction_ids: 17179869186 - } - } } input_output_alias { } @@ -1298,10 +1116,9 @@ gpu_executable { frontend_attributes { map { key: "fingerprint_before_lhs" - value: "5d6b0e05c28f8e5fa1780dcf27c57d60" + value: "7b2e4f79c5524b5cbc566da178aa7fda" } } - payloads: "{\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"native_emitter_backend_config\":{\"type\":\"NATIVE_EMITTER_TYPE_INVALID\",\"unroll_factor\":0},\"operation_queue_id\":\"0\",\"reification_cost\":[]}" payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" } config { @@ -1345,7 +1162,7 @@ gpu_executable { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -1371,7 +1188,6 @@ gpu_executable { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -1384,14 +1200,10 @@ gpu_executable { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -1435,6 +1247,7 @@ gpu_executable { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -1443,7 +1256,7 @@ gpu_executable { xla_cmd_buffer_trace_cache_size: 16 xla_gpu_temp_buffer_use_separate_color: false xla_syntax_sugar_async_ops: false - xla_gpu_autotune_gemm_rtol: 0.1 + xla_gpu_autotune_gemm_rtol: 0.01 xla_enable_command_buffers_during_profiling: false xla_gpu_cudnn_gemm_max_plans: 5 xla_cpu_parallel_codegen_split_count: 32 @@ -1505,7 +1318,6 @@ gpu_executable { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -1514,7 +1326,6 @@ gpu_executable { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -1525,14 +1336,13 @@ gpu_executable { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -1542,21 +1352,32 @@ gpu_executable { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" xla_enable_hlo_sharding_v3: false xla_gpu_log_minmax: false @@ -1567,6 +1388,7 @@ gpu_executable { xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -1586,12 +1408,12 @@ gpu_executable { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -1600,6 +1422,35 @@ gpu_executable { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" } static_device_assignment { replica_count: 2 @@ -1620,95 +1471,65 @@ gpu_executable { } } buffer_assignment { - logical_buffers { - id: 7 - size: 16 - defined_at { - instruction_id: 17179869186 - } - } - logical_buffers { - id: 8 - size: 16 - defined_at { - shape_index: 0 - instruction_id: 17179869186 - } - } logical_buffers { id: 9 - size: 16 - defined_at { - shape_index: 1 - instruction_id: 17179869186 - } - } - logical_buffers { - id: 12 size: 4 defined_at { } } logical_buffers { - id: 13 + id: 10 size: 8 defined_at { - instruction_id: 11 + instruction_id: 9 } } logical_buffers { - id: 14 + id: 11 size: 8 defined_at { instruction_id: 8 } } logical_buffers { - id: 15 + id: 12 size: 16 defined_at { - instruction_id: 9 + instruction_id: 10 } } logical_buffers { - id: 16 + id: 13 size: 16 defined_at { shape_index: 0 - instruction_id: 9 + instruction_id: 10 } } logical_buffers { - id: 17 + id: 14 size: 16 defined_at { shape_index: 1 - instruction_id: 9 - } - } - logical_buffers { - id: 18 - size: 16 - defined_at { instruction_id: 10 } } logical_buffers { - id: 19 + id: 15 size: 16 defined_at { instruction_id: 7 } } logical_buffers { - id: 20 + id: 16 size: 16 defined_at { instruction_id: 6 } } logical_buffers { - id: 21 + id: 17 size: 16 defined_at { instruction_id: 3 @@ -1718,12 +1539,12 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 13 + logical_buffer_id: 10 size: 8 element_type: U32 } assigned { - logical_buffer_id: 19 + logical_buffer_id: 15 size: 16 element_type: U32 } @@ -1733,17 +1554,17 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 8 - size: 16 + logical_buffer_id: 9 + size: 4 element_type: U32 } assigned { - logical_buffer_id: 12 - size: 4 + logical_buffer_id: 13 + size: 16 element_type: U32 } assigned { - logical_buffer_id: 20 + logical_buffer_id: 16 size: 16 element_type: F32 } @@ -1753,7 +1574,7 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 21 + logical_buffer_id: 17 size: 16 element_type: TUPLE } @@ -1761,102 +1582,74 @@ gpu_executable { } buffer_allocations { index: 3 - size: 1544 + size: 520 assigned { - logical_buffer_id: 7 - offset: 1024 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 9 - offset: 1280 - size: 16 - element_type: F32 - } - assigned { - logical_buffer_id: 14 - offset: 1536 + logical_buffer_id: 11 + offset: 512 size: 8 element_type: F32 } assigned { - logical_buffer_id: 15 + logical_buffer_id: 12 size: 16 element_type: TUPLE } assigned { - logical_buffer_id: 16 + logical_buffer_id: 14 offset: 256 size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 17 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 18 - offset: 768 - size: 16 - element_type: TUPLE + element_type: F32 } is_tuple: true } heap_simulator_traces { events { - buffer_id: 14 + buffer_id: 11 computation_name: "test_computation" instruction_name: "loop_broadcast_fusion" } events { - buffer_id: 9 - computation_name: "async_computation" - instruction_name: "allgather.2" + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" } events { kind: FREE - buffer_id: 14 + buffer_id: 11 computation_name: "test_computation" instruction_name: "loop_broadcast_fusion" } events { kind: FREE - buffer_id: 9 - computation_name: "async_computation" - instruction_name: "allgather.2" + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" } whole_module_simulation: true buffer_allocation_index: 3 } peak_buffers { - logical_buffer_ids: 19 + logical_buffer_ids: 15 } peak_buffers { index: 1 - logical_buffer_ids: 20 + logical_buffer_ids: 16 } peak_buffers { index: 2 - logical_buffer_ids: 21 + logical_buffer_ids: 17 } peak_buffers { index: 3 - logical_buffer_ids: 15 - logical_buffer_ids: 16 - logical_buffer_ids: 17 - logical_buffer_ids: 18 - logical_buffer_ids: 7 - logical_buffer_ids: 9 + logical_buffer_ids: 12 + logical_buffer_ids: 11 logical_buffer_ids: 14 } } gpu_compute_capability { cuda_compute_capability { - major: 6 - feature_extension: NONE + major: 9 + feature_extension: ACCELERATED_FEATURES } } module_name: "test" @@ -1887,92 +1680,21 @@ gpu_executable { values { size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 13 - size: 8 - element_type: U32 - } - assigned { - logical_buffer_id: 19 - size: 16 - element_type: U32 - } } values { index: 1 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 8 - size: 16 - element_type: U32 - } - assigned { - logical_buffer_id: 12 - size: 4 - element_type: U32 - } - assigned { - logical_buffer_id: 20 - size: 16 - element_type: F32 - } } values { index: 2 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 21 - size: 16 - element_type: TUPLE - } is_tuple: true } values { index: 3 - size: 1544 - assigned { - logical_buffer_id: 7 - offset: 1024 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 9 - offset: 1280 - size: 16 - element_type: F32 - } - assigned { - logical_buffer_id: 14 - offset: 1536 - size: 8 - element_type: F32 - } - assigned { - logical_buffer_id: 15 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 17 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 18 - offset: 768 - size: 16 - element_type: TUPLE - } + size: 520 is_tuple: true } } @@ -2032,7 +1754,7 @@ gpu_executable { } args { slice { - offset: 1536 + offset: 512 size: 8 buffer_allocation_index: 3 element_type: F32 @@ -2056,7 +1778,7 @@ gpu_executable { name: "loop_broadcast_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\010\000\000\000\000\000\000\300\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000.rel.nv.constant0.loop_broadcast_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.rel.nv.constant0.loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\266\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\007\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000x\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000\007\005\200\007\230L\004\000\027\002\000\000\310\360\342\007\340\367\000\304?\010\003\000\027\005\200\007\230L\002\002\007\000\000\240\324\356\000\004\347\001\000\000(8\346\007\340\374\000t\034\004\004\004\'\005\000\201\030L\005\0007\005\000\010\020L\006\n\'\000\000\000\270\\\361\017\200\374\000\374\037\000\006\004\007\000\000 \334\356\000\017\007\000\000\000\260P\017\000\007\000\000\000\000\343\377\007\000\374\000\200\037\000\017\000\007\377\377\017@\342\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\005\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000E\001\000\000\000\000\000\000\033\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\366\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\007 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\275\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\336\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000*\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000@\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\204\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\206\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\024\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\211\001\000\000\000\000\000\000d\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\244\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\324\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\302\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_broadcast_fusion" @@ -2089,6 +1811,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2136,7 +1859,7 @@ gpu_executable { name: "wrapped_broadcast" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\340\007\000\000\000\000\000\000`\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000.rel.nv.constant0.wrapped_broadcast\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.rel.nv.constant0.wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\246\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\361\000\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000p\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000\007\005\200\007\230L\004\000\027\002\000\000\310\360\342\007\340\367\000\304?\010\003\000\027\005\200\007\230L\002\002\007\000\000\240\324\356\000\004\347\001\000\000(8\346\007\340\374\000\304\037\004\004\004\'\005\000\201\030L\005\0007\005\000\010\020L\002\004\007\000\000 \334\356\352\007\340\377\000\374\037\000\000\017\007\000\000\000\260P\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\361\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0001\001\000\000\000\000\000\000\003\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000J\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\354\002\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\342\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000h\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\202\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\004\000\000\000\000\000\000\200\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000\340\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000x\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\350\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\340\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\261\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\036\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\0000\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\204\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\206\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\004\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000u\001\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\230\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000J\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000t\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\266\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\230\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "wrapped_broadcast" @@ -2169,13 +1892,14 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } } thunks { thunk_info { - profile_annotation: "allgather.2" + profile_annotation: "allgather.3" thunk_id: 4 } all_gather_thunk { @@ -2227,7 +1951,7 @@ gpu_executable { element_count: 2 source_buffer { slice { - offset: 1536 + offset: 512 size: 8 buffer_allocation_index: 3 element_type: F32 @@ -2247,7 +1971,7 @@ gpu_executable { } destination_buffer { slice { - offset: 1280 + offset: 256 size: 16 buffer_allocation_index: 3 element_type: F32 @@ -2317,7 +2041,7 @@ gpu_executable { name: "loop_transpose_fusion_1" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000`\010\000\000\000\000\000\000\340\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000.rel.nv.constant0.loop_transpose_fusion_1\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.rel.nv.constant0.loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\340\343\000\230?\000\001\000\207\000\200\007\230L\005\000\027\002\000\000\310\360\000\005W\000\000\000\020\\\366\007\300\376\000\304\037\010\000\000\'\000\000\000\000\004\000\005\007\360\201\177\337[\003\000\347\001\000\000(8\346\007@\374\000\304\036\000\002\000\007\005\000\201\030L\003\003\027\005\000\010\020L\002\002\007\000\000\240\324\356\361\007\302\374\000\234\037\000\000\005\347\001\000\000(8\004\005\'\005\000\201\030L\005\0007\005\000\010\020L\361\007\341\377\000\374\037\000\002\004\007\000\000 \334\356\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\017\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000O\001\000\000\000\000\000\000\'\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000P\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\001\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\224\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\270\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000`\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\270\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\002\000\000\000\000\000\000(\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000`\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\224\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\303\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000H\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\204\000\000\000\n\000\000\000b\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\206\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\016\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\034\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\223\001\000\000\000\000\000\000n\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\252\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000P\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\264\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\310\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\224\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_transpose_fusion_1" @@ -2350,6 +2074,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2362,7 +2087,7 @@ gpu_executable { custom_kernel_thunk { args { slice { - offset: 1280 + offset: 256 size: 16 buffer_allocation_index: 3 element_type: F32 @@ -2405,7 +2130,7 @@ gpu_executable { name: "loop_transpose_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\010\000\000\000\000\000\000\300\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000.rel.nv.constant0.loop_transpose_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.rel.nv.constant0.loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\266\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\340\343\000\230?\000\001\000\207\000\200\007\230L\005\000\027\002\000\000\310\360\000\005W\000\000\000\020\\\366\007\300\376\000\304\037\010\000\000\'\000\000\000\000\004\000\005\007\360\201\177\337[\003\000\347\001\000\000(8\346\007@\374\000\304\036\000\002\000\007\005\000\201\030L\003\003\027\005\000\010\020L\002\002\007\000\000\240\324\356\361\007\302\374\000\234\037\000\000\005\347\001\000\000(8\004\005\'\005\000\201\030L\005\0007\005\000\010\020L\361\007\341\377\000\374\037\000\002\004\007\000\000 \334\356\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\005\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000E\001\000\000\000\000\000\000\033\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\366\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\275\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\336\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000*\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000@\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\204\000\000\000\n\000\000\000b\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\206\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\016\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\024\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\211\001\000\000\000\000\000\000d\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\244\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\324\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\302\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_transpose_fusion" @@ -2438,6 +2163,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2446,9 +2172,9 @@ gpu_executable { platform_name: "CUDA" cuda_platform_version { cuda_toolkit_version: "12.8.0" - cudnn_version: "9.10.0" + cudnn_version: "9.10.2" cub_version: "3.1.2" } } - buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 1.5KiB\n preallocated temp fragmentation: 240B (15.54%)\n total allocation: 1.6KiB\n total fragmentation: 1.4KiB (91.96%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: async-done\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: (u32[2,1], f32[2,1])\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: ((u32[2,1], f32[2,1]), (u32[2,2], f32[2,2]))\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 8:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 9:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 10:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" } diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/ReduceScatter.pbtxt similarity index 67% rename from third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/ReduceScatter.pbtxt rename to third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/ReduceScatter.pbtxt index 24fe07b3ad1846..89dadf3e45ea87 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/ReduceScatter.pbtxt +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/ReduceScatter.pbtxt @@ -1,9 +1,9 @@ -# proto-file: third_party/tensorflow/compiler/xla/tools/split_proto/split_proto_cli.proto -# proto-message: xla.split_proto_cli.DeserializedSplitExecutableAndOptions - +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable executable_and_options { compile_options { executable_build_options { + device_ordinal: -1 result_layout { element_type: U32 dimensions: 4 @@ -27,7 +27,7 @@ executable_and_options { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -53,7 +53,6 @@ executable_and_options { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -66,14 +65,10 @@ executable_and_options { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -117,6 +112,7 @@ executable_and_options { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -125,7 +121,7 @@ executable_and_options { xla_cmd_buffer_trace_cache_size: 16 xla_gpu_temp_buffer_use_separate_color: false xla_syntax_sugar_async_ops: false - xla_gpu_autotune_gemm_rtol: 0.1 + xla_gpu_autotune_gemm_rtol: 0.01 xla_enable_command_buffers_during_profiling: false xla_gpu_cudnn_gemm_max_plans: 5 xla_cpu_parallel_codegen_split_count: 32 @@ -187,7 +183,6 @@ executable_and_options { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -196,7 +191,6 @@ executable_and_options { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -207,14 +201,13 @@ executable_and_options { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -224,21 +217,32 @@ executable_and_options { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" xla_enable_hlo_sharding_v3: false xla_gpu_log_minmax: false @@ -249,6 +253,7 @@ executable_and_options { xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -268,12 +273,12 @@ executable_and_options { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -282,6 +287,35 @@ executable_and_options { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" } num_replicas: 2 num_partitions: 1 @@ -299,9 +333,9 @@ executable_and_options { allow_spmd_sharding_propagation_to_parameters: false process_count: 1 memory_fitting_level: EFFORT_O2 - slice_size: 2 } } + pjrt_client_name: "PjRtStreamExecutorClient" } gpu_executable { hlo_module_with_config { @@ -320,6 +354,7 @@ gpu_executable { } } metadata { + scheduling_name: "lhs" } frontend_attributes { } @@ -336,6 +371,7 @@ gpu_executable { } } metadata { + scheduling_name: "rhs" } parameter_number: 1 id: 1 @@ -354,6 +390,7 @@ gpu_executable { } } metadata { + scheduling_name: "add.1" } id: 2 operand_ids: 0 @@ -399,9 +436,10 @@ gpu_executable { } } metadata { + scheduling_name: "param_2.2" } parameter_number: 2 - id: 12884901893 + id: 8589934597 frontend_attributes { } statistics_viz { @@ -427,14 +465,14 @@ gpu_executable { } u32s: 0 } - id: 12884901894 + id: 8589934598 frontend_attributes { } statistics_viz { } } instructions { - name: "p.4" + name: "p.2" opcode: "compare" shape { element_type: PRED @@ -443,19 +481,21 @@ gpu_executable { } } metadata { + scheduling_name: "p.2" } - id: 12884901892 - operand_ids: 12884901893 - operand_ids: 12884901894 + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 comparison_direction: "EQ" frontend_attributes { } comparison_type: "UNSIGNED" statistics_viz { } + comparison_order: "TOTAL" } instructions { - name: "pb.4" + name: "pb.2" opcode: "broadcast" shape { element_type: PRED @@ -467,16 +507,17 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "pb.2" } - id: 12884901889 - operand_ids: 12884901892 + id: 8589934593 + operand_ids: 8589934596 frontend_attributes { } statistics_viz { } } instructions { - name: "param_0.4" + name: "param_0.1" opcode: "parameter" shape { element_type: U32 @@ -488,15 +529,16 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.1" } - id: 12884901890 + id: 8589934594 frontend_attributes { } statistics_viz { } } instructions { - name: "param_1.3" + name: "param_1.1" opcode: "parameter" shape { element_type: U32 @@ -508,9 +550,10 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_1.1" } parameter_number: 1 - id: 12884901891 + id: 8589934595 frontend_attributes { } statistics_viz { @@ -529,11 +572,12 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "data.2" } - id: 12884901888 - operand_ids: 12884901889 - operand_ids: 12884901890 - operand_ids: 12884901891 + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 frontend_attributes { } statistics_viz { @@ -573,85 +617,14 @@ gpu_executable { } is_dynamic_dimension: false } - parameter_names: "param_0.4" - parameter_names: "param_1.3" + parameter_names: "param_0.1" + parameter_names: "param_1.1" parameter_names: "param_2.2" } - id: 3 - root_id: 12884901888 + id: 2 + root_id: 8589934592 is_fusion_computation: true } - computations { - name: "async_computation" - instructions { - name: "param_0.5" - opcode: "parameter" - shape { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - metadata { - } - id: 21474836480 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "ars.2" - opcode: "reduce-scatter" - shape { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - metadata { - } - dimensions: 0 - id: 21474836481 - operand_ids: 21474836480 - called_computation_ids: 0 - frontend_attributes { - } - statistics_viz { - } - collective_device_list { - } - } - program_shape { - parameters { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - result { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - parameter_names: "param_0.5" - } - id: 5 - root_id: 21474836481 - } computations { name: "main" instructions { @@ -742,6 +715,7 @@ gpu_executable { } } metadata { + scheduling_name: "id.1" } id: 4294967296 frontend_attributes { @@ -762,64 +736,22 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_select_fusion" } fusion_kind: "kLoop" id: 4294967299 operand_ids: 4294967297 operand_ids: 4294967298 operand_ids: 4294967296 - called_computation_ids: 3 + called_computation_ids: 2 frontend_attributes { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { - name: "reduce-scatter-start" - opcode: "async-start" - shape { - element_type: TUPLE - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - } - tuple_shapes { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - } - metadata { - } - id: 4294967300 - operand_ids: 4294967299 - called_computation_ids: 5 - frontend_attributes { - } - statistics_viz { - } - backend_config_payload { - id: 1 - } - } - instructions { - name: "reduce-scatter-done" - opcode: "async-done" + name: "ars.3" + opcode: "reduce-scatter" shape { element_type: U32 dimensions: 4 @@ -830,13 +762,21 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "reduce-scatter-start" } - id: 4294967301 - operand_ids: 4294967300 + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 frontend_attributes { } statistics_viz { } + collective_device_list { + } + backend_config_payload { + id: 0 + } } program_shape { result { @@ -850,7 +790,7 @@ gpu_executable { } } id: 1 - root_id: 4294967301 + root_id: 4294967300 } host_program_shape { result { @@ -881,14 +821,6 @@ gpu_executable { instruction_ids: 4294967297 instruction_ids: 4294967299 instruction_ids: 4294967300 - instruction_ids: 4294967301 - } - } - sequences { - key: 5 - value { - instruction_ids: 21474836480 - instruction_ids: 21474836481 } } } @@ -907,10 +839,9 @@ gpu_executable { frontend_attributes { map { key: "fingerprint_before_lhs" - value: "f1e11fca63da83b4e8c82bd1944aded5" + value: "c36d8870be58d238935bebefbbc5b827" } } - payloads: "{\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"native_emitter_backend_config\":{\"type\":\"NATIVE_EMITTER_TYPE_INVALID\",\"unroll_factor\":0},\"operation_queue_id\":\"0\",\"reification_cost\":[]}" payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" } config { @@ -942,7 +873,7 @@ gpu_executable { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -968,7 +899,6 @@ gpu_executable { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -981,14 +911,10 @@ gpu_executable { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -1032,6 +958,7 @@ gpu_executable { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -1040,7 +967,7 @@ gpu_executable { xla_cmd_buffer_trace_cache_size: 16 xla_gpu_temp_buffer_use_separate_color: false xla_syntax_sugar_async_ops: false - xla_gpu_autotune_gemm_rtol: 0.1 + xla_gpu_autotune_gemm_rtol: 0.01 xla_enable_command_buffers_during_profiling: false xla_gpu_cudnn_gemm_max_plans: 5 xla_cpu_parallel_codegen_split_count: 32 @@ -1102,7 +1029,6 @@ gpu_executable { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -1111,7 +1037,6 @@ gpu_executable { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -1122,14 +1047,13 @@ gpu_executable { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -1139,21 +1063,32 @@ gpu_executable { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" xla_enable_hlo_sharding_v3: false xla_gpu_log_minmax: false @@ -1164,6 +1099,7 @@ gpu_executable { xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -1183,12 +1119,12 @@ gpu_executable { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -1197,6 +1133,35 @@ gpu_executable { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" } static_device_assignment { replica_count: 2 @@ -1238,58 +1203,43 @@ gpu_executable { } logical_buffers { id: 10 - size: 16 - defined_at { - instruction_id: 21474836481 - } - } - logical_buffers { - id: 11 size: 32 defined_at { instruction_id: 4294967297 } } logical_buffers { - id: 12 + id: 11 size: 32 defined_at { instruction_id: 4294967298 } } logical_buffers { - id: 13 + id: 12 size: 4 defined_at { instruction_id: 4294967296 } } logical_buffers { - id: 14 + id: 13 size: 32 defined_at { instruction_id: 4294967299 } } logical_buffers { - id: 15 + id: 14 size: 16 defined_at { instruction_id: 4294967300 } } - logical_buffers { - id: 16 - size: 8 - defined_at { - shape_index: 0 - instruction_id: 4294967300 - } - } buffer_allocations { size: 32 assigned { - logical_buffer_id: 11 + logical_buffer_id: 10 size: 32 element_type: U32 } @@ -1299,7 +1249,7 @@ gpu_executable { index: 1 size: 32 assigned { - logical_buffer_id: 12 + logical_buffer_id: 11 size: 32 element_type: U32 } @@ -1310,13 +1260,13 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 10 - size: 16 + logical_buffer_id: 12 + size: 4 element_type: U32 } assigned { - logical_buffer_id: 13 - size: 4 + logical_buffer_id: 14 + size: 16 element_type: U32 } } @@ -1335,7 +1285,6 @@ gpu_executable { size: 4 is_thread_local: true assigned { - logical_buffer_id: 1 size: 4 element_type: U32 } @@ -1345,56 +1294,45 @@ gpu_executable { size: 4 is_thread_local: true assigned { + logical_buffer_id: 1 size: 4 element_type: U32 } } buffer_allocations { index: 6 - size: 528 + size: 32 assigned { - logical_buffer_id: 14 + logical_buffer_id: 13 size: 32 element_type: U32 } - assigned { - logical_buffer_id: 15 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 8 - element_type: TUPLE - } } heap_simulator_traces { events { - buffer_id: 14 + buffer_id: 13 computation_name: "main" instruction_name: "loop_select_fusion" } events { kind: FREE - buffer_id: 14 + buffer_id: 13 computation_name: "main" instruction_name: "loop_select_fusion" } whole_module_simulation: true - buffer_allocation_index: 5 + buffer_allocation_index: 3 } peak_buffers { - logical_buffer_ids: 11 + logical_buffer_ids: 10 } peak_buffers { index: 1 - logical_buffer_ids: 12 + logical_buffer_ids: 11 } peak_buffers { index: 2 - logical_buffer_ids: 10 + logical_buffer_ids: 14 } peak_buffers { index: 3 @@ -1402,24 +1340,22 @@ gpu_executable { } peak_buffers { index: 4 - logical_buffer_ids: 1 + logical_buffer_ids: 0 } peak_buffers { index: 5 - logical_buffer_ids: 0 + logical_buffer_ids: 1 } peak_buffers { index: 6 - logical_buffer_ids: 14 - logical_buffer_ids: 16 - logical_buffer_ids: 15 + logical_buffer_ids: 13 } } - binary: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\005\000\000\000\000\000\000\000\003\000\000\000\000\000\000<\005<\000@\0008\000\004\000@\000\t\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.global\000.nv.constant4\000.rel.nv.constant4\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.global\000.nv.constant4\000.rel.nv.constant4\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000buffer_for_c1_1\000buffer_for_c0_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000=\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000y\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\000\000\000\021\000\010\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\230\000\000\000\021\000\010\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\010\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\210\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000y\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\002\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000K\000\000\000\t\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\007\000\000\000\010\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000=\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000h\002\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\003\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000h\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + binary: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\004\000\000\000\000\000\000Z\rZ\000@\0008\000\004\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000a\000\000\000 \240\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\215\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\326\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\021\000\t\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\002\001\000\000\021\000\t\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\270\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\001\000\000\000\000\000\000\022\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000d\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\206\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\002\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\007\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000S\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000(\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0008\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\000\000\000\000\010\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" gpu_compute_capability { cuda_compute_capability { - major: 6 - feature_extension: NONE + major: 9 + feature_extension: ACCELERATED_FEATURES } } module_name: "test" @@ -1437,87 +1373,36 @@ gpu_executable { buffer_allocations { values { size: 32 - assigned { - logical_buffer_id: 11 - size: 32 - element_type: U32 - } is_constant: true } values { index: 1 size: 32 - assigned { - logical_buffer_id: 12 - size: 32 - element_type: U32 - } is_constant: true } values { index: 2 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 10 - size: 16 - element_type: U32 - } - assigned { - logical_buffer_id: 13 - size: 4 - element_type: U32 - } } values { index: 3 size: 4 is_thread_local: true - assigned { - logical_buffer_id: 2 - size: 4 - element_type: U32 - } } values { index: 4 size: 4 is_thread_local: true - assigned { - logical_buffer_id: 1 - size: 4 - element_type: U32 - } } values { index: 5 size: 4 is_thread_local: true - assigned { - size: 4 - element_type: U32 - } } values { index: 6 - size: 528 - assigned { - logical_buffer_id: 14 - size: 32 - element_type: U32 - } - assigned { - logical_buffer_id: 15 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 8 - element_type: TUPLE - } + size: 32 } } output_info_map { @@ -1621,7 +1506,7 @@ gpu_executable { name: "loop_select_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\000\006\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.constant0.loop_select_fusion\000.rel.nv.constant0.loop_select_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.rel.nv.constant0.loop_select_fusion\000.nv.constant0.loop_select_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\252\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\347\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\022\020\t\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004/\010\000\005\000\000\000\007\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001 \000\003\031 \000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\270\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000G\005\200\007\230L\000\000\027\002\000\000\310\360\342\007 \366\000\204\037\000\003\000W\005\200\007\230L\002\002\007\000\000\240\324\356\004\000\007\005\200\007\230L\377\007 \376A\304\037\000\005\000\027\005\200\007\230L\006\000\'\000\000\000H8\000\000\347\001\000\000(8\361\007\251\374\000\230\037\000\004\004\'\005\000\001DK\005\0057\005\000\001DK\004\004g\000\000\200\020\\\342\007 \366\000\230\037\000\005\005\007\000\000\010\020\\\004\004\007\000\000\240\324\356\002\006g\005\000\200\020L\347\007 \376 \374\037\000\003\000w\005\000\010\020L\004\002\007\000\000 \334\356\017\000\007\000\000\000\000\343\377\007\000\374\000\200\037\000\017\000\007\377\377\017@\342\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\366\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0006\001\000\000\000\000\000\000\t\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\002\000\000\000\000\000\000|\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\347\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\205\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000`\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\000\001\000\000\000\000\000\000\003\000\000\000\005\000\000\007 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000`\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\0000\014\000\000\000\000\000\0000\t\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\205\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\264\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\325\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000!\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\0004\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\020\002 \000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\210\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\316\017\000\202{\004\377\000\204\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\212\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\206\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\207\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\000\000\000\000U\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\233\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\004\000\000\000\000\000\000\200\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\271\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\205\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\347\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\0000\002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\0000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\0000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\002\000\000\000\000\000\0000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 4 kernel_name: "loop_select_fusion" @@ -1666,13 +1551,14 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } } thunks { thunk_info { - profile_annotation: "ars.2" + profile_annotation: "ars.3" thunk_id: 3 } reduce_scatter_thunk { @@ -1721,9 +1607,9 @@ gpu_executable { platform_name: "CUDA" cuda_platform_version { cuda_toolkit_version: "12.8.0" - cudnn_version: "9.10.0" + cudnn_version: "9.10.2" cub_version: "3.1.2" } } - buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 528B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 620B\n total fragmentation: 548B (88.39%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: ((u32[8]), u32[4])\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: async-start\n\t\tShape: (u32[8])\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 8:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 9:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" } diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/AllGatherMixedTypes.pbtxt similarity index 60% rename from third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/AllGatherMixedTypes.pbtxt rename to third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/AllGatherMixedTypes.pbtxt index f75e67bfac744d..ca49a6f758af8b 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/AllGatherMixedTypes.pbtxt +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/AllGatherMixedTypes.pbtxt @@ -1,9 +1,9 @@ -# proto-file: third_party/tensorflow/compiler/xla/tools/split_proto/split_proto_cli.proto -# proto-message: xla.split_proto_cli.DeserializedSplitExecutableAndOptions - +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable executable_and_options { compile_options { executable_build_options { + device_ordinal: -1 result_layout { element_type: TUPLE tuple_shapes { @@ -39,7 +39,7 @@ executable_and_options { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -65,7 +65,6 @@ executable_and_options { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -78,14 +77,10 @@ executable_and_options { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -129,6 +124,7 @@ executable_and_options { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -199,7 +195,6 @@ executable_and_options { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -208,7 +203,6 @@ executable_and_options { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -219,14 +213,13 @@ executable_and_options { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -236,31 +229,44 @@ executable_and_options { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" - xla_enable_hlo_sharding_v3: false + xla_enable_hlo_sharding_v3: true xla_gpu_log_minmax: false xla_gpu_print_compilation_stats: false xla_gpu_experimental_enable_conv_fusion: false - xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_experimental_enable_tiling_propagation: true xla_gpu_execution_progress_tracking: 0 xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -280,12 +286,12 @@ executable_and_options { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -294,6 +300,39 @@ executable_and_options { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET } num_replicas: 2 num_partitions: 1 @@ -311,9 +350,9 @@ executable_and_options { allow_spmd_sharding_propagation_to_parameters: false process_count: 1 memory_fitting_level: EFFORT_O2 - slice_size: 2 } } + pjrt_client_name: "PjRtStreamExecutorClient" } gpu_executable { hlo_module_with_config { @@ -338,6 +377,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.1" } id: 4294967298 frontend_attributes { @@ -361,6 +401,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.6.1" } id: 4294967297 operand_ids: 4294967298 @@ -385,6 +426,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "transpose.1.1" } dimensions: 1 dimensions: 0 @@ -444,6 +486,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.3" } id: 8589934594 frontend_attributes { @@ -467,6 +510,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.4.1" } id: 8589934593 operand_ids: 8589934594 @@ -491,6 +535,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "transpose.3" } dimensions: 1 dimensions: 0 @@ -544,6 +589,7 @@ gpu_executable { } } metadata { + scheduling_name: "param_0.5" } id: 12884901890 frontend_attributes { @@ -561,6 +607,7 @@ gpu_executable { } } metadata { + scheduling_name: "p1.1.1" } id: 12884901889 operand_ids: 12884901890 @@ -585,6 +632,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "broadcast.13.1" } id: 12884901888 operand_ids: 12884901889 @@ -618,156 +666,6 @@ gpu_executable { root_id: 12884901888 is_fusion_computation: true } - computations { - name: "async_computation" - instructions { - name: "param_0.6" - opcode: "parameter" - shape { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - metadata { - } - id: 17179869184 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "param_1" - opcode: "parameter" - shape { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - metadata { - } - parameter_number: 1 - id: 17179869185 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "allgather.2" - opcode: "all-gather" - shape { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - metadata { - } - dimensions: 1 - id: 17179869186 - operand_ids: 17179869184 - operand_ids: 17179869185 - frontend_attributes { - } - statistics_viz { - } - collective_device_list { - } - } - program_shape { - parameters { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - parameters { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - result { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - parameter_names: "param_0.6" - parameter_names: "param_1" - } - id: 4 - root_id: 17179869186 - } computations { name: "wrapped_broadcast_computation" instructions { @@ -780,6 +678,7 @@ gpu_executable { } } metadata { + scheduling_name: "param_0.7" } id: 21474836481 frontend_attributes { @@ -803,6 +702,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "broadcast.12.1" } id: 21474836480 operand_ids: 21474836481 @@ -848,6 +748,7 @@ gpu_executable { } } metadata { + scheduling_name: "id.1" } frontend_attributes { } @@ -870,18 +771,16 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "wrapped_broadcast" } fusion_kind: "kLoop" - id: 11 + id: 9 operand_ids: 0 called_computation_ids: 5 frontend_attributes { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "loop_broadcast_fusion" @@ -899,6 +798,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_broadcast_fusion" } fusion_kind: "kLoop" id: 8 @@ -908,87 +808,10 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } - } - instructions { - name: "all-gather-start" - opcode: "async-start" - shape { - element_type: TUPLE - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - } - metadata { - } - id: 9 - operand_ids: 11 - operand_ids: 8 - called_computation_ids: 4 - frontend_attributes { - } - statistics_viz { - } - backend_config_payload { - id: 1 - } } instructions { - name: "all-gather-done" - opcode: "async-done" + name: "allgather.3" + opcode: "all-gather" shape { element_type: TUPLE tuple_shapes { @@ -1017,13 +840,21 @@ gpu_executable { } } metadata { + scheduling_name: "all-gather-start" } + dimensions: 1 id: 10 operand_ids: 9 + operand_ids: 8 frontend_attributes { } statistics_viz { } + collective_device_list { + } + backend_config_payload { + id: 0 + } } instructions { name: "ag0.1" @@ -1041,6 +872,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "ag0.1" } id: 1 operand_ids: 10 @@ -1065,6 +897,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_transpose_fusion.1" } fusion_kind: "kLoop" id: 7 @@ -1074,9 +907,6 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "bitcast.24" @@ -1091,6 +921,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.24" } id: 4 operand_ids: 7 @@ -1115,6 +946,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "ag1.1" } tuple_index: 1 id: 2 @@ -1140,6 +972,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_transpose_fusion" } fusion_kind: "kLoop" id: 6 @@ -1149,9 +982,6 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "bitcast.1.0" @@ -1166,6 +996,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.1.0" } id: 5 operand_ids: 6 @@ -1199,6 +1030,7 @@ gpu_executable { } } metadata { + scheduling_name: "out.1" } id: 3 operand_ids: 4 @@ -1262,26 +1094,17 @@ gpu_executable { value { instruction_ids: 0 instruction_ids: 8 - instruction_ids: 11 instruction_ids: 9 instruction_ids: 10 - instruction_ids: 2 instruction_ids: 1 instruction_ids: 7 - instruction_ids: 4 + instruction_ids: 2 instruction_ids: 6 + instruction_ids: 4 instruction_ids: 5 instruction_ids: 3 } } - sequences { - key: 4 - value { - instruction_ids: 17179869185 - instruction_ids: 17179869184 - instruction_ids: 17179869186 - } - } } input_output_alias { } @@ -1298,11 +1121,10 @@ gpu_executable { frontend_attributes { map { key: "fingerprint_before_lhs" - value: "27beb9c461cc00eca8e7c7d6d313ca37" + value: "7b2e4f79c5524b5cbc566da178aa7fda" } } - payloads: "{\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"native_emitter_backend_config\":{\"type\":\"NATIVE_EMITTER_TYPE_INVALID\",\"unroll_factor\":0},\"operation_queue_id\":\"0\",\"reification_cost\":[]}" - payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" } config { entry_computation_layout { @@ -1345,7 +1167,7 @@ gpu_executable { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -1371,7 +1193,6 @@ gpu_executable { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -1384,14 +1205,10 @@ gpu_executable { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -1435,6 +1252,7 @@ gpu_executable { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -1505,7 +1323,6 @@ gpu_executable { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -1514,7 +1331,6 @@ gpu_executable { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -1525,14 +1341,13 @@ gpu_executable { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -1542,31 +1357,44 @@ gpu_executable { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" - xla_enable_hlo_sharding_v3: false + xla_enable_hlo_sharding_v3: true xla_gpu_log_minmax: false xla_gpu_print_compilation_stats: false xla_gpu_experimental_enable_conv_fusion: false - xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_experimental_enable_tiling_propagation: true xla_gpu_execution_progress_tracking: 0 xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -1586,12 +1414,12 @@ gpu_executable { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -1600,6 +1428,39 @@ gpu_executable { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET } static_device_assignment { replica_count: 2 @@ -1620,95 +1481,65 @@ gpu_executable { } } buffer_assignment { - logical_buffers { - id: 7 - size: 16 - defined_at { - instruction_id: 17179869186 - } - } - logical_buffers { - id: 8 - size: 16 - defined_at { - shape_index: 0 - instruction_id: 17179869186 - } - } logical_buffers { id: 9 - size: 16 - defined_at { - shape_index: 1 - instruction_id: 17179869186 - } - } - logical_buffers { - id: 12 size: 4 defined_at { } } logical_buffers { - id: 13 + id: 10 size: 8 defined_at { - instruction_id: 11 + instruction_id: 9 } } logical_buffers { - id: 14 + id: 11 size: 8 defined_at { instruction_id: 8 } } logical_buffers { - id: 15 + id: 12 size: 16 defined_at { - instruction_id: 9 + instruction_id: 10 } } logical_buffers { - id: 16 + id: 13 size: 16 defined_at { shape_index: 0 - instruction_id: 9 + instruction_id: 10 } } logical_buffers { - id: 17 + id: 14 size: 16 defined_at { shape_index: 1 - instruction_id: 9 - } - } - logical_buffers { - id: 18 - size: 16 - defined_at { instruction_id: 10 } } logical_buffers { - id: 19 + id: 15 size: 16 defined_at { instruction_id: 7 } } logical_buffers { - id: 20 + id: 16 size: 16 defined_at { instruction_id: 6 } } logical_buffers { - id: 21 + id: 17 size: 16 defined_at { instruction_id: 3 @@ -1718,12 +1549,12 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 13 + logical_buffer_id: 10 size: 8 element_type: U32 } assigned { - logical_buffer_id: 19 + logical_buffer_id: 15 size: 16 element_type: U32 } @@ -1733,17 +1564,17 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 8 - size: 16 + logical_buffer_id: 9 + size: 4 element_type: U32 } assigned { - logical_buffer_id: 12 - size: 4 + logical_buffer_id: 13 + size: 16 element_type: U32 } assigned { - logical_buffer_id: 20 + logical_buffer_id: 16 size: 16 element_type: F32 } @@ -1753,7 +1584,7 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 21 + logical_buffer_id: 17 size: 16 element_type: TUPLE } @@ -1761,102 +1592,74 @@ gpu_executable { } buffer_allocations { index: 3 - size: 1544 + size: 520 assigned { - logical_buffer_id: 7 - offset: 1024 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 9 - offset: 1280 - size: 16 - element_type: F32 - } - assigned { - logical_buffer_id: 14 - offset: 1536 + logical_buffer_id: 11 + offset: 512 size: 8 element_type: F32 } assigned { - logical_buffer_id: 15 + logical_buffer_id: 12 size: 16 element_type: TUPLE } assigned { - logical_buffer_id: 16 + logical_buffer_id: 14 offset: 256 size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 17 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 18 - offset: 768 - size: 16 - element_type: TUPLE + element_type: F32 } is_tuple: true } heap_simulator_traces { events { - buffer_id: 14 + buffer_id: 11 computation_name: "test_computation" instruction_name: "loop_broadcast_fusion" } events { - buffer_id: 9 - computation_name: "async_computation" - instruction_name: "allgather.2" + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" } events { kind: FREE - buffer_id: 14 + buffer_id: 11 computation_name: "test_computation" instruction_name: "loop_broadcast_fusion" } events { kind: FREE - buffer_id: 9 - computation_name: "async_computation" - instruction_name: "allgather.2" + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" } whole_module_simulation: true buffer_allocation_index: 3 } peak_buffers { - logical_buffer_ids: 19 + logical_buffer_ids: 15 } peak_buffers { index: 1 - logical_buffer_ids: 20 + logical_buffer_ids: 16 } peak_buffers { index: 2 - logical_buffer_ids: 21 + logical_buffer_ids: 17 } peak_buffers { index: 3 - logical_buffer_ids: 15 - logical_buffer_ids: 16 - logical_buffer_ids: 17 - logical_buffer_ids: 18 - logical_buffer_ids: 7 - logical_buffer_ids: 9 + logical_buffer_ids: 12 + logical_buffer_ids: 11 logical_buffer_ids: 14 } } gpu_compute_capability { cuda_compute_capability { - major: 6 - feature_extension: NONE + major: 9 + feature_extension: ACCELERATED_FEATURES } } module_name: "test" @@ -1887,92 +1690,21 @@ gpu_executable { values { size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 13 - size: 8 - element_type: U32 - } - assigned { - logical_buffer_id: 19 - size: 16 - element_type: U32 - } } values { index: 1 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 8 - size: 16 - element_type: U32 - } - assigned { - logical_buffer_id: 12 - size: 4 - element_type: U32 - } - assigned { - logical_buffer_id: 20 - size: 16 - element_type: F32 - } } values { index: 2 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 21 - size: 16 - element_type: TUPLE - } is_tuple: true } values { index: 3 - size: 1544 - assigned { - logical_buffer_id: 7 - offset: 1024 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 9 - offset: 1280 - size: 16 - element_type: F32 - } - assigned { - logical_buffer_id: 14 - offset: 1536 - size: 8 - element_type: F32 - } - assigned { - logical_buffer_id: 15 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 17 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 18 - offset: 768 - size: 16 - element_type: TUPLE - } + size: 520 is_tuple: true } } @@ -2032,7 +1764,7 @@ gpu_executable { } args { slice { - offset: 1536 + offset: 512 size: 8 buffer_allocation_index: 3 element_type: F32 @@ -2056,7 +1788,7 @@ gpu_executable { name: "loop_broadcast_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\010\000\000\000\000\000\000\300\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000.rel.nv.constant0.loop_broadcast_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.rel.nv.constant0.loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\266\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\007\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000x\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000\007\005\200\007\230L\004\000\027\002\000\000\310\360\342\007\340\367\000\304?\010\003\000\027\005\200\007\230L\002\002\007\000\000\240\324\356\000\004\347\001\000\000(8\346\007\340\374\000t\034\004\004\004\'\005\000\201\030L\005\0007\005\000\010\020L\006\n\'\000\000\000\270\\\361\017\200\374\000\374\037\000\006\004\007\000\000 \334\356\000\017\007\000\000\000\260P\017\000\007\000\000\000\000\343\377\007\000\374\000\200\037\000\017\000\007\377\377\017@\342\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\005\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000E\001\000\000\000\000\000\000\033\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\366\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\007 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\275\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\336\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000*\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000@\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\204\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\206\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\024\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\211\001\000\000\000\000\000\000d\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\244\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\324\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\302\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_broadcast_fusion" @@ -2089,6 +1821,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2136,7 +1869,7 @@ gpu_executable { name: "wrapped_broadcast" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\340\007\000\000\000\000\000\000`\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000.rel.nv.constant0.wrapped_broadcast\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.rel.nv.constant0.wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\246\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\361\000\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000p\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000\007\005\200\007\230L\004\000\027\002\000\000\310\360\342\007\340\367\000\304?\010\003\000\027\005\200\007\230L\002\002\007\000\000\240\324\356\000\004\347\001\000\000(8\346\007\340\374\000\304\037\004\004\004\'\005\000\201\030L\005\0007\005\000\010\020L\002\004\007\000\000 \334\356\352\007\340\377\000\374\037\000\000\017\007\000\000\000\260P\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\361\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0001\001\000\000\000\000\000\000\003\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000J\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\354\002\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\342\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000h\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\202\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\004\000\000\000\000\000\000\200\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000\340\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000x\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\350\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\340\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\261\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\036\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\0000\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\204\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\206\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\004\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000u\001\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\230\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000J\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000t\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\266\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\230\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "wrapped_broadcast" @@ -2169,13 +1902,14 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } } thunks { thunk_info { - profile_annotation: "allgather.2" + profile_annotation: "allgather.3" thunk_id: 4 } all_gather_thunk { @@ -2227,7 +1961,7 @@ gpu_executable { element_count: 2 source_buffer { slice { - offset: 1536 + offset: 512 size: 8 buffer_allocation_index: 3 element_type: F32 @@ -2247,7 +1981,7 @@ gpu_executable { } destination_buffer { slice { - offset: 1280 + offset: 256 size: 16 buffer_allocation_index: 3 element_type: F32 @@ -2317,7 +2051,7 @@ gpu_executable { name: "loop_transpose_fusion_1" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000`\010\000\000\000\000\000\000\340\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000.rel.nv.constant0.loop_transpose_fusion_1\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.rel.nv.constant0.loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\340\343\000\230?\000\001\000\207\000\200\007\230L\005\000\027\002\000\000\310\360\000\005W\000\000\000\020\\\366\007\300\376\000\304\037\010\000\000\'\000\000\000\000\004\000\005\007\360\201\177\337[\003\000\347\001\000\000(8\346\007@\374\000\304\036\000\002\000\007\005\000\201\030L\003\003\027\005\000\010\020L\002\002\007\000\000\240\324\356\361\007\302\374\000\234\037\000\000\005\347\001\000\000(8\004\005\'\005\000\201\030L\005\0007\005\000\010\020L\361\007\341\377\000\374\037\000\002\004\007\000\000 \334\356\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\017\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000O\001\000\000\000\000\000\000\'\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000P\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\001\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\224\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\270\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000`\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\270\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\002\000\000\000\000\000\000(\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000`\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\224\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\303\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000H\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\204\000\000\000\n\000\000\000b\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\206\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\016\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\034\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\223\001\000\000\000\000\000\000n\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\252\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000P\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\264\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\310\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\224\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_transpose_fusion_1" @@ -2350,6 +2084,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2362,7 +2097,7 @@ gpu_executable { custom_kernel_thunk { args { slice { - offset: 1280 + offset: 256 size: 16 buffer_allocation_index: 3 element_type: F32 @@ -2405,7 +2140,7 @@ gpu_executable { name: "loop_transpose_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\010\000\000\000\000\000\000\300\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000.rel.nv.constant0.loop_transpose_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.rel.nv.constant0.loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\266\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\340\343\000\230?\000\001\000\207\000\200\007\230L\005\000\027\002\000\000\310\360\000\005W\000\000\000\020\\\366\007\300\376\000\304\037\010\000\000\'\000\000\000\000\004\000\005\007\360\201\177\337[\003\000\347\001\000\000(8\346\007@\374\000\304\036\000\002\000\007\005\000\201\030L\003\003\027\005\000\010\020L\002\002\007\000\000\240\324\356\361\007\302\374\000\234\037\000\000\005\347\001\000\000(8\004\005\'\005\000\201\030L\005\0007\005\000\010\020L\361\007\341\377\000\374\037\000\002\004\007\000\000 \334\356\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\005\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000E\001\000\000\000\000\000\000\033\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\366\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\275\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\336\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000*\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000@\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\204\000\000\000\n\000\000\000b\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\206\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\016\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\024\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\211\001\000\000\000\000\000\000d\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\244\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\324\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\302\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_transpose_fusion" @@ -2438,6 +2173,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2446,9 +2182,9 @@ gpu_executable { platform_name: "CUDA" cuda_platform_version { cuda_toolkit_version: "12.8.0" - cudnn_version: "9.10.0" + cudnn_version: "9.10.2" cub_version: "3.1.2" } } - buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 1.5KiB\n preallocated temp fragmentation: 240B (15.54%)\n total allocation: 1.6KiB\n total fragmentation: 1.4KiB (91.96%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: async-done\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: (u32[2,1], f32[2,1])\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: ((u32[2,1], f32[2,1]), (u32[2,2], f32[2,2]))\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 8:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 9:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 10:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" } diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/ReduceScatter.pbtxt similarity index 66% rename from third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/ReduceScatter.pbtxt rename to third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/ReduceScatter.pbtxt index dbf790a32fd03a..fcd7a336a45f1c 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/ReduceScatter.pbtxt +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/ReduceScatter.pbtxt @@ -1,9 +1,9 @@ -# proto-file: third_party/tensorflow/compiler/xla/tools/split_proto/split_proto_cli.proto -# proto-message: xla.split_proto_cli.DeserializedSplitExecutableAndOptions - +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable executable_and_options { compile_options { executable_build_options { + device_ordinal: -1 result_layout { element_type: U32 dimensions: 4 @@ -27,7 +27,7 @@ executable_and_options { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -53,7 +53,6 @@ executable_and_options { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -66,14 +65,10 @@ executable_and_options { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -117,6 +112,7 @@ executable_and_options { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -187,7 +183,6 @@ executable_and_options { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -196,7 +191,6 @@ executable_and_options { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -207,14 +201,13 @@ executable_and_options { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -224,31 +217,44 @@ executable_and_options { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" - xla_enable_hlo_sharding_v3: false + xla_enable_hlo_sharding_v3: true xla_gpu_log_minmax: false xla_gpu_print_compilation_stats: false xla_gpu_experimental_enable_conv_fusion: false - xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_experimental_enable_tiling_propagation: true xla_gpu_execution_progress_tracking: 0 xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -268,12 +274,12 @@ executable_and_options { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -282,6 +288,39 @@ executable_and_options { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET } num_replicas: 2 num_partitions: 1 @@ -299,9 +338,9 @@ executable_and_options { allow_spmd_sharding_propagation_to_parameters: false process_count: 1 memory_fitting_level: EFFORT_O2 - slice_size: 2 } } + pjrt_client_name: "PjRtStreamExecutorClient" } gpu_executable { hlo_module_with_config { @@ -320,6 +359,7 @@ gpu_executable { } } metadata { + scheduling_name: "lhs" } frontend_attributes { } @@ -336,6 +376,7 @@ gpu_executable { } } metadata { + scheduling_name: "rhs" } parameter_number: 1 id: 1 @@ -354,6 +395,7 @@ gpu_executable { } } metadata { + scheduling_name: "add.1" } id: 2 operand_ids: 0 @@ -399,9 +441,10 @@ gpu_executable { } } metadata { + scheduling_name: "param_2.2" } parameter_number: 2 - id: 12884901893 + id: 8589934597 frontend_attributes { } statistics_viz { @@ -427,14 +470,14 @@ gpu_executable { } u32s: 0 } - id: 12884901894 + id: 8589934598 frontend_attributes { } statistics_viz { } } instructions { - name: "p.4" + name: "p.2" opcode: "compare" shape { element_type: PRED @@ -443,19 +486,20 @@ gpu_executable { } } metadata { + scheduling_name: "p.2" } - id: 12884901892 - operand_ids: 12884901893 - operand_ids: 12884901894 + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 comparison_direction: "EQ" frontend_attributes { } - comparison_type: "UNSIGNED" statistics_viz { } + comparison_order: "TOTAL" } instructions { - name: "pb.4" + name: "pb.2" opcode: "broadcast" shape { element_type: PRED @@ -467,16 +511,17 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "pb.2" } - id: 12884901889 - operand_ids: 12884901892 + id: 8589934593 + operand_ids: 8589934596 frontend_attributes { } statistics_viz { } } instructions { - name: "param_0.4" + name: "param_0.1" opcode: "parameter" shape { element_type: U32 @@ -488,15 +533,16 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.1" } - id: 12884901890 + id: 8589934594 frontend_attributes { } statistics_viz { } } instructions { - name: "param_1.3" + name: "param_1.1" opcode: "parameter" shape { element_type: U32 @@ -508,9 +554,10 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_1.1" } parameter_number: 1 - id: 12884901891 + id: 8589934595 frontend_attributes { } statistics_viz { @@ -529,11 +576,12 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "data.2" } - id: 12884901888 - operand_ids: 12884901889 - operand_ids: 12884901890 - operand_ids: 12884901891 + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 frontend_attributes { } statistics_viz { @@ -573,85 +621,14 @@ gpu_executable { } is_dynamic_dimension: false } - parameter_names: "param_0.4" - parameter_names: "param_1.3" + parameter_names: "param_0.1" + parameter_names: "param_1.1" parameter_names: "param_2.2" } - id: 3 - root_id: 12884901888 + id: 2 + root_id: 8589934592 is_fusion_computation: true } - computations { - name: "async_computation" - instructions { - name: "param_0.5" - opcode: "parameter" - shape { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - metadata { - } - id: 21474836480 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "ars.2" - opcode: "reduce-scatter" - shape { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - metadata { - } - dimensions: 0 - id: 21474836481 - operand_ids: 21474836480 - called_computation_ids: 0 - frontend_attributes { - } - statistics_viz { - } - collective_device_list { - } - } - program_shape { - parameters { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - result { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - parameter_names: "param_0.5" - } - id: 5 - root_id: 21474836481 - } computations { name: "main" instructions { @@ -742,6 +719,7 @@ gpu_executable { } } metadata { + scheduling_name: "id.1" } id: 4294967296 frontend_attributes { @@ -762,64 +740,22 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_select_fusion" } fusion_kind: "kLoop" id: 4294967299 operand_ids: 4294967297 operand_ids: 4294967298 operand_ids: 4294967296 - called_computation_ids: 3 + called_computation_ids: 2 frontend_attributes { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { - name: "reduce-scatter-start" - opcode: "async-start" - shape { - element_type: TUPLE - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - } - tuple_shapes { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - } - metadata { - } - id: 4294967300 - operand_ids: 4294967299 - called_computation_ids: 5 - frontend_attributes { - } - statistics_viz { - } - backend_config_payload { - id: 1 - } - } - instructions { - name: "reduce-scatter-done" - opcode: "async-done" + name: "ars.3" + opcode: "reduce-scatter" shape { element_type: U32 dimensions: 4 @@ -830,13 +766,21 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "reduce-scatter-start" } - id: 4294967301 - operand_ids: 4294967300 + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 frontend_attributes { } statistics_viz { } + collective_device_list { + } + backend_config_payload { + id: 0 + } } program_shape { result { @@ -850,7 +794,7 @@ gpu_executable { } } id: 1 - root_id: 4294967301 + root_id: 4294967300 } host_program_shape { result { @@ -881,14 +825,6 @@ gpu_executable { instruction_ids: 4294967297 instruction_ids: 4294967299 instruction_ids: 4294967300 - instruction_ids: 4294967301 - } - } - sequences { - key: 5 - value { - instruction_ids: 21474836480 - instruction_ids: 21474836481 } } } @@ -907,11 +843,10 @@ gpu_executable { frontend_attributes { map { key: "fingerprint_before_lhs" - value: "ecc7fa57c3eeb812403576a4054128b5" + value: "c36d8870be58d238935bebefbbc5b827" } } - payloads: "{\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"native_emitter_backend_config\":{\"type\":\"NATIVE_EMITTER_TYPE_INVALID\",\"unroll_factor\":0},\"operation_queue_id\":\"0\",\"reification_cost\":[]}" - payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" } config { entry_computation_layout { @@ -942,7 +877,7 @@ gpu_executable { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -968,7 +903,6 @@ gpu_executable { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -981,14 +915,10 @@ gpu_executable { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -1032,6 +962,7 @@ gpu_executable { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -1102,7 +1033,6 @@ gpu_executable { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -1111,7 +1041,6 @@ gpu_executable { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -1122,14 +1051,13 @@ gpu_executable { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -1139,31 +1067,44 @@ gpu_executable { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" - xla_enable_hlo_sharding_v3: false + xla_enable_hlo_sharding_v3: true xla_gpu_log_minmax: false xla_gpu_print_compilation_stats: false xla_gpu_experimental_enable_conv_fusion: false - xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_experimental_enable_tiling_propagation: true xla_gpu_execution_progress_tracking: 0 xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -1183,12 +1124,12 @@ gpu_executable { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -1197,6 +1138,39 @@ gpu_executable { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET } static_device_assignment { replica_count: 2 @@ -1238,58 +1212,43 @@ gpu_executable { } logical_buffers { id: 10 - size: 16 - defined_at { - instruction_id: 21474836481 - } - } - logical_buffers { - id: 11 size: 32 defined_at { instruction_id: 4294967297 } } logical_buffers { - id: 12 + id: 11 size: 32 defined_at { instruction_id: 4294967298 } } logical_buffers { - id: 13 + id: 12 size: 4 defined_at { instruction_id: 4294967296 } } logical_buffers { - id: 14 + id: 13 size: 32 defined_at { instruction_id: 4294967299 } } logical_buffers { - id: 15 + id: 14 size: 16 defined_at { instruction_id: 4294967300 } } - logical_buffers { - id: 16 - size: 8 - defined_at { - shape_index: 0 - instruction_id: 4294967300 - } - } buffer_allocations { size: 32 assigned { - logical_buffer_id: 11 + logical_buffer_id: 10 size: 32 element_type: U32 } @@ -1299,7 +1258,7 @@ gpu_executable { index: 1 size: 32 assigned { - logical_buffer_id: 12 + logical_buffer_id: 11 size: 32 element_type: U32 } @@ -1310,13 +1269,13 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 10 - size: 16 + logical_buffer_id: 12 + size: 4 element_type: U32 } assigned { - logical_buffer_id: 13 - size: 4 + logical_buffer_id: 14 + size: 16 element_type: U32 } } @@ -1335,7 +1294,6 @@ gpu_executable { size: 4 is_thread_local: true assigned { - logical_buffer_id: 1 size: 4 element_type: U32 } @@ -1345,56 +1303,45 @@ gpu_executable { size: 4 is_thread_local: true assigned { + logical_buffer_id: 1 size: 4 element_type: U32 } } buffer_allocations { index: 6 - size: 528 + size: 32 assigned { - logical_buffer_id: 14 + logical_buffer_id: 13 size: 32 element_type: U32 } - assigned { - logical_buffer_id: 15 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 8 - element_type: TUPLE - } } heap_simulator_traces { events { - buffer_id: 14 + buffer_id: 13 computation_name: "main" instruction_name: "loop_select_fusion" } events { kind: FREE - buffer_id: 14 + buffer_id: 13 computation_name: "main" instruction_name: "loop_select_fusion" } whole_module_simulation: true - buffer_allocation_index: 5 + buffer_allocation_index: 3 } peak_buffers { - logical_buffer_ids: 11 + logical_buffer_ids: 10 } peak_buffers { index: 1 - logical_buffer_ids: 12 + logical_buffer_ids: 11 } peak_buffers { index: 2 - logical_buffer_ids: 10 + logical_buffer_ids: 14 } peak_buffers { index: 3 @@ -1402,24 +1349,22 @@ gpu_executable { } peak_buffers { index: 4 - logical_buffer_ids: 1 + logical_buffer_ids: 0 } peak_buffers { index: 5 - logical_buffer_ids: 0 + logical_buffer_ids: 1 } peak_buffers { index: 6 - logical_buffer_ids: 14 - logical_buffer_ids: 16 - logical_buffer_ids: 15 + logical_buffer_ids: 13 } } - binary: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\005\000\000\000\000\000\000\000\003\000\000\000\000\000\000<\005<\000@\0008\000\004\000@\000\t\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.global\000.nv.constant4\000.rel.nv.constant4\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.global\000.nv.constant4\000.rel.nv.constant4\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000buffer_for_c1_1\000buffer_for_c0_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000=\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000y\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\000\000\000\021\000\010\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\230\000\000\000\021\000\010\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\010\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\210\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000y\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\002\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000K\000\000\000\t\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\007\000\000\000\010\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000=\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000h\002\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\003\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000h\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + binary: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\004\000\000\000\000\000\000Z\rZ\000@\0008\000\004\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000a\000\000\000 \240\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\215\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\326\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\021\000\t\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\002\001\000\000\021\000\t\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\270\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\001\000\000\000\000\000\000\022\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000d\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\206\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\002\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\007\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000S\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000(\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0008\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\000\000\000\000\010\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" gpu_compute_capability { cuda_compute_capability { - major: 6 - feature_extension: NONE + major: 9 + feature_extension: ACCELERATED_FEATURES } } module_name: "test" @@ -1437,87 +1382,36 @@ gpu_executable { buffer_allocations { values { size: 32 - assigned { - logical_buffer_id: 11 - size: 32 - element_type: U32 - } is_constant: true } values { index: 1 size: 32 - assigned { - logical_buffer_id: 12 - size: 32 - element_type: U32 - } is_constant: true } values { index: 2 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 10 - size: 16 - element_type: U32 - } - assigned { - logical_buffer_id: 13 - size: 4 - element_type: U32 - } } values { index: 3 size: 4 is_thread_local: true - assigned { - logical_buffer_id: 2 - size: 4 - element_type: U32 - } } values { index: 4 size: 4 is_thread_local: true - assigned { - logical_buffer_id: 1 - size: 4 - element_type: U32 - } } values { index: 5 size: 4 is_thread_local: true - assigned { - size: 4 - element_type: U32 - } } values { index: 6 - size: 528 - assigned { - logical_buffer_id: 14 - size: 32 - element_type: U32 - } - assigned { - logical_buffer_id: 15 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 8 - element_type: TUPLE - } + size: 32 } } output_info_map { @@ -1621,7 +1515,7 @@ gpu_executable { name: "loop_select_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\000\006\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.constant0.loop_select_fusion\000.rel.nv.constant0.loop_select_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.rel.nv.constant0.loop_select_fusion\000.nv.constant0.loop_select_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\252\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\347\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\022\020\t\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004/\010\000\005\000\000\000\007\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001 \000\003\031 \000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\270\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000G\005\200\007\230L\000\000\027\002\000\000\310\360\342\007 \366\000\204\037\000\003\000W\005\200\007\230L\002\002\007\000\000\240\324\356\004\000\007\005\200\007\230L\377\007 \376A\304\037\000\005\000\027\005\200\007\230L\006\000\'\000\000\000H8\000\000\347\001\000\000(8\361\007\251\374\000\230\037\000\004\004\'\005\000\001DK\005\0057\005\000\001DK\004\004g\000\000\200\020\\\342\007 \366\000\230\037\000\005\005\007\000\000\010\020\\\004\004\007\000\000\240\324\356\002\006g\005\000\200\020L\347\007 \376 \374\037\000\003\000w\005\000\010\020L\004\002\007\000\000 \334\356\017\000\007\000\000\000\000\343\377\007\000\374\000\200\037\000\017\000\007\377\377\017@\342\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\366\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0006\001\000\000\000\000\000\000\t\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\002\000\000\000\000\000\000|\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\347\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\205\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000`\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\000\001\000\000\000\000\000\000\003\000\000\000\005\000\000\007 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000`\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\0000\014\000\000\000\000\000\0000\t\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\205\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\264\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\325\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000!\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\0004\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\020\002 \000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\210\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\316\017\000\202{\004\377\000\204\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\212\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\206\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\207\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\000\000\000\000U\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\233\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\004\000\000\000\000\000\000\200\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\271\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\205\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\347\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\0000\002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\0000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\0000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\002\000\000\000\000\000\0000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 4 kernel_name: "loop_select_fusion" @@ -1666,13 +1560,14 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } } thunks { thunk_info { - profile_annotation: "ars.2" + profile_annotation: "ars.3" thunk_id: 3 } reduce_scatter_thunk { @@ -1721,9 +1616,9 @@ gpu_executable { platform_name: "CUDA" cuda_platform_version { cuda_toolkit_version: "12.8.0" - cudnn_version: "9.10.0" + cudnn_version: "9.10.2" cub_version: "3.1.2" } } - buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 528B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 620B\n total fragmentation: 548B (88.39%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: ((u32[8]), u32[4])\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: async-start\n\t\tShape: (u32[8])\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 8:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 9:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" } diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc index 065a3691bfb7ef..99c74ff514a656 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc @@ -23,13 +23,19 @@ limitations under the License. #include #include +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/ascii.h" #include "absl/strings/match.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "xla/pjrt/pjrt_client.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" +#include "xla/stream_executor/stream_executor.h" #include "xla/tests/aot_interception_pjrt_client.h" #include "xla/tests/hlo_test_base.h" #include "xla/tests/pjrt_client_registry.h" @@ -88,6 +94,36 @@ absl::StatusOr GetUndeclaredOutputsDir() { } // namespace +std::string DetectGpuArchToken() { + if (const char* env_arch = std::getenv("XLA_AOT_GOLDEN_ARCH")) { + return env_arch; + } + absl::StatusOr platform = + stream_executor::PlatformManager::PlatformWithName("CUDA"); + CHECK_OK(platform.status()) + << "AOT golden arch detection failed to get the CUDA platform; set " + "XLA_AOT_GOLDEN_ARCH to override."; + absl::StatusOr executor = + (*platform)->ExecutorForDevice(0); + CHECK_OK(executor.status()) + << "AOT golden arch detection failed to get a device executor; set " + "XLA_AOT_GOLDEN_ARCH to override."; + const std::string name = + absl::AsciiStrToLower((*executor)->GetDeviceDescription().name()); + + // Longest token first: "gb200" contains "b200". gb300 is deliberately + // absent -- it is listed in gpu/BUILD's disabled_backends (b/491194726), so + // no gb300 goldens exist and no gb300 target is generated. + for (const absl::string_view arch : {"gb200", "b200", "h100"}) { + if (absl::StrContains(name, arch)) { + return std::string(arch); + } + } + + LOG(FATAL) << "Unrecognized GPU arch for AOT goldens: " << name + << "; set XLA_AOT_GOLDEN_ARCH to override."; +} + std::string GetExecutablesDirectory(absl::string_view target_name, AOTTestPlatform platform) { return tsl::io::JoinPath( diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h index 5f0b89b9bd71bb..43a6fe9dad2a46 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h @@ -48,6 +48,13 @@ struct AotTestParam { } }; +// Returns the golden-directory token for the GPU this process is running on: +// one of "h100", "b200", "gb200". `XLA_AOT_GOLDEN_ARCH`, when set, is returned +// verbatim and skips device detection; this is what update_goldens.py uses. +// Crashes with LOG(FATAL) if the device cannot be queried or its name matches +// no known token. +std::string DetectGpuArchToken(); + // Returns the path to the executables directory for the given test target and // platform. std::string GetExecutablesDirectory(absl::string_view target_name, From dfa49395468b20c48716c5c21eb757f8f72cd013 Mon Sep 17 00:00:00 2001 From: Sohaib Iftikhar Date: Thu, 24 Sep 2026 09:28:56 -0700 Subject: [PATCH 36/48] [XLA:GPU]: Implement cufunction record for GPU backend. PiperOrigin-RevId: 987558515 --- .../xla/xla/backends/gpu/runtime/BUILD | 25 +++- .../xla/backends/gpu/runtime/record_ffi.cc | 120 ++++++++++++----- .../backends/gpu/runtime/record_ffi_test.cc | 123 ++++++++++++++---- .../gpu/runtime/record_ffi_test_cuda.cc | 46 +++++++ .../gpu/runtime/record_ffi_test_default.cc | 33 +++++ .../gpu/runtime/record_ffi_test_rocm.cc | 46 +++++++ .../xla/xla/stream_executor/command_buffer.h | 16 +++ .../cuda/cuda_command_buffer.cc | 87 ++++++++----- .../cuda/cuda_command_buffer.h | 13 ++ .../stream_executor/gpu/gpu_command_buffer.cc | 25 ++++ .../stream_executor/gpu/gpu_command_buffer.h | 24 ++++ third_party/xla/xla/stream_executor/kernel.h | 11 ++ .../xla/stream_executor/mock_command_buffer.h | 15 +++ .../xla/xla/stream_executor/rocm/BUILD | 1 + .../rocm/rocm_command_buffer.cc | 74 +++++++---- .../rocm/rocm_command_buffer.h | 13 ++ 16 files changed, 557 insertions(+), 115 deletions(-) create mode 100644 third_party/xla/xla/backends/gpu/runtime/record_ffi_test_cuda.cc create mode 100644 third_party/xla/xla/backends/gpu/runtime/record_ffi_test_default.cc create mode 100644 third_party/xla/xla/backends/gpu/runtime/record_ffi_test_rocm.cc diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index c1fdc1ecc51f15..789e090da66f79 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -1,6 +1,8 @@ load("@local_config_cuda//cuda:build_defs.bzl", "cuda_library") +load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured") load("@rules_cc//cc:cc_library.bzl", "cc_library") load("//xla:xla.default.bzl", "xla_cc_test", "xla_internal", "xla_py_proto_library") +load("//xla/stream_executor:build_defs.bzl", "if_cuda_or_rocm_is_configured") load("//xla/tests:build_defs.bzl", "xla_test") load("//xla/tsl:tsl.bzl", "internal_visibility", "nvtx_headers") load("//xla/tsl:tsl.default.bzl", "get_compatible_with_portable") @@ -5539,6 +5541,7 @@ cc_library( "//xla/stream_executor:stream_executor_h", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/debugging:symbolize", + "@com_google_absl//absl/functional:overload", "@com_google_absl//absl/hash", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", @@ -5554,7 +5557,14 @@ cc_library( xla_test( name = "record_ffi_test", - srcs = ["record_ffi_test.cc"], + srcs = ["record_ffi_test.cc"] + if_cuda_is_configured( + ["record_ffi_test_cuda.cc"], + ) + if_rocm_is_configured( + ["record_ffi_test_rocm.cc"], + ) + if_cuda_or_rocm_is_configured( + [], + ["record_ffi_test_default.cc"], + ), backends = ["gpu"], deps = [ ":record_ffi", @@ -5566,9 +5576,9 @@ xla_test( "//xla/ffi:invoke", "//xla/ffi:record_ffi", "//xla/ffi/api:c_api", + "//xla/stream_executor:activate_context", "//xla/stream_executor:command_buffer", "//xla/stream_executor:device_address", - "//xla/stream_executor:kernel_spec", "//xla/stream_executor:platform", "//xla/stream_executor:platform_manager", "//xla/stream_executor:stream", @@ -5576,12 +5586,21 @@ xla_test( "//xla/stream_executor/gpu:gpu_init", "//xla/stream_executor/gpu:gpu_test_kernels_fatbin", "//xla/tsl/platform:test", + "@com_google_absl//absl/base", "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", - ], + ] + if_cuda_is_configured([ + "@local_config_cuda//cuda:cuda_headers", + "//xla/stream_executor/cuda:cuda_status", + ]) + if_rocm_is_configured([ + "@local_config_rocm//rocm:hip", # buildcleaner: keep + "@local_config_rocm//rocm:rocm_headers", + "//xla/stream_executor/rocm:rocm_status", + ]), ) diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi.cc index ef8bf056b05b7e..a423b0881be065 100644 --- a/third_party/xla/xla/backends/gpu/runtime/record_ffi.cc +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi.cc @@ -24,10 +24,12 @@ limitations under the License. #include #include #include +#include #include #include "absl/container/flat_hash_map.h" #include "absl/debugging/symbolize.h" +#include "absl/functional/overload.h" #include "absl/hash/hash.h" #include "absl/log/check.h" #include "absl/log/log.h" @@ -53,6 +55,7 @@ limitations under the License. #include "tsl/platform/mem.h" namespace xla::gpu { + namespace { template @@ -209,18 +212,24 @@ struct FfiLaunchParams { LaunchDimensions launch_dimensions; std::optional cluster_dims; - se::Kernel* kernel; + std::variant kernel; uint32_t shared_mem_bytes; }; template void AbslStringify(Sink& sink, const FfiLaunchParams& params) { + std::string kernel_name = + std::visit(absl::Overload{[](se::Kernel* k) { return GetSymbolName(k); }, + [](const se::NativeKernel& nk) { + return GetSymbolName(nk.device_fn); + }}, + params.kernel); absl::Format(&sink, "FfiLaunchParams(launch_dimensions=%s, cluster_dims=%s, " "kernel=%s, shared_mem_bytes=%u)", params.launch_dimensions.ToString(), params.cluster_dims.value_or(se::ClusterDim{0, 0, 0}).ToString(), - GetSymbolName(params.kernel), params.shared_mem_bytes); + kernel_name, params.shared_mem_bytes); } class FfiKernelCache : public se::CommandBuffer::Resource { @@ -268,22 +277,35 @@ class FfiKernelCache : public se::CommandBuffer::Resource { << ", ptr: " << it->second.get(); return it->second.get(); } - bool is_ptx = (format == XLA_FFI_SourceFormat_PTX); std::string kernel_name(kernel_name_view); - se::KernelLoaderSpec spec = - is_ptx - ? se::KernelLoaderSpec::CreateCudaPtxInMemorySpec( - AsStringView(kernel_data, kernel_size), kernel_name, num_args) - : se::KernelLoaderSpec::CreateCudaCubinInMemorySpec( - AsByteSpan(kernel_data, kernel_size), kernel_name, num_args); - + std::optional spec; + switch (format) { + case XLA_FFI_SourceFormat_FUNCTION_PTR: { + return absl::InvalidArgumentError( + "XLA_FFI_SourceFormat_FUNCTION_PTR does not use FfiKernelCache"); + } + case XLA_FFI_SourceFormat_PTX: { + spec = se::KernelLoaderSpec::CreateCudaPtxInMemorySpec( + AsStringView(kernel_data, kernel_size), kernel_name, num_args); + break; + } + case XLA_FFI_SourceFormat_CUBIN: { + spec = se::KernelLoaderSpec::CreateCudaCubinInMemorySpec( + AsByteSpan(kernel_data, kernel_size), kernel_name, num_args); + break; + } + } + if (!spec.has_value()) { + return absl::InvalidArgumentError( + absl::StrFormat("Unsupported XLA_FFI_SourceFormat: %d", format)); + } ABSL_ASSIGN_OR_RETURN(std::unique_ptr kernel, - executor->LoadKernel(spec)); - - se::Kernel* kernel_ptr = kernel.get(); - kernels_[KernelKey{kernel_name, kernel_data}] = std::move(kernel); + executor->LoadKernel(*spec)); + auto [inserted_it, unused] = kernels_.try_emplace( + KernelKey{std::move(kernel_name), kernel_data}, std::move(kernel)); + se::Kernel* kernel_ptr = inserted_it->second.get(); XLA_VLOG_DEVICE(3, device_ordinal) - << "FfiKernelCache: created kernel: " << kernel_name + << "FfiKernelCache: created kernel: " << kernel_name_view << ", ptr: " << kernel_ptr; return kernel_ptr; } @@ -326,16 +348,31 @@ XLA_FFI_Error* FfiCreateLaunch( const XLA_FFI_KernelArgs* args, const XLA_FFI_Command* const* dependencies, uint32_t num_dependencies, const XLA_FFI_Command** out_command) { se::CommandBuffer* cmd_buffer = ctx->command_buffer; + const bool pdl = ctx->pdl_enabled && (use_pdl != 0); auto* cache = cmd_buffer->GetOrConstructResource(); - auto kernel_or = - cache->GetOrCreateKernel(ctx->executor, kernel_name, kernel_data, - kernel_size, format, args->num_args); - if (!kernel_or.ok()) { - return xla::ffi::CreateError(kernel_or.status()); + std::variant kernel_target; + if (format == XLA_FFI_SourceFormat_FUNCTION_PTR) { + if (kernel_data == nullptr) { + return xla::ffi::CreateError(absl::InvalidArgumentError( + "kernel_data cannot be null for XLA_FFI_SourceFormat_FUNCTION_PTR")); + } + // FFI C API passes kernel_data as const void*, whereas driver function + // handles (e.g. CUfunction / hipFunction_t) are non-const opaque pointers. + // NOLINTNEXTLINE + kernel_target = se::NativeKernel{const_cast(kernel_data), + std::string(kernel_name), pdl}; + } else { + auto kernel_or = + cache->GetOrCreateKernel(ctx->executor, kernel_name, kernel_data, + kernel_size, format, args->num_args); + if (!kernel_or.ok()) { + return xla::ffi::CreateError(kernel_or.status()); + } + se::Kernel* kernel = *kernel_or; + kernel->set_use_pdl(pdl); + kernel_target = kernel; } - se::Kernel* kernel = *kernel_or; - kernel->set_use_pdl(ctx->pdl_enabled && (use_pdl != 0)); FfiKernelArgsPacked packed_args(args->num_args); packed_args.add_shared_bytes(shared_mem_bytes); @@ -366,19 +403,27 @@ XLA_FFI_Error* FfiCreateLaunch( FfiLaunchParams params{ /*.launch_dimensions = */ LaunchDimensions(blocks, threads), /*.cluster_dims =*/cluster_dims, - /*.kernel =*/kernel, + /*.kernel =*/kernel_target, /*.shared_mem_bytes = */ shared_mem_bytes, }; const int32_t device_ordinal = ctx->executor->device_ordinal(); XLA_VLOG_DEVICE(3, device_ordinal) - << "FfiCreateLaunch for kernel: " << kernel_name - << ", use_pdl: " << kernel->use_pdl() + << "FfiCreateLaunch for kernel: " << kernel_name << ", use_pdl: " << pdl << ", num_dependencies passed: " << num_dependencies << ", resolved deps size: " << deps.size() << ", launch_params: " << params; - auto status_or_cmd = cmd_buffer->CreateLaunch(threads, blocks, cluster_dims, - *kernel, packed_args, deps); + auto status_or_cmd = std::visit( + absl::Overload{ + [&](se::Kernel* k) { + return cmd_buffer->CreateLaunch(threads, blocks, cluster_dims, *k, + packed_args, deps); + }, + [&](const se::NativeKernel& nk) { + return cmd_buffer->CreateLaunch(threads, blocks, cluster_dims, nk, + packed_args, deps); + }}, + params.kernel); if (!status_or_cmd.ok()) { return xla::ffi::CreateError(status_or_cmd.status()); @@ -410,11 +455,22 @@ XLA_FFI_Error* FfiUpdateLaunch(XLA_FFI_RecordContext* ctx, if (absl::Status status = PackArgs(args, packed_args); !status.ok()) { return xla::ffi::CreateError(status); } - if (absl::Status status = cmd_buffer->UpdateLaunch( - cmd, params->launch_dimensions.thread_counts_per_block(), - params->launch_dimensions.block_counts(), params->cluster_dims, - *params->kernel, packed_args); - !status.ok()) { + absl::Status status = std::visit( + absl::Overload{ + [&](se::Kernel* k) { + return cmd_buffer->UpdateLaunch( + cmd, params->launch_dimensions.thread_counts_per_block(), + params->launch_dimensions.block_counts(), params->cluster_dims, + *k, packed_args); + }, + [&](const se::NativeKernel& nk) { + return cmd_buffer->UpdateLaunch( + cmd, params->launch_dimensions.thread_counts_per_block(), + params->launch_dimensions.block_counts(), params->cluster_dims, + nk, packed_args); + }}, + params->kernel); + if (!status.ok()) { return xla::ffi::CreateError(status); } diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc index eca91c3fc50cde..caebad7b3c0747 100644 --- a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc @@ -26,6 +26,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/ffi/api/record_api.h" @@ -39,7 +40,6 @@ limitations under the License. #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/gpu/gpu_init.h" #include "xla/stream_executor/gpu/gpu_test_kernels_fatbin.h" -#include "xla/stream_executor/kernel_spec.h" #include "xla/stream_executor/platform.h" #include "xla/stream_executor/platform_manager.h" #include "xla/stream_executor/stream.h" @@ -48,6 +48,11 @@ limitations under the License. #include "xla/xla_data.pb.h" namespace xla::gpu { + +absl::StatusOr LoadNativeFunctionPtr(se::StreamExecutor* executor, + absl::Span binary, + absl::string_view name); + namespace { struct KernelBinary { @@ -132,17 +137,7 @@ struct DeviceMemoryBundle { } }; -TEST(RecordFfiTest, KernelLaunchBoundFfi) { - ASSERT_OK_AND_ASSIGN(auto platform, - stream_executor::PlatformManager::PlatformWithName( - stream_executor::GpuPlatformName())); - ASSERT_OK_AND_ASSIGN(auto* executor, platform->ExecutorForDevice(0)); - const auto* cuda_cc = executor->GetDeviceDescription() - .gpu_compute_capability() - .cuda_compute_capability(); - if (cuda_cc && !cuda_cc->IsAtLeastAmpere()) { - GTEST_SKIP() << "Skipping test for compute capability less than Ampere."; - } +void RunRecordFfiTest(se::StreamExecutor* executor, ffi::Ffi& handler) { ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream()); std::vector a = {1, 2, 3, 4, 5, 6, 7, 8}; @@ -179,15 +174,7 @@ TEST(RecordFfiTest, KernelLaunchBoundFfi) { ffi::CallFrame call_frame = builder.Build(); - std::unique_ptr handler = - ffi::Ffi::BindRecord() - .Ctx>() - .Arg() - .Arg() - .Arg() - .Arg() - .To(RecordFfiHandler); - ASSERT_OK(ffi::Invoke(ffi::GetXlaFfiApi(), *handler, call_frame, + ASSERT_OK(ffi::Invoke(ffi::GetXlaFfiApi(), handler, call_frame, invoke_context, ffi::ExecutionStage::kRecord)); ASSERT_OK(cmd_buffer->Finalize()); ASSERT_OK(cmd_buffer->Submit(stream.get())); @@ -219,7 +206,7 @@ TEST(RecordFfiTest, KernelLaunchBoundFfi) { ffi::InvokeContext update_context; update_context.extension_start = &record_extension.extension_base; ASSERT_OK(cmd_buffer->Update()); // Begin command buffer update. - ASSERT_OK(ffi::Invoke(ffi::GetXlaFfiApi(), *handler, call_frame, + ASSERT_OK(ffi::Invoke(ffi::GetXlaFfiApi(), handler, call_frame, update_context, ffi::ExecutionStage::kRecord)); ASSERT_OK(cmd_buffer->Finalize()); ASSERT_OK(cmd_buffer->Submit(stream.get())); @@ -232,5 +219,97 @@ TEST(RecordFfiTest, KernelLaunchBoundFfi) { EXPECT_EQ(c, expected_update); } +TEST(RecordFfiTest, KernelLaunchBoundFfi) { + ASSERT_OK_AND_ASSIGN(auto platform, + stream_executor::PlatformManager::PlatformWithName( + stream_executor::GpuPlatformName())); + ASSERT_OK_AND_ASSIGN(auto* executor, platform->ExecutorForDevice(0)); + const auto* cuda_cc = executor->GetDeviceDescription() + .gpu_compute_capability() + .cuda_compute_capability(); + if (cuda_cc && !cuda_cc->IsAtLeastAmpere()) { + GTEST_SKIP() << "Skipping test for compute capability less than Ampere."; + } + + std::unique_ptr handler = + ffi::Ffi::BindRecord() + .Ctx>() + .Arg() + .Arg() + .Arg() + .Arg() + .To(RecordFfiHandler); + RunRecordFfiTest(executor, *handler); +} + +TEST(RecordFfiTest, KernelLaunchWithCuFunc) { + ASSERT_OK_AND_ASSIGN(auto platform, + stream_executor::PlatformManager::PlatformWithName( + stream_executor::GpuPlatformName())); + ASSERT_OK_AND_ASSIGN(auto* executor, platform->ExecutorForDevice(0)); + const auto* cuda_cc = executor->GetDeviceDescription() + .gpu_compute_capability() + .cuda_compute_capability(); + if (cuda_cc && !cuda_cc->IsAtLeastAmpere()) { + GTEST_SKIP() << "Skipping test for compute capability less than Ampere."; + } + + ASSERT_OK_AND_ASSIGN(KernelBinary binary, GetKernelSpec()); + absl::StatusOr func_ptr = + LoadNativeFunctionPtr(executor, binary.bytes, "AddI32"); + if (absl::IsUnimplemented(func_ptr.status())) { + GTEST_SKIP() << func_ptr.status().message(); + } + ASSERT_OK(func_ptr.status()); + + std::unique_ptr handler = + ffi::Ffi::BindRecord() + .Ctx>() + .Arg() + .Arg() + .Arg() + .Arg() + .To([func_ptr = *func_ptr]( + ffi::RecordContext record_ctx, ffi::AnyBuffer input_0, + ffi::AnyBuffer input_1, ffi::AnyBuffer buffer_scratch, + ffi::AnyBuffer result) -> absl::Status { + void* in0 = input_0.untyped_data(); + void* in1 = input_1.untyped_data(); + void* scratch = buffer_scratch.untyped_data(); + void* res = result.untyped_data(); + size_t size = 8 * sizeof(int32_t); + + if (record_ctx.action() == ffi::RecordAction::kCreate) { + ABSL_ASSIGN_OR_RETURN(const XLA_FFI_Command* memcpy_d2d, + record_ctx.CreateMemcpyD2D(scratch, in0, size)); + ABSL_RETURN_IF_ERROR( + record_ctx + .CreateLaunch( + "AddI32", func_ptr, /*kernel_size=*/0, + ffi::SourceFormat::kFunctionPtr, + /*launch_dims=*/{{1, 1, 1}, {8, 1, 1}}, + /*shared_mem_bytes=*/1024, + /*uses_pdl=*/false, + std::vector{ + ffi::DevicePointer{scratch}, + ffi::DevicePointer{in1}, ffi::DevicePointer{res}}, + /*dependencies=*/{memcpy_d2d}) + .status()); + } else if (record_ctx.action() == ffi::RecordAction::kUpdate) { + auto cmds = record_ctx.commands(); + TF_RET_CHECK(cmds.size() == 2) << "Expected 2 commands recorded."; + ABSL_RETURN_IF_ERROR( + record_ctx.UpdateMemcpyD2D(cmds[0], scratch, in0, size)); + ABSL_RETURN_IF_ERROR(record_ctx.UpdateLaunch( + cmds[1], + std::vector{ffi::DevicePointer{scratch}, + ffi::DevicePointer{in1}, + ffi::DevicePointer{res}})); + } + return absl::OkStatus(); + }); + RunRecordFfiTest(executor, *handler); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_cuda.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_cuda.cc new file mode 100644 index 00000000000000..856ecf040c1ab4 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_cuda.cc @@ -0,0 +1,46 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include + +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "third_party/gpus/cuda/include/cuda.h" +#include "xla/stream_executor/activate_context.h" +#include "xla/stream_executor/cuda/cuda_status.h" +#include "xla/stream_executor/stream_executor.h" + +namespace xla::gpu { + +absl::StatusOr LoadNativeFunctionPtr( + stream_executor::StreamExecutor* executor, absl::Span binary, + absl::string_view name) { + std::unique_ptr activation = + executor->Activate(); + CUmodule module = nullptr; + ABSL_RETURN_IF_ERROR(stream_executor::cuda::ToStatus( + cuModuleLoadData(&module, binary.data()))); + CUfunction function = nullptr; + std::string kernel_name(name); + ABSL_RETURN_IF_ERROR(stream_executor::cuda::ToStatus( + cuModuleGetFunction(&function, module, kernel_name.c_str()))); + return function; +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_default.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_default.cc new file mode 100644 index 00000000000000..ec47494a4fa92e --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_default.cc @@ -0,0 +1,33 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "xla/stream_executor/stream_executor.h" + +namespace xla::gpu { + +absl::StatusOr LoadNativeFunctionPtr(stream_executor::StreamExecutor*, + absl::Span, + absl::string_view) { + return absl::UnimplementedError( + "Native function pointer loading is not supported on this platform."); +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_rocm.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_rocm.cc new file mode 100644 index 00000000000000..df24728c1442ef --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_rocm.cc @@ -0,0 +1,46 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include + +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "rocm/include/hip/hip_runtime_api.h" +#include "xla/stream_executor/activate_context.h" +#include "xla/stream_executor/rocm/rocm_status.h" // IWYU pragma: keep +#include "xla/stream_executor/stream_executor.h" + +namespace xla::gpu { + +absl::StatusOr LoadNativeFunctionPtr( + stream_executor::StreamExecutor* executor, absl::Span binary, + absl::string_view name) { + std::unique_ptr activation = + executor->Activate(); + hipModule_t module = nullptr; + ABSL_RETURN_IF_ERROR(stream_executor::gpu::ToStatus( + hipModuleLoadData(&module, binary.data()))); + hipFunction_t function = nullptr; + std::string kernel_name(name); + ABSL_RETURN_IF_ERROR(stream_executor::gpu::ToStatus( + hipModuleGetFunction(&function, module, kernel_name.c_str()))); + return function; +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/stream_executor/command_buffer.h b/third_party/xla/xla/stream_executor/command_buffer.h index e10feee73baadc..f33039e255bacb 100644 --- a/third_party/xla/xla/stream_executor/command_buffer.h +++ b/third_party/xla/xla/stream_executor/command_buffer.h @@ -169,12 +169,28 @@ class CommandBuffer { const KernelArgs& args, absl::Span dependencies, StreamPriority priority = StreamPriority::Default) = 0; + // Creates a kernel launch command from an externally-managed native function. + // Nolint because it matches existing CreateLaunch virtual method signature + // precedent for default arguments. + // NOLINTNEXTLINE + virtual absl::StatusOr CreateLaunch( + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args, + absl::Span dependencies, + StreamPriority priority = StreamPriority::Default) = 0; + // Updates a kernel launch command. virtual absl::Status UpdateLaunch( const Command* command, const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgs& args) = 0; + virtual absl::Status UpdateLaunch( + const Command* command, const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) = 0; + // Type-safe wrapper for launching typed kernels. Notice that the order of // arguments is different do disambiguate from the regular launch API. template diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc index a9f5c52482a0b1..1feae5877b0801 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc @@ -722,6 +722,29 @@ absl::StatusOr CudaCommandBuffer::CreateKernelNode( const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) { + const auto& cuda_kernel = static_cast(kernel); + ABSL_RETURN_IF_ERROR(cuda_kernel.UpdateMaxDynamicSharedMemoryBytes( + args.number_of_shared_bytes())); + + std::unique_ptr repacked; + const KernelArgsPackedArrayBase* packed_args = &args; + if (cuda_kernel.args_packing()) { + ABSL_ASSIGN_OR_RETURN(repacked, cuda_kernel.args_packing()(cuda_kernel, args)); + packed_args = repacked.get(); + } + + return CreateKernelNode( + dependencies, priority, threads, blocks, cluster_dims, + NativeKernel{cuda_kernel.gpu_function(), std::string(kernel.name()), + kernel.use_pdl()}, + *packed_args); +} + +absl::StatusOr CudaCommandBuffer::CreateKernelNode( + absl::Span dependencies, StreamPriority priority, + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) { const uint64_t shared_mem_bytes = args.number_of_shared_bytes(); XLA_VLOG_DEVICE(2, stream_exec_->device_ordinal()) @@ -733,19 +756,7 @@ absl::StatusOr CudaCommandBuffer::CreateKernelNode( << "): " << FormatGraphNodeHandles(dependencies); CUgraphNode node_handle = nullptr; - const auto& cuda_kernel = static_cast(kernel); - CUfunction function = cuda_kernel.gpu_function(); - ABSL_RETURN_IF_ERROR( - cuda_kernel.UpdateMaxDynamicSharedMemoryBytes(shared_mem_bytes)); - - std::unique_ptr repacked; - const KernelArgsPackedArrayBase* packed_args; - if (cuda_kernel.args_packing()) { - ABSL_ASSIGN_OR_RETURN(repacked, cuda_kernel.args_packing()(cuda_kernel, args)); - packed_args = repacked.get(); - } else { - packed_args = &args; - } + CUfunction function = static_cast(kernel.device_fn); auto set_params = [&](auto& params) { params.func = function; @@ -756,8 +767,10 @@ absl::StatusOr CudaCommandBuffer::CreateKernelNode( params.blockDimY = threads.y; params.blockDimZ = threads.z; params.sharedMemBytes = shared_mem_bytes; - params.kernelParams = - const_cast(packed_args->argument_addresses().data()); + // CUDA driver API requires void** for kernelParams even though it does not + // mutate the argument pointers. + // NOLINTNEXTLINE + params.kernelParams = const_cast(args.argument_addresses().data()); params.extra = nullptr; }; @@ -845,6 +858,28 @@ absl::Status CudaCommandBuffer::UpdateKernelNode( GraphNodeHandle node_handle, const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) { + const auto& cuda_kernel = static_cast(kernel); + ABSL_RETURN_IF_ERROR(cuda_kernel.UpdateMaxDynamicSharedMemoryBytes( + args.number_of_shared_bytes())); + + std::unique_ptr repacked; + const KernelArgsPackedArrayBase* packed_args = &args; + if (cuda_kernel.args_packing()) { + ABSL_ASSIGN_OR_RETURN(repacked, cuda_kernel.args_packing()(cuda_kernel, args)); + packed_args = repacked.get(); + } + + return UpdateKernelNode( + node_handle, threads, blocks, cluster_dims, + NativeKernel{cuda_kernel.gpu_function(), std::string(kernel.name()), + kernel.use_pdl()}, + *packed_args); +} + +absl::Status CudaCommandBuffer::UpdateKernelNode( + GraphNodeHandle node_handle, const ThreadDim& threads, + const BlockDim& blocks, const std::optional& cluster_dims, + const NativeKernel& kernel, const KernelArgsPackedArrayBase& args) { const uint64_t shared_mem_bytes = args.number_of_shared_bytes(); VLOG(2) << "Set kernel node params " << node_handle << " in graph executable " @@ -855,18 +890,7 @@ absl::Status CudaCommandBuffer::UpdateKernelNode( << "; shmem: " << shared_mem_bytes; CUDA_KERNEL_NODE_PARAMS params{}; - const auto& cuda_kernel = static_cast(kernel); - - std::unique_ptr repacked; - const KernelArgsPackedArrayBase* packed_args; - if (cuda_kernel.args_packing()) { - ABSL_ASSIGN_OR_RETURN(repacked, cuda_kernel.args_packing()(cuda_kernel, args)); - packed_args = repacked.get(); - } else { - packed_args = &args; - } - - CUfunction function = cuda_kernel.gpu_function(); + CUfunction function = static_cast(kernel.device_fn); params.func = function; params.gridDimX = blocks.x; params.gridDimY = blocks.y; @@ -875,13 +899,12 @@ absl::Status CudaCommandBuffer::UpdateKernelNode( params.blockDimY = threads.y; params.blockDimZ = threads.z; params.sharedMemBytes = shared_mem_bytes; - params.kernelParams = - const_cast(packed_args->argument_addresses().data()); + // CUDA driver API requires void** for kernelParams even though it does not + // mutate the argument pointers. + // NOLINTNEXTLINE + params.kernelParams = const_cast(args.argument_addresses().data()); params.extra = nullptr; - ABSL_RETURN_IF_ERROR( - cuda_kernel.UpdateMaxDynamicSharedMemoryBytes(shared_mem_bytes)); - if (cluster_dims.has_value()) { CUlaunchAttributeValue value; value.clusterDim.x = static_cast(cluster_dims->x); diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h index 1145053f600854..48e4ced1630c75 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h +++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h @@ -173,6 +173,12 @@ class CudaCommandBuffer final : public GpuCommandBuffer { const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) override; + absl::StatusOr CreateKernelNode( + absl::Span dependencies, StreamPriority priority, + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) override; + absl::Status UpdateKernelNode(GraphNodeHandle node_handle, const ThreadDim& threads, const BlockDim& blocks, @@ -180,6 +186,13 @@ class CudaCommandBuffer final : public GpuCommandBuffer { const Kernel& kernel, const KernelArgsPackedArrayBase& args) override; + absl::Status UpdateKernelNode(GraphNodeHandle node_handle, + const ThreadDim& threads, + const BlockDim& blocks, + const std::optional& cluster_dims, + const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) override; + absl::StatusOr CreateEmptyNode( absl::Span dependencies) override; diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc index 0a06de913f30ea..e99b249e3c5f11 100644 --- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc @@ -164,6 +164,31 @@ absl::Status GpuCommandBuffer::UpdateLaunchWithPackedArgs( kernel, packed_args); } +absl::StatusOr GpuCommandBuffer::CreateLaunch( + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args, + absl::Span dependencies, StreamPriority priority) { + ABSL_RETURN_IF_ERROR(CheckInState(State::kCreate)); + + ABSL_ASSIGN_OR_RETURN( + GraphNodeHandle handle, + CreateKernelNode(ToGraphNodeDependencies(dependencies), priority, threads, + blocks, cluster_dims, kernel, args)); + + return AppendCommand(GpuCommand{handle}); +} + +absl::Status GpuCommandBuffer::UpdateLaunch( + const Command* command, const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) { + ABSL_RETURN_IF_ERROR(CheckInState(State::kUpdate)); + auto* gpu_command = absl::down_cast(command); + return UpdateKernelNode(gpu_command->handle, threads, blocks, cluster_dims, + kernel, args); +} + absl::StatusOr GpuCommandBuffer::CreateLaunch( const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h index 3053c21b5dcb0f..bcb43b5272c230 100644 --- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h +++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h @@ -135,12 +135,25 @@ class GpuCommandBuffer : public CommandBuffer { const KernelArgs& args, absl::Span dependencies, StreamPriority priority) override; + absl::StatusOr CreateLaunch( + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args, + absl::Span dependencies, + StreamPriority priority) override; + absl::Status UpdateLaunch(const Command* command, const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgs& args) override; + absl::Status UpdateLaunch(const Command* command, const ThreadDim& threads, + const BlockDim& blocks, + const std::optional& cluster_dims, + const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) override; + absl::StatusOr CreateChildCommand( const CommandBuffer& nested, absl::Span dependencies) override; @@ -420,6 +433,12 @@ class GpuCommandBuffer : public CommandBuffer { const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) = 0; + virtual absl::StatusOr CreateKernelNode( + absl::Span dependencies, StreamPriority priority, + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) = 0; + // Updates the kernel launch node with the given parameters. Will return an // error if the given node has not been created as a kernel launch node. virtual absl::Status UpdateKernelNode( @@ -427,6 +446,11 @@ class GpuCommandBuffer : public CommandBuffer { const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) = 0; + virtual absl::Status UpdateKernelNode( + GraphNodeHandle node_handle, const ThreadDim& threads, + const BlockDim& blocks, const std::optional& cluster_dims, + const NativeKernel& kernel, const KernelArgsPackedArrayBase& args) = 0; + virtual absl::StatusOr CreateHostNode( absl::Span dependencies, absl::AnyInvocable callback) = 0; diff --git a/third_party/xla/xla/stream_executor/kernel.h b/third_party/xla/xla/stream_executor/kernel.h index a754b67e243d2f..77797c10c4d042 100644 --- a/third_party/xla/xla/stream_executor/kernel.h +++ b/third_party/xla/xla/stream_executor/kernel.h @@ -92,6 +92,17 @@ namespace stream_executor { // Kernel //===----------------------------------------------------------------------===// +// Platform-agnostic wrapper for an externally-managed native device function +// pointer (e.g. CUfunction or hipFunction_t). +struct NativeKernel { + void* device_fn = nullptr; + std::string kernel_name; + bool pdl = false; + + absl::string_view name() const { return kernel_name; } + bool use_pdl() const { return pdl; } +}; + // A data-parallel kernel (code entity) for launching via the StreamExecutor, // analogous to a void* device function pointer. See TypedKernel for the typed // variant. diff --git a/third_party/xla/xla/stream_executor/mock_command_buffer.h b/third_party/xla/xla/stream_executor/mock_command_buffer.h index 15aae998c0c2c2..1f830c44e9fbb8 100644 --- a/third_party/xla/xla/stream_executor/mock_command_buffer.h +++ b/third_party/xla/xla/stream_executor/mock_command_buffer.h @@ -55,12 +55,27 @@ class MockCommandBuffer : public CommandBuffer { absl::Span dependencies, StreamPriority priority), (override)); + MOCK_METHOD(absl::StatusOr, CreateLaunch, + (const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, + const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args, + absl::Span dependencies, + StreamPriority priority), + (override)); MOCK_METHOD(absl::Status, UpdateLaunch, (const Command* command, const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgs& args), (override)); + MOCK_METHOD(absl::Status, UpdateLaunch, + (const Command* command, const ThreadDim& threads, + const BlockDim& blocks, + const std::optional& cluster_dims, + const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args), + (override)); MOCK_METHOD(absl::StatusOr, CreateChildCommand, (const CommandBuffer& nested, absl::Span dependencies), diff --git a/third_party/xla/xla/stream_executor/rocm/BUILD b/third_party/xla/xla/stream_executor/rocm/BUILD index 0a73c4da4aa0cc..9a87e4ba0e043b 100644 --- a/third_party/xla/xla/stream_executor/rocm/BUILD +++ b/third_party/xla/xla/stream_executor/rocm/BUILD @@ -1233,6 +1233,7 @@ cc_library( "//xla/stream_executor:command_buffer", "//xla/stream_executor:device_address", "//xla/stream_executor:kernel", + "//xla/stream_executor:kernel_args", "//xla/stream_executor:launch_dim", "//xla/stream_executor:platform", "//xla/stream_executor:stream_executor_h", diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc index b81487137e927d..3e8d3aef633e9e 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc @@ -39,6 +39,7 @@ limitations under the License. #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/gpu/gpu_command_buffer.h" #include "xla/stream_executor/kernel.h" +#include "xla/stream_executor/kernel_args.h" #include "xla/stream_executor/launch_dim.h" #include "xla/stream_executor/platform.h" #include "xla/stream_executor/rocm/rocm_kernel.h" @@ -327,6 +328,25 @@ absl::StatusOr RocmCommandBuffer::CreateKernelNode( const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) { + std::unique_ptr repacked; + const KernelArgsPackedArrayBase* packed_args = &args; + if (kernel.args_packing()) { + ABSL_ASSIGN_OR_RETURN(repacked, kernel.args_packing()(kernel, args)); + packed_args = repacked.get(); + } + + return CreateKernelNode( + dependencies, priority, threads, blocks, cluster_dims, + NativeKernel{static_cast(kernel).gpu_function(), + std::string(kernel.name()), kernel.use_pdl()}, + *packed_args); +} + +absl::StatusOr RocmCommandBuffer::CreateKernelNode( + absl::Span dependencies, StreamPriority priority, + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) { const uint64_t shared_mem_bytes = args.number_of_shared_bytes(); VLOG(2) << "Add kernel node to a graph " << graph_ @@ -336,18 +356,8 @@ absl::StatusOr RocmCommandBuffer::CreateKernelNode( << " bdz: " << threads.z << "; shmem: " << shared_mem_bytes << "; deps: " << dependencies.size(); - std::unique_ptr repacked; - const KernelArgsPackedArrayBase* packed_args; - if (kernel.args_packing()) { - ABSL_ASSIGN_OR_RETURN(repacked, kernel.args_packing()(kernel, args)); - packed_args = repacked.get(); - } else { - packed_args = &args; - } - hipKernelNodeParams params{}; - hipFunction_t function = - static_cast(kernel).gpu_function(); + hipFunction_t function = static_cast(kernel.device_fn); params.func = function; params.gridDim.x = blocks.x; params.gridDim.y = blocks.y; @@ -356,8 +366,10 @@ absl::StatusOr RocmCommandBuffer::CreateKernelNode( params.blockDim.y = threads.y; params.blockDim.z = threads.z; params.sharedMemBytes = shared_mem_bytes; - params.kernelParams = - const_cast(packed_args->argument_addresses().data()); + // HIP driver API requires void** for kernelParams even though it does not + // mutate the argument pointers. + // NOLINTNEXTLINE + params.kernelParams = const_cast(args.argument_addresses().data()); params.extra = nullptr; if (shared_mem_bytes != 0) { @@ -383,6 +395,24 @@ absl::Status RocmCommandBuffer::UpdateKernelNode( GraphNodeHandle node_handle, const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) { + std::unique_ptr repacked; + const KernelArgsPackedArrayBase* packed_args = &args; + if (kernel.args_packing()) { + ABSL_ASSIGN_OR_RETURN(repacked, kernel.args_packing()(kernel, args)); + packed_args = repacked.get(); + } + + return UpdateKernelNode( + node_handle, threads, blocks, cluster_dims, + NativeKernel{static_cast(kernel).gpu_function(), + std::string(kernel.name()), kernel.use_pdl()}, + *packed_args); +} + +absl::Status RocmCommandBuffer::UpdateKernelNode( + GraphNodeHandle node_handle, const ThreadDim& threads, + const BlockDim& blocks, const std::optional& cluster_dims, + const NativeKernel& kernel, const KernelArgsPackedArrayBase& args) { const uint64_t shared_mem_bytes = args.number_of_shared_bytes(); VLOG(2) << "Set kernel node params " << node_handle << " in graph executable " @@ -391,18 +421,8 @@ absl::Status RocmCommandBuffer::UpdateKernelNode( << " bdx: " << threads.x << " bdy: " << threads.y << " bdz: " << threads.z << "; shmem: " << shared_mem_bytes; - std::unique_ptr repacked; - const KernelArgsPackedArrayBase* packed_args; - if (kernel.args_packing()) { - ABSL_ASSIGN_OR_RETURN(repacked, kernel.args_packing()(kernel, args)); - packed_args = repacked.get(); - } else { - packed_args = &args; - } - hipKernelNodeParams params{}; - hipFunction_t function = - static_cast(kernel).gpu_function(); + hipFunction_t function = static_cast(kernel.device_fn); params.func = function; params.gridDim.x = blocks.x; params.gridDim.y = blocks.y; @@ -411,8 +431,10 @@ absl::Status RocmCommandBuffer::UpdateKernelNode( params.blockDim.y = threads.y; params.blockDim.z = threads.z; params.sharedMemBytes = shared_mem_bytes; - params.kernelParams = - const_cast(packed_args->argument_addresses().data()); + // HIP driver API requires void** for kernelParams even though it does not + // mutate the argument pointers. + // NOLINTNEXTLINE + params.kernelParams = const_cast(args.argument_addresses().data()); params.extra = nullptr; if (shared_mem_bytes != 0) { diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h index b8e3ac54e6016a..3e2fd7b35d145c 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h +++ b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h @@ -159,6 +159,12 @@ class RocmCommandBuffer : public GpuCommandBuffer { const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) override; + absl::StatusOr CreateKernelNode( + absl::Span dependencies, StreamPriority priority, + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) override; + absl::Status UpdateKernelNode(GraphNodeHandle node_handle, const ThreadDim& threads, const BlockDim& blocks, @@ -166,6 +172,13 @@ class RocmCommandBuffer : public GpuCommandBuffer { const Kernel& kernel, const KernelArgsPackedArrayBase& args) override; + absl::Status UpdateKernelNode(GraphNodeHandle node_handle, + const ThreadDim& threads, + const BlockDim& blocks, + const std::optional& cluster_dims, + const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) override; + absl::StatusOr CreateEmptyNode( absl::Span dependencies) override; From eaedaed010a90b6c3e229a56c0db0b3493d74455 Mon Sep 17 00:00:00 2001 From: Jian Cai Date: Thu, 24 Sep 2026 09:47:01 -0700 Subject: [PATCH 37/48] [XLA][HLO Value Tracking] Maintain OriginalValue consistency across while loop transformations - In WhileUtil::AppendToWhileLoopOriginalValue and WhileUtil::MakeInstructionsLiveIn, update while instruction, while init, body root, body parameter, and condition parameter to keep all loop instruction OriginalValues structurally consistent when widening while loops, and preserve call_hierarchy. - In WhileLoopSimplifier (TryRemoveDeadTupleElements, TryRemoveConstantTupleElements, TryFlattenNestedTuples, and TrySimplifyInductionVariables), propagate and structurally update OriginalValues across all 5 loop boundary instructions and verify compatibility with the new shapes. PiperOrigin-RevId: 987569790 --- third_party/xla/xla/service/BUILD | 1 + .../xla/xla/service/while_loop_simplifier.cc | 141 ++++++++---- .../xla/service/while_loop_simplifier_test.cc | 124 +++++++++-- third_party/xla/xla/service/while_util.cc | 129 +++++++---- .../xla/xla/service/while_util_test.cc | 201 ++++++++++++++++++ third_party/xla/xla/tuple_tree.h | 11 + third_party/xla/xla/tuple_tree_test.cc | 5 + 7 files changed, 515 insertions(+), 97 deletions(-) diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 395738a19efe68..39fcaf0c8e5eea 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -4702,6 +4702,7 @@ xla_cc_test( srcs = ["while_util_test.cc"], deps = [ ":while_util", + "//xla:shape_util", "//xla:util", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", diff --git a/third_party/xla/xla/service/while_loop_simplifier.cc b/third_party/xla/xla/service/while_loop_simplifier.cc index 74c9615b2c254c..aa1cb7b60455cb 100644 --- a/third_party/xla/xla/service/while_loop_simplifier.cc +++ b/third_party/xla/xla/service/while_loop_simplifier.cc @@ -278,6 +278,15 @@ static absl::StatusOr RemoveDeadTupleIndices( CopyOriginalValue(while_init, new_while_init, old_to_new_tuple_idx); CopyOriginalValue(while_op, new_while_op, old_to_new_tuple_idx); + CopyOriginalValue(while_body->parameter_instruction(0), + new_while_op->while_body()->parameter_instruction(0), + old_to_new_tuple_idx); + CopyOriginalValue(while_cond->parameter_instruction(0), + new_while_op->while_condition()->parameter_instruction(0), + old_to_new_tuple_idx); + CopyOriginalValue(while_body_root, + new_while_op->while_body()->root_instruction(), + old_to_new_tuple_idx); // Create a tuple op that recreates the output of the old while op. That is, // we transform to @@ -846,7 +855,9 @@ static absl::StatusOr TryRemoveConstantParams(HloInstruction* while_op) { }; auto add_constant_elems = - [&](HloInstruction* instr) -> std::unique_ptr { + [&](HloInstruction* instr, + const HloInstruction* old_instr = + nullptr) -> std::unique_ptr { CHECK(ShapeUtil::Compatible(instr->shape(), new_while_shape)); std::vector tuple_elems; @@ -861,7 +872,11 @@ static absl::StatusOr TryRemoveConstantParams(HloInstruction* while_op) { ++j; } } - return HloInstruction::CreateTuple(tuple_elems); + auto tuple = HloInstruction::CreateTuple(tuple_elems); + if (old_instr != nullptr) { + tuple->CopyOriginalValue(old_instr); + } + return tuple; }; // Special case: constant_tuple_indices covers the whole while parameter, so @@ -883,17 +898,20 @@ static absl::StatusOr TryRemoveConstantParams(HloInstruction* while_op) { while_cond->CloneWithReplacementPairs({ while_cond->parameter_instruction(0), add_constant_elems(add_new_instr(HloInstruction::CreateParameter( - 0, new_while_shape, - while_cond->parameter_instruction(0)->name()))), + 0, new_while_shape, + while_cond->parameter_instruction(0)->name())), + while_cond->parameter_instruction(0)), }); std::unique_ptr new_while_body = while_body->CloneWithReplacementPairs( { while_body->parameter_instruction(0), - add_constant_elems(add_new_instr(HloInstruction::CreateParameter( - 0, new_while_shape, - while_cond->parameter_instruction(0)->name()))), + add_constant_elems( + add_new_instr(HloInstruction::CreateParameter( + 0, new_while_shape, + while_cond->parameter_instruction(0)->name())), + while_body->parameter_instruction(0)), }, { while_body->root_instruction(), @@ -909,22 +927,29 @@ static absl::StatusOr TryRemoveConstantParams(HloInstruction* while_op) { module->AddEmbeddedComputation(std::move(new_while_cond)), module->AddEmbeddedComputation(std::move(new_while_body)), add_new_instr(remove_constant_elems(while_init)))); - if (while_op->original_value()) { - auto new_ov = std::make_shared(new_while_op->shape()); - int64_t new_i = 0; - for (int i = 0; i < while_shape.tuple_shapes().size(); ++i) { - if (!constant_tuple_indices.count(i)) { - CHECK_OK(new_ov->mutable_tree()->CopyCompatibleSubtreeFrom( - while_op->original_value()->tree(), {i}, {new_i++})); - } + absl::flat_hash_map old_to_new_tuple_idx; + for (int64_t i = 0, new_i = 0; i < while_shape.tuple_shapes().size(); ++i) { + if (!constant_tuple_indices.count(i)) { + old_to_new_tuple_idx[i] = new_i++; } - new_while_op->set_original_value(new_ov); } + CopyOriginalValue(while_init, new_while_op->mutable_operand(0), + old_to_new_tuple_idx); + CopyOriginalValue(while_op, new_while_op, old_to_new_tuple_idx); + CopyOriginalValue(while_body->parameter_instruction(0), + new_while_op->while_body()->parameter_instruction(0), + old_to_new_tuple_idx); + CopyOriginalValue(while_cond->parameter_instruction(0), + new_while_op->while_condition()->parameter_instruction(0), + old_to_new_tuple_idx); + CopyOriginalValue(while_body->root_instruction(), + new_while_op->while_body()->root_instruction(), + old_to_new_tuple_idx); new_while_op->CopyBackendConfigFrom(while_op); CopyFrontendAttributes(while_op, new_while_op); CopyMetadata(while_op, new_while_op); ABSL_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( - while_op, add_constant_elems(new_while_op))); + while_op, add_constant_elems(new_while_op, while_op))); for (auto& instr : new_instrs) { computation->AddInstruction(std::move(instr)); } @@ -1170,11 +1195,11 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { if (HasDisableWhileLoopDceAttr(while_op)) { return false; } - auto flatten_original_value = [&](HloInstruction* old_instr, + auto flatten_original_value = [&](const HloInstruction* old_instr, HloInstruction* new_instr) { if (old_instr->original_value()) { - auto new_original_value = - std::make_shared(new_instr->shape()); + auto new_original_value = std::make_shared( + new_instr->shape(), old_instr->original_value()->call_hierarchy()); int64_t i = 0; for (auto& [shape_index, original_array] : old_instr->original_value()->tree().leaves()) { @@ -1223,7 +1248,8 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { return new_instrs.back().get(); }; - auto nested = [&](HloInstruction* instr) { + auto nested = [&](HloInstruction* instr, + const HloInstruction* old_instr = nullptr) { std::vector gtes; const Shape& flat_shape = instr->shape(); gtes.reserve(flat_shape.tuple_shapes().size()); @@ -1236,6 +1262,9 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { CHECK(ShapeUtil::Compatible(nested_instr->shape(), while_shape)) << ShapeUtil::HumanString(nested_instr->shape()) << " vs " << ShapeUtil::HumanString(while_shape); + if (old_instr != nullptr) { + nested_instr->CopyOriginalValue(old_instr); + } return nested_instr; }; @@ -1249,8 +1278,9 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { while_cond->CloneWithReplacementPairs({ while_cond->parameter_instruction(0), nested(add_new_instr(HloInstruction::CreateParameter( - 0, flattened_shape, - while_cond->parameter_instruction(0)->name()))), + 0, flattened_shape, + while_cond->parameter_instruction(0)->name())), + while_cond->parameter_instruction(0)), }); // Create a new while-body computation, where parameter 0 has a flat shape and @@ -1261,8 +1291,9 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { { while_body->parameter_instruction(0), nested(add_new_instr(HloInstruction::CreateParameter( - 0, flattened_shape, - while_body->parameter_instruction(0)->name()))), + 0, flattened_shape, + while_body->parameter_instruction(0)->name())), + while_body->parameter_instruction(0)), }, { while_body->root_instruction(), @@ -1280,14 +1311,21 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { new_while_op->CopyBackendConfigFrom(while_op); CopyFrontendAttributes(while_op, new_while_op); CopyMetadata(while_op, new_while_op); - ABSL_RETURN_IF_ERROR( - computation->ReplaceWithNewInstruction(while_op, nested(new_while_op))); + flatten_original_value(while_init, new_while_op->mutable_operand(0)); + flatten_original_value(while_op, new_while_op); + flatten_original_value(while_body->parameter_instruction(0), + new_while_op->while_body()->parameter_instruction(0)); + flatten_original_value( + while_cond->parameter_instruction(0), + new_while_op->while_condition()->parameter_instruction(0)); + flatten_original_value(while_body_root, + new_while_op->while_body()->root_instruction()); + ABSL_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( + while_op, nested(new_while_op, while_op))); for (auto& instr : new_instrs) { computation->AddInstruction(std::move(instr)); } - flatten_original_value(while_init, new_while_op->mutable_operand(0)); - flatten_original_value(while_op, new_while_op); return true; } @@ -1416,7 +1454,8 @@ static absl::StatusOr TryMergeInductionVariables( // Converts `instr` into a tuple of the "old" form -- that is, to a tuple with // shape `while_body->shape()` and where the induction variables are "reified" // (i.e. they have value + * ). - auto convert_to_old_form = [&](HloInstruction* instr) { + auto convert_to_old_form = [&](HloInstruction* instr, + const HloInstruction* old_instr = nullptr) { CHECK(ShapeUtil::Compatible(instr->shape(), new_while_shape)); std::vector tuple_elems; for (int i = 0; i < while_shape.tuple_shapes().size(); ++i) { @@ -1433,7 +1472,11 @@ static absl::StatusOr TryMergeInductionVariables( // Copy the original value of the induction variable to its replacement. tuple_elems.back()->CopyOriginalValue(while_body_root->operand(i)); } - return HloInstruction::CreateTuple(tuple_elems); + auto tuple = HloInstruction::CreateTuple(tuple_elems); + if (old_instr != nullptr) { + tuple->CopyOriginalValue(old_instr); + } + return tuple; }; // Converts `root` into a tuple of the "new" form -- that is, to a tuple with @@ -1485,9 +1528,11 @@ static absl::StatusOr TryMergeInductionVariables( std::unique_ptr new_while_cond = while_cond->CloneWithReplacementPairs({ while_cond->parameter_instruction(0), - convert_to_old_form(add_new_instr(HloInstruction::CreateParameter( - 0, new_while_shape, - while_cond->parameter_instruction(0)->name()))), + convert_to_old_form( + add_new_instr(HloInstruction::CreateParameter( + 0, new_while_shape, + while_cond->parameter_instruction(0)->name())), + while_cond->parameter_instruction(0)), }); // Creating the new while body proceeds in two steps. First we convert the @@ -1502,9 +1547,11 @@ static absl::StatusOr TryMergeInductionVariables( HloComputation* temp_new_while_body = module->AddEmbeddedComputation(while_body->CloneWithReplacementPairs({ while_body->parameter_instruction(0), - convert_to_old_form(add_new_instr(HloInstruction::CreateParameter( - 0, new_while_shape, - while_body->parameter_instruction(0)->name()))), + convert_to_old_form( + add_new_instr(HloInstruction::CreateParameter( + 0, new_while_shape, + while_body->parameter_instruction(0)->name())), + while_body->parameter_instruction(0)), })); std::unique_ptr new_while_body = temp_new_while_body->CloneWithReplacementPairs({ @@ -1531,13 +1578,29 @@ static absl::StatusOr TryMergeInductionVariables( if (auto original_value = while_op->original_value()) { new_while->set_original_value(original_value); } + if (auto original_value = + while_op->while_body()->parameter_instruction(0)->original_value()) { + new_while->while_body()->parameter_instruction(0)->set_original_value( + original_value); + } + if (auto original_value = while_op->while_condition() + ->parameter_instruction(0) + ->original_value()) { + new_while->while_condition()->parameter_instruction(0)->set_original_value( + original_value); + } + if (auto original_value = + while_op->while_body()->root_instruction()->original_value()) { + new_while->while_body()->root_instruction()->set_original_value( + original_value); + } if (added_trip_counter) { - AppendToWhileLoopOriginalValue(new_while, {}); + AppendToWhileLoopOriginalValue(new_while, {nullptr}); } CopyFrontendAttributes(while_op, new_while); CopyMetadata(while_op, new_while); ABSL_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( - while_op, convert_to_old_form(new_while))); + while_op, convert_to_old_form(new_while, while_op))); for (auto& instr : new_instrs) { computation->AddInstruction(std::move(instr)); } diff --git a/third_party/xla/xla/service/while_loop_simplifier_test.cc b/third_party/xla/xla/service/while_loop_simplifier_test.cc index 60aa04044f0146..37d7f9a1a3c853 100644 --- a/third_party/xla/xla/service/while_loop_simplifier_test.cc +++ b/third_party/xla/xla/service/while_loop_simplifier_test.cc @@ -1615,15 +1615,15 @@ TEST_F(WhileLoopSimplifierTest, RemoveConstantFromLoopCarryWithOriginalValue) { const std::string hlo_string = R"( HloModule Test Body { - param = (s32[1], s32[2], s32[3]) parameter(0) + param = (s32[1], s32[2], s32[3]) parameter(0), origin={({"bp0"},{"bp1"},{"bp2"})} a = s32[1] get-tuple-element(param), index=0 a.1 = s32[1] add(a, a) b = s32[2] constant({1,1}) c = s32[3] constant({10,10,10}) - ROOT tuple = (s32[1], s32[2], s32[3]) tuple(a.1, b, c) + ROOT tuple = (s32[1], s32[2], s32[3]) tuple(a.1, b, c), origin={({"br0"},{"br1"},{"br2"})} } Cond { - param = (s32[1], s32[2], s32[3]) parameter(0) + param = (s32[1], s32[2], s32[3]) parameter(0), origin={({"cp0"},{"cp1"},{"cp2"})} a = s32[1] get-tuple-element(param), index=0 b = s32[2] get-tuple-element(param), index=1 c = s32[3] get-tuple-element(param), index=2 @@ -1646,6 +1646,30 @@ TEST_F(WhileLoopSimplifierTest, RemoveConstantFromLoopCarryWithOriginalValue) { HloInstruction* while_init = while_instr->while_init(); ASSERT_NE(while_init->original_value(), nullptr); EXPECT_EQ(while_init->original_value()->ToString(), R"(({"a"}, {"c"}))"); + ASSERT_NE( + while_instr->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"bp0"}, {"bp2"}))"); + ASSERT_NE(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"cp0"}, {"cp2"}))"); + ASSERT_NE(while_instr->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->root_instruction() + ->original_value() + ->ToString(), + R"(({"br0"}, {"br2"}))"); HloInstruction* root_instr = m->entry_computation()->root_instruction(); ASSERT_NE(root_instr->original_value(), nullptr); EXPECT_EQ(root_instr->original_value()->ToString(), @@ -1698,7 +1722,7 @@ TEST_F(WhileLoopSimplifierTest, RemoveDeadTupleIndicesWithOriginalValue) { HloModule dus %while.body (arg_tuple: (f32[3], f32[2], f32[2], f32[3], s32[])) -> (f32[3], f32[2], f32[2], f32[3], s32[]) { - %arg_tuple = (f32[3], f32[2], f32[2], f32[3], s32[]) parameter(0) + %arg_tuple = (f32[3], f32[2], f32[2], f32[3], s32[]) parameter(0), origin={({"bp0"}, {"bp1"}, {"bp2"}, {"bp3"}, {"bp4"})} %get-tuple-element.0 = f32[3] get-tuple-element(%arg_tuple), index=0 %get-tuple-element.1 = f32[2] get-tuple-element(%arg_tuple), index=1 %get-tuple-element.2 = f32[2] get-tuple-element(%arg_tuple), index=2 @@ -1710,11 +1734,11 @@ TEST_F(WhileLoopSimplifierTest, RemoveDeadTupleIndicesWithOriginalValue) { %dynamic-update-slice.0 = f32[3] dynamic-update-slice(%get-tuple-element.0, %constant.v0, s32[] %constant.1) %dynamic-update-slice.3 = f32[3] dynamic-update-slice(%get-tuple-element.3, %constant.v0, s32[] %constant.1) %add = add(s32[] %get-tuple-element.4, s32[] %constant.1) - ROOT %tuple = tuple(%dynamic-update-slice.0, %get-tuple-element.1, %get-tuple-element.2, %dynamic-update-slice.3, %add) + ROOT %tuple = tuple(%dynamic-update-slice.0, %get-tuple-element.1, %get-tuple-element.2, %dynamic-update-slice.3, %add), origin={({"br0"}, {"br1"}, {"br2"}, {"br3"}, {"br4"})} } %while.condition (arg_tuple.cond:(f32[3], f32[2], f32[2], f32[3], s32[])) -> pred[] { - %arg_tuple.cond = (f32[3], f32[2], f32[2], f32[3], s32[]) parameter(0) + %arg_tuple.cond = (f32[3], f32[2], f32[2], f32[3], s32[]) parameter(0), origin={({"cp0"}, {"cp1"}, {"cp2"}, {"cp3"}, {"cp4"})} %get-tuple-element.cond = s32[] get-tuple-element(%arg_tuple.cond), index=4 %constant.3 = s32[] constant(3) ROOT %compare = pred[] compare(s32[] %get-tuple-element.cond, s32[] %constant.3), direction=LT @@ -1744,21 +1768,45 @@ ENTRY %main (arg.0: f32[3], arg.1: f32[2]) -> (f32[3], f32[2], f32[2], f32[3]) { ASSERT_NE(while_init->original_value(), nullptr); EXPECT_EQ(while_init->original_value()->ToString(), R"(({"arg.0"}, {"arg.1"}, {"constant.0"}))"); + ASSERT_NE( + while_instr->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"bp0"}, {"bp1"}, {"bp4"}))"); + ASSERT_NE(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"cp0"}, {"cp1"}, {"cp4"}))"); + ASSERT_NE(while_instr->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->root_instruction() + ->original_value() + ->ToString(), + R"(({"br0"}, {"br1"}, {"br4"}))"); } TEST_F(WhileLoopSimplifierTest, FlattenNestedTupleWithOriginalValue) { const std::string hlo_string = R"( HloModule Test Body { - param = ((s32[1]), (s32[2], s32[3], (s32[4]))) parameter(0) + param = ((s32[1]), (s32[2], s32[3], (s32[4]))) parameter(0), origin={(({"bp0"}), ({"bp1"}, {"bp2"}, ({"bp3"})))} ta = (s32[1]) get-tuple-element(param), index=0 a = s32[1] get-tuple-element(ta), index=0 a.1 = s32[1] add(a, a) tbcd = (s32[2], s32[3], (s32[4])) get-tuple-element(param), index=1 - ROOT tuple = ((s32[1]), (s32[2], s32[3], (s32[4]))) tuple(ta, tbcd) + ROOT tuple = ((s32[1]), (s32[2], s32[3], (s32[4]))) tuple(ta, tbcd), origin={(({"br0"}), ({"br1"}, {"br2"}, ({"br3"})))} } Cond { - param = ((s32[1]), (s32[2], s32[3], (s32[4]))) parameter(0) + param = ((s32[1]), (s32[2], s32[3], (s32[4]))) parameter(0), origin={(({"cp0"}), ({"cp1"}, {"cp2"}, ({"cp3"})))} ROOT cond = pred[] constant(true) } ENTRY Loop { @@ -1773,7 +1821,7 @@ TEST_F(WhileLoopSimplifierTest, FlattenNestedTupleWithOriginalValue) { {"b"}, {"c"}, ({"d"})))} ROOT while = ((s32[1]), (s32[2], s32[3], (s32[4]))) while(init), condition=Cond, body=Body, origin={(({"while.116" {0}}), ( - {"while.116" {1}}, {"while.116" {2}}, ({"while.116" {3}})))} + {"while.116" {1}}, {"while.116" {2}}, ({"while.116" {3}}))),["while#$"]} })"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); @@ -1783,17 +1831,41 @@ TEST_F(WhileLoopSimplifierTest, FlattenNestedTupleWithOriginalValue) { ASSERT_NE(while_instr->original_value(), nullptr); EXPECT_EQ( while_instr->original_value()->ToString(), - R"(({"while.116" {0}}, {"while.116" {1}}, {"while.116" {2}}, {"while.116" {3}}))"); + R"(({"while.116" {0}}, {"while.116" {1}}, {"while.116" {2}}, {"while.116" {3}}),["while#$"])"); HloInstruction* while_init = while_instr->while_init(); ASSERT_NE(while_init->original_value(), nullptr); EXPECT_EQ(while_init->original_value()->ToString(), R"(({"a"}, {"b"}, {"c"}, {"d"}))"); + ASSERT_NE( + while_instr->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"bp0"}, {"bp1"}, {"bp2"}, {"bp3"}))"); + ASSERT_NE(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"cp0"}, {"cp1"}, {"cp2"}, {"cp3"}))"); + ASSERT_NE(while_instr->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->root_instruction() + ->original_value() + ->ToString(), + R"(({"br0"}, {"br1"}, {"br2"}, {"br3"}))"); } const char* const kSimpleMergeInductionVariablesModuleWithOriginalValue = R"( HloModule Test Body { - param = (TYPE[], TYPE[], TYPE[]) parameter(0) + param = (TYPE[], TYPE[], TYPE[]) parameter(0), origin={({"bp0"}, {"bp1"}, {"bp2"})} a = TYPE[] get-tuple-element(param), index=0 one = TYPE[] constant(1) @@ -1805,10 +1877,10 @@ const char* const kSimpleMergeInductionVariablesModuleWithOriginalValue = R"( c = TYPE[] add(a, b) - ROOT tuple = (TYPE[], TYPE[], TYPE[]) tuple(a1,b1,c) + ROOT tuple = (TYPE[], TYPE[], TYPE[]) tuple(a1,b1,c), origin={({"br0"}, {"br1"}, {"br2"})} } Cond { - param = (TYPE[], TYPE[], TYPE[]) parameter(0) + param = (TYPE[], TYPE[], TYPE[]) parameter(0), origin={({"cp0"}, {"cp1"}, {"cp2"})} a = TYPE[] get-tuple-element(param), index=0 b = TYPE[] get-tuple-element(param), index=1 sum = TYPE[] power(a, b) @@ -1843,6 +1915,30 @@ TEST_F(WhileLoopSimplifierTest, MergeInductionVariablesWithOriginalValue) { ASSERT_NE(while_init->original_value(), nullptr); EXPECT_EQ(while_init->original_value()->ToString(), R"(({"a"}, {"b"}, {"c"}, {}))"); + ASSERT_NE( + while_instr->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"bp0"}, {"bp1"}, {"bp2"}, {}))"); + ASSERT_NE(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"cp0"}, {"cp1"}, {"cp2"}, {}))"); + ASSERT_NE(while_instr->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->root_instruction() + ->original_value() + ->ToString(), + R"(({"br0"}, {"br1"}, {"br2"}, {}))"); const HloInstruction* add = module->entry_computation()->root_instruction()->operand(0); const HloInstruction* induction_var_0 = add->operand(0); diff --git a/third_party/xla/xla/service/while_util.cc b/third_party/xla/xla/service/while_util.cc index 8e4c69ea77e656..0dc7ebd1b15889 100644 --- a/third_party/xla/xla/service/while_util.cc +++ b/third_party/xla/xla/service/while_util.cc @@ -285,40 +285,33 @@ WhileUtil::MakeInstructionsLiveIn( HloInstruction* new_while = while_instr->AddInstruction( HloInstruction::CreateWhile(new_while_shape, new_while_condition, new_while_body, new_while_init)); - if (while_instr->original_value() != nullptr) { - OriginalValue new_original_value( - new_while_shape, while_instr->original_value()->call_hierarchy()); - for (auto& [shape_index, original_array] : - new_original_value.mutable_original_arrays()) { - // The hoisted instructions are appended to the end of the while - // instruction, so the shape indices smaller than - // `elements_in_old_while_shape` are all the old original arrays that - // need to be propagated. - if (shape_index[0] < elements_in_old_while_shape) { - original_array = - while_instr->original_value()->tree().element(shape_index); - } else { - // If the element is new (i.e., hoisted), fetch its original value from - // the instruction that was hoisted. - int instruction_idx = shape_index[0] - elements_in_old_while_shape; - HloInstruction* instruction = instructions[instruction_idx]; - if (instruction->original_value() != nullptr) { - ShapeIndex shape_index_in_instruction = shape_index; - shape_index_in_instruction.erase(shape_index_in_instruction.begin()); - original_array = instruction->original_value()->original_array( - shape_index_in_instruction); - } - } - } - new_while->set_original_value( - std::make_shared(std::move(new_original_value))); + new_while_init->set_original_value( + while_instr->while_init()->original_value()); + new_while->set_original_value(while_instr->original_value()); + if (while_instr->while_body()->num_parameters() == 1) { + CHECK_EQ(new_while_body->num_parameters(), 1); + new_while_body->parameter_instruction(0)->set_original_value( + while_instr->while_body()->parameter_instruction(0)->original_value()); } + if (while_instr->while_condition()->num_parameters() == 1) { + CHECK_EQ(new_while_condition->num_parameters(), 1); + new_while_condition->parameter_instruction(0)->set_original_value( + while_instr->while_condition() + ->parameter_instruction(0) + ->original_value()); + } + new_while_body->root_instruction()->set_original_value( + while_instr->while_body()->root_instruction()->original_value()); + AppendToWhileLoopOriginalValue( + new_while, HloInstruction::InstructionVector(instructions.begin(), + instructions.end())); // We want to get rid of the old while instruction even if it has side // effecting operations so we do a manual HloComputation::RemoveInstruction // instead of relying on HloComputation::ReplaceInstruction. HloInstruction* replacement_instr = TupleUtil::ExtractPrefix( new_while, while_instr->shape().tuple_shapes().size()); + replacement_instr->CopyOriginalValue(while_instr); ABSL_RETURN_IF_ERROR(new_while->CopyAllControlDepsFrom(while_instr)); ABSL_RETURN_IF_ERROR(while_instr->DropAllControlDeps()); ABSL_RETURN_IF_ERROR(while_instr->ReplaceAllUsesWith(replacement_instr)); @@ -328,12 +321,15 @@ WhileUtil::MakeInstructionsLiveIn( std::vector live_in_instructions; for (int64_t i = elements_in_old_while_shape; i < new_while_shape.tuple_shapes().size(); i++) { - live_in_instructions.push_back(new_while_body->AddInstruction( + HloInstruction* live_in_gte = new_while_body->AddInstruction( HloInstruction::CreateGetTupleElement( instructions[i - elements_in_old_while_shape]->shape(), while_body_param, i), absl::StrCat(new_while_body->name(), ".in.", - i - elements_in_old_while_shape))); + i - elements_in_old_while_shape)); + live_in_gte->CopyOriginalValue( + instructions[i - elements_in_old_while_shape]); + live_in_instructions.push_back(live_in_gte); } WhileUtil::MakeInstructionsLiveInResult result; @@ -650,13 +646,13 @@ absl::Status WhileUtil::IncrementWhileLoopTripCount( void AppendToWhileLoopOriginalValue( HloInstruction* while_instr, const HloInstruction::InstructionVector& new_while_input_tuple_elements) { + absl::flat_hash_set updated_instrs; auto append_to_original_value = [&](HloInstruction* instr, int64_t next_index) { - std::shared_ptr old_original_value = instr->original_value(); - if (old_original_value != nullptr && - old_original_value->IsCompatibleWith(instr->shape())) { + if (instr == nullptr || !updated_instrs.insert(instr).second) { return; } + std::shared_ptr old_original_value = instr->original_value(); // Returns if neither the instruction nor any of its new tuple elements have // an original value. @@ -666,7 +662,8 @@ void AppendToWhileLoopOriginalValue( new_while_input_tuple_elements.end(), [&has_original_value](const HloInstruction* instr) { has_original_value |= - (instr->original_value() != nullptr && + (instr != nullptr && + instr->original_value() != nullptr && !instr->original_value()->IsEmpty()); }); if (!has_original_value) { @@ -674,27 +671,59 @@ void AppendToWhileLoopOriginalValue( } } + std::optional call_hierarchy; + if (old_original_value != nullptr) { + call_hierarchy = old_original_value->call_hierarchy(); + } else if (instr->opcode() == HloOpcode::kWhile) { + call_hierarchy = absl::StrCat(instr->name(), "#$"); + } + std::shared_ptr new_original_value = - std::make_shared(instr->shape()); + std::make_shared(instr->shape(), + std::move(call_hierarchy)); if (old_original_value != nullptr) { if (!old_original_value->IsTuple()) { - new_original_value->mutable_tree()->CopySubtreeFrom( - old_original_value->tree(), {}, {0}); + CHECK_EQ(next_index, 1) + << "Expected next_index == 1 when converting non-tuple while loop " + "to tuple: " + << instr->ToString(); + CHECK_OK(new_original_value->mutable_tree()->CopyCompatibleSubtreeFrom( + old_original_value->tree(), {}, {0})) + << "Incompatible OriginalValue subtree when converting non-tuple " + "while loop to tuple: " + << instr->ToString(); } else { - for (auto& [shape_index, original_array] : old_original_value->tree()) { - *new_original_value->mutable_original_array(shape_index) = - original_array; + int64_t old_tuple_size = old_original_value->tree().num_children(); + CHECK(old_tuple_size == next_index || + old_tuple_size == instr->shape().tuple_shapes().size()) + << "Unexpected old OriginalValue tuple size " << old_tuple_size + << " vs next_index " << next_index << " and widened tuple size " + << instr->shape().tuple_shapes().size() + << " for instruction: " << instr->ToString(); + for (int64_t i = 0; i < old_tuple_size; ++i) { + CHECK_OK( + new_original_value->mutable_tree()->CopyCompatibleSubtreeFrom( + old_original_value->tree(), {i}, {i})) + << "Incompatible OriginalValue subtree at index " << i << ": " + << instr->ToString(); } } } for (int64_t i = 0; i < new_while_input_tuple_elements.size(); ++i) { - if (new_while_input_tuple_elements[i]->original_value() != nullptr) { - new_original_value->mutable_tree()->CopySubtreeFrom( + if (new_while_input_tuple_elements[i] != nullptr && + new_while_input_tuple_elements[i]->original_value() != nullptr) { + CHECK_OK(new_original_value->mutable_tree()->CopyCompatibleSubtreeFrom( new_while_input_tuple_elements[i]->original_value()->tree(), {}, - {next_index + i}); + {next_index + i})) + << "Incompatible OriginalValue subtree for appended while input " + "element " + << i << ": " << new_while_input_tuple_elements[i]->ToString(); } } + CHECK(new_original_value->IsCompatibleWith(instr->shape())) + << "OriginalValue is incompatible with instruction shape: " + << instr->ToString(); instr->set_original_value(new_original_value); }; @@ -711,8 +740,20 @@ void AppendToWhileLoopOriginalValue( while_shape.tuple_shapes().size() - new_while_input_tuple_elements.size(); append_to_original_value(while_instr->while_init(), next_index); append_to_original_value(while_instr, next_index); - append_to_original_value(while_instr->while_body()->root_instruction(), - next_index); + HloInstruction* body_root = while_instr->while_body()->root_instruction(); + if (body_root->shape().IsTuple() && + body_root->shape().tuple_shapes().size() == + while_shape.tuple_shapes().size()) { + append_to_original_value(body_root, next_index); + } + if (while_instr->while_body()->num_parameters() == 1) { + append_to_original_value( + while_instr->while_body()->parameter_instruction(0), next_index); + } + if (while_instr->while_condition()->num_parameters() == 1) { + append_to_original_value( + while_instr->while_condition()->parameter_instruction(0), next_index); + } } bool WhileUtil::IsUpdatedBufferWriteOnly(const HloInstruction* instr) { diff --git a/third_party/xla/xla/service/while_util_test.cc b/third_party/xla/xla/service/while_util_test.cc index 7eefce7a2855b6..5644f9771f4f94 100644 --- a/third_party/xla/xla/service/while_util_test.cc +++ b/third_party/xla/xla/service/while_util_test.cc @@ -26,10 +26,12 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_original_value.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/test.h" #include "xla/hlo/testlib/verified_hlo_module.h" #include "xla/hlo/utils/hlo_matchers.h" +#include "xla/shape.h" #include "xla/tsl/platform/logging.h" #include "xla/util.h" @@ -503,5 +505,204 @@ ENTRY entry { EXPECT_FALSE(WhileUtil::IsUpdatedBufferWriteOnly(dus4)); } +TEST_F(WhileUtilTest, + AppendToWhileLoopOriginalValueUpdatesAllLoopInstructions) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +while_body { + ROOT p_body = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} +} + +while_condition { + p_cond = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} + ROOT result = pred[] constant(true) +} + +ENTRY entry { + p0 = f32[] parameter(0), origin={{"p0"}} + p1 = f32[] parameter(1), origin={{"p1"}} + p2 = f32[] parameter(2), origin={{"p2"}} + init = (f32[], f32[]) tuple(p0, p1), origin={({"p0"}, {"p1"})} + while = (f32[], f32[]) while(init), condition=while_condition, body=while_body, origin={({"p0"}, {"p1"}),["while#$"]} + gte0 = f32[] get-tuple-element(while), index=0 + ROOT root = (f32[]) tuple(gte0) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloInstruction* while_instr = FindInstruction(module.get(), "while"); + HloInstruction* p2 = module->entry_computation()->parameter_instruction(2); + + HloComputation* body = while_instr->while_body(); + HloComputation* cond = while_instr->while_condition(); + + Shape widened_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeScalarShape(F32), ShapeUtil::MakeScalarShape(F32), + ShapeUtil::MakeScalarShape(F32)}); + *while_instr->mutable_shape() = widened_shape; + *body->parameter_instruction(0)->mutable_shape() = widened_shape; + *cond->parameter_instruction(0)->mutable_shape() = widened_shape; + *body->root_instruction()->mutable_shape() = widened_shape; + + HloInstruction* p0 = module->entry_computation()->parameter_instruction(0); + HloInstruction* p1 = module->entry_computation()->parameter_instruction(1); + HloInstruction* new_init = module->entry_computation()->AddInstruction( + HloInstruction::CreateTuple({p0, p1, p2})); + ASSERT_OK(while_instr->ReplaceOperandWithDifferentShape(0, new_init)); + + AppendToWhileLoopOriginalValue(while_instr, {p2}); + + ASSERT_NE(while_instr->original_value(), nullptr); + EXPECT_TRUE(while_instr->original_value()->IsCompatibleWith(widened_shape)); + EXPECT_THAT(while_instr->original_value()->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + ASSERT_TRUE(while_instr->original_value()->call_hierarchy().has_value()); + EXPECT_EQ(*while_instr->original_value()->call_hierarchy(), "while#$"); + + ASSERT_NE(body->parameter_instruction(0)->original_value(), nullptr); + EXPECT_TRUE( + body->parameter_instruction(0)->original_value()->IsCompatibleWith( + widened_shape)); + EXPECT_THAT( + body->parameter_instruction(0)->original_value()->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + + ASSERT_NE(cond->parameter_instruction(0)->original_value(), nullptr); + EXPECT_TRUE( + cond->parameter_instruction(0)->original_value()->IsCompatibleWith( + widened_shape)); + EXPECT_THAT( + cond->parameter_instruction(0)->original_value()->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + + ASSERT_NE(body->root_instruction()->original_value(), nullptr); + EXPECT_TRUE(body->root_instruction()->original_value()->IsCompatibleWith( + widened_shape)); +} + +TEST_F(WhileUtilTest, + AppendToWhileLoopOriginalValueIncompatibleSubtreeCrashes) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +while_body { + ROOT p_body = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} +} + +while_condition { + p_cond = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} + ROOT result = pred[] constant(true) +} + +ENTRY entry { + p0 = f32[] parameter(0), origin={{"p0"}} + p1 = f32[] parameter(1), origin={{"p1"}} + p2 = (f32[], f32[]) parameter(2), origin={({"p2_0"}, {"p2_1"})} + init = (f32[], f32[]) tuple(p0, p1), origin={({"p0"}, {"p1"})} + while = (f32[], f32[]) while(init), condition=while_condition, body=while_body, origin={({"p0"}, {"p1"}),["while#$"]} + gte0 = f32[] get-tuple-element(while), index=0 + ROOT root = (f32[]) tuple(gte0) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloInstruction* while_instr = FindInstruction(module.get(), "while"); + HloInstruction* p2 = module->entry_computation()->parameter_instruction(2); + + EXPECT_DEATH( + { + Shape widened_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeScalarShape(F32), ShapeUtil::MakeScalarShape(F32), + ShapeUtil::MakeScalarShape(F32)}); + *while_instr->mutable_shape() = widened_shape; + AppendToWhileLoopOriginalValue(while_instr, {p2}); + }, + "Incompatible OriginalValue subtree for appended while input element 0"); +} + +TEST_F(WhileUtilTest, MakeInstructionsLiveInPropagatesOriginalValue) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +while_body { + ROOT p_body = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} +} + +while_condition { + p_cond = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} + ROOT result = pred[] constant(true) +} + +ENTRY entry { + p0 = f32[] parameter(0), origin={{"p0"}} + p1 = f32[] parameter(1), origin={{"p1"}} + p2 = f32[] parameter(2), origin={{"p2"}} + init = (f32[], f32[]) tuple(p0, p1), origin={({"p0"}, {"p1"})} + while = (f32[], f32[]) while(init), condition=while_condition, body=while_body, origin={({"p0"}, {"p1"}),["while#$"]} + gte0 = f32[] get-tuple-element(while), index=0 + ROOT root = (f32[]) tuple(gte0) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloInstruction* while_instr = FindInstruction(module.get(), "while"); + HloInstruction* p2 = module->entry_computation()->parameter_instruction(2); + + ASSERT_OK_AND_ASSIGN(WhileUtil::MakeInstructionsLiveInResult result, + WhileUtil::MakeInstructionsLiveIn(while_instr, {p2})); + + HloInstruction* new_while = result.new_while_instr; + ASSERT_NE(new_while->original_value(), nullptr); + EXPECT_TRUE( + new_while->original_value()->IsCompatibleWith(new_while->shape())); + EXPECT_THAT(new_while->original_value()->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + + ASSERT_NE(new_while->while_init()->original_value(), nullptr); + EXPECT_TRUE(new_while->while_init()->original_value()->IsCompatibleWith( + new_while->while_init()->shape())); + EXPECT_THAT(new_while->while_init()->original_value()->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + + ASSERT_NE(new_while->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_TRUE( + new_while->while_body() + ->parameter_instruction(0) + ->original_value() + ->IsCompatibleWith( + new_while->while_body()->parameter_instruction(0)->shape())); + EXPECT_THAT(new_while->while_body() + ->parameter_instruction(0) + ->original_value() + ->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + + ASSERT_NE( + new_while->while_condition()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_TRUE( + new_while->while_condition() + ->parameter_instruction(0) + ->original_value() + ->IsCompatibleWith( + new_while->while_condition()->parameter_instruction(0)->shape())); + + ASSERT_NE(new_while->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_TRUE(new_while->while_body() + ->root_instruction() + ->original_value() + ->IsCompatibleWith( + new_while->while_body()->root_instruction()->shape())); + + ASSERT_NE(result.replacement_instr->original_value(), nullptr); + EXPECT_TRUE(result.replacement_instr->original_value()->IsCompatibleWith( + result.replacement_instr->shape())); + ASSERT_EQ(result.while_body_live_in_values.size(), 1); + ASSERT_NE(result.while_body_live_in_values[0]->original_value(), nullptr); + EXPECT_THAT( + result.while_body_live_in_values[0]->original_value()->original_array({}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/tuple_tree.h b/third_party/xla/xla/tuple_tree.h index 75671f26320a3f..37e7668bb05fb5 100644 --- a/third_party/xla/xla/tuple_tree.h +++ b/third_party/xla/xla/tuple_tree.h @@ -357,6 +357,17 @@ class TupleTree { return entry_or.value()->children_start_id == -1; } + // Returns the number of children of the node at the given index, or 0 if the + // node does not exist or is a leaf. + size_t num_children(ShapeIndexView index = {}) const { + absl::StatusOr entry_or = + index_table_.GetEntry(index); + if (!entry_or.ok() || entry_or.value()->children_start_id == -1) { + return 0; + } + return entry_or.value()->num_children; + } + // Checks if the structure of this TupleTree is compatible with the given // shape. bool IsStructurallyCompatible(const Shape& shape) const { diff --git a/third_party/xla/xla/tuple_tree_test.cc b/third_party/xla/xla/tuple_tree_test.cc index 5741d694424682..353124d76123ae 100644 --- a/third_party/xla/xla/tuple_tree_test.cc +++ b/third_party/xla/xla/tuple_tree_test.cc @@ -46,12 +46,14 @@ namespace { TEST_F(TupleTreeTest, SingleLeafConstructor) { TupleTree tree(42); EXPECT_TRUE(tree.IsLeaf({})); + EXPECT_EQ(tree.num_children(), 0); EXPECT_EQ(tree.element({}), 42); } TEST_F(TupleTreeTest, EmptyConstructor) { TupleTree tree; EXPECT_FALSE(tree.IsLeaf({})); + EXPECT_EQ(tree.num_children(), 0); EXPECT_THAT(tree.nodes(), ElementsAre(Pair(ShapeIndex({}), 0))); EXPECT_THAT(tree.leaves(), ElementsAre()); EXPECT_EQ(tree.element({}), 0); @@ -71,12 +73,15 @@ TEST_F(TupleTreeTest, NodeWithValueAndChildren) { // Node with value and empty children (empty tuple) TupleTree tree1(Node::Tuple(100)); EXPECT_FALSE(tree1.IsLeaf({})); // It's a tuple, not a leaf + EXPECT_EQ(tree1.num_children(), 0); EXPECT_THAT(tree1.nodes(), ElementsAre(Pair(ShapeIndex({}), 100))); EXPECT_THAT(tree1.leaves(), ElementsAre()); // Node with value and non-empty children TupleTree tree2(Node::Tuple(200, {Node::Leaf(1), Node::Leaf(2)})); EXPECT_FALSE(tree2.IsLeaf({})); + EXPECT_EQ(tree2.num_children(), 2); + EXPECT_EQ(tree2.num_children({0}), 0); EXPECT_TRUE(tree2.IsLeaf({0})); EXPECT_EQ(tree2.element({0}), 1); EXPECT_TRUE(tree2.IsLeaf({1})); From a1ae4db3a82863dee36b94de43f5881627a75911 Mon Sep 17 00:00:00 2001 From: Sohaib Iftikhar Date: Thu, 24 Sep 2026 10:26:13 -0700 Subject: [PATCH 38/48] [XLA:GPU]: Fix caching issue with kernels when unrolling while loops When while loops are unrolled the same thunk is called with `Record` multiple times for different loop iterations. Currently `CustomCallRecordState` stored a single sequence of `XLA_FFI_Command*`, so during update only the last recorded sequence of commands was passed to the FFI record API, overwriting the last iteration and leaving earlier iterations of the while loop with stale buffer pointers. Keying the state by the sink command pointer returned on `RecordCreate` (and passed back via `RecordUpdate::command`) so each unrolled iteration updates its own recorded command sequence. PiperOrigin-RevId: 987596254 --- .../backends/gpu/runtime/custom_call_thunk.cc | 35 ++++---- .../gpu/runtime/custom_call_thunk_test.cc | 90 +++++++++++++++++++ 2 files changed, 109 insertions(+), 16 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc index 2a2bc8f7f20a3b..0ecd9b0d940d1e 100644 --- a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc @@ -91,7 +91,9 @@ std::string GetSymbolName(const void* ptr) { } struct CustomCallRecordState : public CommandState { - std::vector commands; + absl::flat_hash_map> + commands; }; // A per-execution state that holds state for prepare and initialize stages. @@ -606,13 +608,14 @@ absl::StatusOr CustomCallThunk::Record( const bool is_record_create = std::holds_alternative(record_action); - const bool is_record_update = - std::holds_alternative(record_action); - if (is_record_update) { // Copy over commands from state to inline storage. - TF_RET_CHECK(state->commands.size() <= kMaxCommands) - << "Too many commands to fit in inline storage"; - std::copy(state->commands.begin(), state->commands.end(), commands_storage); - num_commands = state->commands.size(); + if (const auto* record_update = std::get_if(&record_action)) { + if (auto it = state->commands.find(record_update->command); + it != state->commands.end()) { + TF_RET_CHECK(it->second.size() <= kMaxCommands) + << "Too many commands to fit in inline storage"; + std::copy(it->second.begin(), it->second.end(), commands_storage); + num_commands = it->second.size(); + } } XLA_FFI_RecordAction action_to_pass = is_record_create @@ -679,19 +682,19 @@ absl::StatusOr CustomCallThunk::Record( command_buffer); } - // Save newly recorded commands to state if this is the Create action - // Must be done after returning from the FFI handler. - if (is_record_create) { - state->commands.assign(commands_storage, commands_storage + num_commands); - } - // Return the last command in the chain for dependency tracking. // If more than one command was recorded, and they are independent, a dummy // node must be added to the command graph by the FFI client so that XLA // can track a single dependency for the entire chain. if (num_commands > 0 && commands_storage[num_commands - 1] != nullptr) { - return reinterpret_cast( - commands_storage[num_commands - 1]); + const auto* sink_command = + reinterpret_cast( + commands_storage[num_commands - 1]); + if (is_record_create) { + state->commands[sink_command].assign(commands_storage, + commands_storage + num_commands); + } + return sink_command; } // No commands were recorded. return nullptr; diff --git a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk_test.cc index 63b6b6a03e9daf..459c6b35f04648 100644 --- a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk_test.cc @@ -1175,5 +1175,95 @@ TEST(CustomCallThunkTest, RecordCommandBufferFfiRecordWithEmptyCommand) { EXPECT_EQ(host_update, 90); } +TEST(CustomCallThunkTest, RecordCommandBufferMultipleRecordsInSameBuffer) { + ASSERT_OK_AND_ASSIGN(se::StreamExecutor * executor, GpuExecutor()); + if (executor->GetDeviceDescription().gpu_compute_capability().IsRocm()) { + GTEST_SKIP() << "AddI32 PTX kernel not supported on ROCm."; + } + + CustomCallThunk::OwnedHandlerBundle bundle; + bundle.execute = + ffi::Ffi::BindExecute().To([]() { return absl::OkStatus(); }); + bundle.record = ffi::Ffi::BindRecord() + .Ctx>() + .Arg() + .Arg() + .Ret() + .To(AddI32FfiHandler); + + ASSERT_OK_AND_ASSIGN(auto setup, FfiRecordTestSetup::Create( + executor, std::move(bundle), + "add_i32_ffi_unroll", {10, 20}, {0})); + + RecordTestAlloc alloc_iter1_create(executor); + ASSERT_OK_AND_ASSIGN( + auto slices_iter1_create, + AllocateAndCopy(*setup->stream, alloc_iter1_create, {1, 2}, {0})); + Thunk::ExecuteParams execute_params_iter1_create = + Thunk::ExecuteParams::Create( + ServiceExecutableRunOptions(), *alloc_iter1_create.buffer_allocations, + setup->stream.get(), setup->stream.get(), nullptr, nullptr, nullptr); + + ASSERT_OK_AND_ASSIGN(auto cb, executor->CreateCommandBuffer( + se::CommandBuffer::Mode::kPrimary)); + + // Record two iterations of the same CustomCallThunk into the same command + // buffer using the same CommandStateManager (as unrolled WhileThunk does). + ASSERT_OK_AND_ASSIGN( + const se::CommandBuffer::Command* cmd_0, + setup->thunk->Record(*setup->execute_params, *setup->record_params, + Command::RecordCreate{/*dependencies=*/{}}, + cb.get())); + ASSERT_OK_AND_ASSIGN( + const se::CommandBuffer::Command* cmd_1, + setup->thunk->Record(execute_params_iter1_create, *setup->record_params, + Command::RecordCreate{/*dependencies=*/{cmd_0}}, + cb.get())); + ASSERT_NE(cmd_0, cmd_1); + + ASSERT_OK(cb->Finalize()); + ASSERT_OK(cb->Submit(setup->stream.get())); + ASSERT_OK(setup->stream->BlockHostUntilDone()); + + // Update both recorded nodes with new buffer allocations. + ASSERT_OK(cb->Update()); + RecordTestAlloc alloc_iter0_update(executor); + ASSERT_OK_AND_ASSIGN( + auto slices_iter0_update, + AllocateAndCopy(*setup->stream, alloc_iter0_update, {40, 50}, {0})); + Thunk::ExecuteParams execute_params_iter0_update = + Thunk::ExecuteParams::Create( + ServiceExecutableRunOptions(), *alloc_iter0_update.buffer_allocations, + setup->stream.get(), setup->stream.get(), nullptr, nullptr, nullptr); + + RecordTestAlloc alloc_iter1_update(executor); + ASSERT_OK_AND_ASSIGN( + auto slices_iter1_update, + AllocateAndCopy(*setup->stream, alloc_iter1_update, {100, 200}, {0})); + Thunk::ExecuteParams execute_params_iter1_update = + Thunk::ExecuteParams::Create( + ServiceExecutableRunOptions(), *alloc_iter1_update.buffer_allocations, + setup->stream.get(), setup->stream.get(), nullptr, nullptr, nullptr); + + ASSERT_OK(setup->thunk->Record(execute_params_iter0_update, + *setup->record_params, + Command::RecordUpdate{cmd_0}, cb.get())); + ASSERT_OK(setup->thunk->Record(execute_params_iter1_update, + *setup->record_params, + Command::RecordUpdate{cmd_1}, cb.get())); + ASSERT_OK(cb->Finalize()); + ASSERT_OK(cb->Submit(setup->stream.get())); + ASSERT_OK(setup->stream->BlockHostUntilDone()); + + int32_t out_0 = 0; + int32_t out_1 = 0; + ASSERT_OK(setup->stream->Memcpy(&out_0, alloc_iter0_update.result_dev_ptrs[0], + sizeof(int32_t))); + ASSERT_OK(setup->stream->Memcpy(&out_1, alloc_iter1_update.result_dev_ptrs[0], + sizeof(int32_t))); + EXPECT_EQ(out_0, 90); + EXPECT_EQ(out_1, 300); +} + } // namespace } // namespace xla::gpu From 659f756d1ab50d4ec07797679a037ee693fa74ed Mon Sep 17 00:00:00 2001 From: Yue Sheng Date: Thu, 24 Sep 2026 10:32:07 -0700 Subject: [PATCH 39/48] [Mosaic TPU] Add some fold rules for the transpose op. Reverts 2a3525476e35c62dcd1652be92985b42e9523dfc PiperOrigin-RevId: 987600180 --- .../xla/xla/mosaic/dialect/tpu/tpu_ops.cc | 22 +++++++++++++++++++ .../xla/xla/mosaic/dialect/tpu/tpu_ops.td | 1 + 2 files changed, 23 insertions(+) diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc index 6cf17ae96335fa..09432b11e5d080 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -955,6 +955,28 @@ LogicalResult TransposeOp::verify() { return success(); } +OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) { + if (llvm::is_sorted(getPermutation())) { + return getVector(); + } + if (auto cst = dyn_cast_if_present(adaptor.getVector())) { + return cst.reshape(getType()); + } + if (auto prev_op = getVector().getDefiningOp()) { + ArrayRef prev_perm = prev_op.getPermutation(); + ArrayRef curr_perm = getPermutation(); + CHECK_EQ(prev_perm.size(), curr_perm.size()); + SmallVector composed_perm(curr_perm.size()); + for (size_t i = 0; i < curr_perm.size(); ++i) { + composed_perm[i] = prev_perm[curr_perm[i]]; + } + if (llvm::is_sorted(composed_perm)) { + return prev_op.getVector(); + } + } + return nullptr; +} + LogicalResult MemRefBitcastOp::verify() { auto src_ty = getInput().getType(); auto tgt_ty = getType(); diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td index e23c2715ba515f..1130919108db8c 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td @@ -1921,6 +1921,7 @@ def TPU_TransposeOp : TPU_Op<"transpose", [Pure]> { } }]; let hasVerifier = 1; + let hasFolder = 1; } def TPU_LogOp : TPU_Op<"log"> { From 5442ef5943436b6cd8d3e43673aaaa9f3f08edc0 Mon Sep 17 00:00:00 2001 From: Majid Dadashi Date: Thu, 24 Sep 2026 11:06:35 -0700 Subject: [PATCH 40/48] =?UTF-8?q?Add=20blockwise=20support=C2=A0in=20the?= =?UTF-8?q?=20litert=20converter.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Also includes a quant_spec metadata in tfl.fully_connected to allow encoding custom specs. PiperOrigin-RevId: 987624321 --- RELEASE.md | 3 +- tensorflow/compiler/mlir/lite/BUILD | 2 + .../compiler/mlir/lite/converter_gen.cc | 127 ++++---- .../lite/core/api/flatbuffer_conversions.cc | 5 + .../mlir/lite/core/c/builtin_op_data.h | 12 + .../compiler/mlir/lite/flatbuffer_operator.cc | 76 ++++- tensorflow/compiler/mlir/lite/ir/tfl_ops.cc | 115 +++++++ tensorflow/compiler/mlir/lite/ir/tfl_ops.h | 6 + tensorflow/compiler/mlir/lite/ir/tfl_ops.td | 101 ++++++ .../lite/python/stablehlo_tfl_pipeline.cc | 2 + .../compiler/mlir/lite/schema/schema.fbs | 12 + .../mlir/lite/schema/schema_generated.h | 40 ++- .../mlir/lite/tests/blockwise-ops.mlir | 126 ++++++++ .../mlir/lite/tools/versioning/op_version.cc | 9 + .../lite/tools/versioning/runtime_version.cc | 1 + ...fold_stablehlo_constant_transforms_pass.cc | 277 ++++++++++++++++ .../transforms/large_constant_fold_pass.cc | 161 ++++++++++ .../lower_quant_annotations_helper.h | 304 +++++++++++------- .../lower_quant_annotations_pass.cc | 299 ++++++++++++++++- .../transforms/optimize_batch_matmul_pass.cc | 2 +- .../compiler/mlir/lite/transforms/passes.h | 3 + .../compiler/mlir/lite/transforms/passes.td | 15 + .../lite/core/api/flatbuffer_conversions.cc | 5 + .../delegates/xnnpack/xnnpack_delegate.cc | 5 + .../serialization/option_writer_generator.cc | 6 + 25 files changed, 1537 insertions(+), 177 deletions(-) create mode 100644 tensorflow/compiler/mlir/lite/tests/blockwise-ops.mlir create mode 100644 tensorflow/compiler/mlir/lite/transforms/fold_stablehlo_constant_transforms_pass.cc diff --git a/RELEASE.md b/RELEASE.md index 38afa6c8b11cbe..ff95217c3bb665 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -14,7 +14,8 @@ ### Major Features and Improvements -* +* `tf.lite` + * Adds blockwise quantization support in the converter and `quant_spec` metadata in `FullyConnectedOptions`. ### Bug Fixes and Other Changes diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD index c5da7b7e7139e5..cdb76a3f91a6aa 100644 --- a/tensorflow/compiler/mlir/lite/BUILD +++ b/tensorflow/compiler/mlir/lite/BUILD @@ -1155,6 +1155,7 @@ cc_library( "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", ], ) @@ -1552,6 +1553,7 @@ cc_library( srcs = [ "transforms/decompose_hybrid_quantization.cc", "transforms/default_quant_params.cc", + "transforms/fold_stablehlo_constant_transforms_pass.cc", "transforms/generated_post_quantize.inc", "transforms/generated_quantize.inc", "transforms/lower_quant_annotations_helper.cc", diff --git a/tensorflow/compiler/mlir/lite/converter_gen.cc b/tensorflow/compiler/mlir/lite/converter_gen.cc index b08972c0301f33..df8c3c5d111e06 100644 --- a/tensorflow/compiler/mlir/lite/converter_gen.cc +++ b/tensorflow/compiler/mlir/lite/converter_gen.cc @@ -62,11 +62,11 @@ llvm::cl::opt action( "Generate TFLite runtime verifiers"))); // Returns the associated option name for the given op definition. -static inline std::string GetOperatorOptionName(const Record &def) { +static inline std::string GetOperatorOptionName(const Record& def) { assert(def.getName().starts_with("TFL_") && "unexpected op prefix"); assert(def.getName().ends_with("Op") && "unexpected op suffix"); - auto *custom_option = dyn_cast(def.getValueInit("customOption")); + auto* custom_option = dyn_cast(def.getValueInit("customOption")); std::ostringstream oss; if (custom_option) oss << custom_option->getValue().str(); @@ -90,7 +90,7 @@ static inline bool IsLstmOp(const StringRef op_name) { return op_name.take_back(6) == "LSTMOp"; } -static int HasOptions(const Record &def) { +static int HasOptions(const Record& def) { if (def.getValueAsBit("hasOptions")) { return 1; } @@ -100,13 +100,13 @@ static int HasOptions(const Record &def) { return 0; } -static void EmitOptionBuilders(const RecordKeeper &record_keeper, - const std::vector &defs, - raw_ostream *ostream) { - raw_ostream &os = *ostream; +static void EmitOptionBuilders(const RecordKeeper& record_keeper, + const std::vector& defs, + raw_ostream* ostream) { + raw_ostream& os = *ostream; const auto attr_type = record_keeper.getClass("Attr"); - for (const auto *def : defs) { + for (const auto* def : defs) { const int has_options = HasOptions(*def); // TFLite ops without options are skipped over. if (!has_options) { @@ -125,11 +125,11 @@ static void EmitOptionBuilders(const RecordKeeper &record_keeper, // Construct all the builder option needed. SmallVector options; // Add options due to attributes (not-derived). - auto *arg_values = def->getValueAsDag("arguments"); + auto* arg_values = def->getValueAsDag("arguments"); mlir::tblgen::Operator op(*def); for (unsigned i = 0, e = arg_values->getNumArgs(); i != e; ++i) { auto arg = arg_values->getArg(i); - const auto *arg_def = dyn_cast(arg); + const auto* arg_def = dyn_cast(arg); if (!arg_def) continue; if (arg_def->getDef()->isSubClassOf(attr_type)) { // This binds the name of the attribute in the TD file with the name @@ -156,8 +156,8 @@ static void EmitOptionBuilders(const RecordKeeper &record_keeper, } // Add options due to derived attributes. - for (const auto &val : def->getValues()) { - if (auto *record = dyn_cast(val.getType())) { + for (const auto& val : def->getValues()) { + if (auto* record = dyn_cast(val.getType())) { if (record->isSubClassOf(attr_type)) { if (record->getClasses().size() != 1) { PrintFatalError( @@ -174,7 +174,7 @@ static void EmitOptionBuilders(const RecordKeeper &record_keeper, } os << " tflite::" << tflite_option_name << "Builder b(*fbb);\n"; - for (const auto &option : options) + for (const auto& option : options) os << formatv(" b.add_{0}(std::move({0}));\n", option); os << " return b.Finish();\n}\n"; } @@ -187,11 +187,11 @@ static void EmitOptionBuilders(const RecordKeeper &record_keeper, // arguments that depend on op definitions should be auto-generated and then // operator should be built by the caller because it does not require // auto-generation. -static void EmitOperatorBuilders(const std::vector &defs, - raw_ostream *ostream) { - raw_ostream &os = *ostream; +static void EmitOperatorBuilders(const std::vector& defs, + raw_ostream* ostream) { + raw_ostream& os = *ostream; - for (const auto *def : defs) { + for (const auto* def : defs) { StringRef op_name = def->getName().drop_front(4); const bool has_intermediates = op_name.take_back(6) == "LSTMOp"; @@ -257,7 +257,7 @@ static void EmitOperatorBuilders(const std::vector &defs, } } -static inline std::string GetOperatorName(const Record &def) { +static inline std::string GetOperatorName(const Record& def) { auto name = def.getValueAsString("opName"); // Special case for basic_lstm. if (name == "basic_lstm") { @@ -275,14 +275,14 @@ static inline std::string GetOperatorName(const Record &def) { // // TODO(hinsu): Consider converting this to a static constant associative // container instead of a series of if conditions, if required. -static void EmitGetBuiltinOpCode(const std::vector &defs, - raw_ostream *ostream) { - raw_ostream &os = *ostream; +static void EmitGetBuiltinOpCode(const std::vector& defs, + raw_ostream* ostream) { + raw_ostream& os = *ostream; os << "std::optional " "mlir::GetBuiltinOpCode(mlir::Operation* op) {\n"; - for (const auto *def : defs) { + for (const auto* def : defs) { StringRef op_name = def->getName().drop_front(4); auto operator_name = GetOperatorName(*def); os << " if (isa(op))\n" @@ -304,20 +304,20 @@ static void EmitGetBuiltinOpCode(const std::vector &defs, // ... // return {0, 0}; // } -static void EmitOperandNumbers(const RecordKeeper &record_keeper, - const std::vector &defs, - raw_ostream *ostream) { - raw_ostream &os = *ostream; +static void EmitOperandNumbers(const RecordKeeper& record_keeper, + const std::vector& defs, + raw_ostream* ostream) { + raw_ostream& os = *ostream; const auto attr_type = record_keeper.getClass("Attr"); const auto optional_tensor = record_keeper.getClass("TFL_TensorOfOrNone"); os << "llvm::MinMax mlir::OperandNumbersMinMax(llvm::StringRef op_name) {\n"; - for (const auto *def : defs) { + for (const auto* def : defs) { auto op_name = def->getValueAsString("opName"); int tail_optional_tensor = 0, tensor_number_max = 0; - auto *arg_values = def->getValueAsDag("arguments"); + auto* arg_values = def->getValueAsDag("arguments"); for (int i = 0, e = arg_values->getNumArgs(); i < e; ++i) { auto arg = arg_values->getArg(i); - auto *arg_def = dyn_cast(arg); + auto* arg_def = dyn_cast(arg); if (!arg_def) continue; if (!arg_def->getDef()->isSubClassOf(attr_type)) { tensor_number_max++; @@ -351,9 +351,9 @@ static void EmitOperandNumbers(const RecordKeeper &record_keeper, // const std::vector& intermediates, // flatbuffers::FlatBufferBuilder *fbb, // std::optional debug_metadata_index); -static void EmitBuildOperator(const std::vector &defs, - raw_ostream *ostream) { - raw_ostream &os = *ostream; +static void EmitBuildOperator(const std::vector& defs, + raw_ostream* ostream) { + raw_ostream& os = *ostream; // Signature os << "std::optional>\n" @@ -365,7 +365,7 @@ static void EmitBuildOperator(const std::vector &defs, "flatbuffers::FlatBufferBuilder *fbb," "std::optional debug_metadata_index) {\n"; - for (const auto *def : defs) { + for (const auto* def : defs) { StringRef op_name = def->getName().drop_front(4); // Try to cast to each op case and call the corresponding op builder @@ -391,9 +391,9 @@ static void EmitBuildOperator(const std::vector &defs, // where id is an empty string if builtin_options_id is 1, or builtin_options_id // otherwise. static void EmitBuiltinOptionsToAttributes( - const RecordKeeper &record_keeper, const std::vector &defs, - raw_ostream *ostream, const int builtin_options_id) { - raw_ostream &os = *ostream; + const RecordKeeper& record_keeper, const std::vector& defs, + raw_ostream* ostream, const int builtin_options_id) { + raw_ostream& os = *ostream; const std::string builtin_options_suffix = [&] { switch (builtin_options_id) { @@ -415,7 +415,7 @@ static void EmitBuiltinOptionsToAttributes( "llvm::SmallVectorImpl &attributes) {\n"; const auto attr_type = record_keeper.getClass("Attr"); - for (const auto *def : defs) { + for (const auto* def : defs) { const int has_options = HasOptions(*def); if (has_options != builtin_options_id) { continue; @@ -429,10 +429,10 @@ static void EmitBuiltinOptionsToAttributes( os << formatv(" if(const auto *op = op_union.As{0}()) {{\n", option_name); // We only care about options that are in arguments - auto *arg_values = def->getValueAsDag("arguments"); + auto* arg_values = def->getValueAsDag("arguments"); for (unsigned i = 0, e = arg_values->getNumArgs(); i != e; ++i) { auto arg = arg_values->getArg(i); - const auto *arg_def = dyn_cast(arg); + const auto* arg_def = dyn_cast(arg); if (!arg_def) continue; if (arg_def->getDef()->isSubClassOf(attr_type)) { StringRef arg_name = arg_values->getArgNameStr(i); @@ -449,6 +449,21 @@ static void EmitBuiltinOptionsToAttributes( } } + for (const auto& val : def->getValues()) { + if (auto* record = dyn_cast(val.getType())) { + if (record->isSubClassOf(attr_type) && + record->getClasses().size() == 1 && + record->getClasses()[0]->getName() == "DerivedQuantSpecAttr") { + os << formatv( + " if (auto attr = BuildDerivedQuantSpecAttr(op->{0}, " + "builder))\n" + " attributes.emplace_back(builder.getNamedAttr(" + "mlir::TFL::kQuantSpecAttrName, attr));\n", + val.getName()); + } + } + } + os << " return;\n"; os << " }\n"; } @@ -463,14 +478,14 @@ static void EmitBuiltinOptionsToAttributes( // The function below has a non-constant reference as that is required by LLVM's // TableGenMain. // NOLINTNEXTLINE -static bool OperatorWritersMain(raw_ostream &os, const RecordKeeper &records) { +static bool OperatorWritersMain(raw_ostream& os, const RecordKeeper& records) { emitSourceFileHeader("MLIR TFLite FlatBuffer Builders", os); // Retrieve all the definitions derived from TFL_Op and sort by record name. - std::vector defs = records.getAllDerivedDefinitions("TFL_Op"); + std::vector defs = records.getAllDerivedDefinitions("TFL_Op"); llvm::sort(defs, LessRecord()); - for (const auto *def : defs) { + for (const auto* def : defs) { // TFLite ops in the .td file are expected to follow the naming convention: // TFL_Op. // The generated TFLite op C++ class should be TFL::Op. @@ -501,15 +516,15 @@ static bool OperatorWritersMain(raw_ostream &os, const RecordKeeper &records) { return false; } -static void GenOperandResultVerifier(raw_ostream &os, - llvm::ArrayRef values, +static void GenOperandResultVerifier(raw_ostream& os, + llvm::ArrayRef values, StringRef valueKind) { mlir::tblgen::FmtContext fctx; bool first = true; - for (const auto &static_value : llvm::enumerate(values)) { - auto *definit = llvm::cast(static_value.value()); - auto *val = definit->getDef()->getValue("tflRuntimeTypePredicate"); + for (const auto& static_value : llvm::enumerate(values)) { + auto* definit = llvm::cast(static_value.value()); + auto* val = definit->getDef()->getValue("tflRuntimeTypePredicate"); if (!val) continue; // Create code block on first type to verify. @@ -550,16 +565,16 @@ static void GenOperandResultVerifier(raw_ostream &os, } // NOLINTNEXTLINE -static bool RuntimeVerifierWriterMain(raw_ostream &os, - const RecordKeeper &records) { +static bool RuntimeVerifierWriterMain(raw_ostream& os, + const RecordKeeper& records) { emitSourceFileHeader("MLIR TFLite Runtime Verifiers", os); // Retrieve all the definitions derived from TFL_Op and sort by record name. - std::vector defs = records.getAllDerivedDefinitions("Op"); + std::vector defs = records.getAllDerivedDefinitions("Op"); llvm::sort(defs, LessRecord()); // Iterate through all the ops defined. - for (const auto *def : defs) { + for (const auto* def : defs) { mlir::tblgen::Operator op(*def); if (!op.getTrait("TflRuntimeVerifyOpInterface::Trait")) continue; @@ -571,7 +586,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, verify_ctx.addSubst("_op", "(*op)"); for (int i = 0, e = op.getNumOperands(); i < e; ++i) { - auto &value = op.getOperand(i); + auto& value = op.getOperand(i); // Skip from first variadic operands for now. Else getOperand index used // below doesn't match. if (value.isVariableLength()) break; @@ -579,7 +594,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, verify_ctx.addSubst(value.name, formatv("op->getOperand({0})", i)); } for (int i = 0, e = op.getNumResults(); i < e; ++i) { - auto &value = op.getResult(i); + auto& value = op.getResult(i); // Skip from first variadic results for now. Else getResult index used // below doesn't match. if (value.isVariableLength()) break; @@ -591,7 +606,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, GenOperandResultVerifier(os, def->getValueAsDag("results")->getArgs(), "result"); - for (auto &trait : op.getTraits()) { + for (auto& trait : op.getTraits()) { if (!trait.getDef().isSubClassOf("GenInternalOpTrait")) { continue; } @@ -600,7 +615,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, continue; } - auto *val = trait.getDef().getValue("tflRuntimePredicate"); + auto* val = trait.getDef().getValue("tflRuntimePredicate"); if (!val) continue; auto desc = trait.getDef().getValueAsString("tflRuntimeDescription"); @@ -643,7 +658,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, return false; } -int main(int argc, char **argv) { +int main(int argc, char** argv) { llvm::InitLLVM y(argc, argv); llvm::cl::ParseCommandLineOptions(argc, argv); if (action == ActionType::OpConv) diff --git a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc index 2c4b62b39ec179..7747bc5368564b 100644 --- a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc +++ b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc @@ -1871,6 +1871,11 @@ absl::Status ParseFullyConnected(const Operator* op, schema_params->asymmetric_quantize_inputs(); TFL_FILE_ENSURE_STATUS(ConvertTensorType( schema_params->quantized_bias_type(), ¶ms->quantized_bias_type)); + if (const auto* quant_spec = schema_params->quant_spec()) { + // Borrowed from the model buffer; see TfLiteFullyConnectedParams. + params->quant_spec = quant_spec->data(); + params->quant_spec_size = static_cast(quant_spec->size()); + } switch (schema_params->weights_format()) { case FullyConnectedOptionsWeightsFormat_DEFAULT: params->weights_format = kTfLiteFullyConnectedWeightsFormatDefault; diff --git a/tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h b/tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h index c580bf03cd3f59..c69daeddc44c2f 100644 --- a/tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h +++ b/tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h @@ -208,6 +208,18 @@ typedef struct { // Parameters for FullyConnected version 10 or above. // Used to determine the default value for the quantized bias. TfLiteType quantized_bias_type; + + // Parameters for FullyConnected version 15 or above. + // + // Opaque, flexbuffers-encoded description of a non-standard quantization + // contract; see `FullyConnectedOptions.quant_spec` in schema.fbs. Null when + // the op uses standard FullyConnected quantization semantics. + // + // This points into the model's buffer and is only valid for as long as the + // model is, in the same way as `TfLiteNode::custom_initial_data`. It is not + // owned by this struct. + const uint8_t* quant_spec; + int quant_spec_size; } TfLiteFullyConnectedParams; typedef enum { diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc b/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc index 1262ffa2191be2..5394b3566a7cce 100644 --- a/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc +++ b/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc @@ -27,7 +27,6 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "flatbuffers/buffer.h" // from @flatbuffers #include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers -#include "flatbuffers/flatbuffers.h" // from @flatbuffers #include "flatbuffers/flexbuffers.h" // from @flatbuffers #include "flatbuffers/string.h" // from @flatbuffers #include "flatbuffers/vector.h" // from @flatbuffers @@ -52,11 +51,8 @@ limitations under the License. #include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h" #include "tensorflow/compiler/mlir/lite/schema/schema_utils.h" #include "tensorflow/compiler/mlir/lite/utils/convert_type.h" -#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h" -#include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/status.h" -#include "tsl/platform/status.h" namespace { @@ -252,6 +248,78 @@ static tflite::LSTMKernelType ConvertTFL_LSTMKernelTypeAttrForOptionWriter( llvm_unreachable("invalid lstm_kernel_type in conversion."); } +// Serializes one `quant_spec` value. Only the kinds that the specs actually +// use are supported; anything else is a bug in the pass that built the +// dictionary, and dropping it would silently change the op's meaning. +// +// The key is taken by `const std::string&` rather than `llvm::StringRef` +// because the `flexbuffers::Builder` key overloads call `strlen` on it, and a +// `StringRef` is not guaranteed to be null-terminated. +static void SerializeQuantSpecValue(const std::string& key, + mlir::Attribute value, + flexbuffers::Builder* fbb) { + if (auto str = mlir::dyn_cast(value)) { + fbb->String(key.c_str(), str.getValue().str()); + } else if (auto boolean = mlir::dyn_cast(value)) { + fbb->Bool(key.c_str(), boolean.getValue()); + } else if (auto integer = mlir::dyn_cast(value)) { + fbb->Int(key.c_str(), integer.getInt()); + } else if (auto fp = mlir::dyn_cast(value)) { + fbb->Double(key.c_str(), fp.getValueAsDouble()); + } else { + llvm::report_fatal_error(llvm::Twine("quant_spec key '") + key + + "' has an unsupported attribute kind"); + } +} + +static flatbuffers::Offset> +ConvertDerivedQuantSpecAttrForOptionWriter( + mlir::DictionaryAttr attr, flatbuffers::FlatBufferBuilder* builder) { + if (!attr || attr.empty()) return 0; + + flexbuffers::Builder fbb; + const size_t map_start = fbb.StartMap(); + for (const mlir::NamedAttribute& named : attr) { + SerializeQuantSpecValue(named.getName().str(), named.getValue(), &fbb); + } + fbb.EndMap(map_start); + fbb.Finish(); + return builder->CreateVector(fbb.GetBuffer()); +} + +static mlir::DictionaryAttr BuildDerivedQuantSpecAttr( + const std::vector& quant_spec, mlir::Builder builder) { + if (quant_spec.empty()) return nullptr; + + flexbuffers::Reference root = flexbuffers::GetRoot(quant_spec); + if (!root.IsMap()) return nullptr; + flexbuffers::Map map = root.AsMap(); + flexbuffers::TypedVector keys = map.Keys(); + flexbuffers::Vector values = map.Values(); + + llvm::SmallVector entries; + entries.reserve(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + llvm::StringRef key = keys[i].AsKey(); + flexbuffers::Reference val = values[i]; + mlir::Attribute attr; + if (val.IsString()) { + attr = builder.getStringAttr(val.AsString().str()); + } else if (val.IsBool()) { + attr = builder.getBoolAttr(val.AsBool()); + } else if (val.IsIntOrUint()) { + attr = builder.getI64IntegerAttr(val.AsInt64()); + } else if (val.IsFloat()) { + attr = builder.getF32FloatAttr(val.AsFloat()); + } else { + llvm::report_fatal_error(llvm::Twine("quant_spec key '") + key + + "' has an unsupported flexbuffer type"); + } + entries.push_back(builder.getNamedAttr(key, attr)); + } + return builder.getDictionaryAttr(entries); +} + static mlir::Attribute BuildBoolAttr(bool value, mlir::Builder builder) { return builder.getBoolAttr(value); } diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc index b8527afd3648ba..f286d667cdadc4 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc +++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc @@ -1971,6 +1971,121 @@ LogicalResult BatchMatMulOp::verify() { return success(); } +//===----------------------------------------------------------------------===// +// BlockwiseQuantizeOp / BlockwiseDequantizeOp +//===----------------------------------------------------------------------===// + +namespace { + +// Returns the shape of the scale/zero-point grid implied by tiling `type` with +// `block_shape`, or failure (after emitting a diagnostic on `op`) if the two +// are not compatible. +FailureOr> GetBlockGridShape(Operation* op, + ShapedType type, + ArrayAttr block_shape) { + if (!type.hasStaticShape()) { + return op->emitOpError("expects a statically shaped tensor, got ") << type; + } + ArrayRef shape = type.getShape(); + if (block_shape.size() != shape.size()) { + return op->emitOpError("expects block_shape of rank ") + << shape.size() << ", got " << block_shape; + } + + SmallVector grid_shape; + grid_shape.reserve(shape.size()); + for (auto [dim, dim_size] : llvm::enumerate(shape)) { + auto block_attr = mlir::dyn_cast(block_shape[dim]); + if (!block_attr) { + return op->emitOpError("expects an integer block_shape, got ") + << block_shape; + } + const int64_t block_size = block_attr.getInt(); + if (block_size <= 0) { + return op->emitOpError("expects a positive block_shape, got ") + << block_shape; + } + if (dim_size % block_size != 0) { + return op->emitOpError("expects dimension ") + << dim << " (" << dim_size << ") to be divisible by block_shape[" + << dim << "] (" << block_size << ")"; + } + grid_shape.push_back(dim_size / block_size); + } + return grid_shape; +} + +// Verifies that `value`, if present, is shaped like the block grid. A +// dimension of size 1 is accepted anywhere as a broadcast. +LogicalResult VerifyBlockGridOperand(Operation* op, Value value, + ArrayRef grid_shape, + StringRef name) { + if (!value || mlir::isa(value.getType())) return success(); + auto type = mlir::dyn_cast(value.getType()); + if (!type) return success(); + + if (type.getRank() != static_cast(grid_shape.size())) { + return op->emitOpError("expects ") + << name << " of rank " << grid_shape.size() << ", got " << type; + } + for (auto [dim, dim_size] : llvm::enumerate(type.getShape())) { + if (dim_size != 1 && dim_size != grid_shape[dim]) { + return op->emitOpError("expects ") + << name << " dimension " << dim << " to be 1 or " + << grid_shape[dim] << ", got " << type; + } + } + return success(); +} + +} // namespace + +LogicalResult BlockwiseQuantizeOp::verify() { + auto grid_shape = GetBlockGridShape( + *this, mlir::cast(getInput().getType()), getBlockShape()); + if (failed(grid_shape)) return failure(); + + if (failed(VerifyBlockGridOperand(*this, getScale(), *grid_shape, "scale")) || + failed(VerifyBlockGridOperand(*this, getZeroPoint(), *grid_shape, + "zero_point"))) { + return failure(); + } + + if (getScaleType() != + mlir::cast(getScale().getType()).getElementType()) { + return emitOpError("expects the scale element type to match scale_type (") + << getScaleType() << ")"; + } + return success(); +} + +LogicalResult BlockwiseDequantizeOp::verify() { + auto grid_shape = GetBlockGridShape( + *this, mlir::cast(getInput().getType()), getBlockShape()); + if (failed(grid_shape)) return failure(); + + if (failed( + VerifyBlockGridOperand(*this, getScales(), *grid_shape, "scales")) || + failed(VerifyBlockGridOperand(*this, getZeroPoints(), *grid_shape, + "zero_points"))) { + return failure(); + } + + if (getSymmetric() && getZeroPoints() && + !mlir::isa(getZeroPoints().getType())) { + auto zero_points = mlir::dyn_cast_or_null( + getZeroPoints().getDefiningOp() + ? getZeroPoints().getDefiningOp()->getAttrOfType( + "value") + : nullptr); + if (zero_points && !zero_points.isSplat()) { + return emitOpError( + "expects a per-tensor zero_point when symmetric is set"); + } + } + return success(); +} + //===----------------------------------------------------------------------===// // FullyConnectedOp //===----------------------------------------------------------------------===// diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.h b/tensorflow/compiler/mlir/lite/ir/tfl_ops.h index 89d1a5ed960275..6e51d2334aba99 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.h +++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.h @@ -55,6 +55,12 @@ class ControlType : public Type::TypeBase { static constexpr StringLiteral name = "tfl.control"; }; +// Discardable attribute naming a non-standard quantization contract that the +// operand types cannot express on their own. Read by the `quant_spec` derived +// attribute of `tfl.fully_connected`; see `FullyConnectedOptions.quant_spec` in +// schema.fbs. +inline constexpr StringLiteral kQuantSpecAttrName = "tfl.quant_spec"; + #include "tensorflow/compiler/mlir/lite/ir/tfl_ops_interface.h.inc" } // end namespace TFL diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td index b281b8bee7748d..7a2778dc978f2a 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td +++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td @@ -58,6 +58,13 @@ def TFL_Control: Type($_self)">, "control">, // A type attribute containing the TensorType. def TensorTypeAttr : TypeAttrBase<"TensorType", "Tensor type attribute">; +// The element type of the scale that a blockwise quantize op computes at +// runtime. Only `f32` and `f8E8M0FNU` are meaningful today; the latter +// constrains the scale to a power of two. +def TFL_ScaleTypeAttr : TypeAttrBase<"mlir::Type", "scale element type attribute"> { + let constBuilderCall = "($0 ? ::mlir::TypeAttr::get($0) : ::mlir::TypeAttr())"; +} + //===----------------------------------------------------------------------===// // Derived shape attribute class. //===----------------------------------------------------------------------===// @@ -65,6 +72,11 @@ class DerivedShapeAttr : DerivedAttr<"ArrayRef", body>; class DerivedTFLiteTypeAttr : DerivedAttr<"tflite::TensorType", body, convert>; +// Reads the opaque quantization spec of an op; see +// `FullyConnectedOptions.quant_spec` in schema.fbs. Null when the op uses +// standard quantization semantics. +class DerivedQuantSpecAttr : DerivedAttr<"::mlir::DictionaryAttr", body>; + // TFL Runtime op trait predicate. class TFL_RuntimePredOpTrait : GenInternalOpTrait<"TFLRuntimeOpTrait"> { @@ -1140,6 +1152,21 @@ def TFL_FullyConnectedOp : TFL_Op<"fully_connected", [ OptionalAttr:$asymmetric_quantize_inputs ); + // Names a non-standard quantization contract that the operand types cannot + // express on their own; see `FullyConnectedOptions.quant_spec` in schema.fbs. + // A runtime that does not recognize the named spec must reject the op rather + // than fall back to the standard behavior. + // + // This is a derived attribute reading the discardable `tfl.quant_spec` + // rather than an entry in `arguments`, because adding an argument would + // change the generated builders and break every pass, in this tree and out + // of it, that constructs a fully_connected positionally. + DerivedQuantSpecAttr quant_spec = DerivedQuantSpecAttr<[{ + return (*this)->getAttrOfType<::mlir::DictionaryAttr>( + TFL::kQuantSpecAttrName); + }]>; + + // Depending on the weights format, this op can have one or two outputs. let results = (outs TFL_VariadicTensorOf<[F32, F16, QI8, QUI8, QI16, QUI16]>:$output @@ -4352,6 +4379,80 @@ def TFL_DequantizeOp: TFL_Op<"dequantize", [NoMemoryEffect]> { ]; } +// The two ops below model blockwise (sub-channel) quantization, which +// `!quant.uniform` cannot express. They are deliberately *not* part of the +// TFLite schema: they only exist between `LowerQuantAnnotations` and +// `FuseA4W2DynamicRangeFullyConnected`, which collapses the recognized +// patterns back into a single `tfl.fully_connected`. An op that survives to +// the exporter is a bug, and the exporter reports it as an unsupported op. +// +// Block layout follows the convention used throughout this dialect: the +// tensor is tiled by `block_shape`, one element of the scale/zero-point grid +// per tile, with the grid laid out in the same dimension order as the tensor. +// A grid dimension of size 1 broadcasts across that axis. + +// These two ops exist only inside the MLIR pipeline: they are declared with the +// raw `Op` base rather than `TFL_Op` so that `converter_gen` +// does not try to map them onto a `BuiltinOperator` that the schema does not +// have. `FuseA4W2DynamicRangeFullyConnected` collapses them into a single +// `tfl.fully_connected`; any that survive to export fail there. +def TFL_BlockwiseQuantizeOp : Op { + let summary = "Blockwise dynamic quantize operator"; + + let description = [{ + Computes a per-block scale (and, when asymmetric, zero point) from the + input at runtime and quantizes the input with them. The computed + parameters are returned alongside the quantized tensor so that a matching + `tfl.blockwise_dequantize` can undo the transformation. + + The scale of a block is `max(abs(block)) / (q_max + range_dilation)`, + where `q_max` is the largest representable value of the output storage + type. `range_dilation` widens the denominator so that outliers clip + instead of stretching the whole block; it defaults to 0. When + `scale_type` is `f8E8M0FNU`, the scale is rounded up to the next power of + two. + }]; + + let arguments = (ins + TFL_FpTensor:$input, + I64ArrayAttr:$block_shape, + TFL_ScaleTypeAttr:$scale_type, + BoolAttr:$symmetric, + DefaultValuedAttr:$range_dilation + ); + + let results = (outs + TensorOf<[TFL_I2, TFL_I4, I8, I16]>:$output, + TFL_TensorOf<[F8E8M0FNU, F16, BF16, F32]>:$scale, + TFL_TensorOfOrNone<[TFL_I2, TFL_I4, I8, I16, F32]>:$zero_point + ); + + let hasVerifier = 1; +} + +def TFL_BlockwiseDequantizeOp + : Op { + let summary = "Blockwise dequantize operator"; + + let description = [{ + Reconstructs floating point values from a blockwise quantized tensor as + `scale * (input - zero_point)`. A missing `zero_points` operand means a + zero point of 0. + }]; + + let arguments = (ins + TensorOf<[TFL_I2, TFL_I4, I8, I16]>:$input, + TFL_TensorOf<[F8E8M0FNU, F16, BF16, F32]>:$scales, + TFL_TensorOfOrNone<[TFL_I2, TFL_I4, I8, I16, F32]>:$zero_points, + I64ArrayAttr:$block_shape, + BoolAttr:$symmetric + ); + + let results = (outs TFL_FpTensor:$output); + + let hasVerifier = 1; +} + def TFL_FakeQuantOp : TFL_Op<"fake_quant", [ Pure, QuantizableResult]> { diff --git a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc index d91aed63f18f61..9518f575636555 100644 --- a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc +++ b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc @@ -150,6 +150,8 @@ void AddPipelinePasses(mlir::OpPassManager& pass_manager, // func.call boundaries that cause func.call result type mismatches when // quantized types are introduced. pass_manager.addPass(mlir::createInlinerPass()); + pass_manager.addNestedPass( + mlir::TFL::CreateFoldStablehloConstantTransformsPass()); pass_manager.addPass(mlir::TFL::CreateLowerQuantAnnotationsPass()); pass_manager.addPass(mlir::createSymbolDCEPass()); diff --git a/tensorflow/compiler/mlir/lite/schema/schema.fbs b/tensorflow/compiler/mlir/lite/schema/schema.fbs index b51dce02b63040..44af9948988d5b 100644 --- a/tensorflow/compiler/mlir/lite/schema/schema.fbs +++ b/tensorflow/compiler/mlir/lite/schema/schema.fbs @@ -1023,6 +1023,18 @@ table FullyConnectedOptions { // Parameters for FullyConnected version 11 or above. // When set, quantized_bias_type defines the dtype for both bias and accumulator. quantized_bias_type: TensorType; + + // Parameters for FullyConnected version 15 or above. + // + // An opaque, self-describing description of a non-standard quantization + // contract that the op's tensor-level quantization parameters cannot express + // on their own (for example a dynamically quantized, blockwise activation, or + // a weight grid that is centered between integer values). + // + // The payload is an encoded FlexBuffers map. A runtime that does not + // understand the contract named in the map must reject the op rather than + // fall back to standard quantization semantics. + quant_spec: [ubyte]; } table SoftmaxOptions { diff --git a/tensorflow/compiler/mlir/lite/schema/schema_generated.h b/tensorflow/compiler/mlir/lite/schema/schema_generated.h index cb9f9a479665f9..1bb5ef0f6a6fe6 100755 --- a/tensorflow/compiler/mlir/lite/schema/schema_generated.h +++ b/tensorflow/compiler/mlir/lite/schema/schema_generated.h @@ -9101,6 +9101,7 @@ struct FullyConnectedOptionsT : public ::flatbuffers::NativeTable { bool keep_num_dims = false; bool asymmetric_quantize_inputs = false; tflite::TensorType quantized_bias_type = tflite::TensorType_FLOAT32; + std::vector quant_spec{}; }; struct FullyConnectedOptions FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table { @@ -9111,7 +9112,8 @@ struct FullyConnectedOptions FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Ta VT_WEIGHTS_FORMAT = 6, VT_KEEP_NUM_DIMS = 8, VT_ASYMMETRIC_QUANTIZE_INPUTS = 10, - VT_QUANTIZED_BIAS_TYPE = 12 + VT_QUANTIZED_BIAS_TYPE = 12, + VT_QUANT_SPEC = 14 }; tflite::ActivationFunctionType fused_activation_function() const { return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); @@ -9128,6 +9130,9 @@ struct FullyConnectedOptions FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Ta tflite::TensorType quantized_bias_type() const { return static_cast(GetField(VT_QUANTIZED_BIAS_TYPE, 0)); } + const ::flatbuffers::Vector *quant_spec() const { + return GetPointer *>(VT_QUANT_SPEC); + } template bool Verify(::flatbuffers::VerifierTemplate &verifier) const { return VerifyTableStart(verifier) && @@ -9136,6 +9141,8 @@ struct FullyConnectedOptions FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Ta VerifyField(verifier, VT_KEEP_NUM_DIMS, 1) && VerifyField(verifier, VT_ASYMMETRIC_QUANTIZE_INPUTS, 1) && VerifyField(verifier, VT_QUANTIZED_BIAS_TYPE, 1) && + VerifyOffset(verifier, VT_QUANT_SPEC) && + verifier.VerifyVector(quant_spec()) && verifier.EndTable(); } FullyConnectedOptionsT *UnPack(const ::flatbuffers::resolver_function_t *_resolver = nullptr) const; @@ -9162,6 +9169,9 @@ struct FullyConnectedOptionsBuilder { void add_quantized_bias_type(tflite::TensorType quantized_bias_type) { fbb_.AddElement(FullyConnectedOptions::VT_QUANTIZED_BIAS_TYPE, static_cast(quantized_bias_type), 0); } + void add_quant_spec(::flatbuffers::Offset<::flatbuffers::Vector> quant_spec) { + fbb_.AddOffset(FullyConnectedOptions::VT_QUANT_SPEC, quant_spec); + } explicit FullyConnectedOptionsBuilder(::flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { start_ = fbb_.StartTable(); @@ -9179,8 +9189,10 @@ inline ::flatbuffers::Offset CreateFullyConnectedOptions( tflite::FullyConnectedOptionsWeightsFormat weights_format = tflite::FullyConnectedOptionsWeightsFormat_DEFAULT, bool keep_num_dims = false, bool asymmetric_quantize_inputs = false, - tflite::TensorType quantized_bias_type = tflite::TensorType_FLOAT32) { + tflite::TensorType quantized_bias_type = tflite::TensorType_FLOAT32, + ::flatbuffers::Offset<::flatbuffers::Vector> quant_spec = 0) { FullyConnectedOptionsBuilder builder_(_fbb); + builder_.add_quant_spec(quant_spec); builder_.add_quantized_bias_type(quantized_bias_type); builder_.add_asymmetric_quantize_inputs(asymmetric_quantize_inputs); builder_.add_keep_num_dims(keep_num_dims); @@ -9189,6 +9201,25 @@ inline ::flatbuffers::Offset CreateFullyConnectedOptions( return builder_.Finish(); } +inline ::flatbuffers::Offset CreateFullyConnectedOptionsDirect( + ::flatbuffers::FlatBufferBuilder &_fbb, + tflite::ActivationFunctionType fused_activation_function = tflite::ActivationFunctionType_NONE, + tflite::FullyConnectedOptionsWeightsFormat weights_format = tflite::FullyConnectedOptionsWeightsFormat_DEFAULT, + bool keep_num_dims = false, + bool asymmetric_quantize_inputs = false, + tflite::TensorType quantized_bias_type = tflite::TensorType_FLOAT32, + const std::vector *quant_spec = nullptr) { + auto quant_spec__ = quant_spec ? _fbb.CreateVector(*quant_spec) : 0; + return tflite::CreateFullyConnectedOptions( + _fbb, + fused_activation_function, + weights_format, + keep_num_dims, + asymmetric_quantize_inputs, + quantized_bias_type, + quant_spec__); +} + ::flatbuffers::Offset CreateFullyConnectedOptions(::flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, const ::flatbuffers::rehasher_function_t *_rehasher = nullptr); struct SoftmaxOptionsT : public ::flatbuffers::NativeTable { @@ -18986,6 +19017,7 @@ inline void FullyConnectedOptions::UnPackTo(FullyConnectedOptionsT *_o, const :: { auto _e = keep_num_dims(); _o->keep_num_dims = _e; } { auto _e = asymmetric_quantize_inputs(); _o->asymmetric_quantize_inputs = _e; } { auto _e = quantized_bias_type(); _o->quantized_bias_type = _e; } + { auto _e = quant_spec(); if (_e) { _o->quant_spec.resize(_e->size()); std::copy(_e->begin(), _e->end(), _o->quant_spec.begin()); } } } inline ::flatbuffers::Offset CreateFullyConnectedOptions(::flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, const ::flatbuffers::rehasher_function_t *_rehasher) { @@ -19001,13 +19033,15 @@ inline ::flatbuffers::Offset FullyConnectedOptions::Pack( auto _keep_num_dims = _o->keep_num_dims; auto _asymmetric_quantize_inputs = _o->asymmetric_quantize_inputs; auto _quantized_bias_type = _o->quantized_bias_type; + auto _quant_spec = _o->quant_spec.size() ? _fbb.CreateVector(_o->quant_spec) : 0; return tflite::CreateFullyConnectedOptions( _fbb, _fused_activation_function, _weights_format, _keep_num_dims, _asymmetric_quantize_inputs, - _quantized_bias_type); + _quantized_bias_type, + _quant_spec); } inline SoftmaxOptionsT *SoftmaxOptions::UnPack(const ::flatbuffers::resolver_function_t *_resolver) const { diff --git a/tensorflow/compiler/mlir/lite/tests/blockwise-ops.mlir b/tensorflow/compiler/mlir/lite/tests/blockwise-ops.mlir new file mode 100644 index 00000000000000..b4adc4d2a73926 --- /dev/null +++ b/tensorflow/compiler/mlir/lite/tests/blockwise-ops.mlir @@ -0,0 +1,126 @@ +// Copyright 2026 The TensorFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================== + +// Verifier tests for the blockwise quantize/dequantize ops. +// +// These two ops exist only between `LowerQuantAnnotations` and +// `FuseA4W2DynamicRangeFullyConnected`, so they are deliberately absent from +// the TFLite schema and are not run through `-tfl-runtime-verify`. + +// RUN: litert-opt -split-input-file -verify-diagnostics %s | FileCheck %s + +// CHECK-LABEL: testBlockwiseQuantize +func.func @testBlockwiseQuantize(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // CHECK: "tfl.blockwise_quantize" + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 32], scale_type = f8E8M0FNU, symmetric = true, range_dilation = 1.500000e+00 : f32} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x2xf8E8M0FNU>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +func.func @testBlockwiseQuantizeDynamicShape(%arg0: tensor) -> tensor { + // expected-error @+1 {{expects a statically shaped tensor}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 32], scale_type = f8E8M0FNU, symmetric = true} : (tensor) -> (tensor, tensor, none) + func.return %0#0 : tensor +} + +// ----- + +func.func @testBlockwiseQuantizeBlockShapeRank(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // expected-error @+1 {{expects block_shape of rank 2}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [32], scale_type = f8E8M0FNU, symmetric = true} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x2xf8E8M0FNU>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +func.func @testBlockwiseQuantizeIndivisibleBlock(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // expected-error @+1 {{expects dimension 1 (64) to be divisible by block_shape[1] (7)}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 7], scale_type = f8E8M0FNU, symmetric = true} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x2xf8E8M0FNU>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +func.func @testBlockwiseQuantizeNonPositiveBlock(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // expected-error @+1 {{expects a positive block_shape}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 0], scale_type = f8E8M0FNU, symmetric = true} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x2xf8E8M0FNU>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +func.func @testBlockwiseQuantizeScaleGridMismatch(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // expected-error @+1 {{expects scale dimension 1 to be 1 or 2}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 32], scale_type = f8E8M0FNU, symmetric = true} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x8xf8E8M0FNU>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +func.func @testBlockwiseQuantizeScaleTypeMismatch(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // expected-error @+1 {{expects the scale element type to match scale_type}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 32], scale_type = f8E8M0FNU, symmetric = true} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x2xf32>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +// CHECK-LABEL: testBlockwiseDequantize +func.func @testBlockwiseDequantize(%arg0: tensor<32x64xi2>, %arg1: tensor<32x1xf32>, %arg2: tensor<32x1xf32>) -> tensor<32x64xf32> { + // CHECK: "tfl.blockwise_dequantize" + %0 = "tfl.blockwise_dequantize"(%arg0, %arg1, %arg2) {block_shape = [1, 64], symmetric = false} : (tensor<32x64xi2>, tensor<32x1xf32>, tensor<32x1xf32>) -> tensor<32x64xf32> + func.return %0 : tensor<32x64xf32> +} + +// ----- + +// A missing zero point is how a symmetric dequantize is spelled. +// CHECK-LABEL: testBlockwiseDequantizeNoZeroPoint +func.func @testBlockwiseDequantizeNoZeroPoint(%arg0: tensor<32x64xi2>, %arg1: tensor<32x1xf32>) -> tensor<32x64xf32> { + %none = "tfl.no_value"() {value} : () -> none + // CHECK: "tfl.blockwise_dequantize" + %0 = "tfl.blockwise_dequantize"(%arg0, %arg1, %none) {block_shape = [1, 64], symmetric = true} : (tensor<32x64xi2>, tensor<32x1xf32>, none) -> tensor<32x64xf32> + func.return %0 : tensor<32x64xf32> +} + +// ----- + +func.func @testBlockwiseDequantizeScalesRank(%arg0: tensor<32x64xi2>, %arg1: tensor<32xf32>) -> tensor<32x64xf32> { + %none = "tfl.no_value"() {value} : () -> none + // expected-error @+1 {{expects scales of rank 2}} + %0 = "tfl.blockwise_dequantize"(%arg0, %arg1, %none) {block_shape = [1, 64], symmetric = true} : (tensor<32x64xi2>, tensor<32xf32>, none) -> tensor<32x64xf32> + func.return %0 : tensor<32x64xf32> +} + +// ----- + +func.func @testBlockwiseDequantizeIndivisibleBlock(%arg0: tensor<32x64xi2>, %arg1: tensor<32x1xf32>) -> tensor<32x64xf32> { + %none = "tfl.no_value"() {value} : () -> none + // expected-error @+1 {{expects dimension 1 (64) to be divisible by block_shape[1] (7)}} + %0 = "tfl.blockwise_dequantize"(%arg0, %arg1, %none) {block_shape = [1, 7], symmetric = true} : (tensor<32x64xi2>, tensor<32x1xf32>, none) -> tensor<32x64xf32> + func.return %0 : tensor<32x64xf32> +} + +// ----- + +func.func @testBlockwiseDequantizeSymmetricWithZeroPoints(%arg0: tensor<32x64xi2>, %arg1: tensor<32x1xf32>) -> tensor<32x64xf32> { + %zp = "tfl.pseudo_const"() {value = dense<[[0.0], [-5.000000e-01]]> : tensor<2x1xf32>} : () -> tensor<2x1xf32> + %w = "tfl.pseudo_const"() {value = dense<1> : tensor<2x64xi2>} : () -> tensor<2x64xi2> + %s = "tfl.pseudo_const"() {value = dense<2.500000e-01> : tensor<2x1xf32>} : () -> tensor<2x1xf32> + // expected-error @+1 {{expects a per-tensor zero_point when symmetric is set}} + %0 = "tfl.blockwise_dequantize"(%w, %s, %zp) {block_shape = [1, 64], symmetric = true} : (tensor<2x64xi2>, tensor<2x1xf32>, tensor<2x1xf32>) -> tensor<2x64xf32> + func.return %arg0 : tensor<32x64xi2> +} diff --git a/tensorflow/compiler/mlir/lite/tools/versioning/op_version.cc b/tensorflow/compiler/mlir/lite/tools/versioning/op_version.cc index f4eedaa48f0b9a..d1bd505883c463 100644 --- a/tensorflow/compiler/mlir/lite/tools/versioning/op_version.cc +++ b/tensorflow/compiler/mlir/lite/tools/versioning/op_version.cc @@ -183,6 +183,15 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) { reinterpret_cast(op_sig.builtin_data); TFLITE_DCHECK(fully_connected_params != nullptr); + // A non-empty `quant_spec` describes numerics that a runtime without + // explicit support for that spec cannot reproduce. It must be checked + // before the cases below, because an op carrying a `quant_spec` also + // matches them (an a4w2 op has int2 weights) and would otherwise be + // reported as a version that older runtimes accept. + if (fully_connected_params->quant_spec_size > 0) { + return 15; + } + if (op_sig.inputs.at(1).type == kTfLiteInt2) { return 14; } diff --git a/tensorflow/compiler/mlir/lite/tools/versioning/runtime_version.cc b/tensorflow/compiler/mlir/lite/tools/versioning/runtime_version.cc index 76074657c25fd4..9efeb19cfc2816 100644 --- a/tensorflow/compiler/mlir/lite/tools/versioning/runtime_version.cc +++ b/tensorflow/compiler/mlir/lite/tools/versioning/runtime_version.cc @@ -147,6 +147,7 @@ std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code, {{BuiltinOperator_FULLY_CONNECTED, 12}, "2.17.0"}, {{BuiltinOperator_FULLY_CONNECTED, 13}, "2.18.0"}, {{BuiltinOperator_FULLY_CONNECTED, 14}, "2.21.0"}, + {{BuiltinOperator_FULLY_CONNECTED, 15}, "2.23.0"}, {{BuiltinOperator_GATHER, 1}, "1.6.0"}, {{BuiltinOperator_GATHER, 2}, "1.14.0"}, {{BuiltinOperator_GATHER, 3}, "1.15.0"}, diff --git a/tensorflow/compiler/mlir/lite/transforms/fold_stablehlo_constant_transforms_pass.cc b/tensorflow/compiler/mlir/lite/transforms/fold_stablehlo_constant_transforms_pass.cc new file mode 100644 index 00000000000000..7d0ce95d969d26 --- /dev/null +++ b/tensorflow/compiler/mlir/lite/transforms/fold_stablehlo_constant_transforms_pass.cc @@ -0,0 +1,277 @@ +/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include + +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/FormatVariadic.h" +#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project +#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/IR/AsmState.h" // from @llvm-project +#include "mlir/IR/Attributes.h" // from @llvm-project +#include "mlir/IR/Builders.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project +#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project +#include "mlir/IR/DialectResourceBlobManager.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/Matchers.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/IR/Value.h" // from @llvm-project +#include "mlir/Pass/Pass.h" // from @llvm-project +#include "mlir/Support/LLVM.h" // from @llvm-project +#include "mlir/Support/LogicalResult.h" // from @llvm-project +#include "mlir/Support/TypeID.h" // from @llvm-project +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project +#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo +#include "tensorflow/compiler/mlir/lite/transforms/passes.h" + +namespace mlir { +namespace TFL { +namespace { + +#define GEN_PASS_DEF_FOLDSTABLEHLOCONSTANTTRANSFORMSPASS +#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc" + +static void ComputePermutation( + llvm::ArrayRef perms, llvm::ArrayRef output_shape, + const char* raw_input, const int element_byte_size, + const int64_t current_axis, char*& output_element_addr, + llvm::MutableArrayRef current_input_index, + mlir::ShapedType input_shape_type) { + const int64_t input_axis = perms[current_axis]; + const bool is_last_axis = current_axis == output_shape.size() - 1; + + for (int i = 0; i < output_shape[current_axis]; ++i) { + current_input_index[input_axis] = i; + if (is_last_axis) { + int64_t input_flat_index = ElementsAttr::getFlattenedIndex( + input_shape_type, current_input_index); + const char* input_element_addr = + raw_input + input_flat_index * element_byte_size; + std::memcpy(output_element_addr, input_element_addr, element_byte_size); + output_element_addr += element_byte_size; + } else { + ComputePermutation(perms, output_shape, raw_input, element_byte_size, + current_axis + 1, output_element_addr, + current_input_index, input_shape_type); + } + } +} + +static ElementsAttr GetElementsAttrFromValue(mlir::Value val) { + if (!val) return nullptr; + ElementsAttr attr; + if (matchPattern(val, m_Constant(&attr))) { + return attr; + } + if (auto cst = val.getDefiningOp()) { + return mlir::dyn_cast_or_null(cst.getValue()); + } + if (auto cst = val.getDefiningOp()) { + return mlir::dyn_cast_or_null(cst.getValue()); + } + return nullptr; +} + +// Pattern to fold stablehlo.reshape on constant inputs. +struct FoldStablehloReshapeOpPattern + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(stablehlo::ReshapeOp op, + PatternRewriter& rewriter) const override { + mlir::Value input = op.getOperand(); + ElementsAttr input_attr = GetElementsAttrFromValue(input); + if (!input_attr) return failure(); + + auto result_type = mlir::cast(op.getType()); + if (!result_type.hasStaticShape()) return failure(); + + if (auto dense_elements = mlir::dyn_cast(input_attr)) { + rewriter.replaceOpWithNewOp( + op, dense_elements.reshape(result_type)); + return success(); + } else if (auto dense_resource = + mlir::dyn_cast(input_attr)) { + AsmResourceBlob* blob = dense_resource.getRawHandle().getBlob(); + if (!blob && dense_resource.getRawHandle().getResource()) { + blob = dense_resource.getRawHandle().getResource()->getBlob(); + } + if (!blob) return failure(); + + static std::atomic reshape_counter{0}; + auto key = dense_resource.getRawHandle().getKey(); + std::string new_key = + llvm::formatv("{0}_reshaped_{1}", key, reshape_counter++); + + auto new_blob = mlir::HeapAsmResourceBlob::allocate( + blob->getData().size(), /*align=*/64, true); + std::memcpy(const_cast(new_blob.getData().data()), + blob->getData().data(), blob->getData().size()); + auto new_attr = DenseResourceElementsAttr::get(result_type, new_key, + std::move(new_blob)); + rewriter.replaceOpWithNewOp(op, new_attr); + return success(); + } + return failure(); + } +}; + +// Pattern to fold stablehlo.transpose on constant inputs. +struct FoldStablehloTransposeOpPattern + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(stablehlo::TransposeOp op, + PatternRewriter& rewriter) const override { + mlir::Value input = op.getOperand(); + ElementsAttr input_attr = GetElementsAttrFromValue(input); + if (!input_attr) return failure(); + + auto input_type = mlir::cast(input_attr.getType()); + auto result_type = mlir::cast(op.getType()); + if (!result_type.hasStaticShape() || !input_type.hasStaticShape()) { + return failure(); + } + + if (!input_type.getElementType().isIntOrIndexOrFloat()) { + return failure(); + } + + const int num_dimensions = input_type.getRank(); + if (num_dimensions <= 1) { + rewriter.replaceOp(op, input); + return success(); + } + + llvm::ArrayRef perms = op.getPermutation(); + if (static_cast(perms.size()) != num_dimensions) { + return failure(); + } + + llvm::SmallVector output_shape; + for (int i = 0; i < num_dimensions; ++i) { + output_shape.push_back(input_type.getDimSize(perms[i])); + } + + if (auto dense_elements = mlir::dyn_cast(input_attr)) { + if (dense_elements.isSplat()) { + rewriter.replaceOpWithNewOp( + op, dense_elements.reshape(result_type)); + return success(); + } + + const int element_byte_size = + dense_elements.getElementType().getIntOrFloatBitWidth() / 8; + if (element_byte_size <= 0) return failure(); + + const char* raw_input = dense_elements.getRawData().data(); + llvm::SmallVector current_input_index(num_dimensions, 0); + llvm::SmallVector raw_output_arr( + dense_elements.getRawData().begin(), + dense_elements.getRawData().end()); + char* raw_output = raw_output_arr.data(); + + ComputePermutation(perms, output_shape, raw_input, element_byte_size, + /*current_axis=*/0, raw_output, current_input_index, + input_type); + + if (!DenseElementsAttr::isValidRawBuffer(result_type, raw_output_arr)) { + return failure(); + } + + auto new_attr = + DenseElementsAttr::getFromRawBuffer(result_type, raw_output_arr); + rewriter.replaceOpWithNewOp(op, new_attr); + return success(); + } else if (auto dense_resource = + mlir::dyn_cast(input_attr)) { + AsmResourceBlob* blob = dense_resource.getRawHandle().getBlob(); + if (!blob && dense_resource.getRawHandle().getResource()) { + blob = dense_resource.getRawHandle().getResource()->getBlob(); + } + if (!blob) return failure(); + + const int element_byte_size = + input_type.getElementType().getIntOrFloatBitWidth() / 8; + if (element_byte_size <= 0) return failure(); + + int64_t total_elements = result_type.getNumElements(); + size_t total_bytes = total_elements * element_byte_size; + + auto raw_output_blob = mlir::HeapAsmResourceBlob::allocate( + total_bytes, /*align=*/64, /*dataIsMutable=*/true); + llvm::ArrayRef data = raw_output_blob.getDataAs(); + llvm::MutableArrayRef raw_output_arr(const_cast(data.data()), + data.size()); + char* raw_output = raw_output_arr.data(); + const char* raw_input = blob->getData().data(); + if (!raw_input) return failure(); + + llvm::SmallVector current_input_index(num_dimensions, 0); + ComputePermutation(perms, output_shape, raw_input, element_byte_size, + /*current_axis=*/0, raw_output, current_input_index, + input_type); + + static std::atomic transpose_counter{0}; + std::string new_key = llvm::formatv("stablehlo_transpose_fold_result_{0}", + transpose_counter++); + auto new_attr = DenseResourceElementsAttr::get( + result_type, new_key, std::move(raw_output_blob)); + rewriter.replaceOpWithNewOp(op, new_attr); + return success(); + } + + return failure(); + } +}; + +struct FoldStablehloConstantTransformsPass + : public impl::FoldStablehloConstantTransformsPassBase< + FoldStablehloConstantTransformsPass> { + public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID( + FoldStablehloConstantTransformsPass) + + void runOnOperation() override { + func::FuncOp func = getOperation(); + MLIRContext* context = func.getContext(); + + RewritePatternSet patterns(context); + patterns + .add( + context); + + if (failed(applyPatternsGreedily(func, std::move(patterns)))) { + signalPassFailure(); + } + } +}; + +} // namespace + +std::unique_ptr> +CreateFoldStablehloConstantTransformsPass() { + return std::make_unique(); +} + +} // namespace TFL +} // namespace mlir diff --git a/tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.cc b/tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.cc index 9aae3a93606c55..b09ee0bb3006a4 100644 --- a/tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.cc @@ -26,6 +26,7 @@ limitations under the License. #include "Eigen/Core" // from @eigen_archive #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" #include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project @@ -34,12 +35,15 @@ limitations under the License. #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Matchers.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project #include "mlir/IR/Types.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/Pass/Pass.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Support/LogicalResult.h" // from @llvm-project +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h" #include "tensorflow/compiler/mlir/lite/ir/tfl_ops_in_place_transpose.h" #include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h" @@ -460,10 +464,167 @@ LogicalResult FoldResourceOpPattern(ModuleOp module) { return success(); } +struct PushTransposeThroughBlockwiseDequantize + : public OpRewritePattern { + explicit PushTransposeThroughBlockwiseDequantize(MLIRContext* ctx) + : OpRewritePattern(ctx) {} + + LogicalResult matchAndRewrite(TransposeOp op, + PatternRewriter& rewriter) const override { + auto deq = op.getOperand(0).getDefiningOp(); + if (!deq) return failure(); + + if (!deq.getOutput().hasOneUse()) return failure(); + + ElementsAttr input_attr; + if (!matchPattern(deq.getInput(), m_Constant(&input_attr))) { + if (auto const_op = deq.getInput().getDefiningOp()) { + input_attr = const_op.getValue(); + } else { + return failure(); + } + } + + ElementsAttr perm_attr; + if (!matchPattern(op.getOperand(1), m_Constant(&perm_attr))) { + if (auto const_op = op.getOperand(1).getDefiningOp()) { + perm_attr = const_op.getValue(); + } else { + return failure(); + } + } + auto int_perm_attr = mlir::dyn_cast(perm_attr); + if (!int_perm_attr) return failure(); + + auto input_type = + mlir::dyn_cast(deq.getInput().getType()); + auto output_type = mlir::dyn_cast(op.getType()); + if (!input_type || !output_type) return failure(); + + int rank = input_type.getRank(); + auto perms = llvm::to_vector<4>( + llvm::map_range(int_perm_attr.getValues(), + [](const APInt& val) { return val.getSExtValue(); })); + if (perms.size() != rank) return failure(); + + llvm::SmallDenseSet seen_dims; + for (int64_t p : perms) { + if (p < 0 || p >= rank || !seen_dims.insert(p).second) return failure(); + } + + ArrayAttr new_block_shape_attr = nullptr; + if (auto block_shape_attr = deq.getBlockShapeAttr()) { + if (block_shape_attr.size() != rank) return failure(); + SmallVector new_block_shape(rank); + for (int i = 0; i < rank; ++i) { + new_block_shape[i] = + mlir::cast(block_shape_attr[perms[i]]).getInt(); + } + new_block_shape_attr = rewriter.getI64ArrayAttr(new_block_shape); + } else { + return failure(); + } + + SmallVector transposed_input_shape; + for (int64_t p : perms) { + transposed_input_shape.push_back(input_type.getDimSize(p)); + } + auto transposed_input_type = RankedTensorType::get( + transposed_input_shape, input_type.getElementType()); + Value new_input = rewriter.create( + op.getLoc(), transposed_input_type, deq.getInput(), op.getOperand(1)); + + Value new_scales = deq.getScales(); + if (new_scales && + !(new_scales.getDefiningOp() && + mlir::isa(new_scales.getDefiningOp()))) { + auto scales_type = mlir::dyn_cast(new_scales.getType()); + if (!scales_type || scales_type.getRank() != rank) { + return failure(); + } + SmallVector transposed_scales_shape; + for (int64_t p : perms) { + transposed_scales_shape.push_back(scales_type.getDimSize(p)); + } + auto transposed_scales_type = RankedTensorType::get( + transposed_scales_shape, scales_type.getElementType()); + ElementsAttr scales_attr; + if (matchPattern(new_scales, m_Constant(&scales_attr))) { + if (auto dense_attr = mlir::dyn_cast(scales_attr)) { + if (dense_attr.isSplat()) { + new_scales = rewriter.create( + op.getLoc(), dense_attr.reshape(transposed_scales_type)); + } else { + new_scales = rewriter.create( + op.getLoc(), transposed_scales_type, new_scales, + op.getOperand(1)); + } + } else { + new_scales = + rewriter.create(op.getLoc(), transposed_scales_type, + new_scales, op.getOperand(1)); + } + } else { + new_scales = rewriter.create( + op.getLoc(), transposed_scales_type, new_scales, op.getOperand(1)); + } + } + + Value new_zp = deq.getZeroPoints(); + if (new_zp && !(new_zp.getDefiningOp() && + mlir::isa(new_zp.getDefiningOp()))) { + auto zp_type = mlir::dyn_cast(new_zp.getType()); + if (!zp_type || zp_type.getRank() != rank) { + return failure(); + } + SmallVector transposed_zp_shape; + for (int64_t p : perms) { + transposed_zp_shape.push_back(zp_type.getDimSize(p)); + } + auto transposed_zp_type = + RankedTensorType::get(transposed_zp_shape, zp_type.getElementType()); + ElementsAttr zp_attr; + if (matchPattern(new_zp, m_Constant(&zp_attr))) { + if (auto dense_attr = mlir::dyn_cast(zp_attr)) { + if (dense_attr.isSplat()) { + new_zp = rewriter.create( + op.getLoc(), dense_attr.reshape(transposed_zp_type)); + } else { + new_zp = rewriter.create( + op.getLoc(), transposed_zp_type, new_zp, op.getOperand(1)); + } + } else { + new_zp = rewriter.create(op.getLoc(), transposed_zp_type, + new_zp, op.getOperand(1)); + } + } else { + new_zp = rewriter.create(op.getLoc(), transposed_zp_type, + new_zp, op.getOperand(1)); + } + } + + auto new_deq = rewriter.create( + op.getLoc(), output_type, new_input, new_scales, new_zp, + new_block_shape_attr, deq.getSymmetricAttr()); + rewriter.replaceOp(op, new_deq.getOutput()); + return success(); + } +}; + } // namespace void LargeConstantFoldPass::runOnOperation() { ModuleOp module = getOperation(); + MLIRContext* ctx = &getContext(); + + RewritePatternSet patterns(ctx); + patterns.add(ctx); + GreedyRewriteConfig config; + config.enableFolding(false); + if (failed(applyPatternsGreedily(module, std::move(patterns), config))) { + signalPassFailure(); + return; + } if (failed(FoldResourceOpPattern(module))) { signalPassFailure(); } diff --git a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h index a7ebd380742df7..a396accce5c515 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h +++ b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h @@ -18,9 +18,12 @@ limitations under the License. #include +#include "llvm/ADT/StringRef.h" #include "llvm/Support/Casting.h" +#include "mlir/IR/Attributes.h" // from @llvm-project #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project +#include "mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Types.h" // from @llvm-project @@ -30,106 +33,208 @@ limitations under the License. namespace mlir::TFL { -template -LogicalResult FillCompositeParams(CompositeOpType op, - SmallVector& scales, - SmallVector& zero_points, - int& num_bits, bool& is_signed, - bool& is_narrow_range) { - auto scale_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("scale")); - if (scale_attr == nullptr) { - return failure(); - } - for (auto float_attr : scale_attr.template getValues()) { - scales.push_back(float_attr.getValue().convertToDouble()); - } +// Quantization annotations reach this pass as either a `stablehlo.composite` or +// a `stablehlo.custom_call`. The two carry the same information but spell the +// accessors differently, so these overloads let the templates below handle both +// without duplicating the extraction logic. +inline Attribute GetOpAttribute(stablehlo::CompositeOp op, StringRef name) { + return op.getCompositeAttributes().get(name); +} - auto zero_point_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("zero_point")); - if (zero_point_attr == nullptr) { - for (int i = 0; i < scales.size(); ++i) { - zero_points.push_back(0); - } - } else if (zero_point_attr.isSplat()) { - for (int i = 0; i < scales.size(); ++i) { - zero_points.push_back( - zero_point_attr.template getSplatValue().getInt()); - } - } else { - for (IntegerAttr zp : zero_point_attr.template getValues()) { - zero_points.push_back(zp.getInt()); - } - } +inline Attribute GetOpAttribute(stablehlo::CustomCallOp op, StringRef name) { + return op->getAttr(name); +} - auto dtype_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("dtype")); - if (dtype_attr == nullptr) { - return failure(); +inline StringRef GetQuantAnnotationName(stablehlo::CompositeOp op) { + return op.getName(); +} + +inline StringRef GetQuantAnnotationName(stablehlo::CustomCallOp op) { + return op.getCallTargetName(); +} + +// The integer storage of a quantized tensor, as named by the annotation's +// `dtype` attribute. +struct QuantStorage { + Type type; + int num_bits = 0; + bool is_signed = false; + + // Inclusive bounds of the representable range. `narrow_range` drops the most + // negative value so that the range is symmetric around zero. + // + // This is the single definition of the storage range used by the whole + // lowering: the constant folder that quantizes weights at compile time and + // the runtime kernels must agree on it exactly, or the model's numerics + // change depending on which one ran. + int64_t Min(bool narrow_range) const { + if (!is_signed) return 0; + const int64_t min = -(int64_t{1} << (num_bits - 1)); + return narrow_range ? min + 1 : min; + } + int64_t Max(bool narrow_range) const { + if (!is_signed) return (int64_t{1} << num_bits) - 1; + return (int64_t{1} << (num_bits - 1)) - 1; } - auto dtype = dtype_attr.getValue(); +}; +// Parses the annotation's `dtype` string. Fails on an unrecognized name rather +// than defaulting, because every caller uses the result to decide how to round +// and clamp real values. +inline LogicalResult ParseQuantStorage(StringRef dtype, MLIRContext* ctx, + QuantStorage& storage) { if (dtype == "i2") { - num_bits = 2; - is_signed = true; + storage = {IntegerType::get(ctx, 2), 2, /*is_signed=*/true}; } else if (dtype == "i4") { - num_bits = 4; - is_signed = true; + storage = {IntegerType::get(ctx, 4), 4, /*is_signed=*/true}; } else if (dtype == "ui4") { - num_bits = 4; - is_signed = false; + storage = {IntegerType::get(ctx, 4), 4, /*is_signed=*/false}; } else if (dtype == "i8") { - num_bits = 8; - is_signed = true; + storage = {IntegerType::get(ctx, 8), 8, /*is_signed=*/true}; + } else if (dtype == "ui8") { + storage = {IntegerType::get(ctx, 8), 8, /*is_signed=*/false}; } else if (dtype == "i16") { - num_bits = 16; - is_signed = true; + storage = {IntegerType::get(ctx, 16), 16, /*is_signed=*/true}; } else { return failure(); } - auto narrow_range_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("narrow_range")); - if (narrow_range_attr == nullptr) { + return success(); +} + +// Everything a `quant.*` annotation says about a blockwise quantized tensor. +// +// All quantization parameters are annotation attributes rather than operands: +// the annotation op always has exactly one operand, the value being quantized. +struct BlockwiseAnnotation { + SmallVector block_shape; + QuantStorage storage; + bool narrow_range = false; + bool symmetric = false; + // Element type of the scale. `f8E8M0FNU` constrains the scale to a power of + // two; otherwise `f32`. + Type scale_type; + // Widens the quantized range used to derive a dynamic scale, so that + // outliers clip rather than stretching the whole block. 0 when absent. + double range_dilation = 0.0; + // Null for dynamic range quantization, where the scale is derived from the + // tensor at runtime. + ElementsAttr scale; + // Null when the quantization is symmetric. + ElementsAttr zero_point; + + bool IsDynamic() const { return scale == nullptr; } +}; + +// Returns true if `op` carries a `block_shape` attribute, i.e. it describes +// blockwise rather than per-tensor or per-axis quantization. +template +bool IsBlockwiseAnnotation(OpType op) { + return llvm::isa_and_nonnull( + GetOpAttribute(op, "block_shape")); +} + +// Extracts the blockwise quantization parameters of `op`. Only call this when +// `IsBlockwiseAnnotation(op)` holds. +// +// Fails on any attribute that is present but unrecognized. Falling back to a +// default would produce a model that runs but computes something other than +// what the annotation asked for. +template +LogicalResult ParseBlockwiseAnnotation(OpType op, BlockwiseAnnotation& out) { + MLIRContext* ctx = op->getContext(); + + auto block_shape_attr = llvm::dyn_cast_or_null( + GetOpAttribute(op, "block_shape")); + if (block_shape_attr == nullptr) return failure(); + for (const IntegerAttr block : + block_shape_attr.template getValues()) { + out.block_shape.push_back(block.getInt()); + } + + auto dtype_attr = + llvm::dyn_cast_or_null(GetOpAttribute(op, "dtype")); + if (dtype_attr == nullptr) return failure(); + if (failed(ParseQuantStorage(dtype_attr.getValue(), ctx, out.storage))) { return failure(); } - is_narrow_range = narrow_range_attr.getValue(); - return success(); -} + if (auto narrow_range = llvm::dyn_cast_or_null( + GetOpAttribute(op, "narrow_range"))) { + out.narrow_range = narrow_range.getValue(); + } + if (auto symmetric = + llvm::dyn_cast_or_null(GetOpAttribute(op, "symmetric"))) { + out.symmetric = symmetric.getValue(); + } -template -bool IsDrqFakeQuant(CompositeOpType op) { - if (op.getName() != "quant.fake_quant") { - return false; + out.scale_type = Float32Type::get(ctx); + if (Attribute scale_dtype = GetOpAttribute(op, "act_scale_dtype")) { + auto scale_dtype_str = llvm::dyn_cast(scale_dtype); + if (scale_dtype_str == nullptr) return failure(); + if (scale_dtype_str.getValue() == "e8m0") { + out.scale_type = Float8E8M0FNUType::get(ctx); + } else if (scale_dtype_str.getValue() != "fp32" && + scale_dtype_str.getValue() != "float32") { + return failure(); + } } - SmallVector scales; - SmallVector zero_points; - int num_bits; - bool is_signed; - bool is_narrow_range; - if (failed(FillCompositeParams(op, scales, zero_points, num_bits, is_signed, - is_narrow_range))) { - return false; + + if (Attribute dilation = GetOpAttribute(op, "range_dilation")) { + auto dilation_fp = llvm::dyn_cast(dilation); + if (dilation_fp == nullptr) return failure(); + out.range_dilation = dilation_fp.getValueAsDouble(); } - return scales.empty() && zero_points.empty(); + + // An empty `scale` means dynamic range quantization; see `IsDrqFakeQuant`. + auto scale = + llvm::dyn_cast_or_null(GetOpAttribute(op, "scale")); + if (scale != nullptr && !scale.empty()) out.scale = scale; + + auto zero_point = + llvm::dyn_cast_or_null(GetOpAttribute(op, "zero_point")); + if (zero_point != nullptr && !zero_point.empty()) out.zero_point = zero_point; + + return success(); +} + +// Returns true if `op` carries the static `scale` attribute that describes its +// quantization parameters, i.e. the parameters are baked into the annotation. +// +// An annotation without one describes dynamic range quantization (DRQ), where +// the parameters are derived from the tensor at runtime. This is the exact +// condition under which `FillCompositeParams` below fails, so the two must stay +// in sync. +template +bool HasStaticQuantParams(OpType op) { + auto scale_attr = + llvm::dyn_cast_or_null(GetOpAttribute(op, "scale")); + return scale_attr != nullptr && !scale_attr.empty(); } -template <> -inline LogicalResult FillCompositeParams( - stablehlo::CustomCallOp op, SmallVector& scales, - SmallVector& zero_points, int& num_bits, bool& is_signed, - bool& is_narrow_range) { +// Extracts the quantization parameters from a per-tensor or per-axis `quant.*` +// annotation into `scales`, `zero_points`, `num_bits`, `is_signed` and +// `is_narrow_range`. +// +// Fails if any required attribute is missing or unrecognized. In particular it +// fails on DRQ annotations, which carry no static `scale`; callers that need to +// handle those must check `IsDrqFakeQuant`/`HasStaticQuantParams` first. On +// success `scales` is guaranteed non-empty and `zero_points` has the same size. +template +LogicalResult FillCompositeParams(OpType op, SmallVector& scales, + SmallVector& zero_points, + int& num_bits, bool& is_signed, + bool& is_narrow_range) { auto scale_attr = - llvm::dyn_cast_or_null(op->getAttr("scale")); - if (scale_attr == nullptr) { + llvm::dyn_cast_or_null(GetOpAttribute(op, "scale")); + if (scale_attr == nullptr || scale_attr.empty()) { return failure(); } - for (auto float_attr : scale_attr.getValues()) { + for (auto float_attr : scale_attr.template getValues()) { scales.push_back(float_attr.getValue().convertToDouble()); } - auto zero_point_attr = - llvm::dyn_cast_or_null(op->getAttr("zero_point")); + auto zero_point_attr = llvm::dyn_cast_or_null( + GetOpAttribute(op, "zero_point")); if (zero_point_attr == nullptr) { for (int i = 0; i < scales.size(); ++i) { zero_points.push_back(0); @@ -137,40 +242,29 @@ inline LogicalResult FillCompositeParams( } else if (zero_point_attr.isSplat()) { for (int i = 0; i < scales.size(); ++i) { zero_points.push_back( - zero_point_attr.getSplatValue().getInt()); + zero_point_attr.template getSplatValue().getInt()); } } else { - for (IntegerAttr zp : zero_point_attr.getValues()) { + for (IntegerAttr zp : zero_point_attr.template getValues()) { zero_points.push_back(zp.getInt()); } } - auto dtype_attr = llvm::dyn_cast_or_null(op->getAttr("dtype")); + auto dtype_attr = + llvm::dyn_cast_or_null(GetOpAttribute(op, "dtype")); if (dtype_attr == nullptr) { return failure(); } - auto dtype = dtype_attr.getValue(); - - if (dtype == "i2") { - num_bits = 2; - is_signed = true; - } else if (dtype == "i4") { - num_bits = 4; - is_signed = true; - } else if (dtype == "ui4") { - num_bits = 4; - is_signed = false; - } else if (dtype == "i8") { - num_bits = 8; - is_signed = true; - } else if (dtype == "i16") { - num_bits = 16; - is_signed = true; - } else { + QuantStorage storage; + if (failed(ParseQuantStorage(dtype_attr.getValue(), op->getContext(), + storage))) { return failure(); } + num_bits = storage.num_bits; + is_signed = storage.is_signed; + auto narrow_range_attr = - llvm::dyn_cast_or_null(op->getAttr("narrow_range")); + llvm::dyn_cast_or_null(GetOpAttribute(op, "narrow_range")); if (narrow_range_attr == nullptr) { return failure(); } @@ -179,22 +273,12 @@ inline LogicalResult FillCompositeParams( return success(); } -template <> -inline bool IsDrqFakeQuant( - stablehlo::CustomCallOp op) { - if (op.getCallTargetName() != "quant.fake_quant") { - return false; - } - SmallVector scales; - SmallVector zero_points; - int num_bits; - bool is_signed; - bool is_narrow_range; - if (failed(FillCompositeParams(op, scales, zero_points, num_bits, is_signed, - is_narrow_range))) { +template +bool IsDrqFakeQuant(OpType op) { + if (GetQuantAnnotationName(op) != "quant.fake_quant") { return false; } - return scales.empty() && zero_points.empty(); + return !HasStaticQuantParams(op); } LogicalResult GetStorageParams(unsigned num_bits, bool narrow_range, diff --git a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc index 9ecf7dca971e65..9eb737d5f71d70 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc @@ -15,10 +15,14 @@ limitations under the License. // This transformation pass applies quantization on TFLite dialect. +#include +#include +#include #include #include #include +#include "llvm/ADT/STLExtras.h" #include "llvm/Support/Casting.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project // IWYU pragma: keep @@ -28,6 +32,7 @@ limitations under the License. #include "mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/IR/Diagnostics.h" // from @llvm-project #include "mlir/IR/DialectResourceBlobManager.h" // from @llvm-project // IWYU pragma: keep +#include "mlir/IR/Operation.h" // from @llvm-project #include "mlir/IR/PatternMatch.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project @@ -49,6 +54,284 @@ namespace { #define GEN_PASS_DEF_LOWERQUANTANNOTATIONSPASS #include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc" +//===----------------------------------------------------------------------===// +// Blockwise annotation lowering +//===----------------------------------------------------------------------===// + +// Returns the blob backing `attr`, or null if it has none. +AsmResourceBlob* GetResourceBlob(DenseResourceElementsAttr attr) { + if (AsmResourceBlob* blob = attr.getRawHandle().getBlob()) return blob; + if (auto* resource = attr.getRawHandle().getResource()) { + return resource->getBlob(); + } + return nullptr; +} + +// Appends the elements of a resource blob to `out`, reinterpreting the raw +// bytes according to `element_type`. +// +// A resource attribute has no per-element accessor, so the element type has to +// be dispatched on explicitly. Reading an unexpected type as `float` would +// produce plausible-looking garbage, so anything unhandled fails. +template +void AppendBlobAs(AsmResourceBlob* blob, SmallVectorImpl& out) { + for (T value : blob->getDataAs()) + out.push_back(static_cast(value)); +} + +LogicalResult ReadNumericElements(ElementsAttr attr, + SmallVectorImpl& out) { + if (auto dense_fp = mlir::dyn_cast(attr)) { + for (const APFloat& value : dense_fp.getValues()) { + out.push_back(value.convertToDouble()); + } + return success(); + } + if (auto dense_int = mlir::dyn_cast(attr)) { + for (const APInt& value : dense_int.getValues()) { + out.push_back(static_cast(value.getSExtValue())); + } + return success(); + } + + auto dense_resource = mlir::dyn_cast(attr); + if (dense_resource == nullptr) return failure(); + AsmResourceBlob* blob = GetResourceBlob(dense_resource); + if (blob == nullptr) return failure(); + + const Type element_type = dense_resource.getType().getElementType(); + if (element_type.isF32()) { + AppendBlobAs(blob, out); + } else if (element_type.isF64()) { + AppendBlobAs(blob, out); + } else if (element_type.isInteger(8)) { + AppendBlobAs(blob, out); + } else if (element_type.isInteger(16)) { + AppendBlobAs(blob, out); + } else if (element_type.isInteger(32)) { + AppendBlobAs(blob, out); + } else if (element_type.isInteger(64)) { + AppendBlobAs(blob, out); + } else { + return failure(); + } + return success(); +} + +ElementsAttr GetElementsAttrFromValue(Value value) { + if (!value) return nullptr; + mlir::Operation* defining_op = value.getDefiningOp(); + if (defining_op == nullptr) return nullptr; + if (auto reshape = mlir::dyn_cast(defining_op)) { + return GetElementsAttrFromValue(reshape.getOperand()); + } + if (auto reshape = mlir::dyn_cast(defining_op)) { + return GetElementsAttrFromValue(reshape.getInput()); + } + return defining_op->getAttrOfType("value"); +} + +// Maps each element of a tensor of shape `shape` to its block in a grid tiled +// by `block_shape`, and from there to an index into a parameter tensor of shape +// `param_shape`. A parameter dimension of size 1 broadcasts across that axis. +class BlockIndexer { + public: + BlockIndexer(ArrayRef shape, ArrayRef block_shape, + ArrayRef param_shape) + : shape_(shape), block_shape_(block_shape), param_shape_(param_shape) {} + + int64_t ParamIndexOf(int64_t flat_index) const { + int64_t param_index = 0; + int64_t remainder = flat_index; + // Recover the coordinates from the innermost dimension outwards, then fold + // them into the parameter index from the outermost inwards. + SmallVector block_coords(shape_.size()); + for (int dim = shape_.size() - 1; dim >= 0; --dim) { + block_coords[dim] = (remainder % shape_[dim]) / block_shape_[dim]; + remainder /= shape_[dim]; + } + for (size_t dim = 0; dim < shape_.size(); ++dim) { + const int64_t param_dim = param_shape_[dim]; + param_index = + param_index * param_dim + (param_dim == 1 ? 0 : block_coords[dim]); + } + return param_index; + } + + private: + ArrayRef shape_; + ArrayRef block_shape_; + ArrayRef param_shape_; +}; + +// Quantizes a constant at compile time using the blockwise parameters of +// `annotation`, mirroring `tfl.blockwise_dequantize`: +// +// q = clamp(nearbyint(real / scale + zero_point), q_min, q_max) +// +// `nearbyint` rounds halves to even, matching `jnp.round` on the JAX side. The +// clamp bounds come from `QuantStorage`, which the runtime kernels share. +FailureOr QuantizeBlockwise(ElementsAttr real_values, + const BlockwiseAnnotation& annotation, + Location loc) { + auto real_type = mlir::dyn_cast(real_values.getType()); + if (real_type == nullptr || !real_type.hasStaticShape()) { + return emitError(loc, "blockwise weight must have a static shape"); + } + const ArrayRef shape = real_type.getShape(); + if (annotation.block_shape.size() != shape.size()) { + return emitError(loc, "block_shape rank does not match the weight rank"); + } + for (auto [dim, dim_size] : llvm::enumerate(shape)) { + const int64_t block_size = annotation.block_shape[dim]; + if (block_size <= 0 || dim_size % block_size != 0) { + return emitError(loc, "weight dimension ") + << dim << " (" << dim_size << ") is not a multiple of block_shape[" + << dim << "] (" << block_size << ")"; + } + } + + SmallVector values; + SmallVector scales; + if (failed(ReadNumericElements(real_values, values)) || + failed(ReadNumericElements(annotation.scale, scales))) { + return emitError(loc, "unsupported element type in a blockwise weight"); + } + + SmallVector zero_points; + if (annotation.zero_point != nullptr && + failed(ReadNumericElements(annotation.zero_point, zero_points))) { + return emitError(loc, "unsupported element type in a blockwise zero point"); + } + + const BlockIndexer scale_indexer(shape, annotation.block_shape, + annotation.scale.getShapedType().getShape()); + const BlockIndexer zero_point_indexer( + shape, annotation.block_shape, + zero_points.empty() ? ArrayRef() + : annotation.zero_point.getShapedType().getShape()); + + const double q_min = annotation.storage.Min(annotation.narrow_range); + const double q_max = annotation.storage.Max(annotation.narrow_range); + + SmallVector quantized; + quantized.reserve(values.size()); + for (auto [index, value] : llvm::enumerate(values)) { + const int64_t scale_index = scale_indexer.ParamIndexOf(index); + if (scale_index >= static_cast(scales.size())) { + return emitError(loc, "blockwise scale is smaller than its block grid"); + } + const double scale = scales[scale_index]; + if (scale == 0.0) { + return emitError(loc, "blockwise scale contains a zero"); + } + + double zero_point = 0.0; + if (!zero_points.empty()) { + const int64_t index_of = zero_point_indexer.ParamIndexOf(index); + if (index_of >= static_cast(zero_points.size())) { + return emitError(loc, + "blockwise zero point is smaller than its block grid"); + } + zero_point = zero_points[index_of]; + } + + const double q = + std::clamp(std::nearbyint(value / scale + zero_point), q_min, q_max); + quantized.emplace_back(annotation.storage.num_bits, + static_cast(static_cast(q)), + /*isSigned=*/annotation.storage.is_signed); + } + + return cast(DenseElementsAttr::get( + RankedTensorType::get(shape, annotation.storage.type), quantized)); +} + +// Lowers a blockwise `quant.fake_quant` annotation. +// +// Weights carry a static scale and are quantized here, at compile time, into a +// constant feeding a `tfl.blockwise_dequantize`. Activations carry no scale: +// they get a `tfl.blockwise_quantize` that derives one at runtime, feeding the +// matching dequantize. +// +// Both shapes are collapsed back into a single op later, by +// `FuseA4W2DynamicRangeFullyConnected`. +template +LogicalResult LowerBlockwiseFakeQuant(OpType op, PatternRewriter& rewriter) { + BlockwiseAnnotation annotation; + if (failed(ParseBlockwiseAnnotation(op, annotation))) { + return op.emitError( + "blockwise quant.fake_quant has missing or unsupported attributes"); + } + + // The input is the last operand so that dynamically shaped models, which + // prepend shape operands, keep working. + Value input = op.getOperand(op.getNumOperands() - 1); + auto input_type = cast(input.getType()); + auto output_type = cast(op.getType(0)); + const ArrayAttr block_shape = + rewriter.getI64ArrayAttr(annotation.block_shape); + + Value quantized; + Value scale; + Value zero_point; + + if (annotation.IsDynamic()) { + SmallVector grid_shape; + grid_shape.reserve(input_type.getRank()); + for (auto [dim, dim_size] : llvm::enumerate(input_type.getShape())) { + grid_shape.push_back(dim_size / annotation.block_shape[dim]); + } + // A symmetric quantizer has no zero point to report; the op returns none + // rather than a grid of zeros so that the consumer cannot mistake it for + // a meaningful value. + const SmallVector zero_point_shape(input_type.getRank(), 1); + + auto blockwise_quantize = rewriter.create( + op.getLoc(), + /*output=*/ + RankedTensorType::get(input_type.getShape(), annotation.storage.type), + /*scale=*/RankedTensorType::get(grid_shape, annotation.scale_type), + /*zero_point=*/ + annotation.symmetric ? Type(rewriter.getNoneType()) + : Type(RankedTensorType::get( + zero_point_shape, annotation.storage.type)), + input, block_shape, TypeAttr::get(annotation.scale_type), + rewriter.getBoolAttr(annotation.symmetric), + rewriter.getF32FloatAttr(annotation.range_dilation)); + + quantized = blockwise_quantize.getOutput(); + scale = blockwise_quantize.getScale(); + zero_point = blockwise_quantize.getZeroPoint(); + } else { + ElementsAttr real_values = GetElementsAttrFromValue(input); + if (real_values == nullptr) { + return op.emitError( + "blockwise quant.fake_quant has a static scale but a non-constant " + "input"); + } + FailureOr quantized_values = + QuantizeBlockwise(real_values, annotation, op.getLoc()); + if (failed(quantized_values)) return failure(); + + quantized = rewriter.create(op.getLoc(), *quantized_values); + scale = rewriter.create(op.getLoc(), annotation.scale); + if (annotation.zero_point == nullptr) { + zero_point = rewriter.create( + op.getLoc(), rewriter.getNoneType(), rewriter.getUnitAttr()); + } else { + zero_point = + rewriter.create(op.getLoc(), annotation.zero_point); + } + } + + auto blockwise_dequantize = rewriter.create( + op.getLoc(), output_type, quantized, scale, zero_point, block_shape, + rewriter.getBoolAttr(annotation.symmetric)); + rewriter.replaceOp(op, blockwise_dequantize.getOutput()); + return success(); +} + /** * Replaces a quant.quantize composite op with a TFLite quantize op that outputs * a quantized tensor based on the composite op's attributes. @@ -298,7 +581,13 @@ class RewriteFakeQuantCompositeOp public: LogicalResult matchAndRewrite(stablehlo::CompositeOp op, PatternRewriter& rewriter) const final { - if (op.getName() != "quant.fake_quant" || IsDrqFakeQuant(op)) { + if (op.getName() != "quant.fake_quant") { + return failure(); + } + if (IsBlockwiseAnnotation(op)) { + return LowerBlockwiseFakeQuant(op, rewriter); + } + if (IsDrqFakeQuant(op)) { return failure(); } @@ -645,7 +934,13 @@ void LowerQuantAnnotationsPass::runOnOperation() { public: LogicalResult matchAndRewrite(stablehlo::CustomCallOp op, PatternRewriter& rewriter) const final { - if (op.getCallTargetName() != "quant.fake_quant" || IsDrqFakeQuant(op)) { + if (op.getCallTargetName() != "quant.fake_quant") { + return failure(); + } + if (IsBlockwiseAnnotation(op)) { + return LowerBlockwiseFakeQuant(op, rewriter); + } + if (IsDrqFakeQuant(op)) { return failure(); } diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc index f90ba3d574d314..c8e4ce5c6077d9 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc @@ -53,7 +53,7 @@ bool IsFromFoldableChain(mlir::Value value) { while (defining_op) { // `QConstOp` carries its value in an attribute but is neither ConstantLike // nor foldable, so it has to be named rather than left to `m_Constant`. - if (mlir::isa(defining_op) || + if (mlir::isa(defining_op) || defining_op->hasTrait() || matchPattern(defining_op, mlir::m_Constant())) { return true; diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.h b/tensorflow/compiler/mlir/lite/transforms/passes.h index c9e71b5b75d6c5..4c44edc6e4c844 100644 --- a/tensorflow/compiler/mlir/lite/transforms/passes.h +++ b/tensorflow/compiler/mlir/lite/transforms/passes.h @@ -123,6 +123,9 @@ std::unique_ptr> CreateQuantizePass( std::unique_ptr> CreateDefaultQuantizePass(); +std::unique_ptr> +CreateFoldStablehloConstantTransformsPass(); + std::unique_ptr> CreateLowerQuantAnnotationsPass(); // Creates an instance of the TFLite PropagateQParams pass which propagates diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.td b/tensorflow/compiler/mlir/lite/transforms/passes.td index 5aab64ac05b4d9..a82604770e9ad5 100644 --- a/tensorflow/compiler/mlir/lite/transforms/passes.td +++ b/tensorflow/compiler/mlir/lite/transforms/passes.td @@ -354,6 +354,21 @@ def LowerQuantAnnotationsPass : Pass<"tfl-lower-quant-annotations", "mlir::Modul ]; } +def FoldStablehloConstantTransformsPass : Pass<"tfl-fold-stablehlo-constant-transforms", "mlir::func::FuncOp"> { + let summary = "Folds stablehlo.transpose and stablehlo.reshape operations on constants."; + let description = [{ + Unlike the upstream folders, this handles `DenseResourceElementsAttr` + operands, which is what large weights are stored as by the time they + reach this point in the pipeline. + }]; + let constructor = "CreateFoldStablehloConstantTransformsPass()"; + let dependentDialects = [ + "TFL::TensorFlowLiteDialect", + "mlir::arith::ArithDialect", + "stablehlo::StablehloDialect" + ]; +} + def PropagateQParamsPass : Pass<"tfl-propagate-qparams", "mlir::ModuleOp"> { let summary = "Propagates Quantization Parameters (scale and zero point) information through the graph."; let description = [{ diff --git a/tensorflow/lite/core/api/flatbuffer_conversions.cc b/tensorflow/lite/core/api/flatbuffer_conversions.cc index f0d93ff9d2bfc6..ec33d526dc05b0 100644 --- a/tensorflow/lite/core/api/flatbuffer_conversions.cc +++ b/tensorflow/lite/core/api/flatbuffer_conversions.cc @@ -1590,6 +1590,11 @@ TfLiteStatus ParseFullyConnected(const Operator* op, TF_LITE_ENSURE_STATUS( ConvertTensorType(schema_params->quantized_bias_type(), ¶ms->quantized_bias_type, error_reporter)); + if (const auto* quant_spec = schema_params->quant_spec()) { + // Borrowed from the model buffer; see TfLiteFullyConnectedParams. + params->quant_spec = quant_spec->data(); + params->quant_spec_size = static_cast(quant_spec->size()); + } switch (schema_params->weights_format()) { case FullyConnectedOptionsWeightsFormat_DEFAULT: params->weights_format = kTfLiteFullyConnectedWeightsFormatDefault; diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc index 75f78b09a0b108..9a5773b985ae43 100644 --- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc +++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc @@ -2107,6 +2107,11 @@ class Subgraph { node_index); return kTfLiteError; } + if (params->quant_spec != nullptr && params->quant_spec_size > 0) { + TF_LITE_MAYBE_KERNEL_LOG(context, "unsupported quant_spec in node #%d", + node_index); + return kTfLiteError; + } return kTfLiteOk; } diff --git a/tensorflow/lite/tools/serialization/option_writer_generator.cc b/tensorflow/lite/tools/serialization/option_writer_generator.cc index 2ecc872174a202..04ddae4d16a289 100644 --- a/tensorflow/lite/tools/serialization/option_writer_generator.cc +++ b/tensorflow/lite/tools/serialization/option_writer_generator.cc @@ -468,6 +468,12 @@ void GenerateImportForOp(FILE* fp, const std::string& op_name, vector_name = "scatter_dims_to_operand_dims"; vector_size = "num_scatter_dims_to_operand_dims"; vector_type = "int64_t"; + } else if (elem_name == "quant_spec") { + // Opaque flexbuffers blob; see FullyConnectedOptions.quant_spec. + is_vector = true; + vector_name = "quant_spec"; + vector_size = "quant_spec_size"; + vector_type = "uint8_t"; } if (is_vector) { From 0668dfb00c9435f2a0b23c651f6e78b4e3c54ad1 Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Thu, 24 Sep 2026 11:08:44 -0700 Subject: [PATCH 41/48] [XLA:FFI] Make `//xla/ffi:api` an alias of `//xla/ffi/api:api` PiperOrigin-RevId: 987625577 --- third_party/xla/xla/ffi/BUILD | 8 ++------ third_party/xla/xla/ffi/api/BUILD | 2 +- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/third_party/xla/xla/ffi/BUILD b/third_party/xla/xla/ffi/BUILD index c383a0da055c20..e5137056eb2bca 100644 --- a/third_party/xla/xla/ffi/BUILD +++ b/third_party/xla/xla/ffi/BUILD @@ -31,14 +31,10 @@ package_group( ], ) -cc_library( +alias( name = "api", - hdrs = ["//xla/ffi/api:api_headers"], + actual = "//xla/ffi/api:api", visibility = [":ffi_internal"], - deps = [ - "//xla/ffi/api:c_api", - "@com_google_absl//absl/strings:string_view", - ], ) cc_library( diff --git a/third_party/xla/xla/ffi/api/BUILD b/third_party/xla/xla/ffi/api/BUILD index 287edfde68c27f..23c62b1f8aa99e 100644 --- a/third_party/xla/xla/ffi/api/BUILD +++ b/third_party/xla/xla/ffi/api/BUILD @@ -81,7 +81,7 @@ filegroup( cc_library( name = "api", hdrs = [":api_headers"], - visibility = ["//visibility:private"], + visibility = ["//xla/ffi:__pkg__"], deps = [ ":c_api", ], From 8bebdffc7970d1900870ec0c223b245e641784e5 Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Thu, 24 Sep 2026 11:13:15 -0700 Subject: [PATCH 42/48] [PJRT:CPU] Use separate bounded thread pools for compilation, execution, and transfers. - Introduce dedicated bounded thread pools in `PjRtCpuRawClient` for compilation (`compile_thread_pool_`) and execution (`execute_work_runner_`), separate from the general `async_work_runner_` used for buffer transfers and linearization/delinearization. - Route `CompileInternal` to `compile_thread_pool_` and `CpuPjRtRawLoadedExecutable::Execute` to `execute_work_runner_`. Separating execution, compilation, and transfer thread pools prevents multi-device collective launches from starving antecedent input buffer transfers on `async_work_runner_`, while keeping all three pools bounded so bursts of D2H/H2D buffer transfers remain throttled. PiperOrigin-RevId: 987628676 --- third_party/xla/xla/pjrt/cpu/BUILD | 1 - third_party/xla/xla/pjrt/cpu/cpu_client.cc | 16 ++++++++++------ third_party/xla/xla/pjrt/cpu/cpu_client.h | 10 ++++++++++ 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/third_party/xla/xla/pjrt/cpu/BUILD b/third_party/xla/xla/pjrt/cpu/BUILD index 8103e9d1f988c6..bb92234b2867c7 100644 --- a/third_party/xla/xla/pjrt/cpu/BUILD +++ b/third_party/xla/xla/pjrt/cpu/BUILD @@ -248,7 +248,6 @@ cc_library( "//xla/service/cpu:cpu_executable", "//xla/service/cpu:cpu_executable_run_options", "//xla/service/cpu:cpu_xfeed", - "//xla/service/cpu:executable_proto_cc", "//xla/service/llvm_ir:llvm_command_line_options", "//xla/stream_executor:device_address", "//xla/tsl/concurrency:async_value", diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc index e30a0fe49f7877..0873c87b8560c9 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc @@ -71,6 +71,7 @@ limitations under the License. #include "xla/pjrt/cpu/cpu_async_execution_tracker.h" #include "xla/pjrt/cpu/cpu_device_memory.h" #include "xla/pjrt/cpu/cpu_event.h" +#include "xla/pjrt/cpu/execution_stream_event_map.h" #include "xla/pjrt/cpu/raw_buffer.h" #include "xla/pjrt/device_event.h" #include "xla/pjrt/device_event_utils.h" @@ -107,7 +108,6 @@ limitations under the License. #include "xla/service/cpu/cpu_executable.h" #include "xla/service/cpu/cpu_executable_run_options.h" #include "xla/service/cpu/cpu_xfeed.h" -#include "xla/service/cpu/executable.pb.h" #include "xla/service/device_assignment.h" #include "xla/service/dump.h" #include "xla/service/executable.h" @@ -129,7 +129,6 @@ limitations under the License. #include "xla/tsl/platform/threadpool.h" #include "xla/util.h" #include "xla/xla.pb.h" -#include "xla/xla_data.pb.h" #include "tsl/platform/denormal.h" #include "tsl/platform/fingerprint.h" #include "tsl/platform/setround.h" @@ -460,10 +459,15 @@ PjRtCpuRawClient::PjRtCpuRawClient( eigen_intraop_pool_->NumThreads()) : std::unique_ptr()), custom_intraop_device_(intra_op_device), + compile_thread_pool_(std::make_unique( + tsl::Env::Default(), GetThreadOptions(), "XLACompile", num_threads)), + execute_work_runner_(std::make_unique( + tsl::Env::Default(), "XLAExecute", num_threads, GetThreadOptions())), async_work_runner_(std::make_unique( - tsl::Env::Default(), "XLAPjRtCpuClient", num_threads)) {} + tsl::Env::Default(), "XLAPjRtCpuClient", num_threads, + GetThreadOptions())) {} -PjRtCpuRawClient::~PjRtCpuRawClient() {} +PjRtCpuRawClient::~PjRtCpuRawClient() = default; PjRtPluginAttributes GetDefaultCpuPluginAttributes() { PjRtPluginAttributes attrs; @@ -994,7 +998,7 @@ PjRtCpuRawClient::CompileInternal( params.layout_canonicalization_callback = std::move(layout_canonicalization_callback); params.num_threads = num_threads; - params.compile_thread_pool = async_work_runner()->thread_pool(); + params.compile_thread_pool = compile_thread_pool_.get(); params.aot_options = aot_options; params.process_index = process_index; params.collectives_exists = (collectives() != nullptr); @@ -1697,7 +1701,7 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( run_id_.ToInt(), std::move(ready_on_exit).Release()); PjRtDeviceEventSpan events_ref(input_deps); xla::ExecuteWhenReady( - events_ref, raw_client->async_work_runner(), + events_ref, raw_client->execute_work_runner(), [cpu_executable, buffer_alloc = std::move(buffer_alloc), buffer_alloc_and_copy = std::move(buffer_alloc_and_copy), execute_thunks = std::move(execute_thunks), diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h index e249866e775ada..c467ebe46834be 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h @@ -122,6 +122,14 @@ class PjRtCpuRawClient : public PjRtRawClient { return async_work_runner_.get(); } + ThreadPoolAsyncWorkRunner* execute_work_runner() const { + return execute_work_runner_.get(); + } + + tsl::thread::ThreadPool* compile_thread_pool() const { + return compile_thread_pool_.get(); + } + tsl::thread::ThreadPool* eigen_intraop_pool() const { return eigen_intraop_pool_.get(); } @@ -292,6 +300,8 @@ class PjRtCpuRawClient : public PjRtRawClient { std::unique_ptr eigen_intraop_pool_; std::unique_ptr eigen_intraop_device_; const Eigen::ThreadPoolDevice* custom_intraop_device_ = nullptr; + std::unique_ptr compile_thread_pool_; + std::unique_ptr execute_work_runner_; std::unique_ptr async_work_runner_; }; From 7a33d5986cb57524fb0a2e955f3f7b4a9d1aad61 Mon Sep 17 00:00:00 2001 From: Zaim Codi Date: Thu, 24 Sep 2026 11:16:25 -0700 Subject: [PATCH 43/48] Add first-class kShuffle instruction and builder API across HLO and platforms. Introduce `HloOpcode::kShuffle` (`HloShuffleInstruction`), which moves the elements of its operand around along a set of `dimensions`. A single opcode covers the whole family of data-movement patterns (rotate, reverse, permute, ...), which keeps the semantic intent available to compiler analysis, layout assignment, and target-specific lowerings instead of forcing an expansion into slice and concatenate during graph construction. The pattern to apply is selected by `xla.ShuffleMode`, a `oneof` pairing each mode with exactly the attributes that parameterize it, so a mode cannot be combined with another mode's attributes and adding a mode adds no fields to `HloInstructionProto`. Switches over `mode()` are exhaustive and have no `default`, so a new mode does not compile until every consumer has handled it. Only `rotate` is implemented here. It takes one `shifts` entry per shuffled dimension and rotates the elements to the left. Specifically: * **HLO IR & Parser**: `HloShuffleInstruction` carries `dimensions` and a `ShuffleMode`, with text serialization and parsing. The text form is `shuffle(x), dimensions={0,1}, mode=rotate, shifts={2,5}`, where the name of a mode is the name of its field in `ShuffleMode`. * **Shape Inference & Verifier**: `ShapeInference::InferShuffleShape` and the verifier reject duplicate or out-of-range `dimensions`, a missing mode, and mode attributes inconsistent with the mode (for `rotate`, `dimensions` and `shifts` of unequal size). * **XlaBuilder API**: `xla::Shuffle` constructs shuffles from symbolic handles. The new `shuffle` utility provides `MakeRotateMode` to build a rotate mode and `NormalizeShift` to map an arbitrary shift onto `[0, dim_size)`. * **HLO-to-MHLO / StableHLO Translation**: `HloFunctionImporter` imports `kShuffle`, lowering each rotated dimension sequentially into `stablehlo.slice` and `stablehlo.concatenate` (`concat(slice(shift:), slice(0:shift))`). * **Pass Integration**: Register `kShuffle` across core compiler passes including instruction fusion, layout assignment, and sharding propagation. PiperOrigin-RevId: 987630885 --- third_party/xla/docs/operation_semantics.md | 43 ++++ third_party/xla/xla/BUILD | 23 ++ .../backends/gpu/codegen/triton/support.cc | 1 + .../gpu/codegen/triton/support_test.cc | 1 + third_party/xla/xla/hlo/builder/BUILD | 1 + .../xla/xla/hlo/builder/xla_builder.cc | 27 +++ third_party/xla/xla/hlo/builder/xla_builder.h | 17 ++ .../xla/xla/hlo/builder/xla_builder_test.cc | 15 ++ third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h | 1 + .../xla/hlo/ir/dfs_hlo_visitor_with_default.h | 3 + third_party/xla/xla/hlo/ir/hlo_instruction.cc | 38 ++++ third_party/xla/xla/hlo/ir/hlo_instruction.h | 9 + .../xla/xla/hlo/ir/hlo_instructions.cc | 53 +++++ third_party/xla/xla/hlo/ir/hlo_instructions.h | 42 ++++ third_party/xla/xla/hlo/ir/hlo_opcode.h | 1 + third_party/xla/xla/hlo/parser/BUILD | 1 + third_party/xla/xla/hlo/parser/hlo_parser.cc | 66 ++++++ .../xla/xla/hlo/parser/hlo_parser_test.cc | 201 ++++++++++++++++++ .../tools/hlo_diff/hlo_diff_summary_test.cc | 2 +- .../xla/xla/hlo/translate/hlo_to_mhlo/BUILD | 2 - .../hlo_to_mhlo/hlo_function_importer.cc | 29 +-- third_party/xla/xla/hlo/utils/hlo_matchers.h | 1 + third_party/xla/xla/python/_hlo.pyi | 34 +-- third_party/xla/xla/service/BUILD | 3 + .../xla/service/cpu/cpu_instruction_fusion.cc | 1 + third_party/xla/xla/service/float_support.cc | 1 + third_party/xla/xla/service/hlo.proto | 5 +- .../xla/xla/service/hlo_cost_analysis.cc | 4 + .../xla/xla/service/hlo_cost_analysis.h | 1 + .../xla/xla/service/hlo_graph_dumper.cc | 1 + .../xla/xla/service/hlo_instruction_test.cc | 15 ++ third_party/xla/xla/service/hlo_verifier.cc | 8 + third_party/xla/xla/service/hlo_verifier.h | 1 + .../xla/xla/service/hlo_verifier_test.cc | 70 ++++++ .../xla/xla/service/instruction_fusion.cc | 1 + .../xla/xla/service/layout_assignment.cc | 1 + .../xla/xla/service/shape_inference.cc | 33 +++ third_party/xla/xla/service/shape_inference.h | 6 + .../xla/xla/service/shape_inference_test.cc | 53 +++++ .../xla/xla/service/sharding_propagation.cc | 1 + third_party/xla/xla/shuffle.cc | 39 ++++ third_party/xla/xla/shuffle.h | 41 ++++ third_party/xla/xla/shuffle_test.cc | 74 +++++++ third_party/xla/xla/xla_data.proto | 29 +++ 44 files changed, 965 insertions(+), 34 deletions(-) create mode 100644 third_party/xla/xla/shuffle.cc create mode 100644 third_party/xla/xla/shuffle.h create mode 100644 third_party/xla/xla/shuffle_test.cc diff --git a/third_party/xla/docs/operation_semantics.md b/third_party/xla/docs/operation_semantics.md index 84f0a0cb5c177d..ffc81985ea8642 100644 --- a/third_party/xla/docs/operation_semantics.md +++ b/third_party/xla/docs/operation_semantics.md @@ -5170,6 +5170,49 @@ in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - shift_right_logical](https://openxla.org/stablehlo/spec#shift_right_logical). +## Shuffle + +See also +[`XlaBuilder::Shuffle`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). + +**`Shuffle(operand, dimensions, mode)`** + +Arguments | Type | Semantics +------------- | ------------------- | ------------------------------------- +`operand` | `XlaOp` | array of type T +`dimensions` | `ArraySlice` | dimensions to shuffle +`mode` | `ShuffleMode` | the shuffle pattern to apply, together with the attributes of that pattern + +Shuffles the elements of the `operand` array along the specified `dimensions`, +generating an output array of the same shape. Because a shuffle only moves +elements around, it never changes the shape of the operand. + +The `mode` selects which pattern is applied, and carries the attributes of that +pattern: + +Mode | Attributes | Semantics +-------- | ---------- | --------------------------------------------------- +`rotate` | `shifts` | Rotates the elements to the left by `shifts` + +### `rotate` + +Rotates the elements of the `operand` array along the specified `dimensions` to +the left by the corresponding `shifts`, generating an output array of the same +shape. Shifts outside of `[0, size(dimension))` are supported, and wrap around +cyclically. Each element of the output array at a multidimensional index is +retrieved from the operand array at a transformed index. The multidimensional +index is transformed by rotating the index along each specified dimension (i.e., +for a dimension of size N with a corresponding shift S, the element at output +index i is taken from operand index (i + S) % N). + +1-dimensional example: + +```cpp +let a = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; +Shuffle(a, {0}, shuffle::Rotate({4})) +// Result: {4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3} +``` + ## Sign See also diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index 05dc9d4612307c..b720b59902e79c 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -1192,6 +1192,29 @@ cc_library( ], ) +cc_library( + name = "shuffle", + srcs = ["shuffle.cc"], + hdrs = ["shuffle.h"], + visibility = ["//visibility:public"], + deps = [ + ":xla_data_proto_cc", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/types:span", + ], +) + +xla_cc_test( + name = "shuffle_test", + srcs = ["shuffle_test.cc"], + deps = [ + ":shuffle", + ":xla_data_proto_cc", + "//xla/tsl/platform:test_main", + "@com_google_googletest//:gtest", + ], +) + cc_library( name = "window_util", srcs = ["window_util.cc"], diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc index 42d2deeca92700..c9d4e9d9d3a6ad 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc @@ -937,6 +937,7 @@ bool IsTritonUnsupportedOpcode(HloOpcode opcode) { case HloOpcode::kScatter: case HloOpcode::kSelectAndScatter: case HloOpcode::kSetDimensionSize: + case HloOpcode::kShuffle: return true; default: return false; diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc index e30edaf4702c5f..8ee5570c30a039 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc @@ -3737,6 +3737,7 @@ constexpr std::array kUnsupportedOps = { HloOpcode::kScatter, HloOpcode::kSelectAndScatter, HloOpcode::kSetDimensionSize, + HloOpcode::kShuffle, // go/keep-sorted end // clang-format on }; diff --git a/third_party/xla/xla/hlo/builder/BUILD b/third_party/xla/xla/hlo/builder/BUILD index 9087f57cfc2ca3..8eb58ff87a003e 100644 --- a/third_party/xla/xla/hlo/builder/BUILD +++ b/third_party/xla/xla/hlo/builder/BUILD @@ -209,6 +209,7 @@ xla_cc_test( "//xla:comparison_util", "//xla:debug_options_flags", "//xla:shape_util", + "//xla:shuffle", "//xla:tuple_tree", "//xla:util", "//xla:xla_data_proto_cc", diff --git a/third_party/xla/xla/hlo/builder/xla_builder.cc b/third_party/xla/xla/hlo/builder/xla_builder.cc index a672fa261726b7..7072aaf18ec7a0 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.cc +++ b/third_party/xla/xla/hlo/builder/xla_builder.cc @@ -3126,6 +3126,28 @@ absl::StatusOr XlaBuilder::RevInternal( return AddInstruction(std::move(instr), HloOpcode::kReverse, {operand}); } +XlaOp XlaBuilder::Shuffle(XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode) { + return ReportErrorOrReturn([&]() -> absl::StatusOr { + ABSL_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); + ABSL_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferShuffleShape( + *operand_shape, dimensions, mode)); + return ShuffleInternal(shape, operand, dimensions, mode); + }); +} + +absl::StatusOr XlaBuilder::ShuffleInternal( + const Shape& shape, XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode) { + HloInstructionProto instr; + *instr.mutable_shape() = shape.ToProto(); + for (int64_t dim : dimensions) { + instr.add_dimensions(dim); + } + *instr.mutable_shuffle_mode() = mode; + return AddInstruction(std::move(instr), HloOpcode::kShuffle, {operand}); +} + XlaOp XlaBuilder::Sort(absl::Span operands, XlaComputationId comparator, int64_t dimension, bool is_stable) { @@ -6672,6 +6694,11 @@ XlaOp Rev(const XlaOp operand, absl::Span dimensions) { return operand.builder()->Rev(operand, dimensions); } +XlaOp Shuffle(const XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode) { + return operand.builder()->Shuffle(operand, dimensions, mode); +} + XlaOp Sort(absl::Span operands, const XlaComputation& comparator, int64_t dimension, bool is_stable) { return Sort(operands, operands[0].builder()->AddSubComputation(comparator), diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h index 3c0e937a8c1088..d91e6592ebae4a 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.h +++ b/third_party/xla/xla/hlo/builder/xla_builder.h @@ -1111,6 +1111,12 @@ class XlaBuilder { virtual absl::StatusOr RevInternal( const Shape& shape, XlaOp operand, absl::Span dimensions); + XlaOp Shuffle(XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode); + virtual absl::StatusOr ShuffleInternal( + const Shape& shape, XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode); + XlaOp Sort(absl::Span operands, XlaComputationId comparator, int64_t dimension = -1, bool is_stable = false); virtual absl::StatusOr SortInternal(const Shape& shape, @@ -1991,6 +1997,8 @@ class XlaBuilder { friend XlaOp Neg(XlaOp operand); friend XlaOp Transpose(XlaOp operand, absl::Span permutation); friend XlaOp Rev(XlaOp operand, absl::Span dimensions); + friend XlaOp Shuffle(XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode); friend XlaOp Sort(absl::Span operands, const XlaComputation& comparator, int64_t dimension, bool is_stable); @@ -3413,6 +3421,15 @@ XlaOp Transpose(XlaOp operand, absl::Span permutation); // is moved to index dimension_size - 1 - i). XlaOp Rev(XlaOp operand, absl::Span dimensions); +// Enqueues a shuffle instruction onto the computation. The elements are +// shuffled along the given dimensions following the pattern selected by `mode`, +// which also carries the attributes of that mode: +// - rotate: the elements are (left) rotated by `shifts` along the given +// dimensions (i.e. the element at index i in the output is taken from index +// (i + shifts[j]) % dimension_size[j] of the operand). +XlaOp Shuffle(XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode); + // Enqueues a sort instruction onto the computation, using 'comparator' for // comparisons. 'comparator' needs to define a strict weak order. 'is_stable' // determines whether the stable sorting should be used. diff --git a/third_party/xla/xla/hlo/builder/xla_builder_test.cc b/third_party/xla/xla/hlo/builder/xla_builder_test.cc index 4aa44b97be8114..05b30faa5a41e8 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder_test.cc +++ b/third_party/xla/xla/hlo/builder/xla_builder_test.cc @@ -59,6 +59,7 @@ limitations under the License. #include "xla/service/pattern_matcher.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/shuffle.h" #include "xla/tsl/platform/statusor.h" #include "xla/tuple_tree.h" #include "xla/util.h" @@ -3586,6 +3587,20 @@ TEST(XlaBuilderTest, UnboundedReverse) { GmockMatch(m::Op().WithShapeEqualTo(&expected))); } +TEST(XlaBuilderTest, UnboundedShuffle) { + XlaBuilder b(TestName()); + ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 10]")); + ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[?, 10]")); + + Shuffle(Parameter(&b, 0, operand, "operand"), /*dimensions=*/{0, 1}, + shuffle::Rotate(/*shifts=*/{1, 3})); + ASSERT_OK_AND_ASSIGN(const std::unique_ptr module, + BuildHloModule(b)); + + EXPECT_THAT(GetRoot(*module), + GmockMatch(m::Op().WithShapeEqualTo(&expected))); +} + TEST(XlaBuilderTest, UnboundedRngBitGenerator) { XlaBuilder b(TestName()); TF_ASSERT_OK_AND_ASSIGN(const Shape initial_state, ParseShape("u32[?, 10]")); diff --git a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h index 592af8fb1d8553..efd83fd73ba18e 100644 --- a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h +++ b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h @@ -318,6 +318,7 @@ class DfsHloVisitorBase { virtual absl::Status HandleScan(HloInstructionPtr hlo) = 0; virtual absl::Status HandleScatter(HloInstructionPtr hlo) = 0; virtual absl::Status HandleSelectAndScatter(HloInstructionPtr hlo) = 0; + virtual absl::Status HandleShuffle(HloInstructionPtr hlo) = 0; virtual absl::Status HandleSlice(HloInstructionPtr hlo) = 0; virtual absl::Status HandleSort(HloInstructionPtr hlo) = 0; virtual absl::Status HandleTranspose(HloInstructionPtr hlo) = 0; diff --git a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h index 24177f47f7be54..b691eb7ed027a0 100644 --- a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h +++ b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h @@ -170,6 +170,9 @@ class DfsHloVisitorWithDefaultBase absl::Status HandleReverse(HloInstructionPtr reverse) override { return DefaultAction(reverse); } + absl::Status HandleShuffle(HloInstructionPtr shuffle) override { + return DefaultAction(shuffle); + } absl::Status HandleSort(HloInstructionPtr sort) override { return DefaultAction(sort); } diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index d21bebe569b702..273ace08d474d4 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -588,6 +588,13 @@ absl::StatusOr> HloInstruction::CreateFromProto( std::vector(proto.dimensions().begin(), proto.dimensions().end())); break; + case HloOpcode::kShuffle: + instruction = + CreateShuffle(shape, operands(0), + std::vector(proto.dimensions().begin(), + proto.dimensions().end()), + proto.shuffle_mode()); + break; case HloOpcode::kConcatenate: TF_RET_CHECK(proto.dimensions().size() == 1) << "Concatenate instruction should have 1 dimension but sees " @@ -2157,6 +2164,13 @@ HloInstruction::CreateCollectivePermuteStart( return std::make_unique(shape, operand, dimensions); } +/* static */ std::unique_ptr HloInstruction::CreateShuffle( + const Shape& shape, HloInstruction* operand, + absl::Span dimensions, const ShuffleMode& mode) { + return std::make_unique(shape, operand, dimensions, + mode); +} + /* static */ std::unique_ptr HloInstruction::CreateAfterAll( absl::Span operands) { CHECK(!operands.empty()); @@ -2855,6 +2869,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kRecv: case HloOpcode::kRecvDone: case HloOpcode::kReverse: + case HloOpcode::kShuffle: case HloOpcode::kConcatenate: case HloOpcode::kReduce: case HloOpcode::kTranspose: @@ -3570,6 +3585,7 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kTriangularSolve: case HloOpcode::kCholesky: case HloOpcode::kTopK: + case HloOpcode::kShuffle: LOG(FATAL) << "Base class impl called for opcode with subclass: " << opcode(); } @@ -5102,6 +5118,8 @@ absl::Status HloInstruction::Visit( return visitor->HandleTranspose(this); case HloOpcode::kReverse: return visitor->HandleReverse(this); + case HloOpcode::kShuffle: + return visitor->HandleShuffle(this); case HloOpcode::kReducePrecision: return visitor->HandleReducePrecision(this); case HloOpcode::kSlice: @@ -5421,6 +5439,7 @@ static UseKind OperandElementUse(const HloInstruction& instr, case HloOpcode::kSlice: case HloOpcode::kTranspose: case HloOpcode::kGather: + case HloOpcode::kShuffle: return UseKind::kUse; case HloOpcode::kPad: // Pad reuses the padding value but not the padded array elements. @@ -5608,6 +5627,15 @@ std::string PrecisionToString(const PrecisionConfig::Precision& precision) { return absl::AsciiStrToLower(PrecisionConfig::Precision_Name(precision)); } +std::string ShuffleModeToString(ShuffleMode::ModeCase shuffle_mode) { + switch (shuffle_mode) { + case ShuffleMode::kRotate: + return "rotate"; + case ShuffleMode::MODE_NOT_SET: + return "invalid"; + } +} + template void AbslStringify(Sink& sink, const ResultAccuracy::Tolerance& tolerance) { absl::Format(&sink, "tolerance={atol=%v,rtol=%v,ulps=%v}", tolerance.atol(), @@ -5870,6 +5898,16 @@ absl::StatusOr StringToPrecision( "precision"); } +absl::StatusOr StringToShuffleMode( + absl::string_view mode) { + if (mode == "rotate") { + return ShuffleMode::ModeCase::kRotate; + } + return InvalidArgument("Unknown shuffle mode: %s", mode); + // return StringToEnum(mode, ShuffleModeToString, + // "shuffle mode"); +} + absl::StatusOr StringToResultAccuracy( absl::string_view name) { return StringToEnum(name, ResultAccuracyToString, diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h index aa10a9def8c674..a899406d8b6142 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.h +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h @@ -1214,6 +1214,12 @@ class HloInstruction { const Shape& shape, HloInstruction* operand, absl::Span dimensions); + // Creates a shuffle instruction, which shuffles the elements of `operand` + // along the given dimensions following the pattern selected by `mode`. + static std::unique_ptr CreateShuffle( + const Shape& shape, HloInstruction* operand, + absl::Span dimensions, const ShuffleMode& mode); + // Creates a Afterall instruction used for joining or creating new values of // token type which thread through side-effecting operations. Operands must // all be tokens, calls without operands generates a token. @@ -2955,6 +2961,7 @@ std::string ResultAccuracyToleranceToString( std::string RandomAlgorithmToString(const RandomAlgorithm& algorithm); std::string RandomDistributionToString(const RandomDistribution& distribution); std::string PrecisionToString(const PrecisionConfig::Precision& precision); +std::string ShuffleModeToString(ShuffleMode::ModeCase shuffle_mode); std::string ResultAccuracyToString(ResultAccuracy::Mode accuracy_mode); std::string AlgorithmToString(const PrecisionConfig::Algorithm& algorithm); std::string DotDimensionNumbersToString(const DotDimensionNumbers& dnums); @@ -2973,6 +2980,8 @@ absl::StatusOr StringToRandomDistribution( const std::string& name); absl::StatusOr StringToPrecision( const std::string& name); +absl::StatusOr StringToShuffleMode( + absl::string_view mode); absl::StatusOr StringToAlgorithm( const std::string& name); absl::StatusOr StringToResultAccuracy( diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc index ce9d59cefa5a12..4071f0c416a838 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc @@ -1984,6 +1984,59 @@ std::unique_ptr HloReduceInstruction::CloneWithNewOperandsImpl( dimensions(), to_apply()); } +HloShuffleInstruction::HloShuffleInstruction( + const Shape& shape, HloInstruction* operand, + absl::Span dimensions, const ShuffleMode& mode) + : HloDimensionsInstruction(HloOpcode::kShuffle, shape, dimensions), + mode_(mode) { + AppendOperand(operand); +} + +bool HloShuffleInstruction::IdenticalSlowPath( + const HloInstruction& other, + absl::FunctionRef + eq_computations) const { + const auto& casted_other = static_cast(other); + // Shuffle results are determined by the shuffled dimensions and by the mode + // together with its attributes. + return dimensions() == casted_other.dimensions() && + protobuf_util::HaveSameSerialization(mode_, casted_other.mode_); +} + +std::unique_ptr HloShuffleInstruction::CloneWithNewOperandsImpl( + const Shape& shape, absl::Span new_operands, + HloCloneContext* context) const { + CHECK_EQ(new_operands.size(), 1); + return std::make_unique(shape, new_operands[0], + dimensions(), mode_); +} + +void HloShuffleInstruction::PrintExtraAttributesImpl( + AttributePrinter& printer, const HloPrintOptions& options) const { + HloDimensionsInstruction::PrintExtraAttributesImpl(printer, options); + printer.Next([this](Printer* printer) { + printer->Append("mode="); + printer->Append(ShuffleModeToString(mode())); + }); + // Each mode's inner attributes. + switch (mode()) { + case ShuffleMode::kRotate: + printer.Next([this](Printer* printer) { + printer->Append("shifts={"); + AppendJoin(printer, rotate().shifts(), ","); + printer->Append("}"); + }); + break; + case ShuffleMode::MODE_NOT_SET: + break; + } +} + +void HloShuffleInstruction::ToProto(HloInstructionProto* proto) const { + HloDimensionsInstruction::ToProto(proto); + *proto->mutable_shuffle_mode() = mode_; +} + HloScanInstruction::HloScanInstruction(const Shape& shape, absl::Span inputs, absl::Span inits, diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h index 054e036d666b7d..f67a1714749161 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.h +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h @@ -72,6 +72,7 @@ class HloDimensionsInstruction : public HloInstruction { case HloOpcode::kConcatenate: case HloOpcode::kReduce: case HloOpcode::kReverse: + case HloOpcode::kShuffle: case HloOpcode::kSort: case HloOpcode::kTranspose: return true; @@ -1178,6 +1179,47 @@ class HloReverseInstruction : public HloDimensionsInstruction { HloCloneContext* context) const override; }; +class HloShuffleInstruction : public HloDimensionsInstruction { + public: + explicit HloShuffleInstruction(const Shape& shape, HloInstruction* operand, + absl::Span dimensions, + const ShuffleMode& mode); + + static bool ClassOf(const HloInstruction* hlo) { + return hlo->opcode() == HloOpcode::kShuffle; + } + + const ShuffleMode& shuffle_mode() const { return mode_; } + ShuffleMode* mutable_shuffle_mode() { return &mode_; } + ShuffleMode::ModeCase mode() const { return mode_.mode_case(); } + void ToProto(HloInstructionProto* proto) const override; + + // Accessors for the inner attributes of each mode. + const ShuffleMode::Rotate& rotate() const { + CHECK(mode_.has_rotate()); + return mode_.rotate(); + } + ShuffleMode::Rotate* mutable_rotate() { + CHECK(mode_.has_rotate()); + return mode_.mutable_rotate(); + } + + private: + std::unique_ptr CloneWithNewOperandsImpl( + const Shape& shape, absl::Span new_operands, + HloCloneContext* context) const override; + + bool IdenticalSlowPath( + const HloInstruction& other, + absl::FunctionRef + eq_computations) const override; + + void PrintExtraAttributesImpl(AttributePrinter& printer, + const HloPrintOptions& options) const override; + + ShuffleMode mode_; +}; + class HloConcatenateInstruction : public HloDimensionsInstruction { public: explicit HloConcatenateInstruction(const Shape& shape, diff --git a/third_party/xla/xla/hlo/ir/hlo_opcode.h b/third_party/xla/xla/hlo/ir/hlo_opcode.h index b25434781f4158..9da5dc0342a2d7 100644 --- a/third_party/xla/xla/hlo/ir/hlo_opcode.h +++ b/third_party/xla/xla/hlo/ir/hlo_opcode.h @@ -165,6 +165,7 @@ namespace xla { V(kShiftLeft, "shift-left", 2) \ V(kShiftRightArithmetic, "shift-right-arithmetic", 2) \ V(kShiftRightLogical, "shift-right-logical", 2) \ + V(kShuffle, "shuffle", 1) \ V(kSign, "sign", 1) \ V(kSin, "sine", 1) \ V(kSinh, "sinh", 1) \ diff --git a/third_party/xla/xla/hlo/parser/BUILD b/third_party/xla/xla/hlo/parser/BUILD index 6b90eb2a45b9d3..1d0bbd1a6dd6b4 100644 --- a/third_party/xla/xla/hlo/parser/BUILD +++ b/third_party/xla/xla/hlo/parser/BUILD @@ -33,6 +33,7 @@ cc_library( "//xla:literal_util", "//xla:shape_layout", "//xla:shape_util", + "//xla:shuffle", "//xla:tuple_tree", "//xla:types", "//xla:util", diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc index 518f3aee55f2c2..268eedd5fbd1e9 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc @@ -87,6 +87,7 @@ limitations under the License. #include "xla/shape.h" #include "xla/shape_layout.h" #include "xla/shape_util.h" +#include "xla/shuffle.h" #include "xla/tsl/lib/gtl/map_util.h" #include "xla/tsl/platform/logging.h" #include "xla/tuple_tree.h" @@ -200,6 +201,7 @@ bool CanInferShape(HloOpcode code) { case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: case HloOpcode::kShiftRightLogical: + case HloOpcode::kShuffle: case HloOpcode::kSign: case HloOpcode::kSin: case HloOpcode::kSinh: @@ -344,6 +346,7 @@ class HloParserImpl : public HloParser { kShapeList, kEnum, kRandomAlgorithm, + kShuffleMode, kPrecisionAlgorithm, kResultAccuracyType, kAliasing, @@ -636,6 +639,7 @@ class HloParserImpl : public HloParser { bool ParseRandomDistribution(RandomDistribution* result); bool ParseConvKind(ConvolutionKind* result); bool ParseRandomAlgorithm(RandomAlgorithm* result); + bool ParseShuffleMode(ShuffleMode::ModeCase* result); bool ParsePrecision(PrecisionConfig::Precision* result); bool ParseAlgorithm(PrecisionConfig::Algorithm* result); bool ParseResultAccuracyType(ResultAccuracy::Mode* result); @@ -3137,6 +3141,43 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT return builder->AddInstruction( HloInstruction::CreateReverse(*shape, operands[0], *dimensions)); } + case HloOpcode::kShuffle: { + optional> dimensions; + attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, + &dimensions}; + optional mode; + attrs["mode"] = {/*required=*/true, AttrTy::kShuffleMode, &mode}; + // Attributes below are specific to a subset of the modes, so they are + // optional here and checked against the mode afterwards. + optional> shifts; + attrs["shifts"] = {/*required=*/false, AttrTy::kBracedInt64List, &shifts}; + if ((!preset_operands && + !ParseOperands(&operands, builder, /*expected_size=*/1)) || + !ParseAttributes(attrs, allow_attributes, shape)) { + return nullptr; + } + ShuffleMode shuffle_mode; + switch (*mode) { + case ShuffleMode::kRotate: + if (!shifts.has_value()) { + TokenError("expects shifts for a shuffle in rotate mode"); + return nullptr; + } + shuffle_mode = shuffle::Rotate(*shifts); + break; + case ShuffleMode::MODE_NOT_SET: + TokenError("expects a shuffle mode"); + return nullptr; + } + if (!maybe_infer_shape([&] { + return ShapeInference::InferShuffleShape(operands[0]->shape(), + *dimensions, shuffle_mode); + })) { + return nullptr; + } + return builder->AddInstruction(HloInstruction::CreateShuffle( + *shape, operands[0], *dimensions, shuffle_mode)); + } case HloOpcode::kSelectAndScatter: { optional select; attrs["select"] = {/*required=*/true, AttrTy::kHloComputation, &select}; @@ -6416,6 +6457,15 @@ bool HloParserImpl::ParseAttributeHelper( static_cast*>(attr_out_ptr)->emplace(result); return true; } + case AttrTy::kShuffleMode: { + ShuffleMode::ModeCase result; + if (!ParseShuffleMode(&result)) { + return false; + } + static_cast*>(attr_out_ptr) + ->emplace(result); + return true; + } case AttrTy::kPrecisionAlgorithm: { PrecisionConfig::Algorithm result; if (!ParseAlgorithm(&result)) { @@ -8560,6 +8610,22 @@ bool HloParserImpl::ParseRandomAlgorithm(RandomAlgorithm* result) { return true; } +bool HloParserImpl::ParseShuffleMode(ShuffleMode::ModeCase* result) { + VLOG(kDebugLevel) << "ParseShuffleMode"; + if (lexer_.GetKind() != TokKind::kIdent) { + return TokenError("expects shuffle mode"); + } + std::string val = lexer_.GetStrVal(); + auto status_or_result = StringToShuffleMode(val); + if (!status_or_result.ok()) { + return TokenError(StrFormat("expects shuffle mode but sees: %s, error: %s", + val, status_or_result.status().message())); + } + *result = status_or_result.value(); + lexer_.Lex(); + return true; +} + bool HloParserImpl::ParsePrecision(PrecisionConfig::Precision* result) { VLOG(kDebugLevel) << "ParsePrecision"; if (lexer_.GetKind() != TokKind::kIdent) { diff --git a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc index 9ff45bd6026930..6502d1f3d52d87 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc @@ -775,6 +775,18 @@ ENTRY %Reverse4DFloatArrayOnDim01.v2 () -> f32[4,3,2,1] { ROOT %reverse = f32[4,3,2,1]{0,1,2,3} reverse(f32[4,3,2,1]{0,1,2,3} %constant), dimensions={0,1} } +)" +}, +// shuffle(constant, mode=rotate) +{ +"ShuffleRotate2D", +R"(HloModule ShuffleRotate2DFloatArrayOnDim01_module, entry_computation_layout={()->f32[4,3]{1,0}} + +ENTRY %ShuffleRotate2DFloatArrayOnDim01.v2 () -> f32[4,3] { + %constant = f32[4,3]{1,0} constant({ { 1, 2, 3 }, { 4, 5, 6 }, { 7, 8, 9 }, { 10, 11, 12 } }) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %constant), dimensions={0,1}, mode=rotate, shifts={1,2} +} + )" }, // concat @@ -8720,5 +8732,194 @@ ENTRY %main (p0: f32[]) -> f32[] { "raw_computation_config"); EXPECT_EQ(roundtrip_module->ToString(), printed); } + +TEST_F(HloParserTest, ShuffleUnsupportedMode) { + const std::string original = R"(HloModule shuffle_unsupported_mode + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=transpose +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "expects shuffle mode but sees: transpose, error: Unknown " + "shuffle mode: transpose"); +} + +TEST_F(HloParserTest, ShuffleModeIsCaseSensitive) { + const std::string original = R"(HloModule shuffle_uppercase_mode + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=ROTATE, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "expects shuffle mode but sees: ROTATE, error: Unknown " + "shuffle mode: ROTATE"); +} + +TEST_F(HloParserTest, ShuffleModeIsNotAnIdentifier) { + const std::string original = R"(HloModule shuffle_non_identifier_mode + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=0, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), "expects shuffle mode"); +} + +TEST_F(HloParserTest, ShuffleMissingMode) { + const std::string original = R"(HloModule shuffle_missing_mode + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "attribute mode is expected but not seen"); +} + +TEST_F(HloParserTest, ShuffleMissingDimensions) { + const std::string original = R"(HloModule shuffle_missing_dimensions + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), mode=rotate, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "attribute dimensions is expected but not seen"); +} + +// The `shifts` attribute is specific to the rotate mode, so it must be present +// if and only if the mode is rotate. +TEST_F(HloParserTest, ShuffleRotateWithoutShifts) { + const std::string original = R"(HloModule shuffle_rotate_without_shifts + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=rotate +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "expects shifts for a shuffle in rotate mode"); +} + +TEST_F(HloParserTest, ShuffleUnknownAttribute) { + const std::string original = R"(HloModule shuffle_unknown_attribute + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=rotate, shifts={1,2}, strides={1,1} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "unexpected attribute \"strides\""); +} + +TEST_F(HloParserTest, ShuffleRepeatedAttribute) { + const std::string original = R"(HloModule shuffle_repeated_attribute + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=rotate, shifts={1,2}, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), "attribute shifts already exists"); +} + +TEST_F(HloParserTest, ShuffleWrongOperandsSize) { + const std::string original = R"(HloModule shuffle_wrong_operands_size + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p, f32[4,3]{1,0} %p), dimensions={0,1}, mode=rotate, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), "expects 1 operands"); +} + +// The tests below omit the shape of the shuffle, which makes the parser infer +// it and report the shape inference errors that the mode violates. +TEST_F(HloParserTest, ShuffleShiftsSizeMismatchesDimensionsSize) { + const std::string original = R"(HloModule shuffle_shifts_size_mismatch + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=rotate, shifts={1} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "failed to infer shape for opcode: shuffle, error: " + "dimensions and shifts must have the same size, got 2 and 1"); +} + +TEST_F(HloParserTest, ShuffleNoDimensions) { + const std::string original = R"(HloModule shuffle_no_dimensions + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = shuffle(f32[4,3]{1,0} %p), dimensions={}, mode=rotate, shifts={} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "A shuffle must shuffle at least one dimension."); +} + +TEST_F(HloParserTest, ShuffleDuplicateDimensions) { + const std::string original = R"(HloModule shuffle_duplicate_dimensions + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = shuffle(f32[4,3]{1,0} %p), dimensions={0,0}, mode=rotate, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "A dimension number is duplicated in shuffle."); +} + +TEST_F(HloParserTest, ShuffleDimensionOutOfBounds) { + const std::string original = R"(HloModule shuffle_dimension_out_of_bounds + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = shuffle(f32[4,3]{1,0} %p), dimensions={0,2}, mode=rotate, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr( + result.status().message(), + "One of the shuffle dimensions (2) is out-of-bounds in shape f32[4,3]"); +} + +TEST_F(HloParserTest, ShuffleOperandIsNotAnArray) { + const std::string original = R"(HloModule shuffle_non_array_operand + +ENTRY %entry (p: f32[4,3]) -> (f32[4,3]) { + %p = f32[4,3]{1,0} parameter(0) + %tuple = (f32[4,3]{1,0}) tuple(f32[4,3]{1,0} %p) + ROOT %shuffle = shuffle((f32[4,3]{1,0}) %tuple), dimensions={0}, mode=rotate, shifts={1} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "Expected array argument for operand of shuffle"); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_summary_test.cc b/third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_summary_test.cc index 9002ae7455f966..93aa4055b6f1bf 100644 --- a/third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_summary_test.cc +++ b/third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_summary_test.cc @@ -506,7 +506,7 @@ subtract.0 = f32[] subtract(constant.0, constant.1) /*left_unmatched_instruction_count=*/3, /*right_unmatched_instruction_count=*/0)), FieldsAre( - /*fingerprint=*/2080109884842586091U, + /*fingerprint=*/15551884552649592524U, /*computation_groups=*/ UnorderedElementsAre(FieldsAre( /*left_computations=*/IsEmpty(), diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD index cdf5463c660cee..d6b5175a155620 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD @@ -129,8 +129,6 @@ cc_library( "//xla/mlir_hlo", "//xla/mlir_hlo:unregistered_attributes", "//xla/service:hlo_proto_cc", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc index b180a3292b3748..9890ec597b9e52 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc @@ -85,8 +85,6 @@ limitations under the License. #include "xla/service/hlo.pb.h" #include "xla/shape_util.h" #include "xla/status_macros.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -1891,6 +1889,9 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( } return op.getOperation(); } + // TODO(b/565684884): Add conversion when Shuffle is supported in StableHLO. + case HloOpcode::kShuffle: + return InvalidArgument("Shuffle is not supported in StableHLO."); case HloOpcode::kRng: { auto shape = mlir::stablehlo::ConstantOp::create( *func_builder, loc, @@ -2141,8 +2142,8 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( &instruction->precision_config(), builder_))); // If the element types of the operands for convolution are different, - // insert a convert op to convert the operands to the common element type - // while preserving the values. + // insert a convert op to convert the operands to the common element + // type while preserving the values. auto lhs = operands[0]; auto rhs = operands[1]; auto lhs_element_type = instruction->operand(0)->shape().element_type(); @@ -2200,10 +2201,10 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( } case HloOpcode::kAdd: { - // HLO add ops on PRED elements are actually boolean or, but MHLO dialect - // AddOps on i1 are just addition with overflow; so, we have to implement - // the special behavior of HLO add ops on PRED here by creating an - // arith::OrIOp instead. + // HLO add ops on PRED elements are actually boolean or, but MHLO + // dialect AddOps on i1 are just addition with overflow; so, we have to + // implement the special behavior of HLO add ops on PRED here by + // creating an arith::OrIOp instead. if (instruction->shape().element_type() == PRED) { return mlir::stablehlo::OrOp::create(*func_builder, loc, result_type, operands, attributes) @@ -2227,8 +2228,8 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( } case HloOpcode::kConvert: { - // Convert to boolean is special, it requires a comparison to 0 instead of - // a truncation to i1, otherwise it is a 1-1 translation. + // Convert to boolean is special, it requires a comparison to 0 instead + // of a truncation to i1, otherwise it is a 1-1 translation. auto ranked_type = mlir::dyn_cast(result_type); mlir::IntegerType integer_type = (ranked_type) @@ -2279,10 +2280,10 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( "kind", mlir::mhlo::DomainKindAttr::get(func_builder->getContext(), *domain_kind))); - // In XLA, DomainMetadata is open-world, but in the proto, it is hardcoded - // to be ShardingMetadata. Thankfully, the only other implementation of - // DomainMetadata is OpName, which is generally used for debugging and - // never for compiling production models. + // In XLA, DomainMetadata is open-world, but in the proto, it is + // hardcoded to be ShardingMetadata. Thankfully, the only other + // implementation of DomainMetadata is OpName, which is generally used + // for debugging and never for compiling production models. // // Since this is hardcoded as such in the proto, we must follow suit. auto exit_metadata = ShardingMetadata::ToShardingMetadata( diff --git a/third_party/xla/xla/hlo/utils/hlo_matchers.h b/third_party/xla/xla/hlo/utils/hlo_matchers.h index c413dd80d45c7a..14b18c53d9bc14 100644 --- a/third_party/xla/xla/hlo/utils/hlo_matchers.h +++ b/third_party/xla/xla/hlo/utils/hlo_matchers.h @@ -387,6 +387,7 @@ HLO_MATCHER(SetDimensionSize); HLO_MATCHER(ShiftLeft); HLO_MATCHER(ShiftRightArithmetic); HLO_MATCHER(ShiftRightLogical); +HLO_MATCHER(Shuffle); HLO_MATCHER(Sign); HLO_MATCHER(Sinh); HLO_MATCHER(Slice); diff --git a/third_party/xla/xla/python/_hlo.pyi b/third_party/xla/xla/python/_hlo.pyi index f208055c5ae3ba..5b1505f46a492e 100644 --- a/third_party/xla/xla/python/_hlo.pyi +++ b/third_party/xla/xla/python/_hlo.pyi @@ -321,37 +321,39 @@ class HloOpcode(enum.Enum): kShiftRightLogical = 117 - kSign = 118 + kShuffle = 118 - kSin = 119 + kSign = 119 - kSinh = 120 + kSin = 120 - kSlice = 121 + kSinh = 121 - kSort = 122 + kSlice = 122 - kSqrt = 123 + kSort = 123 - kStochasticConvert = 124 + kSqrt = 124 - kSubtract = 125 + kStochasticConvert = 125 - kTan = 126 + kSubtract = 126 - kTanh = 127 + kTan = 127 - kTopK = 128 + kTanh = 128 - kTranspose = 129 + kTopK = 129 - kTriangularSolve = 130 + kTranspose = 130 - kTuple = 131 + kTriangularSolve = 131 - kWhile = 132 + kTuple = 132 - kXor = 133 + kWhile = 133 + + kXor = 134 class Layout: @overload diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 39fcaf0c8e5eea..45d7e446bcae49 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -551,6 +551,7 @@ xla_cc_test( ":hlo_proto_cc", ":shape_inference", "//xla:shape_util", + "//xla:shuffle", "//xla:xla_data_proto_cc", "//xla/hlo/builder:padding", "//xla/hlo/ir:hlo", @@ -561,6 +562,7 @@ xla_cc_test( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", @@ -818,6 +820,7 @@ xla_cc_test( "//xla:comparison_util", "//xla:literal_util", "//xla:shape_util", + "//xla:shuffle", "//xla:util", "//xla:window_util", "//xla:xla_data_proto_cc", diff --git a/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc b/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc index d254d0ea4e144f..c34b0af2b71fd8 100644 --- a/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc +++ b/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc @@ -221,6 +221,7 @@ bool CpuInstructionFusion::IsExpensive(const HloInstruction& instruction) { case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: case HloOpcode::kShiftRightLogical: + case HloOpcode::kShuffle: case HloOpcode::kSlice: case HloOpcode::kStochasticConvert: case HloOpcode::kSubtract: diff --git a/third_party/xla/xla/service/float_support.cc b/third_party/xla/xla/service/float_support.cc index 59a6be3ea620f7..c1854a183e49fe 100644 --- a/third_party/xla/xla/service/float_support.cc +++ b/third_party/xla/xla/service/float_support.cc @@ -133,6 +133,7 @@ bool FloatSupport::EffectiveOperandPrecisionIsOutputPrecision( case HloOpcode::kRaggedAllToAll: case HloOpcode::kReshape: case HloOpcode::kReverse: + case HloOpcode::kShuffle: case HloOpcode::kSlice: case HloOpcode::kSort: case HloOpcode::kTranspose: diff --git a/third_party/xla/xla/service/hlo.proto b/third_party/xla/xla/service/hlo.proto index b3646dd5b84100..ea708ce2b8dc4c 100644 --- a/third_party/xla/xla/service/hlo.proto +++ b/third_party/xla/xla/service/hlo.proto @@ -128,7 +128,7 @@ enum CustomCallApiVersion { } // Serialization of HloInstruction. -// Next ID: 103 +// Next ID: 104 message HloInstructionProto { reserved 10; reserved "parameter_name"; @@ -442,6 +442,9 @@ message HloInstructionProto { // Comparison order (TOTAL, PARTIAL). string comparison_order = 102; + + // The mode of a kShuffle and the attributes of that mode. + xla.ShuffleMode shuffle_mode = 103; } // Serialization of HloComputation. diff --git a/third_party/xla/xla/service/hlo_cost_analysis.cc b/third_party/xla/xla/service/hlo_cost_analysis.cc index 11e310dfde8836..c646a5eae4a9a7 100644 --- a/third_party/xla/xla/service/hlo_cost_analysis.cc +++ b/third_party/xla/xla/service/hlo_cost_analysis.cc @@ -380,6 +380,10 @@ absl::Status HloCostAnalysis::HandleReverse(const HloInstruction*) { return absl::OkStatus(); } +absl::Status HloCostAnalysis::HandleShuffle(const HloInstruction*) { + return absl::OkStatus(); +} + absl::Status HloCostAnalysis::HandleSlice(const HloInstruction* slice) { const int64_t output_shape_size = GetShapeSize(slice->shape()); diff --git a/third_party/xla/xla/service/hlo_cost_analysis.h b/third_party/xla/xla/service/hlo_cost_analysis.h index cd97ea8832cc3a..ca809d4b6779ed 100644 --- a/third_party/xla/xla/service/hlo_cost_analysis.h +++ b/third_party/xla/xla/service/hlo_cost_analysis.h @@ -515,6 +515,7 @@ class HloCostAnalysis : public ConstDfsHloVisitor { absl::Status HandleRngGetAndUpdateState( const HloInstruction* random) override; absl::Status HandleReverse(const HloInstruction* reverse) override; + absl::Status HandleShuffle(const HloInstruction* shuffle) override; absl::Status HandleSort(const HloInstruction* sort) override; absl::Status HandleParameter(const HloInstruction* parameter) override; absl::Status HandleReduce(const HloInstruction* reduce) override; diff --git a/third_party/xla/xla/service/hlo_graph_dumper.cc b/third_party/xla/xla/service/hlo_graph_dumper.cc index ab7e51f833f368..55675bb8934fe6 100644 --- a/third_party/xla/xla/service/hlo_graph_dumper.cc +++ b/third_party/xla/xla/service/hlo_graph_dumper.cc @@ -1245,6 +1245,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kReshape: case HloOpcode::kDynamicReshape: case HloOpcode::kReverse: + case HloOpcode::kShuffle: case HloOpcode::kTranspose: // These data-movement ops can be expensive; emphasize them. (Yes, even // concat can be expensive, at least on GPU, as it can create warp diff --git a/third_party/xla/xla/service/hlo_instruction_test.cc b/third_party/xla/xla/service/hlo_instruction_test.cc index 9d343dfc322867..8480299330cd5f 100644 --- a/third_party/xla/xla/service/hlo_instruction_test.cc +++ b/third_party/xla/xla/service/hlo_instruction_test.cc @@ -52,6 +52,7 @@ limitations under the License. #include "xla/service/pattern_matcher.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/shuffle.h" #include "xla/tsl/util/proto/proto_matchers.h" #include "xla/util.h" #include "xla/window_util.h" @@ -1216,6 +1217,20 @@ TEST_F(HloInstructionTest, IdenticalInstructions) { Identical(*HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op1), *HloInstruction::CreateUnary(shape, HloOpcode::kNegate, op1))); + // Shuffle. + EXPECT_TRUE(Identical(*HloInstruction::CreateShuffle(shape, op1, {0, 1}, + shuffle::Rotate({2, 3})), + *HloInstruction::CreateShuffle( + shape, op1, {0, 1}, shuffle::Rotate({2, 3})))); + EXPECT_FALSE(Identical(*HloInstruction::CreateShuffle( + shape, op1, {0, 1}, shuffle::Rotate({2, 3})), + *HloInstruction::CreateShuffle( + shape, op1, {0, 1}, shuffle::Rotate({3, 2})))); + EXPECT_FALSE(Identical(*HloInstruction::CreateShuffle( + shape, op1, {0, 1}, shuffle::Rotate({2, 3})), + *HloInstruction::CreateShuffle( + shape, op1, {1, 0}, shuffle::Rotate({2, 3})))); + // Tuples. EXPECT_TRUE(Identical(*HloInstruction::CreateTuple({op1, op2}), *HloInstruction::CreateTuple({op1, op2}))); diff --git a/third_party/xla/xla/service/hlo_verifier.cc b/third_party/xla/xla/service/hlo_verifier.cc index 431308d01822f7..0e532bfea16442 100644 --- a/third_party/xla/xla/service/hlo_verifier.cc +++ b/third_party/xla/xla/service/hlo_verifier.cc @@ -1190,6 +1190,14 @@ absl::Status ShapeVerifier::HandleReverse(HloInstruction* reverse) { reverse->dimensions())); } +absl::Status ShapeVerifier::HandleShuffle(HloInstruction* shuffle) { + HloShuffleInstruction* shuffle_instr = Cast(shuffle); + return CheckShape( + shuffle, ShapeInference::InferShuffleShape( + shuffle_instr->operand(0)->shape(), + shuffle_instr->dimensions(), shuffle_instr->shuffle_mode())); +} + absl::Status ShapeVerifier::HandleTopK(HloInstruction* hlo) { return CheckShape( hlo, ShapeInference::InferTopKShape(hlo->operand(0)->shape(), diff --git a/third_party/xla/xla/service/hlo_verifier.h b/third_party/xla/xla/service/hlo_verifier.h index 61f0d3e9ca8cf8..9af83c220ea9d3 100644 --- a/third_party/xla/xla/service/hlo_verifier.h +++ b/third_party/xla/xla/service/hlo_verifier.h @@ -253,6 +253,7 @@ class ShapeVerifier : public DfsHloVisitor { absl::Status HandleRngBitGenerator(HloInstruction*) override; absl::Status HandleRngGetAndUpdateState(HloInstruction*) override; absl::Status HandleReverse(HloInstruction* reverse) override; + absl::Status HandleShuffle(HloInstruction* shuffle) override; absl::Status HandleSort(HloInstruction* hlo) override; absl::Status HandleTopK(HloInstruction* hlo) override; absl::Status HandleConstant(HloInstruction* constant) override; diff --git a/third_party/xla/xla/service/hlo_verifier_test.cc b/third_party/xla/xla/service/hlo_verifier_test.cc index 60edc3cae66153..b459ed7ad29e8f 100644 --- a/third_party/xla/xla/service/hlo_verifier_test.cc +++ b/third_party/xla/xla/service/hlo_verifier_test.cc @@ -4403,6 +4403,76 @@ ENTRY TopK { ASSERT_TRUE(status.ok()); } +TEST_F(HloVerifierTest, ShuffleRotateOK) { + constexpr absl::string_view kHlo = R"( +HloModule shuffle, entry_computation_layout={(f32[10,20]{1,0})->f32[10,20]{1,0}} + +ENTRY Shuffle { + x = f32[10,20]{1,0} parameter(0) + ROOT shuffle = f32[10,20]{1,0} shuffle(x), dimensions={0,1}, mode=rotate, shifts={2,5} +} + )"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(kHlo)); + const absl::StatusOr result = verifier().Run(module.get()); + + EXPECT_THAT(result, absl_testing::IsOk()); +} + +TEST_F(HloVerifierTest, ShuffleInvalidDimensions) { + constexpr absl::string_view kHlo = R"( +HloModule shuffle, entry_computation_layout={(f32[10,20]{1,0})->f32[10,20]{1,0}} + +ENTRY Shuffle { + x = f32[10,20]{1,0} parameter(0) + ROOT shuffle = f32[10,20]{1,0} shuffle(x), dimensions={0,2}, mode=rotate, shifts={2,5} +} + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(kHlo)); + const absl::StatusOr result = verifier().Run(module.get()); + + EXPECT_THAT(result, StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("out-of-bounds"))); +} + +TEST_F(HloVerifierTest, ShuffleRotateMismatchedShiftsSize) { + constexpr absl::string_view kHlo = R"( +HloModule shuffle, entry_computation_layout={(f32[10,20]{1,0})->f32[10,20]{1,0}} + +ENTRY Shuffle { + x = f32[10,20]{1,0} parameter(0) + ROOT shuffle = f32[10,20]{1,0} shuffle(x), dimensions={0,1}, mode=rotate, shifts={2} +} + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(kHlo)); + const absl::StatusOr result = verifier().Run(module.get()); + + EXPECT_THAT(result, + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("dimensions and shifts must have the same " + "size"))); +} + +TEST_F(HloVerifierTest, ShuffleDuplicatedDimensions) { + constexpr absl::string_view kHlo = R"( +HloModule shuffle, entry_computation_layout={(f32[10,20]{1,0})->f32[10,20]{1,0}} + +ENTRY Shuffle { + x = f32[10,20]{1,0} parameter(0) + ROOT shuffle = f32[10,20]{1,0} shuffle(x), dimensions={0,0}, mode=rotate, shifts={2,5} +} + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(kHlo)); + const absl::StatusOr result = verifier().Run(module.get()); + + EXPECT_THAT(result, StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("duplicated"))); +} + TEST_F(HloVerifierTest, InputLayoutMismatchIgnored) { // Note: The mismatch is between the entry_computation_layout and the layout // of parameter(1). diff --git a/third_party/xla/xla/service/instruction_fusion.cc b/third_party/xla/xla/service/instruction_fusion.cc index 7c900b9629562c..b32a0ee94944f0 100644 --- a/third_party/xla/xla/service/instruction_fusion.cc +++ b/third_party/xla/xla/service/instruction_fusion.cc @@ -132,6 +132,7 @@ bool IsAlwaysDuplicable(const HloInstruction& instruction) { case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: case HloOpcode::kShiftRightLogical: + case HloOpcode::kShuffle: case HloOpcode::kSlice: case HloOpcode::kStochasticConvert: case HloOpcode::kSubtract: diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc index 425e2f46e154bd..88fa67ba987317 100644 --- a/third_party/xla/xla/service/layout_assignment.cc +++ b/third_party/xla/xla/service/layout_assignment.cc @@ -3820,6 +3820,7 @@ bool LayoutAssignment::InstructionCanChangeLayout( case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: case HloOpcode::kShiftRightLogical: + case HloOpcode::kShuffle: case HloOpcode::kSign: case HloOpcode::kSin: case HloOpcode::kSinh: diff --git a/third_party/xla/xla/service/shape_inference.cc b/third_party/xla/xla/service/shape_inference.cc index 86ed2b98e42859..ed0006136a6cd3 100644 --- a/third_party/xla/xla/service/shape_inference.cc +++ b/third_party/xla/xla/service/shape_inference.cc @@ -3706,6 +3706,39 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { return operand_shape; } +/*static */ absl::StatusOr ShapeInference::InferShuffleShape( + const Shape& operand_shape, absl::Span dimensions, + const ShuffleMode& mode) { + ABSL_RETURN_IF_ERROR(ExpectArray(operand_shape, "operand of shuffle")); + if (dimensions.empty()) { + return InvalidArgument("A shuffle must shuffle at least one dimension."); + } + if (!AllUnique(dimensions)) { + return InvalidArgument("A dimension number is duplicated in shuffle."); + } + for (int64_t dimension : dimensions) { + if (dimension < 0 || dimension >= operand_shape.dimensions().size()) { + return InvalidArgument( + "One of the shuffle dimensions (%d) is out-of-bounds in shape %s.", + dimension, ShapeUtil::HumanString(operand_shape)); + } + } + // Constraints on the attributes that are specific to the shuffle mode. + switch (mode.mode_case()) { + case ShuffleMode::kRotate: + if (dimensions.size() != mode.rotate().shifts().size()) { + return InvalidArgument( + "dimensions and shifts must have the same size, got %d and %d.", + dimensions.size(), mode.rotate().shifts().size()); + } + break; + case ShuffleMode::MODE_NOT_SET: + return InvalidArgument("A shuffle must specify a mode."); + } + // A shuffle only moves elements around, so the shape is preserved. + return operand_shape; +} + /* static */ absl::StatusOr ShapeInference::InferGetTupleElementShape( const Shape& arg, int64_t index) { if (!arg.IsTuple()) { diff --git a/third_party/xla/xla/service/shape_inference.h b/third_party/xla/xla/service/shape_inference.h index 15e71e840a09c2..c5566066296781 100644 --- a/third_party/xla/xla/service/shape_inference.h +++ b/third_party/xla/xla/service/shape_inference.h @@ -248,6 +248,12 @@ class ShapeInference { static absl::StatusOr InferReverseShape( const Shape& operand_shape, absl::Span dimensions); + // Infers the shape of a shuffle instruction, which rearranges the operand's + // elements along the given dimensions according to `mode`. + static absl::StatusOr InferShuffleShape( + const Shape& operand_shape, absl::Span dimensions, + const ShuffleMode& mode); + // Infers the shape produced by a slice operation spanning from the starts to // the limits in the original shape's dimensions. // diff --git a/third_party/xla/xla/service/shape_inference_test.cc b/third_party/xla/xla/service/shape_inference_test.cc index 0db74cc0433cf8..c8c75fa2673d26 100644 --- a/third_party/xla/xla/service/shape_inference_test.cc +++ b/third_party/xla/xla/service/shape_inference_test.cc @@ -28,6 +28,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/log/check.h" #include "absl/log/log.h" +#include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/strings/str_join.h" #include "absl/strings/string_view.h" @@ -42,6 +43,7 @@ limitations under the License. #include "xla/service/hlo.pb.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/shuffle.h" #include "xla/xla_data.pb.h" #include "tsl/platform/errors.h" @@ -3073,6 +3075,57 @@ TEST_F(ShapeInferenceTest, ReverseInvalidDimension) { HasSubstr("Expected array argument")); } +TEST_F(ShapeInferenceTest, ShuffleRotate) { + const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25}); + + ASSERT_THAT(ShapeInference::InferShuffleShape(input_shape, {0, 1}, + shuffle::Rotate({2, 5})), + absl_testing::IsOkAndHolds(input_shape)); +} + +TEST_F(ShapeInferenceTest, ShuffleInvalidDimension) { + const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25}); + + ASSERT_THAT(ShapeInference::InferShuffleShape(input_shape, {0, 2}, + shuffle::Rotate({2, 5})), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("out-of-bounds"))); + + ASSERT_THAT(ShapeInference::InferShuffleShape(input_shape, {0, -1}, + shuffle::Rotate({2, 5})), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("out-of-bounds"))); + + ASSERT_THAT(ShapeInference::InferShuffleShape(input_shape, {0, 0}, + shuffle::Rotate({2, 5})), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("duplicated"))); + + ASSERT_THAT(ShapeInference::InferShuffleShape(input_shape, {0, 1}, + shuffle::Rotate({2})), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("dimensions and shifts must " + "have the same size"))); +} + +TEST_F(ShapeInferenceTest, ShuffleWithoutMode) { + const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25}); + + ASSERT_THAT( + ShapeInference::InferShuffleShape(input_shape, {0, 1}, ShuffleMode()), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("must specify a mode"))); +} + +TEST_F(ShapeInferenceTest, ShuffleWithoutDimensions) { + const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25}); + + ASSERT_THAT( + ShapeInference::InferShuffleShape(input_shape, {}, shuffle::Rotate({})), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("must shuffle at least one dimension"))); +} + TEST_F(ShapeInferenceTest, Call) { const absl::StatusOr inferred_shape0 = ShapeInference::InferCallShape({}, ShapeUtil::MakeProgramShape({}, f32_)); diff --git a/third_party/xla/xla/service/sharding_propagation.cc b/third_party/xla/xla/service/sharding_propagation.cc index 26e6d65e72ae87..e4a1749f858baa 100644 --- a/third_party/xla/xla/service/sharding_propagation.cc +++ b/third_party/xla/xla/service/sharding_propagation.cc @@ -207,6 +207,7 @@ const HloInstruction* PickRepresentativeOperand( case HloOpcode::kOptimizationBarrier: case HloOpcode::kRaggedAllToAll: case HloOpcode::kReverse: + case HloOpcode::kShuffle: case HloOpcode::kSlice: case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: diff --git a/third_party/xla/xla/shuffle.cc b/third_party/xla/xla/shuffle.cc new file mode 100644 index 00000000000000..49c0e573a6f3ee --- /dev/null +++ b/third_party/xla/xla/shuffle.cc @@ -0,0 +1,39 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/shuffle.h" + +#include + +#include "absl/log/check.h" +#include "absl/types/span.h" +#include "xla/xla_data.pb.h" + +namespace xla { +namespace shuffle { + +ShuffleMode Rotate(absl::Span shifts) { + ShuffleMode mode; + mode.mutable_rotate()->mutable_shifts()->Assign(shifts.begin(), shifts.end()); + return mode; +} + +int64_t NormalizeShift(int64_t shift, int64_t dim_size) { + CHECK_GE(dim_size, 0) << "a dimension cannot have a negative size"; + return dim_size == 0 ? 0 : ((shift % dim_size) + dim_size) % dim_size; +} + +} // namespace shuffle +} // namespace xla diff --git a/third_party/xla/xla/shuffle.h b/third_party/xla/xla/shuffle.h new file mode 100644 index 00000000000000..87bb9b48c92ef5 --- /dev/null +++ b/third_party/xla/xla/shuffle.h @@ -0,0 +1,41 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_SHUFFLE_H_ +#define XLA_SHUFFLE_H_ + +#include + +#include "absl/types/span.h" +#include "xla/xla_data.pb.h" + +namespace xla { +namespace shuffle { + +// The mode of a shuffle is the pattern in which it moves the elements of its +// shuffled dimensions, and the attributes that parameterize the pattern. + +// Returns the mode of a shuffle that rotates each shuffled dimension to the +// left by the corresponding entry of `shifts`. +ShuffleMode Rotate(absl::Span shifts); + +// Returns the normalized shift in `[0, dim_size)`. Returns 0 if `dim_size` is +// 0, because every rotation of an empty dimension is a no-op. +int64_t NormalizeShift(int64_t shift, int64_t dim_size); + +} // namespace shuffle +} // namespace xla + +#endif // XLA_SHUFFLE_H_ diff --git a/third_party/xla/xla/shuffle_test.cc b/third_party/xla/xla/shuffle_test.cc new file mode 100644 index 00000000000000..151b30d1fa336b --- /dev/null +++ b/third_party/xla/xla/shuffle_test.cc @@ -0,0 +1,74 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/shuffle.h" + +#include +#include +#include "xla/xla_data.pb.h" + +namespace xla { +namespace shuffle { +namespace { + +using ::testing::ElementsAre; +using ::testing::IsEmpty; + +TEST(ShuffleTest, RotateSetsTheShifts) { + const ShuffleMode mode = Rotate({2, 3}); + + EXPECT_EQ(mode.mode_case(), ShuffleMode::kRotate); + EXPECT_THAT(mode.rotate().shifts(), ElementsAre(2, 3)); +} + +TEST(ShuffleTest, RotateWithoutShiftsIsStillInRotateMode) { + const ShuffleMode mode = Rotate({}); + + EXPECT_EQ(mode.mode_case(), ShuffleMode::kRotate); + EXPECT_THAT(mode.rotate().shifts(), IsEmpty()); +} + +TEST(ShuffleTest, NormalizeShiftKeepsShiftsThatAreAlreadyInRange) { + EXPECT_EQ(NormalizeShift(0, 4), 0); + EXPECT_EQ(NormalizeShift(1, 4), 1); + EXPECT_EQ(NormalizeShift(3, 4), 3); +} + +TEST(ShuffleTest, NormalizeShiftWrapsShiftsAround) { + EXPECT_EQ(NormalizeShift(4, 4), 0); + EXPECT_EQ(NormalizeShift(5, 4), 1); + EXPECT_EQ(NormalizeShift(9, 4), 1); +} + +TEST(ShuffleTest, NormalizeShiftWrapsNegativeShiftsAround) { + EXPECT_EQ(NormalizeShift(-1, 4), 3); + EXPECT_EQ(NormalizeShift(-4, 4), 0); + EXPECT_EQ(NormalizeShift(-5, 4), 3); +} + +TEST(ShuffleTest, NormalizeShiftOfADimensionOfSizeOneIsANoOp) { + EXPECT_EQ(NormalizeShift(3, 1), 0); + EXPECT_EQ(NormalizeShift(-3, 1), 0); +} + +TEST(ShuffleTest, NormalizeShiftOfAnEmptyDimensionIsANoOp) { + EXPECT_EQ(NormalizeShift(0, 0), 0); + EXPECT_EQ(NormalizeShift(3, 0), 0); + EXPECT_EQ(NormalizeShift(-3, 0), 0); +} + +} // namespace +} // namespace shuffle +} // namespace xla diff --git a/third_party/xla/xla/xla_data.proto b/third_party/xla/xla/xla_data.proto index a8465a05da565d..9cd793c6e2a362 100644 --- a/third_party/xla/xla/xla_data.proto +++ b/third_party/xla/xla/xla_data.proto @@ -965,6 +965,35 @@ enum RandomAlgorithm { // Next: 2 } +// The mode of a `shuffle` instruction together with the attributes that +// parameterize that mode. A shuffle moves the elements of its operand around +// along its shuffled dimensions, and the mode is the pattern of that movement. +// Every mode owns a message holding exactly the attributes that the mode +// needs, and the active case of the `mode` oneof selects the mode; a mode +// without attributes owns an empty message. +// +// The name of a case of the `mode` oneof is also the name of the mode in the +// HLO text format, e.g. `mode=rotate`. +message ShuffleMode { + // Rotates the elements along each shuffled dimension to the left, wrapping + // the elements that fall off the front around to the back, e.g. rotating + // `[1, 2, 3, 4]` by a shift of 1 yields `[2, 3, 4, 1]`. + message Rotate { + // The shift of each rotated dimension, in the order in which the + // dimensions appear in `dimensions`. A shift of `s` moves the + // element at index `i` of a dimension of size `d` to index + // `(i - s) mod d`, so `s` is the index of the element that ends up first. + // Shifts outside of `[0, d)`, including negative ones, wrap around. + repeated int64 shifts = 1; + } + + oneof mode { + Rotate rotate = 1; + } + + // Next: 2 +} + message TriangularSolveOptions { // If true, solves ax = b. If false, solves xa = b. bool left_side = 1; From 26e6f35a311249108322861d92d9be6b90c1e1f9 Mon Sep 17 00:00:00 2001 From: Yulia Baturina Date: Thu, 24 Sep 2026 11:28:41 -0700 Subject: [PATCH 44/48] Update shardy commit. PiperOrigin-RevId: 987639736 --- .../xla/third_party/shardy/temporary.patch | 15 --------------- third_party/xla/third_party/shardy/workspace.bzl | 4 ++-- 2 files changed, 2 insertions(+), 17 deletions(-) diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch index 38c78a67e426b3..95a2a325ed502a 100644 --- a/third_party/xla/third_party/shardy/temporary.patch +++ b/third_party/xla/third_party/shardy/temporary.patch @@ -192,18 +192,3 @@ index 509398da..a9462de1 100644 ++ %fadd = fadd double %fmul26, %fmul29 ++ ret double %fadd ++} -diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl -index 08380e9e..d929da09 100644 ---- a/third_party/llvm/workspace.bzl -+++ b/third_party/llvm/workspace.bzl -@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive") - - def repo(name): - """Imports LLVM.""" -- LLVM_COMMIT = "0038be81b8b5361112faf82e69d4590364ee81b9" -- LLVM_SHA256 = "6675de5b328d62a6bbbef606720ab110d32c55e4645d4fa503e1cfab15bfe878" -+ LLVM_COMMIT = "a22291abb68744cb0b70978eb506a7c5592a3309" -+ LLVM_SHA256 = "cf19b82f0dfbef1e62afe520ba4ee6de4f4823524d7804490e88d815f822c8e9" - - tf_http_archive( - name = name, diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl index ce3f7fdc207531..dd328941030d9f 100644 --- a/third_party/xla/third_party/shardy/workspace.bzl +++ b/third_party/xla/third_party/shardy/workspace.bzl @@ -18,8 +18,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def repo(): - SHARDY_COMMIT = "7105c6a0e1d29141406a120b111e5525ff7b2e64" - SHARDY_SHA256 = "5f307f4e3acd7f7b7780bb9c7bf92a6cfe4108067f3976d80f43bce9c89c8f5a" + SHARDY_COMMIT = "08e079b1c5f5b0a3a0c696abc5c56b6ac436b835" + SHARDY_SHA256 = "b3de5417edb836ddcda6c1800bb0a5782c58110ff9567675a82ada9562d1033b" tf_http_archive( name = "shardy", From e3f6b576262bd9001928f1c5b2effeb55426ddee Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Thu, 24 Sep 2026 11:31:10 -0700 Subject: [PATCH 45/48] Replace IFRT mock usage with real client in //third_party/tensorflow/core/tfrt/ifrt:ifrt_device_utils_test PiperOrigin-RevId: 987641321 --- tensorflow/core/tfrt/ifrt/BUILD | 47 +++++--- .../core/tfrt/ifrt/ifrt_device_utils_test.cc | 110 +++++++++--------- 2 files changed, 81 insertions(+), 76 deletions(-) diff --git a/tensorflow/core/tfrt/ifrt/BUILD b/tensorflow/core/tfrt/ifrt/BUILD index 2dbf203aef7204..ead62b4e02a110 100644 --- a/tensorflow/core/tfrt/ifrt/BUILD +++ b/tensorflow/core/tfrt/ifrt/BUILD @@ -103,25 +103,34 @@ cc_library( ], ) -tf_cc_test( - name = "ifrt_device_utils_test", - srcs = [ - "ifrt_device_utils_test.cc", - ], - tags = ["no_oss"], - deps = [ - ":ifrt_device_utils", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_matchers", - "@com_google_googletest//:gtest_main", - "@xla//xla/python/ifrt", - "@xla//xla/python/ifrt:attribute_map", - "@xla//xla/python/ifrt:mock", - "@xla//xla/service:device_assignment", - "@xla//xla/tsl/platform:status_matchers", - "@xla//xla/tsl/platform:statusor", - ], -) +# copybara:uncomment_begin(PjRt-IFRT TPU client is not available in OSS) +# # TODO(siqiaowu): ifrt_device_utils_test couples with TPU assumptions ("coords" +# # and "core_on_chip" device attributes) and requires a TPU test client, which is +# # not available in OSS. +# tf_cc_test( +# name = "ifrt_device_utils_test", +# srcs = [ +# "ifrt_device_utils_test.cc", +# ], +# tags = [ +# "no_oss", +# "requires-tpu", +# ], +# deps = [ +# ":ifrt_device_utils", +# "//learning/brain/research/ifrt/pjrt_ifrt/pjrt_tpu:pjrt_tpu_client_test_lib", +# "@com_google_googletest//:gtest_main", +# "@com_google_absl//absl/status", +# "@com_google_absl//absl/status:status_matchers", +# "@xla//xla/python/ifrt", +# "@xla//xla/python/ifrt:attribute_map", +# "@xla//xla/python/ifrt:test_util", +# "@xla//xla/service:device_assignment", +# "@xla//xla/tsl/platform:status_matchers", +# "@xla//xla/tsl/platform:statusor", +# ], +# ) +# copybara:uncomment_end cc_library( name = "ifrt_serving_executable", diff --git a/tensorflow/core/tfrt/ifrt/ifrt_device_utils_test.cc b/tensorflow/core/tfrt/ifrt/ifrt_device_utils_test.cc index 49f62d74cdece6..871af37f23f540 100644 --- a/tensorflow/core/tfrt/ifrt/ifrt_device_utils_test.cc +++ b/tensorflow/core/tfrt/ifrt/ifrt_device_utils_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/core/tfrt/ifrt/ifrt_device_utils.h" +#include #include #include #include @@ -24,8 +25,9 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "xla/python/ifrt/attribute_map.h" +#include "xla/python/ifrt/client.h" #include "xla/python/ifrt/device.h" -#include "xla/python/ifrt/mock.h" +#include "xla/python/ifrt/test_util.h" #include "xla/service/device_assignment.h" #include "xla/tsl/platform/status_matchers.h" #include "xla/tsl/platform/statusor.h" @@ -34,99 +36,93 @@ namespace tensorflow { namespace ifrt_serving { namespace { using ::testing::ElementsAre; -using ::testing::Return; -using ::testing::ReturnRef; static constexpr int kNumReplicas = 1; static constexpr int kNumCoresPerReplica = 2; // Intentionally have more devices than kNumReplicas * kNumCoresPerReplica for // testing purposes. -static constexpr int kNumDevices = 4; -static constexpr int kDeviceIdOffset = 8; +static constexpr int kNumDevices = 2; class IfrtDeviceUtilsTest : public ::testing::Test { protected: void SetUp() override { - mocked_devices_.reserve(kNumDevices); - devices_.reserve(kNumDevices); - for (int i = 0; i < kNumDevices; ++i) { - mocked_devices_.push_back(std::make_unique()); - ON_CALL(*mocked_devices_[i], Attributes()) - .WillByDefault(ReturnRef(device_attributes_maps_[i])); - ON_CALL(*mocked_devices_[i], Id()) - .WillByDefault(Return(xla::ifrt::DeviceId(kDeviceIdOffset + i))); - ON_CALL(client_, LookupDevice(xla::ifrt::DeviceId(kDeviceIdOffset + i))) - .WillByDefault(Return(mocked_devices_[i].get())); - - devices_.push_back(mocked_devices_[i].get()); - }; - - ON_CALL(client_, devices()).WillByDefault(Return(devices_)); - - // Default use the last two devices. - xla::DeviceAssignment assignment(kNumReplicas, kNumCoresPerReplica); - assignment(0, 0) = kDeviceIdOffset + 2; - assignment(0, 1) = kDeviceIdOffset + 3; - - ON_CALL(client_, - GetDefaultDeviceAssignment(kNumReplicas, kNumCoresPerReplica)) - .WillByDefault(Return(assignment)); + ASSERT_OK_AND_ASSIGN(client_, xla::ifrt::test_util::GetClient()); + for (auto* device : client_->devices()) { + ASSERT_OK_AND_ASSIGN( + auto coords, + device->Attributes().Get>("coords")); + ASSERT_OK_AND_ASSIGN(auto core, + device->Attributes().Get("core_on_chip")); + ASSERT_EQ(coords.size(), 3); + devices_.push_back(device); + device_coords_.push_back( + {static_cast(coords[0]), static_cast(coords[1]), + static_cast(coords[2]), static_cast(core)}); + } + ASSERT_GE(devices_.size(), kNumDevices); } - xla::ifrt::MockClient client_; - std::vector> mocked_devices_; - + std::shared_ptr client_; std::vector devices_; - std::vector device_attributes_maps_ = { - xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{ - {"coords", xla::ifrt::AttributeMap::Int64ListValue({1, 0, 0})}, - {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(0)}}), - xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{ - {"coords", xla::ifrt::AttributeMap::Int64ListValue({1, 0, 0})}, - {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(1)}}), - xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{ - {"coords", xla::ifrt::AttributeMap::Int64ListValue({2, 0, 0})}, - {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(0)}}), - xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{ - {"coords", xla::ifrt::AttributeMap::Int64ListValue({2, 0, 0})}, - {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(1)}}), - }; + std::vector> device_coords_; }; TEST_F(IfrtDeviceUtilsTest, Basic) { - std::vector device_assignment_attr = {1, 0, 0, 1, 1, 0, 0, 0}; + std::vector device_assignment_attr; + device_assignment_attr.insert(device_assignment_attr.end(), + device_coords_[1].begin(), + device_coords_[1].end()); + device_assignment_attr.insert(device_assignment_attr.end(), + device_coords_[0].begin(), + device_coords_[0].end()); TF_ASSERT_OK_AND_ASSIGN( auto devices_from_attribute, - GetAssignedIfrtDevices(client_, kNumReplicas, kNumCoresPerReplica, + GetAssignedIfrtDevices(*client_, kNumReplicas, kNumCoresPerReplica, device_assignment_attr)); EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[1], devices_[0])); } -TEST_F(IfrtDeviceUtilsTest, SeparateXCoordinates) { - std::vector device_assignment_attr = {1, 0, 0, 1, 2, 0, 0, 0}; +TEST_F(IfrtDeviceUtilsTest, InvertCoordinates) { + std::vector device_assignment_attr; + device_assignment_attr.insert(device_assignment_attr.end(), + device_coords_[0].begin(), + device_coords_[0].end()); + device_assignment_attr.insert(device_assignment_attr.end(), + device_coords_[1].begin(), + device_coords_[1].end()); TF_ASSERT_OK_AND_ASSIGN( auto devices_from_attribute, - GetAssignedIfrtDevices(client_, kNumReplicas, kNumCoresPerReplica, + GetAssignedIfrtDevices(*client_, kNumReplicas, kNumCoresPerReplica, device_assignment_attr)); - EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[1], devices_[2])); + EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[0], devices_[1])); } TEST_F(IfrtDeviceUtilsTest, EmptyDeviceAssignmentShallReturnDefault) { + TF_ASSERT_OK_AND_ASSIGN( + xla::DeviceAssignment default_assignment, + client_->GetDefaultDeviceAssignment(kNumReplicas, kNumCoresPerReplica)); + TF_ASSERT_OK_AND_ASSIGN( + xla::ifrt::Device * dev0, + client_->LookupDevice(xla::ifrt::DeviceId(default_assignment(0, 0)))); + TF_ASSERT_OK_AND_ASSIGN( + xla::ifrt::Device * dev1, + client_->LookupDevice(xla::ifrt::DeviceId(default_assignment(0, 1)))); + TF_ASSERT_OK_AND_ASSIGN( auto devices_from_attribute, - GetAssignedIfrtDevices(client_, kNumReplicas, kNumCoresPerReplica, + GetAssignedIfrtDevices(*client_, kNumReplicas, kNumCoresPerReplica, std::nullopt)); - EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[2], devices_[3])); + EXPECT_THAT(devices_from_attribute, ElementsAre(dev0, dev1)); } TEST_F(IfrtDeviceUtilsTest, MismatchCoordinatesShallFail) { - std::vector device_assignment_attr = {1, 0, 0, 1, 3, 0, 0, 0}; - auto status = GetAssignedIfrtDevices(client_, 1, 2, device_assignment_attr); + std::vector device_assignment_attr = {999, 999, 999, 999, + 999, 999, 999, 998}; + auto status = GetAssignedIfrtDevices(*client_, 1, 2, device_assignment_attr); EXPECT_THAT(status, absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition)); } } // namespace - } // namespace ifrt_serving } // namespace tensorflow From fadb2ca7b5c17ba3b9380f32487cef62eebcf2c9 Mon Sep 17 00:00:00 2001 From: Zaim Codi Date: Thu, 24 Sep 2026 14:09:55 -0700 Subject: [PATCH 46/48] Add algebraic simplification logic for kShuffle::kRotate. Extend the HLO optimization pipeline with algebraic simplifications for `HloOpcode::kShuffle` in `kRotate` mode, allowing compiler pipelines to simplify rotations. Specifically, this commit introduces: * **AlgebraicSimplifier (`algebraic_simplifier.cc`)**: Add simplification rules for `kShuffle::kRotate` instructions - Eliminating identity rotations (`norm_shift == 0` , empty shift dimensions, splat dimension) - Folding/combining chained `kShuffle::kRotate` operations. - Canonicalize shifts to be in `[0, dim_size)` and order dimensions in ascending order * **Testing**: Add comprehensive unit tests in `AlgebraicSimplifierTest`. PiperOrigin-RevId: 987747347 --- .../xla/xla/hlo/transforms/simplifiers/BUILD | 2 + .../simplifiers/algebraic_simplifier.cc | 80 ++++++++++++++++ .../simplifiers/algebraic_simplifier.h | 2 + .../simplifiers/algebraic_simplifier_test.cc | 92 ++++++++++++++++++- 4 files changed, 175 insertions(+), 1 deletion(-) diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD index dae4b21fdff4b6..7fd054c9e50849 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD +++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD @@ -486,6 +486,7 @@ cc_library( "//xla:literal_util", "//xla:permutation_util", "//xla:shape_util", + "//xla:shuffle", "//xla:status_macros", "//xla:util", "//xla:window_util", @@ -508,6 +509,7 @@ cc_library( "//xla/service:shape_inference", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/container:btree", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/container:inlined_vector", diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc index 97edf369ef7ba6..303b102cd71dde 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -32,6 +32,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/base/nullability.h" +#include "absl/container/btree_map.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/container/inlined_vector.h" @@ -77,6 +78,7 @@ limitations under the License. #include "xla/service/shape_inference.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/shuffle.h" #include "xla/status_macros.h" #include "xla/util.h" #include "xla/window_util.h" @@ -7098,6 +7100,84 @@ absl::Status AlgebraicSimplifierVisitor::HandleReverse( return absl::OkStatus(); } +absl::Status AlgebraicSimplifierVisitor::HandleShuffle(HloInstruction* hlo) { + auto* shuffle = Cast(hlo); + + switch (shuffle->mode()) { + case ShuffleMode::kRotate: { + // Accumulate shifts per dimension (in ascending order of dim). + absl::btree_map combined_shifts; + for (int64_t i = 0; i < shuffle->dimensions().size(); ++i) { + combined_shifts[shuffle->dimensions()[i]] += + shuffle->rotate().shifts(i); + } + + // rotate(rotate(x, {d}, {m}), {d}, {n}) ==> rotate(x, {d}, {m + n}) + HloInstruction* base_operand = shuffle->mutable_operand(0); + while (base_operand->opcode() == HloOpcode::kShuffle && + Cast(base_operand)->mode() == + ShuffleMode::kRotate) { + auto* inner_shuffle = Cast(base_operand); + for (int64_t i = 0; i < inner_shuffle->dimensions().size(); ++i) { + combined_shifts[inner_shuffle->dimensions()[i]] += + inner_shuffle->rotate().shifts(i); + } + base_operand = inner_shuffle->mutable_operand(0); + } + + // Helper to check if base_operand is a splat along dim. + auto is_dim_splat = [&](int64_t dim) -> bool { + if (base_operand->IsConstant() && + base_operand->literal().IsAllFirst()) { + return true; + } + if (base_operand->opcode() == HloOpcode::kBroadcast) { + return !absl::c_linear_search(base_operand->dimensions(), dim); + } + return false; + }; + + // Canonicalize shifts & dimensions, remove no-op dims. + DimensionVector new_dimensions; + DimensionVector new_shifts; + for (const auto& [dim, total_shift] : combined_shifts) { + int64_t dim_size = shuffle->shape().dimensions(dim); + // No-op: size 1 dim + if (dim_size <= 1) { + continue; + } + // No-op: splat along dim + if (is_dim_splat(dim)) { + continue; + } + int64_t norm_shift = shuffle::NormalizeShift(total_shift, dim_size); + // No-op: shift == 0 + if (norm_shift == 0) { + continue; + } + new_dimensions.push_back(dim); + new_shifts.push_back(norm_shift); + } + + // Replace if changed. + if (new_dimensions.empty()) { + return ReplaceInstruction(shuffle, base_operand); + } + if (base_operand != shuffle->operand(0) || + new_dimensions != shuffle->dimensions() || + !absl::c_equal(new_shifts, shuffle->rotate().shifts())) { + auto new_shuffle = HloInstruction::CreateShuffle( + shuffle->shape(), base_operand, new_dimensions, + shuffle::Rotate(new_shifts)); + return ReplaceWithNewInstruction(shuffle, std::move(new_shuffle)); + } + return absl::OkStatus(); + } + default: + return absl::OkStatus(); + } +} + absl::StatusOr AlgebraicSimplifierVisitor::TrySimplifyScalarSlice( HloInstruction* slice) { // Only try to do this for effective scalars. We could do the same for slicing diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h index 6ea0d835cef53a..7878708c3f5f0d 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h @@ -592,6 +592,8 @@ class AlgebraicSimplifierVisitor : public DfsHloRewriteVisitor { absl::Status HandleRsqrt(HloInstruction* rsqrt) override; + absl::Status HandleShuffle(HloInstruction* hlo) override; + absl::Status HandleSlice(HloInstruction* slice) override; absl::Status HandleSqrt(HloInstruction* sqrt) override; diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc index b439bf6ed0d202..bf5cc9ea7f9b67 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc @@ -15087,6 +15087,96 @@ TEST_F(AlgebraicSimplifierTest, DynamicSliceOfDynamicSlice) { m::ConstantScalar(6)))))); } +TEST_F(AlgebraicSimplifierTest, FusesShuffleRotates) { + constexpr absl::string_view kModuleStr = R"( + HloModule shuffle_module + + ENTRY main { + p0 = f32[5]{0} parameter(0) + r1 = f32[5]{0} shuffle(p0), dimensions={0}, mode=rotate, shifts={2} + ROOT r2 = f32[5]{0} shuffle(r1), dimensions={0}, mode=rotate, shifts={1} + } + )"; + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_TRUE(changed); + + auto* root = m->entry_computation()->root_instruction(); + ASSERT_EQ(root->opcode(), HloOpcode::kShuffle); + auto* shuffle = Cast(root); + EXPECT_EQ(shuffle->operand(0)->opcode(), HloOpcode::kParameter); + ASSERT_EQ(shuffle->dimensions().size(), 1); + EXPECT_EQ(shuffle->dimensions()[0], 0); + EXPECT_EQ(shuffle->rotate().shifts(0), 3); +} + +TEST_F(AlgebraicSimplifierTest, RemovesShuffleRotateOnSplatBroadcast) { + constexpr absl::string_view kModuleStr = R"( + HloModule shuffle_module + + ENTRY main { + c = f32[] constant(1.0) + b = f32[5,5]{1,0} broadcast(c), dimensions={} + ROOT r = f32[5,5]{1,0} shuffle(b), dimensions={0}, mode=rotate, shifts={2} + } + )"; + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_TRUE(changed); + + auto* root = m->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kBroadcast); +} + +TEST_F(AlgebraicSimplifierTest, RemovesNoopDimAndCanonicalizesShift) { + constexpr absl::string_view kModuleStr = R"( + HloModule shuffle_module + + ENTRY main { + p0 = f32[5,5]{1,0} parameter(0) + ROOT r = f32[5,5]{1,0} shuffle(p0), dimensions={1,0}, mode=rotate, + shifts={0,-1} + } + )"; + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_TRUE(changed); + + auto* root = m->entry_computation()->root_instruction(); + ASSERT_EQ(root->opcode(), HloOpcode::kShuffle); + auto* shuffle = Cast(root); + ASSERT_EQ(shuffle->dimensions().size(), 1); + EXPECT_EQ(shuffle->dimensions()[0], 0); + EXPECT_EQ(shuffle->rotate().shifts(0), 4); +} + +TEST_F(AlgebraicSimplifierTest, CanonicalizesMultiDimShuffleRotate) { + constexpr absl::string_view kModuleStr = R"( + HloModule shuffle_module + + ENTRY main { + p0 = f32[5,6]{1,0} parameter(0) + ROOT r = f32[5,6]{1,0} shuffle(p0), dimensions={1,0}, mode=rotate, + shifts={3,-2} + } + )"; + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_TRUE(changed); + + auto* root = m->entry_computation()->root_instruction(); + ASSERT_EQ(root->opcode(), HloOpcode::kShuffle); + auto* shuffle = Cast(root); + ASSERT_EQ(shuffle->dimensions().size(), 2); + EXPECT_EQ(shuffle->dimensions()[0], 0); + EXPECT_EQ(shuffle->rotate().shifts(0), 3); + EXPECT_EQ(shuffle->dimensions()[1], 1); + EXPECT_EQ(shuffle->rotate().shifts(1), 3); +} + } // namespace } // namespace xla -// Trivial comment to force rebuild From 4fa7af70003017504a6c399b148bc9f8cf4f141b Mon Sep 17 00:00:00 2001 From: Bixia Zheng Date: Thu, 24 Sep 2026 14:23:02 -0700 Subject: [PATCH 47/48] Clear sharding annotations on entry parameters at the end of GSPMD partitioning. After GSPMD generates device local, it clears sharding except for two cases: (1) Send/Recv instructions and (2) Parameter instructions. Case (1) is due to the fact that GSPMD doesn't handle the single device sharding on Send/Recv instructions, and it leaves the sharding for later passses to handle it. Case (2) is due to the fact that there are many places in the compilation pipeline that find the module input sharding from the sharding on parameter instructions, even though the information is also available in the module spmd_parameters_shardings metadata. This CL makes GSPMD also clear sharding on parameter instructions, and updates downstream users to read parameter shardings from the module's spmd_parameters_shardings metadata. PiperOrigin-RevId: 987756729 --- .../hlo/analysis/hlo_replication_analysis.cc | 20 ++++-- .../analysis/hlo_replication_analysis_test.cc | 30 +++++++++ third_party/xla/xla/service/gpu/BUILD | 2 - .../gpu/auto_sharding_gpu_compiler_test.cc | 11 ++-- .../service/propagate_original_value_test.cc | 2 +- .../xla/xla/service/spmd/spmd_partitioner.cc | 19 ------ .../xla/service/spmd/spmd_partitioner_test.cc | 61 ++++++++++++------- .../xla/xla/tools/tests/hlo_expand_test.cc | 8 +-- 8 files changed, 94 insertions(+), 59 deletions(-) diff --git a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc index 0f1e1400a6d4da..5443f988a4991e 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc @@ -499,11 +499,19 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { HloReplication::UniqueOnAllDevices()); std::unique_ptr> sharding_tree = nullptr; - if (cross_partition_spmd_ && param->has_sharding()) { - ABSL_ASSIGN_OR_RETURN(auto result, - param->sharding().AsShapeTree(param->shape())); - sharding_tree = - std::make_unique>(std::move(result)); + if (cross_partition_spmd_) { + const HloSharding* sharding = nullptr; + if (module_->has_spmd_parameters_shardings()) { + CHECK_LT(i, module_->spmd_parameters_shardings().size()); + sharding = &module_->spmd_parameters_shardings()[i]; + } else if (param->has_sharding()) { + sharding = ¶m->sharding(); + } + if (sharding != nullptr) { + ABSL_ASSIGN_OR_RETURN(auto result, sharding->AsShapeTree(param->shape())); + sharding_tree = + std::make_unique>(std::move(result)); + } } const auto& replication = param->parameter_replicated_at_leaf_buffers(); @@ -513,7 +521,7 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { if (!ShapeUtil::IsLeafIndex(param->shape(), index)) { return absl::OkStatus(); } - if (cross_partition_spmd_ && param->has_sharding()) { + if (sharding_tree != nullptr) { // In cross-partition spmd mode, set parameter replication status // based on the parameter's sharding. *shape_tree.mutable_element(index) = diff --git a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc index 1b249702307944..6c601626141e72 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include #include #include "absl/types/span.h" #include "xla/hlo/ir/hlo_computation.h" @@ -223,6 +224,35 @@ ENTRY entry { FindInstruction(module.get(), "all-reduce-different-operand"), {})); } +TEST_F(HloReplicationAnalysisTest, + CrossPartitionSpmdWithModuleParameterShardings) { + const std::string module_str = R"( +HloModule CrossPartitionSpmdWithModuleParameterShardings + +ENTRY entry { + param = (f32[4096,4096]{1,0}, f32[4096,4096]{1,0}) + parameter(0), sharding={{maximal device=0}, {replicated}} + gte0 = f32[4096,4096]{1,0} get-tuple-element(param), index=0 + gte1 = f32[4096,4096]{1,0} get-tuple-element(param), index=1 + ROOT add = f32[4096,4096]{1,0} add(gte0, gte1) +} +)"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/4)); + HloInstruction* param = module->entry_computation()->parameter_instruction(0); + module->set_spmd_parameters_shardings({param->sharding()}); + param->clear_sharding(); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr analysis, + HloReplicationAnalysis::Run(module.get(), /*cross_partition_spmd=*/true)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "gte0"), {})); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "gte1"), {})); +} + TEST_F(HloReplicationAnalysisTest, NestedCall) { const std::string module_str = R"( HloModule NestedCall diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 8fdb54cd4a0841..394417048f9d21 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -2428,9 +2428,7 @@ xla_test( deps = [ "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", - "//xla/hlo/testlib:pattern_matcher_gmock", "//xla/service:hlo_module_config", - "//xla/service:pattern_matcher", "//xla/tests:xla_internal_test_main", "//xla/tests/restricted:hlo_test_base_legacy", "@com_google_absl//absl/log", diff --git a/third_party/xla/xla/service/gpu/auto_sharding_gpu_compiler_test.cc b/third_party/xla/xla/service/gpu/auto_sharding_gpu_compiler_test.cc index 512bdd512ec446..d689536b66936b 100644 --- a/third_party/xla/xla/service/gpu/auto_sharding_gpu_compiler_test.cc +++ b/third_party/xla/xla/service/gpu/auto_sharding_gpu_compiler_test.cc @@ -16,14 +16,12 @@ limitations under the License. #include #include -#include #include #include "absl/log/log.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/testlib/pattern_matcher_gmock.h" +#include "xla/hlo/ir/hlo_sharding.h" #include "xla/service/hlo_module_config.h" -#include "xla/service/pattern_matcher.h" #include "xla/tests/restricted/hlo_test_base_legacy.h" #include "xla/xla_data.pb.h" #include "tsl/platform/logging.h" @@ -32,8 +30,6 @@ namespace xla { namespace gpu { namespace { -namespace m = ::xla::match; - class AutoShardingTest : public HloTestBaseLegacy { protected: const char* const dot_hlo_string_ = R"( @@ -68,7 +64,10 @@ TEST_F(AutoShardingTest, MatMulWithoutAutosharding) { auto* instruction = compiled_module->entry_computation()->parameter_instruction(0); VLOG(2) << instruction->ToString(); - EXPECT_THAT(instruction, GmockMatch(m::Op().WithSharding("{replicated}"))); + EXPECT_FALSE(instruction->has_sharding()); + ASSERT_TRUE(compiled_module->has_spmd_parameters_shardings()); + EXPECT_EQ(compiled_module->spmd_parameters_shardings()[0], + HloSharding::Replicate()); } } // namespace diff --git a/third_party/xla/xla/service/propagate_original_value_test.cc b/third_party/xla/xla/service/propagate_original_value_test.cc index 7dc0b8c8cdfbfd..124026e64b664d 100644 --- a/third_party/xla/xla/service/propagate_original_value_test.cc +++ b/third_party/xla/xla/service/propagate_original_value_test.cc @@ -222,7 +222,7 @@ TEST_F(OriginalValueRecoveryTableTest, constexpr absl::string_view hlo_string = R"hlo( // CHECK-NOT: origin_recovery_table // CHECK: ENTRY %[[COMPUTATION:.*]] (param: f32[4]) -> f32[4] -// CHECK: parameter(0), sharding={replicated}, origin={{[{]}}{"param_origin"} +// CHECK: parameter(0), origin={{[{]}}{"param_origin"} // CHECK: negate(%param), origin={{[{]}}{"negate_origin"} HloModule test, entry_computation_layout={(f32[4]{0})->f32[4]{0}}, num_partitions=2 diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc index a866832ef944b1..4974a9b69ae604 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc @@ -232,10 +232,6 @@ bool ShouldKeepSharding(const HloInstruction* hlo, DynCast(hlo) != nullptr) { return true; } - if (hlo->opcode() == HloOpcode::kParameter && - hlo->parent() == hlo->GetModule()->entry_computation()) { - return true; - } if (keep_valid_shardings && hlo->has_sharding()) { // SPMD partitioner can generate invalid shardings since sharding is // meaningless after this pass. We only keep valid shardings to avoid @@ -253,13 +249,6 @@ absl::Status ClearShardingAttributes( auto has_unreduced_axes = [](const HloInstruction* hlo) -> bool { return hlo->frontend_attributes().map().contains(sdy::kHasUnreducedAxes); }; - for (int64_t i = 0; i < module->entry_computation()->num_parameters(); ++i) { - // Recover the unreduced sharding for parameters. - auto param = module->entry_computation()->parameter_instruction(i); - if (has_unreduced_axes(param)) { - param->set_sharding(module->spmd_parameters_shardings()[i]); - } - } const bool keep_shardings_after_spmd = module->config().debug_options().xla_keep_shardings_after_spmd(); for (HloComputation* computation : module->computations(execution_threads)) { @@ -7408,13 +7397,6 @@ absl::StatusOr SpmdPartitioner::RunImpl( ABSL_RETURN_IF_ERROR(ClearShardingAttributes( module, num_replicas() * num_partitions(), execution_threads)); - auto entry_root = module->entry_computation()->root_instruction(); - if (entry_root->has_sharding()) { - HloSharding final_sharding = - ResolveReductionOpForSharding(entry_root, entry_root->sharding()); - entry_root->set_sharding(std::move(final_sharding)); - } - if (changed) { HloPassPipeline pass("spmd-cleanup"); pass.AddPass(/*remove_cross_partition_collective_ops=*/true); @@ -7423,7 +7405,6 @@ absl::StatusOr SpmdPartitioner::RunImpl( pass.AddPass(/*is_layout_sensitive=*/false); ABSL_RETURN_IF_ERROR(pass.Run(module, execution_threads).status()); } - return changed; } diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc index 024a14599209d8..c9deeba34249bc 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc @@ -11034,10 +11034,12 @@ ENTRY entry { EXPECT_THAT(module->entry_computation()->root_instruction(), op::Tuple(sum_full, sum_manual)); - EXPECT_THAT(module->entry_computation()->parameter_instruction(0), - op::Sharding("{{devices=[2,1]<=[2]},{manual},{manual}}")); - EXPECT_THAT(module->entry_computation()->parameter_instruction(1), - op::Sharding("{{manual},{devices=[2,1]<=[2]},{manual}}")); + EXPECT_THAT(module->entry_computation()->parameter_instructions(), + Each(op::NoSharding())); + EXPECT_THAT( + module->spmd_parameters_shardings(), + ElementsAre(*ParseSharding("{{devices=[2,1]<=[2]},{manual},{manual}}"), + *ParseSharding("{{manual},{devices=[2,1]<=[2]},{manual}}"))); } TEST_P(SpmdPartitioningTest, NestedManual) { @@ -18471,9 +18473,12 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - // Check that unreduced HloSharding is preserved after the pass. + // Check that unreduced HloSharding is preserved on the module and cleared on + // the parameter instructions after the pass. EXPECT_THAT(module->entry_computation()->parameter_instructions(), - Each(op::Sharding("{unreduced}"))); + Each(op::NoSharding())); + EXPECT_THAT(module->spmd_parameters_shardings(), + Each(*ParseSharding("{unreduced}"))); } TEST_P(SpmdPartitioningTest, SubgroupUnreducedParam) { @@ -18488,10 +18493,13 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/4)); VLOG(1) << module->ToString(); - // Check that unreduced HloSharding is preserved after the pass. - EXPECT_THAT( - module->entry_computation()->parameter_instructions(), - Each(op::Sharding("{devices=[1,2,2]<=[4] last_tile_dims={unreduced}}"))); + // Check that unreduced HloSharding is preserved on the module and cleared on + // the parameter instructions after the pass. + EXPECT_THAT(module->entry_computation()->parameter_instructions(), + Each(op::NoSharding())); + EXPECT_THAT(module->spmd_parameters_shardings(), + Each(*ParseSharding( + "{devices=[1,2,2]<=[4] last_tile_dims={unreduced}}"))); } TEST_P(SpmdPartitioningTest, SubgroupUnreducedDot) { @@ -19188,9 +19196,12 @@ ENTRY entry { })"; ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); - // Check that unreduced HloSharding is preserved after the pass. - for (auto* param : module->entry_computation()->parameter_instructions()) { - EXPECT_THAT(param->sharding().ToString(), HasSubstr("unreduced")); + // Check that unreduced HloSharding is preserved on the module and cleared on + // the parameter instructions after the pass. + EXPECT_THAT(module->entry_computation()->parameter_instructions(), + Each(op::NoSharding())); + for (const auto& sharding : module->spmd_parameters_shardings()) { + EXPECT_THAT(sharding.ToString(), HasSubstr("unreduced")); } } @@ -19204,10 +19215,13 @@ ENTRY entry { })"; ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/4)); - // Check that unreduced HloSharding is preserved after the pass. + // Check that unreduced HloSharding is preserved on the module and cleared on + // the parameter instructions after the pass. + EXPECT_THAT(module->entry_computation()->parameter_instructions(), + Each(op::NoSharding())); EXPECT_THAT( - module->entry_computation()->parameter_instructions(), - Each(op::Sharding("{mesh['x'=2,'y'=2] [{?},{'x'}], unreduced={'y'}}"))); + module->spmd_parameters_shardings(), + Each(*ParseSharding("{mesh['x'=2,'y'=2] [{?},{'x'}], unreduced={'y'}}"))); } TEST_F(SpmdPartitioningV3Test, TupleSubgroupUnreducedParamV3) { @@ -19225,8 +19239,10 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); EXPECT_THAT(module->entry_computation()->parameter_instructions(), - Each(op::Sharding("{{mesh['x'=2] [], unreduced={'x'}}, " - "{mesh['x'=2] [], unreduced={'x'}}}"))); + Each(op::NoSharding())); + EXPECT_THAT(module->spmd_parameters_shardings(), + Each(*ParseSharding("{{mesh['x'=2] [], unreduced={'x'}}, " + "{mesh['x'=2] [], unreduced={'x'}}}"))); } TEST_F(SpmdPartitioningV3Test, PatternMatchMergeNamedSharding) { @@ -19490,10 +19506,13 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/4)); VLOG(1) << module->ToString(); - // Check that unreduced HloSharding is preserved after the pass. + // Check that unreduced HloSharding is preserved on the module and cleared on + // the parameter instructions after the pass. + EXPECT_THAT(module->entry_computation()->parameter_instructions(), + Each(op::NoSharding())); EXPECT_THAT( - module->entry_computation()->parameter_instructions(), - Each(op::Sharding("{mesh['a'=2,'b'=2], [], unreduced=max{'a','b'}}"))); + module->spmd_parameters_shardings(), + Each(*ParseSharding("{mesh['a'=2,'b'=2], [], unreduced=max{'a','b'}}"))); } // Verifies that a true 1D scatter-conflict on a reduction dimension (e.g., diff --git a/third_party/xla/xla/tools/tests/hlo_expand_test.cc b/third_party/xla/xla/tools/tests/hlo_expand_test.cc index a38dd367957162..7365f107b8ba94 100644 --- a/third_party/xla/xla/tools/tests/hlo_expand_test.cc +++ b/third_party/xla/xla/tools/tests/hlo_expand_test.cc @@ -88,9 +88,9 @@ TEST_F(HloExpandTest, SpmdHloWithoutHloShardingV3) { R"(HloModule module, entry_computation_layout={(f32[24,64]{1,0}, f32[39296,64]{1,0})->f32[24,19648]{1,0}}, num_partitions=2 ENTRY %entry_spmd (param: f32[24,64], param.1: f32[39296,64]) -> f32[24,19648] { - %param = f32[24,64]{1,0} parameter(0), sharding={replicated} + %param = f32[24,64]{1,0} parameter(0) %lhs.copy.1 = f32[24,64]{1,0} copy(%param) - %param.1 = f32[39296,64]{1,0} parameter(1), sharding={replicated} + %param.1 = f32[39296,64]{1,0} parameter(1) %constant = s32[2]{0} constant({0, 19648}) %partition-id = u32[] partition-id() %dynamic-slice = s32[1]{0} dynamic-slice(%constant, %partition-id), dynamic_slice_sizes={1} @@ -118,9 +118,9 @@ TEST_F(HloExpandTest, SpmdHloWithHloShardingV3) { R"(HloModule module, entry_computation_layout={(f32[24,64]{1,0}, f32[39296,64]{1,0})->f32[24,19648]{1,0}}, num_partitions=2 ENTRY %entry_spmd (param: f32[24,64], param.1: f32[39296,64]) -> f32[24,19648] { - %param = f32[24,64]{1,0} parameter(0), sharding={mesh[], replicated} + %param = f32[24,64]{1,0} parameter(0) %lhs.copy.1 = f32[24,64]{1,0} copy(%param) - %param.1 = f32[39296,64]{1,0} parameter(1), sharding={mesh[], replicated} + %param.1 = f32[39296,64]{1,0} parameter(1) %constant = s32[2]{0} constant({0, 19648}) %partition-id = u32[] partition-id() %dynamic-slice = s32[1]{0} dynamic-slice(%constant, %partition-id), dynamic_slice_sizes={1} From b5d8801b1e0db562ced3f7f940c65f25d7aea453 Mon Sep 17 00:00:00 2001 From: Jian Cai Date: Thu, 24 Sep 2026 14:48:55 -0700 Subject: [PATCH 48/48] [XLA][HLO Value Tracking] Maintain OriginalValue consistency in MemorySpaceAssignment Root Cause: `MemorySpaceAssignment::Process()` rewrites the HLO graph in two sequential loops over all `Allocation` objects: 1. Loop 1 calls the virtual method `Allocation::Process()` (`CopyAllocation::Process()`, `PinnedAllocation::Process()`, and `ParentAllocation::Process()`). 2. Loop 2 calls the virtual method `Allocation::PostProcess()` (`ParentAllocation::PostProcess()`). Several HLO transformations across these two loops failed to maintain or propagate `OriginalValue`: 1. While-loop widening in `ParentAllocation::Process()` (Loop 1) and `ParentAllocation::PostProcess()` (Loop 2): When `ParentAllocation` passes a buffer from a parent computation into a `while` loop, it widens the `while` loop tuple from N elements to N+1 elements at `new_tuple_index`: - `ParentAllocation::Process()` widens the `while` loop inputs by creating a new input tuple (`new_while_operand = TupleUtil::ReplaceTupleWith(...)`) and mutating the shapes of `calling_instruction_` (the `kWhile` instruction), `while_condition->parameter_instruction(0)`, and `while_body->parameter_instruction(0)` in-place. Because the in-place mutated instructions kept their old N-element `OriginalValue` trees while their shapes became (N+1)-tuples, `HloVerifier` failed with shape mismatches. Meanwhile, `new_while_operand` was created with `original_value() == nullptr`. - To keep downstream users of the `while` loop valid, `ParentAllocation::Process()` creates `tuple_with_old_shape = TupleUtil::ExtractPrefix(calling_instruction_, new_tuple_index)` and replaces uses of `calling_instruction_` with `tuple_with_old_shape`. `tuple_with_old_shape` and its `GetTupleElement` operands were created with `original_value() == nullptr`, dropping value tracking for all users after the loop. - `ParentAllocation::PostProcess()` runs in Loop 2 (after all Loop 1 allocations inside `while_body` have finished) to widen the `while_body` root tuple (`new_while_body_root = TupleUtil::ReplaceTupleWith(added_element, old_body_root, ...)`), which also left `new_while_body_root` with `original_value() == nullptr` and dropped the appended element's `OriginalValue` at `{new_tuple_index}`. 2. Tuple replacement and async copy creation in `CopyAllocation::Process()` and `Allocation::UpdateUses()` (Loop 1): - `CopyAllocation::Process()` inserts `copy-start` and `copy-done` instructions to move buffers between default memory (HBM) and alternate memory (VMEM), but left `copy_done_` with `original_value() == nullptr` instead of copying `producing_instruction`'s `OriginalValue`. - `CopyAllocation::Process()`, `PinnedAllocation::Process()`, and `ParentAllocation::Process()` all call the helper `Allocation::UpdateUses()` during Loop 1 to rewire consumer instructions to use `producing_instruction`. When a consumer is a `kTuple` instruction (such as `while_body->root_instruction()` or `while_init`), `Allocation::UpdateUses()` calls `TupleUtil::ReplaceTupleWith`, which constructs a brand-new `kTuple` instruction with `original_value() == nullptr`. Consequently, any Loop 1 allocation updating an element of `while_body->root_instruction()` wiped out its `OriginalValue` to `nullptr` before Loop 2 (`ParentAllocation::PostProcess()`) even ran. What This CL Fixes: - Add `TupleTree::num_children(ShapeIndexView index = {})` in `tuple_tree.h` to query the number of children of a tuple node directly from the index table. - Introduce `AppendToTupleOriginalValue(src_tuple, dest_tuple, appended_instr)` in `allocation.cc`, which uses `CopyOriginalValue` from `hlo_original_value_util.h` to map existing tuple elements from `src_tuple` to the widened `dest_tuple` and copies `appended_instr->original_value()->tree()` into the last element (`dest_tuple->shape().tuple_shapes().size() - 1`) via `CopyCompatibleSubtreeFrom`. - In `ParentAllocation::Process()`, before widening `calling_instruction_` in-place, copy its N-element `OriginalValue` to `tuple_with_old_shape` via `CopyOriginalValue` and propagate subtrees to its `GetTupleElement` operands. Then widen `new_while_operand` with `AppendToTupleOriginalValue` and update `calling_instruction_`, `while_condition` parameter 0, and `while_body` parameter 0 via `AppendToWhileLoopOriginalValue`. - In `AppendToWhileLoopOriginalValue` (`while_util.cc`), only update `while_body->root_instruction()` if its shape has already been widened to match `while_instr->shape()`, supporting two-phase while-loop widening where `while_body` root widening is deferred to `ParentAllocation::PostProcess()`. - In `Allocation::AddGetTupleElements()`, propagate the `OriginalValue` subtree from `defining_position().instruction` at `defining_position().index` to the newly created `GetTupleElement` instruction. - In `ParentAllocation::PostProcess()`, call `AppendToTupleOriginalValue(old_body_root, new_while_body_root, added_element)` so the widened `while_body` root tuple preserves both the body root's `OriginalValue` and the pass-through element's `OriginalValue`. - In `CopyAllocation::Process()`, propagate `producing_instruction`'s `OriginalValue` to `copy_done_` (and any intermediate bitcast). - In `Allocation::UpdateUses()`, when `TupleUtil::ReplaceTupleWith` replaces a `kTuple` instruction, clone `tuple_inst`'s `OriginalValue` onto `replacement_instruction` and update `{use.operand_index}` from `producing_instruction->original_value()`. PiperOrigin-RevId: 987774950 --- .../memory_space_assignment/allocation.cc | 1502 +++++++++++++++++ .../allocation_test.cc | 393 +++++ 2 files changed, 1895 insertions(+) create mode 100644 third_party/xla/xla/service/memory_space_assignment/allocation.cc create mode 100644 third_party/xla/xla/service/memory_space_assignment/allocation_test.cc diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation.cc b/third_party/xla/xla/service/memory_space_assignment/allocation.cc new file mode 100644 index 00000000000000..7e060854083011 --- /dev/null +++ b/third_party/xla/xla/service/memory_space_assignment/allocation.cc @@ -0,0 +1,1502 @@ +/* Copyright 2024 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/service/memory_space_assignment/allocation.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/base/casts.h" +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/functional/function_ref.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "xla/hlo/analysis/hlo_alias_analysis.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_original_value.h" +#include "xla/hlo/ir/hlo_original_value_util.h" +#include "xla/hlo/utils/hlo_live_range.h" +#include "xla/layout.h" +#include "xla/layout_util.h" +#include "xla/service/heap_simulator/allocation_block.h" +#include "xla/service/heap_simulator/heap_simulator.h" +#include "xla/service/hlo_value.h" +#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" +#include "xla/service/memory_space_assignment/slice.h" +#include "xla/service/time_utils.h" +#include "xla/service/tuple_util.h" +#include "xla/service/while_util.h" +#include "xla/shape.h" +#include "xla/shape_util.h" +#include "xla/status_macros.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/statusor.h" +#include "xla/util.h" +#include "xla/xla_data.pb.h" +#include "tsl/platform/casts.h" +#include "tsl/platform/errors.h" +#include "tsl/platform/statusor.h" + +namespace xla::memory_space_assignment { +namespace { + +std::string UsesToString(const std::vector& uses) { + if (uses.empty()) { + return "none"; + } + std::vector uses_str; + uses_str.reserve(uses.size()); + for (const auto& use : uses) { + uses_str.push_back(use.ToString()); + } + return absl::StrJoin(uses_str, ","); +} + +void AppendToTupleOriginalValue(const HloInstruction* src_tuple, + HloInstruction* dest_tuple, + const HloInstruction* appended_instr) { + if (src_tuple == nullptr || dest_tuple == nullptr || + !dest_tuple->shape().IsTuple() || + dest_tuple->shape().tuple_shapes().empty()) { + return; + } + bool has_appended_ov = + appended_instr != nullptr && appended_instr->original_value() != nullptr; + if (src_tuple->original_value() == nullptr && !has_appended_ov) { + return; + } + int64_t new_tuple_index = + static_cast(dest_tuple->shape().tuple_shapes().size()) - 1; + if (src_tuple->original_value() != nullptr) { + int64_t src_size = src_tuple->original_value()->tree().num_children(); + int64_t num_to_copy = std::min(src_size, new_tuple_index); + absl::flat_hash_map old_to_new_tuple_idx; + for (int64_t i = 0; i < num_to_copy; ++i) { + old_to_new_tuple_idx[i] = i; + } + CopyOriginalValue(src_tuple, dest_tuple, old_to_new_tuple_idx); + } else { + dest_tuple->set_original_value( + std::make_shared(dest_tuple->shape())); + } + if (has_appended_ov && dest_tuple->original_value() != nullptr) { + CHECK_OK( + dest_tuple->original_value()->mutable_tree()->CopyCompatibleSubtreeFrom( + appended_instr->original_value()->tree(), {}, {new_tuple_index})) + << "Incompatible OriginalValue subtree for appended element at index " + << new_tuple_index << ": " << appended_instr->ToString(); + } + CHECK(dest_tuple->original_value()->IsCompatibleWith(dest_tuple->shape())) + << "Updated OriginalValue is incompatible with shape: " + << dest_tuple->ToString(); +} + +// Helper function to compute the start time for a SlicedCopyAllocation. +int64_t GetSlicedCopyAllocationExclusiveStartTime( + const std::vector& + slice_decisions_sorted_by_exclusive_start_time) { + if (slice_decisions_sorted_by_exclusive_start_time.empty()) { + return -1; + } + + return slice_decisions_sorted_by_exclusive_start_time.front() + .exclusive_start_time; +} + +// Helper function to compute the underlying Allocation chunk for a +// SlicedCopyAllocation. +std::optional GetSlicedCopyAllocationChunk( + const std::vector& slice_decisions_sorted_by_start_time) { + if (slice_decisions_sorted_by_start_time.empty()) { + return std::nullopt; + } + auto offset_cmp = [](const SliceDecision& lhs, const SliceDecision& rhs) { + return lhs.chunk.offset < rhs.chunk.offset; + }; + auto end_cmp = [](const SliceDecision& lhs, const SliceDecision& rhs) { + return lhs.chunk.chunk_end() < rhs.chunk.chunk_end(); + }; + return HeapSimulator::Chunk::FromOffsetEnd( + std::min_element(slice_decisions_sorted_by_start_time.begin(), + slice_decisions_sorted_by_start_time.end(), offset_cmp) + ->chunk.offset, + std::max_element(slice_decisions_sorted_by_start_time.begin(), + slice_decisions_sorted_by_start_time.end(), end_cmp) + ->chunk.chunk_end()); +} + +} // namespace + +std::string MemorySpaceToString(MemorySpace memory_space) { + switch (memory_space) { + case MemorySpace::kDefault: + return "def"; + case MemorySpace::kAlternate: + return "alt"; + } +} + +std::optional Allocation::cross_program_prefetch_index() const { + return cross_program_prefetch_index_; +} + +HeapSimulator::Chunk Allocation::chunk() const { + CHECK(chunk_.has_value()); + return *chunk_; +} + +void Allocation::set_offset(int64_t offset) { + CHECK(chunk_.has_value()); + *chunk_ = HeapSimulator::Chunk::FromOffsetSize(offset, chunk_->size); +} + +bool Allocation::is_in_alternate_mem() const { + return memory_space_ == MemorySpace::kAlternate; +} + +bool Allocation::is_in_default_mem() const { + return memory_space_ == MemorySpace::kDefault; +} + +void Allocation::RemoveUse(HloUse use) { + uses_.erase(std::remove_if(uses_.begin(), uses_.end(), + [=](const auto& u) { return u == use; }), + uses_.end()); +} + +void Allocation::AddUse(HloUse use) { + HloInstruction* operand = + use.instruction->mutable_operand(use.operand_number); + // If the use is a tuple, look inside the tuple to find the actual use. + for (int64_t index : use.operand_index) { + if (operand->opcode() != HloOpcode::kTuple) { + break; + } + operand = operand->mutable_operand(index); + } + + // Look beyond GetTupleElement(Tuple()) pattern for any bitcasts. + std::function get_simplified_operand; + get_simplified_operand = [&](HloInstruction* instruction) { + while (instruction->opcode() == HloOpcode::kGetTupleElement) { + HloInstruction* operand = + get_simplified_operand(instruction->mutable_operand(0)); + if (operand->opcode() == HloOpcode::kTuple) { + instruction = operand->mutable_operand(instruction->tuple_index()); + } else { + return instruction; + } + } + return instruction; + }; + operand = get_simplified_operand(operand); + + uses_.push_back(use); +} + +absl::Status Allocation::UpdateUses(HloComputation* computation, + HloInstruction* producing_instruction, + const BitcastSplitFn& bitcast_split_fn, + const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + for (const HloUse& use : uses()) { + HloInstruction* replacement_instruction = producing_instruction; + const Shape& operand_shape = + use.instruction->operand(use.operand_number)->shape(); + if (operand_shape.IsTuple()) { + HloInstruction* tuple_inst = + use.instruction->mutable_operand(use.operand_number); + + // Trace back to original instructions in the schedule to use + // HloAliasAnalysis safely. + auto find_original_leaf = [&](const HloInstruction* inst, + ShapeIndex index) + -> std::pair { + while (!hlo_live_range.instruction_schedule().contains(inst)) { + if (inst->opcode() == HloOpcode::kTuple) { + if (index.empty()) { + break; + } + int64_t op_idx = index[0]; + inst = inst->operand(op_idx); + index = ShapeIndex(index.begin() + 1, index.end()); + } else if (inst->opcode() == HloOpcode::kGetTupleElement) { + ShapeIndex new_index; + new_index.push_back(inst->tuple_index()); + for (int i : index) { + new_index.push_back(i); + } + inst = inst->operand(0); + index = new_index; + } else if (inst->opcode() == HloOpcode::kBitcast || + inst->opcode() == HloOpcode::kCopy || + inst->opcode() == HloOpcode::kCopyDone || + inst->opcode() == HloOpcode::kCopyStart) { + inst = inst->operand(0); + } else { + break; + } + } + return {inst, index}; + }; + + bool should_skip_reconstruction = false; + + // If the producing instruction shares the same buffer with the operand, + // then we skip the reconstruction. + auto is_aliasing_operand = [&](int operand_index) { + ShapeIndex sub_index; + if (use.operand_index.size() > 1) { + sub_index = ShapeIndex(use.operand_index.begin() + 1, + use.operand_index.end()); + } + auto [traced_operand, traced_index] = + find_original_leaf(tuple_inst->operand(operand_index), sub_index); + return hlo_live_range.instruction_schedule().contains(traced_operand) && + hlo_live_range.instruction_schedule().contains(tuple_inst) && + alias_analysis.GetUniqueBufferAt(traced_operand, traced_index) + .id() == + alias_analysis + .GetUniqueBufferAt(tuple_inst, use.operand_index) + .id(); + }; + + for (int operand_index = 0; operand_index < tuple_inst->operand_count(); + ++operand_index) { + if (tuple_inst->opcode() == HloOpcode::kTuple && + !use.operand_index.empty() && + operand_index != use.operand_index.front()) { + continue; + } + if (tuple_inst->operand(operand_index) == producing_instruction && + is_aliasing_operand(operand_index)) { + should_skip_reconstruction = true; + break; + } + } + + if (should_skip_reconstruction) { + VLOG(4) << "Skipping tuple reconstruction for " << tuple_inst->name() + << " at index " << use.operand_index.ToString(); + replacement_instruction = tuple_inst; + } else { + ABSL_ASSIGN_OR_RETURN( + replacement_instruction, + TupleUtil::ReplaceTupleWith(producing_instruction, tuple_inst, + use.operand_index)); + if (replacement_instruction != tuple_inst) { + replacement_instruction->CopyOriginalValue(tuple_inst, + /*clone=*/true); + if (replacement_instruction->original_value() != nullptr && + producing_instruction->original_value() != nullptr) { + CHECK_OK(replacement_instruction->original_value() + ->mutable_tree() + ->CopyCompatibleSubtreeFrom( + producing_instruction->original_value()->tree(), + {}, use.operand_index)); + } + } + } + } else if (!Shape::Equal().IgnoreSplitConfigInLayout()( + operand_shape, producing_instruction->shape())) { + // When processing allocations, we treat bitcasts as trivial positions and + // do not create allocations for them. We insert bitcasts after copies, to + // account for the fact that we don't have an allocation for the bitcast. + VLOG(4) << "Old shape = " << operand_shape.ToString() + << ", new shape = " << producing_instruction->shape().ToString() + << "; inserting a bitcast."; + replacement_instruction = computation->AddInstruction( + HloInstruction::CreateBitcast(operand_shape, producing_instruction)); + replacement_instruction->CopyOriginalValue(producing_instruction); + if (mutable_split_shape().has_value() && + producing_instruction->shape().layout().split_configs().size() != 0) { + ABSL_ASSIGN_OR_RETURN( + int64_t split_dim, + bitcast_split_fn( + replacement_instruction, + mutable_split_shape()->layout().split_configs(0).dimension())); + SplitConfig split_config( + split_dim, + {mutable_split_shape()->layout().split_configs(0).split_indices( + 0)}); + replacement_instruction->mutable_shape() + ->mutable_layout() + ->add_split_configs(split_config); + } + } + ABSL_RETURN_IF_ERROR(use.instruction->ReplaceOperandWith( + use.operand_number, replacement_instruction)); + } + return absl::OkStatus(); +} + +bool Allocation::is_copy_like_allocation() const { + return is_copy_allocation() || is_sliced_copy_allocation(); +} + +HloInstruction* Allocation::AddGetTupleElements() const { + CHECK_NE(defining_position().instruction, nullptr); + + Shape shape = defining_position().shape(); + CHECK(shape.IsArray()) << "Allocation shape is not an array. Shape = " + << shape.ToString() + << " position = " << defining_position().shape(); + HloInstruction* result = TupleUtil::AddGetTupleElements(defining_position()); + if (result != defining_position().instruction && + result->original_value() == nullptr && + defining_position().instruction->original_value() != nullptr) { + std::shared_ptr src_ov = + defining_position().instruction->original_value(); + const auto& src_tree = src_ov->tree(); + if (src_tree.find(defining_position().index) != src_tree.end()) { + auto new_ov = std::make_shared(result->shape()); + CHECK_OK(new_ov->mutable_tree()->CopyCompatibleSubtreeFrom( + src_tree, defining_position().index, {})); + if (!new_ov->IsEmpty()) { + result->set_original_value(std::move(new_ov)); + } + } + } + return result; +} + +Allocation::Allocation(HloPosition defining_position, MemorySpace memory_space, + std::optional chunk, + int64_t start_time, int64_t end_time, + std::optional cross_program_prefetch_index) + : original_defining_position_(std::move(defining_position)), + memory_space_(memory_space), + chunk_(chunk), + start_time_(start_time), + end_time_(end_time), + cross_program_prefetch_index_(cross_program_prefetch_index), + split_shape_(std::nullopt) {} + +HloPosition Allocation::original_defining_position() const { + return original_defining_position_; +} + +void Allocation::set_original_defining_position(HloPosition defining_position) { + original_defining_position_ = std::move(defining_position); +} + +bool Allocation::base_is_equal(const Allocation& other) const { + return defining_position() == other.defining_position() && + uses() == other.uses() && memory_space() == other.memory_space() && + chunk() == other.chunk() && start_time() == other.start_time() && + end_time() == other.end_time() && + earliest_available_time() == other.earliest_available_time(); +} + +PinnedAllocation::PinnedAllocation(HloPosition defining_position, + MemorySpace memory_space, + std::optional chunk, + int64_t start_time, int64_t end_time) + : Allocation(std::move(defining_position), memory_space, chunk, start_time, + end_time, + /*cross_program_prefetch_index=*/std::nullopt) {} + +HloPosition PinnedAllocation::defining_position() const { + return original_defining_position(); +} + +bool PinnedAllocation::operator==(const PinnedAllocation& other) const { + return this->base_is_equal(static_cast(other)); +} + +bool MirroredAllocation::operator==(const MirroredAllocation& other) const { + return this->base_is_equal(static_cast(other)); +} + +bool ParentAllocation::operator==(const ParentAllocation& other) const { + return this->base_is_equal(static_cast(other)); +} + +bool PinnedAllocation::operator==(const Allocation& other) const { + const PinnedAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +absl::Status PinnedAllocation::Process(const BitcastSplitFn& bitcast_split_fn, + const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + HloInstruction* producing_instruction = AddGetTupleElements(); + HloComputation* computation = producing_instruction->parent(); + return UpdateUses(computation, producing_instruction, bitcast_split_fn, + hlo_live_range, alias_analysis); +} + +std::string PinnedAllocation::ToString() const { + std::string memory_space_str = MemorySpaceToString(memory_space()); + std::optional chunk = maybe_chunk(); + if (chunk) { + absl::StrAppend(&memory_space_str, " (off: ", chunk->offset, + ", size: ", chunk->size, ")"); + } + return absl::StrCat("PinnedAllocation in ", memory_space_str, " defined at ", + original_defining_position().ToString(), + ", start_time:", start_time(), ", end_time:", end_time(), + ", uses: ", UsesToString(uses())); +} + +void PinnedAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void PinnedAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + needed_allocations.insert(this); +} + +ReservedAllocation::ReservedAllocation(HloPosition defining_position, + HeapSimulator::Chunk chunk, + int64_t start_time, int64_t end_time) + : Allocation(std::move(defining_position), MemorySpace::kAlternate, chunk, + /*start_time=*/start_time, /*end_time=*/end_time, + /*cross_program_prefetch_index=*/std::nullopt), + reserved_(true) {} + +HloPosition ReservedAllocation::defining_position() const { + return original_defining_position(); +} + +absl::Status ReservedAllocation::Process( + const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + return absl::OkStatus(); +} + +void ReservedAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void ReservedAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + needed_allocations.insert(this); +} + +std::string ReservedAllocation::ToString() const { + std::string memory_space_str = MemorySpaceToString(memory_space()); + std::optional chunk = maybe_chunk(); + if (chunk) { + absl::StrAppend(&memory_space_str, " chunk: ", chunk->ToString()); + } + return absl::StrCat("ReservedAllocation defined in ", memory_space_str, + " defined at ", original_defining_position().ToString(), + ", start_time: ", start_time(), + ", end_time: ", end_time(), " reserved: ", reserved_); +} + +bool ReservedAllocation::operator==(const Allocation& other) const { + const ReservedAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +bool ReservedAllocation::operator==(const ReservedAllocation& other) const { + return this->base_is_equal(static_cast(other)) && + reserved_ == other.reserved_; +} + +CopyAllocation::CopyAllocation( + Allocation& prev_allocation, MemorySpace memory_space, + std::optional chunk, + int64_t copy_start_schedule_after_time, + int64_t copy_done_schedule_before_time, int64_t end_time, + std::optional cross_program_prefetch_index, + HloInstruction* sync_mem_op, HloInstruction* async_mem_op_start, + HloInstruction* async_mem_op_done, int64_t source_operand_index) + : Allocation( + /*defining_position=*/{nullptr, {}}, memory_space, chunk, + // Allocation uses an inclusive start time + ExclusiveToInclusiveStartTime(copy_start_schedule_after_time), + end_time, cross_program_prefetch_index), + prev_allocation_(prev_allocation), + copy_start_schedule_after_(copy_start_schedule_after_time), + copy_done_schedule_before_(copy_done_schedule_before_time), + copy_start_(async_mem_op_start), + copy_done_(async_mem_op_done), + sync_mem_op_(sync_mem_op), + source_operand_index_(source_operand_index) {} + +int64_t CopyAllocation::earliest_available_time() const { + return copy_done_schedule_before_; +} + +absl::Status CopyAllocation::Process(const BitcastSplitFn& bitcast_split_fn, + const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + // If the copy start and copy done instructions are already present in the + // original schedule we just need to schedule them in correct positions. There + // is no need to create new instructions. + if (copy_start_ != nullptr && copy_done_ != nullptr) { + // Update the allocation position with the copy complete instruction, so + // that if there are further copies from it, they can find the correct + // position. + set_original_defining_position(HloPosition{copy_start_, {0}}); + return absl::OkStatus(); + } + + // Copy allocations need to insert asynchronous copy nodes. + Shape shape = defining_position().shape(); + HloInstruction* producing_instruction = AddGetTupleElements(); + HloComputation* computation = producing_instruction->parent(); + VLOG(3) << "Processing copy allocation: " << ToString(); + if (sync_mem_op_ != nullptr && sync_mem_op_->opcode() != HloOpcode::kCopy) { + if (sync_mem_op_->opcode() == HloOpcode::kSlice || + sync_mem_op_->opcode() == HloOpcode::kDynamicSlice || + sync_mem_op_->IsCustomFusion()) { + ABSL_ASSIGN_OR_RETURN(copy_done_, + computation->CreateAsyncInstructions( + sync_mem_op_, {ShapeUtil::MakeShape(S32, {})}, + HloInstruction::kMainExecutionThread, false)); + } else { + return Internal("Sync mem op is not a copy, slice, or dynamic slice."); + } + copy_start_ = copy_done_->mutable_operand(0); + // If the shape of the copy start operand is not compatible with the + // shape of the producing instruction, we insert a bitcast to make them + // compatible. + if (!ShapeUtil::CompatibleIgnoringFpPrecision( + producing_instruction->shape(), + copy_start_->operand(source_operand_index_)->shape())) { + HloInstruction* bitcast = + computation->AddInstruction(HloInstruction::CreateBitcast( + copy_start_->operand(source_operand_index_)->shape(), + producing_instruction)); + bitcast->CopyOriginalValue(producing_instruction); + producing_instruction = bitcast; + } + ABSL_RETURN_IF_ERROR(copy_start_->ReplaceOperandWith(source_operand_index_, + producing_instruction)); + } else { + Shape dest_shape = shape; + if (memory_space() == MemorySpace::kDefault) { + dest_shape.mutable_layout()->clear_split_configs(); + } + copy_start_ = computation->AddInstruction(HloInstruction::CreateCopyStart( + ShapeUtil::MakeTupleShape( + {dest_shape, shape, ShapeUtil::MakeShape(U32, {})}), + producing_instruction, cross_program_prefetch_index())); + copy_done_ = computation->AddInstruction(HloInstruction::CreateUnary( + dest_shape, HloOpcode::kCopyDone, copy_start_)); + } + copy_done_->CopyOriginalValue(producing_instruction); + VLOG(4) << "Created " << copy_start_->name() + << " for copy allocation: " << ToString(); + + // Update the allocation position with the copy complete instruction, so that + // if there are further copies from it, they can find the correct position. + set_original_defining_position(HloPosition{copy_done_, {}}); + return UpdateUses(computation, copy_done_, bitcast_split_fn, hlo_live_range, + alias_analysis); +} + +void CopyAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void CopyAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + if (needed_allocations.insert(this).second) { + prev_allocation_.MarkNeeded(needed_allocations); + } +} + +std::string CopyAllocation::ToString() const { + std::string memory_space_str = MemorySpaceToString(memory_space()); + std::optional chunk = maybe_chunk(); + if (chunk) { + absl::StrAppend(&memory_space_str, " (off: ", chunk->offset, + ", size: ", chunk->size, ")"); + } + return absl::StrCat( + "Copy Allocation in ", memory_space_str, ", start_time:", start_time(), + ", end_time:", end_time(), + ", copy_start_after_time: ", copy_start_schedule_after(), + ", copy_done_before_time: ", copy_done_schedule_before(), + ", uses: ", UsesToString(uses()), + ", existing sync_mem_op: ", sync_mem_op_ ? sync_mem_op_->name() : "none", + ", existing copy_start: ", copy_start_ ? copy_start_->name() : "none", + ", existing copy_done: ", copy_done_ ? copy_done_->name() : "none", + ", from ", prev_allocation_.ToString()); +} + +HloPosition CopyAllocation::defining_position() const { + // Unless explicitly set, the defining position of a copy allocation is + // retrieved from the previous allocation. This is because we don't create + // new CopyStart/CopyDone instructions until later and the position should + // point to the previous (copy or otherwise) allocation's position for the + // original defining position. + HloPosition defining_position = original_defining_position(); + if (defining_position.instruction == nullptr) { + return prev_allocation_.defining_position(); + } + return defining_position; +} + +bool CopyAllocation::operator==(const CopyAllocation& other) const { + return this->base_is_equal(static_cast(other)) && + copy_done_schedule_before() == other.copy_done_schedule_before() && + copy_start_schedule_after() == other.copy_start_schedule_after() && + copy_start() == other.copy_start() && copy_done() == other.copy_done(); +} + +bool CopyAllocation::operator==(const Allocation& other) const { + const CopyAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +void CopyAllocation::set_copy_start_schedule_after( + int64_t copy_start_schedule_after) { + copy_start_schedule_after_ = copy_start_schedule_after; +} + +void CopyAllocation::set_copy_done_schedule_before( + int64_t copy_done_schedule_before) { + copy_done_schedule_before_ = copy_done_schedule_before; +} + +int64_t CopyAllocation::copy_start_schedule_after() const { + return copy_start_schedule_after_; +} + +int64_t CopyAllocation::copy_done_schedule_before() const { + return copy_done_schedule_before_; +} + +SlicedCopyAllocation::SlicedCopyAllocation( + const Allocation& prev_allocation, MemorySpace memory_space, + std::vector slice_decisions_sorted_by_exclusive_start_time, + int64_t copy_done_schedule_before_time, int64_t end_time, + const SlicedPrefetchOptions& sliced_prefetch_options, + absl::FunctionRef get_equivalent_s8_shape_fn, + HloInstruction* sync_mem_op) + : Allocation( + /*defining_position=*/{nullptr, {}}, memory_space, + GetSlicedCopyAllocationChunk( + slice_decisions_sorted_by_exclusive_start_time), + // Allocation uses an inclusive start time + ExclusiveToInclusiveStartTime( + GetSlicedCopyAllocationExclusiveStartTime( + slice_decisions_sorted_by_exclusive_start_time)), + end_time, + /*cross_program_prefetch_index=*/std::nullopt), + original_shape_to_slice_(prev_allocation.defining_position().shape()), + prev_allocation_(prev_allocation), + sliced_prefetch_options_(sliced_prefetch_options), + get_equivalent_s8_shape_fn_(get_equivalent_s8_shape_fn), + sync_mem_op_(sync_mem_op) { + CHECK_GE(slice_decisions_sorted_by_exclusive_start_time.size(), 2); + slice_details_sorted_by_exclusive_start_time_.reserve( + slice_decisions_sorted_by_exclusive_start_time.size()); + for (SliceDecision& decision : + slice_decisions_sorted_by_exclusive_start_time) { + int64_t copy_done_schedule_after_time = decision.exclusive_start_time; + slice_details_sorted_by_exclusive_start_time_.push_back(SliceDetail{ + std::move(decision), + copy_done_schedule_after_time, + copy_done_schedule_before_time, + /*copy_start=*/nullptr, + /*copy_done=*/nullptr, + }); + } +} + +absl::Status SlicedCopyAllocation::Process( + const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + Shape shape = defining_position().shape(); + HloInstruction* producing_instruction = AddGetTupleElements(); + + // Calling Process() over the previous allocation might + // have modified the defining position, and hence the shape that was used when + // we computed the slices. In cases where the shape has changed, we insert a + // bitcast, so slice instructions operate on the originally sliced shape. + // + // Note, these bitcasts are being inserted in the same cases that + // UpdateUses() is inserting bitcasts, except we are + // inserting the bitcasts before the copy, instead of after the copy. + if (!Shape::Equal().IgnoreMemorySpaceInLayout()(shape, + original_shape_to_slice_)) { + int64_t new_memory_space = shape.layout().memory_space(); + shape = original_shape_to_slice_; + shape.mutable_layout()->set_memory_space(new_memory_space); + producing_instruction = producing_instruction->parent()->AddInstruction( + HloInstruction::CreateBitcast(shape, producing_instruction)); + } + + HloComputation* computation = producing_instruction->parent(); + std::vector slice_dones; + slice_dones.reserve(slice_details_sorted_by_exclusive_start_time_.size()); + + // If we are trying to make all slices a uniform size, we bitcast the + // producing instruction to an array of bytes, so it is easy to slice into any + // size. + Shape slice_shape = shape; + if (IsUniformSliceSizingEnabled(sliced_prefetch_options_)) { + slice_shape = get_equivalent_s8_shape_fn_(shape); + producing_instruction = producing_instruction->parent()->AddInstruction( + HloInstruction::CreateBitcast(slice_shape, producing_instruction)); + } + + // Sliced copy allocations need to insert asynchronous copy nodes. + for (SliceDetail& slice_detail : + slice_details_sorted_by_exclusive_start_time_) { + ABSL_RETURN_IF_ERROR(slice_detail.CreateAsyncSlice( + slice_shape, *producing_instruction, *computation)); + VLOG(4) << "Created " << slice_detail.copy_start->name() + << " for sliced copy allocation: " << ToString(); + slice_dones.push_back(slice_detail.copy_done); + } + + ABSL_RETURN_IF_ERROR(CreateBitcastConcat(shape, slice_dones)); + + // If we bitcast to an array of bytes above, the result of the concatenated + // slices will also be an array of bytes. Thus, we need to cast the + // concatenation back to the original shape. + if (IsUniformSliceSizingEnabled(sliced_prefetch_options_)) { + concat_ = concat_->parent()->AddInstruction( + HloInstruction::CreateBitcast(shape, concat_)); + } + + // Update the allocation position with the copy complete instruction, so that + // if there are further copies from it, they can find the correct position. + set_original_defining_position(HloPosition{concat_, {}}); + return UpdateUses(computation, concat_, bitcast_split_fn, hlo_live_range, + alias_analysis); +} + +void SlicedCopyAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void SlicedCopyAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + if (needed_allocations.insert(this).second) { + prev_allocation_.MarkNeeded(needed_allocations); + } +} + +HloPosition SlicedCopyAllocation::defining_position() const { + // Unless explicitly set, the defining position of a sliced copy allocation is + // retrieved from the previous allocation. This is because we don't create + // new CopyStart/CopyDone instructions until later and the position should + // point to the previous (copy or otherwise) allocation's position for the + // original defining position. + HloPosition defining_position = original_defining_position(); + if (defining_position.instruction == nullptr) { + return prev_allocation_.defining_position(); + } + return defining_position; +} + +int64_t SlicedCopyAllocation::earliest_available_time() const { + return slice_details_sorted_by_start_time().back().copy_done_before_time; +} + +std::vector SlicedCopyAllocation::SliceOffsetsSortedByStartTime() + const { + std::vector offsets; + offsets.reserve(slice_details_sorted_by_exclusive_start_time_.size()); + + for (const SliceDetail& slice_detail : + slice_details_sorted_by_exclusive_start_time_) { + offsets.push_back(slice_detail.slice_decision.chunk.offset); + } + + return offsets; +} + +void SlicedCopyAllocation::AddDiffToAllSliceOffsets(int64_t diff) { + for (SliceDetail& slice_detail : + slice_details_sorted_by_exclusive_start_time_) { + HeapSimulator::Chunk& chunk = slice_detail.slice_decision.chunk; + chunk = + HeapSimulator::Chunk::FromOffsetSize(chunk.offset + diff, chunk.size); + } +} + +void SlicedCopyAllocation::ImportRepackedSliceData( + const SlicedAllocationData& data) { + int num_slices = slice_details_sorted_by_exclusive_start_time_.size(); + CHECK_EQ(data.slices_sorted_by_offset.size(), num_slices); + + std::vector slice_details_sorted_by_offset; + slice_details_sorted_by_offset.reserve(num_slices); + for (SliceDetail& slice_detail : + slice_details_sorted_by_exclusive_start_time_) { + slice_details_sorted_by_offset.push_back(&slice_detail); + } + absl::c_sort(slice_details_sorted_by_offset, [](const SliceDetail* lhs, + const SliceDetail* rhs) { + return lhs->slice_decision.chunk.offset < rhs->slice_decision.chunk.offset; + }); + + for (int i = 0; i < num_slices; ++i) { + SliceDetail* slice_detail = slice_details_sorted_by_offset[i]; + HeapSimulator::Chunk& chunk = slice_detail->slice_decision.chunk; + const AllocatedSlice& repacked_slice_data = data.slices_sorted_by_offset[i]; + chunk = HeapSimulator::Chunk::FromOffsetSize(repacked_slice_data.offset, + chunk.size); + slice_detail->copy_start_after_time = + repacked_slice_data.inclusive_start_time - 1; + slice_detail->slice_decision.exclusive_start_time = + InclusiveToExclusiveStartTime(repacked_slice_data.inclusive_start_time); + } + + absl::c_sort(slice_details_sorted_by_exclusive_start_time_, + [](const SliceDetail& lhs, const SliceDetail& rhs) { + return std::make_tuple(lhs.copy_start_after_time, + lhs.slice_decision.chunk.offset) < + std::make_tuple(rhs.copy_start_after_time, + rhs.slice_decision.chunk.offset); + }); +} + +const std::vector& +SlicedCopyAllocation::slice_details_sorted_by_start_time() const { + return slice_details_sorted_by_exclusive_start_time_; +} + +std::vector& +SlicedCopyAllocation::mutable_slice_details_sorted_by_start_time() { + return slice_details_sorted_by_exclusive_start_time_; +} + +bool SlicedCopyAllocation::operator==(const SlicedCopyAllocation& other) const { + return this->base_is_equal(static_cast(other)) && + slice_details_sorted_by_exclusive_start_time_ == + other.slice_details_sorted_by_exclusive_start_time_ && + concat_ == other.concat_; +} + +std::string SlicedCopyAllocation::ToString() const { + std::string memory_space_str = MemorySpaceToString(memory_space()); + if (memory_space() == MemorySpace::kAlternate) { + absl::StrAppend(&memory_space_str, " (off: ", maybe_chunk()->offset, ")"); + } + return absl::StrCat( + "Sliced Copy Allocation in ", memory_space_str, + ", start_time:", start_time(), ", end_time:", end_time(), + ", first_slice_copy_start_after_time: ", + slice_details_sorted_by_start_time().front().copy_start_after_time, + ", last_slice_copy_done_before_time: ", + slice_details_sorted_by_start_time().back().copy_done_before_time, + ", uses: ", UsesToString(uses()), + ", sync_mem_op: ", sync_mem_op_ ? sync_mem_op_->name() : "none", + ", from ", prev_allocation_.ToString()); +} + +absl::Status SlicedCopyAllocation::CreateBitcastConcat( + const Shape& shape, absl::Span slices) { + CHECK(!slices.empty()); + concat_ = + slices.front()->parent()->AddInstruction(HloInstruction::CreateCustomCall( + shape, slices, + xla::memory_space_assignment::kConcatBitcastCustomCall)); + return absl::OkStatus(); +} + +std::string SlicedCopyAllocation::SliceDetail::ToString() const { + return absl::StrCat("{ slice_decision: ", slice_decision.ToString(), + ", copy_start_after_time: ", copy_start_after_time, + ", copy_done_before_time: ", copy_done_before_time, " }"); +} + +std::tuple +SliceDetailToTuple(const SlicedCopyAllocation::SliceDetail& slice_detail) { + return std::make_tuple(std::ref(slice_detail.slice_decision), + slice_detail.copy_start_after_time, + slice_detail.copy_done_before_time, + slice_detail.copy_start, slice_detail.copy_done); +} + +bool SlicedCopyAllocation::SliceDetail::operator==( + const SliceDetail& other) const { + return SliceDetailToTuple(*this) == SliceDetailToTuple(other); +} + +absl::Status SlicedCopyAllocation::SliceDetail::CreateAsyncSlice( + const Shape& original_shape, HloInstruction& producer, + HloComputation& parent) { + if (original_shape.dimensions().size() != + slice_decision.sizing.slice_params.size()) { + return FailedPrecondition( + "%s", absl::StrCat("The number of SlicedCopyAllocation parameters ", + slice_decision.sizing.slice_params.size(), + " does not match the rank ", + original_shape.dimensions().size(), + " of the tensor we are slicing.")); + } + + std::vector start_indices; + start_indices.reserve(slice_decision.sizing.slice_params.size()); + std::vector limit_indices; + limit_indices.reserve(slice_decision.sizing.slice_params.size()); + std::vector strides; + strides.reserve(slice_decision.sizing.slice_params.size()); + + for (int i = 0; i < slice_decision.sizing.slice_params.size(); ++i) { + const SliceParam& slice_param = slice_decision.sizing.slice_params[i]; + start_indices.push_back(slice_param.start_inclusive); + limit_indices.push_back(slice_param.end_exclusive); + strides.push_back(1); + const int64_t new_dim = + slice_param.end_exclusive - slice_param.start_inclusive; + if (new_dim <= 0) { + return FailedPrecondition( + "%s", absl::StrCat("SlicedCopyAllocation new dimension size is ", + new_dim, ", expected something > 0.")); + } + if (original_shape.dimensions(i) < new_dim) { + return FailedPrecondition( + "%s", + absl::StrCat("SlicedCopyAllocation sliced dimension size ", new_dim, + " is bigger than its original dimension size of ", + original_shape.dimensions(i), ".")); + } + } + + HloInstruction* slice = parent.AddInstruction( + HloInstruction::CreateSlice(slice_decision.sizing.slice_shape, &producer, + start_indices, limit_indices, strides)); + ABSL_ASSIGN_OR_RETURN(copy_done, parent.CreateAsyncInstructions( + slice, {ShapeUtil::MakeShape(S32, {})})); + copy_start = copy_done->mutable_operand(0); + + return absl::OkStatus(); +} + +bool SlicedCopyAllocation::operator==(const Allocation& other) const { + const SlicedCopyAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +MirroredAllocation::MirroredAllocation(const Allocation& original_allocation, + int64_t time) + : Allocation(original_allocation.defining_position(), + original_allocation.memory_space(), + original_allocation.maybe_chunk(), + /*start_time=*/time, + /*end_time=*/time, + /*cross_program_prefetch_index=*/std::nullopt), + defining_position_(std::nullopt), + original_allocation_(original_allocation) {} + +MirroredAllocation::MirroredAllocation(HloPosition defining_position, + const Allocation& original_allocation, + int64_t start_time, int64_t end_time) + : Allocation(original_allocation.defining_position(), + original_allocation.memory_space(), + original_allocation.maybe_chunk(), + /*start_time=*/start_time, + /*end_time=*/end_time, + /*cross_program_prefetch_index=*/std::nullopt), + defining_position_(defining_position), + original_allocation_(original_allocation) {} + +HloPosition MirroredAllocation::defining_position() const { + if (defining_position_.has_value()) { + return defining_position_.value(); + } + return original_allocation_.defining_position(); +} + +std::string MirroredAllocation::ToString() const { + std::string memory_space_str = MemorySpaceToString(memory_space()); + std::optional chunk = maybe_chunk(); + if (chunk) { + absl::StrAppend(&memory_space_str, " (off: ", chunk->offset, + ", size: ", chunk->size, ")"); + } + return absl::StrCat( + "MirroredAllocation in ", memory_space_str, " defined at ", + defining_position().ToString(), ", start_time:", start_time(), + ", end_time:", end_time(), ", uses: ", UsesToString(uses()), + ", for original allocation: ", original_allocation_.ToString()); +} + +absl::Status MirroredAllocation::Process( + const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + set_original_defining_position(defining_position()); + HloInstruction* producing_instruction = AddGetTupleElements(); + HloComputation* computation = producing_instruction->parent(); + return UpdateUses(computation, producing_instruction, bitcast_split_fn, + hlo_live_range, alias_analysis); +} + +void MirroredAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void MirroredAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + if (needed_allocations.insert(this).second) { + original_allocation_.MarkNeeded(needed_allocations); + } +} + +bool MirroredAllocation::operator==(const Allocation& other) const { + const MirroredAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other) && + defining_position() == other.defining_position(); +} + +ParentAllocation::ParentAllocation(const Allocation& original_allocation, + HloInstruction* calling_instruction, + HloPosition position, int64_t time) + : Allocation(std::move(position), MemorySpace::kDefault, + original_allocation.maybe_chunk(), + /*start_time=*/time, + /*end_time=*/time, + /*cross_program_prefetch_index=*/std::nullopt), + original_allocation_(original_allocation), + calling_instruction_(calling_instruction) {} + +HloPosition ParentAllocation::defining_position() const { + return original_defining_position(); +} + +std::string ParentAllocation::ToString() const { + return absl::StrCat("Parent Allocation mirrored at ", + original_defining_position().ToString(), ", originally ", + original_allocation_.ToString()); +} + +absl::Status ParentAllocation::Process(const BitcastSplitFn& bitcast_split_fn, + const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + // Add an additional parameter to the while HLO with a reference to the buffer + // in the default memory space. + HloInstruction* producing_instruction = + original_allocation_.AddGetTupleElements(); + int new_tuple_index = calling_instruction_->shape().tuple_shapes().size(); + + HloInstruction* old_while_operand = calling_instruction_->mutable_operand(0); + ABSL_ASSIGN_OR_RETURN( + HloInstruction * new_while_operand, + TupleUtil::ReplaceTupleWith(producing_instruction, old_while_operand, + {new_tuple_index})); + AppendToTupleOriginalValue(old_while_operand, new_while_operand, + producing_instruction); + // Also replace the while op with a tuple that has the old shape. Note that we + // need to first take a snapshot of the users before calling ExtractPrefix + // since ExtractPrefix introduces additional gte users. + std::vector while_users = calling_instruction_->users(); + HloInstruction* tuple_with_old_shape = + TupleUtil::ExtractPrefix(calling_instruction_, new_tuple_index); + if (std::shared_ptr src_ov = + calling_instruction_->original_value(); + src_ov != nullptr) { + tuple_with_old_shape->CopyOriginalValue(calling_instruction_, + /*clone=*/true); + const auto& src_tree = src_ov->tree(); + for (int64_t i = 0; i < new_tuple_index; ++i) { + HloInstruction* gte = tuple_with_old_shape->mutable_operand(i); + if (src_tree.find({i}) != src_tree.end()) { + auto gte_ov = std::make_shared(gte->shape()); + CHECK_OK(gte_ov->mutable_tree()->CopyCompatibleSubtreeFrom(src_tree, + {i}, {})); + if (!gte_ov->IsEmpty()) { + gte->set_original_value(std::move(gte_ov)); + } + } + } + } + ABSL_RETURN_IF_ERROR(calling_instruction_->ReplaceOperandWithDifferentShape( + 0, new_while_operand)); + *calling_instruction_->mutable_shape() = new_while_operand->shape(); + *calling_instruction_->while_condition() + ->parameter_instruction(0) + ->mutable_shape() = new_while_operand->shape(); + *calling_instruction_->while_body() + ->parameter_instruction(0) + ->mutable_shape() = new_while_operand->shape(); + ABSL_RETURN_IF_ERROR(calling_instruction_->ReplaceAllUsesWithDifferentShape( + while_users, tuple_with_old_shape)); + + AppendToWhileLoopOriginalValue(calling_instruction_, {producing_instruction}); + HloPosition defining_position = original_defining_position(); + defining_position.index = {new_tuple_index}; + set_original_defining_position(defining_position); + + HloInstruction* final_instruction = AddGetTupleElements(); + HloComputation* computation = final_instruction->parent(); + return UpdateUses(computation, final_instruction, bitcast_split_fn, + hlo_live_range, alias_analysis); +} + +absl::Status ParentAllocation::PostProcess() { + // Update the root of the while body with the new parameter. The reason why we + // need a separate post-process for this is because other allocations may have + // while body root as a use, so they would update the old root instead of the + // new root. Doing the post-process step later ensures the root has been + // updated with other changes, and we can safely add the additional parameter. + HloComputation* while_body = calling_instruction_->while_body(); + HloInstruction* old_body_root = while_body->root_instruction(); + HloInstruction* added_element = AddGetTupleElements(); + ABSL_ASSIGN_OR_RETURN( + HloInstruction * new_while_body_root, + TupleUtil::ReplaceTupleWith(added_element, old_body_root, + original_defining_position().index)); + while_body->set_root_instruction(new_while_body_root, + /*accept_different_shape=*/true); + AppendToTupleOriginalValue(old_body_root, new_while_body_root, added_element); + return absl::OkStatus(); +} + +void ParentAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + // Parent allocations are only needed if they have any uses or if there is a + // copy allocation that copies this value (in that case, the copy allocation + // will call this allocation's MarkNeeded function). + if (!has_no_uses()) { + MarkNeeded(needed_allocations); + } +} + +void ParentAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + if (needed_allocations.insert(this).second) { + original_allocation_.MarkNeeded(needed_allocations); + } +} + +bool ParentAllocation::operator==(const Allocation& other) const { + const ParentAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +WindowPrefetchedAllocation::WindowPrefetchedAllocation( + Allocation& prev_allocation, HloUse use, const HeapSimulator::Chunk& chunk, + int64_t prefetch_start_schedule_after_time, + int64_t prefetch_done_schedule_before_time, const Options& options) + : Allocation( + {nullptr, {}}, MemorySpace::kAlternate, chunk, + ExclusiveToInclusiveStartTime(prefetch_start_schedule_after_time), + InclusiveToExclusiveEndTime(prefetch_done_schedule_before_time), + /*cross_program_prefetch_index=*/std::nullopt), + options_(options), + defining_position_(prev_allocation.defining_position()), + use_(use), + prefetch_start_schedule_after_(prefetch_start_schedule_after_time), + prefetch_done_schedule_before_(prefetch_done_schedule_before_time), + bytes_(chunk.size) {} + +HloPosition WindowPrefetchedAllocation::defining_position() const { + if (defining_position_.instruction != nullptr) { + return defining_position_; + } + return original_defining_position(); +} + +int64_t WindowPrefetchedAllocation::earliest_available_time() const { + return prefetch_done_schedule_before_; +} + +absl::Status WindowPrefetchedAllocation::InsertWindowPrefetchInstruction( + HloInstruction* producing_instruction, HloInstruction* use_instruction, + HloComputation* computation) { + // Derive the shape for window buffer. + Shape buffer_shape = ShapeUtil::MakeShape(U8, {options_.bytes}); + Layout layout = LayoutUtil::MakeLayout({0}); + layout.set_memory_space(options_.alternate_memory_space); + *buffer_shape.mutable_layout() = layout; + // Sync flag shape + Shape sflag_shape = ShapeUtil::MakeShape(S32, {2}); + // Output shape of the WindowPrefetch op. + Shape output_shape = ShapeUtil::MakeTupleShape({buffer_shape, sflag_shape}); + + // Insert WindowPrefetch op. + HloInstruction* custom_call = + computation->AddInstruction(HloInstruction::CreateCustomCall( + output_shape, {producing_instruction}, "WindowPrefetch")); + HloInstruction* get_buffer = computation->AddInstruction( + HloInstruction::CreateGetTupleElement(buffer_shape, custom_call, 0)); + HloInstruction* get_sflag = computation->AddInstruction( + HloInstruction::CreateGetTupleElement(sflag_shape, custom_call, 1)); + use_instruction->AppendOperand(get_buffer); + use_instruction->AppendOperand(get_sflag); + + // Add backend config for the WindowPrefetch op. + int64_t parameter_index = + use_instruction->operand_index(producing_instruction); + int64_t buffer_index = use_instruction->operand_index(get_buffer); + custom_call->set_raw_backend_config_string(absl::StrCat( + "parameter_index=", parameter_index, " buffer_index=", buffer_index)); + + // The buffer's defining position is the get_tuple_element instruction. + prefetch_instruction_ = get_buffer; + return absl::OkStatus(); +} + +absl::Status WindowPrefetchedAllocation::Process( + const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + HloInstruction* use_instruction = use_.instruction; + int64_t use_operand = use_instruction->operand_count(); + HloInstruction* producing_instruction = + use_instruction->mutable_operand(use_.operand_number); + HloComputation* computation = producing_instruction->parent(); + CHECK_EQ(use_instruction->opcode(), HloOpcode::kFusion); + + ABSL_RETURN_IF_ERROR(InsertWindowPrefetchInstruction( + producing_instruction, use_instruction, computation)); + + // Notify the backend that an operand has been appended as a window prefetch + // buffer. + options_.notify_operand_appended_fn(use_instruction, use_.operand_number, + use_operand); + + // Set the original defining position to the window prefetch instruction. + set_original_defining_position(defining_position_); + defining_position_ = {prefetch_instruction_, {}}; + AddUse(HloUse{use_instruction, use_operand}); + return absl::OkStatus(); +} + +void WindowPrefetchedAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void WindowPrefetchedAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + needed_allocations.insert(this); +} + +std::string WindowPrefetchedAllocation::ToString() const { + return absl::StrCat("WindowPrefetched Allocation"); +} + +bool WindowPrefetchedAllocation::operator==( + const WindowPrefetchedAllocation& other) const { + return this->base_is_equal(static_cast(other)) && + prefetch_done_schedule_before() == + other.prefetch_done_schedule_before() && + prefetch_start_schedule_after() == + other.prefetch_start_schedule_after() && + prefetch() == other.prefetch() && bytes_ == other.bytes_; +} + +bool WindowPrefetchedAllocation::operator==(const Allocation& other) const { + const WindowPrefetchedAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +ScopedAllocation::ScopedAllocation(HeapSimulator::Chunk chunk, + int64_t allocation_time, + HloInstruction* defining_instruction, + bool is_post_module) + : Allocation(/*defining_position=*/{defining_instruction, {}}, + MemorySpace::kAlternate, chunk, + /*start_time=*/allocation_time, + /*end_time=*/allocation_time, + /*cross_program_prefetch_index=*/std::nullopt), + is_post_module_(is_post_module) {} + +HloPosition ScopedAllocation::defining_position() const { + return original_defining_position(); +} + +absl::Status ScopedAllocation::Process(const BitcastSplitFn& bitcast_split_fn, + const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + return absl::OkStatus(); +} + +void ScopedAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void ScopedAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + needed_allocations.insert(this); +} + +std::string ScopedAllocation::ToString() const { + std::string name = ""; + if (!is_post_module_) { + name = defining_position().instruction->name(); + } + return absl::StrCat("Scoped Allocation defined for ", name, + ", allocation_time: ", start_time(), + ", chunk: ", chunk().ToString()); +} + +bool ScopedAllocation::operator==(const Allocation& other) const { + const ScopedAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +bool ScopedAllocation::operator==(const ScopedAllocation& other) const { + return this->base_is_equal(static_cast(other)) && + is_post_module_ == other.is_post_module_; +} + +std::tuple GetAllocationSortTuple( + const std::unique_ptr& allocation) { + int64_t scheduled_on_or_before = allocation->start_time(); + int64_t scheduled_on_or_after = allocation->start_time(); + if (allocation->is_copy_allocation()) { + auto copy_allocation = absl::down_cast(allocation.get()); + scheduled_on_or_before = copy_allocation->copy_done_schedule_before(); + scheduled_on_or_after = copy_allocation->copy_start_schedule_after(); + } + return std::forward_as_tuple(scheduled_on_or_before, + !allocation->is_copy_allocation(), + scheduled_on_or_after); +} + +void SortAllocationSequence(AllocationSequence& allocations) { + absl::c_sort(allocations, [](const std::unique_ptr& lhs, + const std::unique_ptr& rhs) { + return GetAllocationSortTuple(lhs) < GetAllocationSortTuple(rhs); + }); +} + +std::string AllocationSequenceToString(AllocationSequence& allocations, + bool sort_allocations) { + if (sort_allocations) { + SortAllocationSequence(allocations); + } + std::string allocations_str = "\n"; + for (const std::unique_ptr& allocation : allocations) { + absl::StrAppend(&allocations_str, allocation->ToString(), "\n"); + } + return allocations_str; +} + +std::vector GetAllocationSequenceInRawPointers( + AllocationSequence& allocations) { + std::vector allocations_in_raw_pointers; + for (const std::unique_ptr& allocation : allocations) { + allocations_in_raw_pointers.push_back(allocation.get()); + } + return allocations_in_raw_pointers; +} + +namespace { + +struct AllocationSummary { + static void Add(const Allocation& allocation, + std::vector& data) { + if (!allocation.is_sliced_copy_allocation()) { + data.push_back( + AllocationSummary{allocation.chunk(), allocation.start_time(), + allocation.end_time(), GetName(allocation)}); + return; + } + const SlicedCopyAllocation& sliced_copy_allocation = + dynamic_cast(allocation); + for (size_t i = 0; + i < sliced_copy_allocation.slice_details_sorted_by_start_time().size(); + ++i) { + const SlicedCopyAllocation::SliceDetail& slice_detail = + sliced_copy_allocation.slice_details_sorted_by_start_time()[i]; + data.push_back(AllocationSummary{ + slice_detail.slice_decision.chunk, + ExclusiveToInclusiveStartTime( + slice_detail.slice_decision.exclusive_start_time), + sliced_copy_allocation.end_time(), GetName(allocation, i)}); + } + } + + static std::string GetName(const Allocation& allocation, + std::optional slice_index = std::nullopt) { + std::vector attributes; + if (slice_index.has_value()) { + attributes.push_back(absl::StrCat("slice ", *slice_index)); + } + if (allocation.is_scoped_allocation()) { + attributes.push_back("scoped"); + } + if (allocation.cross_program_prefetch_index().has_value()) { + attributes.push_back("cross-program prefetch"); + } + if (attributes.empty()) { + return allocation.defining_position().ToString(); + } + return absl::StrCat(allocation.defining_position().ToString(), " (", + absl::StrJoin(attributes, ", "), ")"); + } + + HeapSimulator::Chunk chunk; + int64_t start_time_inclusive; + int64_t end_time_inclusive; + std::string name; +}; + +} // namespace + +void AllocationSequenceDebugging::LogAltMemAllocationsAt( + const AllocationSequence& allocations, int64_t time) { + std::vector data_vector; + for (const std::unique_ptr& allocation : allocations) { + if (allocation->start_time() <= time && allocation->end_time() >= time && + allocation->is_in_alternate_mem()) { + AllocationSummary::Add(*allocation, data_vector); + } + } + absl::c_sort(data_vector, + [](const AllocationSummary& a, const AllocationSummary& b) { + if (a.chunk.offset != b.chunk.offset) { + return a.chunk.offset < b.chunk.offset; + } + return a.start_time_inclusive < b.start_time_inclusive; + }); + LOG(INFO) << "Live allocations in alternate mem at instruction time " << time + << " (before MSA alters the graph):"; + for (const AllocationSummary& data : data_vector) { + LOG(INFO) << "Alt mem allocation in chunk " << data.chunk.ToString() + << " during [" << data.start_time_inclusive << "," + << data.end_time_inclusive << "], holding " << data.name; + } +} + +} // namespace xla::memory_space_assignment diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_test.cc b/third_party/xla/xla/service/memory_space_assignment/allocation_test.cc new file mode 100644 index 00000000000000..4f8339e21c078f --- /dev/null +++ b/third_party/xla/xla/service/memory_space_assignment/allocation_test.cc @@ -0,0 +1,393 @@ +/* Copyright 2024 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/service/memory_space_assignment/allocation.h" + +#include +#include +#include + +#include +#include +#include "absl/strings/string_view.h" +#include "xla/hlo/analysis/alias_info.h" +#include "xla/hlo/analysis/hlo_alias_analysis.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_schedule.h" +#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" +#include "xla/hlo/utils/hlo_live_range.h" +#include "xla/service/heap_simulator/heap_simulator.h" +#include "xla/service/hlo_value.h" +#include "xla/shape_util.h" +#include "xla/xla_data.pb.h" + +namespace xla::memory_space_assignment { +namespace { + +class AllocationTest : public HloHardwareIndependentTestBase { + protected: + void RunAnalysis(HloModule* module, + const std::vector& inst_names, + std::unique_ptr& live_range, + std::unique_ptr& alias_analysis) { + HloSchedule schedule(module); + HloInstructionSequence sequence; + for (auto name : inst_names) { + sequence.push_back(FindInstruction(module, name)); + } + schedule.set_sequence(module->entry_computation(), sequence); + + AliasInfo alias_info; + ASSERT_OK_AND_ASSIGN(alias_analysis, + HloAliasAnalysis::Run(module, &alias_info)); + ASSERT_OK_AND_ASSIGN(live_range, + HloLiveRange::Run(schedule, *alias_analysis, + module->entry_computation())); + } +}; + +TEST_F(AllocationTest, CopyAllocationProcessSimple) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + p0 = f32[2,3]{1,0} parameter(0) + p1 = f32[2,3]{1,0} parameter(1) + p1_negate = f32[2,3]{1,0} negate(p1) + add = f32[2,3]{1,0} add(p0, p1_negate) + ROOT tuple = tuple(add, p0) +} + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + std::unique_ptr hlo_live_range; + std::unique_ptr alias_analysis; + RunAnalysis(module.get(), {"p0", "p1", "p1_negate", "add", "tuple"}, + hlo_live_range, alias_analysis); + // HloComputation* computation = module->entry_computation(); + HloInstruction* add = FindInstruction(module.get(), "add"); + HloInstruction* p1_negate = FindInstruction(module.get(), "p1_negate"); + + HeapSimulator::Chunk p1_negate_chunk = + HeapSimulator::Chunk::FromOffsetSize(0, 24); + + PinnedAllocation p1_negate_pinned(HloPosition{p1_negate, {}}, + MemorySpace::kDefault, p1_negate_chunk, + /*start_time=*/0, + /*end_time=*/5); + CopyAllocation copy_allocation(p1_negate_pinned, MemorySpace::kAlternate, + std::nullopt, + /*copy_start_schedule_after_time=*/2, + /*copy_done_schedule_before_time=*/3, + /*end_time=*/5, std::nullopt, + /*sync_mem_op=*/nullptr); + + // Use the correct instruction and operand numbers for the add instruction + copy_allocation.AddUse(HloUse{add, 1}); // Use of p1_negate in add + BitcastSplitFn split_fn = nullptr; + ASSERT_OK( + copy_allocation.Process(split_fn, *hlo_live_range, *alias_analysis)); + + // Check copy_start and copy_done instructions. + HloInstruction* copy_start = copy_allocation.copy_start(); + ASSERT_NE(copy_start, nullptr); + EXPECT_EQ(copy_start->opcode(), HloOpcode::kCopyStart); + EXPECT_EQ(copy_start->operand(0), p1_negate); + + HloInstruction* copy_done = copy_allocation.copy_done(); + ASSERT_NE(copy_done, nullptr); + EXPECT_EQ(copy_done->opcode(), HloOpcode::kCopyDone); + EXPECT_EQ(copy_done->operand(0), copy_start); + + // Check that uses are updated. + EXPECT_EQ(add->operand(1), copy_done); + + // Check defining position + EXPECT_EQ(copy_allocation.defining_position().instruction, copy_done); +} + +TEST_F(AllocationTest, EvictedSplitShape) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + p0 = f32[2,3]{1,0} parameter(0) + p1 = f32[2,3]{1,0} parameter(1) + p1_negate = f32[2,3]{1,0:S(1)SC(0:1)} negate(p1) + add = f32[2,3]{1,0} add(p0, p1_negate) + ROOT tuple = tuple(add, p0) +} + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + std::unique_ptr hlo_live_range; + std::unique_ptr alias_analysis; + RunAnalysis(module.get(), {"p0", "p1", "p1_negate", "add", "tuple"}, + hlo_live_range, alias_analysis); + // HloComputation* computation = module->entry_computation(); + HloInstruction* add = FindInstruction(module.get(), "add"); + HloInstruction* p1_negate = FindInstruction(module.get(), "p1_negate"); + + HeapSimulator::Chunk p1_negate_chunk = + HeapSimulator::Chunk::FromOffsetSize(0, 24); + + PinnedAllocation p1_negate_pinned(HloPosition{p1_negate, {}}, + MemorySpace::kAlternate, p1_negate_chunk, + /*start_time=*/0, + /*end_time=*/5); + CopyAllocation copy_allocation(p1_negate_pinned, MemorySpace::kDefault, + std::nullopt, + /*copy_start_schedule_after_time=*/2, + /*copy_done_schedule_before_time=*/3, + /*end_time=*/5, std::nullopt, + /*sync_mem_op=*/nullptr); + + // Use the correct instruction and operand numbers for the add instruction + copy_allocation.AddUse(HloUse{add, 1}); // Use of p1_negate in add + BitcastSplitFn split_fn = nullptr; + ASSERT_OK( + copy_allocation.Process(split_fn, *hlo_live_range, *alias_analysis)); + + // Check copy_start and copy_done instructions. + HloInstruction* copy_start = copy_allocation.copy_start(); + ASSERT_NE(copy_start, nullptr); + EXPECT_EQ(copy_start->opcode(), HloOpcode::kCopyStart); + EXPECT_EQ(copy_start->operand(0), p1_negate); + + HloInstruction* copy_done = copy_allocation.copy_done(); + ASSERT_NE(copy_done, nullptr); + EXPECT_EQ(copy_done->opcode(), HloOpcode::kCopyDone); + EXPECT_EQ(copy_done->operand(0), copy_start); + EXPECT_EQ(copy_done->shape().layout().split_configs().size(), 0); + + // Check that uses are updated. + EXPECT_EQ(add->operand(1), copy_done); + + // Check defining position + EXPECT_EQ(copy_allocation.defining_position().instruction, copy_done); +} + +TEST_F(AllocationTest, CopyAllocationProcessReplaceSyncSlice) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + p0 = f32[1,3]{1,0} parameter(0) + p1 = f32[2,3]{1,0} parameter(1) + p1_negate = f32[2,3]{1,0} negate(p1) + slice = f32[1,3]{1,0} slice(p1_negate), slice={[0:1], [0:3]} + add = f32[1,3]{1,0} add(p0, slice) + ROOT tuple = tuple(add, p0) +} + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + std::unique_ptr hlo_live_range; + std::unique_ptr alias_analysis; + RunAnalysis(module.get(), {"p0", "p1", "p1_negate", "slice", "add", "tuple"}, + hlo_live_range, alias_analysis); + // HloComputation* computation = module->entry_computation(); + HloInstruction* add = FindInstruction(module.get(), "add"); + HloInstruction* p1_negate = FindInstruction(module.get(), "p1_negate"); + HloInstruction* slice = FindInstruction(module.get(), "slice"); + + HeapSimulator::Chunk p1_negate_chunk = + HeapSimulator::Chunk::FromOffsetSize(0, 24); + + PinnedAllocation p1_negate_pinned(HloPosition{p1_negate, {}}, + MemorySpace::kAlternate, p1_negate_chunk, + /*start_time=*/0, + /*end_time=*/5); + CopyAllocation copy_allocation(p1_negate_pinned, MemorySpace::kAlternate, + std::nullopt, + /*copy_start_schedule_after_time=*/2, + /*copy_done_schedule_before_time=*/3, + /*end_time=*/5, std::nullopt, + /*sync_mem_op=*/slice); + + // Use the correct instruction and operand numbers for the add instruction + copy_allocation.AddUse(HloUse{add, 1}); // Use of p1_negate in add + BitcastSplitFn split_fn = nullptr; + ASSERT_OK( + copy_allocation.Process(split_fn, *hlo_live_range, *alias_analysis)); + + // Check copy_start and copy_done instructions. + HloInstruction* slice_start = copy_allocation.copy_start(); + ASSERT_NE(slice_start, nullptr); + EXPECT_EQ(slice_start->opcode(), HloOpcode::kAsyncStart); + EXPECT_EQ(slice_start->operand(0), p1_negate); + + HloInstruction* slice_done = copy_allocation.copy_done(); + ASSERT_NE(slice_done, nullptr); + EXPECT_EQ(slice_done->opcode(), HloOpcode::kAsyncDone); + EXPECT_EQ(slice_done->operand(0), slice_start); + + // Check the shapes. + EXPECT_EQ(slice_done->shape(), slice->shape()); + + // Check that uses are updated. + EXPECT_EQ(add->operand(1), slice_done); + + // Check defining position + EXPECT_EQ(copy_allocation.defining_position().instruction, slice_done); +} + +TEST_F(AllocationTest, SkipTupleReconstructionForAsyncCollective) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + new_buffer = f32[2,3]{1,0} parameter(0) + cp-start = (f32[2,3]{1,0}, f32[2,3]{1,0}, u32[], u32[]) collective-permute-start(new_buffer), channel_id=1, source_target_pairs={{0,1}} + cp-done = f32[2,3]{1,0} collective-permute-done(cp-start) + ROOT tuple = tuple(cp-done) +} + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + std::unique_ptr hlo_live_range; + std::unique_ptr alias_analysis; + RunAnalysis(module.get(), {"new_buffer", "cp-start", "cp-done", "tuple"}, + hlo_live_range, alias_analysis); + + HloInstruction* cp_start = FindInstruction(module.get(), "cp-start"); + HloInstruction* cp_done = FindInstruction(module.get(), "cp-done"); + HloInstruction* new_buffer = FindInstruction(module.get(), "new_buffer"); + + HeapSimulator::Chunk chunk = HeapSimulator::Chunk::FromOffsetSize(0, 24); + + PinnedAllocation pinned(HloPosition{new_buffer, {}}, MemorySpace::kAlternate, + chunk, 0, 5); + pinned.AddUse(HloUse{cp_done, 0, {0}}); + + BitcastSplitFn split_fn = nullptr; + ASSERT_OK(pinned.Process(split_fn, *hlo_live_range, *alias_analysis)); + + EXPECT_EQ(cp_done->operand(0), cp_start); +} + +TEST_F(AllocationTest, MirroredAllocationDelegatesChunkToOriginalAllocation) { + HloComputation::Builder builder("entry"); + HloInstruction* p0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {2, 3}), "p0")); + PinnedAllocation original_allocation( + HloPosition{p0, {}}, MemorySpace::kAlternate, + HeapSimulator::Chunk::FromOffsetSize(-1, 64), + /*start_time=*/0, /*end_time=*/5); + MirroredAllocation mirrored_allocation(original_allocation, /*time=*/2); + original_allocation.set_offset(128); + + const Allocation& alloc = mirrored_allocation; + ASSERT_TRUE(alloc.maybe_chunk().has_value()); + EXPECT_EQ(alloc.maybe_chunk()->offset, 128); + EXPECT_EQ(alloc.chunk().offset, 128); +} + +TEST_F(AllocationTest, ParentAllocationPropagatesOriginalValue) { + absl::string_view hlo_string = R"( +HloModule module + +while_body { + param = (f32[2,3]{1,0}, f32[2,3]{1,0}) parameter(0), origin={({"bp0"}, {"bp1"})} + gte0 = f32[2,3]{1,0} get-tuple-element(param), index=0 + gte1 = f32[2,3]{1,0} get-tuple-element(param), index=1 + add = f32[2,3]{1,0} add(gte0, gte1) + ROOT root = (f32[2,3]{1,0}, f32[2,3]{1,0}) tuple(add, gte1), origin={({"br0"}, {"br1"})} +} + +while_condition { + param = (f32[2,3]{1,0}, f32[2,3]{1,0}) parameter(0), origin={({"cp0"}, {"cp1"})} + ROOT cond = pred[] constant(true) +} + +ENTRY entry { + p0 = f32[2,3]{1,0} parameter(0), origin={{"p0"}} + p1 = f32[2,3]{1,0} parameter(1), origin={{"p1"}} + p2 = f32[2,3]{1,0} parameter(2), origin={{"p2"}} + init = (f32[2,3]{1,0}, f32[2,3]{1,0}) tuple(p0, p1), origin={({"p0"}, {"p1"})} + while = (f32[2,3]{1,0}, f32[2,3]{1,0}) while(init), condition=while_condition, body=while_body, origin={({"w0"}, {"w1"}),["while#$"]} + entry_gte0 = f32[2,3]{1,0} get-tuple-element(while), index=0 + ROOT out = (f32[2,3]{1,0}) tuple(entry_gte0) +} + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + std::unique_ptr hlo_live_range; + std::unique_ptr alias_analysis; + RunAnalysis(module.get(), + {"p0", "p1", "p2", "init", "while", "entry_gte0", "out"}, + hlo_live_range, alias_analysis); + + HloInstruction* p2 = FindInstruction(module.get(), "p2"); + HloInstruction* while_instr = FindInstruction(module.get(), "while"); + HloInstruction* entry_gte0 = FindInstruction(module.get(), "entry_gte0"); + + HeapSimulator::Chunk chunk = HeapSimulator::Chunk::FromOffsetSize(0, 24); + PinnedAllocation p2_pinned(HloPosition{p2, {}}, MemorySpace::kDefault, chunk, + /*start_time=*/0, /*end_time=*/10); + + ParentAllocation parent_alloc( + p2_pinned, while_instr, + HloPosition{while_instr->while_body()->parameter_instruction(0), {}}, + /*time=*/5); + + BitcastSplitFn split_fn = nullptr; + ASSERT_OK(parent_alloc.Process(split_fn, *hlo_live_range, *alias_analysis)); + ASSERT_OK(parent_alloc.PostProcess()); + + ASSERT_NE(while_instr->while_init()->original_value(), nullptr); + EXPECT_EQ(while_instr->while_init()->original_value()->ToString(), + R"(({"p0"}, {"p1"}, {"p2"}))"); + + ASSERT_NE(while_instr->original_value(), nullptr); + EXPECT_EQ(while_instr->original_value()->ToString(), + R"(({"w0"}, {"w1"}, {"p2"}),["while#$"])"); + + ASSERT_NE( + while_instr->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"bp0"}, {"bp1"}, {"p2"}))"); + + ASSERT_NE(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"cp0"}, {"cp1"}, {"p2"}))"); + + ASSERT_NE(while_instr->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->root_instruction() + ->original_value() + ->ToString(), + R"(({"br0"}, {"br1"}, {"p2"}))"); + + const HloInstruction* tuple_with_old_shape = entry_gte0->operand(0); + ASSERT_NE(tuple_with_old_shape->original_value(), nullptr); + EXPECT_EQ(tuple_with_old_shape->original_value()->ToString(), + R"(({"w0"}, {"w1"}),["while#$"])"); +} + +} // namespace +} // namespace xla::memory_space_assignment