MEP 8: sort版メディアンフィルタのスクリプト検討

Created: 2022-12-07 14:55:23

ヒストグラムを使うフィルタはレジスタを大量に消費するのでGPGPU的では無い。 sortを使ったもっとGPGPUっぽい実装を考える。

sortがある前提でどういう実装ができるかを考える。

weightの総和の要素数の配列を埋めていく、と考えて、各indexから対応するix, iyが求まればいいんだよな。以下のweightのケースを考える。

1, 2, 1
2, 3, 2
1, 2, 1

この場合、各ピクセルの値をこの重みだけ繰り返してsortしたメディアンを求めたい。

繰り返されたピクセルの総数は4+7+4=15個。逆に考えてこの15個の要素を埋めると考えると、0から14を動くindexに対して、それぞれどのix, iyかを求める方法があれば、繰り返した配列が作れる。

えーと、(0, 0) (0, 1) (0, 1) (0, 2) (1, 0) (1, 0) … と並べてやれば、あとはそれをindexで引けば良いのか。この配列をどうやって作ればいいんだろう。なんかweightのcumsumから出そうな気がするが…cumsumが上回る最初のindexを探して、そのindexをweight.extent(0)のモジュロでixが、割るとiyが出るな。

  @bounds(9)
  def wcumsum |i| {
     let [ix, iy] = [i%3, i/3]
     weight(ix, iy)
  }

  mut! trans<wcumsum>.cumsum(dim=0)

として作ったwcumsumに対して、

  @bounds(15)
  def wmat |i| {
     def i3 by reduce<wcumsum>.find_first_index |_, val| { i < val }

find_first_indexがdef by reduceなのは MEP 6: find_first_index を参照の事。

で元のピクセルのix, iyを出す事ができるか。

     let [ix, iy] = [i3%weight.extent(0), i3/weight.extent(0)]

だから重みだけピクセルを繰り返した配列は以下のように作れる。

  @bounds(15)
  def wmat |i| {
     def i3 by reduce<wcumsum>.find_first_index |_, val| { i < val }
     let [a, r, g, b] = input(ix+x, iy+y)
     [a, r, g, b]
  }

ここで、15はハードコードするしか無い。例えば以下のように求める事はできるのだが、

let wsum = weight.sum |x, y, val| { val }

wmatはローカルの配列なので定数である必要がある。厳密にはコード生成の時に値を埋め込む事も可能だけれど、それは将来コード生成だけを先にしてバイナリで持っておくのを不可能にするので、あまりそういう将来に禍根を残しそうな機能は入れたくない。

以上から、とりあえず以下のように実装ができる。

# 1, 2, 1
# 2, 3, 2
# 1, 2, 1
@bounds(3, 3)
def weight |x, y| {
  3 - abs(x-1) - abs(y-1)
}

# xとyがmedianを求められるように、
# x: 1からw-1
# y: 1からh-1
# の範囲で計算する。
@bounds(input.extent(0)-2, input.extent(1)-2)
def median |x, y| {

  @bounds(9)
  def wcumsum |i| {
     let [ix, iy] = [i%3, i/3]
     weight(ix, iy)
  }

  mut! trans<wcumsum>.cumsum!(dim=0)

  # @bounds(wsum) まだIMMしかサポートしてない。ここは見直しても良い気がする。
  @bounds(15)
  def wmat |i| {
     def i3 by reduce<wcumsum>.find_first_index(dim=0) |_, val | { i < val }
     let [ix, iy] = [i3%weight.extent(0), i3/weight.extent(0)]
     let [a, r, g, b] = input(ix+x, iy+y)
     [a, r, g, b]
  }
  
  mut! trans<wmat>.sort!(dim=0)

  let [ma, mr, mg, mb] = wmat(wmat.extent(0)/2)
  argb( ma, mr, mg, mb )
}

def result |x, y| {
  ifel( x == 0 || y == 0 || x == input.extent(0)-1 || y == input.extent(1) -1,
input(x, y),
median(x-1, y-1))
}

wcumsumを求めるのはちょっとトリッキーに思う。

以下の文は、ようするにreshapeで一次元にしてcumsumを求めたいというもの。

  @bounds(9)
  def wcumsum |i| {
     let [ix, iy] = [i%3, i/3]
     weight(ix, iy)
  }

  mut! trans<wcumsum>.cumsum!(dim=0)

より直接的には以下のように書きたい。

  let weight_1d = weight.reshape(9)
  def wcumsum by trans<weight_1d>.cumsum(dim=0)

このように書くには、以下の２つが必要。

reshapeが必要。しかも結果をletで変数に入れられる
def byにtransを対応する必要がある（今はdef byはreduceだけで、transはmut!の自身を変更するもののみ）

1を考えるに、IRElemとしてshape情報と元のテンソルの情報を持てば良さそう。これはlowerでだいたい問題なく扱える気はする。 2もそのうちやりたい気はするので、こういうふうに書けるようにしたい気はする。ただ今の所他の書き方で書ける事は後回しにして実際にいろんなフィルタを動かしてNYIを潰していきたい。

reshapeについての検討メモ

バッファ名とboundsの名前を区別する必要がありそう。パース時にBoundsInfoに登録しつつReferenceにも登録すれば良いか？ Referenceは少し注意がいるかも。これまで名前からテンソルを引いていた所がScopeから探す必要があるかも。

IRElemのReshapedTensorを作る
テストケースを書く
Lowerで処理（boundsをletにするなど）