[ET-VK] Fix negative dim in normalize_to_dim_index

jorgep31415 · jorgep31415 · commit 8b63ea8d9422 · 2024-09-05T15:44:59.000-07:00
Our `cat` implementation currently fails for negative dim. Differential Revision: [D62270925](https://our.internmc.facebook.com/intern/diff/D62270925/) [ghstack-poisoned]
diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h b/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h
@@ -32,7 +32,8 @@ constexpr DimIndex kChannel4D = DimIndex::DIM_3RD_LAST;
 constexpr DimIndex kBatch4D = DimIndex::DIM_4TH_LAST;
 
 inline DimIndex normalize_to_dim_index(const api::vTensor& v_in, int32_t dim) {
-  return static_cast<DimIndex>(dim - v_in.dim());
+  return dim < 0 ? static_cast<DimIndex>(dim)
+                 : static_cast<DimIndex>(dim - v_in.dim());
 }
 
 /*

Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,8 @@ constexpr DimIndex kChannel4D = DimIndex::DIM_3RD_LAST;`
`32`	`32`	`constexpr DimIndex kBatch4D = DimIndex::DIM_4TH_LAST;`
`33`	`33`
`34`	`34`	`inline DimIndex normalize_to_dim_index(const api::vTensor& v_in, int32_t dim) {`
`35`		`- return static_cast<DimIndex>(dim - v_in.dim());`
	`35`	`+ return dim < 0 ? static_cast<DimIndex>(dim)`
	`36`	`+ : static_cast<DimIndex>(dim - v_in.dim());`
`36`	`37`	`}`
`37`	`38`
`38`	`39`	`/*`