[ET-VK] Modernize argmax and argmin

ssjia · SS-JIA · commit 7e9445b78d00 · 2026-03-17T21:52:40.000-04:00
Modernize argmax and argmin to support ANY_STORAGE via the add_reduce_per_row_node dispatch path. Buffer shader uses BufferMetadata with indexing.glslh. Custom workgroup sizing retained for cooperative row-reduction algorithm with shared memory. Pull Request resolved: #18058 ghstack-source-id: 353546687 @exported-using-ghexport Differential Revision: [D95970165](https://our.internmc.facebook.com/intern/diff/D95970165/)
diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py
@@ -705,7 +705,7 @@ def register_reduce_cpp_ops():
 )
 def register_argreduce_cpp_ops():
     return OpFeatures(
-        inputs_storage=utils.ANY_TEXTURE,
+        inputs_storage=utils.ANY_STORAGE,
         inputs_dtypes=utils.FP_T,
         outputs_dtypes=utils.INT_T,
         supports_resize=True,

Original file line number	Diff line number	Diff line change
`@@ -705,7 +705,7 @@ def register_reduce_cpp_ops():`
`705`	`705`	`)`
`706`	`706`	`def register_argreduce_cpp_ops():`
`707`	`707`	`return OpFeatures(`
`708`		`- inputs_storage=utils.ANY_TEXTURE,`
	`708`	`+ inputs_storage=utils.ANY_STORAGE,`
`709`	`709`	`inputs_dtypes=utils.FP_T,`
`710`	`710`	`outputs_dtypes=utils.INT_T,`
`711`	`711`	`supports_resize=True,`