Skip to content

Commit

Permalink
Add "PassUtils" for common functions to avoid duplicate definitions (#…
Browse files Browse the repository at this point in the history
…2547)

[back-porting to
gh](openxla/xla@5a07f58)
  • Loading branch information
abhigunj committed Sep 16, 2024
1 parent 7f88b35 commit e05a320
Show file tree
Hide file tree
Showing 11 changed files with 111 additions and 96 deletions.
2 changes: 2 additions & 0 deletions BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -1096,6 +1096,7 @@ cc_library(
srcs = [
"stablehlo/transforms/ChloLegalizeToStablehlo.cpp",
"stablehlo/transforms/PassPipelines.cpp",
"stablehlo/transforms/PassUtils.cpp",
"stablehlo/transforms/ShapeLegalizeToStablehlo.cpp",
"stablehlo/transforms/StablehloAggressiveFolder.cpp",
"stablehlo/transforms/StablehloAggressiveSimplification.cpp",
Expand All @@ -1115,6 +1116,7 @@ cc_library(
],
hdrs = [
"stablehlo/transforms/MapStablehloToVhlo.h",
"stablehlo/transforms/PassUtils.h",
"stablehlo/transforms/Passes.h",
"stablehlo/transforms/StablehloRefineShapes.h",
],
Expand Down
42 changes: 7 additions & 35 deletions stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -3029,28 +3029,11 @@ func.func @next_after_f32(%x: tensor<2xf32>, %y: tensor<2xf32>) -> tensor<2xf32>

// -----

// CHECK-LABEL: @tan_f16
// CHECK-SAME: (%[[ARG:.*]]: tensor<f16>)
func.func @tan_f16(%arg : tensor<f16>) -> tensor<f16> {
// %[[TMP_0:.*]] = stablehlo.convert [[ARG]] : (tensor<f16>) -> tensor<f32>
// %[[TMP_1:.*]] = stablehlo.sine %[[TMP_0]]
// %[[TMP_2:.*]] = stablehlo.cosine %[[TMP_0]]
// %[[TMP_3:.*]] = stablehlo.divide %[[TMP_1]], %[[TMP_2]]
// %[[TMP_4:.*]] = stablehlo.convert %[[TMP_3]] : (tensor<f32>) -> tensor<f16>
// return %[[TMP_4]] : tensor<f16>
%1 = chlo.tan %arg : tensor<f16> -> tensor<f16>
func.return %1 : tensor<f16>
}

// -----

// CHECK-LABEL: @tan_f32
// CHECK-SAME: (%[[ARG:.*]]: tensor<f32>)
func.func @tan_f32(%arg : tensor<f32>) -> tensor<f32> {
// %[[TMP_0:.*]] = stablehlo.sine %[[ARG]]
// %[[TMP_1:.*]] = stablehlo.cosine %[[ARG]]
// %[[TMP_2:.*]] = stablehlo.divide %[[TMP_0]], %[[TMP_1]]
// return %[[TMP_2]] : tensor<f32>
// CHECK: %[[TMP_0:.*]] = stablehlo.tan %[[ARG]] : tensor<f32>
// CHECK: return %[[TMP_0]] : tensor<f32>
%1 = chlo.tan %arg : tensor<f32> -> tensor<f32>
func.return %1 : tensor<f32>
}
Expand All @@ -3060,22 +3043,11 @@ func.func @tan_f32(%arg : tensor<f32>) -> tensor<f32> {
// CHECK-LABEL: @tan_complexf32
// CHECK-SAME: %[[ARG0:.+]]: tensor<1xf32>, %[[ARG1:.+]]: tensor<1xf32>
func.func @tan_complexf32(%arg0 : tensor<1xf32>, %arg1 : tensor<1xf32>) -> (tensor<1xf32>, tensor<1xf32>) {
// CHECK: %[[COMPLEX:.+]] = stablehlo.complex %[[ARG0]], %[[ARG1]] : tensor<1xcomplex<f32>>
// CHECK: %[[REAL:.+]] = stablehlo.real %[[COMPLEX]] : (tensor<1xcomplex<f32>>) -> tensor<1xf32>
// CHECK: %[[SINE:.+]] = stablehlo.sine %[[REAL]]
// CHECK: %[[COS:.+]] = stablehlo.cosine %[[REAL]]
// CHECK: %[[TAN:.+]] = stablehlo.divide %[[SINE]], %[[COS]]
// CHECK: %[[IMAG:.+]] = stablehlo.imag %[[COMPLEX]] : (tensor<1xcomplex<f32>>) -> tensor<1xf32>
// CHECK: %[[TANH:.+]] = stablehlo.tanh %[[IMAG]]
// CHECK: %[[NUM:.+]] = stablehlo.complex %[[TAN]], %[[TANH]]
// CHECK: %[[ONE:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<1xf32>
// CHECK: %[[MUL:.+]] = stablehlo.multiply %[[TAN]], %[[TANH]]
// CHECK: %[[NEG:.+]] = stablehlo.negate %[[MUL]]
// CHECK: %[[DEN:.+]] = stablehlo.complex %[[ONE]], %[[NEG]]
// CHECK: %[[RES:.+]] = stablehlo.divide %[[NUM]], %[[DEN]]
// CHECK: %[[REAL:.+]] = stablehlo.real %[[RES]]
// CHECK: %[[IMAG:.+]] = stablehlo.imag %[[RES]]
// CHECK: return %[[REAL]], %[[IMAG]]
// CHECK: %[[TMP_0:.*]] = stablehlo.complex %[[ARG0]], %[[ARG1]] : tensor<1xcomplex<f32>>
// CHECK: %[[TMP_1:.*]] = stablehlo.tan %[[TMP_0]] : tensor<1xcomplex<f32>>
// CHECK: %[[TMP_2:.*]] = stablehlo.real %[[TMP_1]] : (tensor<1xcomplex<f32>>) -> tensor<1xf32>
// CHECK: %[[TMP_3:.*]] = stablehlo.imag %[[TMP_1]] : (tensor<1xcomplex<f32>>) -> tensor<1xf32>
// CHECK: return %[[TMP_2]], %[[TMP_3]] : tensor<1xf32>, tensor<1xf32>
%0 = stablehlo.complex %arg0, %arg1 : tensor<1xcomplex<f32>>
%1 = chlo.tan %0 : tensor<1xcomplex<f32>> -> tensor<1xcomplex<f32>>
%2 = stablehlo.real %1 : (tensor<1xcomplex<f32>>) -> tensor<1xf32>
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
// RUN: stablehlo-opt %s -verify-diagnostics -split-input-file -allow-unregistered-dialect --stablehlo-create-compatibility-expander='target=1.0.0' | FileCheck %s --check-prefixes=CHECK
// RUN: stablehlo-opt %s -verify-diagnostics -split-input-file --stablehlo-create-compatibility-expander='target=1.6.0' | FileCheck %s --check-prefixes=CHECK-NO-DOWNGRADE
// RUN: stablehlo-opt %s -verify-diagnostics -split-input-file --stablehlo-create-compatibility-expander='target=1.0.0' --chlo-legalize-to-stablehlo | FileCheck %s --check-prefixes=CHECK
// RUN: stablehlo-opt %s -verify-diagnostics -split-input-file --stablehlo-create-compatibility-expander='target=1.6.0' --chlo-legalize-to-stablehlo | FileCheck %s --check-prefixes=CHECK-NO-DOWNGRADE

// -----

Expand Down
1 change: 1 addition & 0 deletions stablehlo/transforms/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ add_mlir_dialect_library(StablehloPasses
StablehloRefineShapes.cpp
VhloLegalizeToStablehlo.cpp
VhloToVersion.cpp
PassUtils.cpp

DEPENDS
ChloDecompositionPatternsIncGen
Expand Down
21 changes: 2 additions & 19 deletions stablehlo/transforms/ChloDecompositionPatterns.td
Original file line number Diff line number Diff line change
Expand Up @@ -109,26 +109,9 @@ def : Pat<(CHLO_IsNegInfOp NonComplexElementType:$input),
(STABLEHLO_DEFAULT_COMPARISON_TYPE)
)>;

// Express `tan` as
// sine(x) / cosine(x)
def : Pat<(CHLO_TanOp NonComplexElementType:$input),
(StableHLO_DivOp
(StableHLO_SineOp $input),
(StableHLO_CosineOp $input)
)>;


// Express `tan(a + bi)` as
// (tan(a) + i tanh(b)) / (1 - i tan(a) * tanh(b))
def : Pat<(CHLO_TanOp ComplexElementType:$input),
(StableHLO_DivOp
(StableHLO_ComplexOp
(CHLO_TanOp:$tan (StableHLO_RealOp $input)),
(StableHLO_TanhOp:$tanh (StableHLO_ImagOp $input))),
(StableHLO_ComplexOp
(StableHLO_ConstantLike<"1.0"> $tan),
(StableHLO_NegOp (StableHLO_MulOp $tan, $tanh)))
)>;
def : Pat<(CHLO_TanOp $input),
(StableHLO_TanOp $input)>;

def : Pat<(CHLO_ConstantOp $v),
(StableHLO_ConstantOp $v)>;
Expand Down
24 changes: 1 addition & 23 deletions stablehlo/transforms/ChloLegalizeToStablehlo.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
#include "stablehlo/dialect/BroadcastUtils.h"
#include "stablehlo/dialect/ChloOps.h"
#include "stablehlo/dialect/StablehloOps.h"
#include "stablehlo/transforms/PassUtils.h"
#include "stablehlo/transforms/Passes.h"

namespace mlir {
Expand Down Expand Up @@ -171,29 +172,6 @@ static void populateForBroadcastingBinaryOp(MLIRContext *context,
context, args...);
}

template <typename T>
static Value getConstantLike(OpBuilder &b, Location loc, T constant,
Value val) {
Type ty = getElementTypeOrSelf(val.getType());
auto getAttr = [&]() -> Attribute {
if (isa<IntegerType>(ty)) return b.getIntegerAttr(ty, constant);
if (isa<FloatType>(ty)) return b.getFloatAttr(ty, constant);
if (auto complexTy = dyn_cast<ComplexType>(ty)) {
return complex::NumberAttr::get(complexTy, constant, 0);
}
llvm_unreachable("unhandled element type");
};
return b.create<mlir::chlo::ConstantLikeOp>(loc, cast<TypedAttr>(getAttr()),
val);
}

static Value getConstantLike(OpBuilder &b, Location loc,
const APFloat &constant, Value val) {
Type ty = getElementTypeOrSelf(val.getType());
return b.create<mlir::chlo::ConstantLikeOp>(loc, b.getFloatAttr(ty, constant),
val);
}

static Value getConstantLikeMaxFiniteValue(OpBuilder &b, Location loc,
Value val) {
auto ty = cast<FloatType>(getElementTypeOrSelf(val.getType()));
Expand Down
34 changes: 34 additions & 0 deletions stablehlo/transforms/PassUtils.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
/* Copyright 2024 The StableHLO Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "stablehlo/transforms/PassUtils.h"

#include "mlir/IR/Builders.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LLVM.h"
#include "stablehlo/dialect/ChloOps.h"

namespace mlir {
namespace stablehlo {

Value getConstantLike(OpBuilder &b, Location loc, const APFloat &constant,
Value val) {
Type ty = getElementTypeOrSelf(val.getType());
return b.create<mlir::chlo::ConstantLikeOp>(loc, b.getFloatAttr(ty, constant),
val);
}

} // namespace stablehlo
} // namespace mlir
57 changes: 57 additions & 0 deletions stablehlo/transforms/PassUtils.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
/* Copyright 2024 The StableHLO Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#ifndef THIRD_PARTY_STABLEHLO_STABLEHLO_TRANSFORMS_PASS_UTILS_H_
#define THIRD_PARTY_STABLEHLO_STABLEHLO_TRANSFORMS_PASS_UTILS_H_

#include "llvm/Support/ErrorHandling.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LLVM.h"
#include "stablehlo/dialect/ChloOps.h"

namespace mlir {
namespace stablehlo {
// Add utility functions common across passes.

// Creates a chlo::ConstantLikeOp using a splat `constant` of the same shape
// as `val`.
template <typename T>
Value getConstantLike(OpBuilder &b, Location loc, T constant, Value val) {
Type ty = getElementTypeOrSelf(val.getType());
auto getAttr = [&]() -> Attribute {
if (isa<IntegerType>(ty)) return b.getIntegerAttr(ty, constant);
if (isa<FloatType>(ty)) return b.getFloatAttr(ty, constant);
if (auto complexTy = dyn_cast<ComplexType>(ty)) {
return complex::NumberAttr::get(complexTy, constant, 0);
}
llvm_unreachable("unhandled element type");
};
return b.create<mlir::chlo::ConstantLikeOp>(loc, cast<TypedAttr>(getAttr()),
val);
}

// Creates a chlo::ConstantLikeOp using a APFloat splat `constant` of the
// same shape as `val`.
Value getConstantLike(OpBuilder &b, Location loc, const APFloat &constant,
Value val);

} // namespace stablehlo
} // namespace mlir

#endif // THIRD_PARTY_STABLEHLO_STABLEHLO_TRANSFORMS_PASS_UTILS_H_
1 change: 1 addition & 0 deletions stablehlo/transforms/Passes.td
Original file line number Diff line number Diff line change
Expand Up @@ -338,5 +338,6 @@ def StablehloCreateCompatibilityExpanderPass : Pass<"stablehlo-create-compatibil
];
let dependentDialects = [
"mlir::stablehlo::StablehloDialect",
"mlir::chlo::ChloDialect",
];
}
16 changes: 1 addition & 15 deletions stablehlo/transforms/StablehloCreateCompatibilityExpander.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ limitations under the License.
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "stablehlo/dialect/StablehloOps.h"
#include "stablehlo/dialect/Version.h"
#include "stablehlo/transforms/PassUtils.h"
#include "stablehlo/transforms/Passes.h"

namespace mlir {
Expand All @@ -38,21 +39,6 @@ namespace {
// Helpers.
//===----------------------------------------------------------------------===//

// Creates a constant with all ones.
static Value createConstantWithAllOnes(OpBuilder &b, Location loc, Value val) {
if (!isa<FloatType>(getElementTypeOrSelf(val)))
llvm_unreachable("Unsupported element type, expecting float");

auto shapedTy = dyn_cast<mlir::ShapedType>(val.getType());
if (!shapedTy) llvm_unreachable("Unsupported shaped type.");

mlir::DenseElementsAttr elementsAttr =
mlir::DenseElementsAttr::get(shapedTy, 1.0);

return b.create<mlir::stablehlo::ConstantOp>(loc, val.getType(),
elementsAttr);
}

// Check user-specified target version.
vhlo::Version validateTargetVersion(llvm::StringRef versionRef) {
auto failOrVersion = vhlo::Version::fromString(versionRef);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@ def NonComplexElementType : Type<
CPred<"!isa<ComplexType>(cast<ShapedType>($_self).getElementType())">,
"Non-complex element type">;

def createConstantWithAllOnes : NativeCodeCall<"createConstantWithAllOnes($_builder, $_loc, $0)">;
class StableHLO_ConstantLike<string value> : NativeCodeCall<
"::mlir::stablehlo::getConstantLike($_builder, $_loc, " # value # ", $0)">;

// Express `tan` as
// sine(x) / cosine(x)
Expand All @@ -42,6 +43,6 @@ def TanOp_ComplexElementType_CompatiblityExpander : Pat<(StableHLO_TanOp Complex
(StableHLO_TanOp:$tan (StableHLO_RealOp $input)),
(StableHLO_TanhOp:$tanh (StableHLO_ImagOp $input))),
(StableHLO_ComplexOp
(createConstantWithAllOnes $tan),
(StableHLO_ConstantLike<"1.0"> $tan),
(StableHLO_NegOp (StableHLO_MulOp $tan, $tanh)))
)>;

0 comments on commit e05a320

Please sign in to comment.