This is a reasonable place to put this logic - OpenGL loads flatten differently.

However given that you're doing this, something occurs to me. Would it be possible to rewrite Calls inside a GL kernel into some intrinsic (instead of introducing multi-index loads). E.g:

Call(foo, args) -> gl_texture_fetch(foo.buffer, mutated_args, nearest_neighbor)

foo.buffer would be a Variable::make(Handle(), "foo.buffer"), which is the buffer symbol introduced by storage flattening.

Later passes that look for loads (e.g. WhereIsBufferUsed) may have to understand that this counts as a read of the buffer, but that seems simpler than adding multi-index loads everywhere.

Provides would become something like:
gl_texture_store(foo.buffer, location, value);

Intrinsics are Exprs though, so you'd have to wrap it in an Evaluate node to make it a Stmt.

Is this approach feasible? If so it would also be nice to split this pass out into its own file that is called from Lower.cpp. I guess that means you'd have to add the need_buffer_t map as an argument to storage flattening, but that's also reasonable. StorageFlattening by itself seems to never set need_buffer_t, so it doesn't make sense as-is, and seems like it should be an argument (I guess maybe it used to set it if it encountered a Variable named foo.buffer, but now AllocationBoundsInference handles that. This needs cleaning up, but not as part of this pull request).