diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp index c3d42f7..531fe34 100644 --- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp @@ -140,6 +140,7 @@ void ONNXToSpatialPass::runOnOperation() { target.addIllegalOp(); target.addIllegalOp(); target.addIllegalOp(); + target.addIllegalOp(); target.addIllegalOp(); target.addIllegalOp(); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp index 2f03115..200d7a3 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp @@ -6,6 +6,8 @@ #include #include +#include +#include #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" @@ -19,6 +21,85 @@ using namespace mlir; namespace onnx_mlir { namespace { +struct ReduceMeanSemantics { + SmallVector axes; + int64_t keepdims = 1; + bool isIdentity = false; +}; + +static bool isNoneValueLike(Value value) { return isa_and_nonnull(value.getDefiningOp()); } + +static FailureOr> getConstantIntValues(Value value) { + auto denseAttr = dyn_cast_or_null(getHostConstDenseElementsAttr(value)); + if (!denseAttr) + return failure(); + return SmallVector(denseAttr.getValues().begin(), denseAttr.getValues().end()); +} + +static FailureOr> normalizeAxesChecked(ArrayRef axes, int64_t rank) { + SmallVector normalizedAxes; + normalizedAxes.reserve(axes.size()); + for (int64_t axis : axes) { + auto normalizedAxis = normalizeAxisChecked(axis, rank); + if (failed(normalizedAxis)) + return failure(); + normalizedAxes.push_back(*normalizedAxis); + } + llvm::sort(normalizedAxes); + normalizedAxes.erase(std::unique(normalizedAxes.begin(), normalizedAxes.end()), normalizedAxes.end()); + return normalizedAxes; +} + +template +static FailureOr +getReduceMeanSemantics(ReduceMeanOp reduceMeanOp, ReduceMeanOpAdaptor adaptor, int64_t inputRank) { + ReduceMeanSemantics semantics; + semantics.keepdims = reduceMeanOp.getKeepdims(); + + if constexpr (std::is_same_v) { + auto axes = onnx_mlir::normalizeAxesChecked(std::optional(reduceMeanOp.getAxesAttr()), inputRank); + if (failed(axes)) + return failure(); + semantics.axes = std::move(*axes); + return semantics; + } + else { + if (isNoneValueLike(adaptor.getAxes())) { + if (reduceMeanOp.getNoopWithEmptyAxes() != 0) { + semantics.isIdentity = true; + return semantics; + } + + semantics.axes.reserve(inputRank); + for (int64_t axis = 0; axis < inputRank; ++axis) + semantics.axes.push_back(axis); + return semantics; + } + + auto axes = getConstantIntValues(adaptor.getAxes()); + if (failed(axes)) + return failure(); + + if (axes->empty()) { + if (reduceMeanOp.getNoopWithEmptyAxes() != 0) { + semantics.isIdentity = true; + return semantics; + } + + semantics.axes.reserve(inputRank); + for (int64_t axis = 0; axis < inputRank; ++axis) + semantics.axes.push_back(axis); + return semantics; + } + + auto normalizedAxes = normalizeAxesChecked(*axes, inputRank); + if (failed(normalizedAxes)) + return failure(); + semantics.axes = std::move(*normalizedAxes); + return semantics; + } +} + static SmallVector buildReducedAxesMask(ArrayRef axes, int64_t rank) { SmallVector reducedAxes(rank, false); for (int64_t axis : axes) { @@ -251,11 +332,13 @@ static Value squeezeReducedAxes(Value keepdimsValue, return squeezeCompute.getResult(0); } -struct ReduceMeanToSpatialCompute : OpConversionPattern { - using OpConversionPattern::OpConversionPattern; +template +struct ReduceMeanToSpatialCompute : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + using Adaptor = typename ReduceMeanOp::Adaptor; - LogicalResult matchAndRewrite(ONNXReduceMeanV13Op reduceMeanOp, - ONNXReduceMeanV13OpAdaptor adaptor, + LogicalResult matchAndRewrite(ReduceMeanOp reduceMeanOp, + Adaptor adaptor, ConversionPatternRewriter& rewriter) const override { auto inputType = dyn_cast(adaptor.getData().getType()); auto resultType = dyn_cast(reduceMeanOp.getReduced().getType()); @@ -266,10 +349,18 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern { return success(); } - auto axes = normalizeAxesChecked(std::optional(reduceMeanOp.getAxesAttr()), inputType.getRank()); - if (failed(axes)) - return failure(); - SmallVector reducedAxes = buildReducedAxesMask(*axes, inputType.getRank()); + auto semantics = getReduceMeanSemantics(reduceMeanOp, adaptor, inputType.getRank()); + if (failed(semantics)) + return rewriter.notifyMatchFailure(reduceMeanOp, "requires compile-time constant, in-range ReduceMean axes"); + if (semantics->isIdentity) { + if (inputType != resultType) + return rewriter.notifyMatchFailure( + reduceMeanOp, "noop_with_empty_axes identity requires the result type to match the input type"); + rewriter.replaceOp(reduceMeanOp, adaptor.getData()); + return success(); + } + + SmallVector reducedAxes = buildReducedAxesMask(semantics->axes, inputType.getRank()); if (reducedAxes.empty() && inputType.getRank() != 0) return failure(); @@ -289,7 +380,7 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern { Value reducedKeepdims = buildKeepdimsFromLanePackedBatch(*lanePackedKeepdims, keepdimsType, compactKeptType, reducedAxes, rewriter, loc); - if (reduceMeanOp.getKeepdims() != 0) { + if (semantics->keepdims != 0) { rewriter.replaceOp(reduceMeanOp, reducedKeepdims); return success(); } @@ -303,7 +394,7 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern { } // namespace void populateReduceMeanPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { - patterns.add(ctx); + patterns.add, ReduceMeanToSpatialCompute>(ctx); } } // namespace onnx_mlir diff --git a/validation/.gitignore b/validation/.gitignore index ad1da66..dd52fc0 100644 --- a/validation/.gitignore +++ b/validation/.gitignore @@ -8,3 +8,7 @@ networks/**/outputs networks/**/raptor networks/**/runner networks/**/simulation +networks/**/real_image_val +networks/**/*.png +networks/**/*.jpg +networks/**/*.csv diff --git a/validation/networks/resnet/resnet18_torchvision.onnx b/validation/networks/resnet/resnet18_torchvision.onnx new file mode 100644 index 0000000..7d91dbd Binary files /dev/null and b/validation/networks/resnet/resnet18_torchvision.onnx differ diff --git a/validation/operations/gen_tests.py b/validation/operations/gen_tests.py index ad107b3..4700f5f 100644 --- a/validation/operations/gen_tests.py +++ b/validation/operations/gen_tests.py @@ -1053,6 +1053,92 @@ def reducemean_large_dimension_1024(): save_model(model, "reduce_mean/large_dimension_1024", "reduce_mean_large_dimension_1024.onnx") +def make_legacy_reducemean_model(name, shape, output_shape, directory, filename, *, axes, keepdims=1, + noop_with_empty_axes=0): + """Create an opset-18 ReduceMean model that lowers to ONNXReduceMeanOp.""" + X = helper.make_tensor_value_info("X", TensorProto.FLOAT, shape) + Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, output_shape) + + initializers = [] + node_inputs = ["X", ""] + if axes is not None: + initializers.append(make_int64_initializer("axes", axes)) + node_inputs = ["X", "axes"] + + node = helper.make_node("ReduceMean", node_inputs, ["Y"], + keepdims=keepdims, noop_with_empty_axes=noop_with_empty_axes) + graph = helper.make_graph([node], name, [X], [Y], initializer=initializers) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + save_model(model, directory, filename) + + +def reducemean_legacy_axis1_keepdims_1(): + """Opset-18 ReduceMean over one positive axis, preserving rank.""" + make_legacy_reducemean_model("reducemean_legacy_axis1_keepdims_1", + [2, 3, 4], [2, 1, 4], + "reduce_mean/legacy_axis1_keepdims_1", + "reduce_mean_legacy_axis1_keepdims_1.onnx", + axes=[1], keepdims=1) + + +def reducemean_legacy_axis1_keepdims_0(): + """Opset-18 ReduceMean over one positive axis, dropping the reduced axis.""" + make_legacy_reducemean_model("reducemean_legacy_axis1_keepdims_0", + [2, 3, 4], [2, 4], + "reduce_mean/legacy_axis1_keepdims_0", + "reduce_mean_legacy_axis1_keepdims_0.onnx", + axes=[1], keepdims=0) + + +def reducemean_legacy_axes_1_2_keepdims_1(): + """Opset-18 ReduceMean over multiple positive axes.""" + make_legacy_reducemean_model("reducemean_legacy_axes_1_2_keepdims_1", + [2, 3, 4], [2, 1, 1], + "reduce_mean/legacy_axes_1_2_keepdims_1", + "reduce_mean_legacy_axes_1_2_keepdims_1.onnx", + axes=[1, 2], keepdims=1) + + +def reducemean_legacy_negative_axis(): + """Opset-18 ReduceMean using a negative axis.""" + make_legacy_reducemean_model("reducemean_legacy_negative_axis", + [2, 3, 4], [2, 3, 1], + "reduce_mean/legacy_negative_axis", + "reduce_mean_legacy_negative_axis.onnx", + axes=[-1], keepdims=1) + + +def reducemean_legacy_reduce_all_keepdims_1(): + """Opset-18 ReduceMean over all axes with the optional axes input omitted.""" + make_legacy_reducemean_model("reducemean_legacy_reduce_all_keepdims_1", + [2, 3, 4], [1, 1, 1], + "reduce_mean/legacy_reduce_all_keepdims_1", + "reduce_mean_legacy_reduce_all_keepdims_1.onnx", + axes=None, keepdims=1) + + +def reducemean_legacy_empty_axes_noop(): + """Opset-18 ReduceMean with empty axes and noop_with_empty_axes enabled.""" + X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, 4]) + Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [3, 4]) + axes = make_int64_initializer("axes", []) + reduce = helper.make_node("ReduceMean", ["X", "axes"], ["R"], + keepdims=1, noop_with_empty_axes=1) + relu = helper.make_node("Relu", ["R"], ["Y"]) + graph = helper.make_graph([reduce, relu], "reducemean_legacy_empty_axes_noop", [X], [Y], initializer=[axes]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + save_model(model, "reduce_mean/legacy_empty_axes_noop", "reduce_mean_legacy_empty_axes_noop.onnx") + + +def reducemean_legacy_nchw_spatial(): + """Opset-18 ReduceMean over H and W on an NCHW tensor.""" + make_legacy_reducemean_model("reducemean_legacy_nchw_spatial", + [1, 3, 5, 5], [1, 3, 1, 1], + "reduce_mean/legacy_nchw_spatial", + "reduce_mean_legacy_nchw_spatial.onnx", + axes=[2, 3], keepdims=1) + + # --------------------------------------------------------------------------- # Relu tests # --------------------------------------------------------------------------- @@ -1974,6 +2060,13 @@ if __name__ == "__main__": reducemean_4d_spatial_keepdims_0() reducemean_channel_axis_nchw() reducemean_large_dimension_1024() + reducemean_legacy_axis1_keepdims_1() + reducemean_legacy_axis1_keepdims_0() + reducemean_legacy_axes_1_2_keepdims_1() + reducemean_legacy_negative_axis() + reducemean_legacy_reduce_all_keepdims_1() + reducemean_legacy_empty_axes_noop() + reducemean_legacy_nchw_spatial() print("\nGenerating Relu tests:") relu_basic() diff --git a/validation/operations/reduce_mean/legacy_axes_1_2_keepdims_1/reduce_mean_legacy_axes_1_2_keepdims_1.onnx b/validation/operations/reduce_mean/legacy_axes_1_2_keepdims_1/reduce_mean_legacy_axes_1_2_keepdims_1.onnx new file mode 100644 index 0000000..5362868 Binary files /dev/null and b/validation/operations/reduce_mean/legacy_axes_1_2_keepdims_1/reduce_mean_legacy_axes_1_2_keepdims_1.onnx differ diff --git a/validation/operations/reduce_mean/legacy_axis1_keepdims_0/reduce_mean_legacy_axis1_keepdims_0.onnx b/validation/operations/reduce_mean/legacy_axis1_keepdims_0/reduce_mean_legacy_axis1_keepdims_0.onnx new file mode 100644 index 0000000..7f66f5f Binary files /dev/null and b/validation/operations/reduce_mean/legacy_axis1_keepdims_0/reduce_mean_legacy_axis1_keepdims_0.onnx differ diff --git a/validation/operations/reduce_mean/legacy_axis1_keepdims_1/reduce_mean_legacy_axis1_keepdims_1.onnx b/validation/operations/reduce_mean/legacy_axis1_keepdims_1/reduce_mean_legacy_axis1_keepdims_1.onnx new file mode 100644 index 0000000..98f3cce Binary files /dev/null and b/validation/operations/reduce_mean/legacy_axis1_keepdims_1/reduce_mean_legacy_axis1_keepdims_1.onnx differ diff --git a/validation/operations/reduce_mean/legacy_empty_axes_noop/reduce_mean_legacy_empty_axes_noop.onnx b/validation/operations/reduce_mean/legacy_empty_axes_noop/reduce_mean_legacy_empty_axes_noop.onnx new file mode 100644 index 0000000..e17969b Binary files /dev/null and b/validation/operations/reduce_mean/legacy_empty_axes_noop/reduce_mean_legacy_empty_axes_noop.onnx differ diff --git a/validation/operations/reduce_mean/legacy_nchw_spatial/reduce_mean_legacy_nchw_spatial.onnx b/validation/operations/reduce_mean/legacy_nchw_spatial/reduce_mean_legacy_nchw_spatial.onnx new file mode 100644 index 0000000..3efbd4c Binary files /dev/null and b/validation/operations/reduce_mean/legacy_nchw_spatial/reduce_mean_legacy_nchw_spatial.onnx differ diff --git a/validation/operations/reduce_mean/legacy_negative_axis/reduce_mean_legacy_negative_axis.onnx b/validation/operations/reduce_mean/legacy_negative_axis/reduce_mean_legacy_negative_axis.onnx new file mode 100644 index 0000000..d04ab26 Binary files /dev/null and b/validation/operations/reduce_mean/legacy_negative_axis/reduce_mean_legacy_negative_axis.onnx differ diff --git a/validation/operations/reduce_mean/legacy_reduce_all_keepdims_1/reduce_mean_legacy_reduce_all_keepdims_1.onnx b/validation/operations/reduce_mean/legacy_reduce_all_keepdims_1/reduce_mean_legacy_reduce_all_keepdims_1.onnx new file mode 100644 index 0000000..527b4de Binary files /dev/null and b/validation/operations/reduce_mean/legacy_reduce_all_keepdims_1/reduce_mean_legacy_reduce_all_keepdims_1.onnx differ