How to describe a lookup operation like gather instruction

[daleidalei]

Hi, I have a little problem describing a lookup operation similar to the gather instruction on a new hardware platform,I wrote such a compute and shedule as follows：

class DmaGenerator : public Halide::Generator<DmaGenerator> {
public:
    Input<Buffer<int32_t>> input{"input", 2};
    Input<Buffer<int32_t>> index{"index", 2};
    Output<Buffer<int32_t>> output{"output", 2};
    Var x {"x"};
    Var y {"y"};
    void generate() {
        input_buffer(x, y) = input(x, y);// dma operation
	index_buffer(x, y) = index(x, y);// dma operation
	Expr xCoord = clamp(index_buffer(x, y), 0, input.width() - 1);
	Expr yCoord = clamp(index_buffer(x, y), 0, input.height() - 1);
	lookup(x, y) = input_buffer(xCoord, yCoord);
	output(x, y) = lookup(x, y);
    }
	
	void schedule()
	{
		input.dim(0).set_bounds(0, 480);
                input.dim(1).set_bounds(0, 360);
		index.dim(0).set_bounds(0, 480);
                index.dim(1).set_bounds(0, 360);
		
		input_buffer.store_in(MemoryType::Buffer);
                index_buffer.store_in(MemoryType::Buffer);
                lookup.store_in(MemoryType::Buffer);
		
		int vectorize_factor = 64;
		int tile_w = 64;
		int tile_h = 32;
		Var xOuter("xOuter"), xInner("xInner");
		Var yOuter("yOuter"), yInner("yInner");
                output.tile(x, y, xOuter, yOuter, xInner, yInner, tile_w, tile_h);
        
                output.compute_root();
		index_buffer.compute_at(output, xOuter).align_storage(x, 16);
		input_buffer.compute_at(output, xOuter).align_storage(x, 16);
                lookup.compute_at(output, xOuter).vectorize(x, vectorize_factor, TailStrategy::GuardWithIf).align_storage(x, 16);
	}
	
private:
        Func input_buffer {"input_buffer"};
	Func index_buffer {"index_buffer"};
	Func lookup {"lookup"};
};

The generated partial IR looks like this：

    allocate index_buffer[int32 * 64 * 32] in Buffer
    produce index_buffer {
     for (index_buffer.s0.y, output.s0.y.yInner.base, 32) {
      for (index_buffer.s0.x, output.s0.x.xInner.base, 64) {
       index_buffer[((index_buffer.s0.y - output.s0.y.yInner.base)*64) + (index_buffer.s0.x - output.s0.x.xInner.base)] = index[(index.stride.1*index_buffer.s0.y) + index_buffer.s0.x]
      }
     }
    }
    allocate input_buffer[int32 * 480 * 360] in Buffer
    produce input_buffer {
     for (input_buffer.s0.y, 0, 360) {
      for (input_buffer.s0.x, 0, 480) {
       input_buffer[(input_buffer.s0.y*480) + input_buffer.s0.x] = input[(input.stride.1*input_buffer.s0.y) + input_buffer.s0.x]
      }
     }
    }
    allocate lookup[int32 * 64 * 32] in UBuffer

The problem I'm currently having is that the size of the input_buffer allocation is not what I expected, I expected the buffer allocated by input_buffer to be consistent with the tile size, because the soc on-chip buffer is relatively small, and such a large buffer cannot be stored. The reason for this is that the value of index_buffer can be guaranteed by the input, x in index_buffer will not exceed tile_w and y will not exceed tile_h, so it is enough for input_buffer to allocate tile_w*tile_h size buffer, If the upper bounds in clamp are tile_w and tile_h, this description does not conform to the mathematical definition, and the subsequent comparison with the cpu output will also be inconsistent.
How to express a mathematical definition than can achieve my expectations?
Looking forward for your reply, best wishes!