From 88187b784bde5239c730b919d01fe85baa32ba88 Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Fri, 23 Jun 2023 19:04:13 +0100
Subject: [PATCH] Also optimize the contiguous case for the binary cuda
 kernels.

---
 kernels/src/binary_op_macros.cuh | 60 +++++++++++++++++++++++++-------
 1 file changed, 48 insertions(+), 12 deletions(-)

diff --git a/kernels/src/binary_op_macros.cuh b/kernels/src/binary_op_macros.cuh
index 60531b4d..219ee09c 100644
--- a/kernels/src/binary_op_macros.cuh
+++ b/kernels/src/binary_op_macros.cuh
@@ -12,18 +12,54 @@ extern "C" __global__ void FN_NAME( \
     const size_t *dims = dims_and_strides; \
     const size_t *lhs_strides = dims_and_strides + 1 * num_dims; \
     const size_t *rhs_strides = dims_and_strides + 2 * num_dims; \
-    for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
-        unsigned int tmp_i = i; \
-        unsigned int lhs_i = 0; \
-        unsigned int rhs_i = 0; \
-        for (int d = num_dims - 1; d >= 0; d--) { \
-            unsigned int i_dim = tmp_i % dims[d]; \
-            lhs_i += i_dim * lhs_strides[d]; \
-            rhs_i += i_dim * rhs_strides[d]; \
-            tmp_i /= dims[d]; \
+    bool lhs_cont = is_contiguous(num_dims, dims, lhs_strides); \
+    bool rhs_cont = is_contiguous(num_dims, dims, rhs_strides); \
+    if (lhs_cont && rhs_cont) { \
+        for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
+            TYPENAME x = lhs ? lhs[i] : out[i]; \
+            TYPENAME y = rhs ? rhs[i] : out[i]; \
+            out[i] = FUNC; \
+        } \
+    } else if (lhs_cont) { \
+        for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
+            unsigned int tmp_i = i; \
+            unsigned int rhs_i = 0; \
+            for (int d = num_dims - 1; d >= 0; d--) { \
+                unsigned int i_dim = tmp_i % dims[d]; \
+                rhs_i += i_dim * rhs_strides[d]; \
+                tmp_i /= dims[d]; \
+            } \
+            TYPENAME x = lhs ? lhs[i] : out[i]; \
+            TYPENAME y = rhs ? rhs[rhs_i] : out[i]; \
+            out[i] = FUNC; \
+        } \
+    } else if (rhs_cont) { \
+        for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
+            unsigned int tmp_i = i; \
+            unsigned int lhs_i = 0; \
+            for (int d = num_dims - 1; d >= 0; d--) { \
+                unsigned int i_dim = tmp_i % dims[d]; \
+                lhs_i += i_dim * lhs_strides[d]; \
+                tmp_i /= dims[d]; \
+            } \
+            TYPENAME x = lhs ? lhs[lhs_i] : out[i]; \
+            TYPENAME y = rhs ? rhs[i] : out[i]; \
+            out[i] = FUNC; \
+        } \
+    } else { \
+        for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
+            unsigned int tmp_i = i; \
+            unsigned int lhs_i = 0; \
+            unsigned int rhs_i = 0; \
+            for (int d = num_dims - 1; d >= 0; d--) { \
+                unsigned int i_dim = tmp_i % dims[d]; \
+                lhs_i += i_dim * lhs_strides[d]; \
+                rhs_i += i_dim * rhs_strides[d]; \
+                tmp_i /= dims[d]; \
+            } \
+            TYPENAME x = lhs ? lhs[lhs_i] : out[i]; \
+            TYPENAME y = rhs ? rhs[rhs_i] : out[i]; \
+            out[i] = FUNC; \
         } \
-        TYPENAME x = lhs ? lhs[lhs_i] : out[i]; \
-        TYPENAME y = rhs ? rhs[rhs_i] : out[i]; \
-        out[i] = FUNC; \
     } \
 } \