mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
ebfe631ed8
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/36952 Differential Revision: D21139939 Test Plan: Imported from OSS Pulled By: ZolotukhinM fbshipit-source-id: a6605c0d6ccbb049ce27e6cdcc8fd8d2ebc057e3
111 lines
3.4 KiB
C++
111 lines
3.4 KiB
C++
#include <test/cpp/tensorexpr/test_base.h>
|
|
#include <torch/csrc/jit/ir/ir.h>
|
|
#include <torch/csrc/jit/ir/irparser.h>
|
|
#include <torch/csrc/jit/tensorexpr/buffer.h>
|
|
#include <torch/csrc/jit/tensorexpr/kernel.h>
|
|
#include <torch/csrc/jit/tensorexpr/loopnest.h>
|
|
#include <torch/csrc/jit/tensorexpr/tensor.h>
|
|
#include <torch/torch.h>
|
|
#include <cmath>
|
|
#include <sstream>
|
|
#include <stdexcept>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
|
|
using namespace torch::indexing;
|
|
using namespace torch::jit;
|
|
using namespace torch::jit::tensorexpr;
|
|
|
|
void testKernel_1() {
|
|
KernelScope kernel_scope;
|
|
|
|
const auto graph_string = R"IR(
|
|
graph(%0 : Float(5:3,3:1),
|
|
%1 : Float(5:3,3:1)):
|
|
%2 : Float(5:3,3:1) = aten::mul(%0, %1)
|
|
%3 : Float(5:3,3:1) = aten::mul(%0, %2)
|
|
return (%3))IR";
|
|
auto graph = std::make_shared<Graph>();
|
|
parseIR(graph_string, &*graph);
|
|
|
|
auto a = at::rand({5, 3}, TensorOptions(kCPU).dtype(at::kFloat));
|
|
auto b = at::rand({5, 3}, TensorOptions(kCPU).dtype(at::kFloat));
|
|
auto o = at::zeros({5, 3}, TensorOptions(kCPU).dtype(at::kFloat));
|
|
auto ref = a * (a * b);
|
|
TensorExprKernel k(graph);
|
|
std::vector<at::Tensor> inputs = {a, b};
|
|
Stmt* s = k.getStmtForInputs(fmap<IValue>(inputs));
|
|
// TODO: verify stmt
|
|
|
|
std::vector<IValue> stack = fmap<IValue>(inputs);
|
|
k.run(stack);
|
|
o = stack[0].toTensor();
|
|
for (size_t i = 0; i < 5 * 3; i++) {
|
|
CHECK_EQ(((float*)o.data_ptr())[i], ((float*)ref.data_ptr())[i]);
|
|
}
|
|
}
|
|
|
|
void testKernel_2() {
|
|
KernelScope kernel_scope;
|
|
|
|
const auto graph_string = R"IR(
|
|
graph(%0 : Float(5:3,3:1),
|
|
%1 : Float(5:1,3:5)):
|
|
%2 : Float(5:3,3:1) = aten::mul(%0, %1)
|
|
%3 : Float(5:3,3:1) = aten::mul(%0, %2)
|
|
return (%3))IR";
|
|
auto graph = std::make_shared<Graph>();
|
|
parseIR(graph_string, &*graph);
|
|
|
|
auto a = at::rand({5, 3}, TensorOptions(kCPU).dtype(at::kFloat));
|
|
auto b =
|
|
at::rand({3, 5}, TensorOptions(kCPU).dtype(at::kFloat)).transpose(0, 1);
|
|
auto o = at::zeros({5, 3}, TensorOptions(kCPU).dtype(at::kFloat));
|
|
auto ref = a * (a * b);
|
|
TensorExprKernel k(graph);
|
|
std::vector<at::Tensor> inputs = {a, b};
|
|
Stmt* s = k.getStmtForInputs(fmap<IValue>(inputs));
|
|
// TODO: verify stmt
|
|
|
|
std::vector<IValue> stack = fmap<IValue>(inputs);
|
|
k.run(stack);
|
|
o = stack[0].toTensor();
|
|
for (size_t i = 0; i < 5 * 3; i++) {
|
|
CHECK_EQ(((float*)o.data_ptr())[i], ((float*)ref.data_ptr())[i]);
|
|
}
|
|
}
|
|
|
|
void testKernel_3() {
|
|
KernelScope kernel_scope;
|
|
|
|
const auto graph_string = R"IR(
|
|
graph(%0 : Float(5:3,3:1),
|
|
%1 : Float(5:12,3:2)):
|
|
%2 : Float(5:3,3:1) = aten::mul(%0, %1)
|
|
%3 : Float(5:3,3:1) = aten::mul(%0, %2)
|
|
return (%3))IR";
|
|
auto graph = std::make_shared<Graph>();
|
|
parseIR(graph_string, &*graph);
|
|
|
|
auto a = at::rand({5, 3}, TensorOptions(kCPU).dtype(at::kFloat));
|
|
auto b = at::rand({10, 6}, TensorOptions(kCPU).dtype(at::kFloat))
|
|
.index({Slice(None, None, 2), Slice(None, None, 2)});
|
|
auto o = at::zeros({5, 3}, TensorOptions(kCPU).dtype(at::kFloat));
|
|
auto ref = a * (a * b);
|
|
TensorExprKernel k(graph);
|
|
std::vector<at::Tensor> inputs = {a, b};
|
|
Stmt* s = k.getStmtForInputs(fmap<IValue>(inputs));
|
|
// TODO: verify stmt
|
|
|
|
std::vector<IValue> stack = fmap<IValue>(inputs);
|
|
k.run(stack);
|
|
o = stack[0].toTensor();
|
|
for (size_t i = 0; i < 5 * 3; i++) {
|
|
CHECK_EQ(((float*)o.data_ptr())[i], ((float*)ref.data_ptr())[i]);
|
|
}
|
|
}
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|