CUTLASS 3.2 (#1024)

* CUTLASS 3.2

test/unit/nvrtc/thread/.gitignore

@@ -0,0 +1 @@
+nvrtc_config.hpp

test/unit/nvrtc/thread/CMakeLists.txt

@@ -26,10 +26,15 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+configure_file(nvrtc_config.in nvrtc_config.hpp)
+
 cutlass_test_unit_add_executable(
   cutlass_test_unit_nvrtc_thread
-  gemm_nvrtc.cu
+  nvrtc_gemm.cu
+  nvrtc_contraction.cu
   testbed.h
-  )
+)
 
 target_link_libraries(cutlass_test_unit_nvrtc_thread PRIVATE cutlass_nvrtc)
+
+target_include_directories(cutlass_test_unit_nvrtc_thread PRIVATE ${CMAKE_CURRENT_BINARY_DIR})

test/unit/nvrtc/thread/nvrtc_config.in

@@ -0,0 +1,3 @@
+#pragma once
+
+#define CUDA_INCLUDE_DIR "@CUDA_TOOLKIT_ROOT_DIR@/include"

test/unit/nvrtc/thread/nvrtc_contraction.cu

@@ -0,0 +1,66 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Unit tests for GETT
+*/
+
+#include <cuda.h>
+#include <gtest/gtest.h>
+
+#include "testbed.h"
+
+#include "nvrtc_config.hpp"
+
+#ifndef CUDA_INCLUDE_DIR
+static_assert(0, "CUDA include path is not defined");
+#endif
+
+TEST(SM90_nvrtc_kernel, Contraction) {
+  static const char* nvrtc_opts[] = {
+    "-w",
+    "-default-device",
+    "-std=c++17",
+    "-arch=sm_90",
+    "-I" CUDA_INCLUDE_DIR,
+  };
+
+  EXPECT_TRUE(test::nvrtc::thread::TestbedKernel::compile(
+    "nvrtc::thread::ContractionKernel<"
+        "cutlass::bfloat16_t, cutlass::bfloat16_t, cutlass::bfloat16_t,"
+        "cute::Shape<cute::Shape<cute::_8, cute::_8>, cute::Shape<cute::_16, cute::_8>, cute::Shape<cute::_8, cute::_8>>,"
+        "cute::Shape<cute::_1, cute::_2, cute::_1>,"
+        "true, true,"
+        "10, 10, 10, 10>::Kernel",
+    { nvrtc_opts, nvrtc_opts + 5 }
+  ));
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

test/unit/nvrtc/thread/testbed.h

@@ -35,12 +35,15 @@
 #pragma once
 
 #include <iostream>
+#include <cstdio>
+#include <vector>
 
 #include "cutlass/gemm/thread/mma.h"
 #include "../kernel/thread/testbed_kernel.h"
 
 #include "cutlass/util/host_tensor.h"
 #include "cutlass/util/tensor_view_io.h"
+#include "cutlass/trace.h"
 
 #include "cutlass/util/reference/host/tensor_copy.h"
 #include "cutlass/util/reference/host/tensor_fill.h"
@@ -58,6 +61,78 @@ namespace test {
 namespace nvrtc {
 namespace thread {
 
+#define NVRTC_RETURN_IF_ERROR(api)                    \
+  do {                                                \
+    nvrtcResult _result = api;                        \
+    if (_result != NVRTC_SUCCESS) {                   \
+      CUTLASS_TRACE_HOST("Nvrtc error: " << _result); \
+      return false;                                   \
+    }                                                 \
+  } while(0)
+
+inline const char * cuda_source_fmt = R"""(
+
+#include "kernel/thread/contraction.hpp"
+
+using Operator = %s;
+
+extern "C" __global__ void global_entry(__grid_constant__ Operator::Params const params) {
+  extern __shared__ char smem[];
+
+  Operator op;
+  op(params, smem);
+}
+
+)""";
+
+struct TestbedKernel {
+  static bool compile(std::string const &kernel, std::vector<const char *> const &opts) {
+    int sz = std::snprintf(nullptr, 0, cuda_source_fmt, kernel.c_str());
+    std::vector<char> cuda_source(sz + 1);
+    std::snprintf(&cuda_source[0], cuda_source.size(), cuda_source_fmt, kernel.c_str());
+
+    nvrtcProgram program;
+    NVRTC_RETURN_IF_ERROR(
+        nvrtcCreateProgram(
+            &program,
+            cuda_source.data(),
+            nullptr,
+            static_cast<int32_t>(cutlass::nvrtc::kCutlassHeaderCount),
+            cutlass::nvrtc::kCutlassHeaders,
+            cutlass::nvrtc::kCutlassHeaderNames)
+    );
+
+    nvrtcResult compile_result = 
+        nvrtcCompileProgram(
+            program, 
+            static_cast<int32_t>(opts.size()), 
+            opts.data());
+
+    size_t log_size;
+    NVRTC_RETURN_IF_ERROR(
+        nvrtcGetProgramLogSize(program, &log_size)
+    );
+
+    if (log_size > 1) {
+      auto log = std::make_unique<char[]>(log_size);
+
+      NVRTC_RETURN_IF_ERROR(
+          nvrtcGetProgramLog(program, log.get())
+      );
+                
+      std::cout << log.get() << std::endl;
+    }
+
+    NVRTC_RETURN_IF_ERROR(compile_result);
+
+    NVRTC_RETURN_IF_ERROR(
+        nvrtcDestroyProgram(&program)
+    );
+
+    return true;
+  }
+};
+
 /// Structure to compute the matrix product
 template <
   /// Size of the Gemm problem - concept: gemm::GemmShape<>