libagx: promote math to use AGX address mode

we want to fit into the 64 + ext() << #n pattern to let us fuse address
arithmetic into our loads, so rework some libagx addressing to better match that

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/31964>

src/asahi/lib/shaders/geometry.cl

@@ -665,7 +665,8 @@ uintptr_t
 libagx_vertex_output_address(uintptr_t buffer, uint64_t mask, uint vtx,
                              gl_varying_slot location)
 {
-   return buffer + libagx_tcs_in_offs(vtx, location, mask);
+   /* Written like this to let address arithmetic work */
+   return buffer + ((uintptr_t)libagx_tcs_in_offs_el(vtx, location, mask)) * 16;
 }
 
 uintptr_t

src/asahi/lib/shaders/geometry.h

@@ -217,14 +217,21 @@ AGX_STATIC_ASSERT(sizeof(struct agx_geometry_params) == 82 * 4);
  * TODO: compact.
  */
 static inline uint
-libagx_tcs_in_offs(uint vtx, gl_varying_slot location,
-                   uint64_t crosslane_vs_out_mask)
+libagx_tcs_in_offs_el(uint vtx, gl_varying_slot location,
+                      uint64_t crosslane_vs_out_mask)
 {
    uint base = vtx * libagx_popcount(crosslane_vs_out_mask);
    uint offs = libagx_popcount(crosslane_vs_out_mask &
                                (((uint64_t)(1) << location) - 1));
 
-   return (base + offs) * 16;
+   return base + offs;
+}
+
+static inline uint
+libagx_tcs_in_offs(uint vtx, gl_varying_slot location,
+                   uint64_t crosslane_vs_out_mask)
+{
+   return libagx_tcs_in_offs_el(vtx, location, crosslane_vs_out_mask) * 16;
 }
 
 static inline uint
@@ -247,34 +254,50 @@ libagx_tcs_in_size(uint32_t vertices_in_patch, uint64_t crosslane_vs_out_mask)
  * Bounding boxes are ignored.
  */
 static inline uint
-libagx_tcs_out_offs(uint vtx_id, gl_varying_slot location, uint nr_patch_out,
-                    uint64_t vtx_out_mask)
+libagx_tcs_out_offs_el(uint vtx_id, gl_varying_slot location, uint nr_patch_out,
+                       uint64_t vtx_out_mask)
 {
    uint off = 0;
    if (location == VARYING_SLOT_TESS_LEVEL_OUTER)
       return off;
 
-   off += 4 * sizeof(float);
+   off += 4;
    if (location == VARYING_SLOT_TESS_LEVEL_INNER)
       return off;
 
-   off += 2 * sizeof(float);
+   off += 2;
    if (location >= VARYING_SLOT_PATCH0)
-      return off + (16 * (location - VARYING_SLOT_PATCH0));
+      return off + (4 * (location - VARYING_SLOT_PATCH0));
 
    /* Anything else is a per-vtx output */
-   off += 16 * nr_patch_out;
-   off += 16 * vtx_id * libagx_popcount(vtx_out_mask);
+   off += 4 * nr_patch_out;
+   off += 4 * vtx_id * libagx_popcount(vtx_out_mask);
 
    uint idx = libagx_popcount(vtx_out_mask & (((uint64_t)(1) << location) - 1));
-   return off + (16 * idx);
+   return off + (4 * idx);
+}
+
+static inline uint
+libagx_tcs_out_offs(uint vtx_id, gl_varying_slot location, uint nr_patch_out,
+                    uint64_t vtx_out_mask)
+{
+   return libagx_tcs_out_offs_el(vtx_id, location, nr_patch_out, vtx_out_mask) *
+          4;
+}
+
+static inline uint
+libagx_tcs_out_stride_el(uint nr_patch_out, uint out_patch_size,
+                         uint64_t vtx_out_mask)
+{
+   return libagx_tcs_out_offs_el(out_patch_size, 0, nr_patch_out, vtx_out_mask);
 }
 
 static inline uint
 libagx_tcs_out_stride(uint nr_patch_out, uint out_patch_size,
                       uint64_t vtx_out_mask)
 {
-   return libagx_tcs_out_offs(out_patch_size, 0, nr_patch_out, vtx_out_mask);
+   return libagx_tcs_out_stride_el(nr_patch_out, out_patch_size, vtx_out_mask) *
+          4;
 }
 
 /* In a tess eval shader, stride for hw vertex ID */

src/asahi/lib/shaders/tessellation.cl

@@ -71,13 +71,16 @@ libagx_tcs_out_address(constant struct libagx_tess_args *p, uint patch_id,
                        uint vtx_id, gl_varying_slot location, uint nr_patch_out,
                        uint out_patch_size, uint64_t vtx_out_mask)
 {
-   uint stride =
-      libagx_tcs_out_stride(nr_patch_out, out_patch_size, vtx_out_mask);
+   uint stride_el =
+      libagx_tcs_out_stride_el(nr_patch_out, out_patch_size, vtx_out_mask);
 
-   uint offs =
-      libagx_tcs_out_offs(vtx_id, location, nr_patch_out, vtx_out_mask);
+   uint offs_el =
+      libagx_tcs_out_offs_el(vtx_id, location, nr_patch_out, vtx_out_mask);
 
-   return (uintptr_t)(p->tcs_buffer) + (patch_id * stride) + offs;
+   offs_el += patch_id * stride_el;
+
+   /* Written to match the AGX addressing mode */
+   return (uintptr_t)(p->tcs_buffer) + (((uintptr_t)offs_el) << 2);
 }
 
 static uint

src/asahi/lib/shaders/texture.cl

@@ -222,7 +222,12 @@ libagx_texture_load_rgb32(constant struct agx_texture_packed *ptr, uint coord,
                           bool is_float)
 {
    agx_unpack(NULL, ptr, TEXTURE, d);
-   constant uint3 *data = (constant uint3 *)(d.address + 12 * coord);
+
+   /* This is carefully written to let us do the * 3 with a 32-bit operation but
+    * still use the free 64-bit add-extend-shift for the rest.
+    */
+   uint64_t addr = d.address + ((uint64_t)(coord * 3)) * 4;
+   constant uint3 *data = (constant uint3 *)addr;
 
    return (uint4)(*data, is_float ? as_uint(1.0f) : 1);
 }