add support for softmax, resize, split, gather

src/PIM/Dialect/Spatial/Spatial.td

@@ -272,6 +272,22 @@ def SpatSigmoidOp : SpatOp<"sigmoid", []> {
   }];
 }
 
+def SpatSoftmaxOp : SpatOp<"softmax", []> {
+  let summary = "Softmax over the full input tensor slice";
+
+  let arguments = (ins
+    SpatTensor:$input
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let assemblyFormat = [{
+    `(` $input `)` attr-dict `:` type($input) `->` type($output)
+  }];
+}
+
 def SpatReluOp : SpatOp<"relu", []> {
   let summary = "Element-wise ReLU activation";
 

src/PIM/Dialect/Spatial/Transforms/MergeComputeNode/MergeComputeNodePass.cpp

@@ -202,9 +202,9 @@ private:
         rewriter.clone(op, mapper);
     }
 
-    for (auto users : oldWeightedCompute->getUsers())
-      if (auto funcRet = dyn_cast<func::ReturnOp>(users))
-        funcRet.setOperand(0, newWeightedCompute.getResult(0));
+    for (auto& use : llvm::make_early_inc_range(oldWeightedCompute->getUses()))
+      if (isa<func::ReturnOp>(use.getOwner()))
+        use.assign(newWeightedCompute.getResult(0));
 
     oldToNewComputeMap.insert({oldWeightedCompute, newWeightedCompute});
     return {cast<SpatWeightedCompute>(newWeightedCompute), computeValueResults};