Merge branch 'int8_miopen_call' into int8_quantize

requirements.txt

 google/protobuf -DCMAKE_POSITION_INDEPENDENT_CODE=On
 RadeonOpenCompute/rocm-cmake@42f6740 --build
 ROCmSoftwarePlatform/rocBLAS@30a992ae02fda568688bcd190edd5e277d6674d9
-ROCmSoftwarePlatform/MIOpen@1.7.0
+ROCmSoftwarePlatform/MIOpen@1.8.0
 blaze,https://bitbucket.org/blaze-lib/blaze/get/f0755dea0e03.tar.gz -X header -DHEADER_DIR=blaze
 half,https://github.com/pfultz2/half/archive/1.12.0.tar.gz -X header -H sha256:0a08660b68abb176ebc2a0cdf8de46e3182a7f46c66443bb80dbfaaec98cf969
 pybind/pybind11@v2.2.4 -DPYBIND11_TEST=Off --build

src/include/migraphx/argument.hpp

@@ -36,7 +36,7 @@ struct argument : raw_data<argument>
     }
 
     /// Provides a raw pointer to the data
-    std::function<char*()> data;
+    std::function<char*()> data = nullptr;
 
     /// Whether data is available
     bool empty() const { return not data; }

src/include/migraphx/op/quant_convolution.hpp

+#ifndef MIGRAPHX_GUARD_OPERATORS_QUANT_CONVOLUTION_HPP
+#define MIGRAPHX_GUARD_OPERATORS_QUANT_CONVOLUTION_HPP
+
+#include <array>
+#include <migraphx/op/common.hpp>
+#include <migraphx/operation.hpp>
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/stringutils.hpp>
+#include <migraphx/streamutils.hpp>
+#include <migraphx/literal.hpp>
+#include <migraphx/shape_for_each.hpp>
+#include <migraphx/config.hpp>
+#include <cmath>
+#include <utility>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace op {
+
+struct quant_convolution
+{
+    std::array<std::size_t, 2> padding  = {{0, 0}};
+    std::array<std::size_t, 2> stride   = {{1, 1}};
+    std::array<std::size_t, 2> dilation = {{1, 1}};
+
+    padding_mode_t padding_mode = default_;
+    int group                   = 1;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.padding, "padding"),
+                    f(self.stride, "stride"),
+                    f(self.dilation, "dilation"),
+                    f(self.padding_mode, "padding_mode"),
+                    f(self.group, "group"));
+    }
+
+    std::string name() const { return "quant_convolution"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs, *this}.has(2).same_type().same_ndims().only_dims(4);
+
+        const shape& input   = inputs.at(0);
+        const shape& weights = inputs.at(1);
+        auto t               = input.type();
+
+        // all input type must be int8_type and output is float_type
+        if(t != shape::int8_type)
+        {
+            MIGRAPHX_THROW("QUANT_CONVOLUTION: only accept input and weights of type int8_t");
+        }
+        t = shape::float_type;
+
+        if(padding_mode == default_)
+        {
+            return {t,
+                    {
+                        input.lens()[0],
+                        weights.lens()[0],
+                        std::size_t(std::max<std::ptrdiff_t>(
+                            1,
+                            (input.lens()[2] - (1 + dilation[0] * (weights.lens()[2] - 1)) +
+                             2 * padding[0]) /
+                                    stride[0] +
+                                1)),
+                        std::size_t(std::max<std::ptrdiff_t>(
+                            1,
+                            (input.lens()[3] - (1 + dilation[1] * (weights.lens()[3] - 1)) +
+                             2 * padding[1]) /
+                                    stride[1] +
+                                1)),
+                    }};
+        }
+        else if(padding_mode == same)
+        {
+            return {t,
+                    {input.lens()[0],
+                     weights.lens()[0],
+                     static_cast<std::size_t>(
+                         std::ceil(static_cast<double>(input.lens()[2]) / stride[0])),
+                     static_cast<std::size_t>(
+                         std::ceil(static_cast<double>(input.lens()[3]) / stride[1]))}};
+        }
+        else if(padding_mode == valid)
+        {
+            return {
+                t,
+                {input.lens()[0],
+                 weights.lens()[0],
+                 static_cast<std::size_t>(std::ceil(
+                     static_cast<double>(input.lens()[2] - weights.lens()[2] + 1) / stride[0])),
+                 static_cast<std::size_t>(std::ceil(
+                     static_cast<double>(input.lens()[3] - weights.lens()[3] + 1) / stride[1]))}};
+        }
+        else
+        {
+            MIGRAPHX_THROW("Invalid padding mode");
+        }
+    }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/op/quant_dot.hpp

