ReduceMean + resnet

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -140,6 +140,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXResizeOp>();
   target.addIllegalOp<ONNXSliceOp>();
   target.addIllegalOp<ONNXLRNOp>();
+  target.addIllegalOp<ONNXReduceMeanOp>();
   target.addIllegalOp<ONNXReduceMeanV13Op>();
   target.addIllegalOp<ONNXSplitOp>();
 

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp

@@ -6,6 +6,8 @@
 
 #include <algorithm>
 #include <numeric>
+#include <optional>
+#include <type_traits>
 
 #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<int64_t> axes;
+  int64_t keepdims = 1;
+  bool isIdentity = false;
+};
+
+static bool isNoneValueLike(Value value) { return isa_and_nonnull<ONNXNoneOp>(value.getDefiningOp()); }
+
+static FailureOr<SmallVector<int64_t>> getConstantIntValues(Value value) {
+  auto denseAttr = dyn_cast_or_null<DenseIntElementsAttr>(getHostConstDenseElementsAttr(value));
+  if (!denseAttr)
+    return failure();
+  return SmallVector<int64_t>(denseAttr.getValues<int64_t>().begin(), denseAttr.getValues<int64_t>().end());
+}
+
+static FailureOr<SmallVector<int64_t>> normalizeAxesChecked(ArrayRef<int64_t> axes, int64_t rank) {
+  SmallVector<int64_t> 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 <typename ReduceMeanOp, typename ReduceMeanOpAdaptor>
+static FailureOr<ReduceMeanSemantics>
+getReduceMeanSemantics(ReduceMeanOp reduceMeanOp, ReduceMeanOpAdaptor adaptor, int64_t inputRank) {
+  ReduceMeanSemantics semantics;
+  semantics.keepdims = reduceMeanOp.getKeepdims();
+
+  if constexpr (std::is_same_v<ReduceMeanOp, ONNXReduceMeanV13Op>) {
+    auto axes = onnx_mlir::normalizeAxesChecked(std::optional<ArrayAttr>(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<bool> buildReducedAxesMask(ArrayRef<int64_t> axes, int64_t rank) {
   SmallVector<bool> reducedAxes(rank, false);
   for (int64_t axis : axes) {
@@ -251,11 +332,13 @@ static Value squeezeReducedAxes(Value keepdimsValue,
   return squeezeCompute.getResult(0);
 }
 
-struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
-  using OpConversionPattern::OpConversionPattern;
+template <typename ReduceMeanOp>
+struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {
+  using OpConversionPattern<ReduceMeanOp>::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<RankedTensorType>(adaptor.getData().getType());
     auto resultType = dyn_cast<RankedTensorType>(reduceMeanOp.getReduced().getType());
@@ -266,10 +349,18 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
       return success();
     }
 
-    auto axes = normalizeAxesChecked(std::optional<ArrayAttr>(reduceMeanOp.getAxesAttr()), inputType.getRank());
-    if (failed(axes))
-      return failure();
-    SmallVector<bool> 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<bool> reducedAxes = buildReducedAxesMask(semantics->axes, inputType.getRank());
     if (reducedAxes.empty() && inputType.getRank() != 0)
       return failure();
 
@@ -289,7 +380,7 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
     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<ONNXReduceMeanV13Op> {
 } // namespace
 
 void populateReduceMeanPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.add<ReduceMeanToSpatialCompute>(ctx);
+  patterns.add<ReduceMeanToSpatialCompute<ONNXReduceMeanV13Op>, ReduceMeanToSpatialCompute<ONNXReduceMeanOp>>(ctx);
 }
 
 } // namespace onnx_mlir

validation/.gitignore

@@ -8,3 +8,7 @@ networks/**/outputs
 networks/**/raptor
 networks/**/runner
 networks/**/simulation
+networks/**/real_image_val
+networks/**/*.png
+networks/**/*.jpg
+networks/**/*.csv

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()