diff --git a/src/asahi/isa/AGX2.xml b/src/asahi/isa/AGX2.xml
index c63196454c9..a9e757091cc 100644
--- a/src/asahi/isa/AGX2.xml
+++ b/src/asahi/isa/AGX2.xml
@@ -274,6 +274,11 @@
     <property name="size64"/>
   </encoding>
 
+  <encoding name="ureg64_shift2">
+    <immediate name="uniform" bit="0:6" shift="2"/>
+    <property name="size64"/>
+  </encoding>
+
   <encoding name="exreg16">
     <immediate name="reg" bit="0:7"/>
   </encoding>
@@ -283,6 +288,20 @@
     <property name="size32"/>
   </encoding>
 
+  <encoding name="heap_offs">
+    <case>
+      <exact bit="31">0</exact>
+      <immediate name="reg" bit="0:8"/>
+      <property name="size32"/>
+    </case>
+
+    <case>
+      <exact bit="31">1</exact>
+      <immediate name="imm" bit="0:30"/>
+      <property name="size32"/>
+    </case>
+  </encoding>
+
   <encoding name="atomic_dest">
     <case>
       <exact bit="8">1</exact>
@@ -1363,16 +1382,16 @@
   <ins name="tex_state_store" length="8" display="load">
     <exact bit="0:7 20">111101101</exact>
     <dest kind="ts" bit="8:14"/>
-    <src kind="ureg64" bit="58:63"/>
-    <src kind="exreg32" bit="27:33"/>
+    <src kind="ureg64_shift2" bit="59:63"/>
+    <src kind="heap_offs" bit="26:56 58"/>
     <modifier name="group" kind="group" bit="16"/>
   </ins>
 
   <ins name="sampler_state_store" length="8" display="load">
     <exact bit="0:7 20">110101101</exact>
     <dest kind="ss" bit="8:14"/>
-    <src kind="ureg64" bit="58:63"/>
-    <src kind="exreg32" bit="27:33"/>
+    <src kind="ureg64_shift2" bit="59:63"/>
+    <src kind="heap_offs" bit="26:56 58"/>
     <modifier name="group" kind="group" bit="16"/>
   </ins>
 
diff --git a/src/asahi/isa/test/cases.txt b/src/asahi/isa/test/cases.txt
index f07b5050617..e017f3f0183 100644
--- a/src/asahi/isa/test/cases.txt
+++ b/src/asahi/isa/test/cases.txt
@@ -491,6 +491,7 @@ a191014000000000     ldcf     r4_r5_r6, cf1
 a191024000000000     ldcf     r4_r5_r6, cf2
 a191034000000000     ldcf     r4_r5_r6, cf3
 ad01100800000010     load     ss1, du4, r1, a
+ad01116001000024     load     ss1, du8, 88, b
 b1800080004a00000a00 imgwblk  r0l, ts0, r0, 0, 2d, u8norm
 b1848480024a00000a00 imgwblk  r1l, ts2, r2, 0, explicit_coords, 2d, u8norm
 b1880680004880000a00 imgwblk  r2l, ts0, r3, 0, 2d_ms_array, u8norm
@@ -536,6 +537,7 @@ d500688c10500c00     atomic   _, r6_r7, du12, r3.sx, cmpxchg, coherent1, a
 d501248c10500600     atomic   _, r3, du10, r1.sx, imax, coherent1, a
 d508288c10d01000     atomic   r1, r8_r9, du12, r1.sx, cmpxchg, coherent1, a
 d509a48c10d00800     atomic   r1, r4, du10, r5.sx, imax, coherent1, a
+ed00100001000024     load     ts0, du8, 64, a
 ed01100800000000     load     ts1, du0, r1, a
 ed02100800000008     load     ts2, du2, r1, a
 ed02100800000018     load     ts2, du6, r1, a