+#ifndef MIGRAPHX_GUARD_OPERATORS_QUANT_DOT_HPP
+#define MIGRAPHX_GUARD_OPERATORS_QUANT_DOT_HPP
+
+#include <array>
+#include <migraphx/operation.hpp>
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/stringutils.hpp>
+#include <migraphx/streamutils.hpp>
+#include <migraphx/literal.hpp>
+#include <migraphx/shape_for_each.hpp>
+#include <migraphx/config.hpp>
+#include <cmath>
+#include <utility>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace op {
+
+struct quant_dot
+{
+    int32_t alpha = 1;
+    int32_t beta  = 1;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(as_number(self.alpha), "alpha"), f(as_number(self.beta), "beta"));
+    }
+
+    std::string name() const { return "quant_dot"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{{inputs.at(0), inputs.at(1)}, *this}.same_type();
+        const shape& a = inputs.at(0);
+        const shape& b = inputs.at(1);
+        auto t         = a.type();
+        if(t != shape::int8_type)
+        {
+            MIGRAPHX_THROW("QUANT_DOT: only support data type int8_t");
+        }
+
+        if(!std::all_of(inputs.begin(), inputs.end(), [](auto s) { return s.lens().size() >= 2; }))
+        {
+            MIGRAPHX_THROW("QUANT_DOT: dot only accept 2 or more dims operands");
+        }
+
+        // only handle the case that the batch size of a and b are the same
+        if(!std::equal(
+               a.lens().rbegin() + 2, a.lens().rend(), b.lens().rbegin() + 2, b.lens().rend()))
+        {
+            MIGRAPHX_THROW("QUANT_DOT: batch size of A and B mismatch: {" +
+                           to_string_range(a.lens()) + "} x {" + to_string_range(b.lens()) + "}");
+        }
+
+        std::size_t dim_0 = a.lens().size() - 2;
+        std::size_t dim_1 = a.lens().size() - 1;
+        if(a.lens()[dim_1] != b.lens()[dim_0])
+        {
+            MIGRAPHX_THROW("QUANT_DOT: inner dimensions do not match: {" +
+                           to_string_range(a.lens()) + "} x {" + to_string_range(b.lens()) + "}");
+        }
+
+        // k be multiple of 4
+        if((a.lens()[dim_1] % 4) != 0)
+        {
+            MIGRAPHX_THROW("QUANT_DOT: size of A {" + to_string_range(a.lens()) + "} and B {" +
+                           to_string_range(b.lens()) + "} must be multiple of 4 for int8 type");
+        }
+
+        auto out_lens   = a.lens();
+        out_lens[dim_1] = b.lens()[dim_1];
+        if(inputs.size() == 3 && out_lens != inputs.at(2).lens())
+        {
+            MIGRAPHX_THROW("QUANT_DOT: dimension mismatch, operand C: {" +
+                           to_string_range(inputs.at(2).lens()) +
+                           "}, cannot add to operand A * B: {" + to_string_range(out_lens) + "}");
+        }
+
+        if(inputs.size() == 3 && inputs.at(2).type() != shape::int32_type)
+        {
+            MIGRAPHX_THROW("QUANT_DOT: operand C type must be int32");
+        }
+
+        return {shape::int32_type, out_lens};
+    }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/operators.hpp

@@ -42,6 +42,8 @@
 #include <migraphx/op/outline.hpp>
 #include <migraphx/op/pad.hpp>
 #include <migraphx/op/pooling.hpp>
+#include <migraphx/op/quant_convolution.hpp>
+#include <migraphx/op/quant_dot.hpp>
 #include <migraphx/op/relu.hpp>
 #include <migraphx/op/reshape.hpp>
 #include <migraphx/op/rnn.hpp>

src/include/migraphx/tensor_view.hpp

@@ -12,6 +12,14 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
+template <class T>
+T as_number(T x)
+{
+    return x;
+}
+inline int32_t as_number(int8_t x) { return static_cast<int32_t>(x); }
+inline uint32_t as_number(uint8_t x) { return static_cast<uint32_t>(x); }
+
 template <class T>
 struct tensor_view
 {
@@ -130,10 +138,10 @@ struct tensor_view
     {
         if(!x.empty())
         {
-            os << x.front();
+            os << as_number(x.front());
             for(std::size_t i = 1; i < x.m_shape.elements(); i++)
             {
-                os << ", " << x.m_data[x.m_shape.index(i)];
+                os << ", " << as_number(x.m_data[x.m_shape.index(i)]);
             }
         }
         return os;

src/targets/cpu/gemm.cpp

@@ -44,13 +44,9 @@ struct is_fast_gemm_type<float> : std::true_type
 {
 };
 
-template <class T>
-void migemm_impl(tensor_view<T> cmat,
-                 tensor_view<T> amat,
-                 tensor_view<T> bmat,
-                 float alpha,
-                 float beta,
-                 std::true_type)
+template <class T, class F>
+void migemm_impl(
+    tensor_view<T> cmat, tensor_view<T> amat, tensor_view<T> bmat, F alpha, F beta, std::true_type)
 {
     visit_mat(amat, [&](const auto& a) {
         visit_mat(bmat, [&](const auto& b) {
@@ -66,13 +62,9 @@ void migemm_impl(tensor_view<T> cmat,
     });
 }
 
-template <class T>
-void migemm_impl(tensor_view<T> cmat,
-                 tensor_view<T> amat,
-                 tensor_view<T> bmat,
-                 float alpha,
-                 float beta,
-                 std::false_type)
+template <class T, class F>
+void migemm_impl(
+    tensor_view<T> cmat, tensor_view<T> amat, tensor_view<T> bmat, F alpha, F beta, std::false_type)
 {
     std::size_t n_dims = cmat.get_shape().lens().size();
     std::size_t dim_0  = n_dims - 2;
@@ -95,9 +87,8 @@ void migemm_impl(tensor_view<T> cmat,
     });
 }
 
-template <class T>
-void migemm_impl(
-    tensor_view<T> cmat, tensor_view<T> amat, tensor_view<T> bmat, float alpha, float beta)
+template <class T, class F>
+void migemm_impl(tensor_view<T> cmat, tensor_view<T> amat, tensor_view<T> bmat, F alpha, F beta)
 {
     auto lens = amat.get_shape().lens();
     bool batch_mul =
@@ -113,13 +104,29 @@ void migemm_impl(
     }
 }
 
-void migemm(
-    const argument& c_arg, const argument& a_arg, const argument& b_arg, float alpha, float beta)
+template <class F>
+void migemm_tpl(
+    const argument& c_arg, const argument& a_arg, const argument& b_arg, F alpha, F beta)
 {
     visit_all(c_arg, a_arg, b_arg)(
         [&](auto cmat, auto amat, auto bmat) { migemm_impl(cmat, amat, bmat, alpha, beta); });
 }
 
+void migemm(
+    const argument& c_arg, const argument& a_arg, const argument& b_arg, float alpha, float beta)
+{
+    migemm_tpl(c_arg, a_arg, b_arg, alpha, beta);
+}
+
+void migemm(const argument& c_arg,
+            const argument& a_arg,
+            const argument& b_arg,
+            int32_t alpha,
+            int32_t beta)
+{
+    migemm_tpl(c_arg, a_arg, b_arg, alpha, beta);
+}
+
 } // namespace cpu
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/targets/cpu/include/migraphx/cpu/gemm.hpp

@@ -10,6 +10,11 @@ namespace cpu {
 
 void migemm(
     const argument& c_arg, const argument& a_arg, const argument& b_arg, float alpha, float beta);
+void migemm(const argument& c_arg,
+            const argument& a_arg,
+            const argument& b_arg,
+            int32_t alpha,
+            int32_t beta);
 
 } // namespace cpu
 } // namespace MIGRAPHX_INLINE_NS

src/targets/cpu/lowering.cpp

@@ -204,6 +204,61 @@ struct cpu_convolution
     }
 };
 
