Support a complex number type for Matrix values

wigging · May 13, 2024, 2:46pm

I have a generic matrix type that conforms to the FloatingPoint protocol as shown below.

import Accelerate

struct Matrix<T: FloatingPoint> {
    let rows: Int
    let columns: Int
    var values: [T]

    init(_ content: [[T]]) {
        self.rows = content.count
        self.columns = content[0].count
        self.values = content.flatMap { $0 }
    }

    init(rows: Int, columns: Int, values: [T]) {
        self.rows = rows
        self.columns = columns
        self.values = values
    }

    subscript(row: Int, column: Int) -> T {
        get { return values[(row * columns) + column] }
        set { values[(row * columns) + column] = newValue }
    }
}

The matrix supports various mathematical operations for single and double precision. Below is the matrix multiplication operation.

extension Matrix {

    static func * (lhs: Matrix, rhs: Matrix) -> Matrix where T == Double {
        precondition(lhs.columns == rhs.rows, "Number of columns in left matrix must equal number of rows in right matrix")
        let a = lhs.values
        let b = rhs.values
        var c = [Double](repeating: 0.0, count: lhs.rows * rhs.columns)

        let m = lhs.rows     // rows in matrices A and C
        let n = rhs.columns  // columns in matrices B and C
        let k = lhs.columns  // columns in matrix A and rows in matrix B
        let alpha = 1.0
        let beta = 0.0

        // matrix multiplication where C ← αAB + βC
        cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, alpha, &a, k, &b, n, beta, &c, n)

        let mat = Matrix(rows: lhs.rows, columns: rhs.columns, values: c)
        return mat
    }
    
    static func * (lhs: Matrix, rhs: Matrix) -> Matrix where T == Float {
        precondition(lhs.columns == rhs.rows, "Number of columns in left matrix must equal number of rows in right matrix")
        let a = lhs.values
        let b = rhs.values
        var c = [Float](repeating: 0.0, count: lhs.rows * rhs.columns)

        let m = lhs.rows     // rows in matrices A and C
        let n = rhs.columns  // columns in matrices B and C
        let k = lhs.columns  // columns in matrix A and rows in matrix B
        let alpha = Float(1.0)
        let beta = Float(0.0)

        // matrix multiplication where C ← αAB + βC
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, alpha, a, k, b, n, beta, &c, n)

        let mat = Matrix(rows: lhs.rows, columns: rhs.columns, values: c)
        return mat
    }
}

I can create matrices with integer, float, and double values. See below for examples that use integer and double values.

let a = Matrix([[1, 2],
                [3, 4]])

let b = Matrix([[1, 2],
                [3, 4]])

let c = a * b

let a = Matrix([[1.0, 2.0],
                [3.0, 4.0]])

let b = Matrix([[1.0, 2.0],
                [3.0, 4.0]])

let c = a * b

I would like to add support for single and double precision complex numbers which I attempted to do below:

struct Complex<T> {
    let real: T
    let imag: T
}

extension Matrix {

    static func * (lhs: Matrix, rhs: Matrix) -> Matrix where T == Complex<Double> {
        precondition(lhs.columns == rhs.rows, "Number of columns in left matrix must equal number of rows in right matrix")
        let a = lhs.values
        let b = rhs.values
        var c = [Complex](repeating: Complex(real: 0.0, imag: 0.0), count: lhs.rows * rhs.columns)
        
        let m = lhs.rows
        let n = rhs.columns
        let k = lhs.columns
        let alpha = [Complex(real: 1.0, imag: 0.0)]
        let beta = [Complex(real: 1.0, imag: 0.0)]
        
        cblas_zgemm_wrapper(m, n, k, alpha, a, k, b, n, beta, &c, n)
        
        let mat = Matrix(rows: lhs.rows, columns: rhs.columns, values: c)
        return mat
    }
}

This doesn't work because obviously Complex doesn't conform to FloatingPoint. I could just use struct Matrix<T> { ... } but then I lose all the benefits of conforming to FloatingPoint. Is there a way to extend the FloatingPoint protocol with a complex number type? Or is there a different approach I should use for this?

tera · May 13, 2024, 4:02pm

What would they be? From what I see in your example – you are not using FloatingPoint requirement anywhere.

scanon · May 13, 2024, 4:50pm

If you use the Numerics package Complex type (or something like it), it comes with conformance to the Numerics AlgebraicField and ElementaryFunctions protocols, which provide most of the scalar operations you would want to implement arbitrary linear algebra (in practice, if you were building out a full linear algebra library from scratch, you would want to refine these to include a few other things as computational hooks for convenience, but they would be fairly minimal).

You don't have to use Numerics' version of Complex, but you're going to end up building something like it in order to express the constraints that you'll want to have.

wigging · May 13, 2024, 5:21pm

I chose FloatingPoint because it prevents strings from being used as values. You can't do Matrix(rows: 2, columns: 2, values: ["one", "two", "three"]) because String doesn't conform to FloatingPoint. It also allows be to define the matrix values as integers but automatically create a Matrix<Double> from those values. I also chose FloatingPoint because, until now, the Accelerate functions that I have used only accept Float and Double values. So FloatingPoint initially seemed like a good choice. But now I'm trying to do complex matrix and vector operations with Accelerate functions which need a Complex type for the values. Basically, I want to have a Matrix<T> where T would be Int, Float, Double, or Complex.

