diff --ruN a/stablehlo/stablehlo/dialect/Serialization.cpp b/stablehlo/stablehlo/dialect/Serialization.cpp --- stablehlo/stablehlo/dialect/Serialization.cpp +++ stablehlo/stablehlo/dialect/Serialization.cpp @@ -24,6 +24,7 @@ #include "mlir/IR/Location.h" #include "mlir/IR/MLIRContext.h" #include "mlir/IR/OwningOpRef.h" +#include "mlir/IR/Verifier.h" #include "mlir/Parser/Parser.h" #include "mlir/Pass/PassManager.h" #include "mlir/Support/LLVM.h" @@ -43,11 +44,18 @@ bool allowOtherDialects) { MLIRContext* context = module.getContext(); + // Only enable verifier in debug builds. + bool enableVerifier = false; +#ifndef NDEBUG + enableVerifier = true; +#endif + // Convert StableHLO --> VHLO. // If allowOtherDialects is true, we will allow other dialects to be present // in the module, otherwise will fail if there are any other dialects present. { PassManager pm(context); + pm.enableVerifier(enableVerifier); StablehloLegalizeToVhloPassOptions options; options.allowOtherDialects = allowOtherDialects; pm.addPass(stablehlo::createStablehloLegalizeToVhloPass(options)); @@ -61,6 +69,7 @@ // target version failures. { PassManager pm(context); + pm.enableVerifier(enableVerifier); pm.addPass(stablehlo::createVhloToVersionPass({targetVersion.str()})); if (!succeeded(pm.run(module))) { return failure(); @@ -81,7 +90,15 @@ OwningOpRef deserializePortableArtifact(StringRef sourceStr, MLIRContext* context) { context->loadDialect(); - auto module = parseSourceString(sourceStr, context); + + // Only enable verifier in debug builds. + bool enableVerifier = false; +#ifndef NDEBUG + enableVerifier = true; +#endif + + ParserConfig config(context, /*verifyAfterParse=*/enableVerifier); + auto module = parseSourceString(sourceStr, config); if (!module) { emitError(UnknownLoc::get(context)) << "failed to deserialize portable artifact using StableHLO_v" @@ -91,8 +108,13 @@ // Convert VHLO --> VHLO(current) --> StableHLO PassManager pm(context); + pm.enableVerifier(enableVerifier); createStablehloDeserializePipeline(pm); if (!succeeded(pm.run(*module))) { + return nullptr; + } + + if (failed(verify(*module))) { return nullptr; } diff --ruN a/stablehlo/stablehlo/tests/TestUtils.cpp b/stablehlo/stablehlo/tests/TestUtils.cpp --- stablehlo/stablehlo/tests/TestUtils.cpp +++ stablehlo/stablehlo/tests/TestUtils.cpp @@ -68,6 +68,32 @@ } }; +struct BroadcastIfNeededPattern : public RewritePattern { + explicit BroadcastIfNeededPattern(MLIRContext* context) + : RewritePattern("hlo_test_broadcast.broadcast_if_needed", 1, context) {} + LogicalResult matchAndRewrite(Operation* op, + PatternRewriter& rewriter) const override { + if (op->getNumOperands() < 1) return failure(); + Value input = op->getOperand(0); + + SmallVector broadcastDimensions; + if (auto bcastDimsAttr = + op->getAttrOfType("broadcast_dimensions")) { + broadcastDimensions = llvm::to_vector(bcastDimsAttr.asArrayRef()); + } + + auto targetShape = stablehlo::getDimensions(op->getResult(0)); + if (failed(targetShape)) return failure(); + + auto broadcastedVal = stablehlo::broadcastIfNeeded( + rewriter, input, *targetShape, broadcastDimensions); + if (failed(broadcastedVal)) return failure(); + + rewriter.replaceOp(op, *broadcastedVal); + return success(); + } +}; + struct InferReturnTypesPattern : public RewritePattern { explicit InferReturnTypesPattern(MLIRContext* context) : RewritePattern("hlo_test_infer.get_return_types", 1, context) {} @@ -222,6 +248,7 @@ LogicalResult initialize(MLIRContext* context) override { RewritePatternSet patterns(context); patterns.add(context); + patterns.add(context); patterns_ = std::move(patterns); return success(); } diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir --- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir +++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir @@ -5282,6 +5282,32 @@ // ----- +// CHECK-LABEL: func.func @ragged_dot_mode_2_rank_lhs_lt_rhs( +// CHECK-SAME: %[[ARG0:.*]]: tensor<2x3xf32>, +// CHECK-SAME: %[[ARG1:.*]]: tensor<2x3x4xf32>, +// CHECK-SAME: %[[ARG2:.*]]: tensor<2xi64>) -> tensor<2x3x3x4xf32> { +// CHECK: %[[DOT_GENERAL_0:.*]] = stablehlo.dot_general %{{.*}}, %[[ARG1]], contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2x3xf32>, tensor<2x3x4xf32>) -> tensor<3x3x4xf32> +// CHECK: %[[DOT_GENERAL_1:.*]] = stablehlo.dot_general %{{.*}}, %[[ARG1]], contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2x3xf32>, tensor<2x3x4xf32>) -> tensor<3x3x4xf32> +// CHECK: %[[CONCATENATE_0:.*]] = stablehlo.concatenate{{.*}}dim = 0 : (tensor<1x3x3x4xf32>, tensor<1x3x3x4xf32>) -> tensor<2x3x3x4xf32> +// CHECK: return %[[CONCATENATE_0]] : tensor<2x3x3x4xf32> +// CHECK: } +func.func @ragged_dot_mode_2_rank_lhs_lt_rhs(%lhs : tensor<2x3xf32>, %rhs : tensor<2x3x4xf32>, %group_sizes : tensor<2xi64>) -> tensor<2x3x3x4xf32> { + %0 = "chlo.ragged_dot"(%lhs, %rhs, %group_sizes) { + ragged_dot_dimension_numbers = #chlo.ragged_dot< + lhs_batching_dimensions = [], + rhs_batching_dimensions = [], + lhs_contracting_dimensions = [0], + rhs_contracting_dimensions = [0], + lhs_ragged_dimensions = [0], + rhs_group_dimensions = [] + >, + precision_config = [#chlo, #chlo] + } : (tensor<2x3xf32>, tensor<2x3x4xf32>, tensor<2xi64>) -> tensor<2x3x3x4xf32> + func.return %0 : tensor<2x3x3x4xf32> +} + +// ----- + // CHECK-LABEL: func.func @ragged_dot_mode_3( // CHECK-SAME: %[[ARG0:.*]]: tensor<2x3x5xf32>, // CHECK-SAME: %[[ARG1:.*]]: tensor<2x5x7xf32>, diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stablehlo/tests/ops_broadcasting.mlir --- stablehlo/stablehlo/tests/ops_broadcasting.mlir +++ stablehlo/stablehlo/tests/ops_broadcasting.mlir @@ -320,3 +320,14 @@ return %0 : !stablehlo.token } +// ----- + +// Non-numpy explicit broadcast_dimensions: [3, 1] -> [3, 4, 5] +// CHECK-LABEL: func @explicit_broadcast_dims +func.func @explicit_broadcast_dims(%arg0: tensor<3x1xf64>) -> tensor<3x4x5xf64> { + // CHECK: %[[BCAST:.+]] = stablehlo.broadcast_in_dim %arg0, dims = [0, 1] : (tensor<3x1xf64>) -> tensor<3x4x5xf64> + // CHECK-NEXT: return %[[BCAST]] : tensor<3x4x5xf64> + %0 = "hlo_test_broadcast.broadcast_if_needed"(%arg0) {broadcast_dimensions = array} : (tensor<3x1xf64>) -> tensor<3x4x5xf64> + return %0 : tensor<3x4x5xf64> +} + diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp --- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp +++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp @@ -2752,13 +2752,25 @@ LogicalResult matchAndRewrite( mlir::chlo::RaggedDotOp op, OpAdaptor, ConversionPatternRewriter& rewriter) const override { - if (op.getLhs().getType().getRank() < op.getRhs().getType().getRank()) { + chlo::RaggedDotDimensionNumbersAttr raggedDotDimensionNumbers = + op.getRaggedDotDimensionNumbers(); + ArrayRef lhsRaggedDimensions = + raggedDotDimensionNumbers.getLhsRaggedDimensions(); + if (lhsRaggedDimensions.empty()) { + return rewriter.notifyMatchFailure( + op, "lhs_ragged_dimensions must not be empty"); + } + const int64_t lhsRaggedDim = lhsRaggedDimensions[0]; + if (llvm::is_contained( + raggedDotDimensionNumbers.getLhsContractingDimensions(), + lhsRaggedDim)) { + return handleRaggedDotMode2(op, rewriter); + } else if (llvm::is_contained( + raggedDotDimensionNumbers.getLhsBatchingDimensions(), + lhsRaggedDim)) { + return handleRaggedDotMode3(op, rewriter); + } else { return handleRaggedDotMode1(op, rewriter); - } else if (op.getLhs().getType().getRank() < - op.getResult().getType().getRank()) { - return handleRaggedDotMode2(op, rewriter); - } else { - return handleRaggedDotMode3(op, rewriter); } } }; diff --ruN a/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp b/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp --- stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp +++ stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp @@ -224,23 +224,51 @@ mlir::RankedTensorType outputType = getRankedTensorType(shape, inputType.getElementType()); - // Short circuit if no broadcasting is needed. - if (inputType == outputType) return input; - int64_t inputRank = inputType.getRank(); int64_t outputRank = outputType.getRank(); if (inputRank > outputRank) return emitError(loc, "input rank must be <= output rank, got ") << inputRank << " vs " << outputRank; - size_t rankDiff = outputRank - inputRank; + // Construct broadcast dimensions (right-aligned for NumPy-style + // broadcasting). + auto broadcastDimensions = + llvm::to_vector(llvm::seq(outputRank - inputRank, outputRank)); + + return broadcastIfNeeded(builder, input, shape, broadcastDimensions); +} + +FailureOr broadcastIfNeeded(OpBuilder& builder, Value input, + const Dimensions& shape, + ArrayRef broadcastDimensions) { + LLVM_DEBUG(llvm::dbgs() << "[broadcastIfNeeded] Broadcasting input " + << input.getType() << " => " << toString(shape) + << "\n"); + auto loc = input.getLoc(); + mlir::RankedTensorType inputType = + dyn_cast(input.getType()); + if (!inputType) + return emitError(loc, "expected ranked tensor type for broadcast inputs"); + mlir::RankedTensorType outputType = + getRankedTensorType(shape, inputType.getElementType()); + + // Short circuit if no broadcasting is needed. + if (inputType == outputType) return input; + + int64_t inputRank = inputType.getRank(); + int64_t outputRank = outputType.getRank(); + if (inputRank > outputRank) + return emitError(loc, "input rank must be <= output rank, got ") + << inputRank << " vs " << outputRank; + + if (static_cast(broadcastDimensions.size()) != inputRank) + return emitError(loc, "broadcast_dimensions size (") + << broadcastDimensions.size() << ") must match input rank (" + << inputRank << ")"; + auto inputShapeOrFail = getDimensions(input); if (failed(inputShapeOrFail)) return failure(); Dimensions inputShape = std::move(*inputShapeOrFail); - - // Construct broadcast dimensions. - auto broadcastDimensions = - llvm::to_vector(llvm::seq(outputRank - inputRank, outputRank)); // Construct the result type of the broadcast // - If input is static and target shape is static, use static shape. @@ -248,14 +276,30 @@ // - If input is not bounded, but target shape is bounded, broadcast to // the padded shape then call SetDimensionSize to make dynamic. auto bcastShape = shape; + llvm::SmallVector isMapped(outputRank, false); + for (int64_t i = 0; i < inputRank; ++i) { + int64_t resultIdx = broadcastDimensions[i]; + if (resultIdx < 0 || resultIdx >= outputRank) + return emitError(loc, "broadcast_dimensions index ") + << resultIdx << " out of bounds for output rank " << outputRank; + + isMapped[resultIdx] = true; + int64_t inputDimSize = inputShape[i].size; - int64_t resultIdx = i + rankDiff; int64_t resultDimSize = shape[resultIdx].size; if (inputDimSize != 1 && inputDimSize != resultDimSize) return emitError(loc, "Cannot broadcast input: ") << inputType << " to target shape " << toString(shape); + if (inputShape[i].boundOp.has_value() && + !shape[resultIdx].boundOp.has_value()) { + return emitError( + loc, "cannot mix bounded and static dimensions in broadcast: ") + << "input dimension " << i << " is bounded, but target dimension " + << resultIdx << " is static"; + } + if (!inputShape[i].boundOp.has_value() && shape[resultIdx].boundOp.has_value()) { // Use padded shape in broadcast. @@ -263,17 +307,18 @@ } } - // Broadcast to padded size for remaining dimensions. - for (size_t i = 0; i < rankDiff; ++i) { - bcastShape[i] = DimensionInfo{shape[i].size}; + // Broadcast to padded size for remaining unmapped dimensions. + for (int64_t i = 0; i < outputRank; ++i) { + if (!isMapped[i]) { + bcastShape[i] = DimensionInfo{shape[i].size}; + } } // Insert broadcast ops mlir::RankedTensorType bcastType = getRankedTensorType(bcastShape, inputType.getElementType()); - LLVM_DEBUG( - llvm::dbgs() << "[numpyBroadcastIfNeeded] Broadcast to padded type " - << bcastType << "\n"); + LLVM_DEBUG(llvm::dbgs() << "[broadcastIfNeeded] Broadcast to padded type " + << bcastType << "\n"); Value bcastOp = stablehlo::BroadcastInDimOp::create( builder, loc, bcastType, input, broadcastDimensions); if (bcastOp.getType() == outputType) return bcastOp; diff --ruN a/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h b/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h --- stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h +++ stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h @@ -70,6 +70,12 @@ FailureOr numpyBroadcastIfNeeded(OpBuilder& builder, Value input, const Dimensions& shape); +// Apply broadcasting to the given operand using the specified +// broadcast_dimensions, returning an error if the operand is not broadcastable. +FailureOr broadcastIfNeeded(OpBuilder& builder, Value input, + const Dimensions& shape, + ArrayRef broadcastDimensions); + } // namespace stablehlo } // namespace mlir