+struct cpu_quant_convolution
+{
+    op::quant_convolution op;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return migraphx::reflect(self.op, f);
+    }
+
+    std::string name() const { return "cpu::quant_convolution"; }
+    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
+    argument compute(context&, shape output_shape, std::vector<argument> args) const
+    {
+        argument result{output_shape};
+        result.visit([&](auto output) {
+            visit_all(args[0], args[1])([&](auto input, auto weights) {
+                auto in   = input.get_shape().lens();
+                auto in_h = in[2];
+                auto in_w = in[3];
+
+                auto wei   = weights.get_shape().lens();
+                auto wei_n = wei[0];
+                auto wei_c = wei[1];
+                auto wei_h = wei[2];
+                auto wei_w = wei[3];
+
+                par_dfor(output_shape.lens()[0],
+                         output_shape.lens()[1],
+                         output_shape.lens()[2],
+                         output_shape.lens()[3])(
+                    [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
+                        const int start_x  = i * op.stride[0] - op.padding[0];
+                        const int start_y  = j * op.stride[1] - op.padding[1];
+                        const int group_id = w / (wei_n / op.group);
+
+                        float acc = 0;
+                        dfor(wei_c, wei_h, wei_w)([&](std::size_t k, std::size_t x, std::size_t y) {
+                            const int in_x  = start_x + x;
+                            const int in_y  = start_y + y;
+                            const int in_ch = group_id * wei_c + k;
+                            if(in_x >= 0 && in_x < in_h && in_y >= 0 && in_y < in_w)
+                            {
+                                acc += input(o, in_ch, in_x, in_y) * weights(w, k, x, y);
+                            }
+                        });
+                        output(o, w, i, j) = acc;
+                    });
+            });
+        });
+
+        return result;
+    }
+};
+
 struct cpu_im2col
 {
     op::im2col op;
@@ -421,7 +476,7 @@ struct cpu_gemm
     {
         argument result{output_shape};
         // 3 inputs, it is alpha * A * B + beta * C, then
-        // A and B are matrics, and C is broadcastable to A * B
+        // A and B are matrices, and C is of the same shape as A * B
         if(args.size() == 3)
         {
             // no need to consider the value of args[2]
@@ -448,6 +503,73 @@ struct cpu_gemm
     }
 };
 
+struct cpu_quant_gemm
+{
+    op::quant_dot op;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return migraphx::reflect(self.op, f);
+    }
+
+    std::string name() const { return "cpu::quant_dot"; }
+    shape compute_shape(const std::vector<shape>& inputs) const
+    {
+        if(inputs.size() == 3)
+        {
+            auto c_shape = inputs.at(2);
+            check_shapes{{c_shape}}.not_broadcasted();
+        }
+        return op.compute_shape(inputs);
+    }
+
+    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
+    {
+        argument result{output_shape};
+        // 3 inputs, it is alpha * A * B + beta * C, then
+        // A and B are matrices, and C is of the same shape to A * B
+
+        // first, convert the args[0] and args[1] from int8_t to int32_t
+        argument arg_0{{shape::int32_type, {args.at(0).get_shape().lens()}}};
+        argument arg_1{{shape::int32_type, {args.at(1).get_shape().lens()}}};
+        arg_0.visit([&](auto output) {
+            args.at(0).visit(
+                [&](auto input) { std::copy(input.begin(), input.end(), output.begin()); });
+        });
+
+        arg_1.visit([&](auto output) {
+            args.at(1).visit(
+                [&](auto input) { std::copy(input.begin(), input.end(), output.begin()); });
+        });
+
+        if(args.size() == 3)
+        {
+            // no need to consider the value of args[2]
+            if(op.beta == 0)
+            {
+                result.visit([&](auto output) { std::fill(output.begin(), output.end(), 0); });
+            }
+            else
+            {
+                visit_all(result, args[2])([&](auto output, auto input) {
+                    std::copy(input.begin(), input.end(), output.begin());
+                });
+            }
+
+            migemm(result, arg_0, arg_1, op.alpha, op.beta);
+
+            return result;
+        }
+
+        // 2 input arguments
+        int32_t beta = 0;
+        migemm(result, arg_0, arg_1, op.alpha, beta);
+
+        return result;
+    }
+};
+
 struct leaky_relu_op
 {
     op::leaky_relu op;
@@ -652,15 +774,17 @@ struct cpu_apply
     {
         apply_map["batch_norm_inference"] =
             extend_op<cpu_batch_norm_inference, op::batch_norm_inference>();
-        apply_map["convolution"] = extend_op<cpu_convolution, op::convolution>();
-        apply_map["dot"]         = extend_op<cpu_gemm, op::dot>();
-        apply_map["elu"]         = extend_op<cpu_unary<elu_op>, op::elu>();
-        apply_map["im2col"]      = extend_op<cpu_im2col, op::im2col>();
-        apply_map["leaky_relu"]  = extend_op<cpu_unary<leaky_relu_op>, op::leaky_relu>();
-        apply_map["logsoftmax"]  = extend_op<cpu_logsoftmax, op::logsoftmax>();
-        apply_map["lrn"]         = extend_op<cpu_lrn, op::lrn>();
-        apply_map["pad"]         = extend_op<cpu_pad, op::pad>();
-        apply_map["softmax"]     = simple_op<softmax2d>();
+        apply_map["convolution"]       = extend_op<cpu_convolution, op::convolution>();
+        apply_map["dot"]               = extend_op<cpu_gemm, op::dot>();
+        apply_map["quant_dot"]         = extend_op<cpu_quant_gemm, op::quant_dot>();
+        apply_map["quant_convolution"] = extend_op<cpu_quant_convolution, op::quant_convolution>();
+        apply_map["elu"]               = extend_op<cpu_unary<elu_op>, op::elu>();
+        apply_map["im2col"]            = extend_op<cpu_im2col, op::im2col>();
+        apply_map["leaky_relu"]        = extend_op<cpu_unary<leaky_relu_op>, op::leaky_relu>();
+        apply_map["logsoftmax"]        = extend_op<cpu_logsoftmax, op::logsoftmax>();
+        apply_map["lrn"]               = extend_op<cpu_lrn, op::lrn>();
+        apply_map["pad"]               = extend_op<cpu_pad, op::pad>();
+        apply_map["softmax"]           = simple_op<softmax2d>();
     }
 
     void apply()

src/targets/gpu/CMakeLists.txt

@@ -33,6 +33,7 @@ add_library(migraphx_device
     device/pad.cpp
     device/gather.cpp
     device/sub.cpp
+    device/pack.cpp
     device/clip.cpp
 )
 set_target_properties(migraphx_device PROPERTIES EXPORT_NAME device)
@@ -48,8 +49,10 @@ add_library(migraphx_gpu
     target.cpp
     lowering.cpp
     gemm.cpp
+    quant_gemm.cpp
     pooling.cpp
     convolution.cpp
+    quant_convolution.cpp
     softmax.cpp
     logsoftmax.cpp
     contiguous.cpp

src/targets/gpu/device/include/migraphx/gpu/device/tensor.hpp

@@ -61,6 +61,12 @@ struct hip_tensor_descriptor
     {
         std::copy(s.lens().begin(), s.lens().end(), lens);
         std::copy(s.strides().begin(), s.strides().end(), strides);
+        std::vector<std::size_t> vec_idx(s.lens().size());
+        std::iota(vec_idx.begin(), vec_idx.end(), 0);
+        std::sort(vec_idx.begin(), vec_idx.end(), [&](size_t i, size_t j) {
+            return strides[i] > strides[j];
+        });
+        std::copy(vec_idx.begin(), vec_idx.end(), indices);
     }
 
     __device__ __host__ hip_index<NDim> multi(size_t idx) const
@@ -69,11 +75,12 @@ struct hip_tensor_descriptor
         size_t tidx = idx;
         for(size_t is = 0; is < NDim; is++)
         {
-            result[is] = tidx / strides[is];
-            tidx       = tidx % strides[is];
+            result[indices[is]] = tidx / strides[indices[is]];
+            tidx                = tidx % strides[indices[is]];
         }
         return result;
     }
+
     __device__ __host__ size_t linear(hip_index<NDim> s) const
     {
         size_t idx = 0;
@@ -83,6 +90,7 @@ struct hip_tensor_descriptor
     }
     size_t lens[NDim]    = {};
     size_t strides[NDim] = {};
+    size_t indices[NDim] = {};
 };
 
 } // namespace device

