SpatialSubOp

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -124,6 +124,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXMatMulOp>();
   target.addIllegalOp<ONNXTransposeOp>();
   target.addIllegalOp<ONNXAddOp>();
+  target.addIllegalOp<ONNXSubOp>();
   target.addIllegalOp<ONNXDivOp>();
   target.addIllegalOp<ONNXMulOp>();
   target.addIllegalOp<ONNXGemmOp>();

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

@@ -189,6 +189,7 @@ struct DivToSpatialCompute : OpConversionPattern<ONNXDivOp> {
 
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXAddOp, spatial::SpatVAddOp>>(ctx);
+  patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXMulOp, spatial::SpatVMulOp>>(ctx);
   patterns.add<DivToSpatialCompute>(ctx);
 }

src/PIM/Conversion/SpatialToPim/SpatialToPim.td

@@ -27,6 +27,12 @@ def spatToPimVVAdd : Pat<
     (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
 >;
 
+def spatToPimVVSub : Pat<
+  (SpatVSubOp:$srcOpRes $a, $b),
+  (PimVVSubOp $a, $b,
+    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
+>;
+
 def spatToPimVVMul : Pat<
   (SpatVMulOp:$srcOpRes $a, $b),
   (PimVVMulOp $a, $b,

src/PIM/Dialect/Spatial/Spatial.td

@@ -257,6 +257,25 @@ def SpatVAddOp : SpatOp<"vadd", []> {
   }];
 }
 
+def SpatVSubOp : SpatOp<"vsub", []> {
+  let summary = "Element-wise subtraction between two tensors; rhs must match lhs or be 1x1";
+
+  let arguments = (ins
+    SpatTensor:$lhs,
+    SpatTensor:$rhs
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+
+  let assemblyFormat = [{
+    $lhs `,` $rhs attr-dict `:` `(` type($lhs) `,` type($rhs) `)` `->` type($output)
+  }];
+}
+
 def SpatVMulOp : SpatOp<"vmul", []> {
   let summary = "Element-wise multiplication between two tensors; rhs must match lhs or be 1x1";
 

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -254,6 +254,12 @@ LogicalResult SpatVAddOp::verify() {
   return OpTrait::impl::verifySameOperandsAndResultType(*this);
 }
 
+LogicalResult SpatVSubOp::verify() {
+  if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
+    return failure();
+  return OpTrait::impl::verifySameOperandsAndResultType(*this);
+}
+
 LogicalResult SpatVMaxOp::verify() {
   if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
     return failure();

validation/operations/gen_tests.py

@@ -1549,6 +1549,82 @@ def add_leading_dimension_broadcast():
     save_model(model, "add/leading_dimension_broadcast", "add_leading_dimension_broadcast.onnx")
 
 
+# ---------------------------------------------------------------------------
+# Sub tests
+# ---------------------------------------------------------------------------
+
+def sub_basic():
+    """Elementwise Sub on two runtime inputs with identical shapes."""
+    A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [4, 8])
+    B = helper.make_tensor_value_info("B", TensorProto.FLOAT, [4, 8])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8])
+    node = helper.make_node("Sub", ["A", "B"], ["Y"])
+    graph = helper.make_graph([node], "sub_basic", [A, B], [Y])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "sub/basic", "sub_basic.onnx")
+
+
+def sub_broadcast_row():
+    """Elementwise Sub with a broadcast row-vector RHS constant."""
+    A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [4, 8])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8])
+    B = numpy_helper.from_array(np.random.default_rng(103).uniform(-1, 1, (8,)).astype(np.float32), name="B")
+    node = helper.make_node("Sub", ["A", "B"], ["Y"])
+    graph = helper.make_graph([node], "sub_broadcast_row", [A], [Y], initializer=[B])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "sub/broadcast_row", "sub_broadcast_row.onnx")
+
+
+def sub_constant_lhs_broadcast():
+    """Elementwise Sub with a broadcast constant LHS to preserve operand order."""
+    B = helper.make_tensor_value_info("B", TensorProto.FLOAT, [4, 8])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8])
+    A = numpy_helper.from_array(np.random.default_rng(104).uniform(-1, 1, (8,)).astype(np.float32), name="A")
+    node = helper.make_node("Sub", ["A", "B"], ["Y"])
+    graph = helper.make_graph([node], "sub_constant_lhs_broadcast", [B], [Y], initializer=[A])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "sub/constant_lhs_broadcast", "sub_constant_lhs_broadcast.onnx")
+
+
+def sub_after_gemm():
+    """Gemm followed by Sub with a broadcast constant vector."""
+    B, K, N = 4, 64, 32
+    rng = np.random.default_rng(105)
+    W = numpy_helper.from_array(rng.uniform(-1, 1, (K, N)).astype(np.float32), name="W")
+    C = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="C")
+    S = numpy_helper.from_array(rng.uniform(-1, 1, (N,)).astype(np.float32), name="S")
+    A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [B, K])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [B, N])
+    gemm = helper.make_node("Gemm", ["A", "W", "C"], ["G"])
+    sub = helper.make_node("Sub", ["G", "S"], ["Y"])
+    graph = helper.make_graph([gemm, sub], "sub_after_gemm", [A], [Y], initializer=[W, C, S])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "sub/after_gemm", "sub_after_gemm.onnx")
+
+
+def sub_channel_broadcast_1024():
+    """Elementwise Sub with 1024-channel constant broadcasting."""
+    A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [1, 1024, 1, 1])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1024, 1, 1])
+    B = numpy_helper.from_array(
+        np.random.default_rng(106).uniform(-1, 1, (1, 1024, 1, 1)).astype(np.float32), name="B")
+    node = helper.make_node("Sub", ["A", "B"], ["Y"])
+    graph = helper.make_graph([node], "sub_channel_broadcast_1024", [A], [Y], initializer=[B])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "sub/channel_broadcast_1024", "sub_channel_broadcast_1024.onnx")
+
+
+def sub_leading_dimension_broadcast():
+    """Elementwise Sub with trailing-dimension constant broadcasting."""
+    A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [2, 3, 4])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3, 4])
+    B = numpy_helper.from_array(np.random.default_rng(107).uniform(-1, 1, (4,)).astype(np.float32), name="B")
+    node = helper.make_node("Sub", ["A", "B"], ["Y"])
+    graph = helper.make_graph([node], "sub_leading_dimension_broadcast", [A], [Y], initializer=[B])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "sub/leading_dimension_broadcast", "sub_leading_dimension_broadcast.onnx")
+
+
 # ---------------------------------------------------------------------------
 # Mul tests
 # ---------------------------------------------------------------------------