tera · May 13, 2024, 5:25pm

In this case you could define your own protocol:

protocol MatrixElement {}

extension Int: MatrixElement {}
extension Float: MatrixElement {}
extension Double: MatrixElement {}
extension Complex: MatrixElement {}

struct Matrix<Element: MatrixElement> { ... }

Nevin · May 13, 2024, 5:33pm

That sounds like an anti-goal to me. You are actively preventing users from doing useful and meaningful things with your type.

I would recommend defining Matrix<Element> with no constraints on Element. Then use constrained extensions to add features:

extension Matrix where Element: AdditiveArithmetic {
  // implement addition and subtraction here
}

extension Matrix where Element: Numeric {
  // implement scalar multiplication here
}

// etc

tera · May 13, 2024, 5:42pm

Very good point! Matrix of strings (or pretty much anything) could make sense: even if you can't, say, multiply it you could still do some meaningful operations on it, e.g. transpose or subscript or extract a submatrix or compare two matrices for equality.

wigging · May 13, 2024, 5:50pm

I tried this approach which works fine for the complex numbers.

struct Matrix<T> {
    let rows: Int
    let columns: Int
    var values: [T]

    init(_ content: [[T]]) {
        self.rows = content.count
        self.columns = content[0].count
        self.values = content.flatMap { $0 }
    }

    init(rows: Int, columns: Int, values: [T]) {
        self.rows = rows
        self.columns = columns
        self.values = values
    }

    subscript(row: Int, column: Int) -> T {
        get { return values[(row * columns) + column] }
        set { values[(row * columns) + column] = newValue }
    }
}

extension Matrix {
    static func * (lhs: Matrix, rhs: Matrix) -> Matrix where T == Complex<Double> {
        precondition(lhs.columns == rhs.rows, "Number of columns in left matrix must equal number of rows in right matrix")
        let a = lhs.values
        let b = rhs.values
        var c = [Complex](repeating: Complex(real: 0.0, imag: 0.0), count: lhs.rows * rhs.columns)

        let m = lhs.rows     // rows in matrices A and C
        let n = rhs.columns  // columns in matrices B and C
        let k = lhs.columns  // columns in matrix A and rows in matrix B

        let alpha = [Complex(real: 1.0, imag: 0.0)]
        let beta = [Complex(real: 1.0, imag: 0.0)]

        cblas_zgemm_wrapper(m, n, k, alpha, a, k, b, n, beta, &c, n)

        let mat = Matrix(rows: lhs.rows, columns: rhs.columns, values: c)
        return mat
    }
}

But I have an issue with integer values because none of the BLAS and LAPACK routines work with integers. As shown below, I have to convert the integers to Double to use the BLAS function. This doesn't seem efficient for large matrices. Any suggestions on how to handle integers for this situation?

extension Matrix {
    static func * (lhs: Matrix, rhs: Matrix) -> Matrix where T == Int {
        precondition(lhs.columns == rhs.rows, "Number of columns in left matrix must equal number of rows in right matrix")
        var a = lhs.values.compactMap { Double($0) }
        var b = rhs.values.compactMap { Double($0) }
        var c = [Double](repeating: 0.0, count: lhs.rows * rhs.columns)

        let m = lhs.rows     // rows in matrices A and C
        let n = rhs.columns  // columns in matrices B and C
        let k = lhs.columns  // columns in matrix A and rows in matrix B
        let alpha = 1.0
        let beta = 0.0

        // matrix multiplication where C ← αAB + βC
        cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, alpha, &a, k, &b, n, beta, &c, n)

        let mat = Matrix(rows: lhs.rows, columns: rhs.columns, values: c.compactMap { Int($0) } )
        return mat
    }
}

scanon · May 13, 2024, 6:02pm

LAPACK doesn't support integers because the operations that it provides don't make any sense for integers--if you try to factor a matrix with integer entries, you immediately leave the integers, unless the integer values are chosen just so (which happens with probability ~0 for real-world data). You can make sense of much of the operations LAPACK provides over the rationals, and the rest over the reals or complex numbers. Floating-point (and complex floating-point) are the best fixed-size approximations we have to those fields (conventional rational approximation is essentially always worse than floating-point, for any fixed size).

So LAPACK is all written in terms of (complex) floating-point, and BLAS, which exists largely as a tunable base library for LAPACK, supports the same.

You can efficiently implement GEMM (and other level-3 BLAS) for integers because there's reuse of the inputs; you get to amortize the O(n²) conversion cost across O(n³) operations.¹ But there's usually not much reason to because of the aforementioned considerations. The place that integers do come up a lot now is as compressed storage for NN weights, where Int8 and even smaller formats are pretty common (but implemented in dedicated NN library routines, rather than BLAS).

¹ usually you would avoid allocating a complete copy of the matrices as integers by writing an outer loop that iterates over "reasonably size" tiles of A, B and C (maybe a few pages each, depending on page size) and converts just those tiles before calling into GEMM. This gets you enough re-use to amortize most of the conversion cost while requiring only constant workspace. A dedicated implementation might also do the conversion while loading A, B, C into register.

