mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
3a0165da49
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/65551 Previously we had a big switch on Op kind to decide how to lower a given JIT operator to NNC. This PR changes this switch to a hash table lookup. Why? This helps us with at least two things: 1) With this approach we can easily check if we know how to handle a given node in advance - i.e. we can inspect the entire graph and tell whether it's possible to compile it or not without actually trying to do that and dying in the middle. This would allow us to, say, provide user-friendly error messages in AOT workflow. 2) We can switch to use schema instead of op kind to determine correct lowering. Unlike op schema, op kind might be ambigous (see e.g. #64963) and using it instead of schema can lead to bugs. Test Plan: Imported from OSS Reviewed By: navahgar Differential Revision: D31148926 Pulled By: ZolotukhinM fbshipit-source-id: ac12684e2126c899426ef5e4cc1e3f70fa01f704
44 lines
1.3 KiB
C++
44 lines
1.3 KiB
C++
#include <gtest/gtest.h>
|
|
#include <torch/csrc/jit/tensorexpr/eval.h>
|
|
#include <torch/csrc/jit/tensorexpr/loopnest.h>
|
|
#include <torch/csrc/jit/tensorexpr/operators/operators.h>
|
|
#include <torch/torch.h>
|
|
|
|
using namespace torch::jit::tensorexpr;
|
|
|
|
using Tensors = std::vector<Tensor>;
|
|
using Args = std::vector<CodeGen::BufferArg>;
|
|
std::unique_ptr<SimpleIREvaluator> compile(
|
|
const Args& inputs,
|
|
const Tensors& outputs) {
|
|
LoopNest nest({outputs});
|
|
nest.prepareForCodegen();
|
|
nest.simplify();
|
|
auto join = inputs;
|
|
join.insert(join.end(), outputs.begin(), outputs.end());
|
|
return std::make_unique<SimpleIREvaluator>(nest.root_stmt(), join);
|
|
}
|
|
|
|
TEST(Ops, Sum) {
|
|
constexpr int M = 8;
|
|
constexpr int N = 16;
|
|
std::vector<IntList> testDims = {{0}, {1}, {0, 1}};
|
|
std::vector<std::vector<ExprHandle>> outputShapes = {{N}, {M}, {}};
|
|
for (int idx = 0; idx < testDims.size(); idx++) {
|
|
const auto& dims = testDims[idx];
|
|
const auto& outShape = outputShapes[idx];
|
|
|
|
BufHandle a("a", {M, N}, kFloat);
|
|
Tensor b = computeSum({a, dims, false}, outShape, c10::kFloat, at::kCPU);
|
|
auto cg = compile({a}, {b});
|
|
|
|
auto at = at::arange(M * N, at::kFloat).view({M, N});
|
|
auto ref = at::sum(at, dims);
|
|
auto bt = at::empty_like(ref);
|
|
|
|
cg->call({at.data_ptr<float>(), bt.data_ptr<float>()});
|
|
|
|
ASSERT_TRUE(at::allclose(bt, ref));
|
|
}
|
|
}
|