@@ -1844,6 +1920,14 @@ if __name__ == "__main__":
     add_channel_broadcast_1024()
     add_leading_dimension_broadcast()
 
+    print("\nGenerating Sub tests:")
+    sub_basic()
+    sub_broadcast_row()
+    sub_constant_lhs_broadcast()
+    sub_after_gemm()
+    sub_channel_broadcast_1024()
+    sub_leading_dimension_broadcast()
+
     print("\nGenerating Mul tests:")
     mul_basic()
     mul_scalar_constant()

validation/raptor.py

@@ -41,7 +41,7 @@ def _format_command(cmd):
 
 def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
                         crossbar_size, crossbar_count, core_count=None, pim_merge_scheduler="peft",
-                        cwd=None, verbose=False, reporter=None, timeout_sec=None):
+                        pim_memory_report="none", cwd=None, verbose=False, reporter=None, timeout_sec=None):
     # Define the arguments, with the possibility to set crossbar size and count
     args = [
         network_path,
@@ -55,6 +55,8 @@ def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
     ]
     if core_count is not None:
         args.append(f"--core-count={core_count}")
+    if pim_memory_report != "none":
+        args.append(f"--pim-memory-report={pim_memory_report}")
     if verbose:
         args.append("--enable-timing")
 

validation/validate.py

@@ -75,6 +75,8 @@ def main():
                     help="Core count to pass to Raptor. Required for PIM validation.")
     ap.add_argument("--pim-merge-scheduler", choices=("peft"), default="peft",
                     help="Scheduler used by the Spatial merge-compute-nodes pass.")
+    ap.add_argument("--pim-memory-report", choices=("none", "summary", "full"), default="none",
+                    help="Emit a human-readable PIM memory planning report during codegen.")
     ap.add_argument("--command-timeout-seconds", type=float, default=1000000.0,
                     help="Per-subprocess timeout in seconds for compiler, runner, and simulator commands.")
     ap.add_argument("--clean", action="store_true",
@@ -142,7 +144,7 @@ def main():
             result = validate_network(
                 onnx_path, a.raptor_path, a.onnx_include_dir, simulator_dir,
                 crossbar_size=a.crossbar_size, crossbar_count=a.crossbar_count, core_count=a.core_count,
-                pim_merge_scheduler=a.pim_merge_scheduler,
+                pim_merge_scheduler=a.pim_merge_scheduler, pim_memory_report=a.pim_memory_report,
                 command_timeout_seconds=a.command_timeout_seconds,
                 threshold=a.threshold,
                 seed=a.seed,

validation/validate_one.py

@@ -289,7 +289,7 @@ def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1
 
 def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                      simulator_dir, crossbar_size=64, crossbar_count=8, core_count=None,
-                     pim_merge_scheduler="peft", threshold=1e-3,
+                     pim_merge_scheduler="peft", pim_memory_report="none", threshold=1e-3,
                      seed=0, reporter=None, model_index=1, model_total=1, verbose=False,
                      command_timeout_seconds=60.0, mode=MODE_FULL):
     network_onnx_path = Path(network_onnx_path).resolve()
@@ -343,6 +343,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                 pim_pass_timings = compile_with_raptor(
                     network_mlir_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
                     core_count=core_count, pim_merge_scheduler=pim_merge_scheduler,
+                    pim_memory_report=pim_memory_report,
                     cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
                 print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
                 reporter.advance()
@@ -382,6 +383,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             pim_pass_timings = compile_with_raptor(
                 network_mlir_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
                 core_count=core_count, pim_merge_scheduler=pim_merge_scheduler,
+                pim_memory_report=pim_memory_report,
                 cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
             print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
             reporter.advance()