Faster cpu ops (#1434)

* faster binary and cleaner copy * use recursive template for other ops * more cleanup * fix from cleanup * more clean * fix binary * use contiguous iterator * add 3d * nits * fix * fix? * fix * fix rebase
ml-explore · Sep 26, 2024 · 5b6f38d · 5b6f38d
1 parent 0b4a586
commit 5b6f38d
Show file tree

Hide file tree

Showing 12 changed files with 577 additions and 1,334 deletions.
diff --git a/mlx/backend/common/binary.h b/mlx/backend/common/binary.h
diff --git a/mlx/backend/common/binary_two.h b/mlx/backend/common/binary_two.h
diff --git a/mlx/backend/common/common.cpp b/mlx/backend/common/common.cpp
@@ -156,8 +156,7 @@ std::pair<bool, std::vector<size_t>> Reshape::prepare_reshape(
   }
 
   // Firstly let's collapse all the contiguous dimensions of the input
-  auto [shape, _strides] = collapse_contiguous_dims(in);
-  auto& strides = _strides[0];
+  auto [shape, strides] = collapse_contiguous_dims(in);
 
   // If shapes fit exactly in the contiguous dims then no copy is necessary so
   // let's check.