src/targets/gpu/device/pack.cpp

+#include <migraphx/shape.hpp>
+#include <migraphx/argument.hpp>
+#include <migraphx/gpu/device/pack.hpp>
+#include <migraphx/gpu/device/tensor.hpp>
+#include <migraphx/gpu/device/launch.hpp>
+#include <migraphx/gpu/device/types.hpp>
+#include <migraphx/gpu/hip.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+namespace device {
+
+void pack_a(hipStream_t stream, const argument& result, const argument& arg)
+{
+    auto output_shape  = result.get_shape();
+    auto out_lens      = output_shape.lens();
+    auto dim_0         = out_lens.size() - 2;
+    auto dim_1         = out_lens.size() - 1;
+    std::size_t lda    = output_shape.strides()[dim_0];
+    std::size_t m_size = out_lens[dim_0] * out_lens[dim_1];
+    visit_all(result, arg)([&](auto output, auto input) {
+        std::size_t nelements = output_shape.elements();
+        auto* out_ptr         = device_cast(output.data());
+        auto* in_ptr          = device_cast(input.data());
+        visit_tensor_size(out_lens.size(), [&](auto out_dim) {
+            hip_tensor_descriptor<out_dim> desc(output_shape);
+            gs_launch(stream, nelements)([=](auto ii) {
+                const size_t nb    = 4;
+                auto idx           = desc.multi(ii);
+                std::size_t i_m    = idx[dim_1];
+                std::size_t i_k    = idx[dim_0];
+                std::size_t offset = ii / m_size * m_size;
+                out_ptr[i_k % nb + (i_m + (i_k / nb) * lda) * nb + offset] =
+                    in_ptr[i_m + i_k * lda + offset];
+            });
+        });
+    });
+}
+
+void pack_b(hipStream_t stream, const argument& result, const argument& arg)
+{
+    auto output_shape  = result.get_shape();
+    auto out_lens      = output_shape.lens();
+    auto dim_0         = output_shape.lens().size() - 2;
+    auto dim_1         = output_shape.lens().size() - 1;
+    std::size_t ldb    = output_shape.strides()[dim_1];
+    std::size_t m_size = out_lens[dim_0] * out_lens[dim_1];
+    visit_all(result, arg)([&](auto output, auto input) {
+        std::size_t nelements = output_shape.elements();
+        auto* out_ptr         = device_cast(output.data());
+        auto* in_ptr          = device_cast(input.data());
+        visit_tensor_size(out_lens.size(), [&](auto out_dim) {
+            hip_tensor_descriptor<out_dim> desc(output_shape);
+            gs_launch(stream, nelements)([=](auto ii) {
+                const size_t nb    = 4;
+                auto idx           = desc.multi(ii);
+                std::size_t i_n    = idx[dim_0];
+                std::size_t i_k    = idx[dim_1];
+                std::size_t offset = ii / m_size * m_size;
+                out_ptr[i_k % nb + (i_n + (i_k / nb) * ldb) * nb + offset] =
+                    in_ptr[i_n + i_k * ldb + offset];
+            });
+        });
+    });
+}
+
+} // namespace device
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/gpu/gemm.cpp

