-
Notifications
You must be signed in to change notification settings - Fork 29
/
PolyOps.cpp
67 lines (57 loc) · 2.65 KB
/
PolyOps.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
#include "include/Dialect/Poly/IR/PolyOps.h"
#include "llvm/include/llvm/ADT/APInt.h" // from @llvm-project
#include "mlir/include/mlir/IR/Builders.h" // from @llvm-project
#include "mlir/include/mlir/Support/LogicalResult.h" // from @llvm-project
namespace mlir {
namespace heir {
namespace poly {
void FromTensorOp::build(OpBuilder &builder, OperationState &result,
Value input, RingAttr ring) {
TensorType tensorType = dyn_cast<TensorType>(input.getType());
APInt cmod(APINT_BIT_WIDTH, 1);
cmod = cmod << tensorType.getElementTypeBitWidth();
Type resultType = PolyType::get(builder.getContext(), ring);
build(builder, result, resultType, input);
}
LogicalResult FromTensorOp::verify() {
auto tensorShape = getInput().getType().getShape();
auto ring = getOutput().getType().getRing();
auto polyDegree = ring.getIdeal().getDegree();
bool compatible = tensorShape.size() == 1 && tensorShape[0] <= polyDegree;
if (!compatible) {
return emitOpError()
<< "input type " << getInput().getType()
<< " does not match output type " << getOutput().getType()
<< ". The input type must be a tensor of shape [d] where d "
"is at most the degree of the polynomial generator of "
"the output ring's ideal.";
}
APInt coefficientModulus = ring.coefficientModulus();
unsigned cmodBitWidth = coefficientModulus.ceilLogBase2();
unsigned inputBitWidth = getInput().getType().getElementTypeBitWidth();
if (inputBitWidth > cmodBitWidth) {
return emitOpError() << "input tensor element type "
<< getInput().getType().getElementType()
<< " is too large to fit in the coefficients of "
<< getOutput().getType()
<< ". The input tensor's elements must be rescaled"
" to fit before using from_tensor.";
}
return success();
}
LogicalResult ToTensorOp::verify() {
auto tensorShape = getOutput().getType().getShape();
auto polyDegree = getInput().getType().getRing().getIdeal().getDegree();
bool compatible = tensorShape.size() == 1 && tensorShape[0] == polyDegree;
return compatible
? success()
: emitOpError()
<< "input type " << getInput().getType()
<< " does not match output type " << getOutput().getType()
<< ". The input type must be a tensor of shape [d] where d "
"is exactly the degree of the polynomial generator of "
"the output ring's ideal.";
}
} // namespace poly
} // namespace heir
} // namespace mlir