nir/spirv: Simplify matrix loads/stores

Instead of handling all of the complexity at the end, we choose to
decorate types a bit more cleverly.  When we have a row-major matrix
type, we give it the stride of a single vector and give it's array
element type (which represents a column) the actual matrix stride.

Previously, we were using stop_at_matrix and handling everything from
matrix on down as special cases but now we walk the access chain all the
way to the end and then load.  Even though this looks like it may lead
to a significant functional change, it doesn't.  The reason why we
needed to do stop_at_matrix before was to handle row-major properly
since the offsets and strides would be all out-of-order.  Now that row
major matrix types have the small stride on the matrix and the large
stride on the vector, offsetting to a single column of a row-major
matrix works fine.  The load/store code simply picks up on the fact that
the stride isn't the type size and does multiple loads.  The generated
code from these methods should be the same.

Reviewed-by: Connor Abbott <[email protected]>

1 files changed, 28 insertions, 1 deletions

diff --git a/src/compiler/spirv/spirv_to_nir.c b/src/compiler/spirv/spirv_to_nir.c
index 7a98843db77..72a890454ed 100644
--- a/src/compiler/spirv/spirv_to_nir.c
+++ b/src/compiler/spirv/spirv_to_nir.c
@@ -520,7 +520,7 @@ struct_member_decoration_cb(struct vtn_builder *b,
       ctx->type->offsets[member] = dec->literals[0];
       break;
    case SpvDecorationMatrixStride:
-      mutable_matrix_member(b, ctx->type, member)->stride = dec->literals[0];
+      /* Handled as a second pass */
       break;
    case SpvDecorationColMajor:
       break; /* Nothing to do here.  Column-major is the default. */
@@ -571,6 +571,32 @@ struct_member_decoration_cb(struct vtn_builder *b,
    }
 }
 
+/* Matrix strides are handled as a separate pass because we need to know
+ * whether the matrix is row-major or not first.
+ */
+static void
+struct_member_matrix_stride_cb(struct vtn_builder *b,
+                               struct vtn_value *val, int member,
+                               const struct vtn_decoration *dec,
+                               void *void_ctx)
+{
+   if (dec->decoration != SpvDecorationMatrixStride)
+      return;
+   assert(member >= 0);
+
+   struct member_decoration_ctx *ctx = void_ctx;
+
+   struct vtn_type *mat_type = mutable_matrix_member(b, ctx->type, member);
+   if (mat_type->row_major) {
+      mat_type->array_element = vtn_type_copy(b, mat_type->array_element);
+      mat_type->stride = mat_type->array_element->stride;
+      mat_type->array_element->stride = dec->literals[0];
+   } else {
+      assert(mat_type->array_element->stride > 0);
+      mat_type->stride = dec->literals[0];
+   }
+}
+
 static void
 type_decoration_cb(struct vtn_builder *b,
                    struct vtn_value *val, int member,
@@ -807,6 +833,7 @@ vtn_handle_type(struct vtn_builder *b, SpvOp opcode,
       };
 
       vtn_foreach_decoration(b, val, struct_member_decoration_cb, &ctx);
+      vtn_foreach_decoration(b, val, struct_member_matrix_stride_cb, &ctx);
 
       const char *name = val->name ? val->name : "struct";