Move mlir to the gpu and update the test

src/CMakeLists.txt

@@ -34,7 +34,6 @@ add_library(migraphx
     json.cpp
     load_save.cpp
     make_op.cpp
-    mlir.cpp
     module.cpp
     msgpack.cpp
     normalize_attributes.cpp
@@ -209,30 +208,6 @@ target_link_libraries(migraphx PRIVATE msgpackc-cxx)
 # Make this available to the tests
 target_link_libraries(migraphx INTERFACE $<BUILD_INTERFACE:msgpackc-cxx>)
 
-set(MIGRAPHX_ENABLE_MLIR OFF CACHE BOOL "")
-if(MIGRAPHX_ENABLE_MLIR)
-    find_library(MLIRAPI_LIBRARY MLIRMIOpen 
-        PATH_SUFFIXES
-        # Workaournd broken mlir install
-        lib/ lib/lib)
-    # REQUIRED is not supported before cmake 3.18
-    if(NOT MLIRAPI_LIBRARY)
-        message(FATAL_ERROR "libMLIRPublicAPI not found")
-    else()
-        message(STATUS "Build with libMLIRPublicAPI: " ${MLIRAPI_LIBRARY})
-    endif()
-
-    find_path(MLIRAPI_HEADERS NAMES mlir-c/Dialect/MIGraphX.h)
-    # Workaround MLIR broken installation
-    find_path(MLIRAPI_HEADERS2 NAMES mlir-c/Registration.h
-        PATH_SUFFIXES 
-        include/external/include external/include)
-
-    target_compile_definitions(migraphx PRIVATE "-DMIGRAPHX_MLIR")
-    target_include_directories(migraphx SYSTEM PRIVATE ${MLIRAPI_HEADERS} ${MLIRAPI_HEADERS2})
-    target_link_libraries(migraphx PUBLIC ${MLIRAPI_LIBRARY})
-endif()
-
 add_library(migraphx_all_targets INTERFACE)
 target_link_libraries(migraphx_all_targets INTERFACE migraphx_ref)
 

src/targets/gpu/CMakeLists.txt

@@ -144,6 +144,7 @@ add_library(migraphx_gpu
     loop.cpp
     lrn.cpp
     leaky_relu.cpp
+    mlir.cpp
     mlir_conv.cpp
     multinomial.cpp
     nonzero.cpp
@@ -284,6 +285,30 @@ endif()
 message(STATUS "clang-offload-bundler: ${MIGRAPHX_OFFLOADBUNDLER_BIN}")
 message(STATUS "extractkernel: ${MIGRAPHX_EXTRACT_KERNEL}")
 
+set(MIGRAPHX_ENABLE_MLIR OFF CACHE BOOL "")
+if(MIGRAPHX_ENABLE_MLIR)
+    find_library(MLIRAPI_LIBRARY MLIRMIOpen 
+        PATH_SUFFIXES
+        # Workaournd broken mlir install
+        lib/ lib/lib)
+    # REQUIRED is not supported before cmake 3.18
+    if(NOT MLIRAPI_LIBRARY)
+        message(FATAL_ERROR "libMLIRPublicAPI not found")
+    else()
+        message(STATUS "Build with libMLIRPublicAPI: " ${MLIRAPI_LIBRARY})
+    endif()
+
+    find_path(MLIRAPI_HEADERS NAMES mlir-c/Dialect/MIGraphX.h)
+    # Workaround MLIR broken installation
+    find_path(MLIRAPI_HEADERS2 NAMES mlir-c/Registration.h
+        PATH_SUFFIXES 
+        include/external/include external/include)
+
+    target_compile_definitions(migraphx_gpu PRIVATE "-DMIGRAPHX_MLIR")
+    target_include_directories(migraphx_gpu SYSTEM PRIVATE ${MLIRAPI_HEADERS} ${MLIRAPI_HEADERS2})
+    target_link_libraries(migraphx_gpu PUBLIC ${MLIRAPI_LIBRARY})
+endif()
+
 set(MIGRAPHX_ENABLE_MLIR_MIOPEN OFF CACHE BOOL "")
 if(MIGRAPHX_ENABLE_MLIR_MIOPEN)
     find_library(LIBMLIRMIOPEN MLIRMIOpenThin REQUIRED)

src/include/migraphx/mlir.hpp → src/targets/gpu/include/migraphx/gpu/mlir.hpp

-#ifndef MIGRAPHX_GUARD_RTGLIB_MLIR_HPP
-#define MIGRAPHX_GUARD_RTGLIB_MLIR_HPP
+#ifndef MIGRAPHX_GUARD_RTGLIB_GPU_MLIR_HPP
+#define MIGRAPHX_GUARD_RTGLIB_GPU_MLIR_HPP
 
 #include <string>
 #include <migraphx/config.hpp>
@@ -7,9 +7,11 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 struct module;
+namespace gpu {
 
 std::string dump_mlir(const module& m);
 
+} // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx
 

src/mlir.cpp → src/targets/gpu/mlir.cpp

-#include <migraphx/mlir.hpp>
+#include <migraphx/gpu/mlir.hpp>
 
 #include <mlir-c/IR.h>
 #include <mlir-c/BuiltinAttributes.h>
@@ -20,6 +20,7 @@
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
 
 template <class T, class F, F f> // NOLINT
 struct mlir_handle
@@ -69,7 +70,7 @@ struct mlir_handle
     std::unique_ptr<ptr, deleter> handle;
 };
 
-#define MIGRAPHX_MANAGE_MLIR_HANDLE(T, F) migraphx::mlir_handle<T, decltype(&F), &F> // NOLINT
+#define MIGRAPHX_MANAGE_MLIR_HANDLE(T, F) migraphx::gpu::mlir_handle<T, decltype(&F), &F> // NOLINT
 
 using mlir_context           = MIGRAPHX_MANAGE_MLIR_HANDLE(MlirContext, mlirContextDestroy);
 using mlir_module            = MIGRAPHX_MANAGE_MLIR_HANDLE(MlirModule, mlirModuleDestroy);
@@ -431,5 +432,6 @@ std::string dump_mlir(const module& m)
     return mlir_print(&mlirOperationPrint, mod_op);
 }
 
+} // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

test/mlir.cpp → test/gpu/mlir.cpp

-#include <migraphx/mlir.hpp>
+#include <migraphx/gpu/mlir.hpp>
 #include <migraphx/module.hpp>
 #include <migraphx/make_op.hpp>
 #include <migraphx/ranges.hpp>
+#include <migraphx/stringutils.hpp>
 #include <test.hpp>
 
+using migraphx::trim;
+
+std::string encode(std::string s)
+{
+    std::stringstream ss;
+    bool prespace = false;
+    for(auto c:s)
+    {
+        if (std::isspace(c))
+        {
+            if (not prespace)
+                ss << "  ";
+            prespace = true;
+        }
+        else if (std::isprint(c))
+        {
+            ss << c;
+            prespace = false;
+        }
+    }
+    return migraphx::trim(ss.str());
+}
+
 TEST_CASE(conv)
 {
+    const std::string mlir_output = R"__migraphx__(
+module  {
+  func @main(%arg0: tensor<1x8x4x4xf32>, %arg1: tensor<2x8x3x3xf32>) -> tensor<1x2x2x2xf32> {
+    %0 = "migraphx.convolution"(%arg0, %arg1) {dilation = [1 : si64, 1 : si64], group = 1 : si64, padding = [0 : si64, 0 : si64], padding_mode = 0 : si64, stride = [1 : si64, 1 :
+si64]} : (tensor<1x8x4x4xf32>, tensor<2x8x3x3xf32>) -> tensor<1x2x2x2xf32>
+  }
+}
+)__migraphx__";
     migraphx::module m;
     auto x = m.add_parameter("x", {migraphx::shape::float_type, {1, 8, 4, 4}});
     auto w = m.add_parameter("w", {migraphx::shape::float_type, {2, 8, 3, 3}});
     m.add_instruction(migraphx::make_op("convolution"), x, w);
-    auto s = migraphx::dump_mlir(m);
-    std::cout << s << std::endl;
-    EXPECT(migraphx::contains(s, "migraphx.convolution"));
-    EXPECT(not migraphx::contains(s, "migraphx.@param"));
+    auto s = migraphx::gpu::dump_mlir(m);
+    EXPECT(encode(s) == encode(mlir_output));
 }
 
 int main(int argc, const char* argv[]) { test::run(argc, argv); }