From 9e671f561a7966970419ed02bddde44de3f5bf36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A1n=20Castell=C3=B3?= Date: Thu, 29 May 2025 09:44:36 +0200 Subject: [PATCH] rvv backend updated with fp16 intrinsics --- src/exo/platforms/rvv.py | 375 +++++++++++++++++++++++++++++++++++++-- 1 file changed, 362 insertions(+), 13 deletions(-) diff --git a/src/exo/platforms/rvv.py b/src/exo/platforms/rvv.py index 72bb883a7..cfbf7077f 100644 --- a/src/exo/platforms/rvv.py +++ b/src/exo/platforms/rvv.py @@ -1,7 +1,7 @@ from __future__ import annotations from exo import Memory, DRAM, instr - +import os def _is_const_size(sz, c): return sz.isdecimal() and int(sz) == c @@ -12,7 +12,7 @@ def _is_some_const_size(sz): # --------------------------------------------------------------------------- # -# Neon registers +# RVV registers # --------------------------------------------------------------------------- # @@ -29,11 +29,19 @@ def can_read(cls): def alloc(cls, new_name, prim_type, shape, srcinfo): if not shape: raise MemGenError(f"{srcinfo}: RVV vectors are not scalar values") - + # This factor will be used to scale the vector length in case of RVV bigger than 128 bits + factor = 1 # Default to 128 bits + try: + if int(os.environ['RVV_BITS']) > 0: + factor = int(os.environ['RVV_BITS'])/128 + except: + factor = 1 + vec_types = { - "float": (4, "vfloat32m1_t") - } # , "double": (2, "float64x2_t"), "_Float16" : (8, "float16x8_t")} - + "float": (4*factor, "vfloat32m1_t"), + "double": (2*factor, "vfloat64m1_t"), + "_Float16" : (8*factor, "vfloat16m1_t")} + if not prim_type in vec_types.keys(): raise MemGenError(f"{srcinfo}: RVV vectors must be f32 (for now)") @@ -75,12 +83,8 @@ def window(cls, basetyp, baseptr, indices, strides, srcinfo): # f32 RVV intrinsics # --------------------------------------------------------------------------- # -# -# Load, Store, Broadcast, FMAdd, Mul, Add? -# -# float32 - +# RVV intrinsics for load f32 vectors with 128 bits @instr("{dst_data} = __riscv_vle32_v_f32m1(&{src_data},{vl});") def rvv_vld_4xf32(dst: [f32][4] @ RVV, src: [f32][4] @ DRAM, vl: size): assert stride(src, 0) == 1 @@ -91,7 +95,19 @@ def rvv_vld_4xf32(dst: [f32][4] @ RVV, src: [f32][4] @ DRAM, vl: size): for i in seq(0, vl): dst[i] = src[i] +# RVV intrinsics for load f32 vectors with 256 bits +@instr("{dst_data} = __riscv_vle32_v_f32m1(&{src_data},{vl});") +def rvv_vld_8xf32(dst: [f32][8] @ RVV, src: [f32][8] @ DRAM, vl: size): + assert stride(src, 0) == 1 + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = src[i] + +# RVV intrinsics for store f32 vectors with 128 bits @instr("__riscv_vse32_v_f32m1(&{dst_data}, {src_data},{vl});") def rvv_vst_4xf32(dst: [f32][4] @ DRAM, src: [f32][4] @ RVV, vl: size): assert stride(src, 0) == 1 @@ -102,7 +118,18 @@ def rvv_vst_4xf32(dst: [f32][4] @ DRAM, src: [f32][4] @ RVV, vl: size): for i in seq(0, vl): dst[i] = src[i] +# RVV intrinsics for store f32 vectors with 256 bits +@instr("__riscv_vse32_v_f32m1(&{dst_data}, {src_data},{vl});") +def rvv_vst_8xf32(dst: [f32][8] @ DRAM, src: [f32][8] @ RVV, vl: size): + assert stride(src, 0) == 1 + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = src[i] +# RVV intrinsics for broadcast f32 scalar (in buffer) into f32 vector with 128 bits @instr("{dst_data} = __riscv_vfmv_v_f_f32m1({src_data},{vl});") def rvv_broadcast_4xf32(dst: [f32][4] @ RVV, src: [f32][1] @ DRAM, vl: size): assert stride(dst, 0) == 1 @@ -112,7 +139,17 @@ def rvv_broadcast_4xf32(dst: [f32][4] @ RVV, src: [f32][1] @ DRAM, vl: size): for i in seq(0, vl): dst[i] = src[0] +# RVV intrinsics for broadcast f32 scalar (in buffer) into f32 vector with 256 bits +@instr("{dst_data} = __riscv_vfmv_v_f_f32m1({src_data},{vl});") +def rvv_broadcast_8xf32(dst: [f32][8] @ RVV, src: [f32][1] @ DRAM, vl: size): + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = src[0] +# RVV intrinsics for broadcast f32 scalar (variable) into f32 vector with 128 bits @instr("{dst_data} = __riscv_vfmv_v_f_f32m1({src_data},{vl});") def rvv_broadcast_4xf32_scalar(dst: [f32][4] @ RVV, src: f32 @ DRAM, vl: size): assert stride(dst, 0) == 1 @@ -122,7 +159,18 @@ def rvv_broadcast_4xf32_scalar(dst: [f32][4] @ RVV, src: f32 @ DRAM, vl: size): for i in seq(0, vl): dst[i] = src +# RVV intrinsics for broadcast f32 scalar (variable) into f32 vector with 256 bits +@instr("{dst_data} = __riscv_vfmv_v_f_f32m1({src_data},{vl});") +def rvv_broadcast_8xf32_scalar(dst: [f32][8] @ RVV, src: f32 @ DRAM, vl: size): + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = src + +# RVV intrinsics for broadcast f32 0.0 value into f32 vector with 128 bits @instr("{dst_data} = __riscv_vfmv_v_f_f32m1(0.0f,{vl});") def rvv_broadcast_4xf32_0(dst: [f32][4] @ RVV, vl: size): assert stride(dst, 0) == 1 @@ -132,7 +180,17 @@ def rvv_broadcast_4xf32_0(dst: [f32][4] @ RVV, vl: size): for i in seq(0, vl): dst[i] = 0.0 +# RVV intrinsics for broadcast f32 0.0 value into f32 vector with 256 bits +@instr("{dst_data} = __riscv_vfmv_v_f_f32m1(0.0f,{vl});") +def rvv_broadcast_8xf32_0(dst: [f32][8] @ RVV, vl: size): + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = 0.0 +# RVV intrinsics for fused multiply-add of f32 vectors with 128 bits @instr("{dst_data} = __riscv_vfmacc_vv_f32m1({dst_data}, {lhs_data}, {rhs_data},{vl});") def rvv_vfmacc_4xf32_4xf32( dst: [f32][4] @ RVV, lhs: [f32][4] @ RVV, rhs: [f32][4] @ RVV, vl: size @@ -146,8 +204,22 @@ def rvv_vfmacc_4xf32_4xf32( for i in seq(0, vl): dst[i] += lhs[i] * rhs[i] +# RVV intrinsics for fused multiply-add of f32 vectors with 256 bits +@instr("{dst_data} = __riscv_vfmacc_vv_f32m1({dst_data}, {lhs_data}, {rhs_data},{vl});") +def rvv_vfmacc_8xf32_8xf32( + dst: [f32][8] @ RVV, lhs: [f32][8] @ RVV, rhs: [f32][8] @ RVV, vl: size +): + assert stride(dst, 0) == 1 + assert stride(lhs, 0) == 1 + assert stride(rhs, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] += lhs[i] * rhs[i] -@instr("{dst_data} = __riscv_vfmacc_vf_f32m1{dst_data}, {rhs_data}, {lhs_data},{vl});") +# RVV intrinsics for fused multiply-add of f32 vector and scalar with 128 bits +@instr("{dst_data} = __riscv_vfmacc_vf_f32m1({dst_data}, {rhs_data}, {lhs_data},{vl});") def rvv_vfmacc_4xf32_1xf32( dst: [f32][4] @ RVV, lhs: [f32][4] @ RVV, rhs: [f32][1] @ DRAM, vl: size ): @@ -160,8 +232,23 @@ def rvv_vfmacc_4xf32_1xf32( for i in seq(0, vl): dst[i] += lhs[i] * rhs[0] +# RVV intrinsics for fused multiply-add of f32 vector and scalar with 256 bits +@instr("{dst_data} = __riscv_vfmacc_vf_f32m1({dst_data}, {rhs_data}, {lhs_data},{vl});") +def rvv_vfmacc_8xf32_1xf32( + dst: [f32][8] @ RVV, lhs: [f32][8] @ RVV, rhs: [f32][1] @ DRAM, vl: size +): + assert stride(dst, 0) == 1 + assert stride(lhs, 0) == 1 + assert stride(rhs, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] += lhs[i] * rhs[0] + -@instr("{dst_data} = __riscv_vfmacc_vf_f32m1{dst_data}, {lhs_data}, {rhs_data},{vl});") +# RVV intrinsics for fused multiply-add of scalar and f32 vector with 128 bits +@instr("{dst_data} = __riscv_vfmacc_vf_f32m1({dst_data}, {lhs_data}, {rhs_data},{vl});") def rvv_vfmacc_1xf32_4xf32( dst: [f32][4] @ RVV, lhs: [f32][1] @ DRAM, rhs: [f32][4] @ RVV, vl: size ): @@ -173,3 +260,265 @@ def rvv_vfmacc_1xf32_4xf32( for i in seq(0, vl): dst[i] += lhs[0] * rhs[i] + +# RVV intrinsics for fused multiply-add of scalar and f32 vector with 256 bits +@instr("{dst_data} = __riscv_vfmacc_vf_f32m1({dst_data}, {lhs_data}, {rhs_data},{vl});") +def rvv_vfmacc_1xf32_8xf32( + dst: [f32][8] @ RVV, lhs: [f32][1] @ DRAM, rhs: [f32][8] @ RVV, vl: size +): + assert stride(dst, 0) == 1 + assert stride(lhs, 0) == 1 + assert stride(rhs, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] += lhs[0] * rhs[i] + +# RVV intrinsics for gather elements from f32 vector with 128 bits +@instr("{dst_data} = __riscv_vrgather_vx_f32m1({src_data}, {imm}, {vl});") +def rvv_gather_4xf32(dst: [f32][4] @ RVV, src: [f32][4] @ RVV, imm: index, vl: size): + assert stride(dst, 0) == 1 + assert stride(src, 0) == 1 + assert imm >= 0 + assert imm < 4 + assert vl >= 0 + assert vl <= 4 + + for i in seq(0, vl): + dst[i] = src[imm] + +# RVV intrinsics for gather elements from f32 vector with 256 bits +@instr("{dst_data} = __riscv_vrgather_vx_f32m1({src_data}, {imm}, {vl});") +def rvv_gather_8xf32(dst: [f32][8] @ RVV, src: [f32][8] @ RVV, imm: index, vl: size): + assert stride(dst, 0) == 1 + assert stride(src, 0) == 1 + assert imm >= 0 + assert imm < 8 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = src[imm] + +# --------------------------------------------------------------------------- # +# f16 RVV intrinsics +# --------------------------------------------------------------------------- # + +# RVV intrinsics for load f16 vectors with 128 bits +@instr("{dst_data} = __riscv_vle16_v_f16m1(&{src_data},{vl});") +def rvv_vld_8xf16(dst: [f16][8] @ RVV, src: [f16][8] @ DRAM, vl: size): + assert stride(src, 0) == 1 + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = src[i] + +# RVV intrinsics for load f16 vectors with 256 bits +@instr("{dst_data} = __riscv_vle16_v_f16m1(&{src_data},{vl});") +def rvv_vld_16xf16(dst: [f16][16] @ RVV, src: [f16][16] @ DRAM, vl: size): + assert stride(src, 0) == 1 + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 16 + + for i in seq(0, vl): + dst[i] = src[i] + + +# RVV intrinsics for store f16 vectors with 128 bits +@instr("__riscv_vse16_v_f16m1(&{dst_data}, {src_data},{vl});") +def rvv_vst_8xf16(dst: [f16][8] @ DRAM, src: [f16][8] @ RVV, vl: size): + assert stride(src, 0) == 1 + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = src[i] + +# RVV intrinsics for store f16 vectors with 256 bits +@instr("__riscv_vse16_v_f16m1(&{dst_data}, {src_data},{vl});") +def rvv_vst_16xf16(dst: [f16][16] @ DRAM, src: [f16][16] @ RVV, vl: size): + assert stride(src, 0) == 1 + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 16 + + for i in seq(0, vl): + dst[i] = src[i] + +# RVV intrinsics for broadcast f32 scalar (in buffer) into f16 vector with 128 bits +@instr("{dst_data} = __riscv_vfmv_v_f_f16m1({src_data},{vl});") +def rvv_broadcast_8xf16(dst: [f16][8] @ RVV, src: [f16][1] @ DRAM, vl: size): + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = src[0] + +# RVV intrinsics for broadcast f32 scalar (in buffer) into f16 vector with 256 bits +@instr("{dst_data} = __riscv_vfmv_v_f_f16m1({src_data},{vl});") +def rvv_broadcast_16xf16(dst: [f16][16] @ RVV, src: [f16][1] @ DRAM, vl: size): + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 16 + + for i in seq(0, vl): + dst[i] = src[0] + +# RVV intrinsics for broadcast f32 scalar (variable) into f16 vector with 128 bits +@instr("{dst_data} = __riscv_vfmv_v_f_f16m1({src_data},{vl});") +def rvv_broadcast_8xf16_scalar(dst: [f16][8] @ RVV, src: f16 @ DRAM, vl: size): + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = src + +# RVV intrinsics for broadcast f32 scalar (variable) into f16 vector with 256 bits +@instr("{dst_data} = __riscv_vfmv_v_f_f16m1({src_data},{vl});") +def rvv_broadcast_16xf16_scalar(dst: [f16][16] @ RVV, src: f16 @ DRAM, vl: size): + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 16 + + for i in seq(0, vl): + dst[i] = src + +# RVV intrinsics for broadcast f32 0.0 value into f16 vector with 128 bits +@instr("{dst_data} = __riscv_vfmv_v_f_f16m1(0.0f,{vl});") +def rvv_broadcast_8xf16_0(dst: [f16][8] @ RVV, vl: size): + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = 0.0 + +# RVV intrinsics for broadcast f32 0.0 value into f16 vector with 256 bits +@instr("{dst_data} = __riscv_vfmv_v_f_f16m1(0.0f,{vl});") +def rvv_broadcast_16xf16_0(dst: [f16][16] @ RVV, vl: size): + assert stride(dst, 0) == 1 + assert vl >= 0 + assert vl <= 16 + + for i in seq(0, vl): + dst[i] = 0.0 + +# RVV intrinsics for fused multiply-add of f16 vectors with 128 bits +@instr("{dst_data} = __riscv_vfmacc_vv_f16m1({dst_data}, {lhs_data}, {rhs_data},{vl});") +def rvv_vfmacc_8xf16_8xf16( + dst: [f16][8] @ RVV, lhs: [f16][8] @ RVV, rhs: [f16][8] @ RVV, vl: size +): + assert stride(dst, 0) == 1 + assert stride(lhs, 0) == 1 + assert stride(rhs, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] += lhs[i] * rhs[i] + + +# RVV intrinsics for fused multiply-add of f16 vectors with 256 bits +@instr("{dst_data} = __riscv_vfmacc_vv_f16m1({dst_data}, {lhs_data}, {rhs_data},{vl});") +def rvv_vfmacc_16xf16_16xf16( + dst: [f16][16] @ RVV, lhs: [f16][16] @ RVV, rhs: [f16][16] @ RVV, vl: size +): + assert stride(dst, 0) == 1 + assert stride(lhs, 0) == 1 + assert stride(rhs, 0) == 1 + assert vl >= 0 + assert vl <= 16 + + for i in seq(0, vl): + dst[i] += lhs[i] * rhs[i] + + +# RVV intrinsics for fused multiply-add of f16 vector and scalar with 128 bits +@instr("{dst_data} = __riscv_vfmacc_vf_f16m1({dst_data}, {rhs_data}, {lhs_data},{vl});") +def rvv_vfmacc_8xf16_1xf16( + dst: [f16][8] @ RVV, lhs: [f16][8] @ RVV, rhs: [f16][1] @ DRAM, vl: size +): + assert stride(dst, 0) == 1 + assert stride(lhs, 0) == 1 + assert stride(rhs, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] += lhs[i] * rhs[0] + +# RVV intrinsics for fused multiply-add of f16 vector and scalar with 256 bits +@instr("{dst_data} = __riscv_vfmacc_vf_f16m1({dst_data}, {rhs_data}, {lhs_data},{vl});") +def rvv_vfmacc_16xf16_1xf16( + dst: [f16][16] @ RVV, lhs: [f16][16] @ RVV, rhs: [f16][1] @ DRAM, vl: size +): + assert stride(dst, 0) == 1 + assert stride(lhs, 0) == 1 + assert stride(rhs, 0) == 1 + assert vl >= 0 + assert vl <= 16 + + for i in seq(0, vl): + dst[i] += lhs[i] * rhs[0] + +# RVV intrinsics for fused multiply-add of scalar and f16 vector with 128 bits +@instr("{dst_data} = __riscv_vfmacc_vf_f16m1({dst_data}, {lhs_data}, {rhs_data},{vl});") +def rvv_vfmacc_1xf16_8xf16( + dst: [f16][8] @ RVV, lhs: [f16][1] @ DRAM, rhs: [f16][8] @ RVV, vl: size +): + assert stride(dst, 0) == 1 + assert stride(lhs, 0) == 1 + assert stride(rhs, 0) == 1 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] += lhs[0] * rhs[i] + +# RVV intrinsics for fused multiply-add of scalar and f16 vector with 256 bits +@instr("{dst_data} = __riscv_vfmacc_vf_f16m1({dst_data}, {lhs_data}, {rhs_data},{vl});") +def rvv_vfmacc_1xf16_16xf16( + dst: [f16][16] @ RVV, lhs: [f16][1] @ DRAM, rhs: [f16][16] @ RVV, vl: size +): + assert stride(dst, 0) == 1 + assert stride(lhs, 0) == 1 + assert stride(rhs, 0) == 1 + assert vl >= 0 + assert vl <= 16 + + for i in seq(0, vl): + dst[i] += lhs[0] * rhs[i] + + +# RVV intrinsics for gather elements from f16 vector with 128 bits +@instr("{dst_data} = __riscv_vrgather_vx_f16m1({src_data}, {imm}, {vl});") +def rvv_gather_8xf16(dst: [f16][8] @ RVV, src: [f16][8] @ RVV, imm: index, vl: size): + assert stride(dst, 0) == 1 + assert stride(src, 0) == 1 + assert imm >= 0 + assert imm < 8 + assert vl >= 0 + assert vl <= 8 + + for i in seq(0, vl): + dst[i] = src[imm] + +# RVV intrinsics for gather elements from f16 vector with 256 bits +@instr("{dst_data} = __riscv_vrgather_vx_f16m1({src_data}, {imm}, {vl});") +def rvv_gather_16xf16(dst: [f16][16] @ RVV, src: [f16][16] @ RVV, imm: index, vl: size): + assert stride(dst, 0) == 1 + assert stride(src, 0) == 1 + assert imm >= 0 + assert imm < 16 + assert vl >= 0 + assert vl <= 16 + + for i in seq(0, vl): + dst[i] = src[imm] \ No newline at end of file