From 919d13ff6a294881d03ed69cbc443d2fbd6f633f Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 21 Sep 2026 11:09:28 -0700 Subject: [PATCH 1/3] Reverts 1656f20935d9f090a071c9cf00c8a4a9530cadaf PiperOrigin-RevId: 985391075 --- .../xla/third_party/stablehlo/temporary.patch | 1862 ++--------------- 1 file changed, 138 insertions(+), 1724 deletions(-) diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index ec90dce6dc5dc6..b6c9e2d8cb1f8a 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================== # -# This file is automatically generated by the `generate_patch` tool. +# This file is automatically generated by generate_patch tool. # Do not edit directly. diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel --- stablehlo/BUILD.bazel @@ -38,14 +38,14 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel + ], licenses = ["notice"], ) - ++ +# copybara:uncomment_begin +# license( +# name = "license", +# package_name = "stablehlo", +# ) +# copybara:uncomment_end -+ + exports_files([ "LICENSE", # Export to be used by MLIR-free python binding impls @@ -323,34 +323,78 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel deps = [ ":reference_errors", ":reference_types", -@@ -588,7 +626,8 @@ +@@ -588,38 +626,41 @@ hdrs = [ "stablehlo/reference/Errors.h", ], -- strip_include_prefix = ".", + compatible_with = get_compatible_with_portable(), + includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - ], -@@ -602,7 +641,8 @@ - hdrs = [ - "stablehlo/reference/Index.h", - ], -- strip_include_prefix = ".", ++ deps = [ ++ "@llvm-project//llvm:Support", ++ ], ++) ++ ++cc_library( ++ name = "reference_index", ++ srcs = [ ++ "stablehlo/reference/Index.cpp", ++ ], ++ hdrs = [ ++ "stablehlo/reference/Index.h", ++ ], + compatible_with = get_compatible_with_portable(), + includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", -@@ -618,6 +658,7 @@ - hdrs = [ - "stablehlo/reference/NumPy.h", - ], ++ deps = [ ++ "@llvm-project//llvm:Support", ++ "@llvm-project//mlir:IR", ++ "@llvm-project//mlir:Support", ++ ], ++) ++ ++cc_library( ++ name = "reference_numpy", ++ srcs = [ ++ "stablehlo/reference/NumPy.cpp", ++ ], ++ hdrs = [ ++ "stablehlo/reference/NumPy.h", ++ ], + compatible_with = get_compatible_with_portable(), strip_include_prefix = ".", deps = [ +- "@llvm-project//llvm:Support", +- ], +-) +- +-cc_library( +- name = "reference_index", +- srcs = [ +- "stablehlo/reference/Index.cpp", +- ], +- hdrs = [ +- "stablehlo/reference/Index.h", +- ], +- strip_include_prefix = ".", +- deps = [ +- "@llvm-project//llvm:Support", +- "@llvm-project//mlir:IR", +- "@llvm-project//mlir:Support", +- ], +-) +- +-cc_library( +- name = "reference_numpy", +- srcs = [ +- "stablehlo/reference/NumPy.cpp", +- ], +- hdrs = [ +- "stablehlo/reference/NumPy.h", +- ], +- strip_include_prefix = ".", +- deps = [ ":reference_tensor", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", @@ -635,7 +676,8 @@ hdrs = [ "stablehlo/reference/Ops.h", @@ -898,7 +942,7 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel tbl_outs = { "stablehlo/dialect/VhloTypeDefs.cpp.inc": ["-gen-typedef-defs"], "stablehlo/dialect/VhloTypeDefs.h.inc": ["-gen-typedef-decls"], -@@ -1826,227 +1896,21 @@ +@@ -1826,227 +1896,7 @@ ], ) @@ -1108,11 +1152,8 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel - "@llvm-project//mlir:InferTypeOpInterface", - "@llvm-project//mlir:Support", - ], -+test_suite( -+ name = "all_tests", -+ tests = [":stablehlo_ci_tests"], - ) - +-) +- -cc_test( - name = "stablehlo_builder_test", - srcs = ["stablehlo/integrations/cpp/builder/StablehloBuilderTest.cpp"], @@ -1129,19 +1170,10 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel - "@llvm-project//third-party/unittest:gtest", - ], -) -+# copybara:uncomment_begin -+# alias( -+# name = "check_ops", -+# actual = "//stablehlo/tests:check_ops", -+# deprecation = "Use //stablehlo/tests:check_ops instead.", -+# ) -+# -+# alias( -+# name = "test_utils", -+# actual = "//stablehlo/tests:test_utils", -+# deprecation = "Use //stablehlo/tests:test_utils instead.", -+# ) -+# copybara:uncomment_end ++test_suite( ++ name = "all_tests", ++ tests = [":stablehlo_ci_tests"], ++) diff --ruN a/stablehlo/build_tools/BUILD.bazel b/stablehlo/build_tools/BUILD.bazel --- stablehlo/build_tools/BUILD.bazel +++ stablehlo/build_tools/BUILD.bazel @@ -1187,7 +1219,7 @@ diff --ruN a/stablehlo/build_tools/BUILD.bazel b/stablehlo/build_tools/BUILD.baz diff --ruN a/stablehlo/build_tools/lit_test_suite.bzl b/stablehlo/build_tools/lit_test_suite.bzl --- stablehlo/build_tools/lit_test_suite.bzl +++ stablehlo/build_tools/lit_test_suite.bzl -@@ -0,0 +1,71 @@ +@@ -0,0 +1,69 @@ +# Copyright 2026 The StableHLO Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); @@ -1234,22 +1266,20 @@ diff --ruN a/stablehlo/build_tools/lit_test_suite.bzl b/stablehlo/build_tools/li + tests_in_suite = [] + + for src in srcs: -+ test_name = ("%s.test" % src).replace("/", "_") # copybara:comment -+ -+ # copybara:uncomment test_name = "%s.test" % src ++ test_name = "%s.test" % src + lit_test( -+ name = test_name, # copybara:comment -+ # copybara:uncomment name = src, -+ srcs = [src], # copybara:comment -+ # copybara:uncomment src = src, ++ name = test_name, ++ # name = src, # copybara:comment ++ srcs = [src], ++ # src = src, # copybara:comment + args = args, + data = data + native.glob( + ["%s.bc" % src], + # Allow empty since not all `.mlir` tests have `.mlir.bc` files. + allow_empty = True, + ), -+ deps = ["@rules_python//python/runfiles"], # copybara:comment -+ # copybara:uncomment driver = "@llvm-project//mlir:run_lit.sh", ++ deps = ["@rules_python//python/runfiles"], ++ # driver = "//third_party/llvm/llvm-project/mlir:run_lit.sh", # copybara:comment + tags = tags, + size = size_per_test, + ) @@ -1769,7 +1799,7 @@ diff --ruN a/stablehlo/examples/cpp/README.md b/stablehlo/examples/cpp/README.md diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel b/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel --- stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel +++ stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel -@@ -11,14 +11,27 @@ +@@ -11,13 +11,26 @@ # 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. @@ -1787,7 +1817,7 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel b/stablehl + ], licenses = ["notice"], ) - ++ +# copybara:uncomment_begin +# filegroup( +# name = "all_files", @@ -1795,10 +1825,9 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel b/stablehl +# visibility = ["//:__subpackages__"], +# ) +# copybara:uncomment_end -+ + # Equivalent of configure_lit_site_cfg from CMakeLists.txt. expand_template( - name = "lit_site_cfg_py_gen", @@ -39,25 +52,15 @@ template = "lit.site.cfg.py.in", ) @@ -1880,69 +1909,10 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/reduce.mlir b/stablehl // CHECK-SAME: strides = dense<2> : vector<2xi64>} // CHECK-SAME: ins(%[[ARG0]], %[[WINDOW]] : tensor<1x17x17x64xf32>, tensor<3x3xf32>) // CHECK-SAME: outs(%[[FILL]] : tensor<1x8x8x64xf32>) -> tensor<1x8x8x64xf32> -diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h b/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h ---- stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h -+++ stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h -@@ -60,26 +60,26 @@ - unsigned nParallelLoops); - - /// Generates an init sparse tensor. --Value getEmptySparseTensor(OpBuilder &b, Location loc, ShapedType type, -+Value getEmptySparseTensor(OpBuilder& b, Location loc, ShapedType type, - ArrayRef dynSizes); - - /// Generates a tensor.empty op. --Value getEmptyTensor(OpBuilder &b, Location loc, ShapedType type, -+Value getEmptyTensor(OpBuilder& b, Location loc, ShapedType type, - ArrayRef dynSizes); - - /// Generates an empty tensor for the result of the operation, which could be a - /// dense tensor or a sparse tensor. --Value getEmptyTensorFor(OpBuilder &b, Location loc, ShapedType resultType, -- Operation *op, ValueRange operands); -+Value getEmptyTensorFor(OpBuilder& b, Location loc, ShapedType resultType, -+ Operation* op, ValueRange operands); - - /// Ensures a tensor has the same shape (not including the element type) as - /// another. --Value coerceTensorShape(OpBuilder &builder, Location loc, -+Value coerceTensorShape(OpBuilder& builder, Location loc, - TypedValue value, ShapedType targetType); - - /// Fills |tensor| with a zero constant of the matching type. Returns the new - /// value. --Value fillTensorWithZeros(OpBuilder &builder, Location loc, Value tensor); -+Value fillTensorWithZeros(OpBuilder& builder, Location loc, Value tensor); - - /// Sparsifies a (block of) operation(s) that cannot be handled directly - /// by the sparse compiler but has well-known semi-ring semantics. -@@ -94,17 +94,17 @@ - /// } - /// absent={} - /// linalg.yield %result --Value preSparsify(Operation *op, llvm::SmallVector &values, Type rtp, -- OpBuilder *b); -+Value preSparsify(Operation* op, llvm::SmallVector& values, Type rtp, -+ OpBuilder* b); - - /// Finalizes sparse semi-ring construction. --Value postSparsify(Operation *op, Value semiring, Value result, OpBuilder *b); -+Value postSparsify(Operation* op, Value semiring, Value result, OpBuilder* b); - - /// Returns true if all operands are tensors with rank 0. --bool allOperandsAreScalarTensors(Operation *op); -+bool allOperandsAreScalarTensors(Operation* op); - - /// Returns true if parent op is linalg. --bool isInBodyOfLinalgOps(Operation *op); -+bool isInBodyOfLinalgOps(Operation* op); - - /// Extracts integer values from the attribute |elements|. - SmallVector extract1DVector(DenseIntElementsAttr elements); diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel --- stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel +++ stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel -@@ -11,14 +11,27 @@ +@@ -11,13 +11,26 @@ # 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. @@ -1960,7 +1930,7 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/ + ], licenses = ["notice"], ) - ++ +# copybara:uncomment_begin +# filegroup( +# name = "all_files", @@ -1968,10 +1938,9 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/ +# visibility = ["//:__subpackages__"], +# ) +# copybara:uncomment_end -+ + # Equivalent of configure_lit_site_cfg from CMakeLists.txt. expand_template( - name = "lit_site_cfg_py_gen", @@ -39,26 +52,16 @@ template = "lit.site.cfg.py.in", ) @@ -2009,147 +1978,6 @@ diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/ + size_per_test = "small", tags = ["stablehlo_tosa_tests"], ) -diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.cpp b/stablehlo/stablehlo/dialect/ChloBytecode.cpp ---- stablehlo/stablehlo/dialect/ChloBytecode.cpp -+++ stablehlo/stablehlo/dialect/ChloBytecode.cpp -@@ -111,37 +111,37 @@ - /// This class implements the bytecode interface for the CHLO dialect. - class ChloBytecodeInterface : public BytecodeDialectInterface { - public: -- ChloBytecodeInterface(Dialect *dialect) : BytecodeDialectInterface(dialect) {} -+ ChloBytecodeInterface(Dialect* dialect) : BytecodeDialectInterface(dialect) {} - - //===--------------------------------------------------------------------===// - // Attributes - - // These methods are invoked by superclass when an attr from CHLO dialect - // is encountered. -- Attribute readAttribute(DialectBytecodeReader &reader) const override; -+ Attribute readAttribute(DialectBytecodeReader& reader) const override; - LogicalResult writeAttribute(Attribute attr, -- DialectBytecodeWriter &writer) const override; -+ DialectBytecodeWriter& writer) const override; - - // TO ADD ATTRIBUTE: Include a read method for each attribute in CHLO - // Ex: SomeAttr readSomeAttr(DialectBytecodeReader &reader) const; - ComparisonDirectionAttr readComparisonDirectionAttr( -- DialectBytecodeReader &reader) const; -+ DialectBytecodeReader& reader) const; - ComparisonTypeAttr readComparisonTypeAttr( -- DialectBytecodeReader &reader) const; -+ DialectBytecodeReader& reader) const; - - // TO ADD ATTRIBUTE: Include a write method for each attribute in CHLO - // Ex: void write(SomeAttr attr, DialectBytecodeWriter &writer) const; -- void write(ComparisonDirectionAttr attr, DialectBytecodeWriter &writer) const; -- void write(ComparisonTypeAttr attr, DialectBytecodeWriter &writer) const; -+ void write(ComparisonDirectionAttr attr, DialectBytecodeWriter& writer) const; -+ void write(ComparisonTypeAttr attr, DialectBytecodeWriter& writer) const; - - //===--------------------------------------------------------------------===// - // Types - - // These methods are invoked by superclass when a type from CHLO dialect - // is encountered. -- Type readType(DialectBytecodeReader &reader) const override; -+ Type readType(DialectBytecodeReader& reader) const override; - LogicalResult writeType(Type type, -- DialectBytecodeWriter &writer) const override; -+ DialectBytecodeWriter& writer) const override; - - // TO ADD TYPE: Include a read method for each type in CHLO - // Ex: SomeType readSomeType(DialectBytecodeReader &reader) const; -@@ -158,7 +158,7 @@ - - // TO ADD ATTRIBUTE: Update the switch to include a branch for the attr. - Attribute ChloBytecodeInterface::readAttribute( -- DialectBytecodeReader &reader) const { -+ DialectBytecodeReader& reader) const { - uint64_t code; - if (failed(reader.readVarInt(code))) return Attribute(); - switch (code) { -@@ -173,7 +173,7 @@ - } - - ComparisonDirectionAttr ChloBytecodeInterface::readComparisonDirectionAttr( -- DialectBytecodeReader &reader) const { -+ DialectBytecodeReader& reader) const { - LOG_READ_CALL; - return hlo::bytecode::readEnumAttribute( - reader, getContext(), -@@ -181,7 +181,7 @@ - } - - ComparisonTypeAttr ChloBytecodeInterface::readComparisonTypeAttr( -- DialectBytecodeReader &reader) const { -+ DialectBytecodeReader& reader) const { - LOG_READ_CALL; - return hlo::bytecode::readEnumAttribute( - reader, getContext(), -@@ -195,7 +195,7 @@ - // If this method returns failure, the string serialization is used in the - // bytecode. - LogicalResult ChloBytecodeInterface::writeAttribute( -- Attribute attr, DialectBytecodeWriter &writer) const { -+ Attribute attr, DialectBytecodeWriter& writer) const { - return TypeSwitch(attr) - .Case([&](auto attr) { - LOG_WRITE_CALL; -@@ -209,13 +209,13 @@ - } - - void ChloBytecodeInterface::write(ComparisonDirectionAttr attr, -- DialectBytecodeWriter &writer) const { -+ DialectBytecodeWriter& writer) const { - writer.writeVarInt(chlo_encoding::kComparisonDirectionAttr); - hlo::bytecode::writeEnumAttribute(attr, writer); - } - - void ChloBytecodeInterface::write(ComparisonTypeAttr attr, -- DialectBytecodeWriter &writer) const { -+ DialectBytecodeWriter& writer) const { - writer.writeVarInt(chlo_encoding::kComparisonTypeAttr); - hlo::bytecode::writeEnumAttribute(attr, writer); - } -@@ -224,7 +224,7 @@ - // Types: Reader - - // TO ADD TYPE: Update the case selection to include the new type. --Type ChloBytecodeInterface::readType(DialectBytecodeReader &reader) const { -+Type ChloBytecodeInterface::readType(DialectBytecodeReader& reader) const { - uint64_t code; - if (failed(reader.readVarInt(code))) return Type(); - -@@ -240,7 +240,7 @@ - - // TO ADD TYPE: Update the case selection to include the new type. - LogicalResult ChloBytecodeInterface::writeType( -- Type type, DialectBytecodeWriter &writer) const { -+ Type type, DialectBytecodeWriter& writer) const { - return TypeSwitch(type).Default([&](Type) { - LOG_NOT_IMPLEMENTED; - return failure(); -@@ -249,7 +249,7 @@ - - } // namespace - --void addBytecodeInterface(ChloDialect *dialect) { -+void addBytecodeInterface(ChloDialect* dialect) { - dialect->addInterfaces(); - } - } // namespace chlo -diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.h b/stablehlo/stablehlo/dialect/ChloBytecode.h ---- stablehlo/stablehlo/dialect/ChloBytecode.h -+++ stablehlo/stablehlo/dialect/ChloBytecode.h -@@ -21,7 +21,7 @@ - - // Add the interface necessary for encoding and decoding the CHLO dialect - // components in bytecode. --void addBytecodeInterface(ChloDialect *dialect); -+void addBytecodeInterface(ChloDialect* dialect); - } // namespace chlo - } // namespace mlir - diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dialect/ChloEnums.td --- stablehlo/stablehlo/dialect/ChloEnums.td +++ stablehlo/stablehlo/dialect/ChloEnums.td @@ -2203,10 +2031,11 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/diale for (auto [i, type] : llvm::enumerate(adaptor.getInputs().getTypes())) { auto inputType = dyn_cast(type); if (!inputType) { -@@ -839,8 +843,27 @@ +@@ -839,7 +843,26 @@ if (dimSize == ShapedType::kDynamic) { dimSize = inputType.getDimSize(dim); } +- } + if (inputType.getEncoding()) { + prototypeEncoding = inputType.getEncoding(); + } @@ -2221,16 +2050,15 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/diale + } + } + } - } - ++ } ++ + // If the scan dimension size is dynamic, propagate any bound inferred from + // the inputs; otherwise static dimensions do not need a bound. + const int64_t insertBound = + (dimSize == ShapedType::kDynamic) ? dimBound : ShapedType::kDynamic; -+ + for (auto [i, type] : llvm::enumerate(terminator->getOperands().getTypes())) { auto resultType = dyn_cast(type); - if (!resultType) { @@ -849,10 +872,34 @@ } SmallVector shape(resultType.getShape().begin(), @@ -2241,6 +2069,8 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/diale + // carries whose shapes and encodings are preserved as-is. if (i < numOutputs) { shape.insert(std::next(shape.begin(), dim), dimSize); +- } +- inferredReturnShapes.emplace_back(shape, resultType.getElementType()); + SmallVector resultBounds = + llvm::to_vector(hlo::encodingToBounds(encoding)); + // Re-insert bounds along the scan dimension if the result has bounds or @@ -2260,8 +2090,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/diale + encoding = hlo::boundsToEncoding(encoding, resultBounds); + } + } - } -- inferredReturnShapes.emplace_back(shape, resultType.getElementType()); ++ } + inferredReturnShapes.emplace_back(shape, resultType.getElementType(), + encoding); } @@ -2286,30 +2115,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/diale } if (!hlo::isCompatibleForHloTypeInference( argType, bodyBlock.getArgument(i).getType())) { -diff --ruN a/stablehlo/stablehlo/dialect/Register.cpp b/stablehlo/stablehlo/dialect/Register.cpp ---- stablehlo/stablehlo/dialect/Register.cpp -+++ stablehlo/stablehlo/dialect/Register.cpp -@@ -27,7 +27,7 @@ - namespace mlir { - namespace stablehlo { - --void registerAllDialects(mlir::DialectRegistry ®istry) { -+void registerAllDialects(mlir::DialectRegistry& registry) { - // clang-format off - registry.insert transposedSizes(meshAxisNames.size()); +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { + // Grouped axes in the specified order + for (int64_t idx : groupedAxisIndices) { + transposeAxes.push_back(idx); - } - -- SmallVector transposedSizes(meshAxisNames.size()); -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { ++ } ++ + SmallVector transposedSizes(axisSizes.size()); + for (size_t i = 0; i < axisSizes.size(); ++i) { transposedSizes[i] = axisSizes[transposeAxes[i]]; @@ -2416,14 +2223,14 @@ diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stabl int64_t coord = (linearTransposeIdx / transposedStrides[k]) % transposedSizes[k]; originalIndex += coord * originalStrides[transposeAxes[k]]; -@@ -102,10 +99,136 @@ +@@ -102,9 +99,135 @@ } groups.push_back(std::move(group)); } - return groups; } - ++ +// Splits mesh axes based on sub-axis references and computes the corresponding +// indices for the communication axes. +FailureOr computeReindexedAxes(ArrayRef axesInMesh, @@ -2550,10 +2357,9 @@ diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stabl +} + +} // namespace -+ + FailureOr>> flattenReplicaGroupMeshAxes( Attribute meshAttr, ArrayAttr commAxes, Location loc) { - // Reference implementation for flattening a mesh-axes based replica group @@ -120,34 +243,13 @@ if (!mesh) return emitOptionalError(loc, "expected stablehlo.mesh for mesh attribute"); @@ -2605,18 +2411,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stabl } } // namespace stablehlo -diff --ruN a/stablehlo/stablehlo/dialect/StablehloBytecode.h b/stablehlo/stablehlo/dialect/StablehloBytecode.h ---- stablehlo/stablehlo/dialect/StablehloBytecode.h -+++ stablehlo/stablehlo/dialect/StablehloBytecode.h -@@ -21,7 +21,7 @@ - - // Add the interface necessary for encoding and decoding StableHLO dialect - // components in bytecode. --void addBytecodeInterface(StablehloDialect *dialect); -+void addBytecodeInterface(StablehloDialect* dialect); - } // namespace stablehlo - } // namespace mlir - diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo/dialect/StablehloEnums.td --- stablehlo/stablehlo/dialect/StablehloEnums.td +++ stablehlo/stablehlo/dialect/StablehloEnums.td @@ -2697,18 +2491,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo #endif // STABLEHLO_DIALECT_STABLEHLO_ENUMS + -diff --ruN a/stablehlo/stablehlo/dialect/VhloBytecode.h b/stablehlo/stablehlo/dialect/VhloBytecode.h ---- stablehlo/stablehlo/dialect/VhloBytecode.h -+++ stablehlo/stablehlo/dialect/VhloBytecode.h -@@ -21,7 +21,7 @@ - - // Add the interface necessary for encoding and decoding VHLO dialect - // components in bytecode. --void addBytecodeInterface(VhloDialect *dialect); -+void addBytecodeInterface(VhloDialect* dialect); - } // namespace vhlo - } // namespace mlir - diff --ruN a/stablehlo/stablehlo/dialect/VhloEnums.td b/stablehlo/stablehlo/dialect/VhloEnums.td --- stablehlo/stablehlo/dialect/VhloEnums.td +++ stablehlo/stablehlo/dialect/VhloEnums.td @@ -2720,169 +2502,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/VhloEnums.td b/stablehlo/stablehlo/dial let extraClassDeclaration = [{ mlir::vhlo::Version getMinVersion() { return mlir::vhlo::Version(}] # !subst(".", ", ", minVersion) # [{); -diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.h b/stablehlo/stablehlo/dialect/VhloOps.h ---- stablehlo/stablehlo/dialect/VhloOps.h -+++ stablehlo/stablehlo/dialect/VhloOps.h -@@ -31,20 +31,20 @@ - - class VhloDialect : public Dialect { - public: -- explicit VhloDialect(MLIRContext *context); -+ explicit VhloDialect(MLIRContext* context); - static StringRef getDialectNamespace() { return "vhlo"; } - - // Parses a type registered to this dialect. -- Type parseType(DialectAsmParser &parser) const override; -+ Type parseType(DialectAsmParser& parser) const override; - - // Prints a type registered to this dialect. -- void printType(Type type, DialectAsmPrinter &os) const override; -+ void printType(Type type, DialectAsmPrinter& os) const override; - - // Parses an attribute registered to this dialect. -- Attribute parseAttribute(DialectAsmParser &parser, Type type) const override; -+ Attribute parseAttribute(DialectAsmParser& parser, Type type) const override; - - // Prints an attribute registered to this dialect. -- void printAttribute(Attribute attr, DialectAsmPrinter &os) const override; -+ void printAttribute(Attribute attr, DialectAsmPrinter& os) const override; - - private: - // Adds VHLO types to this dialect. -diff --ruN a/stablehlo/stablehlo/integrations/c/ChloAttributes.cpp b/stablehlo/stablehlo/integrations/c/ChloAttributes.cpp ---- stablehlo/stablehlo/integrations/c/ChloAttributes.cpp -+++ stablehlo/stablehlo/integrations/c/ChloAttributes.cpp -@@ -73,12 +73,12 @@ - - MlirAttribute chloRaggedDotDimensionNumbersGet( - MlirContext ctx, intptr_t nLhsBatchingDimensions, -- const int64_t *lhsBatchingDimensions, intptr_t nRhsBatchingDimensions, -- const int64_t *rhsBatchingDimensions, intptr_t nLhsContractingDimensions, -- const int64_t *lhsContractingDimensions, intptr_t nRhsContractingDimensions, -- const int64_t *rhsContractingDimensions, intptr_t nLhsRaggedDimensions, -- const int64_t *lhsRaggedDimensions, intptr_t nRhsGroupDimensions, -- const int64_t *rhsGroupDimensions) { -+ const int64_t* lhsBatchingDimensions, intptr_t nRhsBatchingDimensions, -+ const int64_t* rhsBatchingDimensions, intptr_t nLhsContractingDimensions, -+ const int64_t* lhsContractingDimensions, intptr_t nRhsContractingDimensions, -+ const int64_t* rhsContractingDimensions, intptr_t nLhsRaggedDimensions, -+ const int64_t* lhsRaggedDimensions, intptr_t nRhsGroupDimensions, -+ const int64_t* rhsGroupDimensions) { - return wrap(mlir::chlo::RaggedDotDimensionNumbersAttr::get( - unwrap(ctx), - llvm::ArrayRef(lhsBatchingDimensions, nLhsBatchingDimensions), -diff --ruN a/stablehlo/stablehlo/integrations/c/ChloAttributes.h b/stablehlo/stablehlo/integrations/c/ChloAttributes.h ---- stablehlo/stablehlo/integrations/c/ChloAttributes.h -+++ stablehlo/stablehlo/integrations/c/ChloAttributes.h -@@ -53,16 +53,16 @@ - - MLIR_CAPI_EXPORTED MlirAttribute chloRaggedDotDimensionNumbersGet( - MlirContext ctx, // -- intptr_t nLhsBatchingDimensions, const int64_t *lhsBatchingDimensions, // -- intptr_t nRhsBatchingDimensions, const int64_t *rhsBatchingDimensions, // -+ intptr_t nLhsBatchingDimensions, const int64_t* lhsBatchingDimensions, // -+ intptr_t nRhsBatchingDimensions, const int64_t* rhsBatchingDimensions, // - intptr_t nLhsContractingDimensions, // -- const int64_t *lhsContractingDimensions, // -+ const int64_t* lhsContractingDimensions, // - intptr_t nRhsContractingDimensions, // -- const int64_t *rhsContractingDimensions, // -+ const int64_t* rhsContractingDimensions, // - intptr_t nLhsRaggedDimensions, // -- const int64_t *lhsRaggedDimensions, // -+ const int64_t* lhsRaggedDimensions, // - intptr_t nRhsGroupDimensions, // -- const int64_t *rhsGroupDimensions); -+ const int64_t* rhsGroupDimensions); - - MLIR_CAPI_EXPORTED bool chloAttributeIsARaggedDotDimensionNumbers( - MlirAttribute attr); -diff --ruN a/stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp b/stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp ---- stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp -+++ stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp -@@ -48,18 +48,18 @@ - - void stablehloVersionFromCompatibilityRequirement( - MlirStablehloCompatibilityRequirement requirement, -- MlirStringCallback callback, void *userData) { -+ MlirStringCallback callback, void* userData) { - mlir::detail::CallbackOstream stream(callback, userData); - stream << mlir::vhlo::Version::fromCompatibilityRequirement( - unwrapCompatibilityRequirement(requirement)); - } - --void stablehloGetCurrentVersion(MlirStringCallback callback, void *userData) { -+void stablehloGetCurrentVersion(MlirStringCallback callback, void* userData) { - mlir::detail::CallbackOstream stream(callback, userData); - stream << mlir::stablehlo::getCurrentVersion(); - } - --void stablehloGetMinimumVersion(MlirStringCallback callback, void *userData) { -+void stablehloGetMinimumVersion(MlirStringCallback callback, void* userData) { - mlir::detail::CallbackOstream stream(callback, userData); - stream << mlir::stablehlo::getMinimumVersion(); - } -@@ -67,7 +67,7 @@ - MlirLogicalResult stablehloGetSmallerVersion(MlirStringRef version1, - MlirStringRef version2, - MlirStringCallback callback, -- void *userData) { -+ void* userData) { - mlir::detail::CallbackOstream stream(callback, userData); - auto result = - mlir::stablehlo::getSmallerVersion(unwrap(version1), unwrap(version2)); -@@ -78,7 +78,7 @@ - - MlirLogicalResult stablehloSerializePortableArtifactFromModule( - MlirModule moduleStr, MlirStringRef targetVersion, -- MlirStringCallback callback, void *userData, bool allowOtherDialects) { -+ MlirStringCallback callback, void* userData, bool allowOtherDialects) { - mlir::detail::CallbackOstream stream(callback, userData); - if (failed(mlir::stablehlo::serializePortableArtifact( - unwrap(moduleStr), unwrap(targetVersion), stream, -@@ -89,7 +89,7 @@ - - MlirLogicalResult stablehloSerializePortableArtifactFromStringRef( - MlirStringRef moduleStr, MlirStringRef targetVersion, -- MlirStringCallback callback, void *userData) { -+ MlirStringCallback callback, void* userData) { - mlir::detail::CallbackOstream stream(callback, userData); - if (failed(mlir::stablehlo::serializePortableArtifact( - unwrap(moduleStr), unwrap(targetVersion), stream))) -@@ -98,7 +98,7 @@ - } - - MlirLogicalResult stablehloDeserializePortableArtifact( -- MlirStringRef artifactStr, MlirStringCallback callback, void *userData) { -+ MlirStringRef artifactStr, MlirStringCallback callback, void* userData) { - mlir::detail::CallbackOstream stream(callback, userData); - if (failed(mlir::stablehlo::deserializePortableArtifact(unwrap(artifactStr), - stream))) -diff --ruN a/stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp b/stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp ---- stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp -+++ stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp -@@ -32,9 +32,9 @@ - #include "stablehlo/reference/Configuration.h" - - MlirAttribute stablehloEvalModule(MlirModule module, int nArgs, -- MlirAttribute const *args, -- const char *const probeInstrumentationDir, -- int *errorCode) { -+ MlirAttribute const* args, -+ const char* const probeInstrumentationDir, -+ int* errorCode) { - std::vector inputs; - inputs.reserve(nArgs); - for (int i = 0; i < nArgs; ++i) { -@@ -49,7 +49,7 @@ - return MlirAttribute{nullptr}; - } - std::vector resultsVec; -- for (const auto &result : results.value()) { -+ for (const auto& result : results.value()) { - resultsVec.push_back(wrap(result)); - } - return mlirArrayAttrGet(mlirModuleGetContext(module), resultsVec.size(), diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel b/stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel --- stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel +++ stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel @@ -3192,745 +2811,10 @@ diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/MlirBuilder.cpp b/stab #include "mlir/IR/Location.h" #include "mlir/IR/Value.h" -diff --ruN a/stablehlo/stablehlo/integrations/python/StablehloApi.cpp b/stablehlo/stablehlo/integrations/python/StablehloApi.cpp ---- stablehlo/stablehlo/integrations/python/StablehloApi.cpp -+++ stablehlo/stablehlo/integrations/python/StablehloApi.cpp -@@ -42,15 +42,15 @@ - StringWriterHelper() : ss_(s_) {} - - static MlirStringCallback getMlirStringCallback() { -- return [](MlirStringRef string_ref, void *user_data) { -- auto *helper = static_cast(user_data); -+ return [](MlirStringRef string_ref, void* user_data) { -+ auto* helper = static_cast(user_data); - helper->ss_ << llvm::StringRef(string_ref.data, string_ref.length); - }; - } - -- void *getUserData() { return static_cast(this); } -+ void* getUserData() { return static_cast(this); } - -- const std::string &toString() { -+ const std::string& toString() { - ss_.flush(); - return s_; - } -@@ -60,7 +60,7 @@ - llvm::raw_string_ostream ss_; - }; - --static MlirStringRef toMlirStringRef(const std::string &s) { -+static MlirStringRef toMlirStringRef(const std::string& s) { - return mlirStringRefCreate(s.data(), s.size()); - } - -@@ -68,11 +68,11 @@ - return mlirStringRefCreate(s.data(), s.size()); - } - --static MlirStringRef toMlirStringRef(const nb::bytes &s) { -- return mlirStringRefCreate(static_cast(s.data()), s.size()); -+static MlirStringRef toMlirStringRef(const nb::bytes& s) { -+ return mlirStringRefCreate(static_cast(s.data()), s.size()); - } - --void AddStablehloApi(nb::module_ &m) { -+void AddStablehloApi(nb::module_& m) { - // Portable API is a subset of StableHLO API - AddPortableApi(m); - -@@ -146,8 +146,8 @@ - // - m.def( - "eval_module", -- [](MlirModule module, std::vector &args, -- const std::string &probe_instrumentation_dir) -+ [](MlirModule module, std::vector& args, -+ const std::string& probe_instrumentation_dir) - -> std::vector { - for (auto arg : args) { - if (!mlirAttributeIsADenseElements(arg)) { -@@ -174,7 +174,7 @@ - nb::arg("probe_instrumentation_dir") = ""); - } - --void AddPortableApi(nb::module_ &m) { -+void AddPortableApi(nb::module_& m) { - // - // Utility APIs. - // -@@ -182,8 +182,8 @@ - - m.def( - "get_smaller_version", -- [](const std::string &version1, -- const std::string &version2) -> std::string { -+ [](const std::string& version1, -+ const std::string& version2) -> std::string { - StringWriterHelper accumulator; - if (mlirLogicalResultIsFailure(stablehloGetSmallerVersion( - toMlirStringRef(version1), toMlirStringRef(version2), -@@ -263,7 +263,7 @@ - nb::arg("artifact_str")); - m.def( - "deserialize_portable_artifact_str", -- [](const nb::bytes &artifact) -> nb::bytes { -+ [](const nb::bytes& artifact) -> nb::bytes { - StringWriterHelper accumulator; - if (mlirLogicalResultIsFailure(stablehloDeserializePortableArtifact( - toMlirStringRef(artifact), accumulator.getMlirStringCallback(), -diff --ruN a/stablehlo/stablehlo/reference/Api.cpp b/stablehlo/stablehlo/reference/Api.cpp ---- stablehlo/stablehlo/reference/Api.cpp -+++ stablehlo/stablehlo/reference/Api.cpp -@@ -83,7 +83,7 @@ - class DefaultInterpreterFallback : public InterpreterFallback { - public: - DefaultInterpreterFallback(const InterpreterConfiguration& config) -- : config(config){}; -+ : config(config) {}; - - virtual llvm::Error operator()(Operation& op, Scope& scope, - Process* process) final { -diff --ruN a/stablehlo/stablehlo/reference/Api.h b/stablehlo/stablehlo/reference/Api.h ---- stablehlo/stablehlo/reference/Api.h -+++ stablehlo/stablehlo/reference/Api.h -@@ -36,17 +36,17 @@ - /// builtin kernels are matched. - FailureOr> evalModule( - ModuleOp module, ArrayRef inputs, -- const InterpreterConfiguration &config); -+ const InterpreterConfiguration& config); - - /// This wrapper is intended to be easily used by the StableHLO Python bindings. - // It wraps the InterpreterValue API. - FailureOr> evalModule( - ModuleOp module, ArrayRef inputs, -- const InterpreterConfiguration &config); -+ const InterpreterConfiguration& config); - - /// Parses a StableHLO MLIR text program into a ModuleOp. --FailureOr> parseStablehloModule(const std::string &mlir, -- MLIRContext &context); -+FailureOr> parseStablehloModule(const std::string& mlir, -+ MLIRContext& context); - - } // namespace stablehlo - } // namespace mlir -diff --ruN a/stablehlo/stablehlo/reference/Axes.cpp b/stablehlo/stablehlo/reference/Axes.cpp ---- stablehlo/stablehlo/reference/Axes.cpp -+++ stablehlo/stablehlo/reference/Axes.cpp -@@ -21,7 +21,7 @@ - namespace mlir { - namespace stablehlo { - --raw_ostream &operator<<(raw_ostream &os, const Axes &x) { -+raw_ostream& operator<<(raw_ostream& os, const Axes& x) { - os << "["; - llvm::interleave(x, os, ", "); - os << "]"; -diff --ruN a/stablehlo/stablehlo/reference/Axes.h b/stablehlo/stablehlo/reference/Axes.h ---- stablehlo/stablehlo/reference/Axes.h -+++ stablehlo/stablehlo/reference/Axes.h -@@ -29,8 +29,8 @@ - class Axes : public SmallVector { - public: - Axes() = default; -- Axes(const Axes &other) = default; -- Axes &operator=(const Axes &other) = default; -+ Axes(const Axes& other) = default; -+ Axes& operator=(const Axes& other) = default; - - Axes(std::initializer_list list) : SmallVector(list) {} - explicit Axes(size_t size, int64_t element = 0) -@@ -40,7 +40,7 @@ - : SmallVector(attr.getValues()) {} - }; - --raw_ostream &operator<<(raw_ostream &os, const Axes &x); -+raw_ostream& operator<<(raw_ostream& os, const Axes& x); - - } // namespace stablehlo - } // namespace mlir -diff --ruN a/stablehlo/stablehlo/reference/Configuration.cpp b/stablehlo/stablehlo/reference/Configuration.cpp ---- stablehlo/stablehlo/reference/Configuration.cpp -+++ stablehlo/stablehlo/reference/Configuration.cpp -@@ -25,8 +25,8 @@ - namespace mlir { - namespace stablehlo { - --llvm::Error InterpreterFallback::operator()(Operation &op, Scope &scope, -- Process *process) { -+llvm::Error InterpreterFallback::operator()(Operation& op, Scope& scope, -+ Process* process) { - return stablehlo::invalidArgument("Unsupported op: %s", - debugString(op).c_str()); - } -diff --ruN a/stablehlo/stablehlo/reference/Configuration.h b/stablehlo/stablehlo/reference/Configuration.h ---- stablehlo/stablehlo/reference/Configuration.h -+++ stablehlo/stablehlo/reference/Configuration.h -@@ -30,7 +30,7 @@ - public: - /// Custom op kernels for any user specified ops not found in the StableHLO - /// op dialect or StableHLO interpreter dialect. -- virtual llvm::Error operator()(Operation &op, Scope &scope, Process *process); -+ virtual llvm::Error operator()(Operation& op, Scope& scope, Process* process); - - virtual ~InterpreterFallback() = default; - }; -diff --ruN a/stablehlo/stablehlo/reference/Errors.h b/stablehlo/stablehlo/reference/Errors.h ---- stablehlo/stablehlo/reference/Errors.h -+++ stablehlo/stablehlo/reference/Errors.h -@@ -26,7 +26,7 @@ - /// error using a format string and a variadic number of arguments to the format - /// string. - template --inline llvm::Error invalidArgument(char const *Fmt, const Ts &...Vals) { -+inline llvm::Error invalidArgument(char const* Fmt, const Ts&... Vals) { - return createStringError(llvm::errc::invalid_argument, Fmt, Vals...); - } - -diff --ruN a/stablehlo/stablehlo/reference/Index.cpp b/stablehlo/stablehlo/reference/Index.cpp ---- stablehlo/stablehlo/reference/Index.cpp -+++ stablehlo/stablehlo/reference/Index.cpp -@@ -26,7 +26,7 @@ - namespace mlir { - namespace stablehlo { - --raw_ostream &operator<<(raw_ostream &os, const Sizes &x) { -+raw_ostream& operator<<(raw_ostream& os, const Sizes& x) { - os << "["; - llvm::interleave(x, os, ", "); - os << "]"; -@@ -40,7 +40,7 @@ - return result; - } - --bool Sizes::inBounds(const Sizes &bounds) const { -+bool Sizes::inBounds(const Sizes& bounds) const { - if (size() != bounds.size()) return false; - for (auto [size, bound] : llvm::zip(*this, bounds)) - if (size < 0 || size >= bound) return false; -@@ -59,7 +59,7 @@ - return IndexSpaceIterator(*this); - } - --Sizes operator+(const Sizes &x, const Sizes &y) { -+Sizes operator+(const Sizes& x, const Sizes& y) { - if (x.size() != y.size()) llvm::report_fatal_error("expected same size"); - Sizes result(x.size()); - for (size_t i = 0; i < x.size(); ++i) { -@@ -68,11 +68,11 @@ - return result; - } - --Sizes operator+(const Sizes &x, int64_t y) { return x + Sizes(x.size(), y); } -+Sizes operator+(const Sizes& x, int64_t y) { return x + Sizes(x.size(), y); } - --Sizes operator+(int64_t x, const Sizes &y) { return y + x; } -+Sizes operator+(int64_t x, const Sizes& y) { return y + x; } - --Sizes operator-(const Sizes &x, const Sizes &y) { -+Sizes operator-(const Sizes& x, const Sizes& y) { - if (x.size() != y.size()) llvm::report_fatal_error("expected same size"); - Sizes result(x.size()); - for (size_t i = 0; i < x.size(); ++i) { -@@ -81,11 +81,11 @@ - return result; - } - --Sizes operator-(const Sizes &x, int64_t y) { return x - Sizes(x.size(), y); } -+Sizes operator-(const Sizes& x, int64_t y) { return x - Sizes(x.size(), y); } - --Sizes operator-(int64_t x, const Sizes &y) { return Sizes(y.size(), x) - y; } -+Sizes operator-(int64_t x, const Sizes& y) { return Sizes(y.size(), x) - y; } - --Sizes operator*(const Sizes &x, const Sizes &y) { -+Sizes operator*(const Sizes& x, const Sizes& y) { - if (x.size() != y.size()) llvm::report_fatal_error("expected same size"); - Sizes result(x.size()); - for (size_t i = 0; i < x.size(); ++i) { -@@ -94,23 +94,23 @@ - return result; - } - --Sizes operator*(const Sizes &x, int64_t y) { return x * Sizes(x.size(), y); } -+Sizes operator*(const Sizes& x, int64_t y) { return x * Sizes(x.size(), y); } - --Sizes operator*(int64_t &x, const Sizes &y) { return y + x; } -+Sizes operator*(int64_t& x, const Sizes& y) { return y + x; } - --Sizes clamp(int64_t min, const Sizes &x, int64_t max) { -+Sizes clamp(int64_t min, const Sizes& x, int64_t max) { - return clamp(Sizes(x.size(), min), x, Sizes(x.size(), max)); - } - --Sizes clamp(int64_t min, const Sizes &x, const Sizes &max) { -+Sizes clamp(int64_t min, const Sizes& x, const Sizes& max) { - return clamp(Sizes(x.size(), min), x, max); - } - --Sizes clamp(const Sizes &min, const Sizes &x, int64_t max) { -+Sizes clamp(const Sizes& min, const Sizes& x, int64_t max) { - return clamp(min, x, Sizes(x.size(), max)); - } - --Sizes clamp(const Sizes &min, const Sizes &x, const Sizes &max) { -+Sizes clamp(const Sizes& min, const Sizes& x, const Sizes& max) { - if (min.size() != x.size() || x.size() != max.size()) - llvm::report_fatal_error("expected same size"); - Sizes result(x.size()); -@@ -120,15 +120,15 @@ - return result; - } - --const Index &IndexSpaceIterator::operator*() const { -+const Index& IndexSpaceIterator::operator*() const { - if (!index_) - llvm::report_fatal_error("Dereferencing a past-the-end iterator."); - return *index_; - } - --const Index *IndexSpaceIterator::operator->() const { return &(*index_); } -+const Index* IndexSpaceIterator::operator->() const { return &(*index_); } - --IndexSpaceIterator &IndexSpaceIterator::operator++() { -+IndexSpaceIterator& IndexSpaceIterator::operator++() { - if (!index_) - llvm::report_fatal_error("Incrementing a past-the-end iterator."); - -diff --ruN a/stablehlo/stablehlo/reference/Index.h b/stablehlo/stablehlo/reference/Index.h ---- stablehlo/stablehlo/reference/Index.h -+++ stablehlo/stablehlo/reference/Index.h -@@ -32,8 +32,8 @@ - class Sizes : public SmallVector { - public: - Sizes() = default; -- Sizes(const Sizes &other) = default; -- Sizes &operator=(const Sizes &other) = default; -+ Sizes(const Sizes& other) = default; -+ Sizes& operator=(const Sizes& other) = default; - - Sizes(std::initializer_list list) : SmallVector(list) {} - Sizes(iterator begin, iterator end) : SmallVector(begin, end) {} -@@ -49,66 +49,66 @@ - - /// Checks if an element `e` at kth axis of `this` object follows - /// `0 <= e <= bounds[k]`. -- bool inBounds(const Sizes &bounds) const; -+ bool inBounds(const Sizes& bounds) const; - - /// Iterate over the index space of a Sizes object. - IndexSpaceIterator index_begin() const; - IndexSpaceIterator index_end() const; - }; - --raw_ostream &operator<<(raw_ostream &os, const Sizes &x); -+raw_ostream& operator<<(raw_ostream& os, const Sizes& x); - - /// Overloaded add operator to return `Sizes` object `z` such that - /// `z[k] = x[k] + y[k]` for all axis k. --Sizes operator+(const Sizes &x, const Sizes &y); -+Sizes operator+(const Sizes& x, const Sizes& y); - - /// Overloaded add operator to return `Sizes` object `z` such that - /// `z[k] = x[k] + y` for all axis k. --Sizes operator+(const Sizes &x, int64_t y); -+Sizes operator+(const Sizes& x, int64_t y); - - /// Overloaded add operator to return `Sizes` object `z` such that - /// `z[k] = x + y[k]` for all axis k. --Sizes operator+(int64_t x, const Sizes &y); -+Sizes operator+(int64_t x, const Sizes& y); - - /// Overloaded add operator to return `Sizes` object `z` such that - /// `z[k] = x[k] - y[k]` for all axis k. --Sizes operator-(const Sizes &x, const Sizes &y); -+Sizes operator-(const Sizes& x, const Sizes& y); - - /// Overloaded add operator to return `Sizes` object `z` such that - /// `z[k] = x[k] - y` for all axis k. --Sizes operator-(const Sizes &x, int64_t y); -+Sizes operator-(const Sizes& x, int64_t y); - - /// Overloaded add operator to return `Sizes` object `z` such that - /// `z[k] = x - y[k]` for all axis k. --Sizes operator-(int64_t x, const Sizes &y); -+Sizes operator-(int64_t x, const Sizes& y); - - /// Overloaded add operator to return `Sizes` object `z` such that - /// `z[k] = x[k] * y[k]` for all axis k. --Sizes operator*(const Sizes &x, const Sizes &y); -+Sizes operator*(const Sizes& x, const Sizes& y); - - /// Overloaded add operator to return `Sizes` object `z` such that - /// `z[k] = x[k] * y` for all axis k. --Sizes operator*(const Sizes &x, int64_t y); -+Sizes operator*(const Sizes& x, int64_t y); - - /// Overloaded add operator to return `Sizes` object `z` such that - /// `z[k] = x * y[k]` for all axis k. --Sizes operator*(int64_t x, const Sizes &y); -+Sizes operator*(int64_t x, const Sizes& y); - - /// Clamp operator to return `Sizes` object `z` such that - /// `z[k] = std::min(std::max(x[k], min[k]), max[k])` for all axis k. --Sizes clamp(const Sizes &min, const Sizes &x, const Sizes &max); -+Sizes clamp(const Sizes& min, const Sizes& x, const Sizes& max); - - /// Clamp operator to return `Sizes` object `z` such that - /// `z[k] = std::min(std::max(x[k], min), max)` for all axis k. --Sizes clamp(int64_t min, const Sizes &x, int64_t max); -+Sizes clamp(int64_t min, const Sizes& x, int64_t max); - - /// Clamp operator to return `Sizes` object `z` such that - /// `z[k] = std::min(std::max(x[k], min), max[k])` for all axis k. --Sizes clamp(int64_t min, const Sizes &x, const Sizes &max); -+Sizes clamp(int64_t min, const Sizes& x, const Sizes& max); - - /// Clamp operator to return `Sizes` object `z` such that - /// `z[k] = std::min(std::max(x[k], min[k]), max)` for all axis k. --Sizes clamp(const Sizes &min, const Sizes &x, int64_t max); -+Sizes clamp(const Sizes& min, const Sizes& x, int64_t max); - - /// Represents index of a tensor. - using Index = Sizes; -@@ -132,23 +132,23 @@ - /// At any point in time, the iterator can either reference an actual index - /// or the past-the-end element in the index space. - /// Dereferencing a past-the-end iterator will result in a fatal error. -- const Index &operator*() const; -- const Index *operator->() const; -+ const Index& operator*() const; -+ const Index* operator->() const; - - /// Compare the iterator to another iterator. - /// Two iterators are equal if they have the same underlying shape and - /// reference the same element in the index space. -- bool operator==(const IndexSpaceIterator &it) { -+ bool operator==(const IndexSpaceIterator& it) { - return shape_ == it.shape_ && index_ == it.index_; - } -- bool operator!=(const IndexSpaceIterator &it) { return !(*this == it); } -+ bool operator!=(const IndexSpaceIterator& it) { return !(*this == it); } - - /// Increment to the next index while iterating over the index space - /// of a tensor in lexicographical order. - /// Incrementing past the last index will result in a past-the-end iterator - /// which cannot be dereferenced. Incrementing even further will result in - /// a fatal error. -- IndexSpaceIterator &operator++(); -+ IndexSpaceIterator& operator++(); - IndexSpaceIterator operator++(int); - - private: -diff --ruN a/stablehlo/stablehlo/reference/Process.cpp b/stablehlo/stablehlo/reference/Process.cpp ---- stablehlo/stablehlo/reference/Process.cpp -+++ stablehlo/stablehlo/reference/Process.cpp -@@ -25,7 +25,7 @@ - namespace mlir { - namespace stablehlo { - --Process::Process(ProcessId id, ProcessGrid *grid) : id_(id), grid_(grid) {} -+Process::Process(ProcessId id, ProcessGrid* grid) : id_(id), grid_(grid) {} - - ProcessGroups Process::crossPartition( - SmallVector> partitionGroups) { -diff --ruN a/stablehlo/stablehlo/reference/Process.h b/stablehlo/stablehlo/reference/Process.h ---- stablehlo/stablehlo/reference/Process.h -+++ stablehlo/stablehlo/reference/Process.h -@@ -29,7 +29,7 @@ - public: - /// \name Constructors - /// @{ -- Process(ProcessId id, ProcessGrid *grid); -+ Process(ProcessId id, ProcessGrid* grid); - /// @} - - /// See `ProcessGrid::crossPartition`. -@@ -72,7 +72,7 @@ - - /// See ProcessGrid. The pointer is used to gain access to allow - /// synchronization among participating processes. -- ProcessGrid *grid_; -+ ProcessGrid* grid_; - }; - - } // namespace stablehlo -diff --ruN a/stablehlo/stablehlo/reference/ProcessGrid.cpp b/stablehlo/stablehlo/reference/ProcessGrid.cpp ---- stablehlo/stablehlo/reference/ProcessGrid.cpp -+++ stablehlo/stablehlo/reference/ProcessGrid.cpp -@@ -39,16 +39,16 @@ - // ProcessId. - //===----------------------------------------------------------------------===// - --bool ProcessId::operator!=(const ProcessId &other) const { -+bool ProcessId::operator!=(const ProcessId& other) const { - return !(*this == other); - } - --bool ProcessId::operator<(const ProcessId &other) const { -+bool ProcessId::operator<(const ProcessId& other) const { - return std::pair{replicaId, partitionId} < - std::pair{other.replicaId, other.partitionId}; - } - --bool ProcessId::operator==(const ProcessId &other) const { -+bool ProcessId::operator==(const ProcessId& other) const { - return std::pair{replicaId, partitionId} == - std::pair{other.replicaId, other.partitionId}; - } -@@ -70,7 +70,7 @@ - //===----------------------------------------------------------------------===// - - RendezvousResult::RendezvousResult( -- std::map> const &results) -+ std::map> const& results) - : results_(results) {} - - void RendezvousResult::insert(ProcessId processId, -@@ -86,12 +86,12 @@ - - SmallVector> RendezvousResult::getSortedTensors() const { - return llvm::map_to_vector(results_, -- [](const auto &pair) { return pair.second; }); -+ [](const auto& pair) { return pair.second; }); - } - - bool RendezvousResult::hasMatchingOperandsCount() const { - auto count = results_.begin()->second.size(); -- for (const auto &it : results_) -+ for (const auto& it : results_) - if (count != it.second.size()) return false; - return true; - } -@@ -101,7 +101,7 @@ - //===----------------------------------------------------------------------===// - - template --V &detail::ThreadSafeMap::operator[](const K &key) { -+V& detail::ThreadSafeMap::operator[](const K& key) { - std::lock_guard lock(lock_); - return map_[key]; - } -@@ -133,7 +133,7 @@ - //===----------------------------------------------------------------------===// - - template --detail::ThreadSafeQueue::ThreadSafeQueue(const std::queue &queue) -+detail::ThreadSafeQueue::ThreadSafeQueue(const std::queue& queue) - : queue_(queue) {} - - template -@@ -155,7 +155,7 @@ - //===----------------------------------------------------------------------===// - - ProcessGrid::ProcessGrid(uint32_t numReplicas, uint32_t numPartitions, -- std::queue &infeed) -+ std::queue& infeed) - : numReplicas_(numReplicas), - numPartitions_(numPartitions), - infeed_(infeed) {} -@@ -163,7 +163,7 @@ - ProcessGroups ProcessGrid::crossPartition( - SmallVector> partitionGroups) { - ProcessGroups processGroups; -- for (const auto &partitionGroup : partitionGroups) { -+ for (const auto& partitionGroup : partitionGroups) { - for (uint32_t replicaId = 0; replicaId < numReplicas_; ++replicaId) { - ProcessGroup processGroup; - for (uint32_t partitionId : partitionGroup) -@@ -177,7 +177,7 @@ - ProcessGroups ProcessGrid::crossReplica( - SmallVector> replicaGroups) { - ProcessGroups processGroups; -- for (const auto &replicaGroup : replicaGroups) { -+ for (const auto& replicaGroup : replicaGroups) { - for (uint32_t partitionId = 0; partitionId < numPartitions_; - ++partitionId) { - ProcessGroup processGroup; -@@ -192,7 +192,7 @@ - ProcessGroups ProcessGrid::crossReplicaAndPartition( - SmallVector> replicaGroups) { - ProcessGroups processGroups; -- for (const auto &replicaGroup : replicaGroups) { -+ for (const auto& replicaGroup : replicaGroups) { - ProcessGroup processGroup; - for (uint32_t partitionId = 0; partitionId < numPartitions_; ++partitionId) - for (uint32_t replicaId : replicaGroup) -@@ -205,7 +205,7 @@ - ProcessGroups ProcessGrid::flattenedIds( - SmallVector> flattenedIdGroups) { - ProcessGroups processGroups; -- for (const auto &flattenedIdGroup : flattenedIdGroups) { -+ for (const auto& flattenedIdGroup : flattenedIdGroups) { - ProcessGroup processGroup; - for (auto flattenedId : flattenedIdGroup) { - uint32_t replicaId = flattenedId / numPartitions_; -@@ -251,7 +251,7 @@ - } - - std::pair channelKey(processGroup, channelId); -- auto &state = channels_[channelKey]; -+ auto& state = channels_[channelKey]; - - std::unique_lock lock(state.mutex); - state.values[processId] = SmallVector(operands); -diff --ruN a/stablehlo/stablehlo/reference/ProcessGrid.h b/stablehlo/stablehlo/reference/ProcessGrid.h ---- stablehlo/stablehlo/reference/ProcessGrid.h -+++ stablehlo/stablehlo/reference/ProcessGrid.h -@@ -41,7 +41,7 @@ - class RendezvousResult { - public: - RendezvousResult() = default; -- RendezvousResult(std::map> const &results); -+ RendezvousResult(std::map> const& results); - - /// Iterates through the (ProcessId, SmallVector) map entires and - /// returns a vector of Tensors sorted by ProcessId--(replicaId, partitionId) -@@ -104,7 +104,7 @@ - class ThreadSafeMap { - public: - /// Returns a reference to the data associated with the `key`. -- V &operator[](const K &key); -+ V& operator[](const K& key); - - private: - /// Synchronization primitive used to manage concurrent access to the map. -@@ -143,7 +143,7 @@ - /// \name Constructors - /// @{ - ThreadSafeQueue() = default; -- ThreadSafeQueue(const std::queue &queue); -+ ThreadSafeQueue(const std::queue& queue); - /// @} - - /// Remove the first element of the queue and return it. -@@ -173,16 +173,16 @@ - uint32_t partitionId; - - /// Overloaded inequality operator. -- bool operator!=(const ProcessId &other) const; -+ bool operator!=(const ProcessId& other) const; - - /// The sort order for ProcessId is not defined in StableHLO, and it's - /// internally used in ProcessGrid::rendezvous as part of a sorted key on the - /// map. This operator is conveniently used to help define the ordering since - /// ordering is defined for StableHLO process group. -- bool operator<(const ProcessId &other) const; -+ bool operator<(const ProcessId& other) const; - - /// Overloaded equality operator. -- bool operator==(const ProcessId &other) const; -+ bool operator==(const ProcessId& other) const; - }; - - /// StableHLO `process_group`. -@@ -190,9 +190,9 @@ - public: - using SmallVector::SmallVector; - -- bool operator<(const ProcessGroup &other) const { -- return static_cast(*this) < -- static_cast(other); -+ bool operator<(const ProcessGroup& other) const { -+ return static_cast(*this) < -+ static_cast(other); - } - }; - -@@ -211,7 +211,7 @@ - /// \name Constructors - /// @{ - ProcessGrid(uint32_t numReplicas, uint32_t numPartitions, -- std::queue &infeed); -+ std::queue& infeed); - /// @} - - /// StableHLO `cross_partition` communication strategy. -diff --ruN a/stablehlo/stablehlo/reference/Token.cpp b/stablehlo/stablehlo/reference/Token.cpp ---- stablehlo/stablehlo/reference/Token.cpp -+++ stablehlo/stablehlo/reference/Token.cpp -@@ -23,11 +23,11 @@ - namespace mlir { - namespace stablehlo { - --Token::Token(MLIRContext *context) : type_(TokenType::get(context)) {} -+Token::Token(MLIRContext* context) : type_(TokenType::get(context)) {} - - TokenType Token::getType() const { return type_; } - --void Token::print(raw_ostream &os) const { getType().print(os); } -+void Token::print(raw_ostream& os) const { getType().print(os); } - - void Token::dump() const { print(llvm::errs()); } - -diff --ruN a/stablehlo/stablehlo/reference/Token.h b/stablehlo/stablehlo/reference/Token.h ---- stablehlo/stablehlo/reference/Token.h -+++ stablehlo/stablehlo/reference/Token.h -@@ -28,14 +28,14 @@ - public: - /// \name Constructors - /// @{ -- Token(MLIRContext *context); -+ Token(MLIRContext* context); - /// @} - - /// Returns the type of the Token object. - TokenType getType() const; - - /// Prints Token object. -- void print(raw_ostream &os) const; -+ void print(raw_ostream& os) const; - void dump() const; - - private: -@@ -43,7 +43,7 @@ - }; - - /// Print utilities for Token objects. --inline raw_ostream &operator<<(raw_ostream &os, Token Token) { -+inline raw_ostream& operator<<(raw_ostream& os, Token Token) { - Token.print(os); - return os; - } -diff --ruN a/stablehlo/stablehlo/reference/Value.cpp b/stablehlo/stablehlo/reference/Value.cpp ---- stablehlo/stablehlo/reference/Value.cpp -+++ stablehlo/stablehlo/reference/Value.cpp -@@ -46,7 +46,7 @@ - - TupleType Tuple::getType() const { return type_; } - --void Tuple::print(raw_ostream &os) const { -+void Tuple::print(raw_ostream& os) const { - getType().print(os); - os << " (\n"; - for (size_t i = 0; i < values_.size(); ++i) { -@@ -63,9 +63,9 @@ - // InterpreterValue. - //===----------------------------------------------------------------------===// - --InterpreterValue::InterpreterValue(const Tensor &tensor) : value_(tensor) {} --InterpreterValue::InterpreterValue(const Token &token) : value_(token) {} --InterpreterValue::InterpreterValue(const Tuple &tuple) : value_(tuple) {} -+InterpreterValue::InterpreterValue(const Tensor& tensor) : value_(tensor) {} -+InterpreterValue::InterpreterValue(const Token& token) : value_(token) {} -+InterpreterValue::InterpreterValue(const Tuple& tuple) : value_(tuple) {} - - Tensor InterpreterValue::getTensor() const { - if (!isTensor()) -@@ -108,7 +108,7 @@ - return std::holds_alternative(value_); - } - --void InterpreterValue::print(raw_ostream &os) const { -+void InterpreterValue::print(raw_ostream& os) const { - if (isTensor()) - getTensor().print(os); - else if (isToken()) diff --ruN a/stablehlo/stablehlo/testdata/BUILD.bazel b/stablehlo/stablehlo/testdata/BUILD.bazel --- stablehlo/stablehlo/testdata/BUILD.bazel +++ stablehlo/stablehlo/testdata/BUILD.bazel -@@ -11,14 +11,27 @@ +@@ -11,13 +11,26 @@ # 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. @@ -3948,7 +2832,7 @@ diff --ruN a/stablehlo/stablehlo/testdata/BUILD.bazel b/stablehlo/stablehlo/test + ], licenses = ["notice"], ) - ++ +# copybara:uncomment_begin +# filegroup( +# name = "all_files", @@ -3956,10 +2840,9 @@ diff --ruN a/stablehlo/stablehlo/testdata/BUILD.bazel b/stablehlo/stablehlo/test +# visibility = ["//:__subpackages__"], +# ) +# copybara:uncomment_end -+ + # Equivalent of configure_lit_site_cfg from CMakeLists.txt. expand_template( - name = "lit_site_cfg_py_gen", @@ -39,46 +52,42 @@ template = "lit.site.cfg.py.in", ) @@ -4043,7 +2926,7 @@ diff --ruN a/stablehlo/stablehlo/testdata/BUILD.bazel b/stablehlo/stablehlo/test diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/BUILD.bazel --- stablehlo/stablehlo/tests/BUILD.bazel +++ stablehlo/stablehlo/tests/BUILD.bazel -@@ -11,16 +11,30 @@ +@@ -11,15 +11,29 @@ # 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. @@ -4064,7 +2947,7 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B + ], licenses = ["notice"], ) - ++ +# copybara:uncomment_begin +# filegroup( +# name = "all_files", @@ -4072,10 +2955,9 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B +# visibility = ["//:__subpackages__"], +# ) +# copybara:uncomment_end -+ + cc_library( name = "check_ops", - srcs = [ @@ -29,6 +43,7 @@ hdrs = [ "CheckOps.h", @@ -4112,12 +2994,13 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B gentbl_cc_library( name = "test_utils_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", +- strip_include_prefix = ".", - tbl_outs = {"TestUtils.h.inc": [ - "-gen-pass-decls", - "-name=HloTest", - ]}, ++ compatible_with = get_compatible_with_portable(), ++ strip_include_prefix = ".", + tbl_outs = { + "TestUtils.h.inc": [ + "-gen-pass-decls", @@ -4162,10 +3045,8 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B - ) - for src in glob(["**/*.mlir"]) -] -- --test_suite( +lit_test_suite( - name = "stablehlo_tests", ++ name = "stablehlo_tests", + srcs = glob([ + "*.mlir", + "**/*.mlir", @@ -4179,116 +3060,18 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B + "@llvm-project//llvm:not", + ], + size_per_test = "small", - tags = ["stablehlo_tests"], - ) -+ -+test_suite( ++ tags = ["stablehlo_tests"], ++) + + test_suite( +- name = "stablehlo_tests", +- tags = ["stablehlo_tests"], +-) + name = "all_tests", # alias for `:stablehlo_tests` + tests = [ + "//stablehlo/tests:stablehlo_tests", + ], +) -diff --ruN a/stablehlo/stablehlo/tests/CheckOps.cpp b/stablehlo/stablehlo/tests/CheckOps.cpp ---- stablehlo/stablehlo/tests/CheckOps.cpp -+++ stablehlo/stablehlo/tests/CheckOps.cpp -@@ -72,7 +72,7 @@ - // Check Dialect Constructor - //===----------------------------------------------------------------------===// - --CheckDialect::CheckDialect(MLIRContext *context) -+CheckDialect::CheckDialect(MLIRContext* context) - : Dialect(getDialectNamespace(), context, TypeID::get()) { - addOperations< - #define GET_OP_LIST -@@ -80,13 +80,13 @@ - >(); - } - --llvm::Error evalExpectAlmostEqConstOp(const Tensor &lhs, ElementsAttr value, -+llvm::Error evalExpectAlmostEqConstOp(const Tensor& lhs, ElementsAttr value, - APFloat tolerance) { - auto rhs = makeTensor(cast(value)); - return evalExpectAlmostEqOp(lhs, rhs, tolerance); - } - --llvm::Error evalExpectAlmostEqOp(const Tensor &lhs, const Tensor &rhs, -+llvm::Error evalExpectAlmostEqOp(const Tensor& lhs, const Tensor& rhs, - APFloat tolerance) { - for (auto lhsIt = lhs.index_begin(), rhsIt = rhs.index_begin(); - lhsIt != lhs.index_end(); ++lhsIt, ++rhsIt) -@@ -101,12 +101,12 @@ - return llvm::Error::success(); - } - --llvm::Error evalExpectEqConstOp(const Tensor &lhs, ElementsAttr value) { -+llvm::Error evalExpectEqConstOp(const Tensor& lhs, ElementsAttr value) { - auto rhs = makeTensor(cast(value)); - return evalExpectEqOp(lhs, rhs); - } - --llvm::Error evalExpectEqOp(const Tensor &lhs, const Tensor &rhs) { -+llvm::Error evalExpectEqOp(const Tensor& lhs, const Tensor& rhs) { - for (auto lhsIt = lhs.index_begin(), rhsIt = rhs.index_begin(); - lhsIt != lhs.index_end(); ++lhsIt, ++rhsIt) - if ((lhs.get(*lhsIt) != rhs.get(*rhsIt)).getBooleanValue()) -@@ -146,7 +146,7 @@ - return llvm::errc::bad_address; - } - --llvm::Error evalExpectSerializedEqOp(const Tensor &expected, StringRef probeId, -+llvm::Error evalExpectSerializedEqOp(const Tensor& expected, StringRef probeId, - StringRef probeDir, uint32_t iteration) { - auto serializedMetadata = - getSerializedTensorMetadata(probeId, probeDir, iteration); -@@ -206,7 +206,7 @@ - return std::numeric_limits::max(); - } - --static uint64_t ULPDifference(const Element &e1, const Element &e2) { -+static uint64_t ULPDifference(const Element& e1, const Element& e2) { - // caller is responsible for ensuring that e1, e2 have both the same - // float or complex types - if (isSupportedComplexType(e1.getType())) { -@@ -218,7 +218,7 @@ - return ULPDifference(e1.getFloatValue(), e2.getFloatValue()); - } - --llvm::Error evalExpectCloseOp(const Tensor &actual, const Tensor &expected, -+llvm::Error evalExpectCloseOp(const Tensor& actual, const Tensor& expected, - uint64_t min_ulp_difference, - uint64_t max_ulp_difference) { - auto type = actual.getElementType(); -diff --ruN a/stablehlo/stablehlo/tests/CheckOps.h b/stablehlo/stablehlo/tests/CheckOps.h ---- stablehlo/stablehlo/tests/CheckOps.h -+++ stablehlo/stablehlo/tests/CheckOps.h -@@ -32,20 +32,20 @@ - - class CheckDialect : public Dialect { - public: -- explicit CheckDialect(MLIRContext *context); -+ explicit CheckDialect(MLIRContext* context); - static StringRef getDialectNamespace() { return "check"; } - }; - - // The eval functions for the following ops are used only for test harness. --llvm::Error evalExpectAlmostEqConstOp(const Tensor &lhs, ElementsAttr value, -+llvm::Error evalExpectAlmostEqConstOp(const Tensor& lhs, ElementsAttr value, - APFloat tolerance); --llvm::Error evalExpectAlmostEqOp(const Tensor &lhs, const Tensor &rhs, -+llvm::Error evalExpectAlmostEqOp(const Tensor& lhs, const Tensor& rhs, - APFloat tolerance); --llvm::Error evalExpectEqConstOp(const Tensor &lhs, ElementsAttr value); --llvm::Error evalExpectEqOp(const Tensor &lhs, const Tensor &rhs); --llvm::Error evalExpectSerializedEqOp(const Tensor &expected, StringRef probeId, -+llvm::Error evalExpectEqConstOp(const Tensor& lhs, ElementsAttr value); -+llvm::Error evalExpectEqOp(const Tensor& lhs, const Tensor& rhs); -+llvm::Error evalExpectSerializedEqOp(const Tensor& expected, StringRef probeId, - StringRef probeDir, uint32_t iteration); --llvm::Error evalExpectCloseOp(const Tensor &actual, const Tensor &expected, -+llvm::Error evalExpectCloseOp(const Tensor& actual, const Tensor& expected, - uint64_t min_ulp_difference, - uint64_t max_ulp_difference); - diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir --- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir +++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir @@ -4541,10 +3324,14 @@ diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir --- stablehlo/stablehlo/tests/ops_chlo.mlir +++ stablehlo/stablehlo/tests/ops_chlo.mlir -@@ -489,6 +489,50 @@ - - // ----- - +@@ -485,6 +485,50 @@ + stablehlo.return %1, %2, %1, %2 : tensor<3xf32>, tensor<3xi32>, tensor<3xf32>, tensor<3xi32> + } : (tensor<2x3xf32>, tensor<2x3xi32>, tensor<3xf32>, tensor<3xi32>) -> (tensor<2x3xf32>, tensor<2x3xi32>, tensor<3xf32>, tensor<3xi32>) + func.return %0#0, %0#1 : tensor<2x3xf32>, tensor<2x3xi32> ++} ++ ++// ----- ++ +// CHECK-LABEL: func @scan_bounded_dynamism +func.func @scan_bounded_dynamism(%arg0: tensor>, %arg1: tensor>) -> (tensor>, tensor>) { + // CHECK: chlo.scan @@ -4585,13 +3372,9 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests + stablehlo.return %1, %2, %1, %2 : tensor>, tensor>, tensor>, tensor> + } : (tensor>, tensor>, tensor>, tensor>) -> (tensor>, tensor>, tensor>, tensor>) + func.return %0#0, %0#1, %0#2, %0#3 : tensor>, tensor>, tensor>, tensor> -+} -+ -+// ----- -+ - func.func @scan_explicit_size_mismatch(%arg0: tensor<2x3xf32>, %arg1: tensor<3xf32>) -> tensor<2x3xf32> { - // expected-error @+1 {{'chlo.scan' op invalid scan dimension size of operand 0}} - %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 attributes { scan_dim_size = 3 : i64 } { + } + + // ----- @@ -561,6 +605,30 @@ // ----- @@ -4778,181 +3561,13 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir b %2 = stablehlo.dynamic_iota %1, dim = 0 : (tensor<1xi64>) -> tensor return %2 : tensor } -diff --ruN a/stablehlo/stablehlo/tools/StablehloLspServerMain.cpp b/stablehlo/stablehlo/tools/StablehloLspServerMain.cpp ---- stablehlo/stablehlo/tools/StablehloLspServerMain.cpp -+++ stablehlo/stablehlo/tools/StablehloLspServerMain.cpp -@@ -19,7 +19,7 @@ - #include "mlir/Tools/mlir-lsp-server/MlirLspServerMain.h" - #include "stablehlo/dialect/Register.h" - --int main(int argc, char **argv) { -+int main(int argc, char** argv) { - mlir::DialectRegistry registry; - mlir::registerAllDialects(registry); - mlir::registerAllExtensions(registry); -diff --ruN a/stablehlo/stablehlo/tools/StablehloOptMain.cpp b/stablehlo/stablehlo/tools/StablehloOptMain.cpp ---- stablehlo/stablehlo/tools/StablehloOptMain.cpp -+++ stablehlo/stablehlo/tools/StablehloOptMain.cpp -@@ -28,7 +28,7 @@ - #include "stablehlo/transforms/Passes.h" - #include "stablehlo/transforms/optimization/Passes.h" - --int main(int argc, char **argv) { -+int main(int argc, char** argv) { - mlir::registerAllPasses(); - mlir::hlo::registerAllTestPasses(); - mlir::stablehlo::registerPassPipelines(); -diff --ruN a/stablehlo/stablehlo/tools/StablehloTranslateMain.cpp b/stablehlo/stablehlo/tools/StablehloTranslateMain.cpp ---- stablehlo/stablehlo/tools/StablehloTranslateMain.cpp -+++ stablehlo/stablehlo/tools/StablehloTranslateMain.cpp -@@ -92,7 +92,7 @@ - - namespace { - --stablehlo::Tensor makeBooleanTensor(MLIRContext *context, bool value) { -+stablehlo::Tensor makeBooleanTensor(MLIRContext* context, bool value) { - auto builder = Builder(context); - auto type = RankedTensorType::get({}, builder.getI1Type()); - auto res = DenseElementsAttr::get(type, builder.getBoolAttr(true)); -@@ -104,7 +104,7 @@ - // --args=[dense<1> : tensor<2xi32>, ...], where each dense attribute is - // interpreted as a tensor. - mlir::FailureOr> --parseInterpreterArguments(std::string argsStr, MLIRContext *context) { -+parseInterpreterArguments(std::string argsStr, MLIRContext* context) { - llvm::SmallVector inputs; - auto parseError = [&](llvm::StringRef msg) { - std::string usage = "--args=[dense<1> : tensor<2xi32>, ...]"; -@@ -129,7 +129,7 @@ - } - - llvm::Error evalCustomCallCheckEq(stablehlo::CustomCallOp op, -- stablehlo::Scope &scope) { -+ stablehlo::Scope& scope) { - if (op->getNumOperands() != 2) - return stablehlo::invalidArgument("Unsupported op: %s", - debugString(op).c_str()); -@@ -157,10 +157,10 @@ - : public stablehlo::InterpreterFallback { - public: - StablehloTranslateInterpreterFallback( -- const std::string &probeInstrumentationDir) -+ const std::string& probeInstrumentationDir) - : probeInstrumentationDir(probeInstrumentationDir) {} -- virtual llvm::Error operator()(Operation &op, stablehlo::Scope &scope, -- stablehlo::Process *process) final { -+ virtual llvm::Error operator()(Operation& op, stablehlo::Scope& scope, -+ stablehlo::Process* process) final { - llvm::StringRef funcName = op.getParentOfType().getSymName(); - if (auto customCall = dyn_cast(op)) { - auto callTarget = customCall.getCallTargetName(); -@@ -271,7 +271,7 @@ - - TranslateFromMLIRRegistration interpretRegistration( - "interpret", "Interpreter for StableHLO", -- [](ModuleOp module, raw_ostream &os) -> LogicalResult { -+ [](ModuleOp module, raw_ostream& os) -> LogicalResult { - stablehlo::InterpreterConfiguration config; - config.probeInstrumentationDir = probeOutputDir.getValue(); - config.fallback = std::make_unique( -@@ -286,21 +286,21 @@ - - if (interpreterPrintDense.getValue()) { - DenseElementsAttr denseResult; -- for (auto &result : *results) { -+ for (auto& result : *results) { - denseResult = makeDenseElementsAttr(result.getTensor()); - os << denseResult << '\n'; - } - return success(); - } - -- for (auto &result : *results) { -+ for (auto& result : *results) { - result.print(os); - os << '\n'; - } - - return success(); - }, -- [](DialectRegistry ®istry) { -+ [](DialectRegistry& registry) { - registry.insert(); - registry.insert(); - registry.insert(); -@@ -310,7 +310,7 @@ - - TranslateFromMLIRRegistration serializeRegistration( - "serialize", "Serialize StableHLO program into a portable artifact", -- [](ModuleOp module, raw_ostream &os) -> LogicalResult { -+ [](ModuleOp module, raw_ostream& os) -> LogicalResult { - std::string targetVersion = targetOption.getValue(); - if (targetVersion == "current") - targetVersion = vhlo::Version::getCurrentVersion().toString(); -@@ -325,14 +325,14 @@ - return stablehlo::serializePortableArtifact( - module, targetVersion, os, allowOtherDialectsOption.getValue()); - }, -- [](DialectRegistry ®istry) { -+ [](DialectRegistry& registry) { - mlir::registerAllDialects(registry); - mlir::stablehlo::registerAllDialects(registry); - }); - - TranslateToMLIRRegistration deserializeRegistration( - "deserialize", "Deserialize a portable artifact into a StableHLO program", -- [](llvm::StringRef input, mlir::MLIRContext *context) { -+ [](llvm::StringRef input, mlir::MLIRContext* context) { - if (printStablehloVersion.getValue()) { - auto version = stablehlo::getPortableArtifactVersion(input); - if (failed(version)) { -@@ -346,14 +346,14 @@ - } - return stablehlo::deserializePortableArtifact(input, context); - }, -- [](DialectRegistry ®istry) { -+ [](DialectRegistry& registry) { - mlir::registerAllDialects(registry); - mlir::stablehlo::registerAllDialects(registry); - }); - - } // namespace mlir - --int main(int argc, char **argv) { -+int main(int argc, char** argv) { - return failed( - mlir::mlirTranslateMain(argc, argv, "StableHLO interpreter driver\n")); - } -diff --ruN a/stablehlo/stablehlo/transforms/ChloDecompositionUtils.h b/stablehlo/stablehlo/transforms/ChloDecompositionUtils.h ---- stablehlo/stablehlo/transforms/ChloDecompositionUtils.h -+++ stablehlo/stablehlo/transforms/ChloDecompositionUtils.h -@@ -22,13 +22,13 @@ - - // Utility functions used in the Chlo to stablehlo legalization. - --Value materializeLgamma(OpBuilder &rewriter, Location loc, ValueRange args); -+Value materializeLgamma(OpBuilder& rewriter, Location loc, ValueRange args); - --Value materializeDigamma(OpBuilder &rewriter, Location loc, ValueRange args); -+Value materializeDigamma(OpBuilder& rewriter, Location loc, ValueRange args); - --Value materializeZeta(OpBuilder &rewriter, Location loc, ValueRange args); -+Value materializeZeta(OpBuilder& rewriter, Location loc, ValueRange args); - --Value materializePolygamma(OpBuilder &rewriter, Location loc, ValueRange args); -+Value materializePolygamma(OpBuilder& rewriter, Location loc, ValueRange args); - - } // namespace stablehlo - } // namespace mlir diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp --- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp +++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp -@@ -2930,8 +2930,18 @@ +@@ -2930,8 +2930,19 @@ SmallVector sliceShape(inputType.getShape().begin(), inputType.getShape().end()); sliceShape[dim] = 1; -- auto sliceType = -- RankedTensorType::get(sliceShape, inputType.getElementType()); + Attribute sliceEncoding; + ArrayRef inputBounds = + hlo::encodingToBounds(inputType.getEncoding()); @@ -4963,12 +3578,14 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable + sliceEncoding = + hlo::boundsToEncoding(inputType.getEncoding(), sliceBounds); + } -+ auto sliceType = RankedTensorType::get( -+ sliceShape, inputType.getElementType(), sliceEncoding); + auto sliceType = +- RankedTensorType::get(sliceShape, inputType.getElementType()); ++ RankedTensorType::get(sliceShape, inputType.getElementType(), ++ sliceEncoding); Value slice; if (inputType.hasStaticShape()) { -@@ -2980,9 +2990,18 @@ +@@ -2980,9 +2991,18 @@ for (int64_t d = 0; d < rank; ++d) { if (d != dim) resultShape.push_back(sliceShape[d]); } @@ -4988,7 +3605,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable slice); slicedInputs.push_back(reshaped); } -@@ -3004,9 +3023,20 @@ +@@ -3004,9 +3024,20 @@ SmallVector elementShape(outputType.getShape().begin(), outputType.getShape().end()); elementShape[dim] = 1; @@ -5010,207 +3627,4 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable newElement); int64_t outRank = outputType.getRank(); -diff --ruN a/stablehlo/stablehlo/transforms/PassPipelines.cpp b/stablehlo/stablehlo/transforms/PassPipelines.cpp ---- stablehlo/stablehlo/transforms/PassPipelines.cpp -+++ stablehlo/stablehlo/transforms/PassPipelines.cpp -@@ -24,7 +24,7 @@ - namespace mlir { - namespace stablehlo { - --void createStablehloDeserializePipeline(OpPassManager &pm) { -+void createStablehloDeserializePipeline(OpPassManager& pm) { - // Convert VHLO(version x.y.z) --> VHLO(current). - pm.addPass(stablehlo::createVhloToVersionPass( - {vhlo::Version::getCurrentVersion().toString()})); -@@ -33,14 +33,14 @@ - pm.addPass(stablehlo::createVhloLegalizeToStablehloPass()); - } - --void createChloPreSerializationPipeline(OpPassManager &pm) { -+void createChloPreSerializationPipeline(OpPassManager& pm) { - // CHLO --> StableHLO+Shape --> StableHLO - pm.addPass(stablehlo::createChloLegalizeToStablehloPass()); - pm.addPass(stablehlo::createShapeLegalizeToStablehloPass()); - pm.addPass(mlir::createReconcileUnrealizedCastsPass()); - } - --void createStablehloRemoveDynamismPipeline(OpPassManager &pm, -+void createStablehloRemoveDynamismPipeline(OpPassManager& pm, - TypeRange refinedTypes) { - pm.addPass(stablehlo::createStablehloRefineArgumentsPass(refinedTypes)); - pm.addPass(stablehlo::createStablehloRefineShapesPass()); -@@ -48,7 +48,7 @@ - stablehlo::createStablehloCanonicalizeDynamismPass()); - } - --void createStablehloLowerQuantPipeline(OpPassManager &pm) { -+void createStablehloLowerQuantPipeline(OpPassManager& pm) { - pm.addNestedPass( - stablehlo::createStablehloLegalizeQuantToMathPass()); - pm.addNestedPass( -diff --ruN a/stablehlo/stablehlo/transforms/PassUtils.h b/stablehlo/stablehlo/transforms/PassUtils.h ---- stablehlo/stablehlo/transforms/PassUtils.h -+++ stablehlo/stablehlo/transforms/PassUtils.h -@@ -33,7 +33,7 @@ - // Utility functions common across passes. - - template --Attribute getScalarLike(OpBuilder &b, T constant, Type type) { -+Attribute getScalarLike(OpBuilder& b, T constant, Type type) { - Type element = getElementTypeOrSelf(type); - if (isa(element)) return b.getIntegerAttr(element, constant); - if (isa(element)) return b.getFloatAttr(element, constant); -@@ -47,13 +47,13 @@ - // in `scalar`. - // Returns stablehlo::ConstantOp if value type if static, else returns - // chlo::ConstantLikeOp. --Value getConstantLikeImpl(OpBuilder &b, Location loc, Attribute scalar, -+Value getConstantLikeImpl(OpBuilder& b, Location loc, Attribute scalar, - Value val); - - // Creates a chlo::ConstantLikeOp using a splat `constant` of the same shape - // as `val`. - template --Value getConstantLike(OpBuilder &b, Location loc, T constant, Value val) { -+Value getConstantLike(OpBuilder& b, Location loc, T constant, Value val) { - auto shapedTy = cast(val.getType()); - Attribute scalar = getScalarLike(b, constant, shapedTy); - return getConstantLikeImpl(b, loc, scalar, val); -@@ -63,7 +63,7 @@ - // same shape as `val`. - // The distinction between double and APFloat causes issues so need this - // explicit template specialization. --Value getConstantLike(OpBuilder &b, Location loc, const APFloat &constant, -+Value getConstantLike(OpBuilder& b, Location loc, const APFloat& constant, - Value val); - - // Check if any of the given types are mlir::quant::QuantizedType. -diff --ruN a/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp b/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp ---- stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp -+++ stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp -@@ -24,14 +24,14 @@ - - namespace { - --static Value getConstantLikeMaxFiniteValue(OpBuilder &b, Location loc, -+static Value getConstantLikeMaxFiniteValue(OpBuilder& b, Location loc, - Value val) { - auto ty = cast(getElementTypeOrSelf(val.getType())); - return getConstantLike( - b, loc, llvm::APFloat::getLargest(ty.getFloatSemantics()), val); - } - --static Value getConstantLikeInfValue(OpBuilder &b, Location loc, Value val, -+static Value getConstantLikeInfValue(OpBuilder& b, Location loc, Value val, - bool negative) { - auto ty = cast(getElementTypeOrSelf(val.getType())); - return getConstantLike( -@@ -50,7 +50,7 @@ - StablehloComplexMathExpanderPass>() {} - - public: -- LogicalResult initialize(MLIRContext *context) override { -+ LogicalResult initialize(MLIRContext* context) override { - config.setUseTopDownTraversal(true); - RewritePatternSet patterns_(context); - populateStablehloComplexMathExpanderPatterns(context, &patterns_); -@@ -76,8 +76,8 @@ - - } // namespace - --void populateStablehloComplexMathExpanderPatterns(MLIRContext *context, -- RewritePatternSet *patterns) { -+void populateStablehloComplexMathExpanderPatterns(MLIRContext* context, -+ RewritePatternSet* patterns) { - populateWithGenerated(*patterns); - } - -diff --ruN a/stablehlo/stablehlo/transforms/optimization/Passes.h b/stablehlo/stablehlo/transforms/optimization/Passes.h ---- stablehlo/stablehlo/transforms/optimization/Passes.h -+++ stablehlo/stablehlo/transforms/optimization/Passes.h -@@ -47,33 +47,33 @@ - - /// Collection of canonicalization patterns for StableHLO. - void populateStablehloCanonicalizationPatterns( -- MLIRContext *context, RewritePatternSet *patterns, -- const StablehloAggressiveSimplificationPassOptions &options, -+ MLIRContext* context, RewritePatternSet* patterns, -+ const StablehloAggressiveSimplificationPassOptions& options, - PatternBenefit benefit = 1); - --void populateStablehloCanonicalizationPatterns(MLIRContext *context, -- RewritePatternSet *patterns, -+void populateStablehloCanonicalizationPatterns(MLIRContext* context, -+ RewritePatternSet* patterns, - PatternBenefit benefit = 1); - - /// Collection of folding patterns for StableHLO. - void populateStablehloAggressiveFolderPatterns( -- MLIRContext *context, RewritePatternSet *patterns, -- const StablehloAggressiveFolderPassOptions &options, -+ MLIRContext* context, RewritePatternSet* patterns, -+ const StablehloAggressiveFolderPassOptions& options, - PatternBenefit benefit = 1); - --void populateStablehloAggressiveFolderPatterns(MLIRContext *context, -- RewritePatternSet *patterns, -+void populateStablehloAggressiveFolderPatterns(MLIRContext* context, -+ RewritePatternSet* patterns, - PatternBenefit benefit = 1); - - /// A subset of folding patterns for StableHLO that is necessary for shape - /// refinement. - void populateStablehloShapeFolderPatterns( -- MLIRContext *context, RewritePatternSet *patterns, -- const StablehloAggressiveFolderPassOptions &options, -+ MLIRContext* context, RewritePatternSet* patterns, -+ const StablehloAggressiveFolderPassOptions& options, - PatternBenefit benefit = 1); - --void populateStablehloShapeFolderPatterns(MLIRContext *context, -- RewritePatternSet *patterns, -+void populateStablehloShapeFolderPatterns(MLIRContext* context, -+ RewritePatternSet* patterns, - PatternBenefit benefit = 1); - - /// Some workloads in XLA import StableHLO from HLO. Since there are a few -@@ -81,28 +81,28 @@ - /// set of patterns brings the imported HLO back to a more canonical form - /// without applying a full set of graph simplifications. - void populateStablehloHloImportCanonicalizationPatterns( -- MLIRContext *context, RewritePatternSet *patterns, -- const StablehloAggressiveSimplificationPassOptions &options); -+ MLIRContext* context, RewritePatternSet* patterns, -+ const StablehloAggressiveSimplificationPassOptions& options); - - void populateStablehloHloImportCanonicalizationPatterns( -- MLIRContext *context, RewritePatternSet *patterns); -+ MLIRContext* context, RewritePatternSet* patterns); - - // If the `options` argument is provided, ensure it's not a temporary. - void populateStablehloCanonicalizationPatterns( -- MLIRContext *context, RewritePatternSet *patterns, -- StablehloAggressiveSimplificationPassOptions &&options, -+ MLIRContext* context, RewritePatternSet* patterns, -+ StablehloAggressiveSimplificationPassOptions&& options, - PatternBenefit benefit = 1) = delete; - void populateStablehloShapeFolderPatterns( -- MLIRContext *context, RewritePatternSet *patterns, -- StablehloAggressiveFolderPassOptions &&options, -+ MLIRContext* context, RewritePatternSet* patterns, -+ StablehloAggressiveFolderPassOptions&& options, - PatternBenefit benefit = 1) = delete; - void populateStablehloAggressiveFolderPatterns( -- MLIRContext *context, RewritePatternSet *patterns, -- StablehloAggressiveFolderPassOptions &&options, -+ MLIRContext* context, RewritePatternSet* patterns, -+ StablehloAggressiveFolderPassOptions&& options, - PatternBenefit benefit = 1) = delete; - void populateStablehloHloImportCanonicalizationPatterns( -- MLIRContext *context, RewritePatternSet *patterns, -- StablehloAggressiveSimplificationPassOptions &&options) = delete; -+ MLIRContext* context, RewritePatternSet* patterns, -+ StablehloAggressiveSimplificationPassOptions&& options) = delete; - - } // namespace stablehlo - } // namespace mlir + From 3ca01955a693e27f565fd24bf922cab8c0e8c896 Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Mon, 21 Sep 2026 11:45:05 -0700 Subject: [PATCH 2/3] Extend the native custom-call emitter context with shapes, kernel arguments and typed attributes Writing a native custom-call handler requires hand-rolling a large amount of boilerplate that XLA is in a better position to own, and several of those hand-rolled pieces are easy to get subtly wrong. This change moves that work into `NativeCustomCallEmitterContext`, which gains: - `GetDeviceDescription()`. Handlers previously had to dig the device description out of `GetTargetTopology()->gpu_target_config()`, guarding against a target config that, in this code path, is always present. - `GetResultShapedSlice()` / `GetOperandShapedSlice()`, which return the `ShapedSlice` that thunks actually consume, rather than a bare slice the handler then has to pair up with a shape by hand. - `CreateKernelArguments()`, which builds the conventional kernel argument list. Assembling `emitters::KernelArgument`s by hand leaves `slice_index_` uninitialized, skips alignment and aliasing computation, and requires the author to set `written` correctly - and `written` drives thunk scheduling, so getting it wrong produces races rather than an obvious failure. The implementation resolves buffers through `ThunkEmitter`'s allocation overrides, which a handler could not have done itself. - `GetFfiAttributes()` now returns a typed `ffi::Attributes` instead of an `AttributesMap`, so reading an attribute is `attrs.Get("epsilon")` rather than two layers of hand-written `std::get_if` on nested variants. The demonstrator handler is updated for the new signature. PiperOrigin-RevId: 985413447 --- .../libraries/native_custom_call_thunks/BUILD | 11 ++--- .../native_custom_call_emitter_context.h | 48 +++++++++++++++++- ...ative_custom_call_handler_registry_test.cc | 44 ----------------- .../write_value_thunk_folded/BUILD | 4 +- .../write_value_thunk_folded_handler.cc | 41 +++++----------- third_party/xla/xla/service/gpu/BUILD | 1 + .../xla/xla/service/gpu/thunk_emitter.cc | 49 +++++++++++++++++-- 7 files changed, 109 insertions(+), 89 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/BUILD b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/BUILD index 9b3000857b2076..0de5626a4e7952 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/BUILD +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/BUILD @@ -53,10 +53,14 @@ cc_library( "//xla:shape_util", "//xla:xla_proto_cc", "//xla/backends/gpu/runtime:thunk", - "//xla/ffi:attribute_map", + "//xla/codegen/emitters:kernel_arguments", + "//xla/ffi:attributes", "//xla/service:buffer_assignment", "//xla/service:gpu_topology", + "//xla/service:shaped_slice", + "//xla/stream_executor:device_description", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/types:span", ], ) @@ -98,13 +102,8 @@ xla_cc_test( ":native_custom_call_emitter_context", ":native_custom_call_handler_registration", ":native_custom_call_handler_registry", - "//xla:shape_util", "//xla/backends/gpu/runtime:thunk", - "//xla/backends/gpu/runtime:thunk_id", - "//xla/ffi:attribute_map", "//xla/hlo/ir:hlo", - "//xla/service:buffer_assignment", - "//xla/service:gpu_topology", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_emitter_context.h b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_emitter_context.h index 8cce838f2755ce..5a140b0687f307 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_emitter_context.h +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_emitter_context.h @@ -19,11 +19,16 @@ limitations under the License. #include #include "absl/status/statusor.h" +#include "absl/types/span.h" #include "xla/backends/gpu/runtime/thunk.h" -#include "xla/ffi/attribute_map.h" +#include "xla/codegen/emitters/kernel_arguments.h" +#include "xla/ffi/attributes.h" #include "xla/service/buffer_assignment.h" #include "xla/service/gpu_topology.h" +#include "xla/service/shaped_slice.h" +#include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/stream_executor/device_description.h" #include "xla/xla.pb.h" namespace xla::gpu { @@ -44,17 +49,56 @@ class NativeCustomCallEmitterContext { virtual const GpuTopology& GetTargetTopology() const = 0; + // The device the module is being compiled for. Unlike an FFI handler, which + // can query the stream's device at instantiation time, a native handler runs + // ahead of time and must take the target description from here. + virtual const stream_executor::DeviceDescription& GetDeviceDescription() + const = 0; + virtual const DebugOptions& GetDebugOptions() const = 0; virtual Thunk::ThunkInfo GenerateThunkInfo() const = 0; + // Returns the slice for the result or an operand together with the shape + // backing it. This is the form consumed by most thunks, so prefer it over the + // plain allocation slice accessors below. + virtual absl::StatusOr GetResultShapedSlice( + const ShapeIndex& index) const = 0; + + virtual absl::StatusOr GetOperandShapedSlice( + int64_t operand_index, const ShapeIndex& index) const = 0; + virtual absl::StatusOr GetResultAllocationSlice( const ShapeIndex& index) const = 0; virtual absl::StatusOr GetOperandAllocationSlice( int64_t operand_index, const ShapeIndex& index) const = 0; - virtual absl::StatusOr GetFfiAttributes() const = 0; + // Builds the conventional kernel argument list for the custom call: all + // operands in operand order, then the array leaves of the result shape in + // shape-index order, then `unmanaged_arguments`. + // + // This is the same list that XLA's own kernel emitters build, with alignment, + // aliasing, slice deduplication and the `written` flags filled in. Handlers + // should use it rather than assembling `emitters::KernelArgument`s by hand, + // because those flags feed thunk scheduling and getting them wrong produces + // races that are hard to diagnose. + // + // Use `emitters::KernelArguments::PositionOfOperand` and `PositionOfResult` + // to find where a given buffer ended up; those positions are exactly the + // index space of `stream_executor::KernelArgsPackingSpec` relocations. + absl::StatusOr CreateKernelArguments() const { + return CreateKernelArguments({}); + } + virtual absl::StatusOr CreateKernelArguments( + absl::Span unmanaged_arguments) const = 0; + + // The custom call's backend config, decoded into typed FFI attributes. + // + // The returned object owns its storage; values that alias it + // (`absl::string_view`, `absl::Span`, nested `ffi::Dictionary`) are only + // valid while it is alive. + virtual absl::StatusOr GetFfiAttributes() const = 0; }; } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_registry_test.cc b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_registry_test.cc index 2dd76148acd1e3..b6b35f7fdf9e0e 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_registry_test.cc +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_registry_test.cc @@ -15,7 +15,6 @@ limitations under the License. #include "xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_registry.h" -#include #include #include @@ -26,12 +25,7 @@ limitations under the License. #include "xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_emitter_context.h" #include "xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_registration.h" #include "xla/backends/gpu/runtime/thunk.h" -#include "xla/backends/gpu/runtime/thunk_id.h" -#include "xla/ffi/attribute_map.h" #include "xla/hlo/ir/hlo_instructions.h" -#include "xla/service/buffer_assignment.h" -#include "xla/service/gpu_topology.h" -#include "xla/shape_util.h" namespace xla::gpu { namespace { @@ -82,43 +76,5 @@ TEST(NativeCustomCallHandlerRegistryTest, StatusIs(absl::StatusCode::kInvalidArgument)); } -class MockNativeCustomCallEmitterContext - : public NativeCustomCallEmitterContext { - public: - const GpuTopology& GetTargetTopology() const override { return *topology_; } - const DebugOptions& GetDebugOptions() const override { - return debug_options_; - } - Thunk::ThunkInfo GenerateThunkInfo() const override { - Thunk::ThunkInfo info; - info.thunk_id = ThunkId(42); - return info; - } - absl::StatusOr GetResultAllocationSlice( - const ShapeIndex& index) const override { - return slice_; - } - absl::StatusOr GetOperandAllocationSlice( - int64_t operand_index, const ShapeIndex& index) const override { - return slice_; - } - absl::StatusOr GetFfiAttributes() const override { - return xla::ffi::AttributesMap(); - } - - const GpuTopology* topology_ = nullptr; - DebugOptions debug_options_; - BufferAllocation::Slice slice_; -}; - -TEST(NativeCustomCallEmitterContextTest, VirtualDispatchWorks) { - MockNativeCustomCallEmitterContext mock_ctx; - const NativeCustomCallEmitterContext& ctx = mock_ctx; - - EXPECT_EQ(ctx.GenerateThunkInfo().thunk_id, ThunkId(42)); - EXPECT_OK(ctx.GetResultAllocationSlice({})); - EXPECT_OK(ctx.GetOperandAllocationSlice(0, {})); -} - } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/BUILD b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/BUILD index 48804f49136719..3ecff828da8c20 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/BUILD +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/BUILD @@ -58,14 +58,12 @@ cc_library( "//xla/backends/gpu/runtime:custom_kernel_thunk", "//xla/backends/gpu/runtime:thunk", "//xla/codegen/emitters:kernel_arguments", - "//xla/ffi:attribute_map", + "//xla/ffi:attributes", "//xla/hlo/ir:hlo", - "//xla/service:buffer_assignment", "//xla/stream_executor:kernel_args_packing_spec", "//xla/stream_executor:kernel_spec", "//xla/stream_executor:launch_dim", "//xla/stream_executor/cuda:cudart_kernel_registry", - "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", ], diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_handler.cc b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_handler.cc index 611f1b69fef3cb..b9a5db3b2767a0 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_handler.cc +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_handler.cc @@ -23,10 +23,7 @@ limitations under the License. #include #include -#include -#include -#include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "xla/backends/gpu/codegen/kernels/custom_kernel.h" @@ -36,9 +33,8 @@ limitations under the License. #include "xla/backends/gpu/runtime/custom_kernel_thunk.h" #include "xla/backends/gpu/runtime/thunk.h" #include "xla/codegen/emitters/kernel_arguments.h" -#include "xla/ffi/attribute_map.h" +#include "xla/ffi/attributes.h" #include "xla/hlo/ir/hlo_instructions.h" -#include "xla/service/buffer_assignment.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/cuda/cudart_kernel_registry.h" @@ -58,31 +54,23 @@ absl::StatusOr WriteValueThunkFoldedHandler( const Shape& out_shape = instr.shape(); int64_t num_elements = ShapeUtil::ElementsIn(out_shape); - ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice out_slice, - ctx.GetResultAllocationSlice(/*index=*/{})); - - ABSL_ASSIGN_OR_RETURN(xla::ffi::AttributesMap attrs, ctx.GetFfiAttributes()); - auto it = attrs.find("val"); - if (it == attrs.end()) { - return absl::InvalidArgumentError( - "Expected 'val' attribute in backend_config"); - } - if (!std::holds_alternative(it->second.AsVariant())) { - return absl::InvalidArgumentError("'val' attribute must be a scalar"); - } - const xla::ffi::Scalar& scalar = - std::get(it->second.AsVariant()); - if (!std::holds_alternative(scalar.AsVariant())) { - return absl::InvalidArgumentError("'val' attribute must be an int32_t"); - } - int32_t val = std::get(scalar.AsVariant()); + ABSL_ASSIGN_OR_RETURN(xla::ffi::Attributes attrs, ctx.GetFfiAttributes()); + ABSL_ASSIGN_OR_RETURN(int32_t val, attrs.Get("val")); + + // The conventional argument list for this custom call, which here is just the + // result buffer. Letting the context build it keeps alignment, aliasing and + // the `written` flags consistent with the rest of XLA's kernel emitters. + ABSL_ASSIGN_OR_RETURN(emitters::KernelArguments kernel_args, + ctx.CreateKernelArguments()); + ABSL_ASSIGN_OR_RETURN(int64_t result_position, + kernel_args.PositionOfResult(/*index=*/{})); ABSL_ASSIGN_OR_RETURN(stream_executor::KernelLoaderSpec kernel_spec, stream_executor::cuda::FindCudaRuntimeKernel( stream_executor::cuda::GetWriteValueFoldedKernel())); stream_executor::KernelArgsPackingSpec packing_spec; - packing_spec.AddAddressArgument(0); + packing_spec.AddAddressArgument(result_position); packing_spec.AddConstantArgument(val); kernel_spec.set_kernel_args_packing(std::move(packing_spec)); @@ -94,11 +82,6 @@ absl::StatusOr WriteValueThunkFoldedHandler( /*thread_dims=*/stream_executor::ThreadDim(), /*shared_memory_bytes=*/0); - emitters::KernelArgument out_arg(out_shape, out_slice); - out_arg.set_written(true); - emitters::KernelArguments kernel_args( - std::vector{out_arg}); - return ThunkSequence::Of( ctx.GenerateThunkInfo(), std::move(custom_kernel), kernel_args); } diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index b2737e63898297..2c5c6a1cbff212 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -558,6 +558,7 @@ cc_library( "//xla/codegen/xtile:block_level_parameters", "//xla/core/host_offloading:host_offloading_executable_proto_cc", "//xla/ffi:attribute_map", + "//xla/ffi:attributes", "//xla/hlo/ir:hlo", "//xla/hlo/ir:hlo_instruction_utils", "//xla/mlir/utils:error_util", diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc index b17bfa3db9b318..cab5cf94a3db7d 100644 --- a/third_party/xla/xla/service/gpu/thunk_emitter.cc +++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc @@ -116,6 +116,7 @@ limitations under the License. #include "xla/codegen/xtile/block_level_parameters.h" #include "xla/core/host_offloading/host_offloading_executable.pb.h" #include "xla/ffi/attribute_map.h" +#include "xla/ffi/attributes.h" #include "xla/future.h" #include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" @@ -1179,6 +1180,11 @@ class NativeCustomCallEmitterContextImpl return emitter_.ir_emitter_context_->gpu_topology(); } + const stream_executor::DeviceDescription& GetDeviceDescription() + const override { + return emitter_.ir_emitter_context_->gpu_device_info(); + } + const DebugOptions& GetDebugOptions() const override { return emitter_.ir_emitter_context_->debug_options(); } @@ -1195,11 +1201,36 @@ class NativeCustomCallEmitterContextImpl absl::StatusOr GetOperandAllocationSlice( int64_t operand_index, const ShapeIndex& index) const override { - TF_RET_CHECK(operand_index >= 0 && operand_index < instr_.operand_count()); - return emitter_.GetAllocationSlice(instr_.operand(operand_index), index); + ABSL_ASSIGN_OR_RETURN(const HloInstruction* operand, GetOperand(operand_index)); + return emitter_.GetAllocationSlice(operand, index); + } + + absl::StatusOr GetResultShapedSlice( + const ShapeIndex& index) const override { + return emitter_.GetShapedSliceForHlo(&instr_, index); + } + + absl::StatusOr GetOperandShapedSlice( + int64_t operand_index, const ShapeIndex& index) const override { + ABSL_ASSIGN_OR_RETURN(const HloInstruction* operand, GetOperand(operand_index)); + return emitter_.GetShapedSliceForHlo(operand, index); + } + + absl::StatusOr CreateKernelArguments( + absl::Span unmanaged_arguments) const override { + // Resolve slices through the emitter rather than through the buffer + // assignment, so that allocation overrides installed for this instruction + // are applied. + auto slice_provider = [this](const HloInstruction& instruction, + const ShapeIndex& index) { + return emitter_.GetAllocationSlice(&instruction, index); + }; + return emitters::KernelArguments::Create(slice_provider, + GetDefaultBufferAlignment(), + &instr_, unmanaged_arguments); } - absl::StatusOr GetFfiAttributes() const override { + absl::StatusOr GetFfiAttributes() const override { // Decode the opaque backend config into an FFI attributes map, mirroring // EmitGenericCustomCall. For FFI handlers the backend config must be a // string parsable into an MLIR dictionary attribute. @@ -1210,7 +1241,7 @@ class NativeCustomCallEmitterContextImpl ? backend_config->custom_call_backend_config().attributes() : instr_.raw_backend_config_string(); if (backend_config_str.empty()) { - return xla::ffi::AttributesMap(); + return xla::ffi::Attributes::Create(xla::ffi::AttributesMap()); } mlir::Attribute attr = mlir::parseAttribute( backend_config_str, emitter_.ir_emitter_context_->mlir_context()); @@ -1218,10 +1249,18 @@ class NativeCustomCallEmitterContextImpl TF_RET_CHECK(dict != nullptr) << "Unsupported backend config. Expected a string parsable into a " "dictionary attribute."; - return xla::ffi::BuildAttributesMap(dict); + ABSL_ASSIGN_OR_RETURN(xla::ffi::AttributesMap attributes, + xla::ffi::BuildAttributesMap(dict)); + return xla::ffi::Attributes::Create(std::move(attributes)); } private: + absl::StatusOr GetOperand( + int64_t operand_index) const { + TF_RET_CHECK(operand_index >= 0 && operand_index < instr_.operand_count()); + return instr_.operand(operand_index); + } + const ThunkEmitter& emitter_; const HloCustomCallInstruction& instr_; }; From 9edd58f895c58c3db484a9e39ab2e953baca4bca Mon Sep 17 00:00:00 2001 From: David Dunleavy Date: Mon, 21 Sep 2026 12:04:33 -0700 Subject: [PATCH 3/3] Tag `clang_tidy_diff` as `manual` This frequently shows up in x86 CPU logs on the critical path, but we don't actually test this on CI. PiperOrigin-RevId: 985426107 --- third_party/xla/build_tools/ci/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/third_party/xla/build_tools/ci/BUILD b/third_party/xla/build_tools/ci/BUILD index 3bec8df90239b5..3bd92b827859fb 100644 --- a/third_party/xla/build_tools/ci/BUILD +++ b/third_party/xla/build_tools/ci/BUILD @@ -53,6 +53,7 @@ pytype_strict_binary( ":clang_apply_replacements_bin", ]), main = "clang_tidy_diff.py", + tags = ["manual"], deps = [ ":clang_tidy_diff_lib", ],