[MLIR:HLO] Add window_reversal attribute to convolution attributes.
- Add this attribute to match the corresponding XLA HLO attribute on convolution operations. - A true value indicates a reversal of the corresponding kernel spatial dimension. - Since XLA builder does not support this attribute, use a custom HLO converted to map from mlir::mhlo::ConvOp to XLA. PiperOrigin-RevId: 346891737
This commit is contained in:
parent
3442ac270d
commit
f232da1f9d
|
@ -902,6 +902,7 @@ def HLO_ConvOp : HLO_Op<"convolution", [NoSideEffect]>, BASE_HLO_ConvOp {
|
||||||
ConvolutionAttributes.attributes);
|
ConvolutionAttributes.attributes);
|
||||||
|
|
||||||
let results = (outs HLO_Tensor);
|
let results = (outs HLO_Tensor);
|
||||||
|
let hasCustomHLOConverter = 1;
|
||||||
}
|
}
|
||||||
|
|
||||||
def HLO_CopyOp: HLO_Op<"copy", [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_CopyOp {
|
def HLO_CopyOp: HLO_Op<"copy", [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_CopyOp {
|
||||||
|
|
|
@ -958,6 +958,17 @@ def HLO_PrecisionConfigAttr:
|
||||||
OptionalAttr<
|
OptionalAttr<
|
||||||
TypedArrayAttrBase<HLO_PrecisionAttr, "Precision Config attribute">>;
|
TypedArrayAttrBase<HLO_PrecisionAttr, "Precision Config attribute">>;
|
||||||
|
|
||||||
|
def BoolElementsAttr :
|
||||||
|
ElementsAttrBase<
|
||||||
|
And<[CPred<"$_self.isa<::mlir::DenseIntOrFPElementsAttr>()">,
|
||||||
|
CPred<"$_self.cast<::mlir::DenseIntOrFPElementsAttr>().getType().getElementType().isInteger(1)">]>,
|
||||||
|
"constant boolean vector/tensor attribute"> {
|
||||||
|
let storageType = [{ ::mlir::DenseElementsAttr }];
|
||||||
|
let returnType = [{ ::mlir::DenseElementsAttr }];
|
||||||
|
|
||||||
|
let convertFromStorage = "$_self";
|
||||||
|
}
|
||||||
|
|
||||||
def ConvolutionAttributes {
|
def ConvolutionAttributes {
|
||||||
dag attributes = (ins
|
dag attributes = (ins
|
||||||
// Default value: one for each of the spatial dimension.
|
// Default value: one for each of the spatial dimension.
|
||||||
|
@ -968,6 +979,8 @@ def ConvolutionAttributes {
|
||||||
OptionalAttr<I64ElementsAttr>:$lhs_dilation,
|
OptionalAttr<I64ElementsAttr>:$lhs_dilation,
|
||||||
// Default value: one for each of the spatial dimension.
|
// Default value: one for each of the spatial dimension.
|
||||||
OptionalAttr<I64ElementsAttr>:$rhs_dilation,
|
OptionalAttr<I64ElementsAttr>:$rhs_dilation,
|
||||||
|
// Default value: one for each of the spatial dimension.
|
||||||
|
OptionalAttr<BoolElementsAttr>:$window_reversal,
|
||||||
ConvDimensionNumbers:$dimension_numbers,
|
ConvDimensionNumbers:$dimension_numbers,
|
||||||
I64Attr:$feature_group_count,
|
I64Attr:$feature_group_count,
|
||||||
I64Attr:$batch_group_count,
|
I64Attr:$batch_group_count,
|
||||||
|
@ -983,6 +996,14 @@ class BASE_HLO_ConvOp {
|
||||||
|
|
||||||
See https://www.tensorflow.org/xla/operation_semantics#conv_convolution.
|
See https://www.tensorflow.org/xla/operation_semantics#conv_convolution.
|
||||||
}];
|
}];
|
||||||
|
|
||||||
|
code extraClassDeclaration = [{
|
||||||
|
bool hasWindowReversal() {
|
||||||
|
auto reversal = window_reversalAttr();
|
||||||
|
return reversal && llvm::any_of(reversal.getBoolValues(),
|
||||||
|
[](bool v) { return v; });
|
||||||
|
}
|
||||||
|
}];
|
||||||
}
|
}
|
||||||
|
|
||||||
class BASE_HLO_CopyOp {
|
class BASE_HLO_CopyOp {
|
||||||
|
|
|
@ -243,7 +243,8 @@ struct ConvToLinalgConverter : public OpConversionPattern<lmhlo::ConvOp> {
|
||||||
}
|
}
|
||||||
|
|
||||||
// TODO: LHS dilation for deconvolution not supported yet.
|
// TODO: LHS dilation for deconvolution not supported yet.
|
||||||
if (op.lhs_dilation()) {
|
// TODO(jurahul): Window reversal is not supported yet.
|
||||||
|
if (op.lhs_dilation() || op.hasWindowReversal()) {
|
||||||
return failure();
|
return failure();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -103,7 +103,8 @@ func @conv_backinput(%input : memref<4x5x16x16xf64>, %filter : memref<5x3x7x7xf6
|
||||||
precision_config = [],
|
precision_config = [],
|
||||||
result_scale = 1.000000e+00 : f64,
|
result_scale = 1.000000e+00 : f64,
|
||||||
rhs_dilation = dense<1> : tensor<2xi64>,
|
rhs_dilation = dense<1> : tensor<2xi64>,
|
||||||
window_strides = dense<1> : tensor<2xi64>}
|
window_strides = dense<1> : tensor<2xi64>,
|
||||||
|
window_reversal = dense<true>: tensor<2xi1>}
|
||||||
: (memref<4x5x16x16xf64>, memref<5x3x7x7xf64>, memref<4x3x16x16xf64>, memref<32xui8>) -> ()
|
: (memref<4x5x16x16xf64>, memref<5x3x7x7xf64>, memref<4x3x16x16xf64>, memref<32xui8>) -> ()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue