fix #54: implement linalg matmul and batch_matmul

[Ritsuka314]

fix #54: implement linalg.batch_matmul

[Balint-R]

Surely x.element_type is not y.element_type or x.element_type is not init.element_type is simpler. Just move the code for setting init above this check?

[Ritsuka314]

went with len({x.element_type, y.element_type, init.element_type}) != 1 because if it's shorter, if you are OK with that.

[Balint-R]

Actually this might need to be a strict check: x.shape[1] == y.shape[0] not are_dims_compatible. IIRC the last time I tried, linalg.add on memref<10x?xf32> and memref<10x20xf32> didn't work, since linalg wants the shapes to be exactly the same, not just compatible. If compatible but not identical dimensions do work for matmul, add a testcase for it.

[Ritsuka314]

In my matmul and batch_matmul tests:

        TensorTD = TensorFactory((DYNAMIC, DYNAMIC), SInt64)

            assert (
                f[TensorTD, TensorTD, TensorTD](A.copy(), B.copy(), C.copy())
                == cor_res
            ).all()

As for add, the following test passes

def test_linalg_add_dynamic():
    @compile(dump_mlir=True, dump_mlir_passes=True)
    def f(
        m1: MemRef[UInt32, DYNAMIC, DYNAMIC],
        m2: MemRef[UInt32, DYNAMIC, DYNAMIC],
        m3: MemRef[UInt32, DYNAMIC, DYNAMIC],
    ) -> MemRef[UInt32, DYNAMIC, DYNAMIC]:
        return linalg.add(m1, m2, out=m3)

    n1 = multi_arange((10, 10), np.uint32)
    n2 = multi_arange((10, 10), np.uint32) + 1000
    n3 = multi_arange((10, 10), np.uint32) + 2000
    cor_res = n1 + n2
    res = f(n1, n2, n3)
    assert (res == cor_res).all()
    assert (n3 == cor_res).all()

The MLIR is

module {
  func.func public @f(%arg0: memref<?x?xi32>, %arg1: memref<?x?xi32>, %arg2: memref<?x?xi32>) -> memref<?x?xi32> {
    linalg.elemwise_binary {cast = #linalg.type_fn<cast_unsigned>, fun = #linalg.binary_fn<add>} ins(%arg0, %arg1 : memref<?x?xi32>, memref<?x?xi32>) outs(%arg2 : memref<?x?xi32>)
    return %arg2 : memref<?x?xi32>
  }
}

after ConvertLinalgToLoopPass

module {
  func.func public @f(%arg0: memref<?x?xi32>, %arg1: memref<?x?xi32>, %arg2: memref<?x?xi32>) -> memref<?x?xi32> {
    %c1 = arith.constant 1 : index
    %c0 = arith.constant 0 : index
    %dim = memref.dim %arg0, %c0 : memref<?x?xi32>
    %dim_0 = memref.dim %arg0, %c1 : memref<?x?xi32>
    scf.for %arg3 = %c0 to %dim step %c1 {
      scf.for %arg4 = %c0 to %dim_0 step %c1 {
        %0 = memref.load %arg0[%arg3, %arg4] : memref<?x?xi32>
        %1 = memref.load %arg1[%arg3, %arg4] : memref<?x?xi32>
        %2 = arith.addi %0, %1 : i32
        memref.store %2, %arg2[%arg3, %arg4] : memref<?x?xi32>
      }
    }
    return %arg2 : memref<?x?xi32>
  }
}

[Balint-R]

I meant something like the following:

module {
  func.func public @f(%arg0: memref<?x6xi32>, %arg1: memref<?x?xi32>, %arg2: memref<7x?xi32>){
    linalg.add ins(%arg0, %arg1 : memref<?x6xi32>, memref<?x?xi32>) outs(%arg2 : memref<7x?xi32>)
    return
  }
}

Where a static dimension is compared to a dynamic dimension. But this also seems to work now... Let me open an issue to make the existing linalg ops more lenient.

So the only change you should do is add a testcase with both static and dynamic dimensions.

[Ritsuka314]

added test