PyOpenCL indexing 3D arrays inside kernel code

Question

I am using PyOpenCL to process images in Python and to send a 3D numpy array (height x width x 4) to the kernel. I am having trouble indexing the 3D array inside the kernel code. For now I am only able to copy the whole input array to the output. The current code looks like this, where img is the image with img.shape = (320, 512, 4):

__kernel void part1(__global float* img, __global float* results)
{
    unsigned int x = get_global_id(0);
    unsigned int y = get_global_id(1);
    unsigned int z = get_global_id(2);

    int index = x + 320*y + 320*512*z;

    results[index] = img[index];
}

However, I do not quite understand how this work. For example, how do I index the Python equivalent of img[1, 2, 3] inside this kernel? And further, which index should be used into results for storing some item if I want it to be on the position results[1, 2, 3] in the numpy array when I get the results back to Python?

To run this I am using this Python code:

import pyopencl as cl
import numpy as np

class OpenCL:
def __init__(self):
    self.ctx = cl.create_some_context()
    self.queue = cl.CommandQueue(self.ctx)

def loadProgram(self, filename):
    f = open(filename, 'r')
    fstr = "".join(f.readlines())
    self.program = cl.Program(self.ctx, fstr).build()

def opencl_energy(self, img):
    mf = cl.mem_flags

    self.img = img.astype(np.float32)

    self.img_buf = cl.Buffer(self.ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=self.img)
    self.dest_buf = cl.Buffer(self.ctx, mf.WRITE_ONLY, self.img.nbytes)

    self.program.part1(self.queue, self.img.shape, None, self.img_buf, self.dest_buf)
    c = np.empty_like(self.img)
    cl.enqueue_read_buffer(self.queue, self.dest_buf, c).wait()
    return c

example = OpenCL()
example.loadProgram("get_energy.cl")
image = np.random.rand(320, 512, 4)
image = image.astype(np.float32)
results = example.opencl_energy(image)
print("All items are equal:", (results==image).all())

Answer 1

Update: The OpenCL docs state (in 3.5), that

"Memory objects are categorized into two types: buffer objects, and image objects. A buffer
object stores a one-dimensional collection of elements whereas an image object is used to store a
two- or three- dimensional texture, frame-buffer or image." 

so, a buffer always is linear, or linearized as you can see from my sample below.

import pyopencl as cl
import numpy as np


h_a = np.arange(27).reshape((3,3,3)).astype(np.float32)

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)

mf = cl.mem_flags
d_a  = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=h_a)

prg = cl.Program(ctx, """
__kernel void p(__global const float *d_a) {
  printf("Array element is %f ",d_a[10]);
}
""").build()

prg.p(queue, (1,), None, d_a)

Gives me

"Array element is 10" 

as output. So, the buffer actually is the linearized array. Nevertheless, the naive [x,y,z] approach known from numpy doesn't work that way. Using an 2 or 3-D Image instead of a buffer should work nevertheless.