diff --git a/third_party/xla/third_party/llvm/generated.patch b/third_party/xla/third_party/llvm/generated.patch index 02bd6121e33..46b25d1a689 100644 --- a/third_party/xla/third_party/llvm/generated.patch +++ b/third_party/xla/third_party/llvm/generated.patch @@ -1,28 +1,65 @@ Auto generated patch. Do not edit or delete it, even if empty. -diff -ruN --strip-trailing-cr a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp ---- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp -+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp -@@ -6275,11 +6275,6 @@ - if (VPI.getOpcode() != Instruction::Trunc) - continue; +diff -ruN --strip-trailing-cr a/llvm/test/Transforms/LoopVectorize/WebAssembly/CostModel/loop-02-bb1176.ll b/llvm/test/Transforms/LoopVectorize/WebAssembly/CostModel/loop-02-bb1176.ll +--- a/llvm/test/Transforms/LoopVectorize/WebAssembly/CostModel/loop-02-bb1176.ll ++++ b/llvm/test/Transforms/LoopVectorize/WebAssembly/CostModel/loop-02-bb1176.ll +@@ -1,4 +1,5 @@ + ; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py UTC_ARGS: --filter "LV: Found an estimated cost of [0-9]+(.[0-9]+)? for VF 1 For instruction:\s*%v0 = load" --filter "Cost of [0-9]+(.[0-9]+)? for VF [0-9]+: (INTERLEAVE-GROUP with factor [0-9]+,|EXPRESSION|REPLICATE|WIDEN)" --filter "^ ir<.* = load from index" --filter "LV: Selecting VF:" --filter "Estimated cost per lane" ++; REQUIRES: asserts + ; RUN: opt -mtriple=wasm32 -mattr=+simd128 -passes=loop-vectorize %s --debug-only=loop-vectorize --disable-output < %s 2>&1 | FileCheck %s + ; RUN: opt -mtriple=wasm32 -mattr=+simd128 -passes=loop-vectorize %s --vectorizer-maximize-bandwidth --debug-only=loop-vectorize --disable-output < %s 2>&1 | FileCheck %s --check-prefix=MAX -- // A truncate that is not widened is left to the scalarization decisions -- // made earlier. -- if (vputils::onlyFirstLaneUsed(&VPI)) -- continue; -- - VPValue *Op = VPI.getOperand(0); - auto *WideIV = getOptimizableIVOf(Op, PSE); - if (!WideIV) -@@ -6291,6 +6286,11 @@ - if (Op != WideIV) - continue; +diff -ruN --strip-trailing-cr a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel ++++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +@@ -770,6 +770,26 @@ + ) -+ // A truncate that is not widened is left to the scalarization decisions -+ // made earlier. -+ if (vputils::onlyFirstLaneUsed(&VPI)) -+ continue; + mlir_c_api_cc_library( ++ name = "CAPIOpenACC", ++ srcs = [ ++ "lib/CAPI/Dialect/OpenACC.cpp", ++ "lib/CAPI/Dialect/OpenACCPasses.cpp", ++ ], ++ hdrs = ["include/mlir-c/Dialect/OpenACC.h"], ++ capi_deps = [":CAPIIR"], ++ header_deps = [":OpenACCPassIncGen"], ++ includes = ["include"], ++ deps = [ ++ ":GPUDialect", ++ ":IR", ++ ":OpenACCAllExtensions", ++ ":OpenACCDialect", ++ ":OpenACCTransforms", ++ ":Pass", ++ ], ++) + - // Replacing a free truncate would add an induction update instruction to - // each iteration of the loop. The canonical induction is exempt, as it - // needs an update instruction regardless. ++mlir_c_api_cc_library( + name = "CAPIShape", + srcs = ["lib/CAPI/Dialect/Shape.cpp"], + hdrs = ["include/mlir-c/Dialect/Shape.h"], +@@ -10922,10 +10942,20 @@ + + gentbl_cc_library( + name = "OpenACCPassIncGen", +- tbl_outs = {"include/mlir/Dialect/OpenACC/Transforms/Passes.h.inc": [ +- "-gen-pass-decls", +- "-name=OpenACC", +- ]}, ++ tbl_outs = { ++ "include/mlir/Dialect/OpenACC/Transforms/Passes.h.inc": [ ++ "-gen-pass-decls", ++ "-name=OpenACC", ++ ], ++ "include/mlir/Dialect/OpenACC/Transforms/Passes.capi.h.inc": [ ++ "-gen-pass-capi-header", ++ "--prefix=OpenACC", ++ ], ++ "include/mlir/Dialect/OpenACC/Transforms/Passes.capi.cpp.inc": [ ++ "-gen-pass-capi-impl", ++ "--prefix=OpenACC", ++ ], ++ }, + tblgen = ":mlir-tblgen", + td_file = "include/mlir/Dialect/OpenACC/Transforms/Passes.td", + deps = [":PassBaseTdFiles"], diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl index 8ea094a511c..88c8006ce0c 100644 --- a/third_party/xla/third_party/llvm/workspace.bzl +++ b/third_party/xla/third_party/llvm/workspace.bzl @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive") def repo(name): """Imports LLVM.""" - LLVM_COMMIT = "634f85fc756b6a6a93800831dcb65f3bf6f93086" - LLVM_SHA256 = "2ccb02ee9260da144a7f0c4d9ded3af95a39a9a0ddcabac1bb4d210378f32f80" + LLVM_COMMIT = "d5efcd0ec41d0f425076e3dd738a68d333ad610f" + LLVM_SHA256 = "388a8ffba67c0e9af2d966dd136138ddc222cdbecc237923c18a42fab646b00f" tf_http_archive( name = name, diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch index 5c81a5ba779..e69de29bb2d 100644 --- a/third_party/xla/third_party/shardy/temporary.patch +++ b/third_party/xla/third_party/shardy/temporary.patch @@ -1,3303 +0,0 @@ -diff --git a/shardy/dialect/sdy/ir/attrs.td b/shardy/dialect/sdy/ir/attrs.td -index 36738232..d9327cba 100644 ---- a/shardy/dialect/sdy/ir/attrs.td -+++ b/shardy/dialect/sdy/ir/attrs.td -@@ -26,15 +26,6 @@ include "shardy/dialect/sdy/ir/enums.td" - // NOTE: we use `` in assemblyFormat to avoid whitespaces between literals and - // parameters. - --// A list of axes that a ManualComputationOp is manual on. ManualComputationOp --// doesn't use an array of AxisRefAttr since all axes must be full/can't be --// split. --def Sdy_ManualAxes : ArrayOfAttr { -- let summary = "A list of axes that a ManualComputationOp is manual on"; -- let assemblyFormat = "`{` (`}`) : ($value^ `` `}`)?"; --} -- - // LINT.IfChange - def Sdy_MeshAxis : AttrDef { - let mnemonic = "mesh_axis"; -@@ -1234,6 +1225,43 @@ def Sdy_AxisRefList : ArrayOfAttr { -+ let summary = "A list of axes that a ManualComputationOp is manual on"; -+ let assemblyFormat = "`{` (`}`) : ($value^ `` `}`)?"; -+ let extraClassDeclaration = [{ -+ using iterator = const AxisRefAttr*; -+ iterator begin() const { return getValue().begin(); } -+ iterator end() const { return getValue().end(); } -+ size_t size() const { return getValue().size(); } -+ bool empty() const { return getValue().empty(); } -+ const AxisRefAttr& operator[](size_t index) const { -+ return getValue()[index]; -+ } -+ operator ArrayRef() const { return getValue(); } -+ -+ template >>> -+ static ManualAxesAttr get(MLIRContext* context, RangeT&& axes) { -+ return getFromStringAttrs(context, axes); -+ } -+ -+ template >> -+ static ManualAxesAttr get(MLIRContext* context, -+ std::initializer_list axes) { -+ return getFromStringAttrs(context, axes); -+ } -+ -+ static ManualAxesAttr getFromStringAttrs(MLIRContext* context, -+ ArrayRef axes); -+ -+ SmallVector getStringValue() const; -+ }]; -+} -+ - def Sdy_EmptyAxisList: AttrConstraint< - CPred<"mlir::isa($_self) && " - "mlir::cast($_self).empty()">, -diff --git a/shardy/dialect/sdy/ir/bytecode.td b/shardy/dialect/sdy/ir/bytecode.td -index fb4564ae..e44e3c7a 100644 ---- a/shardy/dialect/sdy/ir/bytecode.td -+++ b/shardy/dialect/sdy/ir/bytecode.td -@@ -42,12 +42,20 @@ def OptionalVarInt : - // ----------------------------------------------------------------------------- - - //===----------------------------------------------------------------------===// --// ManualAxesAttr - assigned index 0 -+// ManualAxesAttr V1 - assigned index 0 - //===----------------------------------------------------------------------===// --def Sdy_ManualAxesAttrBc : DialectAttribute<(attr -- Array:$value -+def Sdy_StringAttrArrayBc : -+ WithParser<"succeeded($_reader.readAttributes($_var))", -+ WithBuilder<"$_args", -+ WithPrinter<"$_writer.writeAttributes($_getter)", -+ WithType<"SmallVector">>>>; -+ -+def Sdy_ManualAxesAttrV1Bc : DialectAttribute<(attr -+ Sdy_StringAttrArrayBc:$string_value - )> { - let cType = "ManualAxesAttr"; -+ let cBuilder = "$_resultType::get(context, $_args)"; -+ let printerPredicate = "llvm::none_of($_val.getValue(), [](AxisRefAttr axis) { return static_cast(axis.getSubAxisInfo()); })"; - } - - //===----------------------------------------------------------------------===// -@@ -247,6 +255,16 @@ def Sdy_AllToAllParamListAttrBc: DialectAttribute<(attr - let cType = "AllToAllParamListAttr"; - } - -+//===----------------------------------------------------------------------===// -+// ManualAxesAttr V2 - assigned index 18 -+//===----------------------------------------------------------------------===// -+def Sdy_ManualAxesAttrV2Bc : DialectAttribute<(attr -+ Array:$value -+)> { -+ let cType = "ManualAxesAttr"; -+ let printerPredicate = "llvm::any_of($_val.getValue(), [](AxisRefAttr axis) { return static_cast(axis.getSubAxisInfo()); })"; -+} -+ - // This enum contains marker codes used to indicate which attribute is - // currently being decoded, and how it should be decoded. The order of these - // codes should generally be unchanged, as any changes will inevitably break -@@ -255,7 +273,7 @@ def SdyDialectTypes : DialectAttributes<"sdy"> { - // Take care when chaging this list. - // If needed to remove use ReservedOrDead - let elems = [ -- /*index 0*/ Sdy_ManualAxesAttrBc, -+ /*index 0*/ Sdy_ManualAxesAttrV1Bc, - /*index 1*/ Sdy_MeshAxisAttrBc, - /*index 2*/ Sdy_MeshAttrBc, - /*index 3*/ Sdy_SubAxisInfoAttrBc, -@@ -272,7 +290,8 @@ def SdyDialectTypes : DialectAttributes<"sdy"> { - /*index 14*/ Sdy_AllToAllParamListAttrBc, - /*index 15*/ Sdy_TensorShardingAttrV2Bc, - /*index 16*/ Sdy_TensorShardingPerValueAttrV2Bc, -- /*index 17*/ Sdy_TensorShardingAttrV3Bc -+ /*index 17*/ Sdy_TensorShardingAttrV3Bc, -+ /*index 18*/ Sdy_ManualAxesAttrV2Bc - ]; - } - #endif // SDY_BYTECODE -diff --git a/shardy/dialect/sdy/ir/canonicalization.cc b/shardy/dialect/sdy/ir/canonicalization.cc -index 78fee0b5..67ed717e 100644 ---- a/shardy/dialect/sdy/ir/canonicalization.cc -+++ b/shardy/dialect/sdy/ir/canonicalization.cc -@@ -124,8 +124,8 @@ class RedundantManualComputationPattern - MeshAttr mesh = - getCommonMesh(inShardings, outShardings, manualComputationOp); - SDY_CHECK(mesh) << "expected inputs and outputs to have a common mesh"; -- for (StringAttr manualAxis : manualComputationOp.getManualAxes()) { -- manualAxesProduct *= mesh.getAxisSize(manualAxis); -+ for (AxisRefAttr manualAxis : manualComputationOp.getManualAxes()) { -+ manualAxesProduct *= manualAxis.getSize(mesh); - } - } - -diff --git a/shardy/dialect/sdy/ir/dialect.cc b/shardy/dialect/sdy/ir/dialect.cc -index 46991917..1f617b29 100644 ---- a/shardy/dialect/sdy/ir/dialect.cc -+++ b/shardy/dialect/sdy/ir/dialect.cc -@@ -559,9 +559,12 @@ std::optional AxisRefAttr::getSuffixWithoutOverlap( - if (!overlaps(other)) { - return *this; - } -+ if (!other.getSubAxisInfo() || (!getSubAxisInfo() && !mesh)) { -+ return std::nullopt; -+ } - - int64_t thisNextPreSize = getNextPreSizeOrFullSize(mesh); -- int64_t otherNextPreSize = other.getNextPreSizeOrFullSize(mesh); -+ int64_t otherNextPreSize = other.getSubAxisInfo().getNextPreSize(); - if (thisNextPreSize <= otherNextPreSize) { - return std::nullopt; - } -@@ -1211,6 +1214,29 @@ AllToAllParamListAttr AllToAllParamListAttr::combineAndSort( - return AllToAllParamListAttr::get(getContext(), combinedParams); - } - -+//===----------------------------------------------------------------------===// -+// ManualAxesAttr -+//===----------------------------------------------------------------------===// -+ -+ManualAxesAttr ManualAxesAttr::getFromStringAttrs(MLIRContext* context, -+ ArrayRef axes) { -+ SmallVector axisRefs; -+ axisRefs.reserve(axes.size()); -+ for (StringAttr axis : axes) { -+ axisRefs.push_back(AxisRefAttr::get(context, axis)); -+ } -+ return get(context, axisRefs); -+} -+ -+SmallVector ManualAxesAttr::getStringValue() const { -+ SmallVector stringAxes; -+ stringAxes.reserve(size()); -+ for (AxisRefAttr axis : getValue()) { -+ stringAxes.push_back(StringAttr::get(getContext(), axis.getName())); -+ } -+ return stringAxes; -+} -+ - //===----------------------------------------------------------------------===// - // ManualComputationOp - //===----------------------------------------------------------------------===// -@@ -1229,21 +1255,39 @@ namespace { - // `newSharding`. - TensorShardingAttr addFreeAxesToManualComputationSharding( - TensorShardingAttr outerManualSharding, TensorShardingAttr newSharding, -- ArrayRef manualAxes) { -+ ArrayRef manualAxes, Operation* op) { - // Remove all existing free axes first before adding possibly extra ones in - // `newSharding`. - TensorShardingAttr returnedSharding = - eraseFreeAxes(outerManualSharding, manualAxes); - -+ MeshAttr mesh; - SmallVector resultDimShardings = llvm::map_to_vector( - llvm::zip_equal(returnedSharding.getDimShardings(), - newSharding.getDimShardings()), -- [](auto pair) { -+ [&](auto pair) { - auto [returnedDimSharding, newDimSharding] = pair; -+ SmallVector combinedAxes( -+ returnedDimSharding.getAxes().begin(), -+ returnedDimSharding.getAxes().end()); -+ if (!combinedAxes.empty() && !newDimSharding.emptyAxes() && -+ combinedAxes.back().canMerge(newDimSharding.getAxes().front())) { -+ if (!mesh) { -+ mesh = outerManualSharding.getMesh(op); -+ } -+ if (mesh) { -+ for (AxisRefAttr axis : newDimSharding.getAxes()) { -+ addAxisOrMerge(combinedAxes, axis, mesh); -+ } -+ return DimensionShardingAttr::get( -+ returnedDimSharding.getContext(), combinedAxes, -+ newDimSharding.getIsClosed(), newDimSharding.getPriority()); -+ } -+ } -+ combinedAxes.append(newDimSharding.getAxes().begin(), -+ newDimSharding.getAxes().end()); - return DimensionShardingAttr::get( -- returnedDimSharding.getContext(), -- llvm::to_vector(llvm::concat( -- returnedDimSharding.getAxes(), newDimSharding.getAxes())), -+ returnedDimSharding.getContext(), combinedAxes, - newDimSharding.getIsClosed(), newDimSharding.getPriority()); - }); - return TensorShardingAttr::get(newSharding.getContext(), -@@ -1280,21 +1324,25 @@ SmallVector - ManualComputationOp::getBlockArgumentEdgeOwnerShardings() { - SmallVector shardings; - shardings.reserve(getInShardings().size()); -+ MeshAttr mesh; - for (TensorShardingAttr sharding : getInShardings().getShardings()) { -- shardings.push_back(eraseManualAxes(sharding, getManualAxes())); -+ if (!mesh) { -+ mesh = sharding.getMesh(*this); -+ } -+ shardings.push_back(eraseManualAxes(sharding, getManualAxes(), mesh)); - } - return shardings; - } - - void ManualComputationOp::setBlockArgumentEdgeOwnerShardings( - ArrayRef shardings) { -- ArrayRef manualAxes = getManualAxes(); -+ ArrayRef manualAxes = getManualAxes(); - setInShardings(llvm::map_to_vector( - llvm::zip_equal(getInShardings().getShardings(), shardings), - [&](auto&& pair) { - auto&& [outerSharding, newSharding] = pair; - return addFreeAxesToManualComputationSharding( -- outerSharding, newSharding, manualAxes); -+ outerSharding, newSharding, manualAxes, *this); - })); - } - -@@ -1329,17 +1377,20 @@ TensorShardingAttr ManualComputationOp::transformTargetSharding( - case DataFlowShardingTransformType::kBeforeEdgePropagation: { - if (auto blockArg = dyn_cast(target)) { - return addFreeAxesToManualComputationSharding( -- getInSharding(blockArg.getArgNumber()), sharding, getManualAxes()); -+ getInSharding(blockArg.getArgNumber()), sharding, getManualAxes(), -+ *this); - } -- return eraseManualAxes(sharding, getManualAxes()); -+ return eraseManualAxes(sharding, getManualAxes(), -+ sharding.getMesh(*this)); - } - case DataFlowShardingTransformType::kAfterEdgePropagation: { - if (isa(target)) { -- return eraseManualAxes(sharding, getManualAxes()); -+ return eraseManualAxes(sharding, getManualAxes(), -+ sharding.getMesh(*this)); - } - return addFreeAxesToManualComputationSharding( - getOutSharding(cast(target).getResultNumber()), sharding, -- getManualAxes()); -+ getManualAxes(), *this); - } - } - llvm_unreachable("received an unexpected target type."); -diff --git a/shardy/dialect/sdy/ir/dialect.h b/shardy/dialect/sdy/ir/dialect.h -index 1c5fa090..54be287d 100644 ---- a/shardy/dialect/sdy/ir/dialect.h -+++ b/shardy/dialect/sdy/ir/dialect.h -@@ -21,10 +21,12 @@ limitations under the License. - #include - #include - #include -+#include - #include - #include - #include - #include -+#include - - #include "llvm/ADT/StringExtras.h" - #include "llvm/ADT/StringRef.h" -diff --git a/shardy/dialect/sdy/ir/ops.td b/shardy/dialect/sdy/ir/ops.td -index 71e4df5e..0d2498d0 100644 ---- a/shardy/dialect/sdy/ir/ops.td -+++ b/shardy/dialect/sdy/ir/ops.td -@@ -207,7 +207,7 @@ def Sdy_ManualComputationOp : Sdy_Op<"manual_computation", - OpBuilder<(ins "TypeRange":$results, "ValueRange":$tensors, - "ArrayRef":$in_shardings, - "ArrayRef":$out_shardings, -- "ArrayRef":$manual_axes), [{ -+ CArg<"ArrayRef", "{}">:$manual_axes), [{ - MLIRContext* ctx = $_builder.getContext(); - return build($_builder, $_state, results, tensors, - TensorShardingPerValueAttr::get(ctx, in_shardings), -@@ -217,6 +217,71 @@ def Sdy_ManualComputationOp : Sdy_Op<"manual_computation", - ]; - - let extraClassDeclaration = [{ -+ template >>> -+ static void build(OpBuilder& odsBuilder, OperationState& odsState, -+ TypeRange results, ValueRange tensors, -+ TensorShardingPerValueAttr in_shardings, -+ TensorShardingPerValueAttr out_shardings, -+ RangeT&& manual_axes) { -+ build(odsBuilder, odsState, results, tensors, in_shardings, out_shardings, -+ ManualAxesAttr::get(odsBuilder.getContext(), -+ ArrayRef(manual_axes))); -+ } -+ -+ template >>> -+ static ManualComputationOp create(OpBuilder& builder, Location location, -+ TypeRange results, ValueRange tensors, -+ TensorShardingPerValueAttr in_shardings, -+ TensorShardingPerValueAttr out_shardings, -+ RangeT&& manual_axes) { -+ return create(builder, location, results, tensors, in_shardings, -+ out_shardings, -+ ManualAxesAttr::get(builder.getContext(), -+ ArrayRef(manual_axes))); -+ } -+ -+ template >>> -+ static void build(OpBuilder& odsBuilder, OperationState& odsState, -+ TypeRange results, ValueRange tensors, -+ ArrayRef in_shardings, -+ ArrayRef out_shardings, -+ RangeT&& manual_axes) { -+ MLIRContext* ctx = odsBuilder.getContext(); -+ build(odsBuilder, odsState, results, tensors, -+ TensorShardingPerValueAttr::get(ctx, in_shardings), -+ TensorShardingPerValueAttr::get(ctx, out_shardings), -+ ManualAxesAttr::get(ctx, ArrayRef(manual_axes))); -+ } -+ -+ template >>> -+ static ManualComputationOp create( -+ OpBuilder& builder, Location location, TypeRange results, -+ ValueRange tensors, ArrayRef in_shardings, -+ ArrayRef out_shardings, RangeT&& manual_axes) { -+ MLIRContext* ctx = builder.getContext(); -+ return create( -+ builder, location, results, tensors, -+ TensorShardingPerValueAttr::get(ctx, in_shardings), -+ TensorShardingPerValueAttr::get(ctx, out_shardings), -+ ManualAxesAttr::get(ctx, ArrayRef(manual_axes))); -+ } -+ -+ template >>> -+ void setManualAxes(RangeT&& manual_axes) { -+ setManualAxesAttr( -+ ManualAxesAttr::get(getContext(), ArrayRef(manual_axes))); -+ } -+ - TensorShardingAttr getInSharding(int64_t operandIndex) { - return getInShardings().getSharding(operandIndex); - } -diff --git a/shardy/dialect/sdy/ir/test/manual_computation_parse_print.mlir b/shardy/dialect/sdy/ir/test/manual_computation_parse_print.mlir -index f0ee677d..a047c97d 100644 ---- a/shardy/dialect/sdy/ir/test/manual_computation_parse_print.mlir -+++ b/shardy/dialect/sdy/ir/test/manual_computation_parse_print.mlir -@@ -244,3 +244,35 @@ func.func @token_with_single_device_sharding( - } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) - return %0#0, %0#1 : !stablehlo.token, tensor<2xi64> - } -+ -+// CHECK-LABEL: func @manual_computation_sub_axis_manual_axes -+func.func @manual_computation_sub_axis_manual_axes(%arg0: tensor<16x32xf32>) -> tensor<16x32xf32> { -+ // CHECK{LITERAL}: sdy.manual_computation(%arg0) in_shardings=[<@meshB, [{"a"}, {}]>] out_shardings=[<@meshB, [{"a":(1)2}, {"a":(2)2}]>] manual_axes={"a":(1)2} (%arg1: tensor<8x32xf32>) { -+ // CHECK-NEXT: %[[BODY_RET:.*]] = stablehlo.add -+ // CHECK-NEXT: sdy.return %[[BODY_RET]] -+ // CHECK-NEXT: } : (tensor<16x32xf32>) -> tensor<16x32xf32> -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@meshB, [{"a"}, {}]>] out_shardings=[<@meshB, [{"a":(1)2}, {"a":(2)2}]>] manual_axes={"a":(1)2} (%arg1: tensor<8x32xf32>) { -+ %1 = stablehlo.add %arg1, %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@meshB, [{"a":(2)2}, {}]>]>} : tensor<8x32xf32> -+ sdy.return %1 : tensor<8x32xf32> -+ } : (tensor<16x32xf32>) -> tensor<16x32xf32> -+ func.return %0: tensor<16x32xf32> -+} -+ -+// CHECK-LABEL: func @manual_computation_nested_disjoint_sub_axes -+func.func @manual_computation_nested_disjoint_sub_axes(%arg0: tensor<16x32xf32>) -> tensor<16x32xf32> { -+ // CHECK{LITERAL}: sdy.manual_computation(%arg0) in_shardings=[<@meshB, [{"a"}, {}]>] out_shardings=[<@meshB, [{"a"}, {}]>] manual_axes={"a":(1)2} (%arg1: tensor<8x32xf32>) { -+ // CHECK-NEXT: %[[INNER:.*]] = sdy.manual_computation(%arg1) -+ // CHECK-SAME{LITERAL}: in_shardings=[<@meshB, [{"a":(2)2}, {}]>] out_shardings=[<@meshB, [{"a":(2)2}, {}]>] manual_axes={"a":(2)2} (%arg2: tensor<4x32xf32>) { -+ // CHECK-NEXT: %[[NESTED_RET:.*]] = stablehlo.add -+ // CHECK-NEXT: sdy.return %[[NESTED_RET]] -+ // CHECK-NEXT: } : (tensor<8x32xf32>) -> tensor<8x32xf32> -+ // CHECK-NEXT: sdy.return %[[INNER]] -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@meshB, [{"a"}, {}]>] out_shardings=[<@meshB, [{"a"}, {}]>] manual_axes={"a":(1)2} (%arg1: tensor<8x32xf32>) { -+ %1 = sdy.manual_computation(%arg1) in_shardings=[<@meshB, [{"a":(2)2}, {}]>] out_shardings=[<@meshB, [{"a":(2)2}, {}]>] manual_axes={"a":(2)2} (%arg2: tensor<4x32xf32>) { -+ %2 = stablehlo.add %arg2, %arg2 : tensor<4x32xf32> -+ sdy.return %2 : tensor<4x32xf32> -+ } : (tensor<8x32xf32>) -> tensor<8x32xf32> -+ sdy.return %1 : tensor<8x32xf32> -+ } : (tensor<16x32xf32>) -> tensor<16x32xf32> -+ func.return %0: tensor<16x32xf32> -+} -diff --git a/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir b/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir -index c2c24800..9c354c72 100644 ---- a/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir -+++ b/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir -@@ -319,3 +319,81 @@ func.func @token_region_result(%arg0: tensor<16x32xf32>, %arg1: !stablehlo.token - } : (tensor<16x32xf32>, !stablehlo.token) -> tensor<16x32xf32> - func.return %0: tensor<16x32xf32> - } -+ -+// ----- -+ -+sdy.mesh @mesh = <["a"=4]> -+ -+func.func @invalid_manual_sub_axis(%arg0: tensor<16x32xf32>) -> tensor<16x32xf32> { -+ // expected-error @+1 {{sub-axis next pre-size 8 doesn't divide the size of the full axis 4: "a":(2)4}} -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a":(2)4} (%arg1: tensor<4x32xf32>) { -+ sdy.return %arg1 : tensor<4x32xf32> -+ } : (tensor<16x32xf32>) -> tensor<16x32xf32> -+ func.return %0: tensor<16x32xf32> -+} -+ -+// ----- -+ -+sdy.mesh @mesh = <["a"=12]> -+ -+func.func @non_coexisting_sub_axis_with_manual_sub_axis(%arg0: tensor<24x32xf32>) -> tensor<24x32xf32> { -+ // expected-error @+1 {{op operand sharding at index 0 has axis "a":(1)6 that cannot coexist with manual axis "a":(1)4}} -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a":(1)6}, {}]>] out_shardings=[<@mesh, [{"a":(1)4}, {}]>] manual_axes={"a":(1)4} (%arg1: tensor<6x32xf32>) { -+ sdy.return %arg1 : tensor<6x32xf32> -+ } : (tensor<24x32xf32>) -> tensor<24x32xf32> -+ func.return %0: tensor<24x32xf32> -+} -+ -+// ----- -+ -+sdy.mesh @mesh = <["a"=8]> -+ -+func.func @partial_non_prefix_overlap_with_manual_sub_axis(%arg0: tensor<16x32xf32>) -> tensor<16x32xf32> { -+ // expected-error @+1 {{op operand sharding at index 0 must have all manual axes come before free axes in its dimension sharding at index 0. Saw manual axis "a":(2)2 after free axis "a":(1)2}} -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a":(1)4}, {}]>] out_shardings=[<@mesh, [{"a":(2)4}, {}]>] manual_axes={"a":(2)4} (%arg1: tensor<4x32xf32>) { -+ sdy.return %arg1 : tensor<4x32xf32> -+ } : (tensor<16x32xf32>) -> tensor<16x32xf32> -+ func.return %0: tensor<16x32xf32> -+} -+ -+// ----- -+ -+sdy.mesh @mesh = <["a"=4]> -+ -+func.func @manual_computation_nested_overlapping_sub_axis(%arg0: tensor<16x32xf32>) -> tensor<8x32xf32> { -+ // expected-note @+1 {{parent bounding this axis as manual}} -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a":(1)2}, {}]>] out_shardings=[<@mesh, [{}, {}], replicated={"a":(1)2}>] manual_axes={"a":(1)2} (%arg1: tensor<8x32xf32>) { -+ // expected-error @+1 {{op operates on axis "a" which is already bound by a parent sdy.manual_computation op}} -+ %1 = sdy.manual_computation(%arg1) in_shardings=[<@mesh, [{}, {}]>] out_shardings=[<@mesh, [{}, {}]>] manual_axes={"a"} (%arg2: tensor<8x32xf32>) { -+ sdy.return %arg2 : tensor<8x32xf32> -+ } : (tensor<8x32xf32>) -> tensor<8x32xf32> -+ sdy.return %1 : tensor<8x32xf32> -+ } : (tensor<16x32xf32>) -> tensor<8x32xf32> -+ func.return %0: tensor<8x32xf32> -+} -+ -+// ----- -+ -+sdy.mesh @mesh = <["a"=4]> -+ -+func.func @manual_computation_body_sharding_overlaps_manual_sub_axis(%arg0: tensor<16x32xf32>) -> tensor<16x32xf32> { -+ // expected-note @+1 {{parent bounding this axis as manual}} -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a":(1)2}, {}]>] out_shardings=[<@mesh, [{"a":(1)2}, {}]>] manual_axes={"a":(1)2} (%arg1: tensor<8x32xf32>) { -+ // expected-error @+1 {{'stablehlo.add' op result - operates on axis "a":(1)2 which is already bound by a parent sdy.manual_computation op}} -+ %1 = stablehlo.add %arg1, %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a":(1)2}, {}]>]>} : tensor<8x32xf32> -+ sdy.return %1 : tensor<8x32xf32> -+ } : (tensor<16x32xf32>) -> tensor<16x32xf32> -+ func.return %0: tensor<16x32xf32> -+} -+ -+// ----- -+ -+sdy.mesh @mesh = <["a"=4]> -+ -+func.func @free_sub_axis_before_manual_sub_axis_in_merged_axis(%arg0: tensor<16x32xf32>) -> tensor<16x32xf32> { -+ // expected-error @+1 {{op operand sharding at index 0 must have all manual axes come before free axes in its dimension sharding at index 0. Saw manual axis "a":(2)2 after free axis "a":(1)2}} -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a":(2)2}, {}]>] manual_axes={"a":(2)2} (%arg1: tensor<8x32xf32>) { -+ sdy.return %arg1 : tensor<8x32xf32> -+ } : (tensor<16x32xf32>) -> tensor<16x32xf32> -+ func.return %0: tensor<16x32xf32> -+} -diff --git a/shardy/dialect/sdy/ir/utils.cc b/shardy/dialect/sdy/ir/utils.cc -index 9b4baef5..027835a6 100644 ---- a/shardy/dialect/sdy/ir/utils.cc -+++ b/shardy/dialect/sdy/ir/utils.cc -@@ -643,8 +643,8 @@ using ManualComputationShardingEraserFn = std::function manualAxes, -- ManualComputationShardingEraserFn shardingEraser) { -+ TensorShardingAttr outerManualSharding, ArrayRef manualAxes, -+ MeshAttr mesh, ManualComputationShardingEraserFn shardingEraser) { - SmallVector newDimShardings; - newDimShardings.reserve(outerManualSharding.getRank()); - for (DimensionShardingAttr dimSharding : -@@ -654,18 +654,10 @@ TensorShardingAttr eraseAxesFromManualComputationSharding( - dimSharding, - getFirstFreeAxisIter(dimAxes, manualAxes) - dimAxes.begin())); - } -- // Grab any replicated/unreduced axes that are not manual axes. Can't use -- // `partition_point` as there is no defined order for replicated/unreduced -- // axes. -- auto isFreeAxis = [&](AxisRefAttr axis) { -- return !llvm::is_contained(manualAxes, axis.getName()); -- }; -- SmallVector newReplicatedAxes; -- llvm::copy_if(outerManualSharding.getReplicatedAxes(), -- std::back_inserter(newReplicatedAxes), isFreeAxis); -- SmallVector newUnreducedAxes; -- llvm::copy_if(outerManualSharding.getUnreducedAxes(), -- std::back_inserter(newUnreducedAxes), isFreeAxis); -+ SmallVector newReplicatedAxes = -+ getAxisSetDiff(outerManualSharding.getReplicatedAxes(), manualAxes, mesh); -+ SmallVector newUnreducedAxes = -+ getAxisSetDiff(outerManualSharding.getUnreducedAxes(), manualAxes, mesh); - return TensorShardingAttr::get( - outerManualSharding.getContext(), outerManualSharding.getMeshOrRef(), - newDimShardings, newReplicatedAxes, newUnreducedAxes, -@@ -676,26 +668,74 @@ TensorShardingAttr eraseAxesFromManualComputationSharding( - } // namespace - - TensorShardingAttr eraseManualAxes(TensorShardingAttr outerManualSharding, -- ArrayRef manualAxes) { -+ ArrayRef manualAxes, -+ MeshAttr mesh) { - if (manualAxes.empty()) { - return outerManualSharding; - } -+ if (!mesh) { -+ mesh = dyn_cast(outerManualSharding.getMeshOrRef()); -+ } - return eraseAxesFromManualComputationSharding( -- outerManualSharding, manualAxes, -- std::mem_fn(&DimensionShardingAttr::dropFrontShardingAxes)); -+ outerManualSharding, manualAxes, mesh, -+ [&](DimensionShardingAttr dimSharding, int64_t firstFreeAxisIndex) { -+ ArrayRef dimAxes = dimSharding.getAxes(); -+ if (firstFreeAxisIndex < static_cast(dimAxes.size())) { -+ AxisRefAttr firstFreeAxis = dimAxes[firstFreeAxisIndex]; -+ for (AxisRefAttr manualAxis : manualAxes) { -+ if (std::optional overlap = -+ firstFreeAxis.getOverlap(manualAxis); -+ overlap && overlap->strictPrefixOf(firstFreeAxis)) { -+ if (std::optional suffix = -+ firstFreeAxis.getSuffixWithoutOverlap(*overlap, mesh)) { -+ SmallVector newAxes; -+ newAxes.reserve(dimAxes.size() - firstFreeAxisIndex); -+ newAxes.push_back(*suffix); -+ newAxes.append(dimAxes.begin() + firstFreeAxisIndex + 1, -+ dimAxes.end()); -+ return DimensionShardingAttr::get( -+ dimSharding.getContext(), newAxes, -+ dimSharding.getIsClosed(), dimSharding.getPriority()); -+ } -+ } -+ } -+ } -+ return dimSharding.dropFrontShardingAxes(firstFreeAxisIndex); -+ }); - } - - TensorShardingAttr eraseFreeAxes(TensorShardingAttr outerManualSharding, -- ArrayRef manualAxes) { -+ ArrayRef manualAxes) { -+ MeshAttr mesh = dyn_cast(outerManualSharding.getMeshOrRef()); - return eraseAxesFromManualComputationSharding( -- outerManualSharding, manualAxes, -- std::mem_fn(&DimensionShardingAttr::takeFrontShardingAxes)); -+ outerManualSharding, manualAxes, mesh, -+ [&](DimensionShardingAttr dimSharding, int64_t firstFreeAxisIndex) { -+ ArrayRef dimAxes = dimSharding.getAxes(); -+ if (firstFreeAxisIndex < static_cast(dimAxes.size())) { -+ AxisRefAttr firstFreeAxis = dimAxes[firstFreeAxisIndex]; -+ for (AxisRefAttr manualAxis : manualAxes) { -+ if (std::optional overlap = -+ firstFreeAxis.getOverlap(manualAxis); -+ overlap && overlap->strictPrefixOf(firstFreeAxis)) { -+ SmallVector newAxes( -+ dimAxes.begin(), dimAxes.begin() + firstFreeAxisIndex); -+ newAxes.push_back(*overlap); -+ return DimensionShardingAttr::get( -+ dimSharding.getContext(), newAxes, dimSharding.getIsClosed(), -+ dimSharding.getPriority()); -+ } -+ } -+ } -+ return dimSharding.takeFrontShardingAxes(firstFreeAxisIndex); -+ }); - } - - ArrayRef::const_iterator getFirstFreeAxisIter( -- ArrayRef dimAxes, ArrayRef manualAxes) { -+ ArrayRef dimAxes, ArrayRef manualAxes) { - return llvm::partition_point(dimAxes, [&manualAxes](AxisRefAttr axis) { -- return llvm::is_contained(manualAxes, axis.getName()); -+ return llvm::any_of(manualAxes, [axis](AxisRefAttr manualAxis) { -+ return manualAxis.contains(axis); -+ }); - }); - } - -diff --git a/shardy/dialect/sdy/ir/utils.h b/shardy/dialect/sdy/ir/utils.h -index b4a43b4f..9efa9a9b 100644 ---- a/shardy/dialect/sdy/ir/utils.h -+++ b/shardy/dialect/sdy/ir/utils.h -@@ -561,7 +561,8 @@ SmallVector getOpenShardingsWithShardingAtIndex( - // `ManualComputationOp` are prefixed with the manual axes. So this removes the - // prefix of manual axes (if any exist) from each dim sharding. - TensorShardingAttr eraseManualAxes(TensorShardingAttr outerManualSharding, -- ArrayRef manualAxes); -+ ArrayRef manualAxes, -+ MeshAttr mesh = nullptr); - - // Removes free axes from the sharding. - // -@@ -569,7 +570,7 @@ TensorShardingAttr eraseManualAxes(TensorShardingAttr outerManualSharding, - // `ManualComputationOp` are prefixed with the manual axes. So this removes the - // suffix of free axes (if any exist) from each dim sharding. - TensorShardingAttr eraseFreeAxes(TensorShardingAttr outerManualSharding, -- ArrayRef manualAxes); -+ ArrayRef manualAxes); - - // `dimAxes` is assumed to be the concatenation of some manual axes and some - // free axes. This returns an iterator to the first free axis in `dimAxes`. -@@ -577,7 +578,7 @@ TensorShardingAttr eraseFreeAxes(TensorShardingAttr outerManualSharding, - // Axes in the range [0, firstFreeAxis) are manual axes, and - // [firstFreeAxis, dimAxes.size()) are free axes. - ArrayRef::const_iterator getFirstFreeAxisIter( -- ArrayRef dimAxes, ArrayRef manualAxes); -+ ArrayRef dimAxes, ArrayRef manualAxes); - - // Returns all axes or sub-axes in `axesA` that are not in `axesB`. - // -diff --git a/shardy/dialect/sdy/ir/verifiers.cc b/shardy/dialect/sdy/ir/verifiers.cc -index 31975fe3..e9e41659 100644 ---- a/shardy/dialect/sdy/ir/verifiers.cc -+++ b/shardy/dialect/sdy/ir/verifiers.cc -@@ -60,9 +60,10 @@ using func::FuncOp; - using ::llvm::SmallDenseMap; - using ::llvm::SmallDenseSet; - --// Mapping between an axis name to the ManualComputationOp whose body is manual --// on. --using ManualAxisToOwner = SmallDenseMap; -+// Mapping between a manual axis to the ManualComputationOp whose body is manual -+// on it. -+using ManualAxisToOwner = -+ SmallVector>; - - using EmitErrorFn = std::function; - -@@ -212,15 +213,15 @@ LogicalResult verifySubAxes(ArrayRef subAxes, StringRef axisName, - } - - // ManualComputations op are allowed to be nested within each other. However, --// they cannot operate on the same manual axes. This function creates a mapping --// from a manual mesh axis name to the corresponding ManualComputationOp that -+// they cannot operate on overlapping manual axes. This function creates a -+// mapping from a manual mesh axis to the corresponding ManualComputationOp that - // operates on it to help with verifying this is the case. - ManualAxisToOwner getParentManualComputationOps(Operation* op) { - ManualAxisToOwner alreadyManualAxes; - auto parent = op->getParentOfType(); - while (parent) { -- for (StringRef axisName : parent.getManualAxes()) { -- alreadyManualAxes[axisName] = parent; -+ for (AxisRefAttr axis : parent.getManualAxes()) { -+ alreadyManualAxes.emplace_back(axis, parent); - } - parent = parent->getParentOfType(); - } -@@ -228,11 +229,11 @@ ManualAxisToOwner getParentManualComputationOps(Operation* op) { - } - - LogicalResult emitBoundAxisInManualComputationError(EmitErrorFn emitError, -- StringRef boundAxis, -+ AxisRefAttr boundAxis, - Location parentLoc) { -- return (emitError("operates on axis \"") -- << boundAxis -- << "\" which is already bound by a parent sdy.manual_computation " -+ return (emitError("operates on axis ") -+ << boundAxis.toString() -+ << " which is already bound by a parent sdy.manual_computation " - "op") - .attachNote(parentLoc) - << "parent bounding this axis as manual"; -@@ -383,12 +384,12 @@ LogicalResult verifyTensorShardingAttr(TensorShardingAttr shardingAttr, - // axis due to a parent ManualComputationOp. - if (!alreadyManualAxes.empty()) { - if (shardingAttr.anyOfAxisRef([&](AxisRefAttr axis) { -- auto manualAxisIt = alreadyManualAxes.find(axis.getName()); -- if (manualAxisIt != alreadyManualAxes.end()) { -- return emitBoundAxisInManualComputationError( -- emitError, axis.getName(), -- manualAxisIt->second->getLoc()) -- .failed(); -+ for (const auto& [manualAxis, parent] : alreadyManualAxes) { -+ if (axis.overlaps(manualAxis)) { -+ return emitBoundAxisInManualComputationError(emitError, axis, -+ parent->getLoc()) -+ .failed(); -+ } - } - return false; - })) { -@@ -893,44 +894,71 @@ namespace { - // If an axis in dimShardingAxes belongs to manualAxes, it's an axis - // the user is doing a manual computation on, thus the ManualComputationOp's - // body will have tensors smaller wrt this manual axis. --int64_t accumulatedManualAxesSize( -- Operation* op, ArrayRef dimShardingAxes, -- const llvm::SmallDenseSet& manualAxes, MeshAttr mesh) { -+int64_t accumulatedManualAxesSize(Operation* op, -+ ArrayRef dimShardingAxes, -+ ArrayRef manualAxes, -+ MeshAttr mesh) { - int64_t axesFactor = 1; - for (AxisRefAttr axisRef : dimShardingAxes) { -- if (manualAxes.contains(axisRef.getName())) { -- axesFactor *= axisRef.getSize(mesh); -+ for (AxisRefAttr manualAxis : manualAxes) { -+ if (std::optional overlap = axisRef.getOverlap(manualAxis)) { -+ axesFactor *= overlap->getSize(mesh); -+ } - } - } - - return axesFactor; - } - --// Returns an iterator to a manual axis that comes after a free axis in `axes`. --// If no such axis exists, returns `axes.end()`. --ArrayRef::iterator findManualAxisAfterFreeAxis( -- ArrayRef axes, -- const llvm::SmallDenseSet& manualAxes) { -- auto isManualAxis = [&](AxisRefAttr axisRef) { -- return manualAxes.contains(axisRef.getName()); -- ; -- }; -- return std::find_if(llvm::find_if_not(axes, isManualAxis), axes.end(), -- isManualAxis); -+// Returns a pair of (manualAxis, freeAxis) if a manual axis comes after a free -+// axis in `axes`, or `std::nullopt` if all manual axes come before free axes. -+// Assumes all axes in `axes` and `manualAxes` have already been verified to -+// coexist. -+std::optional> findManualAxisAfterFreeAxis( -+ ArrayRef axes, ArrayRef manualAxes, -+ MeshAttr mesh) { -+ std::optional lastFreeAxis; -+ for (AxisRefAttr axisRef : axes) { -+ AxisRefAttr curr = axisRef; -+ while (curr) { -+ AxisRefAttr earliestManualAxis; -+ for (AxisRefAttr manualAxis : manualAxes) { -+ if (manualAxis.overlaps(curr) && -+ (!earliestManualAxis || manualAxis < earliestManualAxis)) { -+ earliestManualAxis = manualAxis; -+ } -+ } -+ if (!earliestManualAxis) { -+ lastFreeAxis = curr; -+ break; -+ } -+ AxisRefAttr overlap = *curr.getOverlap(earliestManualAxis); -+ if (lastFreeAxis) { -+ return std::make_pair(overlap, *lastFreeAxis); -+ } -+ if (std::optional freePrefix = -+ curr.getPrefixWithoutOverlap(earliestManualAxis)) { -+ return std::make_pair(overlap, *freePrefix); -+ } -+ curr = curr.getSuffixWithoutOverlap(earliestManualAxis, mesh) -+ .value_or(nullptr); -+ } -+ } -+ return std::nullopt; - } - - // For each set of op values (operands/results) and corresponding sharding for - // each value, verifies: --// 1. The sharding list itself wrt the mesh and `globalTypes`. --// 2. The in/out shardings are valid w.r.t the corresponding global type. --// 3. The number of global and local tensor inputs/outputs of the op region -+// 1. The number of global and local tensor inputs/outputs of the op region - // match. --// 4. The manual axes come before any free axes in each dim sharding. --// 5. The manual axes cannot introduce padding. The dimension size must be -+// 2. All axes in the sharding can coexist with `manualAxes`. -+// 3. The manual axes come before any free axes in each dim sharding. -+// 4. The manual axes cannot introduce padding. The dimension size must be - // divisible by the corresponding manual axes size. --// 6. The global shape and local shapes of the op regions arguments/results -+// 5. The global shape and local shapes of the op regions arguments/results - // match. - // -+// Assumes `shardingPerValueAttr` and `manualAxes` have already been verified. - // `valueKindStr` is a string included in any verification error message - // specifying whether the values we are verifying are the operands or results. - // Pass in "operand" or "result". -@@ -938,26 +966,15 @@ template - LogicalResult verifyManualComputationValue( - ManualComputationOp op, GlobalRangeT globalTypes, LocalRangeT localTypes, - TensorShardingPerValueAttr shardingPerValueAttr, -- const SymbolTable& symbolTable, -- const llvm::SmallDenseSet& manualAxesSet, -+ const SymbolTable& symbolTable, ArrayRef manualAxes, - StringRef valueKindStr) { - if (globalTypes.empty() && localTypes.empty() && - shardingPerValueAttr.getShardings().empty()) { - // Nothing to verify since there are no values. - return success(); - } -- // 1. Verify the sharding list itself wrt the mesh and `globalTypes`. -- // 2. Verify the in/out shardings are valid w.r.t the corresponding global -- // type. -- if (failed(verifyTensorShardingPerValueAttr( -- shardingPerValueAttr, globalTypes, op, symbolTable, -- [op, valueKindStr](StringRef msg) { -- return op->emitOpError(valueKindStr) << " " << msg; -- }))) { -- return failure(); -- } - -- // 3. Verify the number of global and local tensor inputs/outputs of the op -+ // 1. Verify the number of global and local tensor inputs/outputs of the op - // region match. - if (globalTypes.size() != localTypes.size()) { - return op->emitOpError("number of op ") -@@ -970,20 +987,39 @@ LogicalResult verifyManualComputationValue( - for (auto [valueIndex, valueEntry] : llvm::enumerate(llvm::zip_equal( - globalTypes, localTypes, shardingPerValueAttr.getShardings()))) { - auto [globalType, localType, sharding] = valueEntry; -+ MeshAttr mesh = sharding.getMesh(symbolTable); -+ -+ // 2. Verify all axes in the sharding can coexist with `manualAxes`. -+ for (AxisRefAttr manualAxis : manualAxes) { -+ AxisRefAttr nonCoexistingAxis; -+ sharding.anyOfAxisRef([&](AxisRefAttr axisRef) { -+ if (!axisRef.canCoexist(manualAxis)) { -+ nonCoexistingAxis = axisRef; -+ return true; -+ } -+ return false; -+ }); -+ if (nonCoexistingAxis) { -+ return op->emitOpError(valueKindStr) -+ << " sharding at index " << valueIndex << " has axis " -+ << nonCoexistingAxis.toString() -+ << " that cannot coexist with manual axis " -+ << manualAxis.toString(); -+ } -+ } - -- // 4. Verify the manual axes come before any free axes in each dim sharding. -+ // 3. Verify the manual axes come before any free axes in each dim sharding. - for (auto [dim, dimSharding] : - llvm::enumerate(sharding.getDimShardings())) { - ArrayRef axes = dimSharding.getAxes(); -- if (ArrayRef::iterator manualAxisItr = -- findManualAxisAfterFreeAxis(axes, manualAxesSet); -- manualAxisItr != axes.end()) { -+ if (std::optional> manualAndFreeAxes = -+ findManualAxisAfterFreeAxis(axes, manualAxes, mesh)) { - return op->emitOpError(valueKindStr) - << " sharding at index " << valueIndex - << " must have all manual axes come before free axes in its " - << "dimension sharding at index " << dim << ". Saw manual axis " -- << manualAxisItr->toString() << " after free axis " -- << (manualAxisItr - 1)->toString(); -+ << manualAndFreeAxes->first.toString() << " after free axis " -+ << manualAndFreeAxes->second.toString(); - } - } - -@@ -999,13 +1035,12 @@ LogicalResult verifyManualComputationValue( - if (dimensionSize == ShapedType::kDynamic) { - newDimSizes.push_back(ShapedType::kDynamic); - } else { -- // 5. The manual axes cannot introduce padding. The dimension size must -+ // 4. The manual axes cannot introduce padding. The dimension size must - // be divisible by the corresponding manual axes size. - - // Safe to call `getMesh` because the sharding was already verified. -- int64_t manualAxesSize = -- accumulatedManualAxesSize(op, dimSharding.getAxes(), manualAxesSet, -- sharding.getMesh(symbolTable)); -+ int64_t manualAxesSize = accumulatedManualAxesSize( -+ op, dimSharding.getAxes(), manualAxes, mesh); - if (dimensionSize % manualAxesSize != 0) { - return op->emitOpError(valueKindStr) - << " dimension size " << dimensionSize -@@ -1015,7 +1050,7 @@ LogicalResult verifyManualComputationValue( - newDimSizes.push_back(dimensionSize / manualAxesSize); - } - } -- // 6. Verify the global shape and local shapes of the op regions -+ // 5. Verify the global shape and local shapes of the op regions - // arguments/results match. - auto expectedLocalRankedType = - RankedTensorType::get(newDimSizes, globalShapedType.getElementType()); -@@ -1046,51 +1081,96 @@ LogicalResult ManualComputationOp::verifySymbolUses( - SymbolTableCollection& symbolTableCollection) { - const SymbolTable& symbolTable = symbolTableCollection.getSymbolTable( - getOperation()->getParentOfType()); -- llvm::SmallDenseSet manualAxesSet(getManualAxes().begin(), -- getManualAxes().end()); -- if (failed(verifyManualComputationValue( -- *this, getOperandTypes(), getBody().getArgumentTypes(), -- getInShardings(), symbolTable, manualAxesSet, "operand")) || -- failed(verifyManualComputationValue( -- *this, getResultTypes(), getBodyTerminatorOpOperandTypes(*this), -- getOutShardings(), symbolTable, manualAxesSet, "result"))) { -+ auto verifyShardingPerValue = -+ [&](TypeRange globalTypes, TypeRange localTypes, -+ TensorShardingPerValueAttr shardingPerValueAttr, -+ StringRef valueKindStr) -> LogicalResult { -+ if (globalTypes.empty() && localTypes.empty() && -+ shardingPerValueAttr.getShardings().empty()) { -+ return success(); -+ } -+ return verifyTensorShardingPerValueAttr( -+ shardingPerValueAttr, globalTypes, *this, symbolTable, -+ [this, valueKindStr](StringRef msg) { -+ return emitOpError(valueKindStr) << " " << msg; -+ }); -+ }; -+ if (failed(verifyShardingPerValue( -+ getOperandTypes(), getBody().getArgumentTypes(), getInShardings(), -+ "operand")) || -+ failed(verifyShardingPerValue( -+ getResultTypes(), getBodyTerminatorOpOperandTypes(*this), -+ getOutShardings(), "result"))) { - return failure(); - } - -- // We don't fail verification if a manual computation has no inputs or outputs -- // and a non-empty body. However, we may require this when running -- // propagation. -- if (getInShardings().empty() && getOutShardings().empty()) { -- return success(); -- } -+ ArrayRef manualAxes = getManualAxes(); -+ if (!getInShardings().empty() || !getOutShardings().empty()) { -+ MeshAttr mesh = getCommonMesh(getInShardings().getShardings(), -+ getOutShardings().getShardings(), -+ symbolTable); -+ if (!mesh && !manualAxes.empty()) { -+ return emitOpError( -+ "all in and out shardings must be bound to the same mesh or an empty " -+ "mesh."); -+ } - -- // We verify a common mesh here so an invalid mesh name reference will be -- // caught before. -- MeshAttr mesh = getCommonMesh(getInShardings().getShardings(), -- getOutShardings().getShardings(), symbolTable); -- if (!mesh && !getManualAxes().empty()) { -- return emitOpError( -- "all in and out shardings must be bound to the same mesh or an empty " -- "mesh."); -- } -+ for (AxisRefAttr axis : manualAxes) { -+ if (mesh && !mesh.hasAxis(axis.getName())) { -+ return emitOpError("unknown manual axis: \"") << axis.getName() << "\""; -+ } -+ } - -- for (StringAttr axisName : getManualAxes()) { -- if (mesh && !mesh.hasAxis(axisName)) { -- return emitOpError("unknown manual axis: ") << axisName; -+ if (mesh) { -+ SmallDenseMap axisNameToSize = -+ mesh.getAxisNameToSize(); -+ SmallDenseSet seenAxisRefs; -+ SmallDenseMap> axisNameToSubAxes; -+ EmitErrorFn emitError = getEmitErrorFn(*this); -+ if (failed(verifyAxisRefList(manualAxes, axisNameToSize, seenAxisRefs, -+ axisNameToSubAxes, emitError))) { -+ return failure(); -+ } -+ for (auto& [axisName, subAxes] : axisNameToSubAxes) { -+ int64_t axisSize = axisNameToSize[axisName]; -+ llvm::sort(subAxes, AxisRefAttr::getMeshComparator(mesh)); -+ if (failed(verifySubAxes(subAxes, axisName, axisSize, mesh, -+ seenAxisRefs, emitError))) { -+ return failure(); -+ } -+ for (auto [prevSubAxis, subAxis] : -+ llvm::zip(subAxes, ArrayRef(subAxes).drop_front())) { -+ if (prevSubAxis.canMerge(subAxis)) { -+ return emitError("two sub-axes in manual_axes can be merged: ") -+ << prevSubAxis.toString() << ", " << subAxis.toString(); -+ } -+ } -+ } - } - } - -+ if (failed(verifyManualComputationValue( -+ *this, getOperandTypes(), getBody().getArgumentTypes(), -+ getInShardings(), symbolTable, manualAxes, "operand")) || -+ failed(verifyManualComputationValue( -+ *this, getResultTypes(), getBodyTerminatorOpOperandTypes(*this), -+ getOutShardings(), symbolTable, manualAxes, "result"))) { -+ return failure(); -+ } -+ - return success(); - } - - LogicalResult ManualComputationOp::verify() { - ManualAxisToOwner alreadyManualAxes = - getParentManualComputationOps(getOperation()); -- for (StringRef axisName : getManualAxes()) { -- if (alreadyManualAxes.contains(axisName)) { -- return emitBoundAxisInManualComputationError( -- [this](StringRef msg) { return emitOpError(msg); }, axisName, -- alreadyManualAxes[axisName]->getLoc()); -+ for (AxisRefAttr axis : getManualAxes()) { -+ for (const auto& [manualAxis, parent] : alreadyManualAxes) { -+ if (axis.overlaps(manualAxis)) { -+ return emitBoundAxisInManualComputationError( -+ [this](StringRef msg) { return emitOpError(msg); }, axis, -+ parent->getLoc()); -+ } - } - } - -diff --git a/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc b/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc -index a665d2e7..8255038b 100644 ---- a/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc -+++ b/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc -@@ -122,8 +122,8 @@ struct ConversionState { - - Value getDeviceId(Location loc, const ConversionState& conversionState, - ConversionPatternRewriter& rewriter) { -- return sdy::getDeviceId(conversionState.replicaCount, -- conversionState.partitionCount, loc, rewriter); -+ return sdy::getDeviceId(loc, conversionState.replicaCount, -+ conversionState.partitionCount, rewriter); - } - - // Returns a channel handle with handle = getNextChannelId() and type = 1 -@@ -235,8 +235,9 @@ class GlobalToLocalTypeConverter : public TypeConverter { - }; - - Operation* createLocalGenericOp( -- Operation* op, ValueRange operands, ConversionPatternRewriter& rewriter, -- const GlobalToLocalTypeConverter* typeConverter) { -+ Operation* op, ValueRange operands, -+ const GlobalToLocalTypeConverter* typeConverter, -+ ConversionPatternRewriter& rewriter) { - OperationState state(op->getLoc(), op->getName()); - state.addOperands(operands); - state.addTypes(typeConverter->convertResultTypes(op->getResults())); -@@ -270,11 +271,11 @@ Operation* createLocalGenericOp( - // Converts op to its local version by replacing its operands with the already - // converted operands and removing the op from the toConvertOps list. - LogicalResult localizeGenericOp(Operation* op, ValueRange operands, -- ConversionPatternRewriter& rewriter, - const GlobalToLocalTypeConverter* typeConverter, -- ConversionState& conversionState) { -+ ConversionState& conversionState, -+ ConversionPatternRewriter& rewriter) { - Operation* newOp = -- createLocalGenericOp(op, operands, rewriter, typeConverter); -+ createLocalGenericOp(op, operands, typeConverter, rewriter); - rewriter.replaceOp(op, newOp->getResults()); - conversionState.removeToConvertOp(op); - return success(); -@@ -305,9 +306,9 @@ class GenericOpPattern : public ConversionPattern { - return failure(); - } - return localizeGenericOp( -- op, operands, rewriter, -+ op, operands, - static_cast(typeConverter), -- conversionState); -+ conversionState, rewriter); - } - - private: -@@ -335,8 +336,8 @@ class PartitioningOpInterfacePattern - conversionState.removeToConvertOp(op); - return success(); - } -- return localizeGenericOp(op, operands, rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, operands, converter, conversionState, -+ rewriter); - } - - private: -@@ -555,10 +556,10 @@ ValueRange createAllReduceForUnreducedAxes( - // Splits `combinedDim` of `input` into `factors`, transposes each factor - // immediately before its corresponding dimension in `targetDims`, and reshapes - // to merge each factor into that dimension. --Value distributeCombinedDimToTargetDims(Value input, int64_t combinedDim, -+Value distributeCombinedDimToTargetDims(Location loc, Value input, -+ int64_t combinedDim, - ArrayRef factors, - ArrayRef targetDims, -- Location loc, - ConversionPatternRewriter& rewriter) { - auto inputType = cast(input.getType()); - ArrayRef combinedShape = inputType.getShape(); -@@ -666,8 +667,8 @@ class AllGatherOpPattern : public OpConversionPattern { - curInput, /*all_gather_dim=*/firstGatheringDim, replicaGroups, - channelHandle, /*use_global_device_ids=*/true); - curInput = distributeCombinedDimToTargetDims( -- allGather.getResult(0), firstGatheringDim, gatheringFactors, -- gatheringDims, op.getLoc(), rewriter); -+ op.getLoc(), allGather.getResult(0), firstGatheringDim, -+ gatheringFactors, gatheringDims, rewriter); - - conversionState.removeToConvertOp(op); - rewriter.replaceOp(op, curInput); -@@ -846,7 +847,7 @@ Value emitDynamicSliceForAxes(Location loc, Value globalTensor, MeshAttr mesh, - for (int64_t i = 0; i < localResultType.getRank(); ++i) { - ArrayRef axes = slicingAxesPerDim[i].getValue(); - if (axes.empty()) { -- startIndices.push_back(createZeroConstant(rewriter, loc, indexTy)); -+ startIndices.push_back(createZeroConstant(loc, indexTy, rewriter)); - continue; - } - -@@ -879,7 +880,7 @@ SmallVector maskNonLeaderShards( - conversionState, rewriter, deviceId); - if (!nonLeaderVal) { - auto localType = cast(inputs.front().getType()); -- nonLeaderVal = createZeroConstant(rewriter, loc, localType); -+ nonLeaderVal = createZeroConstant(loc, localType, rewriter); - } - return llvm::map_to_vector(inputs, [&](Value input) -> Value { - return stablehlo::SelectOp::create(rewriter, loc, isLeader, input, -@@ -1090,8 +1091,8 @@ class AllToAllOpPattern : public OpConversionPattern { - /*split_count=*/numDevicesPerGroup, replicaGroups, channelHandle); - - Value result = distributeCombinedDimToTargetDims( -- allToAll.getResult(0), minAllToAllDim, numDevicesPerParams, srcDims, -- loc, rewriter); -+ loc, allToAll.getResult(0), minAllToAllDim, numDevicesPerParams, -+ srcDims, rewriter); - - rewriter.replaceOp(op, result); - conversionState.removeToConvertOp(op); -@@ -1609,7 +1610,7 @@ class ShardedToUnreducedOpPattern - if (!localResultType.hasStaticShape()) { - return op.emitOpError("dynamic shape not supported"); - } -- Value zeroTensor = createZeroConstant(rewriter, loc, localResultType); -+ Value zeroTensor = createZeroConstant(loc, localResultType, rewriter); - - auto indexTy = RankedTensorType::get({}, rewriter.getI32Type()); - Value zeroOffset; -@@ -1619,7 +1620,7 @@ class ShardedToUnreducedOpPattern - ArrayRef axes = op.getAxes()[i].getValue(); - if (axes.empty()) { - if (!zeroOffset) { -- zeroOffset = createZeroConstant(rewriter, loc, indexTy); -+ zeroOffset = createZeroConstant(loc, indexTy, rewriter); - } - startIndices.push_back(zeroOffset); - continue; -@@ -1745,8 +1746,8 @@ class StablehloConcatenateOpPattern - - // Since the concat dimension is replicated and others are consistent, - // we can safely perform a local concatenation. -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - private: -@@ -1790,8 +1791,8 @@ class StablehloConvolutionOpPattern - - if (isFullyReplicated(lhsSharding) && isFullyReplicated(rhsSharding) && - isFullyReplicated(resultSharding)) { -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - MeshAttr mesh; -@@ -1944,7 +1945,7 @@ class StablehloDotOpPattern : public OpConversionPattern { - if (!doesOpIntroducePartialResult(op, contractingAxes, "dot", - converter)) { - Operation* localDot = createLocalGenericOp(op, adaptor.getOperands(), -- rewriter, converter); -+ converter, rewriter); - ValueRange results = createAllReduceForUnreducedAxes( - op.getLoc(), localDot->getResults(), contractingAxes, mesh, - lhsSharding.getMeshOrRef(), enableRGV3, conversionState, -@@ -1956,8 +1957,8 @@ class StablehloDotOpPattern : public OpConversionPattern { - } - } - -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - private: -@@ -2021,7 +2022,7 @@ class StablehloDotGeneralOpPattern - if (!doesOpIntroducePartialResult(op, reductionAxes, "dot_general", - converter)) { - Operation* localDotGeneral = createLocalGenericOp( -- op, adaptor.getOperands(), rewriter, converter); -+ op, adaptor.getOperands(), converter, rewriter); - ValueRange results = createAllReduceForUnreducedAxes( - op.getLoc(), localDotGeneral->getResults(), reductionAxes, mesh, - lhsSharding.getMeshOrRef(), enableRGV3, conversionState, -@@ -2033,8 +2034,8 @@ class StablehloDotGeneralOpPattern - } - } - -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - private: -@@ -2059,8 +2060,8 @@ class StablehloDynamicSliceOpPattern - TensorShardingAttr outSharding = converter->getSharding(op.getResult()); - - if (isFullyReplicated(inSharding) && isFullyReplicated(outSharding)) { -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - auto globalInputType = op.getOperand().getType(); -@@ -2107,8 +2108,8 @@ class StablehloDynamicSliceOpPattern - ConversionState& conversionState; - }; - --Value createScalarInt(OpBuilder& b, Location loc, RankedTensorType type, -- int64_t val) { -+Value createScalarInt(Location loc, RankedTensorType type, int64_t val, -+ OpBuilder& b) { - return stablehlo::ConstantOp::create( - b, loc, - DenseIntElementsAttr::get(type, -@@ -2124,12 +2125,12 @@ Value computeOwnershipPredicate(Location loc, Value clampedIndex, Value offset, - RankedTensorType scalarIndexType, - ConversionPatternRewriter& rewriter) { - Value sliceSizeConst = -- createScalarInt(rewriter, loc, scalarIndexType, sliceDim); -+ createScalarInt(loc, scalarIndexType, sliceDim, rewriter); - Value clampedPlusSlice = stablehlo::AddOp::create( - rewriter, loc, scalarIndexType, clampedIndex, sliceSizeConst); - - Value shardSizeConst = -- createScalarInt(rewriter, loc, scalarIndexType, shardDim); -+ createScalarInt(loc, scalarIndexType, shardDim, rewriter); - Value offsetPlusShard = stablehlo::AddOp::create( - rewriter, loc, scalarIndexType, offset, shardSizeConst); - -@@ -2158,9 +2159,9 @@ LocalIndexAndOwnership computeLocalIndexAndOwnershipForDim( - RankedTensorType scalarIndexType = RankedTensorType::get({}, indexEltTy); - - // Clamp the start index to [0, globalDim - sliceDim]. -- Value zeroConst = createScalarInt(rewriter, loc, scalarIndexType, 0); -+ Value zeroConst = createScalarInt(loc, scalarIndexType, 0, rewriter); - Value maxConst = -- createScalarInt(rewriter, loc, scalarIndexType, globalDim - sliceDim); -+ createScalarInt(loc, scalarIndexType, globalDim - sliceDim, rewriter); - Value clampedIndex = stablehlo::ClampOp::create( - rewriter, loc, scalarIndexType, zeroConst, rawIndex, maxConst); - -@@ -2275,8 +2276,8 @@ class StablehloDynamicUpdateSliceOpPattern - TensorShardingAttr operandSharding = - converter->getSharding(op.getOperand()); - if (!operandSharding || isFullyReplicated(operandSharding)) { -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - MeshAttr mesh = operandSharding.getMesh(converter->getSymbolTable()); -@@ -2312,8 +2313,8 @@ class StablehloDynamicUpdateSliceOpPattern - } - } - if (isPartitionedSliceDim.none()) { -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - RankedTensorType localOperandType = -@@ -2376,7 +2377,7 @@ class StablehloReduceOpPattern - if (!reductionAxes.empty() && - !doesOpIntroducePartialResult(op, reductionAxes, "reduce", converter)) { - auto localReduceOp = cast( -- createLocalGenericOp(op, adaptor.getOperands(), rewriter, converter)); -+ createLocalGenericOp(op, adaptor.getOperands(), converter, rewriter)); - if (op.getNumResults() == 1) { - ValueRange results = createAllReduceForUnreducedAxes( - op.getLoc(), localReduceOp.getResults(), localReduceOp.getBody(), -@@ -2433,8 +2434,8 @@ class StablehloReduceOpPattern - return success(); - } - -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - private: -@@ -2482,8 +2483,8 @@ std::pair, SmallVector> getTrivialSliceDimBounds( - Location loc, MeshAttr mesh, TensorShardingAttr operandSharding, - RankedTensorType globalOperandType, RankedTensorType localOperandType, - ArrayRef startIndexMap, ArrayRef trivialSliceDims, -- Type indexEltTy, const ConversionState& conversionState, -- ConversionPatternRewriter& rewriter, bool computeMask) { -+ Type indexEltTy, const ConversionState& conversionState, bool computeMask, -+ ConversionPatternRewriter& rewriter) { - SmallVector minBounds, maxBounds; - minBounds.reserve(startIndexMap.size()); - if (computeMask) { -@@ -2495,7 +2496,7 @@ std::pair, SmallVector> getTrivialSliceDimBounds( - if (!llvm::is_contained(trivialSliceDims, indexedDim)) { - // For non-trivial dims, the bounds cover the entire global range. - minBounds.push_back(createZeroConstant( -- rewriter, loc, RankedTensorType::get({}, indexEltTy))); -+ loc, RankedTensorType::get({}, indexEltTy), rewriter)); - if (computeMask) { - maxBounds.push_back(stablehlo::ConstantOp::create( - rewriter, loc, -@@ -2561,7 +2562,7 @@ Value clampGatherIndices(Location loc, Value indices, - - return stablehlo::ClampOp::create( - rewriter, loc, indicesType, -- createZeroConstant(rewriter, loc, indicesType), indices, -+ createZeroConstant(loc, indicesType, rewriter), indices, - buildIndicesLikeTensor(globalMaxValues)); - } - -@@ -2610,8 +2611,8 @@ std::pair computeLocalIndicesAndMask( - } - auto [minBounds, maxBounds] = getTrivialSliceDimBounds( - loc, mesh, operandSharding, globalOperandType, localOperandType, -- startIndexMap, trivialSliceDims, indexEltTy, conversionState, rewriter, -- computeMask); -+ startIndexMap, trivialSliceDims, indexEltTy, conversionState, -+ computeMask, rewriter); - - // Builds bounds tensors matching the indices tensor type. - auto buildBoundsTensor = [&](ArrayRef bounds) -> Value { -@@ -2685,7 +2686,7 @@ std::pair computeLocalIndicesAndMask( - int64_t shardSize = localOperandType.getDimSize(dim); - Value offset = getDimensionOffset(loc, mesh, axes, shardSize, - conversionState, rewriter); -- Value zero = createZeroConstant(rewriter, loc, offset.getType()); -+ Value zero = createZeroConstant(loc, offset.getType(), rewriter); - Value eqZero = stablehlo::CompareOp::create( - rewriter, loc, offset, zero, stablehlo::ComparisonDirection::EQ); - partitionMask = -@@ -2726,8 +2727,8 @@ class StablehloGatherOpPattern - converter->getSharding(op.getOperand()); - - if (isFullyReplicated(operandSharding)) { -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - MeshAttr mesh = operandSharding.getMesh(converter->getSymbolTable()); -@@ -2842,7 +2843,7 @@ class StablehloGatherOpPattern - // Fill the out-of-bounds positions with the reduction identity. - Value identityVal = - reductionOp == ReductionOp::SUM -- ? createZeroConstant(rewriter, loc, resultType) -+ ? createZeroConstant(loc, resultType, rewriter) - : stablehlo::ConstantOp::create( - rewriter, loc, - DenseElementsAttr::get( -@@ -2963,8 +2964,8 @@ class StablehloRngBitGeneratorOpPattern - static_cast(getTypeConverter()); - TensorShardingAttr dataSharding = converter->getSharding(op.getOutput()); - if (!dataSharding || isFullyReplicated(dataSharding)) { -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - MeshAttr mesh = dataSharding.getMesh(converter->getSymbolTable()); -@@ -3023,7 +3024,7 @@ class StablehloRngBitGeneratorOpPattern - Value hi = stablehlo::AddOp::create( - rewriter, loc, x64CombineOp.getOperand(1), broadcastHash); - adjustedState = createLocalGenericOp(x64CombineOp, ValueRange{lo, hi}, -- rewriter, converter) -+ converter, rewriter) - ->getResult(0); - } else { - adjustedState = -@@ -3093,8 +3094,8 @@ class StablehloPadOpPattern : public OpConversionPattern { - // else the pad op should have been converted to a pad followed by a - // collective to reshard its result. - SDY_CHECK(isFullyReplicated(outSharding)); -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - const SymbolTable& symbolTable = converter->getSymbolTable(); - MeshAttr mesh = inSharding.getMesh(symbolTable); -@@ -3310,7 +3311,7 @@ Value getScatterReductionIdentity(stablehlo::ScatterOp scatter, OpBuilder& b) { - - return llvm::TypeSwitch(reductionOp) - .Case([&](stablehlo::AddOp) { -- return createZeroConstant(b, loc, scalarType); -+ return createZeroConstant(loc, scalarType, b); - }) - .Case([&](stablehlo::AndOp) { - return stablehlo::ConstantOp::create( -@@ -3319,7 +3320,7 @@ Value getScatterReductionIdentity(stablehlo::ScatterOp scatter, OpBuilder& b) { - b.getIntegerAttr(elementType, 1))); - }) - .Case([&](stablehlo::OrOp) { -- return createZeroConstant(b, loc, scalarType); -+ return createZeroConstant(loc, scalarType, b); - }) - .Case([&](stablehlo::MulOp) { - if (isa(elementType)) { -@@ -3381,8 +3382,8 @@ class StablehloScatterOpPattern - TensorShardingAttr inputSharding = converter->getSharding(input); - TensorShardingAttr updateSharding = converter->getSharding(update); - if (isFullyReplicated(inputSharding) && isFullyReplicated(updateSharding)) { -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - const SymbolTable& symbolTable = converter->getSymbolTable(); -@@ -3395,8 +3396,8 @@ class StablehloScatterOpPattern - SmallVector reductionAxes = - getScatterReductionAxes(op, updateSharding, inputSharding, mesh); - if (reductionAxes.empty() && isFullyReplicated(inputSharding)) { -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - Location loc = op.getLoc(); -@@ -3441,8 +3442,8 @@ class StablehloScatterOpPattern - // If no indexed dims are sharded and there is no reduction axes, just - // localize the types. - if (indexedInputDims.empty() && reductionAxes.empty()) { -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - // Adjust indices for sharded dimensions (Disjoint writes are parallel). -@@ -3506,8 +3507,8 @@ class StablehloSliceOpPattern : public OpConversionPattern { - TensorShardingAttr outSharding = converter->getSharding(op.getResult()); - - if (isFullyReplicated(inSharding) && isFullyReplicated(outSharding)) { -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - auto globalInputType = op.getOperand().getType(); -@@ -3650,8 +3651,8 @@ class StablehloWindowedOpPattern : public OpConversionPattern { - } - } - -- return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter, -- conversionState); -+ return localizeGenericOp(op, adaptor.getOperands(), converter, -+ conversionState, rewriter); - } - - private: -diff --git a/shardy/dialect/sdy/transforms/export/export_pipeline.cc b/shardy/dialect/sdy/transforms/export/export_pipeline.cc -index 28a9534e..9be5abc6 100644 ---- a/shardy/dialect/sdy/transforms/export/export_pipeline.cc -+++ b/shardy/dialect/sdy/transforms/export/export_pipeline.cc -@@ -51,6 +51,9 @@ void runShardyPartitioner(OpPassManager& pm, int& dumpIndex, - // Catch the cases where unreduced axes are dropped and cause inconsistencies. - pm.addNestedPass(createVerifyUnreducedAxesPass()); - InsertExplicitReshardsPassOptions passOptions; -+ passOptions.enableFullVersion = -+ options.enableInsertExplicitCollectives || -+ options.partitionerStage != PartitionerStage::kUnspecified; - passOptions.markPartialResultWithUnreducedAxes = - options.markPartialResultWithUnreducedAxes || - options.partitionerStage != PartitionerStage::kUnspecified; -diff --git a/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc b/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc -index 47e31742..cb9261ca 100644 ---- a/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc -+++ b/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc -@@ -59,12 +59,12 @@ enum class PaddingValueKind { kZero, kOne }; - constexpr PaddingValueKind kDefaultPaddingValueKind = PaddingValueKind::kZero; - - // Returns a constant for the given PaddingValueKind. --Value createConstant(OpBuilder& b, Location loc, Type elementType, -- PaddingValueKind kind) { -+Value createConstant(Location loc, Type elementType, PaddingValueKind kind, -+ OpBuilder& b) { - auto type = RankedTensorType::get({}, elementType); - switch (kind) { - case PaddingValueKind::kZero: -- return createZeroConstant(b, loc, type); -+ return createZeroConstant(loc, type, b); - case PaddingValueKind::kOne: - if (auto floatType = dyn_cast(elementType)) { - return stablehlo::ConstantOp::create( -@@ -92,51 +92,6 @@ std::optional getPaddingValueKind(Value paddingValue) { - return std::nullopt; - } - --// Helper methods for commonly created ops (stablehlo.slice and stablehlo.pad). -- --// Slices a tensor to `targetType` along dimensions that need trimming. --Value createSliceOp(OpBuilder& b, Location loc, RankedTensorType targetType, -- Value input, TensorShardingAttr sharding = nullptr) { -- SmallVector starts(targetType.getRank(), 0); -- SmallVector strides(targetType.getRank(), 1); -- SmallVector limits(targetType.getShape().begin(), -- targetType.getShape().end()); -- Value sliceRes = stablehlo::SliceOp::create(b, loc, targetType, input, -- b.getDenseI64ArrayAttr(starts), -- b.getDenseI64ArrayAttr(limits), -- b.getDenseI64ArrayAttr(strides)) -- .getResult(); -- if (sharding) { -- setSharding(sliceRes, sharding); -- } -- return sliceRes; --} -- --// Pads a tensor to `targetType` with `paddingKind` along dimensions that need --// expansion. --Value createPadOp(OpBuilder& b, Location loc, RankedTensorType targetType, -- Value input, TensorShardingAttr sharding = nullptr, -- PaddingValueKind paddingKind = kDefaultPaddingValueKind) { -- auto inputType = cast(input.getType()); -- SmallVector edgePaddingHigh(targetType.getRank(), 0); -- for (int d = 0; d < targetType.getRank(); ++d) { -- edgePaddingHigh[d] = targetType.getDimSize(d) - inputType.getDimSize(d); -- } -- Value padding = -- createConstant(b, loc, targetType.getElementType(), paddingKind); -- SmallVector zeroPadding(targetType.getRank(), 0); -- Value padRes = -- stablehlo::PadOp::create(b, loc, targetType, input, padding, -- b.getDenseI64ArrayAttr(zeroPadding), -- b.getDenseI64ArrayAttr(edgePaddingHigh), -- b.getDenseI64ArrayAttr(zeroPadding)) -- .getResult(); -- if (sharding) { -- setSharding(padRes, sharding); -- } -- return padRes; --} -- - class PaddedTypeConverter : public TypeConverter { - public: - explicit PaddedTypeConverter(const SymbolTable& symbolTable) -@@ -282,12 +237,14 @@ class PaddingCache { - - // Creates a padded value for 'value' with 'paddedType' and 'paddingKind', - // and adds the padded value to the PaddingCache if registeredKind is present. --Value createPaddedValue(RankedTensorType paddedType, Value value, -+Value createPaddedValue(Value value, RankedTensorType paddedType, - PaddingValueKind paddingKind, - std::optional registeredKind, -- OpBuilder& rewriter, PaddingCache& cache) { -- Value padOp = createPadOp(rewriter, value.getLoc(), paddedType, value, -- getSharding(value), paddingKind); -+ PaddingCache& cache, OpBuilder& rewriter) { -+ Value padding = createConstant( -+ value.getLoc(), paddedType.getElementType(), paddingKind, rewriter); -+ Value padOp = padHighSideToType(value.getLoc(), value, paddedType, -+ getSharding(value), rewriter, padding); - if (registeredKind) { - cache.setPadding(padOp, *registeredKind); - } -@@ -333,8 +290,8 @@ SmallVector computeReconciledResultTypes(Operation* op, - // unless dimsToEnforce is provided, in which case only the specified - // dimensions are padded. - Value ensurePaddingWithValue( -- Value inputVal, RankedTensorType origType, Value padVal, OpBuilder& b, -- Location loc, -+ Location loc, Value inputVal, RankedTensorType origType, Value padVal, -+ OpBuilder& b, - std::optional> dimsToEnforce = std::nullopt) { - auto paddedType = cast(inputVal.getType()); - if (origType == paddedType) { -@@ -402,8 +359,8 @@ Value ensurePaddingWithValue( - // PaddingValueKind to determine the padding value instead of a user-provided - // value and may put the result into the PaddingCache. - Value ensurePaddingWithKind( -- Value inputVal, RankedTensorType origType, PaddingValueKind requiredKind, -- OpBuilder& b, Location loc, PaddingCache& cache, -+ Location loc, Value inputVal, RankedTensorType origType, -+ PaddingValueKind requiredKind, PaddingCache& cache, OpBuilder& b, - std::optional> dimsToEnforce = std::nullopt) { - // Return early if no padding is applied or the cached padding already - // matches. -@@ -416,9 +373,9 @@ Value ensurePaddingWithKind( - return inputVal; - } - Value newPaddingScalar = -- createConstant(b, loc, paddedType.getElementType(), requiredKind); -- Value select = ensurePaddingWithValue(inputVal, origType, newPaddingScalar, b, -- loc, dimsToEnforce); -+ createConstant(loc, paddedType.getElementType(), requiredKind, b); -+ Value select = ensurePaddingWithValue(loc, inputVal, origType, -+ newPaddingScalar, b, dimsToEnforce); - if (select != inputVal && !dimsToEnforce) { - cache.setPadding(select, requiredKind); - } -@@ -429,8 +386,8 @@ Value ensurePaddingWithKind( - // Converts op to its local version by replacing its operands with the already - // converted operands. - LogicalResult padGenericOp(Operation* op, ValueRange operands, -- ConversionPatternRewriter& rewriter, -- const PaddedTypeConverter* typeConverter) { -+ const PaddedTypeConverter* typeConverter, -+ ConversionPatternRewriter& rewriter) { - // Compute padded shapes for results. - SmallVector inferredTypes; - if (auto inferTypeOp = dyn_cast(op)) { -@@ -484,9 +441,9 @@ LogicalResult padGenericOp(Operation* op, ValueRange operands, - // Slices/trims `res` back to its original sizes on `trimDims`. - Value trimOutputForDims(Value res, Type origType, ArrayRef trimDims, - TensorShardingAttr outSharding, -- ConversionPatternRewriter& rewriter, - std::optional paddingKind, -- PaddingCache& cache) { -+ PaddingCache& cache, -+ ConversionPatternRewriter& rewriter) { - auto origRanked = dyn_cast(origType); - auto newRanked = dyn_cast(res.getType()); - if (!origRanked || !newRanked) { -@@ -514,7 +471,7 @@ Value trimOutputForDims(Value res, Type origType, ArrayRef trimDims, - auto targetType = - RankedTensorType::get(limitIndices, origRanked.getElementType()); - Value sliceRes = -- createSliceOp(rewriter, res.getLoc(), targetType, res, outSharding); -+ sliceHighSideToType(res.getLoc(), res, targetType, outSharding, rewriter); - if (paddingKind) { - cache.setPadding(sliceRes, *paddingKind); - } -@@ -526,8 +483,8 @@ Value trimOutputForDims(Value res, Type origType, ArrayRef trimDims, - // cache. - Operation* padCollectiveOp(Operation* op, Value input, Value inputOrig, - RankedTensorType paddedInputType, -- ConversionPatternRewriter& rewriter, -- PaddingCache& cache) { -+ PaddingCache& cache, -+ ConversionPatternRewriter& rewriter) { - auto rankedInput = cast(input.getType()); - Value padOp = input; - std::optional paddingKind = cache.getPadding(input); -@@ -543,8 +500,8 @@ Operation* padCollectiveOp(Operation* op, Value input, Value inputOrig, - } else if (paddingKind) { - registeredKind = paddingKind; - } -- padOp = createPaddedValue(paddedInputType, input, constantKind, -- registeredKind, rewriter, cache); -+ padOp = createPaddedValue(input, paddedInputType, constantKind, -+ registeredKind, cache, rewriter); - paddingKind = registeredKind; - } - -@@ -577,8 +534,9 @@ class GenericOpPattern : public ConversionPattern { - hasCustomPadHandling(op)) { - return failure(); - } -- return padGenericOp(op, operands, rewriter, -- static_cast(typeConverter)); -+ return padGenericOp(op, operands, -+ static_cast(typeConverter), -+ rewriter); - } - }; - -@@ -597,7 +555,7 @@ class PartitioningOpInterfacePattern - operands, rewriter, *converter, converter->getSymbolTable())) { - return *res; - } -- return padGenericOp(op, operands, rewriter, converter); -+ return padGenericOp(op, operands, converter, rewriter); - } - }; - -@@ -728,7 +686,7 @@ class AllGatherOpPattern : public OpConversionPattern { - Value res = newOp->getResult(i); - Value trimmed = - trimOutputForDims(res, op->getResult(i).getType(), trimDims, -- getSharding(res), rewriter, paddingKind, cache); -+ getSharding(res), paddingKind, cache, rewriter); - replacements.push_back(trimmed); - } - -@@ -805,7 +763,7 @@ class AllSliceOpPattern : public OpConversionPattern { - getDivisiblePaddedType(rankedInputType, outSharding, symbolTable)); - - Operation* newOp = -- padCollectiveOp(op, input, inputOrig, paddedInputType, rewriter, cache); -+ padCollectiveOp(op, input, inputOrig, paddedInputType, cache, rewriter); - - rewriter.replaceOp(op, newOp->getResults()); - return success(); -@@ -841,7 +799,7 @@ class ReduceScatterOpPattern - getDivisiblePaddedType(rankedInputType, outSharding, symbolTable)); - - Operation* newOp = -- padCollectiveOp(op, input, inputOrig, paddedInputType, rewriter, cache); -+ padCollectiveOp(op, input, inputOrig, paddedInputType, cache, rewriter); - - rewriter.replaceOp(op, newOp->getResults()); - return success(); -@@ -887,13 +845,13 @@ class AllToAllOpPattern : public OpConversionPattern { - - std::optional paddingKind = cache.getPadding(input); - Operation* newOp = -- padCollectiveOp(op, input, inputOrig, paddedInputType, rewriter, cache); -+ padCollectiveOp(op, input, inputOrig, paddedInputType, cache, rewriter); - - Value res = newOp->getResult(0); - - SmallVector trimDims = getTrimDims(op); - res = trimOutputForDims(res, inputOrig.getType(), trimDims, outSharding, -- rewriter, paddingKind, cache); -+ paddingKind, cache, rewriter); - - rewriter.replaceOp(op, {res}); - return success(); -@@ -926,7 +884,7 @@ class StablehloReshapeOpPattern - "should have been resolved by resolve-permutation-factors."); - } - -- return padGenericOp(op, adaptor.getOperands(), rewriter, converter); -+ return padGenericOp(op, adaptor.getOperands(), converter, rewriter); - } - }; - -@@ -946,7 +904,7 @@ class StablehloSliceOpPattern : public OpConversionPattern { - resultType, sharding, converter->getSymbolTable())); - - if (paddedType == resultType) { -- return padGenericOp(op, adaptor.getOperands(), rewriter, converter); -+ return padGenericOp(op, adaptor.getOperands(), converter, rewriter); - } - - ArrayRef paddedShape = paddedType.getShape(); -@@ -1037,9 +995,9 @@ class StablehloPadOpPattern : public OpConversionPattern { - // When high padding is reduced, the input's padding region takes the - // place of the trailing padding value for the pad op. As such, we need - // to ensure the input is padded with the correct value. -- input = ensurePaddingWithValue(input, op.getOperand().getType(), -- adaptor.getOperands()[1], rewriter, -- op.getLoc()); -+ input = ensurePaddingWithValue(op.getLoc(), input, -+ op.getOperand().getType(), -+ adaptor.getOperands()[1], rewriter); - } - - auto newHighAttr = rewriter.getDenseI64ArrayAttr(newHigh); -@@ -1094,13 +1052,13 @@ class StablehloDotGeneralOpPattern - return failure(); - } - Value paddedLhs = ensurePaddingWithKind( -- lhs, lhsOrigType, PaddingValueKind::kZero, rewriter, loc, cache, -+ loc, lhs, lhsOrigType, PaddingValueKind::kZero, cache, rewriter, - dimNums.getLhsContractingDimensions()); - Value paddedRhs = ensurePaddingWithKind( -- rhs, rhsOrigType, PaddingValueKind::kZero, rewriter, loc, cache, -+ loc, rhs, rhsOrigType, PaddingValueKind::kZero, cache, rewriter, - dimNums.getRhsContractingDimensions()); - -- return padGenericOp(op, {paddedLhs, paddedRhs}, rewriter, converter); -+ return padGenericOp(op, {paddedLhs, paddedRhs}, converter, rewriter); - } - - private: -@@ -1144,11 +1102,11 @@ class StablehloConvolutionOpPattern - // pass-through dimensions (like batch or kernel output feature) which are - // left out of lhsEnforceDims and rhsEnforceDims. - Value paddedLhs = -- ensurePaddingWithKind(lhs, lhsOrigType, PaddingValueKind::kZero, -- rewriter, loc, cache, lhsEnforceDims); -+ ensurePaddingWithKind(loc, lhs, lhsOrigType, PaddingValueKind::kZero, -+ cache, rewriter, lhsEnforceDims); - Value paddedRhs = -- ensurePaddingWithKind(rhs, rhsOrigType, PaddingValueKind::kZero, -- rewriter, loc, cache, rhsEnforceDims); -+ ensurePaddingWithKind(loc, rhs, rhsOrigType, PaddingValueKind::kZero, -+ cache, rewriter, rhsEnforceDims); - - int64_t newFeatureGroupCount = op.getFeatureGroupCount(); - int64_t origInputFeatureSize = -@@ -1257,7 +1215,7 @@ class StablehloConvolutionOpPattern - Value finalResult = - trimOutputForDims(newOp->getResult(0), paddedResultShaped, - dimNums.getOutputSpatialDimensions(), outSharding, -- rewriter, PaddingValueKind::kZero, cache); -+ PaddingValueKind::kZero, cache, rewriter); - - rewriter.replaceOp(op, finalResult); - return success(); -@@ -1416,13 +1374,13 @@ class StablehloReduceOpPattern - allOperands.push_back(input); - } else { - allOperands.push_back(ensurePaddingWithValue( -- input, origType, initVal, rewriter, loc, op.getDimensions())); -+ loc, input, origType, initVal, rewriter, op.getDimensions())); - } - } - allOperands.append(adaptor.getInitValues().begin(), - adaptor.getInitValues().end()); - -- return padGenericOp(op, allOperands, rewriter, converter); -+ return padGenericOp(op, allOperands, converter, rewriter); - } - - private: -diff --git a/shardy/dialect/sdy/transforms/export/passes.td b/shardy/dialect/sdy/transforms/export/passes.td -index de4db467..9fe5f023 100644 ---- a/shardy/dialect/sdy/transforms/export/passes.td -+++ b/shardy/dialect/sdy/transforms/export/passes.td -@@ -130,12 +130,12 @@ def InsertExplicitReshardsPass : Pass<"sdy-insert-explicit-reshards", "func::Fun - operation has compatible shardings. - }]; - let dependentDialects = ["mlir::sdy::SdyDialect"]; -- // TODO(b/402429253): remove once the full version has landed and been stable for some time. -+ // TODO(b/402429253): delete when explicit reshards is enabled by default. - // The non-full version allows us to insert explicit reshards on specific - // operations for optimizations. - let options = [ - Option<"enableFullVersion", "enable-full-version", -- "bool", /*default=*/"true", -+ "bool", /*default=*/"false", - "Enable full version.">, - Option<"markPartialResultWithUnreducedAxes", - "mark-partial-result-with-unreduced-axes", -diff --git a/shardy/dialect/sdy/transforms/export/per_instruction_partitioning.cc b/shardy/dialect/sdy/transforms/export/per_instruction_partitioning.cc -index ce34913e..2c6b93f5 100644 ---- a/shardy/dialect/sdy/transforms/export/per_instruction_partitioning.cc -+++ b/shardy/dialect/sdy/transforms/export/per_instruction_partitioning.cc -@@ -247,7 +247,7 @@ void collectTransitiveCallees(Operation* op, const SymbolTable& symbolTable, - struct InstructionShardingInfo { - SmallVector inShardings; - SmallVector outShardings; -- SmallVector manualAxesList; -+ SmallVector manualAxesList; - llvm::DenseSet manualAxesSet; - }; - -@@ -299,19 +299,19 @@ InstructionShardingInfo getInstructionShardingInfo(Operation* op, - - if (meshAttr) { - for (MeshAxisAttr axis : meshAttr.getAxes()) { -- info.manualAxesList.push_back(StringAttr::get(ctx, axis.getName())); -+ info.manualAxesList.push_back(AxisRefAttr::get(ctx, axis.getName())); - info.manualAxesSet.insert(axis.getName()); - } - } else { -- llvm::SmallDenseSet seenAxes; -+ llvm::SmallDenseSet seenAxes; - auto collectAxes = [&](TensorShardingAttr sharding) { - if (!sharding) { - return; - } - sharding.forEachAxisRef([&](AxisRefAttr axisRef) { -- StringAttr axisName = StringAttr::get(ctx, axisRef.getName()); -- if (seenAxes.insert(axisName).second) { -- info.manualAxesList.push_back(axisName); -+ AxisRefAttr fullAxis = AxisRefAttr::get(ctx, axisRef.getName()); -+ if (seenAxes.insert(fullAxis).second) { -+ info.manualAxesList.push_back(fullAxis); - info.manualAxesSet.insert(axisRef.getName()); - } - }); -@@ -470,7 +470,7 @@ SmallVector padIndivisibleLiveInValues( - for (auto [operand, paddedType, sharding] : llvm::zip_equal( - liveInValues, paddedArgTypes, shardingInfo.inShardings)) { - manualOperands.push_back(padHighSideToType( -- rewriter, op->getLoc(), operand, paddedType, sharding, -+ op->getLoc(), operand, paddedType, sharding, rewriter, - /*paddingValue=*/nullptr, /*allowSlicePeephole=*/true)); - } - return manualOperands; -@@ -561,8 +561,8 @@ LogicalResult outlineInstruction(Operation* op, - Value arg = block->getArgument(argIdx); - TensorShardingAttr operandSharding = shardingInfo.inShardings[argIdx]; - if (!isReshard && arg.getType() != origType) { -- arg = sliceHighSideToType(tempRewriter, op->getLoc(), arg, origType, -- operandSharding); -+ arg = sliceHighSideToType(op->getLoc(), arg, origType, operandSharding, -+ tempRewriter); - } - if (auto concatOp = dyn_cast(op)) { - int64_t concatDim = concatOp.getDimension(); -@@ -652,8 +652,8 @@ LogicalResult outlineInstruction(Operation* op, - Value res = clonedAllReduce.getResult(); - Type targetType = paddedResultTypes[0]; - if (res.getType() != targetType) { -- res = padHighSideToType(tempRewriter, op->getLoc(), res, targetType, -- shardingInfo.outShardings[0]); -+ res = padHighSideToType(op->getLoc(), res, targetType, -+ shardingInfo.outShardings[0], tempRewriter); - } - func::ReturnOp::create(tempRewriter, op->getLoc(), ValueRange{res}); - return success(); -@@ -670,9 +670,9 @@ LogicalResult outlineInstruction(Operation* op, - } - Type targetType = paddedResultTypes[i]; - if (res.getType() != targetType) { -- returnValues.push_back(padHighSideToType(tempRewriter, op->getLoc(), res, -- targetType, -- shardingInfo.outShardings[i])); -+ returnValues.push_back(padHighSideToType(op->getLoc(), res, targetType, -+ shardingInfo.outShardings[i], -+ tempRewriter)); - } else { - returnValues.push_back(res); - } -@@ -734,8 +734,8 @@ void sliceIndivisibleResults(Operation* op, ManualComputationOp manualCompOp, - Value origRes = op->getResult(j); - Value manualRes = manualCompOp.getResult(j); - Type origResType = origRes.getType(); -- Value slicedRes = sliceHighSideToType(rewriter, op->getLoc(), manualRes, -- origResType, outShardings[j]); -+ Value slicedRes = sliceHighSideToType(op->getLoc(), manualRes, origResType, -+ outShardings[j], rewriter); - rewriter.replaceAllUsesWith(origRes, slicedRes); - } - } -diff --git a/shardy/dialect/sdy/transforms/export/resolve_permutation_factors.cc b/shardy/dialect/sdy/transforms/export/resolve_permutation_factors.cc -index 107fce53..4e1dac3e 100644 ---- a/shardy/dialect/sdy/transforms/export/resolve_permutation_factors.cc -+++ b/shardy/dialect/sdy/transforms/export/resolve_permutation_factors.cc -@@ -142,7 +142,7 @@ struct DataExchangeInfo { - // The result shape after padding the input for sharding divisibility. - SmallVector divisibleInputShape; - // Contains the axes sharding all sharded dimensions (i.e. dim size > 1). -- llvm::SmallDenseSet manualAxes; -+ llvm::SmallDenseSet manualAxes; - // dimExchanges[i] contains the exchange info for the i-th sharded - // dimension of the sharded tensor operand. - SmallVector dimExchanges; -@@ -256,7 +256,7 @@ int64_t getMeshAxisIndex(ArrayRef axes, StringRef name) { - } - - Value createDeviceId(Location loc, ResolutionState& state) { -- return getDeviceId(state.replicaCount, state.partitionCount, loc, -+ return getDeviceId(loc, state.replicaCount, state.partitionCount, - state.rewriter); - } - -@@ -427,7 +427,7 @@ DeviceOffsetInfo getDeviceOffsetInfo(Location loc, Value partitionId, - stablehlo::AddOp::create(rewriter, loc, offsetInPartition, baseOffsetVal); - - auto zeroConst = -- createZeroConstant(rewriter, loc, RankedTensorType::get({}, i32Ty)); -+ createZeroConstant(loc, RankedTensorType::get({}, i32Ty), rewriter); - - dilatedOffset = - stablehlo::MaxOp::create(rewriter, loc, dilatedOffset, zeroConst); -@@ -554,27 +554,6 @@ MeshOp getOrCreateMesh(Location loc, ModuleOp module, MeshAttr meshAttr, - return emplaceResult.first->second; - } - --// Returns a TensorShardingAttr where all axes in 'axes' are removed. --TensorShardingAttr removeAxesFromSharding( -- TensorShardingAttr sharding, const llvm::SmallDenseSet& axes) { -- SmallVector newDimShardings; -- for (DimensionShardingAttr dimSharding : sharding.getDimShardings()) { -- SmallVector newAxes; -- for (AxisRefAttr axis : dimSharding.getAxes()) { -- if (!axes.contains(axis.getName())) { -- newAxes.push_back(axis); -- } -- } -- newDimShardings.push_back(DimensionShardingAttr::get( -- sharding.getContext(), newAxes, dimSharding.getIsClosed(), -- dimSharding.getPriority())); -- } -- return TensorShardingAttr::get(sharding.getContext(), sharding.getMeshOrRef(), -- newDimShardings, sharding.getReplicatedAxes(), -- sharding.getUnreducedAxes(), -- sharding.getReductionOp()); --} -- - // Returns a TensorShardingAttr where all specific axis references in - // axesToRemove are removed. - TensorShardingAttr removeAxisRefsFromSharding( -@@ -598,6 +577,52 @@ TensorShardingAttr removeAxisRefsFromSharding( - sharding.getReductionOp()); - } - -+// Returns a TensorShardingAttr where all axes overlapping `axes` are removed. -+TensorShardingAttr removeAxesFromSharding( -+ TensorShardingAttr sharding, const llvm::SmallDenseSet& axes, -+ MeshAttr mesh) { -+ SmallVector axesToRemove(axes.begin(), axes.end()); -+ return removeAxisRefsFromSharding(sharding, axesToRemove, mesh); -+} -+ -+int64_t getManualShardCount( -+ DimensionShardingAttr dimSharding, MeshAttr mesh, -+ const llvm::SmallDenseSet& manualAxes) { -+ int64_t count = 1; -+ for (AxisRefAttr axis : dimSharding.getAxes()) { -+ for (AxisRefAttr manualAxis : manualAxes) { -+ if (std::optional overlap = axis.getOverlap(manualAxis)) { -+ count *= overlap->getSize(mesh); -+ } -+ } -+ } -+ return count; -+} -+ -+SmallVector getManualAxesAttrs( -+ MeshAttr mesh, const llvm::SmallDenseSet& manualAxes) { -+ SmallVector manualAxesAttrs(manualAxes.begin(), -+ manualAxes.end()); -+ sortAndMergeAxes(manualAxesAttrs, mesh); -+ return manualAxesAttrs; -+} -+ -+SmallVector getManualAxesInDim( -+ DimensionShardingAttr dimSharding, MeshAttr mesh, -+ const llvm::SmallDenseSet& manualAxes) { -+ SmallVector sortedManualAxes = -+ getManualAxesAttrs(mesh, manualAxes); -+ SmallVector manualAxesInDim; -+ for (AxisRefAttr axisRef : dimSharding.getAxes()) { -+ for (AxisRefAttr manualAxis : sortedManualAxes) { -+ if (std::optional overlap = axisRef.getOverlap(manualAxis)) { -+ manualAxesInDim.push_back(*overlap); -+ } -+ } -+ } -+ return manualAxesInDim; -+} -+ - // ----------------------------------------------------------------------------- - // Pad and slice helpers for dimension alignment. - // ----------------------------------------------------------------------------- -@@ -607,25 +632,25 @@ inline int64_t getPaddedDimSize(int64_t dimSize, int64_t shardCount) { - } - - // Creates a zero scalar constant of element type `elemType`. --inline Value createZeroScalarConstant(Location loc, IRRewriter& rewriter, -- Type elemType) { -- return createZeroConstant(rewriter, loc, RankedTensorType::get({}, elemType)); -+inline Value createZeroScalarConstant(Location loc, Type elemType, -+ IRRewriter& rewriter) { -+ return createZeroConstant(loc, RankedTensorType::get({}, elemType), rewriter); - } - - // Pads the high side of `operand` with `paddingValue` (default zero if null) - // to match the target `paddedShape`. When used to align dimensions for the - // input sharding's own divisibility, the created pad op is a - // communication-free operation. --Value padHighSideToShape(Location loc, IRRewriter& rewriter, Value operand, -+Value padHighSideToShape(Location loc, Value operand, - ArrayRef paddedShape, -- TensorShardingAttr sharding, -+ TensorShardingAttr sharding, IRRewriter& rewriter, - Value paddingValue = nullptr) { - auto type = cast(operand.getType()); - if (paddedShape == type.getShape()) { - return operand; - } - auto paddedType = RankedTensorType::get(paddedShape, type.getElementType()); -- return padHighSideToType(rewriter, loc, operand, paddedType, sharding, -+ return padHighSideToType(loc, operand, paddedType, sharding, rewriter, - paddingValue); - } - -@@ -687,21 +712,17 @@ SmallVector getDataShiftSourceTargetPairs( - // the partitioning space using HALO exchange and support multi-hop shifts. - Value rightShiftData(Location loc, Value input, int64_t dim, int64_t totalShift, - TensorShardingAttr sharding, MeshAttr mesh, -- const llvm::SmallDenseSet& manualAxes, -+ const llvm::SmallDenseSet& manualAxes, - ResolutionState& state) { - auto inputType = cast(input.getType()); - int64_t shardSize = inputType.getDimSize(dim); - int64_t rank = inputType.getRank(); - TensorShardingAttr localSharding = -- removeAxesFromSharding(sharding, manualAxes); -+ removeAxesFromSharding(sharding, manualAxes, mesh); - - // Parse manual axes for the sharded dimension. -- SmallVector manualAxesInDim; -- for (auto axisRef : sharding.getDimShardings()[dim].getAxes()) { -- if (manualAxes.contains(axisRef.getName())) { -- manualAxesInDim.push_back(axisRef); -- } -- } -+ SmallVector manualAxesInDim = -+ getManualAxesInDim(sharding.getDimShardings()[dim], mesh, manualAxes); - - int64_t hops = totalShift / shardSize; - int64_t fraction = totalShift % shardSize; -@@ -714,8 +735,8 @@ Value rightShiftData(Location loc, Value input, int64_t dim, int64_t totalShift, - SmallVector pairs = - getDataShiftSourceTargetPairs(mesh, manualAxesInDim, deviceOffset); - if (pairs.empty()) { -- Value zeroConst = createZeroScalarConstant(loc, state.rewriter, -- inputType.getElementType()); -+ Value zeroConst = createZeroScalarConstant( -+ loc, inputType.getElementType(), state.rewriter); - auto bcast = stablehlo::BroadcastInDimOp::create( - state.rewriter, loc, cast(piece.getType()), - zeroConst, state.rewriter.getDenseI64ArrayAttr({})); -@@ -777,7 +798,7 @@ Value rightShiftData(Location loc, Value input, int64_t dim, int64_t totalShift, - // moves the padding in the operand of a reverse op from high edge to low edge. - Value haloRightShiftData(Location loc, Value input, RankedTensorType origType, - TensorShardingAttr sharding, MeshAttr mesh, -- const llvm::SmallDenseSet& manualAxes, -+ const llvm::SmallDenseSet& manualAxes, - ArrayRef paddedShape, - ArrayRef dimsToShift, - ArrayRef shiftAmounts, -@@ -787,22 +808,13 @@ Value haloRightShiftData(Location loc, Value input, RankedTensorType origType, - localShape.reserve(sharding.getDimShardings().size()); - auto inputType = cast(input.getType()); - for (auto [dim, dimSharding] : llvm::enumerate(sharding.getDimShardings())) { -- int64_t manualFactor = 1; -- for (auto axis : dimSharding.getAxes()) { -- if (manualAxes.contains(axis.getName())) { -- manualFactor *= axis.getSize(mesh); -- } -- } -+ int64_t manualFactor = getManualShardCount(dimSharding, mesh, manualAxes); - localShape.push_back(paddedShape[dim] / manualFactor); - } - - // Get the manual axes for the manual computation. -- SmallVector manualAxesAttrs; -- for (auto axis : mesh.getAxes()) { -- if (manualAxes.contains(axis.getName())) { -- manualAxesAttrs.push_back(state.rewriter.getStringAttr(axis.getName())); -- } -- } -+ SmallVector manualAxesAttrs = -+ getManualAxesAttrs(mesh, manualAxes); - - // Build the manual computation. - auto manualComp = -@@ -826,18 +838,6 @@ Value haloRightShiftData(Location loc, Value input, RankedTensorType origType, - return manualComp.getResult(0); - } - --SmallVector getManualAxesAttrs( -- MeshAttr mesh, const llvm::SmallDenseSet& manualAxes, -- Builder& builder) { -- SmallVector manualAxesAttrs; -- for (auto axis : mesh.getAxes()) { -- if (manualAxes.contains(axis.getName())) { -- manualAxesAttrs.push_back(builder.getStringAttr(axis.getName())); -- } -- } -- return manualAxesAttrs; --} -- - // Returns the given input if sourceTargetPairs is empty, otherwise returns a - // collective-permute operation with the given sourceTargetPairs. - Value mayCollectivePermute(Location loc, Value input, -@@ -869,20 +869,16 @@ Value assembleHaloExchangeBuffer( - Location loc, Value input, int64_t dim, int64_t edgePadSize, - Value paddingValue, int64_t leftHops, int64_t rightHops, - TensorShardingAttr sharding, MeshAttr mesh, -- const llvm::SmallDenseSet& manualAxes, -+ const llvm::SmallDenseSet& manualAxes, - TensorShardingAttr localSharding, ResolutionState& state, - int64_t interiorPad = 0, const DimExchangeInfo* dimExchange = nullptr, - Value idInPartitionGroup = nullptr) { -- SmallVector manualAxesInDim; - DimensionShardingAttr activeDimSharded = - (dimExchange && dimExchange->reshapeInfo) - ? dimExchange->reshapeInfo->getTargetDimSharding() - : sharding.getDimShardings()[dim]; -- for (auto axisRef : activeDimSharded.getAxes()) { -- if (manualAxes.contains(axisRef.getName())) { -- manualAxesInDim.push_back(axisRef); -- } -- } -+ SmallVector manualAxesInDim = -+ getManualAxesInDim(activeDimSharded, mesh, manualAxes); - - SmallVector concatSegments; - concatSegments.reserve(leftHops + 1 + rightHops); -@@ -996,7 +992,7 @@ Value dynamicSliceHaloExchangeBuffer(Location loc, Value haloBuffer, - Value exchangeDimWithDynamicOffset( - Location loc, Value operand, Value paddingValue, int64_t dim, - const DimExchangeInfo& dimExchange, TensorShardingAttr sharding, -- MeshAttr mesh, const llvm::SmallDenseSet& manualAxes, -+ MeshAttr mesh, const llvm::SmallDenseSet& manualAxes, - TensorShardingAttr localSharding, int64_t padLow, int64_t stride, - int64_t sOut, int64_t baseDilation, bool dilatedHaloBuffer, - ResolutionState& state) { -@@ -1056,8 +1052,8 @@ Value exchangeDimWithDynamicOffset( - idInPartitionGroup, state.rewriter.getDenseI64ArrayAttr({1})); - offset = stablehlo::ReshapeOp::create( - state.rewriter, loc, RankedTensorType::get({}, i32Ty), sliceOp); -- zeroConst = createZeroConstant(state.rewriter, loc, -- RankedTensorType::get({}, i32Ty)); -+ zeroConst = createZeroConstant(loc, RankedTensorType::get({}, i32Ty), -+ state.rewriter); - } else { - DeviceOffsetInfo offsetInfo = - getDeviceOffsetInfo(loc, idInPartitionGroup, diffSize, -@@ -1078,12 +1074,12 @@ Value exchangeDimWithDynamicOffset( - Value exchangeDataWithDynamicOffset( - Location loc, Value operand, Value paddingValue, - TensorShardingAttr sharding, MeshAttr mesh, -- const llvm::SmallDenseSet& manualAxes, -+ const llvm::SmallDenseSet& manualAxes, - ArrayRef dimExchanges, ArrayRef edgePaddingLow, - ArrayRef windowStrides, bool dilatedHaloBuffer, - ResolutionState& state, ArrayRef baseDilations = {}) { - TensorShardingAttr localSharding = -- removeAxesFromSharding(sharding, manualAxes); -+ removeAxesFromSharding(sharding, manualAxes, mesh); - int64_t dimExchangeIdx = -1; - for (int64_t dim = 0; dim < sharding.getDimShardings().size(); ++dim) { - if (sharding.getDimShardings()[dim].getShardedSize(mesh) <= 1) { -@@ -1106,17 +1102,6 @@ Value exchangeDataWithDynamicOffset( - return operand; - } - --int64_t getManualShardCount(DimensionShardingAttr dimSharding, MeshAttr mesh, -- const llvm::SmallDenseSet& manualAxes) { -- int64_t count = 1; -- for (AxisRefAttr axis : dimSharding.getAxes()) { -- if (manualAxes.contains(axis.getName())) { -- count *= axis.getSize(mesh); -- } -- } -- return count; --} -- - // Performs HALO exchange over the sharded input inside a newly created - // `sdy.manual_computation` block. - // -@@ -1128,14 +1113,14 @@ Value haloDataExchange( - Location loc, Value divisibleInput, ArrayRef divisibleInputShape, - RankedTensorType origInputType, RankedTensorType origOutputType, - TensorShardingAttr inSharding, TensorShardingAttr outSharding, -- MeshAttr mesh, const llvm::SmallDenseSet& manualAxes, -+ MeshAttr mesh, const llvm::SmallDenseSet& manualAxes, - Value paddingValue, ArrayRef edgePaddingLow, - ArrayRef dimExchanges, const PostProcessFn& postProcess, - ResolutionState& state, ArrayRef windowStrides = {}, - ArrayRef baseDilations = {}) { - // Extract manual axes while preserving their original order in the mesh. -- SmallVector manualAxesAttrs = -- getManualAxesAttrs(mesh, manualAxes, state.rewriter); -+ SmallVector manualAxesAttrs = -+ getManualAxesAttrs(mesh, manualAxes); - - // Compute the local shape for input in the manual block based on - // divisibleInputShape. -@@ -1258,7 +1243,7 @@ std::optional prepareShardedDimForExchange( - return std::nullopt; - } - for (AxisRefAttr axis : dimSharding.getAxes()) { -- info.manualAxes.insert(axis.getName()); -+ info.manualAxes.insert(axis); - } - info.divisibleInputShape[dim] = - getPaddedDimSize(info.divisibleInputShape[dim], shardCount); -@@ -1333,24 +1318,20 @@ DataExchangeInfo getPadDataExchangeInfo(stablehlo::PadOp padOp, - Value handleReplicatedPadDims( - Location loc, Value exchangedLocal, Value localPaddingValue, - RankedTensorType paddedGlobalType, TensorShardingAttr sharding, -- MeshAttr mesh, const llvm::SmallDenseSet& manualAxes, -+ MeshAttr mesh, const llvm::SmallDenseSet& manualAxes, - ArrayRef edgePaddingLow, ArrayRef edgePaddingHigh, - ArrayRef interiorPadding, ResolutionState& state) { - auto exchangedLocalType = cast(exchangedLocal.getType()); - int64_t rank = exchangedLocalType.getRank(); - TensorShardingAttr localSharding = -- removeAxesFromSharding(sharding, manualAxes); -+ removeAxesFromSharding(sharding, manualAxes, mesh); - - SmallVector localResultShape(rank, 0), localLow(rank, 0), - localHigh(rank, 0), localInterior(rank, 0); - bool needsLocalPad = false; - for (int64_t i = 0; i < rank; ++i) { -- int64_t manualFactor = 1; -- for (auto axis : sharding.getDimShardings()[i].getAxes()) { -- if (manualAxes.contains(axis.getName())) { -- manualFactor *= axis.getSize(mesh); -- } -- } -+ int64_t manualFactor = -+ getManualShardCount(sharding.getDimShardings()[i], mesh, manualAxes); - int64_t sOutLocal = paddedGlobalType.getDimSize(i) / manualFactor; - localResultShape[i] = sOutLocal; - -@@ -1430,8 +1411,8 @@ LogicalResult handlePadOp(stablehlo::PadOp padOp, ResolutionState& state) { - rewriter.setInsertionPoint(padOp); - // Align logical shard sizes via high-side padding (Comm-free). - Value divisibleInput = -- padHighSideToShape(loc, rewriter, origInput, info.divisibleInputShape, -- inSharding, padOp.getPaddingValue()); -+ padHighSideToShape(loc, origInput, info.divisibleInputShape, inSharding, -+ rewriter, padOp.getPaddingValue()); - - SmallVector baseDilations = llvm::map_to_vector( - padOp.getInteriorPadding(), [](int64_t val) { return val + 1; }); -@@ -1455,7 +1436,7 @@ LogicalResult handlePadOp(stablehlo::PadOp padOp, ResolutionState& state) { - - // Trim final tensor shape to match expected output shape. - result = -- sliceHighSideToType(rewriter, loc, result, padOp.getType(), inSharding); -+ sliceHighSideToType(loc, result, padOp.getType(), inSharding, rewriter); - - rewriter.replaceOp(padOp, result); - return success(); -@@ -1472,7 +1453,7 @@ struct ReshapeHaloInfo { - struct ReshapeShardingAnalysis { - bool isEquivalent = false; - bool isCommunicationFree = false; -- llvm::SmallDenseSet nonDivisibleAxes; -+ llvm::SmallDenseSet nonDivisibleAxes; - SmallVector paddedInputShape; - SmallVector paddedOutputShape; - SmallVector reshapeGroups; -@@ -1599,16 +1580,11 @@ ReshapeShardingAnalysis ReshapeShardingAnalysis::run( - a.paddedOutputShape); - - if (groupHasIndivisibility) { -- auto collectAxes = [&](int64_t startDim, int64_t lastDim, -- TensorShardingAttr sharding) { -- for (int64_t d = startDim; d < lastDim; ++d) { -- for (AxisRefAttr axis : sharding.getDimShardings()[d].getAxes()) { -- a.nonDivisibleAxes.insert(axis.getName()); -- } -+ for (int64_t d = g.inStartDim; d < g.inLastDim; ++d) { -+ for (AxisRefAttr axis : inSharding.getDimShardings()[d].getAxes()) { -+ a.nonDivisibleAxes.insert(axis); - } -- }; -- collectAxes(g.inStartDim, g.inLastDim, inSharding); -- collectAxes(g.outStartDim, g.outLastDim, outSharding); -+ } - } - } - -@@ -1722,7 +1698,7 @@ DataExchangeInfo getReshapeDataExchangeInfo( - TensorShardingAttr outSharding, MeshAttr mesh, - ArrayRef divisibleInputShape, ArrayRef paddedOutputShape, - const ReshapeHaloInfo& haloInfo, -- const llvm::SmallDenseSet& nonDivisibleAxes) { -+ const llvm::SmallDenseSet& nonDivisibleAxes) { - DataExchangeInfo info; - info.manualAxes = nonDivisibleAxes; - auto origInputType = reshapeOp.getOperand().getType(); -@@ -1814,7 +1790,7 @@ LogicalResult handleReshapeHaloExchange( - ArrayRef reshapeGroups, - ArrayRef divisibleInputShape, ArrayRef paddedOutputShape, - ReshapeHaloInfo& haloInfo, -- const llvm::SmallDenseSet& nonDivisibleAxes, -+ const llvm::SmallDenseSet& nonDivisibleAxes, - ResolutionState& state) { - SDY_CHECK(!haloInfo.activeGroups.empty()); - -@@ -1832,12 +1808,12 @@ LogicalResult handleReshapeHaloExchange( - IRRewriter& rewriter = state.rewriter; - rewriter.setInsertionPoint(reshapeOp); - Value paddingValue = -- createZeroScalarConstant(loc, rewriter, resultType.getElementType()); -+ createZeroScalarConstant(loc, resultType.getElementType(), rewriter); - Value operand = reshapeOp.getOperand(); - auto origInputType = cast(operand.getType()); - SmallVector edgePaddingLow(origInputType.getRank(), 0); -- Value paddedInput = padHighSideToShape(loc, rewriter, operand, -- divisibleInputShape, inSharding); -+ Value paddedInput = padHighSideToShape(loc, operand, divisibleInputShape, -+ inSharding, rewriter); - - auto postProcess = [&](Value exchangedLocal, Value localPaddingValue, - RankedTensorType paddedGlobalResultType, -@@ -1853,7 +1829,7 @@ LogicalResult handleReshapeHaloExchange( - auto localOutputType = - RankedTensorType::get(localOutputShape, resultType.getElementType()); - TensorShardingAttr localOutSharding = -- removeAxesFromSharding(outSharding, info.manualAxes); -+ removeAxesFromSharding(outSharding, info.manualAxes, mesh); - auto localReshape = stablehlo::ReshapeOp::create( - state.rewriter, loc, localOutputType, exchangedLocal); - setSharding(localReshape.getResult(), localOutSharding); -@@ -1864,7 +1840,7 @@ LogicalResult handleReshapeHaloExchange( - loc, paddedInput, divisibleInputShape, origInputType, resultType, - inSharding, outSharding, mesh, info.manualAxes, paddingValue, - edgePaddingLow, info.dimExchanges, postProcess, state); -- result = sliceHighSideToType(rewriter, loc, result, resultType, outSharding); -+ result = sliceHighSideToType(loc, result, resultType, outSharding, rewriter); - rewriter.replaceOp(reshapeOp, result); - return success(); - } -@@ -1937,7 +1913,7 @@ struct ReverseOpInfo { - llvm::SmallVector axesToReverse; - // Axes used to shard dimensions being reversed and also involve in - // indivisible dimensions. -- llvm::SmallDenseSet manualAxes; -+ llvm::SmallDenseSet manualAxes; - // The new shape of the input tensor after all indivisible dimensions are - // padded. - SmallVector paddedShape; -@@ -1965,7 +1941,7 @@ std::optional getReverseOpInfo(stablehlo::ReverseOp reverseOp, - - if (type.getDimSize(dim) % shardCount != 0) { - for (AxisRefAttr axis : dimSharding.getAxes()) { -- info.manualAxes.insert(axis.getName()); -+ info.manualAxes.insert(axis); - } - info.paddedShape[dim] = - getPaddedDimSize(type.getDimSize(dim), shardCount); -@@ -2014,8 +1990,8 @@ LogicalResult handleReverseOp(stablehlo::ReverseOp reverseOp, - // Pad indivisible dimensions on the high edge. - if (!info->manualAxes.empty()) { - state.rewriter.setInsertionPoint(reverseOp); -- input = padHighSideToShape(loc, state.rewriter, operand, info->paddedShape, -- inSharding); -+ input = padHighSideToShape(loc, operand, info->paddedShape, inSharding, -+ state.rewriter); - - SmallVector dimsToShift, shiftAmounts; - int64_t rank = origType.getRank(); -@@ -2061,8 +2037,8 @@ LogicalResult handleReverseOp(stablehlo::ReverseOp reverseOp, - - // Trim off the padding. - if (!info->manualAxes.empty()) { -- Value slicedResult = sliceHighSideToType( -- state.rewriter, loc, reversedResult, origType, inSharding); -+ Value slicedResult = sliceHighSideToType(loc, reversedResult, origType, -+ inSharding, state.rewriter); - state.rewriter.replaceAllUsesExcept(reverseOp.getResult(), slicedResult, - reversedResult.getDefiningOp()); - } else { -@@ -2115,13 +2091,13 @@ DataExchangeInfo getSliceDataExchangeInfo(stablehlo::SliceOp sliceOp, - // Implements the slicing operation on replicated dimensions. - Value handleReplicatedSliceDims( - Location loc, Value exchangedLocal, TensorShardingAttr sharding, -- MeshAttr mesh, const llvm::SmallDenseSet& manualAxes, -+ MeshAttr mesh, const llvm::SmallDenseSet& manualAxes, - ArrayRef startIndices, ArrayRef limitIndices, - ArrayRef strides, ResolutionState& state) { - auto exchangedLocalType = cast(exchangedLocal.getType()); - int64_t rank = exchangedLocalType.getRank(); - TensorShardingAttr localSharding = -- removeAxesFromSharding(sharding, manualAxes); -+ removeAxesFromSharding(sharding, manualAxes, mesh); - - SmallVector localSliceStarts(rank, 0); - SmallVector localSliceLimits = -@@ -2208,11 +2184,11 @@ LogicalResult handleSliceOp(stablehlo::SliceOp sliceOp, - rewriter.setInsertionPoint(sliceOp); - // Align logical shard sizes via high-side padding. - Value divisibleInput = -- padHighSideToShape(loc, rewriter, origInput, info.divisibleInputShape, -- inSharding, /*paddingValue=*/nullptr); -+ padHighSideToShape(loc, origInput, info.divisibleInputShape, inSharding, -+ rewriter, /*paddingValue=*/nullptr); - - Value paddingValue = -- createZeroScalarConstant(loc, rewriter, origInputType.getElementType()); -+ createZeroScalarConstant(loc, origInputType.getElementType(), rewriter); - - int64_t rank = origInputType.getRank(); - SmallVector edgePaddingLow, edgePaddingHigh, interiorPadding; -@@ -2245,7 +2221,7 @@ LogicalResult handleSliceOp(stablehlo::SliceOp sliceOp, - - // Trim final tensor shape to match expected output shape. - result = -- sliceHighSideToType(rewriter, loc, result, sliceOp.getType(), inSharding); -+ sliceHighSideToType(loc, result, sliceOp.getType(), inSharding, rewriter); - - rewriter.replaceOp(sliceOp, result); - return success(); -diff --git a/shardy/dialect/sdy/transforms/export/resolve_single_device_sharding.cc b/shardy/dialect/sdy/transforms/export/resolve_single_device_sharding.cc -index cdce7d73..76727deb 100644 ---- a/shardy/dialect/sdy/transforms/export/resolve_single_device_sharding.cc -+++ b/shardy/dialect/sdy/transforms/export/resolve_single_device_sharding.cc -@@ -101,7 +101,7 @@ Value createDeviceGuard(Location loc, int64_t targetDeviceId, - int64_t replicaCount, int64_t partitionCount, - IRRewriter& rewriter) { - Value currentDeviceId = -- getDeviceId(replicaCount, partitionCount, loc, rewriter); -+ getDeviceId(loc, replicaCount, partitionCount, rewriter); - auto scalarUI32Type = RankedTensorType::get( - {}, rewriter.getIntegerType(32, /*isSigned=*/false)); - SDY_CHECK_GE(targetDeviceId, 0); -@@ -157,7 +157,7 @@ void resolveSingleDeviceOp(Operation* op, const SymbolTable& symbolTable, - isTargetDevice = createDeviceGuard(loc, *targetDeviceId, replicaCount, - partitionCount, rewriter); - } -- Value zeroConst = createZeroConstant(rewriter, loc, rankedType); -+ Value zeroConst = createZeroConstant(loc, rankedType, rewriter); - maskedResults.push_back(stablehlo::SelectOp::create( - rewriter, loc, rankedType, isTargetDevice, result, zeroConst)); - } -@@ -186,7 +186,7 @@ void resolveSingleDeviceOp(Operation* op, const SymbolTable& symbolTable, - SmallVector falseZeros; - for (Type type : resultTypes) { - if (auto rankedType = dyn_cast(type)) { -- Value zeroConst = createZeroConstant(rewriter, loc, rankedType); -+ Value zeroConst = createZeroConstant(loc, rankedType, rewriter); - falseZeros.push_back(zeroConst); - } else if (isa(type)) { - Value tokenConst = stablehlo::CreateTokenOp::create(rewriter, loc); -diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_manual_computation.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_manual_computation.mlir -index 7d7421c3..ae605980 100644 ---- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_manual_computation.mlir -+++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_manual_computation.mlir -@@ -206,3 +206,20 @@ func.func @stablehlo_reduce_scatter(%arg0 : tensor<16x8xf32> {sdy.sharding = #sd - // CHECK-NEXT: return %[[RES]] : tensor<4x8xf32> - return %0 : tensor<16x8xf32> - } -+ -+// CHECK-LABEL: func @sub_axis_manual_and_free_merged -+// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4, [{"x"}, {}]>}) -+// CHECK-SAME: -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4, [{"x"}, {}]>}) { -+func.func @sub_axis_manual_and_free_merged(%arg0: tensor<16x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4, [{"x"}, {}]>}) -+ -> (tensor<16x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4, [{"x"}, {}]>}) { -+ // CHECK-NEXT: %[[TANH:.*]] = stablehlo.tanh %[[ARG0]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{"x":(2)2}, {}]>]>} : tensor<4x8xf32> -+ %0 = sdy.manual_computation(%arg0) -+ in_shardings=[<@mesh_4, [{"x"}, {}]>] -+ out_shardings=[<@mesh_4, [{"x"}, {}]>] -+ manual_axes={"x":(1)2} (%arg1: tensor<8x8xf32>) { -+ %1 = stablehlo.tanh %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{"x":(2)2}, {}]>]>} : tensor<8x8xf32> -+ sdy.return %1 : tensor<8x8xf32> -+ } : (tensor<16x8xf32>) -> tensor<16x8xf32> -+ // CHECK-NEXT: return %[[TANH]] : tensor<4x8xf32> -+ return %0 : tensor<16x8xf32> -+} -diff --git a/shardy/dialect/sdy/transforms/export/test/export_pipeline.mlir b/shardy/dialect/sdy/transforms/export/test/export_pipeline.mlir -index d72ae963..6ce4f38b 100644 ---- a/shardy/dialect/sdy/transforms/export/test/export_pipeline.mlir -+++ b/shardy/dialect/sdy/transforms/export/test/export_pipeline.mlir -@@ -13,7 +13,7 @@ func.func @main( - %arg3: tensor<48xf32>, %arg4: tensor<96xf32>, %arg5: tensor<192xf32>) - -> (tensor<4xf32>, tensor<12xf32>, tensor<24xf32>, - tensor<48xf32>, tensor<96xf32>, tensor<192xf32>) { -- // CHECK: sdy.manual_computation -+ // CHECK-NEXT: sdy.manual_computation(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5) - // CHECK-SAME: in_shardings=[, [{"a"}]>, , [{"a"}]>, - // CHECK-SAME: , [{"a", "b":(1)2}]>, , [{"a", "b"}]>, - // CHECK-SAME: , [{"a", "b", "c":(1)2}]>, , [{"a", "b", "c"}]>] -@@ -21,8 +21,6 @@ func.func @main( - // CHECK-SAME: , [{"a", "b":(1)2}]>, , [{"a", "b"}]>, - // CHECK-SAME: , [{"a", "b", "c":(1)2}]>, , [{"a", "b", "c"}]>] - // CHECK-SAME: manual_axes={"a"} -- // CHECK: %[[RESHARD:.*]] = sdy.reshard %{{.*}} , [{}]> : tensor<192xf32> -- // CHECK: return %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[RESHARD]] - %0:6 = sdy.manual_computation(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5) - in_shardings=[<@mesh3d, [{"a", "b", "c"}]>, <@mesh3d, [{"a", "b", "c"}]>, - <@mesh3d, [{"a", "b", "c"}]>, <@mesh3d, [{"a", "b", "c"}]>, -@@ -74,8 +72,7 @@ func.func @main(%arg0: tensor<8x2xi32>, %arg1: tensor<4x2xi32>) -> tensor<12x2xi - // CHECK-NEXT: %2 = sdy.reshard %1#0 , [{}, {"a"}]> : tensor<8x2xi32> - // CHECK-NEXT: %3 = sdy.reshard %1#1 , [{}, {"a"}]> : tensor<4x2xi32> - // CHECK-NEXT: %4 = stablehlo.concatenate %2, %3, dim = 0 {sdy.sharding = #sdy.sharding_per_value<[, [{}, {"a"}]>]>} -- // CHECK-NEXT: %5 = sdy.reshard %4 , [{}, {}]> : tensor<12x2xi32> -- // CHECK-NEXT: return %5 : tensor<12x2xi32> -+ // CHECK-NEXT: return %4 : tensor<12x2xi32> - %0:2 = call @foo(%arg0, %arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>, <@mesh, [{}, {}]>]>} : (tensor<8x2xi32>, tensor<4x2xi32>) -> (tensor<8x2xi32>, tensor<4x2xi32>) - %1 = stablehlo.concatenate %0#0, %0#1, dim=0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"a"}]>]>} : (tensor<8x2xi32>, tensor<4x2xi32>) -> tensor<12x2xi32> - return %1 : tensor<12x2xi32> -@@ -112,8 +109,6 @@ func.func private @foo(%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh, - - // CHECK-LABEL: func private @bar_0(%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding, [{"b"}]>}) -> tensor<8xf32> { - func.func private @bar_0(%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}]>}) -> tensor<8xf32> { -- // CHECK-NEXT: %0 = sdy.reshard %arg0 , [{}]> : tensor<8xf32> -- // CHECK-NEXT: return %0 : tensor<8xf32> - return %arg0 : tensor<8xf32> - } - -@@ -142,15 +137,11 @@ func.func private @foo(%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh, - - // CHECK-LABEL: func private @bar_0(%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding, [{"b"}]>}) -> tensor<8xf32> { - func.func private @bar_0(%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}]>}) -> tensor<8xf32> { -- // CHECK-NEXT: %0 = sdy.reshard %arg0 , [{}]> : tensor<8xf32> -- // CHECK-NEXT: return %0 : tensor<8xf32> - return %arg0 : tensor<8xf32> - } - - // CHECK-LABEL: func private @bar(%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding, [{"c"}]>}) -> tensor<8xf32> { - func.func private @bar(%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"c"}]>}) -> tensor<8xf32> { -- // CHECK-NEXT: %0 = sdy.reshard %arg0 , [{}]> : tensor<8xf32> -- // CHECK-NEXT: return %0 : tensor<8xf32> - return %arg0 : tensor<8xf32> - } - -@@ -161,11 +152,8 @@ sdy.mesh @mesh = <["a"=2, "b"=2]> - func.func @propagation_barrier_removed_before_explicit_reshards( - %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}) - -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {"b"}]>}) { -- // CHECK-NEXT: %[[RESHARD_0:.*]] = sdy.reshard %arg0 , [{"a"}, {"b"}]> : tensor<8x8xf32> -- // CHECK-NEXT: %[[RESHARD_1:.*]] = sdy.reshard %arg0 , [{"a"}, {"b"}]> : tensor<8x8xf32> -- // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[RESHARD_0]], %[[RESHARD_1]] {sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {"b"}]>]>} : tensor<8x8xf32> -- // CHECK-NEXT: %[[RESHARD_2:.*]] = sdy.reshard %[[ADD]] , [{}, {"b"}]> : tensor<8x8xf32> -- // CHECK-NEXT: return %[[RESHARD_2]] : tensor<8x8xf32> -+ // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]>} : tensor<8x8xf32> -+ // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32> - %0 = sdy.propagation_barrier %arg0 allowed_direction=BACKWARD {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>} : tensor<8x8xf32> - %1 = stablehlo.add %0, %0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>} : tensor<8x8xf32> - return %1 : tensor<8x8xf32> -diff --git a/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir b/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir -index a6c01e12..8ff8f150 100644 ---- a/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir -+++ b/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir -@@ -221,7 +221,7 @@ func.func @pad_single_left_hop( - // REPL: %[[RES:.*]] = sdy.reshard %[[PAD]] <@mesh_a4, [{"a":(2)2}, {"b"}]> : tensor<7x8xi32> - - // HALO: %[[CST:.*]] = stablehlo.constant dense<0> : tensor -- // HALO: %[[MC:.*]] = sdy.manual_computation(%[[ARG0]], %[[CST]]) in_shardings=[<@mesh_a4, [{"a":(2)2}, {"b"}]>, <@mesh_a4, []>] out_shardings=[<@mesh_a4, [{"a":(2)2}, {"b"}]>] manual_axes={"a", "b"} (%[[ARG1:.*]]: tensor<2x4xi32>, %[[ARG2:.*]]: tensor) { -+ // HALO: %[[MC:.*]] = sdy.manual_computation(%[[ARG0]], %[[CST]]) in_shardings=[<@mesh_a4, [{"a":(2)2}, {"b"}]>, <@mesh_a4, []>] out_shardings=[<@mesh_a4, [{"a":(2)2}, {"b"}]>] manual_axes={"a":(2)2, "b"} (%[[ARG1:.*]]: tensor<2x4xi32>, %[[ARG2:.*]]: tensor) { - // HALO: %[[PART_ID:.*]] = stablehlo.partition_id : tensor - // HALO: %[[CONVERT:.*]] = stablehlo.convert %[[PART_ID]] : (tensor) -> tensor - // HALO: %[[C_DIV:.*]] = stablehlo.constant dense<2> : tensor -@@ -1294,7 +1294,7 @@ func.func @reverse_indivisible_multiple_axes( - // HALO-SAME: low = [0, 0], high = [0, 1] - - // HALO: %[[HALO_SHIFTED:.*]] = sdy.manual_computation(%[[PAD]]) -- // HALO-SAME: manual_axes={"a", "c"} -+ // HALO-SAME: manual_axes={"a", "c":(2)2} - // HALO-SAME: (%[[LOCAL_IN:.*]]: tensor<6x2xi32>) { - // HALO-NEXT: %[[CP:.*]] = "stablehlo.collective_permute"(%[[LOCAL_IN]]) - // HALO-SAME{LITERAL}: source_target_pairs = dense<[[0, 1], [1, 8], [2, 3], [3, 10], [4, 5], [5, 12], [6, 7], [7, 14], [8, 9], [10, 11], [12, 13], [14, 15]]> : tensor<12x2xi64> -@@ -1317,6 +1317,45 @@ func.func @reverse_indivisible_multiple_axes( - return %2 : tensor<6x7xi32> - } - -+// CHECK-LABEL: func @reverse_indivisible_passthrough_sub_axis -+// CHECK-SAME: (%[[ARG0:.*]]: tensor<6x8xi32> {sdy.sharding = #sdy.sharding<@mesh_abc, [{"c":(1)2}, {"a", "c":(2)2}]>}) -+// CHECK-SAME: -> tensor<6x7xi32> { -+func.func @reverse_indivisible_passthrough_sub_axis( -+ %arg0: tensor<6x8xi32> {sdy.sharding = #sdy.sharding<@mesh_abc, [{"c":(1)2}, {"a", "c":(2)2}]>}) -+ -> tensor<6x7xi32> { -+ -+ // CHECK: %[[SLICE:.*]] = stablehlo.slice %[[ARG0]] [0:6, 0:7] -+ %0 = stablehlo.slice %arg0 [0:6, 0:7] -+ {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_abc, [{"c":(1)2}, {"a", "c":(2)2}]>]>} -+ : (tensor<6x8xi32>) -> tensor<6x7xi32> -+ -+ // REPL: %[[REPL_IN:.*]] = sdy.reshard %[[SLICE]] <@mesh_abc, [{"c":(1)2}, {}]> -+ // REPL: %[[REPL_REV:.*]] = stablehlo.reverse %[[REPL_IN]], dims = [1] -+ // REPL-SAME: {sdy.sharding = #sdy.sharding_per_value<[<@mesh_abc, [{"c":(1)2}, {}]>]>} -+ // REPL: %[[RES:.*]] = sdy.reshard %[[REPL_REV]] <@mesh_abc, [{"c":(1)2}, {"a", "c":(2)2}]> -+ -+ // HALO: %[[HALO_SHIFTED:.*]] = sdy.manual_computation -+ // HALO-SAME: in_shardings=[<@mesh_abc, [{"c":(1)2}, {"a", "c":(2)2}]>] -+ // HALO-SAME: out_shardings=[<@mesh_abc, [{"c":(1)2}, {"a", "c":(2)2}]>] -+ // HALO-SAME: manual_axes={"a", "c":(2)2} -+ // HALO-SAME: (%[[LOCAL_IN:.*]]: tensor<6x2xi32>) { -+ // HALO-NEXT: %[[CP:.*]] = "stablehlo.collective_permute"(%[[LOCAL_IN]]) -+ // HALO-SAME: sdy.sharding = #sdy.sharding_per_value<[<@mesh_abc, [{"c":(1)2}, {}]>]> -+ // HALO-NEXT: %[[OWN_SLICE:.*]] = stablehlo.slice %[[CP]] [0:6, 1:2] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_abc, [{"c":(1)2}, {}]>]>} -+ // HALO-NEXT: %[[HALO_SLICE:.*]] = stablehlo.slice %[[LOCAL_IN]] [0:6, 0:1] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_abc, [{"c":(1)2}, {}]>]>} -+ // HALO-NEXT: %[[CONCAT:.*]] = stablehlo.concatenate %[[OWN_SLICE]], %[[HALO_SLICE]], dim = 1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_abc, [{"c":(1)2}, {}]>]>} -+ // HALO-NEXT: sdy.return %[[CONCAT]] -+ // HALO-NEXT: } : (tensor<6x8xi32>) -> tensor<6x8xi32> -+ %1 = stablehlo.reverse %0, dims = [1] -+ {sdy.sharding = #sdy.sharding_per_value<[<@mesh_abc, [{"c":(1)2}, {"a", "c":(2)2}]>]>} -+ : tensor<6x7xi32> -+ -+ // CHECK: %[[AG:.*]] = sdy.all_gather [{"c":(1)2}, {"a", "c":(2)2}] %[[RES:.*]] -+ %2 = sdy.all_gather [{"c":(1)2}, {"a", "c":(2)2}] %1 out_sharding=<@mesh_abc, [{}, {}]> : tensor<6x7xi32> -+ // CHECK: return %[[AG]] : tensor<6x7xi32> -+ return %2 : tensor<6x7xi32> -+} -+ - //===----------------------------------------------------------------------===// - // stablehlo.select_and_scatter tests - //===----------------------------------------------------------------------===// -diff --git a/shardy/dialect/sdy/transforms/export/test/verify_unreduced_axes.mlir b/shardy/dialect/sdy/transforms/export/test/verify_unreduced_axes.mlir -index 053b3a88..26eb3edb 100644 ---- a/shardy/dialect/sdy/transforms/export/test/verify_unreduced_axes.mlir -+++ b/shardy/dialect/sdy/transforms/export/test/verify_unreduced_axes.mlir -@@ -399,3 +399,15 @@ func.func @case_dropping_unreduced_into_unannotated_result(%idx: tensor {sd - }) : (tensor) -> tensor<8xf32> - return %0 : tensor<8xf32> - } -+ -+// ----- -+ -+sdy.mesh @mesh = <["x"=4, "y"=2]> -+ -+func.func @manual_computation_merged_unreduced_axis_with_manual_sub_axis(%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}], unreduced={"x"}>}) -> (tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}], unreduced={"x":(1)2}>}) { -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{}], unreduced={"x"}>] out_shardings=[<@mesh, [{}], unreduced={"x":(1)2}>] manual_axes={"x":(1)2} (%in0: tensor<8xf32>) { -+ %1 = sdy.reshard %in0 <@mesh, [{}]> : tensor<8xf32> -+ sdy.return %1 : tensor<8xf32> -+ } : (tensor<8xf32>) -> tensor<8xf32> -+ return %0 : tensor<8xf32> -+} -diff --git a/shardy/dialect/sdy/transforms/export/utils.cc b/shardy/dialect/sdy/transforms/export/utils.cc -index 54e5d6f3..1323f0b1 100644 ---- a/shardy/dialect/sdy/transforms/export/utils.cc -+++ b/shardy/dialect/sdy/transforms/export/utils.cc -@@ -284,7 +284,7 @@ bool usePartitionId(int64_t replicaCount, int64_t partitionCount) { - return replicaCount == 1; - } - --Value getDeviceId(int64_t replicaCount, int64_t partitionCount, Location loc, -+Value getDeviceId(Location loc, int64_t replicaCount, int64_t partitionCount, - OpBuilder& rewriter) { - if (usePartitionId(replicaCount, partitionCount)) { - return stablehlo::PartitionIdOp::create(rewriter, loc); -@@ -627,8 +627,8 @@ bool isCommunicationFreeReshape(stablehlo::ReshapeOp reshapeOp, - inputType, outputType, groups); - } - --Value sliceHighSideToType(OpBuilder& builder, Location loc, Value operand, -- Type targetType, TensorShardingAttr sharding) { -+Value sliceHighSideToType(Location loc, Value operand, Type targetType, -+ TensorShardingAttr sharding, OpBuilder& builder) { - Type currentType = operand.getType(); - if (currentType == targetType) { - return operand; -@@ -648,7 +648,7 @@ Value sliceHighSideToType(OpBuilder& builder, Location loc, Value operand, - return sliced; - } - --Attribute getZeroAttr(OpBuilder& builder, Type type) { -+Attribute getZeroAttr(Type type, OpBuilder& builder) { - if (Attribute attr = builder.getZeroAttr(type)) { - return attr; - } -@@ -674,14 +674,14 @@ Attribute getZeroAttr(OpBuilder& builder, Type type) { - return nullptr; - } - --Value createZeroConstant(OpBuilder& builder, Location loc, Type type) { -- Attribute zeroAttr = getZeroAttr(builder, type); -+Value createZeroConstant(Location loc, Type type, OpBuilder& builder) { -+ Attribute zeroAttr = getZeroAttr(type, builder); - SDY_CHECK(zeroAttr) << "Failed to create zero attribute"; - return stablehlo::ConstantOp::create(builder, loc, zeroAttr); - } - --Value padHighSideToType(OpBuilder& builder, Location loc, Value operand, -- Type targetType, TensorShardingAttr sharding, -+Value padHighSideToType(Location loc, Value operand, Type targetType, -+ TensorShardingAttr sharding, OpBuilder& builder, - Value paddingValue, bool allowSlicePeephole) { - Type currentType = operand.getType(); - if (currentType == targetType) { -@@ -710,7 +710,7 @@ Value padHighSideToType(OpBuilder& builder, Location loc, Value operand, - Value padVal = paddingValue; - if (!padVal) { - auto zeroType = RankedTensorType::get({}, targetRanked.getElementType()); -- padVal = createZeroConstant(builder, loc, zeroType); -+ padVal = createZeroConstant(loc, zeroType, builder); - } - Value padded = stablehlo::PadOp::create( - builder, loc, targetRanked, operand, padVal, -@@ -815,7 +815,7 @@ ElementsAttr padElementsAttr(ElementsAttr elementsAttr, - - OpBuilder builder(paddedType.getContext()); - auto zeroDenseAttr = -- dyn_cast_or_null(getZeroAttr(builder, paddedType)); -+ dyn_cast_or_null(getZeroAttr(paddedType, builder)); - SDY_CHECK(zeroDenseAttr) << "Failed to create zero attribute"; - - int64_t rank = origType.getRank(); -diff --git a/shardy/dialect/sdy/transforms/export/utils.h b/shardy/dialect/sdy/transforms/export/utils.h -index 5737504e..db2513b1 100644 ---- a/shardy/dialect/sdy/transforms/export/utils.h -+++ b/shardy/dialect/sdy/transforms/export/utils.h -@@ -117,7 +117,7 @@ bool usePartitionId(int64_t replicaCount, int64_t partitionCount); - // partition count. - // - // TODO(b/545097355): support replica_count > 1 && partition_count > 1. --Value getDeviceId(int64_t replicaCount, int64_t partitionCount, Location loc, -+Value getDeviceId(Location loc, int64_t replicaCount, int64_t partitionCount, - OpBuilder& rewriter); - - // Returns a channel handle with handle = nextChannelId++ and type = 1 -@@ -252,24 +252,23 @@ bool isCommunicationFreeReshape(stablehlo::ReshapeOp reshapeOp, - // Slices the high side of `operand` starting at zero with unit strides to match - // `targetType`. If `operand` already matches `targetType`, returns `operand`. - // Attaches `sharding` to the created slice op if non-null. --Value sliceHighSideToType(OpBuilder& builder, Location loc, Value operand, -- Type targetType, -- TensorShardingAttr sharding = nullptr); -+Value sliceHighSideToType(Location loc, Value operand, Type targetType, -+ TensorShardingAttr sharding, OpBuilder& builder); - - // Returns a zero attribute for any ShapedType or scalar Type, including real - // integer, real float, and all ComplexType elements (complex, - // complex, etc.). --Attribute getZeroAttr(OpBuilder& builder, Type type); -+Attribute getZeroAttr(Type type, OpBuilder& builder); - - // Creates a stablehlo.constant zero matching `type` (scalar or tensor, real or - // complex). --Value createZeroConstant(OpBuilder& builder, Location loc, Type type); -+Value createZeroConstant(Location loc, Type type, OpBuilder& builder); - - // Zero-pads the high side of `operand` to match `targetType`. If `operand` - // already matches `targetType`, returns `operand`. Attaches `sharding` to the - // created pad op if non-null. --Value padHighSideToType(OpBuilder& builder, Location loc, Value operand, -- Type targetType, TensorShardingAttr sharding = nullptr, -+Value padHighSideToType(Location loc, Value operand, Type targetType, -+ TensorShardingAttr sharding, OpBuilder& builder, - Value paddingValue = nullptr, - bool allowSlicePeephole = false); - -diff --git a/shardy/dialect/sdy/transforms/export/verify_unreduced_axes.cc b/shardy/dialect/sdy/transforms/export/verify_unreduced_axes.cc -index 191c4bf0..5399846f 100644 ---- a/shardy/dialect/sdy/transforms/export/verify_unreduced_axes.cc -+++ b/shardy/dialect/sdy/transforms/export/verify_unreduced_axes.cc -@@ -214,7 +214,8 @@ LogicalResult verifyTerminator(Operation* op, func::FuncOp funcOp) { - if (auto manualCompOp = dyn_cast(parentOp)) { - if (resultSharding) { - resultSharding = -- eraseManualAxes(resultSharding, manualCompOp.getManualAxes()); -+ eraseManualAxes(resultSharding, manualCompOp.getManualAxes(), -+ resultSharding.getMesh(manualCompOp)); - } - } - if (failed( -@@ -279,16 +280,22 @@ LogicalResult verifyNamedComputationOp( - } - - TensorShardingAttr retainOnlyManualUnreducedAxes( -- TensorShardingAttr sharding, ArrayRef manualAxes) { -+ TensorShardingAttr sharding, ArrayRef manualAxes, -+ Operation* op) { - if (!sharding || sharding.getUnreducedAxes().empty()) { - return TensorShardingAttr(); - } - SmallVector manualUnreduced; - for (AxisRefAttr axis : sharding.getUnreducedAxes()) { -- if (llvm::is_contained(manualAxes, axis.getName())) { -- manualUnreduced.push_back(axis); -+ for (AxisRefAttr manualAxis : manualAxes) { -+ if (std::optional overlap = axis.getOverlap(manualAxis)) { -+ manualUnreduced.push_back(*overlap); -+ } - } - } -+ if (MeshAttr mesh = sharding.getMesh(op)) { -+ sortAndMergeAxes(manualUnreduced, mesh); -+ } - return TensorShardingAttr::get(sharding.getContext(), sharding.getMeshOrRef(), - sharding.getDimShardings(), - sharding.getReplicatedAxes(), manualUnreduced, -@@ -297,7 +304,7 @@ TensorShardingAttr retainOnlyManualUnreducedAxes( - - LogicalResult verifyManualComputationOp( - sdy::ManualComputationOp manualComputationOp) { -- ArrayRef manualAxes = manualComputationOp.getManualAxes(); -+ ArrayRef manualAxes = manualComputationOp.getManualAxes(); - // Verify non-manual unreduced axes across the operand -> body argument - // boundary (non-manual outputs are verified by verifyTerminator). - for (int64_t i = 0; i < manualComputationOp.getNumOperands(); ++i) { -@@ -307,8 +314,10 @@ LogicalResult verifyManualComputationOp( - } - TensorShardingAttr operandSharding = getSharding(operand); - TensorShardingAttr operandNonManualSharding = -- operandSharding ? eraseManualAxes(operandSharding, manualAxes) -- : TensorShardingAttr(); -+ operandSharding -+ ? eraseManualAxes(operandSharding, manualAxes, -+ operandSharding.getMesh(manualComputationOp)) -+ : TensorShardingAttr(); - TensorShardingAttr argSharding = - getSharding(manualComputationOp.getBody().getArgument(i)); - if (failed(verifyUnreducedShardingEqual( -@@ -342,15 +351,15 @@ LogicalResult verifyManualComputationOp( - - for (Value operand : manualComputationOp.getOperands()) { - if (!isDefiningOpExplicitReshard(operand)) { -- recordSharding( -- retainOnlyManualUnreducedAxes(getSharding(operand), manualAxes), -- operandShardingsByOp); -+ recordSharding(retainOnlyManualUnreducedAxes( -+ getSharding(operand), manualAxes, manualComputationOp), -+ operandShardingsByOp); - } - } - for (Value result : manualComputationOp.getResults()) { -- recordSharding( -- retainOnlyManualUnreducedAxes(getSharding(result), manualAxes), -- resultShardingsByOp); -+ recordSharding(retainOnlyManualUnreducedAxes( -+ getSharding(result), manualAxes, manualComputationOp), -+ resultShardingsByOp); - } - - for (ReductionOp redOp : seenReductionOps) { -diff --git a/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc b/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc -index 6b368f91..cf118942 100644 ---- a/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc -+++ b/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc -@@ -93,14 +93,7 @@ addUnusedManualAxesToReplicatedAxes(ArrayRef shardings, - void addUnusedManualAxesToReplicatedAxes(ManualComputationOp op, MeshAttr mesh, - Attribute commonMeshOrRef, - const SymbolTable& symbolTable) { -- OpBuilder builder(op); -- -- SmallVector manualAxes; -- manualAxes.reserve(op.getManualAxes().size()); -- llvm::transform(op.getManualAxes(), std::back_inserter(manualAxes), -- [&](StringAttr axis) { -- return AxisRefAttr::get(builder.getContext(), axis); -- }); -+ ArrayRef manualAxes = op.getManualAxes(); - - if (std::optional> newShardings = - addUnusedManualAxesToReplicatedAxes( -@@ -119,13 +112,13 @@ void addUnusedManualAxesToReplicatedAxes(ManualComputationOp op, MeshAttr mesh, - - // Sorts the manual axes in mesh axis order. - void sortManualAxes(ManualComputationOp op, MeshAttr mesh) { -- ArrayRef manualAxes = op.getManualAxes(); -- auto meshComparator = mesh.getAxisNameComparator(); -+ ArrayRef manualAxes = op.getManualAxes(); -+ auto meshComparator = AxisRefAttr::getMeshComparator(mesh); - - if (llvm::is_sorted(manualAxes, meshComparator)) { - return; - } -- SmallVector sortedManualAxes = llvm::to_vector(manualAxes); -+ SmallVector sortedManualAxes = llvm::to_vector(manualAxes); - llvm::sort(sortedManualAxes, meshComparator); - op.setManualAxes(sortedManualAxes); - } -diff --git a/shardy/dialect/sdy/transforms/import/test/manual_axes_cleanup.mlir b/shardy/dialect/sdy/transforms/import/test/manual_axes_cleanup.mlir -index b110649c..2dfc2093 100644 ---- a/shardy/dialect/sdy/transforms/import/test/manual_axes_cleanup.mlir -+++ b/shardy/dialect/sdy/transforms/import/test/manual_axes_cleanup.mlir -@@ -154,6 +154,18 @@ func.func @sub_axes_in_out_shardings(%arg0: tensor<8xf32>) -> tensor<8xf32> { - return %0 : tensor<8xf32> - } - -+// CHECK-LABEL: @sub_axes_in_manual_axes -+func.func @sub_axes_in_manual_axes(%arg0: tensor<8xf32>) -> tensor<8xf32> { -+ // CHECK-NEXT: sdy.manual_computation(%arg0) -+ // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_x_8_y_8, [{"y":(1)2}], replicated={"x":(2)2, "y":(2)2}>] -+ // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_x_8_y_8, [{"x":(2)2}], replicated={"y":(1)4}>] -+ // CHECK-SAME{LITERAL}: manual_axes={"x":(2)2, "y":(1)4} (%arg1: tensor<4xf32>) { -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_x_8_y_8, [{"y":(1)2}]>] out_shardings=[<@mesh_x_8_y_8, [{"x":(2)2}], replicated={"y":(1)2}>] manual_axes={"y":(1)4, "x":(2)2} (%arg1: tensor<4xf32>) { -+ sdy.return %arg1 : tensor<4xf32> -+ } : (tensor<8xf32>) -> tensor<8xf32> -+ return %0 : tensor<8xf32> -+} -+ - // ----- - sdy.mesh @mesh = <["c"=2, "a"=2, "b"=2]> - -diff --git a/shardy/dialect/sdy/transforms/propagation/test/basic_propagation_manual_computation.mlir b/shardy/dialect/sdy/transforms/propagation/test/basic_propagation_manual_computation.mlir -index a6e9c643..669ad924 100644 ---- a/shardy/dialect/sdy/transforms/propagation/test/basic_propagation_manual_computation.mlir -+++ b/shardy/dialect/sdy/transforms/propagation/test/basic_propagation_manual_computation.mlir -@@ -589,3 +589,22 @@ func.func @manual_computation_with_tokens( - } : (!stablehlo.token, tensor<4x4xi64>) -> (!stablehlo.token, tensor<4x4xi64>) - return %0#0, %0#1 : !stablehlo.token, tensor<4x4xi64> - } -+ -+// ----- -+sdy.mesh @mesh = <["a"=4, "b"=2]> -+ -+// CHECK-LABEL: func @manual_computation_sub_axis_manual_merges_with_free_suffix( -+// CHECK-SAME: %arg0: tensor<32x32xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>}) -+// CHECK-SAME: -> (tensor<32x32xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>}) { -+func.func @manual_computation_sub_axis_manual_merges_with_free_suffix(%arg0: tensor<32x32xf32>) -> tensor<32x32xf32> { -+ // CHECK: sdy.manual_computation(%arg0) -+ // CHECK-SAME{LITERAL}: in_shardings=[<@mesh, [{"a", ?}, {?}]>] -+ // CHECK-SAME{LITERAL}: out_shardings=[<@mesh, [{"a", ?}, {?}]>] -+ // CHECK-SAME{LITERAL}: manual_axes={"a":(1)2} (%arg1: tensor<16x32xf32>) { -+ %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a":(1)2, ?}, {?}]>] out_shardings=[<@mesh, [{"a":(1)2, ?}, {?}]>] manual_axes={"a":(1)2} (%arg1: tensor<16x32xf32>) { -+ // CHECK: stablehlo.add %arg1, %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a":(2)2, ?}, {?}]>]>} : tensor<16x32xf32> -+ %1 = stablehlo.add %arg1, %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a":(2)2, ?}, {?}]>]>} : tensor<16x32xf32> -+ sdy.return %1 : tensor<16x32xf32> -+ } : (tensor<32x32xf32>) -> tensor<32x32xf32> -+ func.return %0: tensor<32x32xf32> -+} -diff --git a/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline.mlir b/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline.mlir -index 4119f39c..33c7ae02 100644 ---- a/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline.mlir -+++ b/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline.mlir -@@ -33,8 +33,7 @@ sdy.mesh @mesh = <["a"=2, "b"=2, "c"=2]> - func.func @user_priorities( - %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}p1, {"b"}p1]>}, - %arg1: tensor<8x8xf32>, %arg2: tensor<8x8xf32>, %arg3: tensor<8x8xf32>) -> tensor<8x8xf32> { -- // CHECK: %[[RESHARD:.*]] = sdy.reshard %arg0 <@mesh, [{"c"}, {"b"}]> : tensor<8x8xf32> -- // CHECK-NEXT: %[[ADD_0:.*]] = stablehlo.add %[[RESHARD]], %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"c"}, {"b"}]>]>} -+ // CHECK-NEXT: %[[ADD_0:.*]] = stablehlo.add %arg0, %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"c"}, {"b"}]>]>} - // CHECK-NEXT: %[[ADD_1:.*]] = stablehlo.add %[[ADD_0]], %arg2 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"c"}, {"b"}]>]>} - // CHECK-NEXT: stablehlo.divide %[[ADD_1]], %arg3 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"c"}, {"b"}]>]>} - %0 = stablehlo.add %arg0, %arg1 : tensor<8x8xf32> -@@ -49,10 +48,7 @@ sdy.mesh @mesh = <["a"=2, "b"=2, "c"=2]> - // CHECK-LABEL: func @sharding_constraint_applied - func.func @sharding_constraint_applied(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}) - -> tensor<8x8xf32> { -- // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>} : tensor<8x8xf32> -- // CHECK-NEXT: %[[RESHARD_0:.*]] = sdy.reshard %[[ADD]] <@mesh, [{}, {}]> : tensor<8x8xf32> -- // CHECK-NEXT: %[[RESHARD_1:.*]] = sdy.reshard %[[RESHARD_0]] <@mesh, [{}, {}]> : tensor<8x8xf32> -- // CHECK-NEXT: return %[[RESHARD_1]] : tensor<8x8xf32> -+ // CHECK-NEXT: stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>} - %0 = stablehlo.add %arg0, %arg0 : tensor<8x8xf32> - %1 = sdy.sharding_constraint %0 <@mesh, [{}, {}]> : tensor<8x8xf32> - return %1 : tensor<8x8xf32> -@@ -113,13 +109,11 @@ sdy.mesh @mesh = <["a"=2, "b"=2]> - // CHECK-SAME: ) -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {"b"}]>}, - // CHECK-SAME: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {"b"}]>}) { - func.func @add_extra_sharding_constraint_for_incompatible_group_member_shardings(%arg0: tensor<8x8xf32>) -> (tensor<8x8xf32>, tensor<8x8xf32>) { -- // CHECK: %[[RESHARD_0:.*]] = sdy.reshard %arg0 <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32> -- // CHECK-NEXT: %[[RESHARD_1:.*]] = sdy.reshard %[[RESHARD_0]] <@mesh, [{}, {"b"}]> : tensor<8x8xf32> -- // CHECK-NEXT: %[[RESHARD_2:.*]] = sdy.reshard %[[RESHARD_1]] <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32> -- // CHECK-NEXT: %[[RESHARD_3:.*]] = sdy.reshard %[[RESHARD_2]] <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32> -- // CHECK-NEXT: %[[RESHARD_4:.*]] = sdy.reshard %arg0 <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32> -- // CHECK-NEXT: %[[RESHARD_5:.*]] = sdy.reshard %[[RESHARD_4]] <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32> -- // CHECK-NEXT: return %[[RESHARD_3]], %[[RESHARD_5]] : tensor<8x8xf32>, tensor<8x8xf32> -+ // CHECK-NEXT: %[[RESHARD_0:.*]] = sdy.reshard %arg0 <@mesh, [{}, {"b"}]> -+ // CHECK-NEXT: %[[RESHARD_1:.*]] = sdy.reshard %[[RESHARD_0]] <@mesh, [{"a"}, {"b"}]> -+ // CHECK-NEXT: %[[RESHARD_2:.*]] = sdy.reshard %arg0 <@mesh, [{"a"}, {"b"}]> -+ // CHECK-NEXT: %[[RESHARD_3:.*]] = sdy.reshard %[[RESHARD_2]] <@mesh, [{"a"}, {"b"}]> -+ // CHECK-NEXT: return %[[RESHARD_1]], %[[RESHARD_3]] - %0 = sdy.sharding_constraint %arg0 <@mesh, [{}, {"b", ?}]> : tensor<8x8xf32> - sdy.sharding_group %0 group_id=1183 : tensor<8x8xf32> - %1 = sdy.sharding_constraint %arg0 <@mesh, [{"a"}, {?}]> : tensor<8x8xf32> -@@ -272,8 +266,7 @@ func.func @do_not_propagate_manual_axes_to_manual_computation(%arg0: tensor<8xf3 - // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg1, %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}]>]>} : tensor<8xf32> - // CHECK-NEXT: sdy.return %[[ADD]] : tensor<8xf32> - // CHECK-NEXT: } : (tensor<8xf32>) -> tensor<8xf32> -- // CHECK: %[[RESHARD:.*]] = sdy.reshard %[[MAN_COMP]] <@mesh, [{"a", "b"}]> : tensor<8xf32> -- // CHECK-NEXT: %[[SINE:.*]] = stablehlo.sine %[[RESHARD]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b"}]>]>} : tensor<8xf32> -+ // CHECK-NEXT: %[[SINE:.*]] = stablehlo.sine %[[MAN_COMP]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b"}]>]>} : tensor<8xf32> - // CHECK-NEXT: return %[[SINE]] : tensor<8xf32> - %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}]>] out_shardings=[<@mesh, [{?}]>] manual_axes={"b"} (%arg1: tensor<8xf32>) { - %1 = stablehlo.add %arg1, %arg1 : tensor<8xf32> -@@ -293,11 +286,9 @@ sdy.mesh @mesh_a_2 = <["a"=2]> - // CHECK-SAME: %arg1: tensor<1024x1024xf32> {sdy.sharding = #sdy.sharding<@mesh_a_2, [{"a"}, {}]>} - // CHECK-SAME: -> tensor<4x1024xf32> { - func.func @dot_lhs_from_broadcast_and_large_rhs(%arg0: tensor<4xf32> {sdy.sharding = #sdy.sharding<@mesh_a_2, [{"a", ?}]>}, %arg1: tensor<1024x1024xf32> {sdy.sharding = #sdy.sharding<@mesh_a_2, [{"a", ?}, {?}]>}) -> tensor<4x1024xf32> { -- // CHECK: %[[RESHARD:.*]] = sdy.reshard %arg0 <@mesh_a_2, [{}]> : tensor<4xf32> -- // CHECK-NEXT: %[[BROADCAST_IN_DIM:.*]] = stablehlo.broadcast_in_dim %[[RESHARD]], dims = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_a_2, [{}, {"a"}]>]>} -+ // CHECK-NEXT: %[[BROADCAST_IN_DIM:.*]] = stablehlo.broadcast_in_dim %arg0, dims = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_a_2, [{}, {"a"}]>]>} - // CHECK-NEXT: %[[DOT:.*]] = stablehlo.dot %[[BROADCAST_IN_DIM]], %arg1 : -- // CHECK-NEXT: %[[ALL_REDUCE:.*]] = sdy.all_reduce {"a"} %[[DOT]] out_sharding=<@mesh_a_2, [{}, {}]> : tensor<4x1024xf32> -- // CHECK-NEXT: %[[NEGATE:.*]] = stablehlo.negate %[[ALL_REDUCE]] : -+ // CHECK-NEXT: %[[NEGATE:.*]] = stablehlo.negate %[[DOT]] : - // CHECK-NEXT: return %[[NEGATE]] : tensor<4x1024xf32> - %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<4xf32>) -> tensor<4x1024xf32> - %1 = stablehlo.dot %0, %arg1 : (tensor<4x1024xf32>, tensor<1024x1024xf32>) -> tensor<4x1024xf32> -@@ -321,9 +312,7 @@ func.func @dot_general_with_unreduced_result( - // CHECK-SAME: {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}], unreduced={"b"}>]>} - // CHECK-NEXT: %[[ALL_REDUCE:.*]] = sdy.all_reduce {"b"} %[[DOT_GENERAL]] out_sharding=<@mesh, [{"a"}, {}]> - // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %[[ALL_REDUCE]] <@mesh, [{"a"}, {}]> : tensor<8x16xf32> -- // CHECK-NEXT: %[[RESHARD_LHS:.*]] = sdy.reshard %[[RESHARD]] <@mesh, [{"a"}, {"b"}]> : tensor<8x16xf32> -- // CHECK-NEXT: %[[RESHARD_RHS:.*]] = sdy.reshard %[[RESHARD]] <@mesh, [{"a"}, {"b"}]> : tensor<8x16xf32> -- // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[RESHARD_LHS]], %[[RESHARD_RHS]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {"b"}]>]>} -+ // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[RESHARD]], %[[RESHARD]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {"b"}]>]>} - // CHECK-NEXT: return %[[ADD]] - %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] - {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{?}, {?}], unreduced={"b"}>]>} : -@@ -348,8 +337,7 @@ func.func @dot_general_with_unreduced_result_fully_delayed( - // CHECK-SAME: {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}], unreduced={"b", "c"}>]>} - // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %[[DOT_GENERAL]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}], unreduced={"b", "c"}>]>} - // CHECK-NEXT: %[[ALL_REDUCE:.*]] = sdy.all_reduce {"b", "c"} %[[ADD]] out_sharding=<@mesh, [{"a"}, {}]> -- // CHECK-NEXT: %[[RESHARD_0:.*]] = sdy.reshard %[[ALL_REDUCE]] <@mesh, [{"a"}, {}]> : tensor<8x16xf32> -- // CHECK-NEXT: %[[FINAL_RESHARD:.*]] = sdy.reshard %[[RESHARD_0]] <@mesh, [{"a"}, {"b", "c"}]> : tensor<8x16xf32> -+ // CHECK-NEXT: %[[FINAL_RESHARD:.*]] = sdy.reshard %[[ALL_REDUCE]] <@mesh, [{"a"}, {}]> : tensor<8x16xf32> - // CHECK-NEXT: return %[[FINAL_RESHARD]] - %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] - {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{?}, {?}], unreduced={"b", "c"}>]>} : -@@ -376,8 +364,7 @@ func.func @dot_general_with_unreduced_result_partially_delayed( - // CHECK-NEXT: %[[RESHARD_B:.*]] = sdy.reshard %[[ALL_REDUCE_Y]] <@mesh, [{"a"}, {}], unreduced={"c"}> : tensor<8x16xf32> - // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[RESHARD_B]], %[[RESHARD_B]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}], unreduced={"c"}>]>} - // CHECK-NEXT: %[[ALL_REDUCE_C:.*]] = sdy.all_reduce {"c"} %[[ADD]] out_sharding=<@mesh, [{"a"}, {}]> -- // CHECK-NEXT: %[[RESHARD_0:.*]] = sdy.reshard %[[ALL_REDUCE_C]] <@mesh, [{"a"}, {}]> : tensor<8x16xf32> -- // CHECK-NEXT: %[[FINAL_RESHARD:.*]] = sdy.reshard %[[RESHARD_0]] <@mesh, [{"a"}, {"b", "c"}]> : tensor<8x16xf32> -+ // CHECK-NEXT: %[[FINAL_RESHARD:.*]] = sdy.reshard %[[ALL_REDUCE_C]] <@mesh, [{"a"}, {}]> : tensor<8x16xf32> - // CHECK-NEXT: return %[[FINAL_RESHARD]] - %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] - {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{?}, {?}], unreduced={"b", "c"}>]>} : -@@ -625,9 +612,7 @@ func.func @case_not_func_args_or_directly_returned( - %arg1: tensor<8xi64> {sdy.sharding = #sdy.sharding<@mesh_a_2_b_2_c_2, [{"a"}]>}, - %arg2: tensor<8xi64> {sdy.sharding = #sdy.sharding<@mesh_a_2_b_2_c_2, [{"a", "b"}]>} - ) -> tensor<8xi64> { -- // CHECK: %[[RESHARD_0:.*]] = sdy.reshard %arg1 <@mesh_a_2_b_2_c_2, [{"a", "b"}]> : tensor<8xi64> -- // CHECK: %[[RESHARD_1:.*]] = sdy.reshard %arg1 <@mesh_a_2_b_2_c_2, [{"a", "b"}]> : tensor<8xi64> -- // CHECK-NEXT: stablehlo.add %[[RESHARD_0]], %[[RESHARD_1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_a_2_b_2_c_2, [{"a", "b"}]>]>} : tensor<8xi64> -+ // CHECK: stablehlo.add %arg1, %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_a_2_b_2_c_2, [{"a", "b"}]>]>} : tensor<8xi64> - %0 = stablehlo.add %arg1, %arg1 : tensor<8xi64> - // CHECK: stablehlo.add %arg2, %arg2 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_a_2_b_2_c_2, [{"a", "b"}]>]>} : tensor<8xi64> - %1 = stablehlo.add %arg2, %arg2 : tensor<8xi64> -@@ -1043,9 +1028,8 @@ func.func @main(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - // CHECK-SAME: {sdy.sharding = #sdy.sharding<@mesh_a_2_b_2, [{"a"}, {}]>}) - // CHECK-SAME: -> (tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh_a_2_b_2, [{"b"}, {}]>}) { - // CHECK-NEXT: %[[RESHARD0:.*]] = sdy.reshard %arg0 <@mesh_a_2_b_2, [{"a"}, {}]> : tensor<8x2xi32> --// CHECK-NEXT: %[[RESHARD1:.*]] = sdy.reshard %[[RESHARD0]] <@mesh_a_2_b_2, [{"a"}, {}]> : tensor<8x2xi32> --// CHECK-NEXT: %[[RESHARD2:.*]] = sdy.reshard %[[RESHARD1]] <@mesh_a_2_b_2, [{"b"}, {}]> : tensor<8x2xi32> --// CHECK-NEXT: return %[[RESHARD2]] : tensor<8x2xi32> -+// CHECK-NEXT: %[[RESHARD1:.*]] = sdy.reshard %[[RESHARD0]] <@mesh_a_2_b_2, [{"b"}, {}]> : tensor<8x2xi32> -+// CHECK-NEXT: return %[[RESHARD1]] : tensor<8x2xi32> - // CHECK-NEXT: } - func.func private @foo(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - %0 = sdy.sharding_constraint %arg0 <@mesh_a_2_b_2, [{"a"}, {}]> : tensor<8x2xi32> -@@ -1195,9 +1179,8 @@ func.func private @bar(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - // CHECK-SAME: -> (tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}, {}]>}) { - func.func private @foo(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - // CHECK-NEXT: %[[ABS:.*]] = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>} : tensor<8x2xi32> -- // CHECK-NEXT: %[[RESHARD0:.*]] = sdy.reshard %[[ABS]] <@mesh, [{"a"}, {}]> : tensor<8x2xi32> -- // CHECK-NEXT: %[[RESHARD1:.*]] = sdy.reshard %[[RESHARD0]] <@mesh, [{"b"}, {}]> : tensor<8x2xi32> -- // CHECK-NEXT: %[[CALL:.*]] = call @bar(%[[RESHARD1]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : (tensor<8x2xi32>) -> tensor<8x2xi32> -+ // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %0 <@mesh, [{"b"}, {}]> : tensor<8x2xi32> -+ // CHECK-NEXT: %[[CALL:.*]] = call @bar(%[[RESHARD]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : (tensor<8x2xi32>) -> tensor<8x2xi32> - // CHECK-NEXT: return %[[CALL]] : tensor<8x2xi32> - %0 = stablehlo.abs %arg0 : tensor<8x2xi32> - %1 = sdy.sharding_constraint %0 <@mesh, [{"a"}, {}]> : tensor<8x2xi32> -diff --git a/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline_data_flow_edges.mlir b/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline_data_flow_edges.mlir -index 529f9735..def33952 100644 ---- a/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline_data_flow_edges.mlir -+++ b/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline_data_flow_edges.mlir -@@ -236,9 +236,7 @@ func.func @case_not_func_args_or_directly_returned( - %arg1: tensor<8xi64> {sdy.sharding = #sdy.sharding<@mesh_a_2_b_2_c_2, [{"a"}]>}, - %arg2: tensor<8xi64> {sdy.sharding = #sdy.sharding<@mesh_a_2_b_2_c_2, [{"a", "b"}]>} - ) -> tensor<8xi64> { -- // CHECK: %[[RESHARD_0:.*]] = sdy.reshard %arg1 <@mesh_a_2_b_2_c_2, [{"a", "b"}]> : tensor<8xi64> -- // CHECK: %[[RESHARD_1:.*]] = sdy.reshard %arg1 <@mesh_a_2_b_2_c_2, [{"a", "b"}]> : tensor<8xi64> -- // CHECK-NEXT: stablehlo.add %[[RESHARD_0]], %[[RESHARD_1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_a_2_b_2_c_2, [{"a", "b"}]>]>} : tensor<8xi64> -+ // CHECK: stablehlo.add %arg1, %arg1 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_a_2_b_2_c_2, [{"a", "b"}]>]>} : tensor<8xi64> - %0 = stablehlo.add %arg1, %arg1 : tensor<8xi64> - // CHECK: stablehlo.add %arg2, %arg2 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_a_2_b_2_c_2, [{"a", "b"}]>]>} : tensor<8xi64> - %1 = stablehlo.add %arg2, %arg2 : tensor<8xi64> -@@ -654,9 +652,8 @@ func.func @main(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - // CHECK-SAME: {sdy.sharding = #sdy.sharding<@mesh_a_2_b_2, [{"a"}, {}]>}) - // CHECK-SAME: -> (tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh_a_2_b_2, [{"b"}, {}]>}) { - // CHECK-NEXT: %[[RESHARD0:.*]] = sdy.reshard %arg0 <@mesh_a_2_b_2, [{"a"}, {}]> : tensor<8x2xi32> --// CHECK-NEXT: %[[RESHARD1:.*]] = sdy.reshard %[[RESHARD0]] <@mesh_a_2_b_2, [{"a"}, {}]> : tensor<8x2xi32> --// CHECK-NEXT: %[[RESHARD2:.*]] = sdy.reshard %[[RESHARD1]] <@mesh_a_2_b_2, [{"b"}, {}]> : tensor<8x2xi32> --// CHECK-NEXT: return %[[RESHARD2]] : tensor<8x2xi32> -+// CHECK-NEXT: %[[RESHARD1:.*]] = sdy.reshard %[[RESHARD0]] <@mesh_a_2_b_2, [{"b"}, {}]> : tensor<8x2xi32> -+// CHECK-NEXT: return %[[RESHARD1]] : tensor<8x2xi32> - // CHECK-NEXT: } - func.func private @foo(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - %0 = sdy.sharding_constraint %arg0 <@mesh_a_2_b_2, [{"a"}, {}]> : tensor<8x2xi32> -@@ -806,9 +803,8 @@ func.func private @bar(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - // CHECK-SAME: -> (tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}, {}]>}) { - func.func private @foo(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - // CHECK-NEXT: %[[ABS:.*]] = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>} : tensor<8x2xi32> -- // CHECK-NEXT: %[[RESHARD0:.*]] = sdy.reshard %[[ABS]] <@mesh, [{"a"}, {}]> : tensor<8x2xi32> -- // CHECK-NEXT: %[[RESHARD1:.*]] = sdy.reshard %[[RESHARD0]] <@mesh, [{"b"}, {}]> : tensor<8x2xi32> -- // CHECK-NEXT: %[[CALL:.*]] = call @bar(%[[RESHARD1]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : (tensor<8x2xi32>) -> tensor<8x2xi32> -+ // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %0 <@mesh, [{"b"}, {}]> : tensor<8x2xi32> -+ // CHECK-NEXT: %[[CALL:.*]] = call @bar(%[[RESHARD]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : (tensor<8x2xi32>) -> tensor<8x2xi32> - // CHECK-NEXT: return %[[CALL]] : tensor<8x2xi32> - %0 = stablehlo.abs %arg0 : tensor<8x2xi32> - %1 = sdy.sharding_constraint %0 <@mesh, [{"a"}, {}]> : tensor<8x2xi32> -diff --git a/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline_dedup_functions_fully_true.mlir b/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline_dedup_functions_fully_true.mlir -index 2286a0d0..b1f17674 100644 ---- a/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline_dedup_functions_fully_true.mlir -+++ b/shardy/dialect/sdy/transforms/propagation/test/propagation_pipeline_dedup_functions_fully_true.mlir -@@ -104,9 +104,8 @@ func.func @main(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - // CHECK-SAME: -> (tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}, {}]>}) { - func.func private @foo(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - // CHECK-NEXT: %[[ABS0:.*]] = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>} : tensor<8x2xi32> -- // CHECK-NEXT: %[[RESHARD0:.*]] = sdy.reshard %[[ABS0]] <@mesh, [{"a"}, {}]> : tensor<8x2xi32> -- // CHECK-NEXT: %[[RESHARD1:.*]] = sdy.reshard %[[RESHARD0]] <@mesh, [{"b"}, {}]> : tensor<8x2xi32> -- // CHECK-NEXT: %[[ABS1:.*]] = stablehlo.abs %[[RESHARD1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : tensor<8x2xi32> -+ // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %[[ABS0]] <@mesh, [{"b"}, {}]> : tensor<8x2xi32> -+ // CHECK-NEXT: %[[ABS1:.*]] = stablehlo.abs %[[RESHARD]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : tensor<8x2xi32> - // CHECK-NEXT: %[[CALL0:.*]] = call @bar(%[[ABS1]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : (tensor<8x2xi32>) -> tensor<8x2xi32> - // CHECK-NEXT: %[[CALL1:.*]] = call @bar(%[[ABS1]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : (tensor<8x2xi32>) -> tensor<8x2xi32> - // CHECK-NEXT: return %[[CALL0]] : tensor<8x2xi32> -@@ -191,9 +190,8 @@ func.func @main(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - // CHECK-SAME: -> (tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}, {}]>}) { - func.func private @foo(%arg0: tensor<8x2xi32>) -> tensor<8x2xi32> { - // CHECK-NEXT: %[[ABS0:.*]] = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>} : tensor<8x2xi32> -- // CHECK-NEXT: %[[RESHARD0:.*]] = sdy.reshard %[[ABS0]] <@mesh, [{"a"}, {}]> : tensor<8x2xi32> -- // CHECK-NEXT: %[[RESHARD1:.*]] = sdy.reshard %[[RESHARD0]] <@mesh, [{"b"}, {}]> : tensor<8x2xi32> -- // CHECK-NEXT: %[[ABS1:.*]] = stablehlo.abs %[[RESHARD1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : tensor<8x2xi32> -+ // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %[[ABS0]] <@mesh, [{"b"}, {}]> : tensor<8x2xi32> -+ // CHECK-NEXT: %[[ABS1:.*]] = stablehlo.abs %[[RESHARD]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : tensor<8x2xi32> - // CHECK-NEXT: %[[CALL:.*]] = call @bar(%[[ABS1]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : (tensor<8x2xi32>) -> tensor<8x2xi32> - // CHECK-NEXT: return %[[CALL]] : tensor<8x2xi32> - %0 = stablehlo.abs %arg0 : tensor<8x2xi32> -@@ -268,10 +266,9 @@ sdy.mesh @mesh = <["a"=2, "b"=2]> - // CHECK-SAME: %arg0: tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}) - // CHECK-SAME: -> (tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}, {}]>}, tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}, {}]>}) - func.func @main(%arg0: tensor<8x2xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}) -> (tensor<8x2xi32>, tensor<8x2xi32>) { -- // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>} : tensor<8x2xi32> -- // CHECK-NEXT: %[[RESHARD:.*]] = sdy.reshard %[[ADD]] <@mesh, [{"b"}, {}]> : tensor<8x2xi32> -- // CHECK-NEXT: %[[CALL0:.*]] = call @foo(%[[RESHARD]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : (tensor<8x2xi32>) -> tensor<8x2xi32> -- // CHECK-NEXT: %[[ABS:.*]] = stablehlo.abs %[[RESHARD]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : tensor<8x2xi32> -+ // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : tensor<8x2xi32> -+ // CHECK-NEXT: %[[CALL0:.*]] = call @foo(%[[ADD]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : (tensor<8x2xi32>) -> tensor<8x2xi32> -+ // CHECK-NEXT: %[[ABS:.*]] = stablehlo.abs %[[ADD]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : tensor<8x2xi32> - // CHECK-NEXT: %[[CALL1:.*]] = call @bar(%[[ABS]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}]>]>} : (tensor<8x2xi32>) -> tensor<8x2xi32> - // CHECK-NEXT: return %[[CALL0]], %[[CALL1]] : tensor<8x2xi32>, tensor<8x2xi32> - %0 = stablehlo.add %arg0, %arg0 : tensor<8x2xi32> -@@ -315,13 +312,11 @@ sdy.mesh @mesh = <["a"=2, "b"=2]> - // CHECK-SAME: ) -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {"b"}]>}, - // CHECK-SAME: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {"b"}]>}) { - func.func @add_extra_sharding_constraint_for_incompatible_group_member_shardings(%arg0: tensor<8x8xf32>) -> (tensor<8x8xf32>, tensor<8x8xf32>) { -- // CHECK-NEXT: %[[RESHARD_0:.*]] = sdy.reshard %arg0 <@mesh, [{"a"}, {"b"}]> -- // CHECK-NEXT: %[[RESHARD_1:.*]] = sdy.reshard %[[RESHARD_0]] <@mesh, [{}, {"b"}]> -- // CHECK-NEXT: %[[RESHARD_2:.*]] = sdy.reshard %[[RESHARD_1]] <@mesh, [{"a"}, {"b"}]> -+ // CHECK-NEXT: %[[RESHARD_0:.*]] = sdy.reshard %arg0 <@mesh, [{}, {"b"}]> -+ // CHECK-NEXT: %[[RESHARD_1:.*]] = sdy.reshard %[[RESHARD_0]] <@mesh, [{"a"}, {"b"}]> -+ // CHECK-NEXT: %[[RESHARD_2:.*]] = sdy.reshard %arg0 <@mesh, [{"a"}, {"b"}]> - // CHECK-NEXT: %[[RESHARD_3:.*]] = sdy.reshard %[[RESHARD_2]] <@mesh, [{"a"}, {"b"}]> -- // CHECK-NEXT: %[[RESHARD_4:.*]] = sdy.reshard %arg0 <@mesh, [{"a"}, {"b"}]> -- // CHECK-NEXT: %[[RESHARD_5:.*]] = sdy.reshard %[[RESHARD_4]] <@mesh, [{"a"}, {"b"}]> -- // CHECK-NEXT: return %[[RESHARD_3]], %[[RESHARD_5]] -+ // CHECK-NEXT: return %[[RESHARD_1]], %[[RESHARD_3]] - %0 = sdy.sharding_constraint %arg0 <@mesh, [{}, {"b", ?}]> : tensor<8x8xf32> - sdy.sharding_group %0 group_id=1183 : tensor<8x8xf32> - %1 = sdy.sharding_constraint %arg0 <@mesh, [{"a"}, {?}]> : tensor<8x8xf32> -diff --git a/shardy/integrations/c/attributes.cc b/shardy/integrations/c/attributes.cc -index 56bfd11a..1c75af21 100644 ---- a/shardy/integrations/c/attributes.cc -+++ b/shardy/integrations/c/attributes.cc -@@ -436,8 +436,12 @@ bool sdyAttributeIsAManualAxesAttr(MlirAttribute attr) { - return mlir::isa(unwrap(attr)); - } - --MlirAttribute sdyManualAxesAttrGet( -- MlirContext ctx, intptr_t nAxes, const MlirAttribute* axes) { -+MlirAttribute sdyManualAxesAttrGet(MlirContext ctx, intptr_t nAxes, -+ const MlirAttribute* axes) { -+ if (nAxes > 0 && mlir::isa(unwrap(axes[0]))) { -+ return wrap(sdy::ManualAxesAttr::get( -+ unwrap(ctx), unwrapAttrs(axes, nAxes))); -+ } - return wrap(sdy::ManualAxesAttr::get( - unwrap(ctx), unwrapAttrs(axes, nAxes))); - } -@@ -446,9 +450,8 @@ intptr_t sdyManualAxesAttrGetAxesSize(MlirAttribute attr) { - return unwrapAttr(attr).size(); - } - --MlirStringRef sdyManualAxesAttrGetAxesElem( -- MlirAttribute attr, intptr_t pos) { -- return wrap(unwrapAttr(attr)[pos].getValue()); -+MlirStringRef sdyManualAxesAttrGetAxesElem(MlirAttribute attr, intptr_t pos) { -+ return wrap(unwrapAttr(attr)[pos].getName()); - } - - } // extern "C" diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl index 1284554410c..3d83548ce08 100644 --- a/third_party/xla/third_party/shardy/workspace.bzl +++ b/third_party/xla/third_party/shardy/workspace.bzl @@ -18,8 +18,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def repo(): - SHARDY_COMMIT = "f16d43be0f7c9a7e42c4013180adc5d0594b2535" - SHARDY_SHA256 = "be8607a2279333d707fb97a9c23addefc953fb38e626cdb668f9f49668f13fe1" + SHARDY_COMMIT = "b85d2632dd8a4c76a652d8ab329d7118ad0b35ea" + SHARDY_SHA256 = "20eb2fd15663e1c02e9c580786e8e7d035905a2e659e2ef9e73810a849985db3" tf_http_archive( name = "shardy", diff --git a/third_party/xla/third_party/triton/common/llvm_cl1004097735.patch b/third_party/xla/third_party/triton/common/llvm_cl1004097735.patch new file mode 100755 index 00000000000..e2cbc35cf97 --- /dev/null +++ b/third_party/xla/third_party/triton/common/llvm_cl1004097735.patch @@ -0,0 +1,26 @@ +--- a/lib/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.cpp ++++ b/lib/Conversion/TritonGPUToLLVM/WarpSpecializeUtility.cpp +@@ -369,8 +369,8 @@ + MLIRContext *ctx = b.getContext(); + auto licmMD = LLVM::LoopLICMAttr::get(ctx, b.getBoolAttr(true), {}); + auto loopMD = +- LLVM::LoopAnnotationAttr::get(b.getContext(), {}, {}, {}, {}, {}, licmMD, +- {}, {}, {}, {}, {}, {}, {}, {}, {}); ++ LLVM::LoopAnnotationAttr::get(b.getContext(), {}, {}, {}, {}, {}, {}, ++ licmMD, {}, {}, {}, {}, {}, {}, {}, {}, {}); + latchBr.setLoopAnnotationAttr(loopMD); + } + + +--- a/python/src/ir.cc ++++ b/python/src/ir.cc +@@ -959,7 +959,7 @@ + self.getBuilder().getBoolAttr(true)); + mlir::LLVM::LoopAnnotationAttr la = + mlir::LLVM::LoopAnnotationAttr::get( +- self.getBuilder().getContext(), {}, {}, {}, {}, {}, ++ self.getBuilder().getContext(), {}, {}, {}, {}, {}, {}, + licmAttr, {}, {}, {}, {}, {}, {}, {}, {}, {}); + return la; + }) + diff --git a/third_party/xla/third_party/triton/common/series b/third_party/xla/third_party/triton/common/series index f41d6167743..99074680ef0 100644 --- a/third_party/xla/third_party/triton/common/series +++ b/third_party/xla/third_party/triton/common/series @@ -38,3 +38,4 @@ cxx17_support.patch llvm_cl996124567.patch llvm_cl998243985.patch llvm_cl1000840178.patch +llvm_cl1004097735.patch diff --git a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/add_loop_unroll_flags.cc b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/add_loop_unroll_flags.cc index 4a9bbbc7f5b..c177bd2ebfd 100644 --- a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/add_loop_unroll_flags.cc +++ b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/add_loop_unroll_flags.cc @@ -68,7 +68,7 @@ class AddLoopUnrollFlagsPass /*followupUnrolled=*/nullptr, /*followupRemainder=*/nullptr, /*followupAll=*/nullptr); auto loop_annotation = mlir::LLVM::LoopAnnotationAttr::get( - context, + context, /*id=*/nullptr, /*disableNonforced=*/nullptr, /*vectorize=*/nullptr, /*interleave=*/nullptr, loop_unroll, /*unrollAndJam=*/nullptr, /*licm=*/nullptr, /*distribute=*/nullptr,