@@ -215,6 +215,10 @@ argument miopen_gemm::compute(context& ctx,
         auto to_pointer = [&](auto&& arg) { return to_rocblas_type(as.from(arg.data())); };
         if(num_matrices == 1)
         {
+            // the rocblas_gemm API handles inputs and output matrices as
+            // column-major format. When doing a C = A * B, we actually do
+            // C^T = (B^T) * (A^T). That is the reason we input args[1] as
+            // A and args[0] as B in calling the rocblas_gemm.
             generic_rocblas_gemm(as,
                                  ctx.get_stream().get_rocblas(),
                                  transb ? rocblas_operation_transpose : rocblas_operation_none,

src/targets/gpu/include/migraphx/gpu/device/pack.hpp

+#ifndef MIGRAPHX_GUARD_RTGLIB_DEVICE_PACK_HPP
+#define MIGRAPHX_GUARD_RTGLIB_DEVICE_PACK_HPP
+
+#include <migraphx/argument.hpp>
+#include <migraphx/config.hpp>
+#include <hip/hip_runtime_api.h>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+namespace device {
+
+void pack_a(hipStream_t stream, const argument& result, const argument& arg);
+
+void pack_b(hipStream_t stream, const argument& result, const argument& arg);
+} // namespace device
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

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

@@ -34,11 +34,11 @@ Result make_obj(F f, Ts... xs)
     auto status                = f(&x, xs...);
     Result r{x};
     if(status != miopenStatusSuccess)
-        MIGRAPHX_THROW("MIOpen call failed");
+        MIGRAPHX_THROW("MAKE_OBJ: MIOpen call failed");
     return r;
 }
 
-inline tensor_descriptor make_tensor(const migraphx::shape& s)
+inline tensor_descriptor make_tensor(const migraphx::shape& s, bool pack = false)
 {
     auto t = make_obj<tensor_descriptor>(&miopenCreateTensorDescriptor);
     // Convert to ints
@@ -49,13 +49,31 @@ inline tensor_descriptor make_tensor(const migraphx::shape& s)
         d = miopenFloat;
     else if(s.type() == shape::half_type)
         d = miopenHalf;
+    else if(s.type() == shape::int8_type)
+    {
+        if(pack)
+        {
+            // update the lens and corresponding strides
+            d          = miopenInt8x4;
+            lens[1]    = ((lens[1] + 3) / 4) * 4;
+            strides[0] = strides[1] * lens[1];
+        }
+        else
+        {
+            d = miopenInt8;
+        }
+    }
     else
-        MIGRAPHX_THROW("Unsupported type");
+    {
+        MIGRAPHX_THROW("MAKE_TENSOR: unsupported type");
+    }
     miopenSetTensorDescriptor(t.get(), d, s.lens().size(), lens.data(), strides.data());
+
     return t;
 }
 
-inline convolution_descriptor make_conv(const migraphx::op::convolution& op)
+template <class T>
+inline convolution_descriptor make_conv(const T& op)
 {
     auto c = make_obj<convolution_descriptor>(&miopenCreateConvolutionDescriptor);
     miopenConvolutionMode_t c_mode = miopenConvolution;

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

+#ifndef MIGRAPHX_GUARD_RTGLIB_QUANT_CONVOLUTION_HPP
+#define MIGRAPHX_GUARD_RTGLIB_QUANT_CONVOLUTION_HPP
+
+#include <migraphx/shape.hpp>
+#include <migraphx/op/quant_convolution.hpp>
+#include <migraphx/gpu/miopen.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+struct context;
+
+struct miopen_quant_convolution
+{
+    op::quant_convolution op;
+    shared<convolution_descriptor> cd;
+    miopenConvFwdAlgorithm_t algo{};
+    miopenHandle_t handle = nullptr;
+    argument arg_vec4_x{};
+    argument arg_vec4_w{};
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        // TODO: Add algo
+        return op::quant_convolution::reflect(self.op, f);
+    }
+
+    std::string name() const { return "gpu::quant_convolution"; }
+    shape compute_shape(const std::vector<shape>& inputs) const;
+    argument
+    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    shape compile(context& ctx, const shape& output_shape, std::vector<shape> inputs);
+    void finalize(context& ctx, const shape& output_shape, std::vector<shape> inputs);
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
+
+    private:
+    shape pack_int8_shape(shape& s);
+};
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

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

+#ifndef MIGRAPHX_GUARD_RTGLIB_QUANT_GEMM_HPP
+#define MIGRAPHX_GUARD_RTGLIB_QUANT_GEMM_HPP
+
+#include <migraphx/shape.hpp>
+#include <migraphx/op/quant_dot.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+struct context;
+
+struct miopen_quant_gemm
+{
+    op::quant_dot op;
+    mutable argument arg_a{};
+    mutable argument arg_b{};
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return migraphx::reflect(self.op, f);
+    }
+
+    std::string name() const { return "gpu::quant_gemm"; }
+    shape compute_shape(const std::vector<shape>& inputs) const;
+    argument
+    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
+    {
+        return shapes.size() - 1;
+    }
+};
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/targets/gpu/lowering.cpp

