Fix pea.sparse on IPU, SDK 3.2 (#16)

* Work around the error 'No default opset version defined for domain 'ai.graphcore.pea'' * Add the correct cache keys to StaticSparseMatmul, to prevent spurious cache hits * Move away from ai.graphcore.pea generally, as we don't know when ONNX shape inference logic will cause an error
graphcore-research · Apr 28, 2023 · 08311ee · 08311ee
1 parent f07a396
commit 08311ee
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 10 deletions.
diff --git a/poptorch_experimental_addons/_impl/core.py b/poptorch_experimental_addons/_impl/core.py
@@ -52,7 +52,7 @@ def autograd_proxy(fwd: Tensor, proxy: Tensor) -> Tensor:
         (y,) = poptorch.custom_op(
             [fwd, proxy],
             name="AutogradProxy",
-            domain="ai.graphcore.pea",
+            domain="ai.graphcore",
             domain_version=1,
             example_outputs=[fwd],
         )
@@ -93,7 +93,7 @@ def distance_matrix(tensor1: Tensor, tensor2: Tensor, p: int) -> Tensor:
         (y,) = poptorch.custom_op(
             name=f"L{p}Distance",
             domain_version=1,
-            domain="ai.graphcore.pea",
+            domain="ai.graphcore",
             inputs=[tensor1, tensor2],
             example_outputs=[
                 torch.zeros(

diff --git a/poptorch_experimental_addons/_impl/sparse.py b/poptorch_experimental_addons/_impl/sparse.py
@@ -154,7 +154,7 @@ def block_coo_spmm_ipu(sparse: Tensor, dense: Tensor, mode: str) -> Tensor:
     (y,) = poptorch.custom_op(
         [dense],
         name="StaticSparseMatmul",
-        domain="ai.graphcore.pea",
+        domain="ai.graphcore",
         domain_version=1,
         example_outputs=[
             torch.zeros(output_shape, dtype=dense.dtype, device=dense.device)

diff --git a/poptorch_experimental_addons/cpp/autograd_proxy.cpp b/poptorch_experimental_addons/cpp/autograd_proxy.cpp
@@ -97,8 +97,8 @@ struct Opx : popart::popx::Opx {
     }
 };
 
-const popart::OperatorIdentifier Op::ID = {"ai.graphcore.pea", "AutogradProxy", 1};
-const popart::OperatorIdentifier GradOp::ID = {"ai.graphcore.pea", "AutogradProxyGrad", 1};
+const popart::OperatorIdentifier Op::ID = {"ai.graphcore", "AutogradProxy", 1};
+const popart::OperatorIdentifier GradOp::ID = {"ai.graphcore", "AutogradProxyGrad", 1};
 popart::OpDefinition::DataTypes T = {popart::DataType::FLOAT16, popart::DataType::FLOAT};
 popart::OpCreator<Op> opCreator(
     {{Op::ID,

diff --git a/poptorch_experimental_addons/cpp/distance_matrix.cpp b/poptorch_experimental_addons/cpp/distance_matrix.cpp
@@ -234,10 +234,10 @@ poplar::Tensor l2distancegrad(poplar::Graph& graph,
     return castMaybe(graph, grad, a.elementType(), prog);
 }
 
-const popart::OperatorIdentifier L1DistanceId = {"ai.graphcore.pea", "L1Distance", 1};
-const popart::OperatorIdentifier L2DistanceId = {"ai.graphcore.pea", "L2Distance", 1};
-const popart::OperatorIdentifier L1DistanceGradId = {"ai.graphcore.pea", "L1DistanceGrad", 1};
-const popart::OperatorIdentifier L2DistanceGradId = {"ai.graphcore.pea", "L2DistanceGrad", 1};
+const popart::OperatorIdentifier L1DistanceId = {"ai.graphcore", "L1Distance", 1};
+const popart::OperatorIdentifier L2DistanceId = {"ai.graphcore", "L2Distance", 1};
+const popart::OperatorIdentifier L1DistanceGradId = {"ai.graphcore", "L1DistanceGrad", 1};
+const popart::OperatorIdentifier L2DistanceGradId = {"ai.graphcore", "L2DistanceGrad", 1};
 
 class L1DistanceOp;
 class L1DistanceGradOpx;

diff --git a/poptorch_experimental_addons/cpp/static_spmm.cpp b/poptorch_experimental_addons/cpp/static_spmm.cpp
@@ -15,6 +15,7 @@
 #pragma GCC diagnostic ignored "-Wunused-parameter"
 #include <popart/op.hpp>
 #include <popart/opmanager.hpp>
+#include <popart/opserialiser.hpp>
 #include <popart/popx/opx.hpp>
 #include <popart/popx/opxmanager.hpp>
 #pragma GCC diagnostic pop
@@ -136,6 +137,32 @@ struct CustomOp : popart::Op {
                 ? std::vector<int64_t>{static_cast<int>(matrix.numRows), input.dim(1)}
                 : std::vector<int64_t>{input.dim(0), static_cast<int>(matrix.numColumns)}};
     }
+    void appendAttributes(popart::OpSerialiserBase& os) const final {
+        popart::Op::appendAttributes(os);
+        appendLocalAttributes(os);
+    }
+    void appendOutlineAttributes(popart::OpSerialiserBase& os) const final {
+        popart::Op::appendOutlineAttributes(os);
+        appendLocalAttributes(os);
+    }
+
+   private:
+    template <class T>
+    static std::string vectorToString(const std::vector<T>& v) {
+        std::ostringstream str;
+        std::copy(v.begin(), v.end(), std::ostream_iterator<T>(str, " "));
+        return str.str();
+    }
+    void appendLocalAttributes(popart::OpSerialiserBase& os) const {
+        os.appendAttribute("mode", mode);
+        os.appendAttribute("numRows", matrix.numRows);
+        os.appendAttribute("numColumns", matrix.numColumns);
+        os.appendAttribute("blockSizeRows", matrix.getBlockDimensions()[0]);
+        os.appendAttribute("blockSizeColumns", matrix.getBlockDimensions()[1]);
+        os.appendAttribute("rowIndices", vectorToString(matrix.rowIndices));
+        os.appendAttribute("columnIndices", vectorToString(matrix.columnIndices));
+        os.appendAttribute("nzValues", vectorToString(matrix.nzValues));
+    }
 };
 
 struct CustomOpx : popart::popx::Opx {
@@ -153,7 +180,9 @@ struct CustomOpx : popart::popx::Opx {
     }
 };
 
-const popart::OperatorIdentifier CustomOp::ID = {"ai.graphcore.pea", "StaticSparseMatmul", 1};
+// We cannot use "ai.graphcore.pea", since shape inference tries to call
+// `Ir::getDefaultOpsetVersion` which cannot be extended to custom domains
+const popart::OperatorIdentifier CustomOp::ID = {"ai.graphcore", "StaticSparseMatmul", 1};
 popart::OpDefinition::DataTypes T = {popart::DataType::FLOAT16, popart::DataType::FLOAT};
 popart::OpCreator<CustomOp> opCreator(
     {{CustomOp::ID,