blazingly faster

src/PIM/Dialect/Spatial/Spatial.td

@@ -49,6 +49,7 @@ class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
 }
 
 def SpatGraphCompute : SpatComputeLikeBase<"graph_compute"> {
+  let hasCanonicalizer = 1;
   let extraClassDeclaration = [{
     std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
     std::optional<::mlir::BlockArgument> getInputArgument(unsigned idx);
@@ -186,7 +187,8 @@ def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]
   let summary = "Temporary scheduled payload derivation placeholder";
 
   let arguments = (ins
-    Variadic<SpatTensor>:$sources
+    Variadic<SpatTensor>:$sources,
+    OptionalAttr<I64Attr>:$specialization_count
   );
 
   let results = (outs
@@ -199,6 +201,22 @@ def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]
   let hasCustomAssemblyFormat = 1;
 }
 
+def SpatDeferredSourceSelectOp : SpatOp<"deferred_source_select", []> {
+  let summary = "Select a deferred tensor source with a statically analyzable index";
+
+  let arguments = (ins
+    Index:$selector,
+    Variadic<SpatTensor>:$sources
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 def SpatExtractRowsOp : SpatOp<"extract_rows", []> {
   let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors";