1 files changed, 94 insertions, 2 deletions
diff --git a/cutlass-example/cutlass_example.cu b/cutlass-example/cutlass_example.cu
index d3e8f89..6371164 100644
--- a/cutlass-example/cutlass_example.cu
+++ b/cutlass-example/cutlass_example.cu
@@ -6,11 +6,103 @@
 #include <gemm.h>
 
 int main(int argc, char* argv[]) {
-    typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kColumnMajor,
-                                        cutlass::MatrixLayout::kRowMajor, 
+
+#ifdef WMMA_GEMM_16x16x16_NT
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::Shape<32, 16, 16> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(16, 16, 16);
+#endif
+
+#ifdef WMMA_GEMM_16x16x32_NT
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::Shape<32, 16, 16> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(16, 16, 32);
+#endif
+
+#ifdef WMMA_GEMM_16x16x16_NN
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::Shape<32, 16, 16> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(16, 16, 16);
+#endif
+
+
+#ifdef WMMA_GEMM_16x16x32_NN
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::Shape<32, 16, 16> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(16, 16, 32);
+#endif
+
+#ifdef WMMA_GEMM_16x16x16_TT
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::Shape<32, 16, 16> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(16, 16, 16);
+#endif
+
+#ifdef WMMA_GEMM_16x16x32_TT
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::Shape<32, 16, 16> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(16, 16, 32);
+#endif
+
+#ifdef WMMA_GEMM_16x16x16_TN
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::Shape<32, 16, 16> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(16, 16, 16);
+#endif
+
+#ifdef WMMA_GEMM_16x16x32_TN
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::Shape<32, 16, 16> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(16, 16, 32);
+#endif
+
+#ifdef WMMA_16x16x16_GEMM_256x256x128_NT
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::Shape<32, 128, 128> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(256, 256, 128);
+#endif
+
+#ifdef WMMA_16x16x16_GEMM_256x256x128_NN
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::MatrixLayout::kColumnMajor,
+                                        cutlass::Shape<32, 128, 128> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(256, 256, 128);
+#endif
+
+#ifdef WMMA_16x16x16_GEMM_256x256x128_TT
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::Shape<32, 128, 128> >
+      WmmaGemmTraits;
+  run_gemm<WmmaGemmTraits>(256, 256, 128);
+#endif
+
+#ifdef WMMA_16x16x16_GEMM_256x256x128_TN
+  typedef cutlass::gemm::WmmaGemmTraits<cutlass::MatrixLayout::kRowMajor,
+                                        cutlass::MatrixLayout::kColumnMajor,
                                         cutlass::Shape<32, 128, 128> >
       WmmaGemmTraits;
   run_gemm<WmmaGemmTraits>(256, 256, 128);
+#endif
 
 }