// One wave per output row. Instantiated per projection by linear.rs.
%offset = index.constant 0 : offset
%zero = index.constant 0 : index
%four = index.constant 4 : index
%eight = index.constant 8 : index
%step = index.constant 128 : index
%zf = scalar.constant 0.0 : f32
%width = scalar.constant 32 : i32
%lane = kernel.subgroup.lane.id : index
%wave = kernel.subgroup.id : index
%group = kernel.workgroup.id<x> : index
%n = config.get @hrx.linear.n@P@ : index
%k = config.get @hrx.linear.k : index
%wn = config.get @hrx.linear.wn@P@ : index
%xg = buffer.assume.memory_space<global> %x : buffer
%wg = buffer.assume.memory_space<global> %w@P@ : buffer
%og = buffer.assume.memory_space<global> %out@P@ : buffer
%xv = buffer.view %xg[%offset] : buffer -> view<[%k]xbf16>
%wv = buffer.view %wg[%offset] : buffer -> view<[%wn]xbf16>
%ov = buffer.view %og[%offset] : buffer -> view<[%n]x@TYPE@>
%r0 = index.mul %group, %eight : index
%row = index.add %r0, %wave : index
%valid = index.cmp ult, %row, %n : index
scf.if %valid {
%base = index.mul %row, %k : index
%start = index.mul %lane, %four : index
%km4 = index.sub %k, %four : index
%wm4 = index.sub %wn, %four : index
%s0, %s1, %s2, %s3 = scf.for %col = [%start to %k step %step](%a0 = %zf : f32, %a1 = %zf : f32, %a2 = %zf : f32, %a3 = %zf : f32) -> (f32, f32, f32, f32) {
%xi = index.assume %col [range(%col, 0, 1073741824), le(%col, %km4), mul(%col, 4)] : index
%wi0 = index.add %base, %col : index
%wi = index.assume %wi0 [range(%wi0, 0, 1073741824), le(%wi0, %wm4), mul(%wi0, 4)] : index
%xs = vector.load %xv[%xi] : view<[%k]xbf16> -> vector<4xbf16>
%ws = vector.load %wv[%wi] : view<[%wn]xbf16> -> vector<4xbf16>
%xb0 = vector.extract %xs[0] : vector<4xbf16> -> bf16
%wb0 = vector.extract %ws[0] : vector<4xbf16> -> bf16
%xf0 = scalar.extf %xb0 : bf16 to f32
%wf0 = scalar.extf %wb0 : bf16 to f32
%p0 = scalar.mulf %xf0, %wf0 : f32
%v0 = scalar.addf %a0, %p0 : f32
%xb1 = vector.extract %xs[1] : vector<4xbf16> -> bf16
%wb1 = vector.extract %ws[1] : vector<4xbf16> -> bf16
%xf1 = scalar.extf %xb1 : bf16 to f32
%wf1 = scalar.extf %wb1 : bf16 to f32
%p1 = scalar.mulf %xf1, %wf1 : f32
%v1 = scalar.addf %a1, %p1 : f32
%xb2 = vector.extract %xs[2] : vector<4xbf16> -> bf16
%wb2 = vector.extract %ws[2] : vector<4xbf16> -> bf16
%xf2 = scalar.extf %xb2 : bf16 to f32
%wf2 = scalar.extf %wb2 : bf16 to f32
%p2 = scalar.mulf %xf2, %wf2 : f32
%v2 = scalar.addf %a2, %p2 : f32
%xb3 = vector.extract %xs[3] : vector<4xbf16> -> bf16
%wb3 = vector.extract %ws[3] : vector<4xbf16> -> bf16
%xf3 = scalar.extf %xb3 : bf16 to f32
%wf3 = scalar.extf %wb3 : bf16 to f32
%p3 = scalar.mulf %xf3, %wf3 : f32
%v3 = scalar.addf %a3, %p3 : f32
scf.yield %v0, %v1, %v2, %v3 : f32, f32, f32, f32
}
%s01 = scalar.addf %s0, %s1 : f32
%s23 = scalar.addf %s2, %s3 : f32
%sum = scalar.addf %s01, %s23 : f32
%d16 = scalar.constant 16 : i32
%sh16, %ok16 = kernel.subgroup.shuffle<xor> %sum, %d16, %width : f32, i32, i32
%sum16 = scalar.addf %sum, %sh16 : f32
%d8 = scalar.constant 8 : i32
%sh8, %ok8 = kernel.subgroup.shuffle<xor> %sum16, %d8, %width : f32, i32, i32
%sum8 = scalar.addf %sum16, %sh8 : f32
%d4 = scalar.constant 4 : i32
%sh4, %ok4 = kernel.subgroup.shuffle<xor> %sum8, %d4, %width : f32, i32, i32
%sum4 = scalar.addf %sum8, %sh4 : f32
%d2 = scalar.constant 2 : i32
%sh2, %ok2 = kernel.subgroup.shuffle<xor> %sum4, %d2, %width : f32, i32, i32
%sum2 = scalar.addf %sum4, %sh2 : f32
%d1 = scalar.constant 1 : i32
%sh1, %ok1 = kernel.subgroup.shuffle<xor> %sum2, %d1, %width : f32, i32, i32
%sum1 = scalar.addf %sum2, %sh1 : f32
%leader = index.cmp eq, %lane, %zero : index
scf.if %leader {
%oi = index.assume %row [range(%row, 0, 1073741824), lt(%row, %n)] : index
%result = @CONVERT@
view.store %result, %ov[%oi] : @TYPE@, view<[%n]x@TYPE@>
}
}