[Gemmini Dialect] Gemmini Dialect enhancement on tiled_matmul

backend/include/llvm/IR/IntrinsicsRISCVBuddyExt.td

@@ -18,7 +18,13 @@
 //
 //===----------------------------------------------------------------------===//
 let TargetPrefix = "riscv" in
-def int_riscv_mvin : Intrinsic<[],[llvm_i64_ty, llvm_i64_ty],[]>;
+def int_riscv_mvin : Intrinsic<[], [llvm_i64_ty, llvm_i64_ty],[]>;
+
+let TargetPrefix = "riscv" in
+def int_riscv_mvin2 : Intrinsic<[], [llvm_i64_ty, llvm_i64_ty],[]>;
+
+let TargetPrefix = "riscv" in
+def int_riscv_mvin3 : Intrinsic<[], [llvm_i64_ty, llvm_i64_ty],[]>;
 
 let TargetPrefix = "riscv" in
 def int_riscv_mvout : Intrinsic<[], [llvm_i64_ty, llvm_i64_ty], []>;
@@ -35,6 +41,9 @@ def int_riscv_config_st : Intrinsic<[], [llvm_i64_ty, llvm_i64_ty], []>;
 let TargetPrefix = "riscv" in 
 def int_riscv_config_ex : Intrinsic<[], [llvm_i64_ty, llvm_i64_ty], []>;
 
+let TargetPrefix = "riscv" in 
+def int_riscv_config_norm : Intrinsic<[], [llvm_i64_ty, llvm_i64_ty], []>;
+
 let TargetPrefix = "riscv" in 
 def int_riscv_preload : Intrinsic<[], [llvm_i64_ty, llvm_i64_ty], []>;
 

backend/llvm/lib/Target/RISCV/RISCVInstrInfoBuddyExt.td

@@ -35,6 +35,18 @@ def MVIN : RVInstR<0b0000010, 0b011, OPC_CUSTOM_3, (outs),
   let rd = 0;              
 }
 
+let hasSideEffects = 1, mayLoad = 1, mayStore = 1, Predicates = [HasBuddyExt] in
+def MVIN2 : RVInstR<0b0000001, 0b011, OPC_CUSTOM_3, (outs),
+                   (ins GPR:$rs1, GPR:$rs2), "mvin2","$rs1, $rs2"> {
+  let rd = 0;
+}
+
+let hasSideEffects = 1, mayLoad = 1, mayStore = 1, Predicates = [HasBuddyExt] in
+def MVIN3 : RVInstR<0b0001110, 0b011, OPC_CUSTOM_3, (outs),
+                   (ins GPR:$rs1, GPR:$rs2), "mvin3","$rs1, $rs2"> {
+  let rd = 0;
+}
+
 let hasSideEffects = 1, mayLoad = 1, mayStore = 1, Predicates = [HasBuddyExt] in 
 def MVOUT : RVInstR<0b0000011, 0b011, OPC_CUSTOM_3, (outs),
                     (ins GPR:$rs1, GPR:$rs2), "mvout","$rs1, $rs2">{
@@ -65,6 +77,12 @@ def CONFIG_EX : RVInstR<0b0000000, 0b011, OPC_CUSTOM_3,(outs),
   let rd = 0;
 }
 
