[MLIR][NVVM] Add test cases to check translation of matrix-multiply accumulate ops to the corresponding intrinsics in NVPTX backend

[navdeepkk]

Summary

Add test cases to check translation of the following NVVM ops:-

1.) wmma.m16n16k16.load
2.) wmma.m16n16k1616.store
3.) wmma.m16n16k16.mma

[ftynse]

Test cases must be added in the same commit as the code they are testing.

[bondhugula]

mlir/test/Target/nvvmir.mlir

76

You may want to add a short comment here as to the test.

[navdeepkk]

Closing revision as changes here were moved to D95331.

mlir/test/Target/nvvmir.mlir

llvm.func @nvvm_mma(%a0 : !llvm.vec<2 x half>, %a1 : !llvm.vec<2 x half>,
                     %b0 : !llvm.vec<2 x half>, %b1 : !llvm.vec<2 x half>,
                     %c0 : !llvm.float, %c1 : !llvm.float, %c2 : !llvm.float, %c3 : !llvm.float,
                     %c4 : !llvm.float, %c5 : !llvm.float, %c6 : !llvm.float, %c7 : !llvm.float) {
   // CHECK: call { float, float, float, float, float, float, float, float } @llvm.nvvm.mma.m8n8k4.row.col.f32.f32
   %0 = nvvm.mma.sync %a0, %a1, %b0, %b1, %c0, %c1, %c2, %c3, %c4, %c5, %c6, %c7 {alayout="row", blayout="col"} : (!llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float, !llvm.float) -> !llvm.struct<(float, float, float, float, float, float, float, float)>
   llvm.return %0 : !llvm.struct<(float, float, float, float, float, float, float, float)>
 }
 
+llvm.func @gpu_wmma_load_op(%arg0: !llvm.ptr<i32, 3>, %arg1: !llvm.i32) {
+  // CHECK: call { <2 x half>, <2 x half>, <2 x half>, <2 x half>, <2 x half>, <2 x half>, <2 x half>, <2 x half> } @llvm.nvvm.wmma.m16n16k16.load.a.row.stride.f16.p3i32(i32 addrspace(3)* %{{.*}}, i32 %{{.*}})
+  %0 = nvvm.wmma.m16n16k16.load %arg0, %arg1 {ldm = 32 : i64, operand = "AOp"} : !llvm.ptr<i32, 3>, !llvm.i32 -> !llvm.struct<(vec<2 x half>, vec<2 x half>, vec<2 x half>, vec<2 x half>, vec<2 x half>, vec<2 x half>, vec<2 x half>, vec<2 x half>)>
+
+  llvm.return
+}
+
+llvm.func @gpu_wmma_store_op(%arg0: !llvm.ptr<i32, 3>, %arg1: !llvm.vec<2 x half>,
+                            %arg2: !llvm.vec<2 x half>, %arg3: !llvm.vec<2 x half>,
+                            %arg4: !llvm.vec<2 xhalf>, %arg5: !llvm.i32) {
+  // CHECK: call void @llvm.nvvm.wmma.m16n16k16.store.d.row.stride.f16.p3i32(i32 addrspace(3)* %{{.*}}, <2 x half> {{.*}}, <2 x half> %{{.*}}, <2 x half> %{{.*}}, <2 x half> %{{.*}}, i32 %{{.*}})
+  nvvm.wmma.m16n16k16.store %arg0, %arg1, %arg2, %arg3, %arg4, %arg5 : !llvm.ptr<i32, 3>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.i32
+  llvm.return
+}
+
+llvm.func @gpu_wmma_mma_op(%arg0: !llvm.vec<2 x half>, %arg1: !llvm.vec<2 x half>,
+                        %arg2: !llvm.vec<2 x half>, %arg3: !llvm.vec<2 x half>,
+                        %arg4: !llvm.vec<2 x half>, %arg5: !llvm.vec<2 x half>,
+                        %arg6: !llvm.vec<2 x half>, %arg7: !llvm.vec<2 x half>,
+                        %arg8: !llvm.vec<2 x half>, %arg9: !llvm.vec<2 x half>,
+                        %arg10: !llvm.vec<2 x half>, %arg11: !llvm.vec<2 x half>,
+                        %arg12: !llvm.vec<2 x half>, %arg13: !llvm.vec<2 x half>,
+                        %arg14: !llvm.vec<2 x half>, %arg15: !llvm.vec<2 x half>,
+                        %arg16: !llvm.vec<2 x half>, %arg17: !llvm.vec<2 x half>,
+                        %arg18: !llvm.vec<2 x half>, %arg19: !llvm.vec<2 x half>) {
+  // CHECK: call { <2 x half>, <2 x half>, <2 x half>, <2 x half> } @llvm.nvvm.wmma.m16n16k16.mma.row.row.f16.f16(<2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}}, <2 x half> {{.*}})
+  %0 = nvvm.wmma.m16n16k16.mma %arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg15, %arg16, %arg17, %arg18, %arg19 : !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half>, !llvm.vec<2 x half> -> !llvm.struct<(vec<2 x half>, vec<2 x half>, vec<2 x half>, vec<2 x half>)>
+
+  llvm.return
+}
+
 // This function has the "kernel" attribute attached and should appear in the
 // NVVM annotations after conversion.
 llvm.func @kernel_func() attributes {gpu.kernel} {
   llvm.return
 }
 
 // CHECK:     !nvvm.annotations =
 // CHECK-NOT: {i32 ()* @nvvm_special_regs, !"kernel", i32 1}
 // CHECK:     {void ()* @kernel_func, !"kernel", i32 1}