Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Preparation for Scala 3 migration #201

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions src/main/scala/apps/cameraPipeline.scala
Original file line number Diff line number Diff line change
Expand Up @@ -281,7 +281,7 @@ object cameraPipeline {
(lf32(1.0f) / kelvin - lf32(1.0f / 3200)) / lf32(1.0f / 7000 - 1.0f / 3200)
(
zipND(2)(matrix_3200, matrix_7000) |>
map(map(fun(p => p._1 * alpha + p._2 * (lf32(1.0f) - alpha)))) >>
map(map(fun(p => p.`1` * alpha + p.`2` * (lf32(1.0f) - alpha)))) >>
map(map(fun(v => cast(v * lf32(256.0f)) :: i16))) // Q8.8 fixed point
) |> fun(matrix =>
input |> transpose >>
Expand Down Expand Up @@ -380,11 +380,11 @@ object cameraPipeline {
blur121(u8)(u16)),
unsharp => {
letImage(mapImage(zipImage(plane, unsharp), fun(p =>
(cast(p._1) :: i16) - (cast(p._2) :: i16)
(cast(p.`1`) :: i16) - (cast(p.`2`) :: i16)
)), mask => {
mapImage(zipImage(plane, mask), fun(p =>
u8_sat(i16)((cast(p._1) :: i16)
+ (p._2 * (cast(strength_x32) :: i16)) / li16(32))
u8_sat(i16)((cast(p.`1`) :: i16)
+ (p.`2` * (cast(strength_x32) :: i16)) / li16(32))
)).expr
})
}))
Expand Down
8 changes: 3 additions & 5 deletions src/main/scala/apps/convolution.scala
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,8 @@ object convolution {

private val dotElemWeightsSeq = fun((weights, elem) =>
oclReduceSeqUnroll(AddressSpace.Private)(fun((acc, pair) => {
val pixel = pair._1
val weight = pair._2
val pixel = pair.`1`
val weight = pair.`2`
acc + (pixel * weight)
}))(lf32(0.0f))(zip(join(elem))(weights)))

Expand Down Expand Up @@ -143,9 +143,7 @@ object convolution {
matrix: Array[Array[Float]],
weights: Array[Float]
): (Array[Float], TimeSpan[Time.ms]) = {
val f = k.as[ScalaFunction `(`
Array[Array[Float]] `,` Array[Float]
`)=>` Array[Float]]
val f = k.as[Args `(` Array[Array[Float]] `,` Array[Float], Array[Float]]
Bastacyclop marked this conversation as resolved.
Show resolved Hide resolved
f(localSize, globalSize)(matrix `,` weights)
}
}
52 changes: 26 additions & 26 deletions src/main/scala/apps/gemmTensor.scala
Original file line number Diff line number Diff line change
Expand Up @@ -48,20 +48,20 @@ object gemmTensor {

zip
(bT |> split(nTileFrag))
(aRowsC._2 |> transpose |> split(nTileFrag)) |> // n/nTileFrag.(nTileFrag.k.f16, nTileFrag.mTileFrag.f32)
(aRowsC.`2` |> transpose |> split(nTileFrag)) |> // n/nTileFrag.(nTileFrag.k.f16, nTileFrag.mTileFrag.f32)
mapWarp(0)(fun(bColumnsTCT =>

zip
(transpose(aRowsC._1) |> split(kTileFrag))
(transpose(bColumnsTCT._1) |> split(kTileFrag)) |> // k/kTileFrag.(kTile.mTileFrag.f16 x kTile.nTile.f16)
(transpose(aRowsC.`1`) |> split(kTileFrag))
(transpose(bColumnsTCT.`1`) |> split(kTileFrag)) |> // k/kTileFrag.(kTile.mTileFrag.f16 x kTile.nTile.f16)

oclReduceSeq(AddressSpace.Private)(fun((cTile, abTiles) =>
tensorMMA(
abTiles._1 |> transpose |> asFragment |> toPrivate,
abTiles._2 |> asFragment |> toPrivate,
abTiles.`1` |> transpose |> asFragment |> toPrivate,
abTiles.`2` |> asFragment |> toPrivate,
cTile)))

(bColumnsTCT._2 |>
(bColumnsTCT.`2` |>
transpose |>
asFragment |> toPrivate |>
mapFragment(fun(x => x * (beta / alpha)))) |>
Expand Down Expand Up @@ -107,48 +107,48 @@ object gemmTensor {

zip
(bT |> split(nTileBlock))
(aRowsBlockC._2 |> transpose |> split(nTileBlock)) |>
(aRowsBlockC.`2` |> transpose |> split(nTileBlock)) |>
mapBlock(0)(fun(bColumnsTBlockCT => // (nTileBlock.k.f16, nTileBlock.mTileBlock.f32)

zip
(aRowsBlockC._1 |> split(mTileWarp))
(bColumnsTBlockCT._2 |> transpose |> split(mTileWarp)) |>
(aRowsBlockC.`1` |> split(mTileWarp))
(bColumnsTBlockCT.`2` |> transpose |> split(mTileWarp)) |>
mapThreads(1)(fun(aRowsWarpC => // (mTileWarp.k.f16, mTileWarp.nTileBlock.f32)

zip
(bColumnsTBlockCT._1 |> split(nTileWarp))
(aRowsWarpC._2 |> transpose |> split(nTileWarp)) |>
(bColumnsTBlockCT.`1` |> split(nTileWarp))
(aRowsWarpC.`2` |> transpose |> split(nTileWarp)) |>
mapWarp(0)(fun(bColumnsTWarpCT => // (nTileWarp.k.f16, nTileWarp.mTileWarp.f32)

zip
(aRowsWarpC._1 |> transpose |> split(kTileFrag))
(bColumnsTWarpCT._1 |> transpose |> split(kTileFrag)) |>
(aRowsWarpC.`1` |> transpose |> split(kTileFrag))
(bColumnsTWarpCT.`1` |> transpose |> split(kTileFrag)) |>
// k/kTileFrag.(kTileFrag.mTileWarp.f16, kTileFrag.nTileWarp.f16)

oclReduceSeq(AddressSpace.Private)(fun((cFrags, abTilesWarp) =>

//Load tiles of a-matrix into fragments
let(toPrivate(
abTilesWarp._1 |> transpose |> split(mTileFrag) |>
abTilesWarp.`1` |> transpose |> split(mTileFrag) |>
mapSeqUnroll(fun(aFragTile =>
aFragTile |> asFragment))))
be(aFrags => // mTileWarp/mTileFrag.WmmaAMatrix

//Load tiles of b-matrix into fragments
let(toPrivate(
abTilesWarp._2 |> transpose |> split(nTileFrag) |>
abTilesWarp.`2` |> transpose |> split(nTileFrag) |>
mapSeqUnroll(fun(bFragTileT =>
bFragTileT |> transpose |> asFragment))))
be(bFrags => // nTileWarp/nTileFrag.WmmaBMatrix

//Do matrix multiplication and accumulate with tensor cores
zip(aFrags)(cFrags) |>
mapSeqUnroll(fun(acFrags =>
zip(bFrags)(acFrags._2) |>
zip(bFrags)(acFrags.`2`) |>
mapSeqUnroll(fun(bcFrags =>
tensorMMA(acFrags._1, bcFrags._1, bcFrags._2)))))))))
tensorMMA(acFrags.`1`, bcFrags.`1`, bcFrags.`2`)))))))))

(bColumnsTWarpCT._2 |> transpose |> split(mTileFrag) |>
(bColumnsTWarpCT.`2` |> transpose |> split(mTileFrag) |>
mapSeq(fun(cTiles =>
cTiles |> transpose |> split(nTileFrag) |>
mapSeq(fun(cTileFragT =>
Expand Down Expand Up @@ -270,14 +270,14 @@ object gemmTensor {
oclReduceSeq(AddressSpace.Private)(fun((cFragsBlock, aTbTileBlock) =>

//Load aTile and bTile to shared memory
let(aTbTileBlock._1 |>
let(aTbTileBlock.`1` |>
transpose |>
copyMatrix(config.mTileBlock, config.kTileBlock, 8) |>
toSharedWithPadding(config.kTileBlock, 8))
be(aTile =>

//Load bTile transposed to shared memory
let(aTbTileBlock._2 |>
let(aTbTileBlock.`2` |>
transpose |>
copyMatrix(config.nTileBlock, config.kTileBlock, 8) |>
toSharedWithPadding(config.kTileBlock, 8))
Expand All @@ -289,9 +289,9 @@ object gemmTensor {

mapWarp(fun(abWarpC =>
warpMMA(
abWarpC._1._1,
abWarpC._1._2,
abWarpC._2 |> split(config.nNumberOfFragsWarp)) |>
abWarpC.`1`.`1`,
abWarpC.`1`.`2`,
abWarpC.`2` |> split(config.nNumberOfFragsWarp)) |>

mapSeq(mapSeq(fun(x => x))))) |>
join |>
Expand Down Expand Up @@ -464,9 +464,9 @@ object gemmTensor {

blockGEMM(
alpha, beta,
aRowsBlockBColumnBlockCTileBlock._1._1, //aRowsBlockBColumnBlockCTileBlock._1._1
aRowsBlockBColumnBlockCTileBlock._1._2, //aRowsBlockBColumnBlockCTileBlock._1._2
aRowsBlockBColumnBlockCTileBlock._2) |> //mTileBlock.nTilcblock.f32
aRowsBlockBColumnBlockCTileBlock.`1`.`1`, //aRowsBlockBColumnBlockCTileBlock.`1`.`1`
aRowsBlockBColumnBlockCTileBlock.`1`.`2`, //aRowsBlockBColumnBlockCTileBlock.`1`.`2`
aRowsBlockBColumnBlockCTileBlock.`2`) |> //mTileBlock.nTilcblock.f32
transpose)) |> //m/mTileBlock*n/nTileBlock.nTileBlock.mTilcblock.f32
join |> //m/mTileBlock*n.mTilcblock.f32
split(n) |> //m/mTileBlock.n.mTilcblock.f32
Expand Down
28 changes: 14 additions & 14 deletions src/main/scala/apps/gemv.scala
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import HighLevelConstructs.reorderWithStride

object gemv {
// we can use implicit type parameters and type annotations to specify the function type of mult
val mult = impl{ dt: DataType => fun(x => x._1 * x._2) :: ((dt x dt) ->: dt) }
val mult = impl{ dt: DataType => fun(x => x.`1` * x.`2`) :: ((dt x dt) ->: dt) }
val add = fun(x => fun(y => x + y))
val scal = impl { n: Nat =>
fun(xs => fun(a =>
Expand All @@ -30,15 +30,15 @@ object gemv {
(m`.`f32)
)((mat, xs, ys, alpha, beta) =>
zip(map(fun(row => alpha * dot(row, xs)))(mat))(scal(ys, beta)) |>
map(fun(x => x._1 + x._2))
map(fun(x => x.`1` + x.`2`))
))

val gemvSequential = depFun((n: Nat, m: Nat) => fun(
(m`.`n`.`f32) ->: (n`.`f32) ->: (m`.`f32) ->: f32 ->: f32 ->:
(m`.`f32)
)((mat, xs, ys, alpha, beta) =>
toMem(zip(mapSeq(fun(row => alpha * dotSeq(row, xs)))(mat))(scalSeq(ys, beta))) |>
mapSeq(fun(x => x._1 + x._2))
mapSeq(fun(x => x.`1` + x.`2`))
))

object ocl {
Expand Down Expand Up @@ -81,12 +81,12 @@ object gemv {
)((mat, xs, ys, alpha, beta) =>
zip(mat)(ys) |>
mapWorkGroup(fun(t =>
zip(xs)(t._1) |>
zip(xs)(t.`1`) |>
split(n) |>
toLocalFun(mapLocal(
reduceSeq(fun(a => fun(x => mult(x) + a)))(lf32(0.0f))
)) |>
mapLocal(fun(x => (alpha * x) + (t._2 * beta)))
mapLocal(fun(x => (alpha * x) + (t.`2` * beta)))
)) |> join
))

Expand All @@ -95,15 +95,15 @@ object gemv {
)((mat, xs, ys, alpha, beta) =>
zip(mat)(ys) |>
mapWorkGroup(fun(t =>
zip(xs)(t._1) |>
zip(xs)(t.`1`) |>
reorderWithStride(128) |>
split(n /^ 128) |>
toLocalFun(mapLocal(
reduceSeq(fun(a => fun(x => mult(x) + a)))(lf32(0.0f))
)) |>
split(128) |>
toLocalFun(mapLocal(reduceSeq(add)(lf32(0.0f)))) |>
mapLocal(fun(x => (alpha * x) + (t._2 * beta)))
mapLocal(fun(x => (alpha * x) + (t.`2` * beta)))
)) |> join
))

Expand All @@ -113,14 +113,14 @@ object gemv {
)((mat, xs, ys, alpha, beta) =>
zip(mat)(ys) |>
mapWorkGroup(fun(t =>
zip(xs)(t._1) |>
zip(xs)(t.`1`) |>
reorderWithStride(128) |>
split(n /^ 128) |>
toLocalFun(mapLocal(
reduceSeq(fun(a => fun(x => mult(x) + a)))(lf32(0.0f))
)) |>
toLocalFun(reduceSeq(add)(lf32(0.0f))) |>
fun(x => (alpha * x) + (t._2 * beta))
fun(x => (alpha * x) + (t.`2` * beta))
))
))
}
Expand All @@ -134,10 +134,10 @@ object gemv {
)((mat, xs, ys, alpha, beta) =>
zip(mat)(ys) |>
mapPar(fun(t =>
zip(xs)(t._1) |>
zip(xs)(t.`1`) |>
split(n) |>
toMemFun(mapSeq(reduceSeq(fun(a => fun(x => mult(x) + a)))(lf32(0.0f)))) |>
mapSeq(fun(x => (alpha * x) + (t._2 * beta)))
mapSeq(fun(x => (alpha * x) + (t.`2` * beta)))
)) |> join
))
}
Expand Down Expand Up @@ -208,10 +208,10 @@ object gemv {
val localSize = cgo17_localSize
val globalSize = GlobalSize(M)

val run = kernel.as[ScalaFunction `(`
val run = kernel.as[Args `(`
Int `,` Int `,` Array[Array[Float]] `,`
Array[Float] `,` Array[Float] `,` Float `,` Float
`)=>` Array[Float]]
Array[Float] `,` Array[Float] `,` Float `,` Float,
Array[Float]]
run(localSize, globalSize)(N `,` M `,` mat `,` xs `,` ys `,` alpha `,` beta)
}
}
Loading