OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 00:19:58 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

[XLA][codegen] Migrate triton operations that have shared dialect lowerings are implemented for.

Browse Source 
These were missed in previous commits.
Addresses transpose and bitcast.

PiperOrigin-RevId: 826158776
This commit is contained in:
Karlo Basioli 2025-10-30 13:17:07 -07:00 committed by TensorFlower Gardener
parent 1424c4f739
commit f4ebf9d47d
3 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
third_party/xla/xla/backends/gpu/codegen/triton/BUILD vendored
View File

@ -137,8 +137,8 @@ cc_library(
"@llvm-project//mlir:IR",
"@llvm-project//mlir:LLVMDialect",
"@llvm-project//mlir:MathDialect",
"@llvm-project//mlir:NVVMDialect",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TensorDialect",
"@triton//:TritonDialects",
],
)
@ -428,6 +428,7 @@ cc_library(
"@llvm-project//mlir:MathDialect",
"@llvm-project//mlir:Support",
"@local_tsl//tsl/platform:tensor_float_32_hdr_lib",
"@stablehlo//:stablehlo_ops",
"@triton//:TritonDialects",
],
)

5
third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms.cc vendored
View File

@ -36,6 +36,7 @@ limitations under the License.
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LLVM.h"
#include "stablehlo/dialect/StablehloOps.h"
#include "xla/backends/gpu/codegen/triton/emitter_helpers.h"
#include "xla/codegen/emitter_loc_op_builder.h"
#include "xla/hlo/ir/hlo_instruction.h"
@ -160,8 +161,8 @@ absl::StatusOr<Value> ScaledDot(EmitterLocOpBuilder b,
Value rhs_scale;
if (rhs_dot_elem_type != ttir::ScaleDotElemType::BF16) {
rhs_scale = Bitcast(b, operands.rhs_scale, b.getI8Type());
rhs_scale =
b.create<ttir::TransOp>(rhs_scale, mlir::ArrayRef<int32_t>{1, 0});
rhs_scale = b.create<mlir::stablehlo::TransposeOp>(
rhs_scale, b.getDenseI64ArrayAttr({1, 0}));
}
// make type with the same shape as the scale but with i8 type

3
third_party/xla/xla/backends/gpu/codegen/triton/emitter_helpers.cc vendored
View File

@ -35,6 +35,7 @@ limitations under the License.
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
@ -501,7 +502,7 @@ absl::StatusOr<Value> EmitElementwise(EmitterLocOpBuilder& b,
mh::ComparisonDirection::NE),
inputs[1], inputs[2]);
case HloOpcode::kReducePrecision:
return mh::reducePrecision<mt::BitcastOp>(
return mh::reducePrecision<mlir::tensor::BitcastOp>(
b.getLoc(), inputs[0], hlo.exponent_bits(), hlo.mantissa_bits(), &b);
default:
return absl::InvalidArgumentError(