tera · May 13, 2024, 6:13pm

I'd do this:

struct Matrix<Element> { ... }

extension Matrix where Element == Int {
    static func * (lhs: Matrix, rhs: Matrix) -> Matrix { ... }
}
extension Matrix where Element == Float {
    static func * (lhs: Matrix, rhs: Matrix) -> Matrix { ... }
}
extension Matrix where Element == Double {
    static func * (lhs: Matrix, rhs: Matrix) -> Matrix { ... }
}

Orthogonal to this is the question of what to do inside each:

convert to float/double and do dgemm (BTW, this could still be faster as to convert from integers (and convert the result from doubles back to integers) takes O(n^2) and the matrix multiplication itself takes O(n^3)).
use GPU (MetalPerformanceShaders)
Use a manual implementation

The fastest method could depend upon matrix dimensions (e.g. for small matrices (3) would be the winner, for large integer matrices (2) would be the winner, etc). I'd test with different matrix sizes, and then figure out the relevant thresholds (per type) to use to choose between 1 ... 3.

func methodToUse(type: Element, rows: Int, columns: Int) -> MethodToUse { ... }

extension Matrix where Element == Double {
    static func * (lhs: Matrix, rhs: Matrix) -> Matrix {
        switch methodToUse(type: .double, rows: rows, columns: columns) {
            case .gemm: cblas_dgemm(...)
            case .gpu: ... use MPSMatrixMultiplication(...)
            case .manual: ... use manual implementation
        }
    }
}

wigging · May 13, 2024, 6:18pm

I think your first item here (1.) is what I did in the previous example using compactMap.

wigging · May 13, 2024, 6:35pm

Thank you for the BLAS and LAPACK explanation. I had no idea that it would be so involved to do these operations with integers. In the original example I posted, where Matrix<T: FloatingPoint>, I could multiply two matrices of integers and get a Matrix<Double> back. Is the FloatingPoint protocol automatically converting the integers to double?

tera · May 13, 2024, 6:43pm

You were probably using integer literals, not integers...

func foo<T: FloatingPoint>(_ v: T) {}
foo(42) // uses Double here as you you wrote foo(42.0)
let i = 42
foo(i) // 🛑 Error

Which, BTW, would be different with the above suggested MatrixElement protocol:

protocol MatrixElement {}
extension Float: MatrixElement {}
extension Double: MatrixElement {}
extension Int: MatrixElement {}

func bar<T: MatrixElement>(_ v: T) {}
bar(42)
let j = 42
bar(j) // ✅... You have to provide a valid implementation for integers as well as other types

wigging · May 13, 2024, 7:04pm

The FloatingPoint protocol inherits from ExpressibleByIntegerLiteral, is that how I was able to define the matrix with integer like syntax? If I do something like this:

static func * (lhs: Matrix, rhs: Matrix) -> Matrix where T: ExpressibleByIntegerLiteral {
let a = lhs.values
let b = rhs.values
...
}

Then how do I get the left and right matrix values (lhs and rhs) as doubles?

tera · May 13, 2024, 7:46pm

I think you are looking for this:

extension Matrix where Element == Int {
    static func * (lhs: Matrix, rhs: Matrix) -> Matrix {
        let result = Matrix<Double>(rows: lhs.rows, columns: lhs.columns, values: lhs.values.map { Double($0) }) *
            Matrix<Double>(rows: rhs.rows, columns: rhs.columns, values: rhs.values.map { Double($0) })
        return Matrix<Int>(rows: result.rows, columns: result.columns, values: result.values.map { Int($0) })
    }
}

wigging · May 13, 2024, 8:35pm

I have more questions about using integers but I'll ask them in a separate post since the topic of this post is about complex numbers. Anyway, I have good idea on how to proceed with the complex matrix values. Thanks again for the help.

tera · May 14, 2024, 11:08am

Good.

May I ask something in return:

this is from cblas_dgemm docs. Could I set β = 0 and C instead of A? to get to the:

C←αCB

This would be super handy to implement *= (multiply in place) operation.

wigging · May 14, 2024, 12:33pm

The matrix dimensions for cblas_dgemm need to be compatible to perform the multiplication. For example, if matrix A is m x n and matrix B is n x p then matrix C must have dimension m x p. I think C and B would have to be the same dimension to do in-place multiplication C ← αCB otherwise the size of C would change.

tera · May 14, 2024, 12:38pm

Ah, yep, now that you said that that's obvious

I wonder if a * b is the best API choice in this case then. More performant would be an API that doesn't need to allocate the result:

multiply(a, b, &c)

where c needs to be of the correct size.

wigging · May 14, 2024, 12:50pm

I like using a * b because it resembles the math syntax. Actually, I wonder if a × b could be used too? Anyway, both c = a * b or multiply(a, b, &c) require three matrices to be created a, b, and c. I don't see how multiply(a, b, &c) would avoid this. For a * b, the c matrix is created in the function, for multiply(a, b, &c) the c matrix is created outside the function.