vx-lang / vx-lang/Vx

Checker computes the matmul result shape: [m,k] @ [k,n] types as [m,n], and a k-mismatch is an error

Open
#397 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

core-lang enhancement
Dominant language
Rust
Stars
14
Forks
2
Avg merge
13h 13m
Merged PRs (30d)
70

Description

Step 4 of hiraditya/Vx.1#390, split out so its dependents can point at it.

check_binary's matmul arm (src/hir/check/operators.rs) checks the element types (E7002) and the ranks (E7001), then returns Type::Tensor(el_ty_l, vec![], top_l) — dims-less — for every matmul. Two things are missing:

  1. The inner dimensions are never compared: [2,3] @ [4,5] type-checks today. The mismatch surfaces later or never (the AST path sizes the result with a runtime memref.dim; the flat path declines it).
  2. The result type carries no shape, so everything downstream of a matmul loses static-shape information the operands had.

The checker should compute [m,n] from [m,k] @ [k,n] when both operands carry static dims, error on k mismatch, and keep the current dims-less answer only when an operand is genuinely dynamic.

The flat path already re-derives exactly this in its own matmul arm (flatten.rs, from hiraditya/Vx.1#390) — duplicated logic that exists only because the checked type is dims-less. Once the checker computes it, that arm reads the type instead.

Unblocks hiraditya/Vx.1#391: the declared-vs-computed shape check for c = a @ b needs the computed shape to exist, and the 6 corpus programs declining as "a matmul assigned to a tensor that already exists" sit behind that.

Example: both holes, at HEAD

1. The inner-dimension mismatch compiles clean.

fn main() -> i32 {
  let mut a : Tensor<f32> = Tensor<f32>([2, 3]);
  let mut b : Tensor<f32> = Tensor<f32>([4, 5]);
  a[0][0] = 1.0;
  b[0][0] = 1.0;
  let c = a @ b;   // [2,3] @ [4,5]: k is 3 on one side and 4 on the other
  return 0;
}

vxc --action emit-mlir exits 0. The flat path declines it (under the reason "a matmul whose operands are not two statically shaped matrices" -- the operands are statically shaped, the string just also covers the mismatch; the checker should own this rejection). The AST path then sizes the result at run time and emits the pair anyway:

%dim   = memref.dim %a, %c0 : memref<?x?xf32>   // m = 2
%dim_6 = memref.dim %b, %c1 : memref<?x?xf32>   // n = 5
// alloc [2,5]; linalg.matmul over a 3-vs-4 inner dimension

Nothing ever compares the 3 and the 4. The program ships and the multiply walks mismatched extents.

2. A wrong declared shape is accepted silently -- the dims-less answer makes the check vacuous.

fn main() -> i32 {
  let mut a : Tensor<f32> = Tensor<f32>([2, 3]);
  let mut b : Tensor<f32> = Tensor<f32>([3, 4]);
  a[0][0] = 1.0;
  b[0][0] = 1.0;
  let c : Tensor<f32, [9, 9]> = a @ b;   // result is [2,4]; the annotation lies
  return 0;
}

Also exits 0, on both backends. The mechanism: a @ b checks to Type::Tensor(F32, vec![], top), and is_assignable(Tensor[9,9], Tensor[]) skips the dimension comparison entirely when the source has no dims -- so the annotation is never compared against anything. The flat path's emitted IR shows what c really is:

%alloc_6 = memref.alloc() : memref<2x4xf32>    // c, declared Tensor<f32, [9, 9]>
linalg.matmul ins(%alloc, %alloc_0 : memref<2x3xf32>, memref<3x4xf32>)
              outs(%alloc_6 : memref<2x4xf32>)

Every later use that trusts the [9, 9] annotation is lied to. This is the same vacuity hiraditya/Vx.1#391 needs closed before c = a @ b into an existing buffer can check the destination's shape.

After the fix, program 1 is a type error at the @ (inner dimensions 3 vs 4), program 2 is a declaration mismatch ([9, 9] vs the computed [2, 4]), and both backends read the [2, 4] off the checked type instead of re-deriving it -- the flat arm from hiraditya/Vx.1#390 shrinks to a read, and the AST path's memref-string splitting goes away for the static case.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in src/hir/check/operators.rs at check_binary's matmul arm, then inspect the corresponding matmul handling in flatten.rs and the AST path. Make static [m,k] @ [k,n] produce [m,n], reject inner-dimension mismatches, preserve a dims-less result only for genuinely dynamic operands, and have both backends consume the checked shape.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.