Merge remote-tracking branch 'origin/gemm-row-col' into ck-gemm-row-col

src/simplify_algebra.cpp

@@ -326,7 +326,11 @@ struct find_double_add_lit_broadcast
 
 struct find_inner_broadcast
 {
-    auto matcher() const { return pointwise(match::all_of[match::inputs()](match::broadcast())); }
+    auto matcher() const
+    {
+        return pointwise(match::all_of[match::inputs()](match::broadcast()))(
+            match::none_of(match::name("layout")));
+    }
 
     void apply(module& m, const match::matcher_result& r) const
     {

src/targets/gpu/CMakeLists.txt

@@ -109,6 +109,7 @@ add_library(migraphx_gpu
     perfdb.cpp
     pooling.cpp
     reverse.cpp
+    rewrite_ops.cpp
     rnn_variable_seq_lens.cpp
     rocblas.cpp
     scatter.cpp

src/targets/gpu/include/migraphx/gpu/rewrite_ops.hpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#ifndef MIGRAPHX_GUARD_GPU_REWRITE_OPS_HPP
+#define MIGRAPHX_GUARD_GPU_REWRITE_OPS_HPP
+
+#include <migraphx/config.hpp>
+#include <string>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct module;
+
+namespace gpu {
+
+struct rewrite_ops
+{
+    std::string name() const { return "gpu::rewrite_ops"; }
+    void apply(module& m) const;
+};
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif // MIGRAPHX_GUARD_GPU_REWRITE_OPS_HPP

src/targets/gpu/rewrite_ops.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/gpu/rewrite_ops.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/matcher.hpp>
+#include <migraphx/permutation.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+namespace {
+
+MIGRAPHX_PRED_MATCHER(col_matrix, instruction_ref ins)
+{
+    if(not ins->get_shape().transposed())
+        return false;
+    if(ins->get_shape().ndim() < 2)
+        return false;
+    auto perm = find_permutation(ins->get_shape());
+    auto n    = perm.size() - 1;
+    return perm[n] == n - 1 and perm[n - 1] == n;
+}
+
+MIGRAPHX_PRED_MATCHER(broadcast_matrix_dims, instruction_ref ins)
+{
+    if(not ins->get_shape().broadcasted())
+        return false;
+    if(ins->get_shape().ndim() < 2)
+        return false;
+    return std::any_of(ins->get_shape().lens().rbegin(),
+                       ins->get_shape().lens().rend() + 2,
+                       [](auto i) { return i == 0; });
+}
+
+struct find_dot_const
+{
+    auto matcher() const
+    {
+        return match::name("dot")(match::arg(1)(
+            match::is_constant(),
+            match::none_of(col_matrix(), broadcast_matrix_dims()),
+            match::skip_broadcasts(match::any().bind("w"))))(match::none_of(match::is_constant()));
+    }
+
+    void apply(module& m, const match::matcher_result& r) const
+    {
+        auto ins = r.result;
+        auto w   = r.instructions["w"];
+        if(w->get_shape().ndim() < 2)
+            return;
+        auto perm = find_permutation(w->get_shape());
+        auto n    = perm.size() - 1;
+        std::swap(perm[n], perm[n - 1]);
+        auto wl = m.insert_instruction(std::next(w), make_op("layout", {{"permutation", perm}}), w);
+        m.replace_instruction(w, wl);
+    }
+};
+
+} // namespace
+
+void rewrite_ops::apply(module& m) const { match::find_matches(m, find_dot_const{}); }
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/gpu/target.cpp

@@ -43,6 +43,7 @@
 #include <migraphx/register_target.hpp>
 #include <migraphx/replace_allocate.hpp>
 #include <migraphx/rewrite_gelu.hpp>
+#include <migraphx/gpu/rewrite_ops.hpp>
 #include <migraphx/rewrite_pooling.hpp>
 #include <migraphx/rewrite_quantization.hpp>
 #include <migraphx/rewrite_rnn.hpp>
@@ -130,6 +131,8 @@ std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_opti
         simplify_algebra{},
         prefuse_ops{},
         dead_code_elimination{},
+        rewrite_ops{},
+        dead_code_elimination{},
         auto_contiguous{},
         simplify_reshapes{},
         propagate_constant{},