@@ -14,6 +14,7 @@
 #include <migraphx/gpu/rocblas.hpp>
 #include <migraphx/gpu/context.hpp>
 #include <migraphx/gpu/convolution.hpp>
+#include <migraphx/gpu/quant_convolution.hpp>
 #include <migraphx/gpu/contiguous.hpp>
 #include <migraphx/gpu/relu.hpp>
 #include <migraphx/gpu/sigmoid.hpp>
@@ -41,6 +42,7 @@
 #include <migraphx/gpu/batchnorm.hpp>
 #include <migraphx/gpu/pooling.hpp>
 #include <migraphx/gpu/gemm.hpp>
+#include <migraphx/gpu/quant_gemm.hpp>
 #include <migraphx/gpu/concat.hpp>
 #include <migraphx/gpu/pad.hpp>
 #include <migraphx/gpu/gather.hpp>
@@ -97,6 +99,7 @@ struct miopen_apply
         add_generic_op<hip_min>("min");
 
         add_extend_op<miopen_gemm, op::dot>("dot");
+        add_extend_op<miopen_quant_gemm, op::quant_dot>("quant_dot");
         add_extend_op<miopen_contiguous, op::contiguous>("contiguous");
         add_extend_op<hip_concat, op::concat>("concat");
         add_extend_op<miopen_softmax, op::softmax>("softmax");
@@ -108,6 +111,7 @@ struct miopen_apply
 
         add_lrn_op();
         add_convolution_op();
+        add_quant_convolution_op();
         add_pooling_op();
         add_batch_norm_inference_op();
     }
@@ -154,6 +158,22 @@ struct miopen_apply
         });
     }
 