+let Predicates = [HasBuddyExt] in 
+def CONFIG_NORM : RVInstR<0b0000000, 0b011, OPC_CUSTOM_3,(outs),
+                    (ins GPR:$rs1, GPR:$rs2), "config_norm", "$rs1, $rs2"> {
+  let rd = 0;
+}
+
 let hasSideEffects = 1, mayLoad = 1, mayStore =1, Predicates = [HasBuddyExt] in 
 def PRELOAD : RVInstR<0b0000110, 0b011,OPC_CUSTOM_3,(outs),
                     (ins GPR:$rs1, GPR:$rs2), "preload", "$rs1, $rs2">{
@@ -164,6 +182,12 @@ def LOOP_CONV_WS_CONFIG6 : RVInstR<0b0010101, 0b011, OPC_CUSTOM_3, (outs),
 let Predicates = [HasBuddyExt] in 
 def : Pat<(int_riscv_mvin GPR:$rs1, GPR:$rs2), (MVIN GPR:$rs1, GPR:$rs2)>;
 
+let Predicates = [HasBuddyExt] in
+def : Pat<(int_riscv_mvin2 GPR:$rs1, GPR:$rs2), (MVIN2 GPR:$rs1, GPR:$rs2)>;
+
+let Predicates = [HasBuddyExt] in
+def : Pat<(int_riscv_mvin3 GPR:$rs1, GPR:$rs2), (MVIN3 GPR:$rs1, GPR:$rs2)>;
+
 let Predicates = [HasBuddyExt] in 
 def : Pat<(int_riscv_mvout GPR:$rs1, GPR:$rs2), (MVOUT GPR:$rs1, GPR:$rs2)>;
 
@@ -179,6 +203,9 @@ def : Pat<(int_riscv_config_st GPR:$rs1, GPR:$rs2), (CONFIG_ST GPR:$rs1, GPR:$rs
 let Predicates = [HasBuddyExt] in 
 def : Pat<(int_riscv_config_ex GPR:$rs1, GPR:$rs2), (CONFIG_EX GPR:$rs1, GPR:$rs2)>;
 
+let Predicates = [HasBuddyExt] in
+def : Pat<(int_riscv_config_norm GPR:$rs1, GPR:$rs2), (CONFIG_NORM GPR:$rs1, GPR:$rs2)>;
+
 let Predicates = [HasBuddyExt] in 
 def : Pat<(int_riscv_preload GPR:$rs1, GPR:$rs2), (PRELOAD GPR:$rs1, GPR:$rs2)>;
 

examples/GemminiDialect/ciface.mlir

@@ -127,53 +127,53 @@ func.func @linalg_conv6(%arg0 : memref<1x1x256x256xi8>, %arg1 : memref<1x1x13x13
 func.func @gemmini_conv1(%input: memref<1x256x256x1xi8>, %weights: memref<9x1xi8>, %bias: memref<1xi32>, %output: memref<64516x1xi8>) {
   %outdim = arith.constant 254 : i64 
   %kernelDim = arith.constant 3 : i64
-  gemmini.tile_conv %input %weights %bias %output %outdim %kernelDim {stride = 1} : 
-  memref<1x256x256x1xi8> memref<9x1xi8> memref<1xi32> memref<64516x1xi8> i64 i64
+  gemmini.tile_conv %input %weights %bias %output %outdim %outdim %kernelDim {stride = 1} :
+  memref<1x256x256x1xi8> memref<9x1xi8> memref<1xi32> memref<64516x1xi8> i64 i64 i64
   return 
 }
 
 // CHECK: llvm.func @_mlir_ciface_gemmini_conv2
 func.func @gemmini_conv2(%input: memref<1x256x256x1xi8>, %weights: memref<25x1xi8>, %bias: memref<1xi32>, %output: memref<63504x1xi8>) {
   %outdim = arith.constant 252 : i64 
   %kernelDim = arith.constant 5 : i64
-  gemmini.tile_conv %input %weights %bias %output %outdim %kernelDim {stride = 1} : 
-  memref<1x256x256x1xi8> memref<25x1xi8> memref<1xi32> memref<63504x1xi8> i64 i64
+  gemmini.tile_conv %input %weights %bias %output %outdim %outdim %kernelDim {stride = 1} :
+  memref<1x256x256x1xi8> memref<25x1xi8> memref<1xi32> memref<63504x1xi8> i64 i64 i64
   return 
 }
 
 // CHECK: llvm.func @_mlir_ciface_gemmini_conv3
 func.func @gemmini_conv3(%input: memref<1x256x256x1xi8>, %weights: memref<49x1xi8>, %bias: memref<1xi32>, %output: memref<62500x1xi8>) {
   %outdim = arith.constant 250 : i64 
   %kernelDim = arith.constant 7 : i64
-  gemmini.tile_conv %input %weights %bias %output %outdim %kernelDim {stride = 1} : 
-  memref<1x256x256x1xi8> memref<49x1xi8> memref<1xi32> memref<62500x1xi8> i64 i64
+  gemmini.tile_conv %input %weights %bias %output %outdim %outdim %kernelDim {stride = 1} :
+  memref<1x256x256x1xi8> memref<49x1xi8> memref<1xi32> memref<62500x1xi8> i64 i64 i64
   return 
 }
 
 // CHECK: llvm.func @_mlir_ciface_gemmini_conv4
 func.func @gemmini_conv4(%input: memref<1x256x256x1xi8>, %weights: memref<81x1xi8>, %bias: memref<1xi32>, %output: memref<61504x1xi8>) {
   %outdim = arith.constant 248 : i64 
   %kernelDim = arith.constant 9 : i64
-  gemmini.tile_conv %input %weights %bias %output %outdim %kernelDim {stride = 1} : 
-  memref<1x256x256x1xi8> memref<81x1xi8> memref<1xi32> memref<61504x1xi8> i64 i64
+  gemmini.tile_conv %input %weights %bias %output %outdim %outdim %kernelDim {stride = 1} :
+  memref<1x256x256x1xi8> memref<81x1xi8> memref<1xi32> memref<61504x1xi8> i64 i64 i64
   return 
 }
 
 // CHECK: llvm.func @_mlir_ciface_gemmini_conv5
 func.func @gemmini_conv5(%input: memref<1x256x256x1xi8>, %weights: memref<121x1xi8>, %bias: memref<1xi32>, %output: memref<60516x1xi8>) {
   %outdim = arith.constant 246 : i64 
   %kernelDim = arith.constant 11 : i64
-  gemmini.tile_conv %input %weights %bias %output %outdim %kernelDim {stride = 1} : 
-  memref<1x256x256x1xi8> memref<121x1xi8> memref<1xi32> memref<60516x1xi8> i64 i64
+  gemmini.tile_conv %input %weights %bias %output %outdim %outdim %kernelDim {stride = 1} :
+  memref<1x256x256x1xi8> memref<121x1xi8> memref<1xi32> memref<60516x1xi8> i64 i64 i64
   return 
 }
 
 // CHECK: llvm.func @_mlir_ciface_gemmini_conv6
 func.func @gemmini_conv6(%input: memref<1x256x256x1xi8>, %weights: memref<169x1xi8>, %bias: memref<1xi32>, %output: memref<59536x1xi8>) {
   %outdim = arith.constant 244 : i64 
   %kernelDim = arith.constant 13 : i64
-  gemmini.tile_conv %input %weights %bias %output %outdim %kernelDim {stride = 1} : 
-  memref<1x256x256x1xi8> memref<169x1xi8> memref<1xi32> memref<59536x1xi8> i64 i64
+  gemmini.tile_conv %input %weights %bias %output %outdim %outdim %kernelDim {stride = 1} :
+  memref<1x256x256x1xi8> memref<169x1xi8> memref<1xi32> memref<59536x1xi8> i64 i64 i64
   return 
 }
 

examples/GemminiDialect/makefile

@@ -76,6 +76,24 @@ tile-matmul-run:
 	@riscv64-unknown-linux-gnu-gcc log.o -O2 -static -o a.out
 	@spike --extension=gemmini pk a.out
 
+tile-matmul-os-run:
+	@${BUDDY_OPT} ./tile-matmul-os.mlir -lower-gemmini | \
+	${BUDDY_TRANSLATE} --buddy-to-llvmir | \
+	${BUDDY_LLC} -filetype=obj -mtriple=riscv64 \
+		-mattr=+buddyext,+D -float-abi=hard \
+		-o log.o
+	@riscv64-unknown-linux-gnu-gcc log.o -O2 -static -o a.out
+	@spike --extension=gemmini pk a.out
+
+tile-matmul-ws-igelu-run:
+	@${BUDDY_OPT} ./tile-matmul-ws-igelu.mlir -lower-gemmini | \
+	${BUDDY_TRANSLATE} --buddy-to-llvmir | \
+	${BUDDY_LLC} -filetype=obj -mtriple=riscv64 \
+		-mattr=+buddyext,+D -float-abi=hard \
+		-o log.o
+	@riscv64-unknown-linux-gnu-gcc log.o -O2 -static -o a.out
+	@spike --extension=gemmini pk a.out
+
 tile-conv-run:
 	@${BUDDY_OPT} ./tile-conv.mlir -lower-gemmini | \
 	${BUDDY_TRANSLATE} --buddy-to-llvmir | \
@@ -85,6 +103,15 @@ tile-conv-run:
 	@riscv64-unknown-linux-gnu-gcc log.o -O2 -static -o a.out
 	@spike --extension=gemmini pk a.out
 
+tile-rect-conv-run:
+	@${BUDDY_OPT} ./tile-rect-conv.mlir -lower-gemmini | \
+	${BUDDY_TRANSLATE} --buddy-to-llvmir | \
+	${BUDDY_LLC} -filetype=obj -mtriple=riscv64 \
+		-mattr=+buddyext,+D -float-abi=hard \
+		-o log.o
+	@riscv64-unknown-linux-gnu-gcc log.o -O2 -static -o a.out
+	@spike --extension=gemmini pk a.out
+
 gemmini-linalg-matmul-run:
 	@${BUDDY_OPT} ./matmul.mlir \
 		-convert-linalg-to-gemmini \

examples/GemminiDialect/tile-conv.mlir

@@ -32,8 +32,8 @@ func.func @main() -> i64 {
   // CHECK: "gemmini.intr.loop_conv_ws_config6"
   // CHECK: "gemmini.intr.loop_conv_ws"
   // CHECK: "gemmini.intr.flush"
-  gemmini.tile_conv %input %weight %bias %output %3 %3 {stride = 1}: 
-  memref<1x5x5x1xi8> memref<9x2xi8> memref<2xi32> memref<9x2xi8> i64 i64
+  gemmini.tile_conv %input %weight %bias %output %3 %3 %3 {stride = 1}:
+  memref<1x5x5x1xi8> memref<9x2xi8> memref<2xi32> memref<9x2xi8> i64 i64 i64
   gemmini.print %output : memref<9x2xi8>
   return %0 : i64
 }

examples/GemminiDialect/tile-matmul-os.mlir

@@ -0,0 +1,36 @@
+// RUN: buddy-opt %s \
+// RUN:     --lower-gemmini | \
+// RUN: FileCheck %s
+
+func.func @main() -> i8 {
+  %i0 = arith.constant 0 : i8
+  %i1I8 = arith.constant 1 : i8
+  %i2I8 = arith.constant 2 : i8
+  %i2I32 = arith.constant 2 : i32
+  %c0 = arith.constant 0 : index 
+  %c1 = arith.constant 1 : index 
+  %aArray = memref.alloc() {alignment = 16} : memref<64x64xi8>
+  %bArray = memref.alloc() {alignment = 16}: memref<64x64xi8>
+  %cArray = memref.alloc() {alignment = 16}: memref<64x64xi8>
+  %dArray = memref.alloc() {alignment = 64} : memref<64x64xi32>
+  %dim = memref.dim %aArray, %c0 : memref<64x64xi8> 
+  scf.for %i = %c0 to %dim step %c1 {
+    scf.for %j = %c0 to %dim step %c1 {
+      memref.store %i1I8, %aArray[%i, %j] : memref<64x64xi8> 
+      memref.store %i1I8, %bArray[%i, %j] : memref<64x64xi8> 
+      memref.store %i2I32, %dArray[%i, %j] : memref<64x64xi32>
+    }
+  }
+  gemmini.print %aArray : memref<64x64xi8>
+  gemmini.print %bArray : memref<64x64xi8>
+  gemmini.print %dArray : memref<64x64xi32>
+  // CHECK: "gemmini.intr.config_ld"
+  // CHECK: "gemmini.intr.mvin"
+  // CHECK: "gemmini.intr.preload"
+  // CHECK: "gemmini.intr.compute_preloaded"
+  // CHECK: "gemmini.intr.compute_accumulated"
+  // CHECK: "gemmini.intr.mvout"
+  gemmini.tile_matmul  %aArray %bArray %cArray %dArray {dataflow=0} : memref<64x64xi8> memref<64x64xi8> memref<64x64xi8> memref <64x64xi32>
+  gemmini.print %cArray : memref<64x64xi8>
+  return %i0 : i8
+}

examples/GemminiDialect/tile-matmul-ws-igelu.mlir

@@ -0,0 +1,46 @@
+// RUN: buddy-opt %s \
+// RUN:     --lower-gemmini | \
+// RUN: FileCheck %s
+
+memref.global "private" @g1 : memref<3x3xi8> = dense<[[1, 0, 0], [1, -1, 1], [-1, 0, 1]]>
+memref.global "private" @g2 : memref<3x3xi8> = dense<[[1, -1, 0], [1, 0, -1], [-1, -1, 0]]>
+
+func.func @main() -> i8 {
+  %i0 = arith.constant 0 : i8
+  %i1I8 = arith.constant 1 : i8
+  %minus1 = arith.constant -2 : i8
+  %i2I8 = arith.constant 2 : i8
+  %i2I32 = arith.constant 2 : i32
+  %dI32 = arith.constant 0 : i32
+  %c0 = arith.constant 0 : index 
+  %c1 = arith.constant 1 : index
+  %aArray = memref.get_global @g1 : memref<3x3xi8>
+  %bArray = memref.get_global @g2 : memref<3x3xi8>
+  %cArray = memref.alloc()  : memref<3x3xi8>
+  %dArray = memref.alloc()  : memref<3x3xi32>
+  %dim_I = memref.dim %aArray, %c0 : memref<3x3xi8>
+  %dim_J = memref.dim %bArray, %c1 : memref<3x3xi8>
+  %dim_K = memref.dim %aArray, %c1 : memref<3x3xi8>
+
+  scf.for %i3 = %c0 to %dim_I step %c1 {
+    scf.for %j3 = %c0 to %dim_J step %c1 {
+      memref.store %dI32, %dArray[%i3, %j3] : memref<3x3xi32>
+    }
+  }
+  gemmini.print %aArray : memref<3x3xi8>
+  gemmini.print %bArray : memref<3x3xi8>
+  // CHECK: "gemmini.intr.config_ex"
+  // CHECK: "gemmini.intr.config_st"
+  // CHECK: "gemmini.intr.config_ld"
+  // CHECK: "gemmini.intr.config_norm"
+  // CHECK: "gemmini.intr.loop_ws_config_bounds"
+  // CHECK: "gemmini.intr.loop_ws_config_addrs_ab"
+  // CHECK: "gemmini.intr.loop_ws_config_addrs_dc"
+  // CHECK: "gemmini.intr.loop_ws_config_strides_ab"
+  // CHECK: "gemmini.intr.loop_ws_config_strides_dc"
+  // CHECK: "gemmini.intr.loop_ws"
+  // CHECk: "gemmini.intr.flush"
+  gemmini.tile_matmul %aArray %bArray %cArray %dArray {dataflow=1, act=3, bertScale=0.8:f32}: memref<3x3xi8> memref<3x3xi8> memref<3x3xi8> memref<3x3xi32>
+  gemmini.print %cArray : memref<3x3xi8>
+  return %i0 : i8
+}

examples/GemminiDialect/tile-rect-conv.mlir

@@ -0,0 +1,41 @@
+// RUN: buddy-opt %s \
+// RUN:     --lower-gemmini | \
+// RUN: FileCheck %s
+
+// batchSize = 1 inputRowDim = 5 inputColDim inChannels = 2
+memref.global "private" @input : memref<1x5x10x1xi8> = dense<[[[[1], [0], [-1], [0], [1], [1], [0], [-1], [0], [1]],
+                                                              [[1], [0], [-1], [0], [1], [1], [0], [-1], [0], [1]],
+                                                              [[1], [0], [-1], [0], [1], [1], [0], [-1], [0], [1]],
+                                                              [[1], [0], [-1], [0], [1], [1], [0], [-1], [0], [1]],
+                                                              [[1], [0], [-1], [0], [1], [1], [0], [-1], [0], [1]]]]>
+
+// outChannels = 2 kernelDim = 3 inChannels = 1
+memref.global "private" @weight : memref<9x2xi8> = dense<[[1, 2], [1, 2], [1, 2],
+                                                          [1, 2], [1, 2], [1, 2],
+                                                          [1, 2], [1, 2], [1, 2]]>
+
+// outChannels = 2
+memref.global "private" @bias : memref<2xi32> = dense<[1,1]>
+
+func.func @main() -> i64 {
+  %0 = arith.constant 0 : i64
+  %3 = arith.constant 3 : i64
+  %8 = arith.constant 8 : i64
+  %input = memref.get_global @input : memref<1x5x10x1xi8>
+  %weight = memref.get_global @weight : memref<9x2xi8>
+  %bias = memref.get_global @bias : memref<2xi32>
+  %output = memref.alloc() : memref<24x2xi8>
+  // CHECK: "gemmini.intr.config_st"
+  // CHECK: "gemmini.intr.config_ex"
+  // CHECK: "gemmini.intr.config_ld"
+  // CHECK: "gemmini.intr.mvin3"
+  // CHECK: "gemmini.intr.mvin"
+  // CHECK: "gemmini.intr.mvin2"
+  // CHECK: "gemmini.intr.preload"
+  // CHECK: "gemmini.intr.compute_preloaded"
+  // CHECK: "gemmini.intr.compute_accumulated"
+  gemmini.tile_conv %input %weight %bias %output %3 %8 %3 {stride = 1}:
+  memref<1x5x10x1xi8> memref<9x2xi8> memref<2xi32> memref<24x2xi8> i64 i64 i64
+  gemmini.print %output : memref<24x2xi8>
+  return %0 : i64
+}