Debugging transpose operator for now

src/include/rtg/operators.hpp

@@ -227,6 +227,46 @@ struct activation
     }
 };
 
+struct transpose
+{
+    std::vector<int64_t> dims;
+    std::string name() const { return "transpose"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs}.has(1);
+        auto input         = inputs.at(0);
+        auto input_size    = input.lens();
+        auto input_strides = input.strides();
+        auto t             = input.type();
+        if (dims.size() != input_size.size()) {
+            RTG_THROW("Permutation has wrong number of axes");
+        }
+
+        // DEBUG
+        for (int i = 0; i < dims.size(); i++) {
+            std::cout << dims[i] << std::endl;
+        }
+        std::cout << std::endl;
+
+        std::vector<int64_t> axes(dims.size());
+        std::iota(axes.begin(), axes.end(), 0);
+        if (!std::is_permutation(axes.begin(), axes.end(), dims.begin())) {
+            RTG_THROW("Invalid permutation");
+        }
+        std::vector<size_t> output_size(input_size.size());
+        std::vector<size_t> output_strides(input_size.size());
+        for (int i = 0; i < output_size.size(); i++) {
+            output_size[i] = input_size[dims[i]];
+            output_strides[i] = input_strides[dims[i]];
+            //std::cout << input_size[i] << "    " << output_size[i] << std::endl;
+            //std::cout << input_strides[i] << "    " << output_strides[i] << std::endl;
+        }
+        std::cout << std::endl;
+        return {t, output_size, output_strides};
+    }
+    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
+};
+
 struct reshape
 {
     std::vector<int64_t> dims;

src/targets/cpu/cpu_target.cpp

@@ -64,13 +64,31 @@ struct cpu_transpose
     {
         argument result{output_shape};
         visit_all(result, args[0])([&](auto output, auto input) {
-            dfor(output_shape.lens()[0],
-                 output_shape.lens()[0],
-                 output_shape.lens()[0],
-                 output_shape.lens()[0]
+            using value_type = typename decltype(input)::value_type;
+            value_type* ptr = static_cast<value_type*>(output.data());
+            auto nb = output_shape.lens()[0];
+            auto nc = output_shape.lens()[1]; 
+            auto nh = output_shape.lens()[2]; 
+            auto nw = output_shape.lens()[3];
+            for (int kk = 0; kk < 4; kk++) {
+                std::cout << "cpu_transpose: " << output_shape.lens()[kk] << "    " << output_shape.strides()[kk] << std::endl;
+            }
+            
+            for (int b = 0; b < nb; b++) {
+                for (int c = 0; c < nc; c++) {
+                    for (int i = 0; i < nh; i++) {
+                        for (int j = 0; j < nw; j++) {
+                            *ptr++ = input(b,c,i,j);
+                             std::cout << input(b,c,i,j) << "  ";
+                        }
+                    }
+                }
+            }
+            std::cout << std::endl;
         });
+        return result;
     }
-}
+};
 
 struct cpu_reshape
 {
@@ -371,6 +389,10 @@ struct cpu_apply
             {
                 apply_reshape(it);
             }
+            else if(it->op.name() == "transpose")
+            {
+                apply_transpose(it);
+            }
             else if(it->op.name() == "activation")
             {
                 apply_activation(it);
@@ -432,6 +454,12 @@ struct cpu_apply
         prog->replace_instruction(ins, cpu_reshape{op}, ins->arguments);
     }
 
+    void apply_transpose(instruction_ref ins)
+    {
+        auto&& op = any_cast<transpose>(ins->op);
+        prog->replace_instruction(ins, cpu_transpose{op}, ins->arguments);
+    }
+
     void apply_activation(instruction_ref ins)
     {
         auto&& op = any_cast<activation>(ins->op);

test/cpu_ops_test.cpp

@@ -392,6 +392,26 @@ void conv2d_padding_stride_test()
     EXPECT(test::verify_range(results_vector, s));
 }
 
+void transpose_test()
+{
+    rtg::shape a_shape{rtg::shape::float_type, {1,2,2,3}};
+    std::vector<float> data(12);
+    std::iota(data.begin(), data.end(), 0);
+
+    rtg::program p;
+    auto l = p.add_literal(rtg::literal{a_shape, data});
+    std::vector<int64_t> perm = {0,3,1,2};
+    p.add_instruction(rtg::transpose{perm}, l);
+    p.compile(rtg::cpu::cpu_target{});
+    auto result = p.eval({});
+    std::vector<float> results_vector(12);
+    result.visit([&] (auto output){ results_vector.assign(output.begin(), output.end()); });
+    float tol = 1e-6;
+    for (int i = 0; i < results_vector.size(); i++) {
+      std::cout << results_vector[i] << std::endl;
+    }
+}
+
 int main()
 {
     exp_test();
@@ -400,6 +420,7 @@ int main()
     tan_test();
     gemm_test();
     reshape_test();
+    transpose_test();
     softmax_test();
     conv2d_test();
     conv2d_padding_test();