+    void add_quant_convolution_op()
+    {
+        apply_map.emplace("quant_convolution", [=](instruction_ref ins) {
+            auto&& op = any_cast<op::quant_convolution>(ins->get_operator());
+
+            auto conv = miopen_quant_convolution{op, make_conv(op)};
+            auto ws   = conv.compile(ctx, ins->get_shape(), to_shapes(ins->inputs()));
+
+            auto workspace = insert_allocation(ins, ws, "workspace");
+            auto output    = insert_allocation(ins, ins->get_shape());
+
+            return prog->replace_instruction(
+                ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
+        });
+    }
+
     void add_pooling_op()
     {
         apply_map.emplace("pooling", [=](instruction_ref ins) {

src/targets/gpu/quant_convolution.cpp

+#include <migraphx/gpu/quant_convolution.hpp>
+#include <migraphx/gpu/context.hpp>
+#include <migraphx/generate.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+shape miopen_quant_convolution::compute_shape(const std::vector<shape>& inputs) const
+{
+    check_shapes{inputs, *this}.has(4).standard();
+    return op.compute_shape({inputs.at(0), inputs.at(1)});
+}
+argument miopen_quant_convolution::compute(context& ctx,
+                                           const shape& output_shape,
+                                           const std::vector<argument>& args) const
+{
+    auto x_desc      = make_tensor(args[0].get_shape());
+    auto x_desc_vec4 = make_tensor(args[0].get_shape(), true);
+    auto w_desc      = make_tensor(args[1].get_shape());
+    auto w_desc_vec4 = make_tensor(args[1].get_shape(), true);
+    auto y_desc      = make_tensor(output_shape);
+
+    float alpha = 1;
+    float beta  = 0;
+
+    // pack input to vec4 format
+    auto status = miopenTransformTensor(ctx.get_stream().get_miopen(),
+                                        &alpha,
+                                        x_desc.get(),
+                                        args[0].implicit(),
+                                        &beta,
+                                        x_desc_vec4.get(),
+                                        arg_vec4_x.implicit());
+    if(status != miopenStatusSuccess)
+    {
+        MIGRAPHX_THROW("QUANT_CONVOLUTION: transform input tensor failed");
+    }
+
+    // pack input to vec4 format
+    status = miopenTransformTensor(ctx.get_stream().get_miopen(),
+                                   &alpha,
+                                   w_desc.get(),
+                                   args[1].implicit(),
+                                   &beta,
+                                   w_desc_vec4.get(),
+                                   arg_vec4_w.implicit());
+    if(status != miopenStatusSuccess)
+    {
+        MIGRAPHX_THROW("QUANT_CONVOLUTION: transform weight tensor failed");
+    }
+
+    status = miopenConvolutionForward(ctx.get_stream().get_miopen(),
+                                      &alpha,
+                                      x_desc_vec4.get(),
+                                      arg_vec4_x.implicit(),
+                                      w_desc_vec4.get(),
+                                      arg_vec4_w.implicit(),
+                                      cd.get(),
+                                      algo,
+                                      &beta,
+                                      y_desc.get(),
+                                      args[3].implicit(),
+                                      args[2].implicit(),
+                                      args[2].get_shape().bytes());
+    if(status != miopenStatusSuccess)
+    {
+        MIGRAPHX_THROW("QUANT_CONVOLUTION: run convolution forward failed");
+    }
+    return args[3];
+}
+
+shape miopen_quant_convolution::compile(context& ctx,
+                                        const shape& output_shape,
+                                        std::vector<shape> inputs)
+{
+    shape workspace_shape{};
+    auto x_desc = make_tensor(inputs[0], true);
+    auto w_desc = make_tensor(inputs[1], true);
+    auto y_desc = make_tensor(output_shape);
+
+    std::size_t workspace_size = 0;
+    miopenConvolutionForwardGetWorkSpaceSize(ctx.get_stream().get_miopen(),
+                                             w_desc.get(),
+                                             x_desc.get(),
+                                             cd.get(),
+                                             y_desc.get(),
+                                             &workspace_size);
+    workspace_shape = shape{shape::int8_type, {workspace_size}};
+
+    arg_vec4_x     = to_gpu(generate_argument(pack_int8_shape(inputs[0])));
+    arg_vec4_w     = to_gpu(generate_argument(pack_int8_shape(inputs[1])));
+    auto y         = allocate_gpu(output_shape);
+    auto workspace = allocate_gpu(workspace_shape);
+
+    int algo_count = 1;
+    miopenConvAlgoPerf_t perf;
+    auto status = miopenFindConvolutionForwardAlgorithm(ctx.get_stream().get_miopen(),
+                                                        x_desc.get(),
+                                                        arg_vec4_x.implicit(),
+                                                        w_desc.get(),
+                                                        arg_vec4_x.implicit(),
+                                                        cd.get(),
+                                                        y_desc.get(),
+                                                        y.implicit(),
+                                                        1,
+                                                        &algo_count,
+                                                        &perf,
+                                                        workspace.implicit(),
+                                                        workspace_size,
+                                                        false);
+    if(status != miopenStatusSuccess)
+    {
+        MIGRAPHX_THROW("QUANT_CONVOLUTION: find convolution failed");
+    }
+    handle = ctx.get_stream().get_miopen();
+    algo   = perf.fwd_algo;
+    return shape{shape::int8_type, {perf.memory}};
+}
+
+void miopen_quant_convolution::finalize(context& ctx,
+                                        const shape& output_shape,
+                                        std::vector<shape> inputs)
+{
+    if(handle == ctx.get_stream().get_miopen())
+        return;
+    // Check that workspace hasn't changed
+    auto size = inputs.at(2).bytes();
+    auto ws   = compile(ctx, output_shape, std::move(inputs));
+    if(ws.bytes() > size)
+        MIGRAPHX_THROW("Workspace has changed during finalization.");
+}
+
+shape miopen_quant_convolution::pack_int8_shape(shape& s)
+{
+    if(s.type() != shape::int8_type)
+    {
+        MIGRAPHX_THROW("PACK_INT8_SHAPE: only process int8_type");
+    }
+
+    auto lens    = s.lens();
+    auto strides = s.strides();
+    lens[1]      = (lens[1] + 3) / 4 * 4;
+    strides[0]   = strides[1] * lens[1];
+
+    return {s.type(), lens, strides};
+}
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/gpu/quant_gemm.cpp

+#include <migraphx/gpu/quant_gemm.hpp>
+#include <migraphx/gpu/device/pack.hpp>
+#include <migraphx/gpu/context.hpp>
+#include <migraphx/generate.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+template <class... Ts>
+rocblas_status generic_rocblas_gemm_ex(Ts&&... xs)
+{
+    return rocblas_gemm_ex(std::forward<Ts>(xs)...);
+}
+
+template <class... Ts>
+rocblas_status generic_rocblas_batched_gemm_ex(Ts&&... xs)
+{
+    return rocblas_gemm_strided_batched_ex(std::forward<Ts>(xs)...);
+}
+
+template <class T>
+struct compute_rocblas_type
+{
+    using type = T;
+};
+
+template <class T>
+struct compute_rocblas_type<const T>
+{
+    using type = const typename compute_rocblas_type<T>::type;
+};
+
+template <>
+struct compute_rocblas_type<half>
+{
+    using type = rocblas_half;
+};
+
+template <class T>
+using rb_type = typename compute_rocblas_type<T>::type;
+
+template <class T>
+rb_type<T> to_rocblas_type(T x)
+{
+    return reinterpret_cast<const rb_type<T>&>(x);
+}
+
+template <class T>
+rb_type<T>* to_rocblas_type(T* x)
+{
+    return reinterpret_cast<rb_type<T>*>(x);
+}
+
+shape miopen_quant_gemm::compute_shape(const std::vector<shape>& inputs) const
+{
+    std::vector<shape> input_shapes(inputs);
+    input_shapes.pop_back();
+    check_shapes{input_shapes}.not_broadcasted();
+    return op.compute_shape(input_shapes);
+}
+
+argument miopen_quant_gemm::compute(context& ctx,
+                                    const shape& output_shape,
+                                    const std::vector<argument>& args) const
+{
+    bool transa     = args[0].get_shape().transposed();
+    bool transb     = args[1].get_shape().transposed();
+    auto n_dim      = output_shape.lens().size();
+    auto dim_1      = n_dim - 1;
+    auto dim_0      = n_dim - 2;
+    rocblas_int lda = args[0].get_shape().strides()[transa ? dim_1 : dim_0];
+    rocblas_int ldb = args[1].get_shape().strides()[transb ? dim_1 : dim_0];
+    rocblas_int ldc = args[2].get_shape().strides()[dim_0];
+
+    if(!transb)
+    {
+        if(arg_b.empty())
+        {
+            arg_b = allocate_gpu(args[1].get_shape());
+        }
+        device::pack_a(ctx.get_stream().get(), arg_b, args[1]);
+    }
+
+    // need to pack A in this scenario, use the algorithm to pack B in the
+    // comment of the API
+    if(transa)
+    {
+        if(arg_a.empty())
+        {
+            arg_a = allocate_gpu(args.at(0).get_shape());
+        }
+        device::pack_b(ctx.get_stream().get(), arg_a, args[0]);
+    }
+
+    bool is_3inputs = (args.size() == 4);
+    int32_t beta    = 0;
+    if(is_3inputs)
+    {
+        beta = op.beta;
+    }
+
+    auto a_lens = args[0].get_shape().lens();
+    auto b_lens = args[1].get_shape().lens();
+    output_shape.visit_type([&](auto as) {
+        auto alpha_r    = to_rocblas_type(as(op.alpha));
+        auto beta_r     = to_rocblas_type(as(beta));
+        auto out_lens   = output_shape.lens();
+        rocblas_int m   = out_lens[dim_0];
+        rocblas_int n   = out_lens[dim_1];
+        rocblas_int k   = args[0].get_shape().lens()[dim_1];
+        auto to_pointer = [&](auto&& arg) { return to_rocblas_type(as.from(arg.data())); };
+        assert(k % 4 == 0);
+
+        auto num_matrices = std::accumulate(
+            out_lens.rbegin() + 2, out_lens.rend(), std::size_t{1}, std::multiplies<std::size_t>());
+        if(num_matrices == 1)
+        {
+            // the rocblas_gemm API handles inputs and output matrices as
+            // column-major format. When doing a C = A * B, we actually do
+            // C^T = (B^T) * (A^T). That is the reason we input args[1] as
+            // A and args[0] as B in calling the rocblas_gemm.
+            generic_rocblas_gemm_ex(ctx.get_stream().get_rocblas(),
+                                    transb ? rocblas_operation_transpose : rocblas_operation_none,
+                                    transa ? rocblas_operation_transpose : rocblas_operation_none,
+                                    n,
+                                    m,
+                                    k,
+                                    &alpha_r,
+                                    (!transb) ? to_pointer(arg_b) : to_pointer(args.at(1)),
+                                    rocblas_datatype_i8_r,
+                                    ldb,
+                                    transa ? to_pointer(arg_a) : to_pointer(args.at(0)),
+                                    rocblas_datatype_i8_r,
+                                    lda,
+                                    &beta_r,
+                                    to_pointer(args[2]),
+                                    rocblas_datatype_i32_r,
+                                    ldc,
+                                    is_3inputs ? to_pointer(args.at(3)) : to_pointer(args[2]),
+                                    rocblas_datatype_i32_r,
+                                    ldc,
+                                    rocblas_datatype_i32_r,
+                                    rocblas_gemm_algo_standard,
+                                    0,
+                                    0,
+                                    nullptr,
+                                    nullptr);
+        }
+        else
+        {
+            generic_rocblas_batched_gemm_ex(
+                ctx.get_stream().get_rocblas(),
+                transb ? rocblas_operation_transpose : rocblas_operation_none,
+                transa ? rocblas_operation_transpose : rocblas_operation_none,
+                n,
+                m,
+                k,
+                &alpha_r,
+                (!transb) ? to_pointer(arg_b) : to_pointer(args.at(1)),
+                rocblas_datatype_i8_r,
+                ldb,
+                k * n,
+                transa ? to_pointer(arg_a) : to_pointer(args.at(0)),
+                rocblas_datatype_i8_r,
+                lda,
+                m * k,
+                &beta_r,
+                to_pointer(args[2]),
+                rocblas_datatype_i32_r,
+                ldc,
+                m * n,
+                is_3inputs ? to_pointer(args.at(3)) : to_pointer(args[2]),
+                rocblas_datatype_i32_r,
+                ldc,
+                m * n,
+                num_matrices,
+                rocblas_datatype_i32_r,
+                rocblas_gemm_algo_standard,
+                0,
+                0,
+                nullptr,
+                nullptr);
+        }
+    });
+
+    return is_3inputs ? args.at(3) : args[2];
+}
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx