From d4ba9787cd1918a1e73c0c917983fa8d629f5c52 Mon Sep 17 00:00:00 2001 From: Christian Sigg Date: Fri, 28 Aug 2026 23:35:42 -0700 Subject: [PATCH 1/4] Add HLO to TensorIR conversion and lit tests. Implement conversion from HLO fusion computations to the TensorIR dialect, supported operations analysis, the `fusion_to_tensorir` test tool, and lit test suites. PiperOrigin-RevId: 972986533 --- third_party/xla/MODULE.bazel | 2 + .../third_party/extensions/third_party.bzl | 4 + .../xla/backends/gpu/codegen/tensor_ir/BUILD | 57 + .../gpu/codegen/tensor_ir/conversion.cc | 1040 +++++++++++++++++ .../gpu/codegen/tensor_ir/conversion.h | 40 + .../backends/gpu/codegen/tensor_ir/support.cc | 274 +++++ .../backends/gpu/codegen/tensor_ir/support.h | 38 + .../gpu/codegen/tensor_ir/tests/BUILD | 24 + .../gpu/codegen/tensor_ir/tests/dot.hlo | 132 +++ .../tensor_ir/tests/elementwise_binary.hlo | 142 +++ .../tensor_ir/tests/elementwise_compare.hlo | 143 +++ .../tensor_ir/tests/elementwise_convert.hlo | 129 ++ .../tensor_ir/tests/elementwise_ternary.hlo | 28 + .../tensor_ir/tests/elementwise_unary.hlo | 207 ++++ .../tensor_ir/tests/layout_bitcast.hlo | 128 ++ .../tensor_ir/tests/layout_broadcast.hlo | 63 + .../codegen/tensor_ir/tests/layout_other.hlo | 23 + .../gpu/codegen/tensor_ir/tests/other_ops.hlo | 45 + .../tensor_ir/tests/reduce_compare.hlo | 101 ++ .../codegen/tensor_ir/tests/reduce_float.hlo | 188 +++ .../tensor_ir/tests/reduce_int_signed.hlo | 214 ++++ .../tensor_ir/tests/reduce_int_unsigned.hlo | 216 ++++ .../xla/xla/backends/gpu/codegen/tools/BUILD | 29 + .../gpu/codegen/tools/fusion_to_tensorir.cc | 108 ++ 24 files changed, 3375 insertions(+) create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.cc create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.h create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.h create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/dot.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_binary.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_compare.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_convert.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_ternary.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_unary.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_bitcast.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_broadcast.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_other.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/other_ops.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_compare.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_float.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_signed.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_unsigned.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/tools/fusion_to_tensorir.cc diff --git a/third_party/xla/MODULE.bazel b/third_party/xla/MODULE.bazel index 14ed536f9cf367..64b3d29e143192 100644 --- a/third_party/xla/MODULE.bazel +++ b/third_party/xla/MODULE.bazel @@ -108,6 +108,7 @@ use_repo( "KleidiAI", "XNNPACK", "cpuinfo", + "cuda_tile", "cudnn_frontend_archive", "cutlass_archive", "cutlass_cutedsl_runtime", @@ -147,6 +148,7 @@ use_repo( "slinky", "spdlog", "stablehlo", + "tensor_ir", "transformer_engine", "triton", "uv", diff --git a/third_party/xla/third_party/extensions/third_party.bzl b/third_party/xla/third_party/extensions/third_party.bzl index 3735e75bf24dda..be31c8f18eca14 100644 --- a/third_party/xla/third_party/extensions/third_party.bzl +++ b/third_party/xla/third_party/extensions/third_party.bzl @@ -19,6 +19,7 @@ load("//third_party/benchmark:workspace.bzl", benchmark = "repo") load("//third_party/brotli:workspace.bzl", brotli = "repo") load("//third_party/compute_library:workspace.bzl", compute_library = "repo") load("//third_party/cpuinfo:workspace.bzl", cpuinfo = "repo") +load("//third_party/cuda_tile:workspace.bzl", cuda_tile = "repo") load("//third_party/cudnn_frontend:workspace.bzl", cudnn_frontend = "repo") load("//third_party/cutlass:workspace.bzl", cutlass = "repo") load("//third_party/cutlass_cutedsl_runtime:workspace.bzl", cutlass_cutedsl_runtime = "repo") @@ -59,6 +60,7 @@ load("//third_party/slinky:workspace.bzl", slinky = "repo") load("//third_party/spdlog:workspace.bzl", spdlog = "repo") load("//third_party/sqlite:workspace.bzl", sqlite = "repo") load("//third_party/stablehlo:workspace.bzl", stablehlo = "repo") +load("//third_party/tensor_ir:workspace.bzl", tensor_ir = "repo") load("//third_party/tensorrt:workspace.bzl", tensorrt = "repo") load("//third_party/transformer_engine:workspace.bzl", transformer_engine = "repo") load("//third_party/triton:workspace.bzl", triton = "repo") @@ -72,6 +74,7 @@ def _third_party_ext_impl(mctx): # @unused brotli() # Needed for Riegeli compute_library() cpuinfo() + cuda_tile() cudnn_frontend() cutlass() cutlass_cutedsl_runtime() @@ -112,6 +115,7 @@ def _third_party_ext_impl(mctx): # @unused spdlog() sqlite() stablehlo() + tensor_ir() tensorrt() transformer_engine() triton() diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD new file mode 100644 index 00000000000000..b340d8ccf06c4a --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD @@ -0,0 +1,57 @@ +load("@rules_cc//cc:cc_library.bzl", "cc_library") + +package( + # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], + default_visibility = [":friends"], + licenses = ["notice"], +) + +package_group( + name = "friends", + includes = [ + "//xla:friends", + ], +) + +cc_library( + name = "conversion", + srcs = ["conversion.cc"], + hdrs = ["conversion.h"], + deps = [ + "//xla:comparison_util", + "//xla:literal", + "//xla:literal_util", + "//xla:shape_util", + "//xla:util", + "//xla:xla_data_proto_cc", + "//xla/hlo/ir:hlo", + "//xla/service/llvm_ir:llvm_util", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:SideEffectInterfaces", + "@llvm-project//mlir:Support", + "@tensor_ir//:NVTensorIRDialect", + ], +) + +cc_library( + name = "support", + srcs = ["support.cc"], + hdrs = ["support.h"], + deps = [ + "//xla:comparison_util", + "//xla:xla_data_proto_cc", + "//xla/hlo/ir:hlo", + "//xla/service:decision", + "@com_google_absl//absl/strings", + ], +) diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.cc new file mode 100644 index 00000000000000..8dda36ef5d50cf --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.cc @@ -0,0 +1,1040 @@ +/* Copyright 2026 The OpenXLA Authors. + +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 "xla/backends/gpu/codegen/tensor_ir/conversion.h" + +#include +#include +#include +#include + +#include "tensor_ir/Dialect/TensorIR.h" +#include "tensor_ir/Dialect/TensorIRAttrs.h" +#include "absl/algorithm/container.h" +#include "absl/cleanup/cleanup.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" +#include "absl/types/span.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Casting.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributeInterfaces.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/Location.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/IR/Verifier.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Support/LLVM.h" +#include "xla/comparison_util.h" +#include "xla/hlo/ir/hlo_casting_utils.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/layout_util.h" +#include "xla/literal.h" +#include "xla/literal_util.h" +#include "xla/primitive_util.h" +#include "xla/service/llvm_ir/llvm_util.h" +#include "xla/shape.h" +#include "xla/shape_util.h" +#include "xla/util.h" +#include "xla/xla_data.pb.h" + +// NOLINTBEGIN(clang-diagnostic-pre-c++20-compat) + +namespace xla::gpu::tensor_ir { +namespace { + +namespace arith = ::mlir::arith; +namespace tir = ::mlir::nv_tensor_ir; + +absl::StatusOr ConvertReductionInstruction( + const HloInstruction& source, mlir::ValueRange operands, + mlir::Block& target); + +// Creates MLIR location from HLO instruction metadata. +mlir::Location GetLocationFromInstruction(const HloInstruction& source, + mlir::MLIRContext* context) { + const OpMetadata& metadata = source.metadata(); + + llvm::SmallVector locations; + if (!metadata.op_name().empty()) { + locations.push_back( + mlir::NameLoc::get(mlir::StringAttr::get(context, metadata.op_name()))); + } + if (!metadata.source_file().empty()) { + locations.push_back(mlir::FileLineColRange::get( + context, metadata.source_file(), metadata.source_line(), + metadata.source_column(), metadata.source_end_line(), + metadata.source_end_column())); + } + + if (locations.empty()) { + return mlir::UnknownLoc::get(context); + } + if (locations.size() == 1) { + return locations.front(); + } + return mlir::FusedLoc::get(context, locations); +} + +// Creates MLIR type from HLO primitive type. +absl::StatusOr GetElementType(PrimitiveType type, + mlir::Builder& builder) { + switch (type) { + case PrimitiveType::PRED: + return builder.getI1Type(); + case PrimitiveType::S8: + case PrimitiveType::S16: + case PrimitiveType::S32: + case PrimitiveType::S64: + return builder.getIntegerType(primitive_util::BitWidth(type), + /*isSigned=*/true); + case PrimitiveType::U8: + case PrimitiveType::U16: + case PrimitiveType::U32: + case PrimitiveType::U64: + return builder.getIntegerType(primitive_util::BitWidth(type), + /*isSigned=*/false); + case PrimitiveType::F16: + return builder.getF16Type(); + case PrimitiveType::BF16: + return builder.getBF16Type(); + case PrimitiveType::F32: + return builder.getF32Type(); + case PrimitiveType::F64: + return builder.getF64Type(); + default: + return absl::InvalidArgumentError( + absl::StrCat("Unsupported primitive type: ", type)); + } +} + +// Converts a scalar literal to MLIR attribute. +absl::StatusOr GetScalarAttribute( + const LiteralBase& literal, mlir::Builder& builder, + bool use_signless_integer_type) { + ABSL_ASSIGN_OR_RETURN(auto element_type, + GetElementType(literal.shape().element_type(), builder)); + if (element_type.isFloat()) { + if (auto cst = literal.GetAsDouble({})) { + return mlir::FloatAttr::get(element_type, *cst); + } + } else { + if (auto cst = literal.GetIntegralAsS64({})) { + if (use_signless_integer_type) { + element_type = + builder.getIntegerType(element_type.getIntOrFloatBitWidth()); + } + return mlir::IntegerAttr::get(element_type, *cst); + } + } + return absl::InvalidArgumentError( + absl::StrCat("Unsupported scalar literal: ", literal.ToString())); +} + +// Creates MLIR tensor type from HLO shape. +absl::StatusOr GetTensorType(const Shape& shape, + mlir::Builder& builder) { + ABSL_ASSIGN_OR_RETURN(auto element_type, + GetElementType(shape.element_type(), builder)); + // Promote 0-D scalar shapes to 1-D tensors of size 1 for TensorIR. + if (shape.dimensions().empty()) { + return mlir::RankedTensorType::get({1}, element_type); + } + return mlir::RankedTensorType::get(shape.dimensions(), element_type); +} + +// Builds a constant operation with given value and shape. +absl::StatusOr BuildConstant(const LiteralBase& literal, + const Shape& shape, + mlir::ImplicitLocOpBuilder& builder) { + ABSL_ASSIGN_OR_RETURN(auto tensor_type, GetTensorType(shape, builder)); + auto dense_attr = primitive_util::ArrayTypeSwitch( + [&](auto type) -> mlir::DenseElementsAttr { + using NativeT = primitive_util::NativeTypeOf; + return mlir::DenseElementsAttr::get(tensor_type, + literal.GetFirstElement()); + }, + literal.shape().element_type()); + return tir::ConstantOp::create(builder, dense_attr); +} + +absl::StatusOr BuildFloatConstant( + double value, const Shape& shape, mlir::ImplicitLocOpBuilder& builder) { + return BuildConstant(LiteralUtil::CreateR0(shape.element_type(), value), + shape, builder); +} + +// Builds a reshape operation from HLO bitcast instruction. +absl::StatusOr BuildBitcast(const HloInstruction& source, + mlir::Value operand, + mlir::ImplicitLocOpBuilder& builder) { + mlir::Value result = operand; + + // Calculate the operand/result shapes with default layout. + Shape operand_normal_shape = + ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( + source.operand(0)->shape()); + Shape result_normal_shape = + ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( + source.shape()); + + // If the operand doesn't have a default layout, transpose it. + // Example: [A,B,C]{2,0,1} -> [B,A,C]{2,1,0} uses permutation (1,0,2) + if (source.operand(0)->shape() != operand_normal_shape) { + auto permutation = llvm::to_vector( + llvm::reverse(source.operand(0)->shape().layout().minor_to_major())); + ABSL_ASSIGN_OR_RETURN(auto transpose_type, + GetTensorType(operand_normal_shape, builder)); + VLOG(3) << "Transposing operand: " + << llvm_ir::DumpToString(result.getType()) << " to " + << llvm_ir::DumpToString(transpose_type); + result = + tir::TransposeOp::create(builder, transpose_type, result, permutation); + } + + // If the normalized shapes are not equal, create a reshape operation. + if (operand_normal_shape != result_normal_shape) { + ABSL_ASSIGN_OR_RETURN(auto reshape_type, + GetTensorType(result_normal_shape, builder)); + VLOG(3) << "Reshaping operand: " << llvm_ir::DumpToString(result.getType()) + << " to " << llvm_ir::DumpToString(reshape_type); + result = tir::ReshapeOp::create(builder, reshape_type, result); + } + + // If the result doesn't have a default layout, transpose it. + // Example: [A,B,C]{2,1,0} -> [C,A,B]{0,2,1} uses permutation (2,0,1) + if (source.shape() != result_normal_shape) { + llvm::SmallVector permutation(source.shape().dimensions().size()); + for (auto [idx, pos] : llvm::enumerate( + llvm::reverse(source.shape().layout().minor_to_major()))) { + permutation[pos] = idx; + } + ABSL_ASSIGN_OR_RETURN(auto transpose_type, + GetTensorType(source.shape(), builder)); + VLOG(3) << "Transposing operand: " + << llvm_ir::DumpToString(result.getType()) << " to " + << llvm_ir::DumpToString(transpose_type); + result = + tir::TransposeOp::create(builder, transpose_type, result, permutation); + } + + return result; +} + +// Builds a broadcast operation from HLO broadcast instruction. +absl::StatusOr BuildBroadcast( + const HloBroadcastInstruction& source, mlir::Value operand, + mlir::ImplicitLocOpBuilder& builder) { + auto dimensions = source.dimensions(); + + // If the operand is a scalar constant, return shaped constant. + if (const auto* constant = DynCast(source.operand(0)); + constant != nullptr && constant->shape().dimensions().empty()) { + VLOG(3) << "Broadcasting scalar constant: " + << constant->literal().ToString() << " to " << source.shape(); + return BuildConstant(constant->literal(), source.shape(), builder); + } + + // If the source dimensions are not sorted, a transpose is needed. + // https://openxla.org/stablehlo/spec#broadcast_in_dim + if (!absl::c_is_sorted(dimensions)) { + // Get the transpose permutation by sorting the dimensions. + llvm::SmallVector permutation(dimensions.size()); + std::iota(permutation.begin(), permutation.end(), 0); + absl::c_sort(permutation, [&](int64_t a, int64_t b) { + return dimensions[a] < dimensions[b]; + }); + + // Create the transpose operation. + llvm::SmallVector permuted_dims; + permuted_dims.reserve(dimensions.size()); + for (int64_t i : permutation) { + permuted_dims.push_back(source.operand(0)->shape().dimensions(i)); + } + mlir::ShapedType transpose_type = + llvm::cast(operand.getType()).clone(permuted_dims); + VLOG(3) << "Transposing operand: " + << llvm_ir::DumpToString(operand.getType()) << " to " + << llvm_ir::DumpToString(transpose_type); + operand = + tir::TransposeOp::create(builder, transpose_type, operand, permutation); + } + + // Reshape operand, add unit dimensions for the broadcast. + llvm::SmallVector pre_broadcast_dims; + for (int i = 0, n = source.shape().dimensions().size(); i < n; ++i) { + auto it = absl::c_find(dimensions, i); + if (it != dimensions.end()) { + pre_broadcast_dims.push_back( + source.operand(0)->shape().dimensions(it - dimensions.begin())); + } else { + pre_broadcast_dims.push_back(1); + } + } + mlir::ShapedType reshape_type = + llvm::cast(operand.getType()).clone(pre_broadcast_dims); + VLOG(3) << "Reshaping operand: " << llvm_ir::DumpToString(operand.getType()) + << " to " << llvm_ir::DumpToString(reshape_type); + operand = tir::ReshapeOp::create(builder, reshape_type, operand); + + // Create the broadcast operation. + mlir::ShapedType broadcast_type = + reshape_type.clone(source.shape().dimensions()); + VLOG(3) << "Broadcasting operand: " + << llvm_ir::DumpToString(operand.getType()) << " to " + << llvm_ir::DumpToString(broadcast_type); + return tir::BroadcastOp::create(builder, broadcast_type, operand); +} + +// Builds a dot operation from HLO dot instruction. +absl::StatusOr BuildDot(const HloDotInstruction& source, + mlir::Value lhs, mlir::Value rhs, + mlir::ImplicitLocOpBuilder& builder) { + // Calculate LHS transpose permutation. + int lhs_rank = source.operand(0)->shape().dimensions().size(); + const auto& lhs_batch = source.dot_dimension_numbers().lhs_batch_dimensions(); + const auto& lhs_contracting = + source.dot_dimension_numbers().lhs_contracting_dimensions(); + auto lhs_non_contracting = + GetNonContractingDims(lhs_rank, lhs_contracting, lhs_batch); + + llvm::SmallVector lhs_permutation(lhs_batch.begin(), + lhs_batch.end()); + lhs_permutation.append(lhs_non_contracting.begin(), + lhs_non_contracting.end()); + lhs_permutation.append(lhs_contracting.begin(), lhs_contracting.end()); + + // Transpose LHS, if needed. + llvm::SmallVector lhs_permuted_dims; + for (int i = 0; i < lhs_rank; ++i) { + lhs_permuted_dims.push_back( + source.operand(0)->shape().dimensions(lhs_permutation[i])); + } + + if (!absl::c_is_sorted(lhs_permutation)) { + mlir::ShapedType transpose_type = + llvm::cast(lhs.getType()).clone(lhs_permuted_dims); + VLOG(3) << "Transposing LHS operand: " + << llvm_ir::DumpToString(lhs.getType()) << " to " + << llvm_ir::DumpToString(transpose_type); + lhs = + tir::TransposeOp::create(builder, transpose_type, lhs, lhs_permutation); + } + + // Calculate LHS dot operand shape. + int64_t batch_size = + Product(absl::MakeSpan(lhs_permuted_dims).first(lhs_batch.size())); + int64_t contracting_size = + Product(absl::MakeSpan(lhs_permuted_dims).last(lhs_contracting.size())); + int64_t lhs_non_contracting_size = + Product(absl::MakeSpan(lhs_permuted_dims) + .subspan(lhs_batch.size(), lhs_non_contracting.size())); + + llvm::SmallVector lhs_expected_dims; + if (!lhs_batch.empty()) { + lhs_expected_dims.push_back(batch_size); + } + lhs_expected_dims.push_back(lhs_non_contracting_size); + lhs_expected_dims.push_back(contracting_size); + + // Reshape LHS, if needed. + auto lhs_reshape_type = + llvm::cast(lhs.getType()).clone(lhs_expected_dims); + if (lhs.getType() != lhs_reshape_type) { + VLOG(3) << "Reshaping LHS operand: " << llvm_ir::DumpToString(lhs.getType()) + << " to " << llvm_ir::DumpToString(lhs_reshape_type); + lhs = tir::ReshapeOp::create(builder, lhs_reshape_type, lhs); + } + + // Calculate RHS transpose permutation. + int rhs_rank = source.operand(1)->shape().dimensions().size(); + const auto& rhs_batch = source.dot_dimension_numbers().rhs_batch_dimensions(); + const auto& rhs_contracting = + source.dot_dimension_numbers().rhs_contracting_dimensions(); + auto rhs_non_contracting = + GetNonContractingDims(rhs_rank, rhs_contracting, rhs_batch); + + llvm::SmallVector rhs_permutation(rhs_batch.begin(), + rhs_batch.end()); + rhs_permutation.append(rhs_contracting.begin(), rhs_contracting.end()); + rhs_permutation.append(rhs_non_contracting.begin(), + rhs_non_contracting.end()); + + // Transpose RHS, if needed. + llvm::SmallVector rhs_permuted_dims; + for (int i = 0; i < rhs_rank; ++i) { + rhs_permuted_dims.push_back( + source.operand(1)->shape().dimensions(rhs_permutation[i])); + } + + if (!absl::c_is_sorted(rhs_permutation)) { + mlir::ShapedType transpose_type = + llvm::cast(rhs.getType()).clone(rhs_permuted_dims); + VLOG(3) << "Transposing RHS operand: " + << llvm_ir::DumpToString(rhs.getType()) << " to " + << llvm_ir::DumpToString(transpose_type); + rhs = + tir::TransposeOp::create(builder, transpose_type, rhs, rhs_permutation); + } + + // Calculate RHS dot operand shape. + int64_t rhs_non_contracting_size = Product( + absl::MakeSpan(rhs_permuted_dims).last(rhs_non_contracting.size())); + + llvm::SmallVector rhs_expected_dims; + if (!rhs_batch.empty()) { + rhs_expected_dims.push_back(batch_size); + } + rhs_expected_dims.push_back(contracting_size); + rhs_expected_dims.push_back(rhs_non_contracting_size); + + // Reshape RHS, if needed. + auto rhs_reshape_type = + llvm::cast(rhs.getType()).clone(rhs_expected_dims); + if (rhs.getType() != rhs_reshape_type) { + VLOG(3) << "Reshaping RHS operand: " << llvm_ir::DumpToString(rhs.getType()) + << " to " << llvm_ir::DumpToString(rhs_reshape_type); + rhs = tir::ReshapeOp::create(builder, rhs_reshape_type, rhs); + } + + // Calculate matmul result shape. + ABSL_ASSIGN_OR_RETURN(auto result_type, GetTensorType(source.shape(), builder)); + + llvm::SmallVector matmul_shape; + if (!lhs_batch.empty()) { + matmul_shape.push_back(batch_size); + } + matmul_shape.push_back(lhs_non_contracting_size); + matmul_shape.push_back(rhs_non_contracting_size); + mlir::ShapedType matmul_type = result_type.clone(matmul_shape); + + // Create the matmul operation. + VLOG(3) << "Creating matmul operation: " << source.name(); + auto matmul_op = tir::MatmulOp::create(builder, matmul_type, lhs, rhs); + mlir::Value result = matmul_op.getResult(); + + // Align the result to the source shape. + if (result_type != matmul_type) { + VLOG(3) << "Reshaping result: " << llvm_ir::DumpToString(matmul_type) + << " to " << llvm_ir::DumpToString(result_type); + result = tir::ReshapeOp::create(builder, result_type, result); + } + return result; +} + +// Returns the shaped type for reduction output with unit dimensions for reduced +// axes. +mlir::ShapedType GetReduceOutputType(mlir::RankedTensorType result_type, + const Shape& operand_shape, + const Shape& result_shape, + absl::Span reduce_dims) { + llvm::SmallVector output_shape; + output_shape.reserve(operand_shape.dimensions().size()); + for (int i = 0, p = 0; i < operand_shape.dimensions().size(); ++i) { + if (!absl::c_contains(reduce_dims, i)) { + output_shape.push_back(result_shape.dimensions(p++)); + } else { + output_shape.push_back(1); + } + } + return result_type.clone(output_shape); +} + +// Builds a reduction operation from HLO reduce instruction. +absl::StatusOr BuildReduce(const HloReduceInstruction& source, + mlir::ValueRange operands, + mlir::ImplicitLocOpBuilder& builder) { + llvm::SmallVector initial_values; + for (const HloInstruction* init_instr : source.init_values()) { + const auto* constant = Cast(init_instr); + ABSL_ASSIGN_OR_RETURN(auto initial_value, + GetScalarAttribute(constant->literal(), builder, + /*use_signless_integer_type=*/false)); + initial_values.push_back(initial_value); + } + + llvm::SmallVector reduce_dims(source.dimensions().begin(), + source.dimensions().end()); + ABSL_ASSIGN_OR_RETURN(auto result_type, GetTensorType(source.shape(), builder)); + mlir::ShapedType output_type = GetReduceOutputType( + result_type, source.operand(0)->shape(), source.shape(), reduce_dims); + + VLOG(3) << "Creating reduction operation: " << source.name(); + auto reduce_op = tir::ReduceUDOp::create( + builder, output_type, operands.take_front(source.input_count()), + reduce_dims, builder.getArrayAttr(initial_values)); + + mlir::Block* body = builder.createBlock(&reduce_op.getRegion()); + for (int i = 0; i < 2; ++i) { + for (mlir::Attribute init : initial_values) { + auto arg_type = llvm::cast(init).getType(); + if (arg_type.isInteger()) { + arg_type = builder.getIntegerType(arg_type.getIntOrFloatBitWidth()); + } + body->addArgument(arg_type, reduce_op.getLoc()); + } + } + + const HloComputation* computation = source.to_apply(); + VLOG(3) << "Converting reduction HLO computation: " << computation->name(); + + llvm::DenseMap converted; + for (const HloInstruction* instruction : + computation->MakeInstructionPostOrder()) { + llvm::SmallVector new_operands; + for (const auto& operand : instruction->operands()) { + new_operands.push_back(converted[operand]); + } + ABSL_ASSIGN_OR_RETURN( + converted[instruction], + ConvertReductionInstruction(*instruction, new_operands, *body)); + } + + mlir::Value result = converted[computation->root_instruction()]; + tir::YieldOp::create(builder, result); + builder.setInsertionPointAfter(reduce_op); + + VLOG(3) << "Reshaping result: " << llvm_ir::DumpToString(output_type) + << " to " << llvm_ir::DumpToString(result_type); + return tir::ReshapeOp::create(builder, result_type, reduce_op.getResult(0)); +} + +// Returns the reduction type, if the body is a supported single-op reduction. +std::optional GetReductionType( + const HloReduceInstruction& reduce) { + // Reduction must have exactly two parameters. + const HloComputation* comp = reduce.to_apply(); + if (comp->num_parameters() != 2) { + return std::nullopt; + } + const HloInstruction* param0 = comp->parameter_instruction(0); + const HloInstruction* param1 = comp->parameter_instruction(1); + + // Reduction root must use the parameters as operands. + const HloInstruction* root = comp->root_instruction(); + if (root->operand_count() != 2 || root->operand(0) != param0 || + root->operand(1) != param1) { + return std::nullopt; + } + + // Reduction initial value must be a constant. + const auto* init_value = + DynCast(reduce.init_values().front()); + if (init_value == nullptr) { + return std::nullopt; + } + + if (root->opcode() == HloOpcode::kAdd && init_value->literal().IsAll(0)) { + return tir::ReductionMode::add; + } + + if (root->opcode() == HloOpcode::kMultiply && + init_value->literal().IsAll(1)) { + return tir::ReductionMode::mul; + } + + return std::nullopt; +} + +// Builds a reduction operation for a reduction type. +absl::StatusOr BuildSimpleReduce( + const HloReduceInstruction& source, tir::ReductionMode mode, + mlir::Value operand, mlir::ImplicitLocOpBuilder& builder) { + llvm::SmallVector reduce_dims(source.dimensions().begin(), + source.dimensions().end()); + ABSL_ASSIGN_OR_RETURN(auto result_type, GetTensorType(source.shape(), builder)); + mlir::ShapedType output_type = GetReduceOutputType( + result_type, source.operand(0)->shape(), source.shape(), reduce_dims); + + VLOG(3) << "Creating reduction operation: " << source.name(); + auto reduce_op = + tir::ReduceOp::create(builder, output_type, operand, reduce_dims, mode); + + VLOG(3) << "Reshaping result: " << llvm_ir::DumpToString(output_type) + << " to " << llvm_ir::DumpToString(result_type); + return tir::ReshapeOp::create(builder, result_type, reduce_op.getResult()); +} + +// Maps HLO compare instruction to TensorIR comparator. +tir::Comparator GetComparator(const HloCompareInstruction& source) { + bool is_float = primitive_util::IsFloatingPointType( + source.operand(0)->shape().element_type()); + if (is_float) { + switch (source.direction()) { + case ComparisonDirection::kEq: + return tir::Comparator::oeq; + case ComparisonDirection::kNe: + return tir::Comparator::une; + case ComparisonDirection::kGt: + return tir::Comparator::ogt; + case ComparisonDirection::kGe: + return tir::Comparator::oge; + case ComparisonDirection::kLt: + return tir::Comparator::olt; + case ComparisonDirection::kLe: + return tir::Comparator::ole; + } + } else { + switch (source.direction()) { + case ComparisonDirection::kEq: + return tir::Comparator::eq; + case ComparisonDirection::kNe: + return tir::Comparator::neq; + case ComparisonDirection::kGt: + return tir::Comparator::gt; + case ComparisonDirection::kGe: + return tir::Comparator::ge; + case ComparisonDirection::kLt: + return tir::Comparator::lt; + case ComparisonDirection::kLe: + return tir::Comparator::le; + } + } +} + +// Maps HLO compare instruction to `arith` comparator (floating-point). +arith::CmpFPredicate GetArithFloatComparator( + const HloCompareInstruction& source) { + switch (source.direction()) { + case ComparisonDirection::kEq: + return arith::CmpFPredicate::OEQ; + case ComparisonDirection::kNe: + return arith::CmpFPredicate::UNE; + case ComparisonDirection::kGt: + return arith::CmpFPredicate::OGT; + case ComparisonDirection::kGe: + return arith::CmpFPredicate::OGE; + case ComparisonDirection::kLt: + return arith::CmpFPredicate::OLT; + case ComparisonDirection::kLe: + return arith::CmpFPredicate::OLE; + } +} + +// Maps HLO compare instruction to `arith` comparator (integer). +arith::CmpIPredicate GetArithIntComparator( + const HloCompareInstruction& source) { + bool is_signed = primitive_util::IsSignedIntegralType( + source.operand(0)->shape().element_type()); + switch (source.direction()) { + case ComparisonDirection::kEq: + return arith::CmpIPredicate::eq; + case ComparisonDirection::kNe: + return arith::CmpIPredicate::ne; + case ComparisonDirection::kGt: + return is_signed ? arith::CmpIPredicate::sgt : arith::CmpIPredicate::ugt; + case ComparisonDirection::kGe: + return is_signed ? arith::CmpIPredicate::sge : arith::CmpIPredicate::uge; + case ComparisonDirection::kLt: + return is_signed ? arith::CmpIPredicate::slt : arith::CmpIPredicate::ult; + case ComparisonDirection::kLe: + return is_signed ? arith::CmpIPredicate::sle : arith::CmpIPredicate::ule; + } +} + +// Template for creating `arith` operations. +template +mlir::Value CreateArithOp(Args&&... args) { + return Op::create(std::forward(args)...).getResult(); +} + +// Creates strides attribute (helper). +mlir::NamedAttribute CreateStridesAttributeFromShape(mlir::MLIRContext* context, + const Shape& shape) { + // Calculate strides from the shape and the layout. + llvm::SmallVector strides(shape.dimensions().size()); + int64_t stride = 1; + for (int dim : shape.layout().minor_to_major()) { + strides[dim] = stride; + stride *= shape.dimensions(dim); + } + + // Create the attribute with strides formatted as a string. + auto strides_fmt = absl::StrCat("(", absl::StrJoin(strides, ","), ")"); + auto strides_attr = mlir::StringAttr::get(context, strides_fmt); + return {tir::TensorIRDialect::getStrideAttrName(), strides_attr}; +} + +mlir::DictionaryAttr CreateStridesDictionaryAttribute( + mlir::MLIRContext* context, const Shape& shape) { + if (LayoutUtil::IsMonotonicWithDim0Major(shape.layout())) { + return mlir::DictionaryAttr::get(context); + } + auto attr = CreateStridesAttributeFromShape(context, shape); + return mlir::DictionaryAttr::get(context, {attr}); +} + +// Converts HLO instruction `source` with already-converted `operands` values +// to TensorIR ops and insert them at the end of `target` block. +absl::StatusOr ConvertFusionInstruction( + const HloInstruction& source, mlir::ValueRange operands, + mlir::Block& target) { + mlir::MLIRContext* context = target.getParent()->getContext(); + VLOG(3) << "Converting HLO instruction: " << source.ToString(); + + mlir::Location location = GetLocationFromInstruction(source, context); + mlir::ImplicitLocOpBuilder builder(location, context); + builder.setInsertionPointToEnd(&target); + + switch (source.opcode()) { + // Unary elementwise operations. + case HloOpcode::kAbs: + return tir::AbsOp::create(builder, operands); + case HloOpcode::kCeil: + return tir::CeilOp::create(builder, operands); + case HloOpcode::kConvert: { + ABSL_ASSIGN_OR_RETURN(auto convert_type, + GetTensorType(source.shape(), builder)); + return tir::ConvertOp::create(builder, convert_type, operands[0]); + } + case HloOpcode::kCos: + return tir::CosOp::create(builder, operands); + case HloOpcode::kErf: + return tir::ErfOp::create(builder, operands); + case HloOpcode::kExp: + return tir::ExpOp::create(builder, operands); + case HloOpcode::kExpm1: { + // NOTE: Decomposing Expm1(x) to Exp(x) - 1.0 can result in a severe loss + // of precision for values of x close to 0. + ABSL_ASSIGN_OR_RETURN(auto one, + BuildFloatConstant(1.0, source.shape(), builder)); + return tir::SubOp::create(builder, + tir::ExpOp::create(builder, operands[0]), one); + } + case HloOpcode::kFloor: + return tir::FloorOp::create(builder, operands); + case HloOpcode::kLog: + return tir::LogOp::create(builder, operands); + case HloOpcode::kLog1p: { + // NOTE: Decomposing Log1p(x) to Log(x + 1.0) can result in a severe loss + // of precision for values of x close to 0. + ABSL_ASSIGN_OR_RETURN(auto one, + BuildFloatConstant(1.0, source.shape(), builder)); + return tir::LogOp::create(builder, + tir::AddOp::create(builder, operands[0], one)); + } + case HloOpcode::kNot: + return tir::LogicalNotOp::create(builder, operands); + case HloOpcode::kNegate: + return tir::NegOp::create(builder, operands); + case HloOpcode::kRsqrt: + return tir::RsqrtOp::create(builder, operands); + case HloOpcode::kSin: + return tir::SinOp::create(builder, operands); + case HloOpcode::kSqrt: + return tir::SqrtOp::create(builder, operands); + case HloOpcode::kTan: + return tir::TanOp::create(builder, operands); + case HloOpcode::kTanh: + return tir::TanhFwdOp::create(builder, operands); + + // Binary elementwise operations. + case HloOpcode::kAdd: + return tir::AddOp::create(builder, operands); + case HloOpcode::kAtan2: + return tir::Atan2Op::create(builder, operands); + case HloOpcode::kCompare: { + tir::Comparator comparator = + GetComparator(*Cast(&source)); + return tir::CmpOp::create(builder, comparator, operands[0], operands[1]); + } + case HloOpcode::kDivide: + return tir::DivOp::create(builder, operands); + case HloOpcode::kMaximum: + return tir::MaxOp::create(builder, operands); + case HloOpcode::kMinimum: + return tir::MinOp::create(builder, operands); + case HloOpcode::kMultiply: + return tir::MulOp::create(builder, operands); + case HloOpcode::kPower: + return tir::PowOp::create(builder, operands); + case HloOpcode::kRemainder: + return tir::RemOp::create(builder, operands); + case HloOpcode::kSubtract: + return tir::SubOp::create(builder, operands); + case HloOpcode::kAnd: + return tir::LogicalAndOp::create(builder, operands); + case HloOpcode::kOr: + return tir::LogicalOrOp::create(builder, operands); + + // Ternary elementwise operations. + case HloOpcode::kSelect: + return tir::BinarySelectOp::create(builder, operands); + case HloOpcode::kClamp: + return tir::MinOp::create( + builder, operands[2], + tir::MaxOp::create(builder, operands[0], operands[1])); + + // Layout modification operations. + case HloOpcode::kReshape: + case HloOpcode::kBitcast: + return BuildBitcast(source, operands[0], builder); + case HloOpcode::kBroadcast: + return BuildBroadcast(*Cast(&source), + operands[0], builder); + case HloOpcode::kSlice: { + ABSL_ASSIGN_OR_RETURN(auto slice_type, GetTensorType(source.shape(), builder)); + return tir::SliceOp::create(builder, slice_type, operands[0], + source.slice_starts(), source.slice_limits(), + source.slice_strides()); + } + case HloOpcode::kTranspose: { + auto transpose = Cast(&source); + ABSL_ASSIGN_OR_RETURN(auto transpose_type, + GetTensorType(source.shape(), builder)); + return tir::TransposeOp::create(builder, transpose_type, operands[0], + transpose->dimensions()); + } + + // Reduction operations. + case HloOpcode::kDot: + return BuildDot(*Cast(&source), operands[0], + operands[1], builder); + case HloOpcode::kReduce: { + auto reduce = Cast(&source); + auto type = GetReductionType(*reduce); + return type.has_value() + ? BuildSimpleReduce(*reduce, *type, operands[0], builder) + : BuildReduce(*reduce, operands, builder); + } + + // Miscellaneous operations. + case HloOpcode::kParameter: + return target.getArgument(source.parameter_number()); + case HloOpcode::kConstant: + return BuildConstant(source.literal(), source.shape(), builder); + case HloOpcode::kIota: { + auto iota = Cast(&source); + ABSL_ASSIGN_OR_RETURN(auto iota_type, GetTensorType(source.shape(), builder)); + return tir::IotaOp::create(builder, iota_type, iota->iota_dimension(), + /*dynamic_sizes=*/{}); + } + case HloOpcode::kConcatenate: { + ABSL_ASSIGN_OR_RETURN(auto concatenate_type, + GetTensorType(source.shape(), builder)); + return tir::ConcatenateOp::create(builder, concatenate_type, operands, + source.concatenate_dimension()); + } + + default: + return absl::UnimplementedError(absl::StrCat( + "Unsupported instruction: ", HloOpcodeString(source.opcode()))); + } +} + +// Converts HLO instruction `source` in the context of a reduction computation +// with already-converted `operands` values to TensorIR ops and insert them at +// the end of `target` block. +absl::StatusOr ConvertReductionInstruction( + const HloInstruction& source, mlir::ValueRange operands, + mlir::Block& target) { + mlir::MLIRContext* context = target.getParent()->getContext(); + VLOG(3) << "Converting HLO instruction: " << source.ToString(); + + mlir::Location location = GetLocationFromInstruction(source, context); + mlir::ImplicitLocOpBuilder builder(location, context); + builder.setInsertionPointToEnd(&target); + + // `arith` dialect (used in reduction body) has separate instructions for + // floating-point and signed/unsigned integer types. + const HloInstruction* type_source = + source.operand_count() != 0 ? source.operand(0) : &source; + PrimitiveType element_type = type_source->shape().element_type(); + bool is_float = primitive_util::IsFloatingPointType(element_type); + bool is_signed = primitive_util::IsSignedIntegralType(element_type); + + switch (source.opcode()) { + // Binary elementwise operations. + case HloOpcode::kAdd: + return is_float ? CreateArithOp(builder, operands) + : CreateArithOp(builder, operands); + case HloOpcode::kCompare: { + const auto& compare = *Cast(&source); + return is_float ? CreateArithOp( + builder, GetArithFloatComparator(compare), + operands[0], operands[1]) + : CreateArithOp( + builder, GetArithIntComparator(compare), + operands[0], operands[1]); + } + case HloOpcode::kMaximum: + return is_float ? CreateArithOp(builder, operands) + : is_signed ? CreateArithOp(builder, operands) + : CreateArithOp(builder, operands); + case HloOpcode::kMinimum: + return is_float ? CreateArithOp(builder, operands) + : is_signed ? CreateArithOp(builder, operands) + : CreateArithOp(builder, operands); + case HloOpcode::kMultiply: + return is_float ? CreateArithOp(builder, operands) + : CreateArithOp(builder, operands); + case HloOpcode::kAnd: + return CreateArithOp(builder, operands); + case HloOpcode::kOr: + return CreateArithOp(builder, operands); + case HloOpcode::kXor: + return CreateArithOp(builder, operands); + + // Ternary elementwise operations. + case HloOpcode::kSelect: + return CreateArithOp(builder, operands); + case HloOpcode::kClamp: + return is_float ? CreateArithOp( + builder, operands[2], + CreateArithOp( + builder, operands[0], operands[1])) + : is_signed ? CreateArithOp( + builder, operands[2], + CreateArithOp( + builder, operands[0], operands[1])) + : CreateArithOp( + builder, operands[2], + CreateArithOp( + builder, operands[0], operands[1])); + + // Miscellaneous operations. + case HloOpcode::kParameter: + return target.getArgument(source.parameter_number()); + case HloOpcode::kConstant: { + const auto& constant = *Cast(&source); + ABSL_ASSIGN_OR_RETURN(auto value, + GetScalarAttribute(constant.literal(), builder, + /*use_signless_integer_type=*/true)); + return arith::ConstantOp::create(builder, value); + } + + default: + return absl::UnimplementedError(absl::StrCat( + "Unsupported instruction: ", HloOpcodeString(source.opcode()))); + } +} + +} // namespace + +// Converts an HLO fusion computation into a TensorIR `GraphOp` appended +// to the body of `target`. If conversion fails, `target` is unmodified. +absl::StatusOr ConvertFusionComputation( + const HloComputation& source, mlir::ModuleOp target) { + mlir::MLIRContext* context = target.getContext(); + VLOG(3) << "Converting HLO computation: " << source.name(); + + mlir::ImplicitLocOpBuilder builder(mlir::UnknownLoc::get(context), context); + builder.setInsertionPointToEnd(target.getBody()); + + llvm::SmallVector input_types; + llvm::SmallVector input_attrs; + for (const auto& parameter : source.parameter_instructions()) { + ABSL_ASSIGN_OR_RETURN(auto parameter_type, + GetTensorType(parameter->shape(), builder)); + input_types.push_back(parameter_type); + input_attrs.push_back( + CreateStridesDictionaryAttribute(context, parameter->shape())); + } + + mlir::ArrayAttr arg_attrs; + if (llvm::any_of(input_attrs, [](mlir::Attribute attr) { + return !llvm::cast(attr).empty(); + })) { + arg_attrs = mlir::ArrayAttr::get(context, input_attrs); + } + + llvm::SmallVector output_types; + ABSL_ASSIGN_OR_RETURN(auto output_type, + GetTensorType(source.root_instruction()->shape(), builder)); + output_types.push_back(output_type); + + mlir::ArrayAttr res_attrs; + auto output_attr = CreateStridesDictionaryAttribute( + context, source.root_instruction()->shape()); + if (!output_attr.empty()) { + res_attrs = mlir::ArrayAttr::get(context, {output_attr}); + } + + auto function_type = + mlir::FunctionType::get(context, input_types, output_types); + VLOG(3) << "Function type: " << llvm_ir::DumpToString(function_type); + + auto graph_op = tir::GraphOp::create(builder, source.name(), function_type, + arg_attrs, res_attrs); + + absl::Cleanup cleanup = [&] { graph_op.erase(); }; + + mlir::Block* body = builder.createBlock(&graph_op.getRegion()); + for (mlir::Type type : input_types) { + body->addArgument(type, target.getLoc()); + } + + llvm::DenseMap converted; + for (const HloInstruction* instruction : source.MakeInstructionPostOrder()) { + llvm::SmallVector operands; + operands.reserve(instruction->operand_count()); + for (const auto& operand : instruction->operands()) { + operands.push_back(converted[operand]); + } + ABSL_ASSIGN_OR_RETURN( + converted[instruction], + ConvertFusionInstruction(*instruction, operands, *graph_op.getBody())); + } + + mlir::Value result = converted[source.root_instruction()]; + tir::ResultsOp::create(builder, result); + + // Remove dead operations (scalar constants). + for (mlir::Operation& op : llvm::make_early_inc_range(graph_op.getOps())) { + if (mlir::isOpTriviallyDead(&op)) { + op.erase(); + } + } + + if (mlir::failed(mlir::verify(graph_op))) { + return absl::InternalError(absl::StrCat( + "Verification failed for TensorIR graph: ", source.name())); + } + + std::move(cleanup).Cancel(); + return graph_op; +} + +// Creates a new MLIR module containing the converted TensorIR graph for +// `source`. +absl::StatusOr> ConvertFusionComputation( + const HloComputation& source, mlir::MLIRContext* context) { + mlir::Location location = mlir::UnknownLoc::get(context); + mlir::OwningOpRef module = + llvm_ir::CreateMlirModuleOp(location); + ABSL_RETURN_IF_ERROR(ConvertFusionComputation(source, *module).status()); + return module; +} + +} // namespace xla::gpu::tensor_ir + +// NOLINTEND(clang-diagnostic-pre-c++20-compat) diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.h new file mode 100644 index 00000000000000..eb0ffc1f2f4e2d --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.h @@ -0,0 +1,40 @@ +/* Copyright 2026 The OpenXLA Authors. + +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 XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_CONVERSION_H_ +#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_CONVERSION_H_ + +#include "tensor_ir/Dialect/TensorIR.h" +#include "absl/status/statusor.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" +#include "xla/hlo/ir/hlo_computation.h" + +namespace xla::gpu::tensor_ir { + +// Converts an HLO fusion computation into a TensorIR `GraphOp` appended +// to the body of `target`. If conversion fails, `target` is unmodified. +absl::StatusOr ConvertFusionComputation( + const HloComputation& source, mlir::ModuleOp target); + +// Creates a new MLIR module containing the converted TensorIR graph for +// `source`. +absl::StatusOr> ConvertFusionComputation( + const HloComputation& source, mlir::MLIRContext* context); + +} // namespace xla::gpu::tensor_ir + +#endif // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_CONVERSION_H_ diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc new file mode 100644 index 00000000000000..9503fafab09fc6 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc @@ -0,0 +1,274 @@ +/* Copyright 2026 The OpenXLA Authors. + +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 "xla/backends/gpu/codegen/tensor_ir/support.h" + +#include "absl/strings/str_cat.h" +#include "xla/comparison_util.h" +#include "xla/hlo/ir/hlo_casting_utils.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/primitive_util.h" +#include "xla/xla_data.pb.h" + +namespace xla::gpu::tensor_ir { +namespace { + +bool IsSupportedPrimitiveType(PrimitiveType type) { + switch (type) { + case PrimitiveType::PRED: + case PrimitiveType::S8: + case PrimitiveType::S16: + case PrimitiveType::S32: + case PrimitiveType::S64: + case PrimitiveType::U8: + case PrimitiveType::U16: + case PrimitiveType::U32: + case PrimitiveType::U64: + case PrimitiveType::F16: + case PrimitiveType::BF16: + case PrimitiveType::F32: + case PrimitiveType::F64: { + return true; + } + + default: { + return false; + } + } +} + +bool IsSupportedFusionOpcode(HloOpcode opcode) { + switch (opcode) { + // Unary elementwise operations. + case HloOpcode::kAbs: + case HloOpcode::kCeil: + case HloOpcode::kConvert: + case HloOpcode::kCos: + case HloOpcode::kErf: + case HloOpcode::kExp: + case HloOpcode::kExpm1: + case HloOpcode::kFloor: + case HloOpcode::kLog: + case HloOpcode::kLog1p: + case HloOpcode::kNot: + case HloOpcode::kNegate: + case HloOpcode::kRsqrt: + case HloOpcode::kSin: + case HloOpcode::kSqrt: + case HloOpcode::kTan: + case HloOpcode::kTanh: + + // Binary elementwise operations. + case HloOpcode::kAdd: + case HloOpcode::kAtan2: + case HloOpcode::kCompare: + case HloOpcode::kDivide: + case HloOpcode::kMaximum: + case HloOpcode::kMinimum: + case HloOpcode::kMultiply: + case HloOpcode::kPower: + case HloOpcode::kRemainder: + case HloOpcode::kSubtract: + case HloOpcode::kAnd: + case HloOpcode::kOr: + + // Ternary elementwise operations. + case HloOpcode::kSelect: + case HloOpcode::kClamp: + + // Layout modification operations. + case HloOpcode::kReshape: + case HloOpcode::kBitcast: + case HloOpcode::kBroadcast: + case HloOpcode::kSlice: + case HloOpcode::kTranspose: + + // Reduction operations. + case HloOpcode::kDot: + case HloOpcode::kReduce: + + // Miscellaneous operations. + case HloOpcode::kParameter: + case HloOpcode::kConstant: + case HloOpcode::kIota: + case HloOpcode::kConcatenate: { + return true; + } + + default: { + return false; + } + } +} + +bool IsSupportedReductionOpcode(HloOpcode opcode) { + switch (opcode) { + // Binary elementwise operations. + case HloOpcode::kAdd: + case HloOpcode::kCompare: + case HloOpcode::kMaximum: + case HloOpcode::kMinimum: + case HloOpcode::kMultiply: + case HloOpcode::kAnd: + case HloOpcode::kOr: + case HloOpcode::kXor: + + // Ternary elementwise operations. + case HloOpcode::kSelect: + case HloOpcode::kClamp: + + // Miscellaneous operations. + case HloOpcode::kParameter: + case HloOpcode::kConstant: { + return true; + } + + default: { + return false; + } + } +} + +CodegenDecision CheckCompareInstruction(const HloInstruction& instr) { + auto compare = Cast(&instr); + if (compare->order() == ComparisonOrder::kTotal && + primitive_util::IsFloatingPointType( + compare->operand(0)->shape().element_type())) { + return CodegenDecision::Forbid( + "Total order comparison is not supported for floating-point types"); + } + return CodegenDecision::Allow(); +} + +CodegenDecision IsSupportedReductionInstruction(const HloInstruction& instr) { + if (!instr.shape().IsArray()) { + return CodegenDecision::Forbid(absl::StrCat("Unsupported non-array shape: ", + instr.shape().ToString())); + } + if (!IsSupportedPrimitiveType(instr.shape().element_type())) { + return CodegenDecision::Forbid( + absl::StrCat("Unsupported element type: ", + primitive_util::LowercasePrimitiveTypeName( + instr.shape().element_type()))); + } + if (!IsSupportedReductionOpcode(instr.opcode())) { + return CodegenDecision::Forbid( + absl::StrCat("Unsupported reduction instruction: ", + HloOpcodeString(instr.opcode()))); + } + if (instr.opcode() == HloOpcode::kCompare) { + auto decision = CheckCompareInstruction(instr); + if (!decision.IsAllowed()) { + return decision; + } + } + return CodegenDecision::Allow(); +} + +} // namespace + +CodegenDecision IsSupportedFusionComputation(const HloComputation& comp) { + for (const HloInstruction* instruction : comp.instructions()) { + auto decision = IsInstructionSupportedForFusion(*instruction); + if (!decision.IsAllowed()) { + return decision; + } + } + return CodegenDecision::Allow(); +} + +CodegenDecision IsInstructionSupportedForFusion(const HloInstruction& instr) { + if (instr.opcode() == HloOpcode::kFusion) { + return IsSupportedFusionComputation( + *instr.fused_instructions_computation()); + } + + if (!instr.shape().IsArray()) { + return CodegenDecision::Forbid(absl::StrCat("Unsupported non-array shape: ", + instr.shape().ToString())); + } + + if (!IsSupportedPrimitiveType(instr.shape().element_type())) { + return CodegenDecision::Forbid( + absl::StrCat("Unsupported element type: ", + primitive_util::LowercasePrimitiveTypeName( + instr.shape().element_type()))); + } + if (!IsSupportedFusionOpcode(instr.opcode())) { + return CodegenDecision::Forbid(absl::StrCat( + "Unsupported instruction: ", HloOpcodeString(instr.opcode()))); + } + + switch (instr.opcode()) { + case HloOpcode::kReduce: { + auto reduce = Cast(&instr); + if (reduce->input_count() != 1) { + return CodegenDecision::Forbid( + absl::StrCat("Unsupported variadic reduction: ", instr.name())); + } + for (const HloInstruction* init_value : reduce->init_values()) { + if (DynCast(init_value) == nullptr) { + return CodegenDecision::Forbid(absl::StrCat( + "Unsupported reduction initial value: ", init_value->name())); + } + } + for (const HloInstruction* inner : instr.to_apply()->instructions()) { + auto decision = IsSupportedReductionInstruction(*inner); + if (!decision.IsAllowed()) { + return decision; + } + } + break; + } + + case HloOpcode::kConstant: { + // Only scalar constants are supported. + if (!instr.shape().dimensions().empty()) { + return CodegenDecision::Forbid( + absl::StrCat("Unsupported non-scalar constant: ", instr.name())); + } + break; + } + + case HloOpcode::kAnd: + case HloOpcode::kOr: + case HloOpcode::kNot: { + // Verify that logical operations are used on boolean types. + if (instr.shape().element_type() != PrimitiveType::PRED) { + return CodegenDecision::Forbid(absl::StrCat( + "Unsupported type for logical operation: ", instr.name())); + } + break; + } + + case HloOpcode::kCompare: { + auto decision = CheckCompareInstruction(instr); + if (!decision.IsAllowed()) { + return decision; + } + break; + } + + default: { + break; + } + } + return CodegenDecision::Allow(); +} + +} // namespace xla::gpu::tensor_ir diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.h new file mode 100644 index 00000000000000..88f3b25dc9245a --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.h @@ -0,0 +1,38 @@ +/* Copyright 2026 The OpenXLA Authors. + +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 XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_SUPPORT_H_ +#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_SUPPORT_H_ + +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/service/decision.h" + +namespace xla::gpu::tensor_ir { + +using CodegenDecision = Decision; + +// Returns `Decision::Allow` if the fusion computation is supported by the +// TensorIR fusion emitter. +CodegenDecision IsSupportedFusionComputation(const HloComputation& comp); + +// Returns `Decision::Allow` if the given instruction is supported by the +// TensorIR fusion emitter. If `instr` is a fusion instruction, verifies the +// fused computation. +CodegenDecision IsInstructionSupportedForFusion(const HloInstruction& instr); + +} // namespace xla::gpu::tensor_ir + +#endif // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_SUPPORT_H_ diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD new file mode 100644 index 00000000000000..01bef89604ce0c --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD @@ -0,0 +1,24 @@ +load("//xla:lit.bzl", "lit_test_suite") + +package( + # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], + licenses = ["notice"], +) + +lit_test_suite( + name = "lit_tests", + srcs = glob(["**/*.hlo"]), + cfg = "//xla:lit.cfg.py", + default_tags = [ + "cuda-only", + "gpu", + ], + tags = [ + "cuda-only", + "gpu", + ], + tools = [ + "//xla/backends/gpu/codegen/tools:fusion_to_tensorir", + "@llvm-project//llvm:FileCheck", + ], +) diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/dot.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/dot.hlo new file mode 100644 index 00000000000000..96997353b6fd04 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/dot.hlo @@ -0,0 +1,132 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_dot_no_batch +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x32xf32>, %[[ARG1:.*]]: tensor<32x16xf32> +// CHECK: %[[OUT:.*]] = matmul(%[[ARG0]], %[[ARG1]]) : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32> +// CHECK: results %[[OUT]] + +test_dot_no_batch { + %p0 = f32[8,32] parameter(0) + %p1 = f32[32,16] parameter(1) + %out = f32[8,16] dot(%p0, %p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} +} + +// ----- + +// CHECK-LABEL: @test_dot_with_batch +// CHECK-SAME: %[[ARG0:.*]]: tensor<4x8x32xf32>, %[[ARG1:.*]]: tensor<4x32x16xf32> +// CHECK: %[[OUT:.*]] = matmul(%[[ARG0]], %[[ARG1]]) : (tensor<4x8x32xf32>, tensor<4x32x16xf32>) -> tensor<4x8x16xf32> +// CHECK: results %[[OUT]] + +test_dot_with_batch { + %p0 = f32[4,8,32] parameter(0) + %p1 = f32[4,32,16] parameter(1) + %out = f32[4,8,16] dot(%p0, %p1), lhs_contracting_dims={2}, rhs_contracting_dims={1}, lhs_batch_dims={0}, rhs_batch_dims={0} +} + +// ----- + +// CHECK-LABEL: @test_dot_transpose_lhs +// CHECK-SAME: %[[ARG0:.*]]: tensor<32x4x8xf32>, %[[ARG1:.*]]: tensor<4x32x16xf32> +// CHECK: %[[LHS:.*]] = transpose %[[ARG0]] permutation = [1, 2, 0] : tensor<32x4x8xf32> -> tensor<4x8x32xf32> +// CHECK: %[[OUT:.*]] = matmul(%[[LHS]], %[[ARG1]]) : (tensor<4x8x32xf32>, tensor<4x32x16xf32>) -> tensor<4x8x16xf32> +// CHECK: results %[[OUT]] + +test_dot_transpose_lhs { + %p0 = f32[32,4,8] parameter(0) + %p1 = f32[4,32,16] parameter(1) + %out = f32[4,8,16] dot(%p0, %p1), lhs_contracting_dims={0}, rhs_contracting_dims={1}, lhs_batch_dims={1}, rhs_batch_dims={0} +} + +// ----- + +// CHECK-LABEL: @test_dot_transpose_rhs +// CHECK-SAME: %[[ARG0:.*]]: tensor<4x8x32xf32>, %[[ARG1:.*]]: tensor<32x4x16xf32> +// CHECK: %[[RHS:.*]] = transpose %[[ARG1]] permutation = [1, 0, 2] : tensor<32x4x16xf32> -> tensor<4x32x16xf32> +// CHECK: %[[OUT:.*]] = matmul(%[[ARG0]], %[[RHS]]) : (tensor<4x8x32xf32>, tensor<4x32x16xf32>) -> tensor<4x8x16xf32> +// CHECK: results %[[OUT]] + +test_dot_transpose_rhs { + %p0 = f32[4,8,32] parameter(0) + %p1 = f32[32,4,16] parameter(1) + %out = f32[4,8,16] dot(%p0, %p1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, lhs_batch_dims={0}, rhs_batch_dims={1} +} + +// ----- + +// CHECK-LABEL: @test_dot_reshape_batch +// CHECK-SAME: %[[ARG0:.*]]: tensor<2x4x8x32xf32>, %[[ARG1:.*]]: tensor<2x4x32x16xf32> +// CHECK: %[[LHS:.*]] = reshape %[[ARG0]] : tensor<2x4x8x32xf32> -> tensor<8x8x32xf32> +// CHECK: %[[RHS:.*]] = reshape %[[ARG1]] : tensor<2x4x32x16xf32> -> tensor<8x32x16xf32> +// CHECK: %[[MATMUL:.*]] = matmul(%[[LHS]], %[[RHS]]) : (tensor<8x8x32xf32>, tensor<8x32x16xf32>) -> tensor<8x8x16xf32> +// CHECK: %[[OUT:.*]] = reshape %[[MATMUL]] : tensor<8x8x16xf32> -> tensor<2x4x8x16xf32> +// CHECK: results %[[OUT]] + +test_dot_reshape_batch { + %p0 = f32[2,4,8,32] parameter(0) + %p1 = f32[2,4,32,16] parameter(1) + %out = f32[2,4,8,16] dot(%p0, %p1), lhs_contracting_dims={3}, rhs_contracting_dims={2}, lhs_batch_dims={0,1}, rhs_batch_dims={0,1} +} + +// ----- + +// CHECK-LABEL: @test_dot_reshape_contracting +// CHECK-SAME: %[[ARG0:.*]]: tensor<4x8x32x2xf32>, %[[ARG1:.*]]: tensor<4x32x2x16xf32> +// CHECK: %[[LHS:.*]] = reshape %[[ARG0]] : tensor<4x8x32x2xf32> -> tensor<4x8x64xf32> +// CHECK: %[[RHS:.*]] = reshape %[[ARG1]] : tensor<4x32x2x16xf32> -> tensor<4x64x16xf32> +// CHECK: %[[OUT:.*]] = matmul(%[[LHS]], %[[RHS]]) : (tensor<4x8x64xf32>, tensor<4x64x16xf32>) -> tensor<4x8x16xf32> +// CHECK: results %[[OUT]] + +test_dot_reshape_contracting { + %p0 = f32[4,8,32,2] parameter(0) + %p1 = f32[4,32,2,16] parameter(1) + %out = f32[4,8,16] dot(%p0, %p1), lhs_contracting_dims={2,3}, rhs_contracting_dims={1,2}, lhs_batch_dims={0}, rhs_batch_dims={0} +} + +// ----- + +// CHECK-LABEL: @test_dot_reshape_lhs_noncontracting +// CHECK-SAME: %[[ARG0:.*]]: tensor<4x8x8x32xf32>, %[[ARG1:.*]]: tensor<4x32x16xf32> +// CHECK: %[[LHS:.*]] = reshape %[[ARG0]] : tensor<4x8x8x32xf32> -> tensor<4x64x32xf32> +// CHECK: %[[MATMUL:.*]] = matmul(%[[LHS]], %[[ARG1]]) : (tensor<4x64x32xf32>, tensor<4x32x16xf32>) -> tensor<4x64x16xf32> +// CHECK: %[[OUT:.*]] = reshape %[[MATMUL]] : tensor<4x64x16xf32> -> tensor<4x8x8x16xf32> +// CHECK: results %[[OUT]] + +test_dot_reshape_lhs_noncontracting { + %p0 = f32[4,8,8,32] parameter(0) + %p1 = f32[4,32,16] parameter(1) + %out = f32[4,8,8,16] dot(%p0, %p1), lhs_contracting_dims={3}, rhs_contracting_dims={1}, lhs_batch_dims={0}, rhs_batch_dims={0} +} + +// ----- + +// CHECK-LABEL: @test_dot_reshape_rhs_noncontracting +// CHECK-SAME: %[[ARG0:.*]]: tensor<4x8x32xf32>, %[[ARG1:.*]]: tensor<4x32x16x16xf32> +// CHECK: %[[RHS:.*]] = reshape %[[ARG1]] : tensor<4x32x16x16xf32> -> tensor<4x32x256xf32> +// CHECK: %[[MATMUL:.*]] = matmul(%[[ARG0]], %[[RHS]]) : (tensor<4x8x32xf32>, tensor<4x32x256xf32>) -> tensor<4x8x256xf32> +// CHECK: %[[OUT:.*]] = reshape %[[MATMUL]] : tensor<4x8x256xf32> -> tensor<4x8x16x16xf32> +// CHECK: results %[[OUT]] + +test_dot_reshape_rhs_noncontracting { + %p0 = f32[4,8,32] parameter(0) + %p1 = f32[4,32,16,16] parameter(1) + %out = f32[4,8,16,16] dot(%p0, %p1), lhs_contracting_dims={2}, rhs_contracting_dims={1}, lhs_batch_dims={0}, rhs_batch_dims={0} +} + +// ----- + +// CHECK-LABEL: @test_dot_transpose_reshape_all +// CHECK-SAME: %[[ARG0:.*]]: tensor<9x8x5x4x3x2xf32>, %[[ARG1:.*]]: tensor<2x6x8x3x7x9xf32> +// CHECK: %[[T1:.*]] = transpose %[[ARG0]] permutation = [5, 4, 2, 3, 1, 0] : tensor<9x8x5x4x3x2xf32> -> tensor<2x3x5x4x8x9xf32> +// CHECK: %[[LHS:.*]] = reshape %[[T1]] : tensor<2x3x5x4x8x9xf32> -> tensor<6x20x72xf32> +// CHECK: %[[T2:.*]] = transpose %[[ARG1]] permutation = [0, 3, 2, 5, 1, 4] : tensor<2x6x8x3x7x9xf32> -> tensor<2x3x8x9x6x7xf32> +// CHECK: %[[RHS:.*]] = reshape %[[T2]] : tensor<2x3x8x9x6x7xf32> -> tensor<6x72x42xf32> +// CHECK: %[[MATMUL:.*]] = matmul(%[[LHS]], %[[RHS]]) : (tensor<6x20x72xf32>, tensor<6x72x42xf32>) -> tensor<6x20x42xf32> +// CHECK: %[[OUT:.*]] = reshape %[[MATMUL]] : tensor<6x20x42xf32> -> tensor<2x3x5x4x6x7xf32> +// CHECK: results %[[OUT]] + +test_dot_transpose_reshape_all { + %p0 = f32[9,8,5,4,3,2] parameter(0) + %p1 = f32[2,6,8,3,7,9] parameter(1) + %out = f32[2,3,5,4,6,7] dot(%p0, %p1), lhs_contracting_dims={1,0}, rhs_contracting_dims={2,5}, lhs_batch_dims={5,4}, rhs_batch_dims={0,3} +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_binary.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_binary.hlo new file mode 100644 index 00000000000000..22af29c2a17fed --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_binary.hlo @@ -0,0 +1,142 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_add +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = add %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_add { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = f32[8,16] add(%p0, %p1) +} + +// ----- + +// CHECK-LABEL: @test_atan2 +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = atan2 %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_atan2 { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = f32[8,16] atan2(%p0, %p1) +} + +// ----- + +// CHECK-LABEL: @test_div +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = div %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_div { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = f32[8,16] divide(%p0, %p1) +} + +// ----- + +// CHECK-LABEL: @test_max +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = max %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_max { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = f32[8,16] maximum(%p0, %p1) +} + +// ----- + +// CHECK-LABEL: @test_min +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = min %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_min { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = f32[8,16] minimum(%p0, %p1) +} + +// ----- + +// CHECK-LABEL: @test_mul +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = mul %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_mul { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = f32[8,16] multiply(%p0, %p1) +} + +// ----- + +// CHECK-LABEL: @test_pow +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = pow %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_pow { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = f32[8,16] power(%p0, %p1) +} + +// ----- + +// CHECK-LABEL: @test_rem +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = rem %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_rem { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = f32[8,16] remainder(%p0, %p1) +} + +// ----- + +// CHECK-LABEL: @test_sub +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = sub %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_sub { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = f32[8,16] subtract(%p0, %p1) +} + +// ----- + +// CHECK-LABEL: @test_and +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xi1>, %[[ARG1:.*]]: tensor<8x16xi1> +// CHECK: %[[OUT:.*]] = and %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_and { + %p0 = pred[8,16] parameter(0) + %p1 = pred[8,16] parameter(1) + %out = pred[8,16] and(%p0, %p1) +} + +// ----- + +// CHECK-LABEL: @test_or +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xi1>, %[[ARG1:.*]]: tensor<8x16xi1> +// CHECK: %[[OUT:.*]] = or %[[ARG0]], %[[ARG1]] +// CHECK: results %[[OUT]] + +test_or { + %p0 = pred[8,16] parameter(0) + %p1 = pred[8,16] parameter(1) + %out = pred[8,16] or(%p0, %p1) +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_compare.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_compare.hlo new file mode 100644 index 00000000000000..ce9304cf22031e --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_compare.hlo @@ -0,0 +1,143 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_float_eq +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: cmp %[[ARG0]] oeq %[[ARG1]] + +test_float_eq { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=EQ +} + +// ----- + +// CHECK-LABEL: @test_float_ne +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: cmp %[[ARG0]] une %[[ARG1]] + +test_float_ne { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=NE +} + +// ----- + +// CHECK-LABEL: @test_float_gt +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: cmp %[[ARG0]] ogt %[[ARG1]] + +test_float_gt { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=GT +} + +// ----- + +// CHECK-LABEL: @test_float_ge +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: cmp %[[ARG0]] oge %[[ARG1]] + +test_float_ge { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=GE +} + +// ----- + +// CHECK-LABEL: @test_float_lt +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: cmp %[[ARG0]] olt %[[ARG1]] + +test_float_lt { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=LT +} + +// ----- + +// CHECK-LABEL: @test_float_le +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32> +// CHECK: cmp %[[ARG0]] ole %[[ARG1]] + +test_float_le { + %p0 = f32[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=LE +} + +// ----- + +// CHECK-LABEL: @test_integer_eq +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xsi32>, %[[ARG1:.*]]: tensor<8x16xsi32> +// CHECK: cmp %[[ARG0]] eq %[[ARG1]] + +test_integer_eq { + %p0 = s32[8,16] parameter(0) + %p1 = s32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=EQ +} + +// ----- + +// CHECK-LABEL: @test_integer_ne +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xsi32>, %[[ARG1:.*]]: tensor<8x16xsi32> +// CHECK: cmp %[[ARG0]] neq %[[ARG1]] + +test_integer_ne { + %p0 = s32[8,16] parameter(0) + %p1 = s32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=NE +} + +// ----- + +// CHECK-LABEL: @test_integer_lt +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xsi32>, %[[ARG1:.*]]: tensor<8x16xsi32> +// CHECK: cmp %[[ARG0]] lt %[[ARG1]] + +test_integer_lt { + %p0 = s32[8,16] parameter(0) + %p1 = s32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=LT +} + +// ----- + +// CHECK-LABEL: @test_integer_le +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xsi32>, %[[ARG1:.*]]: tensor<8x16xsi32> +// CHECK: cmp %[[ARG0]] le %[[ARG1]] + +test_integer_le { + %p0 = s32[8,16] parameter(0) + %p1 = s32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=LE +} + +// ----- + +// CHECK-LABEL: @test_integer_gt +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xsi32>, %[[ARG1:.*]]: tensor<8x16xsi32> +// CHECK: cmp %[[ARG0]] gt %[[ARG1]] + +test_integer_gt { + %p0 = s32[8,16] parameter(0) + %p1 = s32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=GT +} + +// ----- + +// CHECK-LABEL: @test_integer_ge +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xsi32>, %[[ARG1:.*]]: tensor<8x16xsi32> +// CHECK: cmp %[[ARG0]] ge %[[ARG1]] + +test_integer_ge { + %p0 = s32[8,16] parameter(0) + %p1 = s32[8,16] parameter(1) + %out = pred[8,16] compare(%p0, %p1), direction=GE +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_convert.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_convert.hlo new file mode 100644 index 00000000000000..813fbe9c3efee0 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_convert.hlo @@ -0,0 +1,129 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_float_extend +// CHECK: convert %{{.*}} : tensor<8x16xf32> -> tensor<8x16xf64> + +test_float_extend { + %p0 = f32[8,16] parameter(0) + %out = f64[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_float_truncate +// CHECK: convert %{{.*}} : tensor<8x16xf32> -> tensor<8x16xf16> + +test_float_truncate { + %p0 = f32[8,16] parameter(0) + %out = f16[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_float_to_signed +// CHECK: convert %{{.*}} : tensor<8x16xf32> -> tensor<8x16xsi32> + +test_float_to_signed { + %p0 = f32[8,16] parameter(0) + %out = s32[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_float_to_unsigned +// CHECK: convert %{{.*}} : tensor<8x16xf32> -> tensor<8x16xui32> + +test_float_to_unsigned { + %p0 = f32[8,16] parameter(0) + %out = u32[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_float_to_pred +// CHECK: convert %{{.*}} : tensor<8x16xf32> -> tensor<8x16xi1> + +test_float_to_pred { + %p0 = f32[8,16] parameter(0) + %out = pred[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_signed_extend +// CHECK: convert %{{.*}} : tensor<8x16xsi16> -> tensor<8x16xsi64> + +test_signed_extend { + %p0 = s16[8,16] parameter(0) + %out = s64[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_signed_truncate +// CHECK: convert %{{.*}} : tensor<8x16xsi16> -> tensor<8x16xsi8> + +test_signed_truncate { + %p0 = s16[8,16] parameter(0) + %out = s8[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_signed_to_float +// CHECK: convert %{{.*}} : tensor<8x16xsi16> -> tensor<8x16xbf16> + +test_signed_to_float { + %p0 = s16[8,16] parameter(0) + %out = bf16[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_signed_to_pred +// CHECK: convert %{{.*}} : tensor<8x16xsi16> -> tensor<8x16xi1> + +test_signed_to_pred { + %p0 = s16[8,16] parameter(0) + %out = pred[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_unsigned_extend +// CHECK: convert %{{.*}} : tensor<8x16xui16> -> tensor<8x16xui64> + +test_unsigned_extend { + %p0 = u16[8,16] parameter(0) + %out = u64[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_unsigned_truncate +// CHECK: convert %{{.*}} : tensor<8x16xui16> -> tensor<8x16xui8> + +test_unsigned_truncate { + %p0 = u16[8,16] parameter(0) + %out = u8[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_unsigned_to_float +// CHECK: convert %{{.*}} : tensor<8x16xui16> -> tensor<8x16xbf16> + +test_unsigned_to_float { + %p0 = u16[8,16] parameter(0) + %out = bf16[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_unsigned_to_pred +// CHECK: convert %{{.*}} : tensor<8x16xui16> -> tensor<8x16xi1> + +test_unsigned_to_pred { + %p0 = u16[8,16] parameter(0) + %out = pred[8,16] convert(%p0) +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_ternary.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_ternary.hlo new file mode 100644 index 00000000000000..4251883b0f1f0c --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_ternary.hlo @@ -0,0 +1,28 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_select +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xi1>, %[[ARG1:.*]]: tensor<8x16xf32>, %[[ARG2:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = binary_select %[[ARG0]], %[[ARG1]], %[[ARG2]] +// CHECK: results %[[OUT]] + +test_select { + %p0 = pred[8,16] parameter(0) + %p1 = f32[8,16] parameter(1) + %p2 = f32[8,16] parameter(2) + %out = f32[8,16] select(%p0, %p1, %p2) +} + +// ----- + +// CHECK-LABEL: @test_clamp +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32>, %[[ARG2:.*]]: tensor<8x16xf32> +// CHECK: %[[MAX:.*]] = max %[[ARG1]], %[[ARG0]] +// CHECK: %[[OUT:.*]] = min %[[ARG2]], %[[MAX]] +// CHECK: results %[[OUT]] + +test_clamp { + %p0 = f32[8,16] parameter(0) + %min = f32[8,16] parameter(1) + %max = f32[8,16] parameter(2) + %out = f32[8,16] clamp(%min, %p0, %max) +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_unary.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_unary.hlo new file mode 100644 index 00000000000000..4c904091ac8edd --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_unary.hlo @@ -0,0 +1,207 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_abs +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = abs %[[ARG0]] +// CHECK: results %[[OUT]] + +test_abs { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] abs(%p0) +} + +// ----- + +// CHECK-LABEL: @test_ceil +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = ceil %[[ARG0]] +// CHECK: results %[[OUT]] + +test_ceil { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] ceil(%p0) +} + +// ----- + +// CHECK-LABEL: @test_convert +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = convert %[[ARG0]] : tensor<8x16xf32> -> tensor<8x16xbf16> +// CHECK: results %[[OUT]] + +test_convert { + %p0 = f32[8,16] parameter(0) + %out = bf16[8,16] convert(%p0) +} + +// ----- + +// CHECK-LABEL: @test_cos +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = cos %[[ARG0]] +// CHECK: results %[[OUT]] + +test_cos { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] cosine(%p0) +} + +// ----- + +// CHECK-LABEL: @test_erf +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = erf %[[ARG0]] +// CHECK: results %[[OUT]] + +test_erf { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] erf(%p0) +} + +// ----- + +// CHECK-LABEL: @test_exp +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = exp %[[ARG0]] +// CHECK: results %[[OUT]] + +test_exp { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] exponential(%p0) +} + +// ----- + +// CHECK-LABEL: @test_expm1 +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK-DAG: %[[ONE:.*]] = constant dense<1.000000e+00> +// CHECK-DAG: %[[EXP:.*]] = exp %[[ARG0]] +// CHECK: %[[OUT:.*]] = sub %[[EXP]], %[[ONE]] +// CHECK: results %[[OUT]] + +test_expm1 { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] exponential-minus-one(%p0) +} + +// ----- + +// CHECK-LABEL: @test_floor +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = floor %[[ARG0]] +// CHECK: results %[[OUT]] + +test_floor { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] floor(%p0) +} + +// ----- + +// CHECK-LABEL: @test_log +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = log %[[ARG0]] +// CHECK: results %[[OUT]] + +test_log { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] log(%p0) +} + +// ----- + +// CHECK-LABEL: @test_log1p +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK-DAG: %[[ONE:.*]] = constant dense<1.000000e+00> +// CHECK-DAG: %[[ADD:.*]] = add %[[ARG0]], %[[ONE]] +// CHECK: %[[OUT:.*]] = log %[[ADD]] +// CHECK: results %[[OUT]] + +test_log1p { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] log-plus-one(%p0) +} + +// ----- + +// CHECK-LABEL: @test_not +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xi1> +// CHECK: %[[OUT:.*]] = not %[[ARG0]] +// CHECK: results %[[OUT]] + +test_not { + %p0 = pred[8,16] parameter(0) + %out = pred[8,16] not(%p0) +} + +// ----- + +// CHECK-LABEL: @test_neg +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = neg %[[ARG0]] +// CHECK: results %[[OUT]] + +test_neg { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] negate(%p0) +} + +// ----- + +// CHECK-LABEL: @test_rsqrt +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = rsqrt %[[ARG0]] +// CHECK: results %[[OUT]] + +test_rsqrt { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] rsqrt(%p0) +} + +// ----- + +// CHECK-LABEL: @test_sin +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = sin %[[ARG0]] +// CHECK: results %[[OUT]] + +test_sin { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] sine(%p0) +} + +// ----- + +// CHECK-LABEL: @test_sqrt +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = sqrt %[[ARG0]] +// CHECK: results %[[OUT]] + +test_sqrt { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] sqrt(%p0) +} + +// ----- + +// CHECK-LABEL: @test_tan +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = tan %[[ARG0]] +// CHECK: results %[[OUT]] + +test_tan { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] tan(%p0) +} + +// ----- + +// CHECK-LABEL: @test_tanh +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = tanh_fwd %[[ARG0]] +// CHECK: results %[[OUT]] + +test_tanh { + %p0 = f32[8,16] parameter(0) + %out = f32[8,16] tanh(%p0) +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_bitcast.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_bitcast.hlo new file mode 100644 index 00000000000000..079ed0680b168c --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_bitcast.hlo @@ -0,0 +1,128 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_reshape_split +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = reshape %[[ARG0]] : tensor<8x16xf32> -> tensor<8x4x4xf32> +// CHECK: results %[[OUT]] + +test_reshape_split { + %p0 = f32[8,16] parameter(0) + %out = f32[8,4,4] bitcast(%p0) +} + +// ----- + +// CHECK-LABEL: @test_reshape_join +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = reshape %[[ARG0]] : tensor<8x16xf32> -> tensor<128xf32> +// CHECK: results %[[OUT]] + +test_reshape_join { + %p0 = f32[8,16] parameter(0) + %out = f32[128] bitcast(%p0) +} + +// ----- + +// CHECK-LABEL: @test_reshape_mixed +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = reshape %[[ARG0]] : tensor<8x16xf32> -> tensor<32x4xf32> +// CHECK: results %[[OUT]] + +test_reshape_mixed { + %p0 = f32[8,16] parameter(0) + %out = f32[32,4] bitcast(%p0) +} + +// ----- + +// CHECK-LABEL: @test_input_layout +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16x32xf32> {nv_tensor_ir.stride = "(32,256,1)"} +// CHECK: %[[OUT:.*]] = transpose %[[ARG0]] permutation = [1, 0, 2] : tensor<8x16x32xf32> -> tensor<16x8x32xf32> +// CHECK: results %[[OUT]] + +test_input_layout { + %p0 = f32[8,16,32]{2,0,1} parameter(0) + %out = f32[16,8,32]{2,1,0} bitcast(%p0) +} + +// ----- + +// CHECK-LABEL: @test_output_layout +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16x32xf32> +// CHECK-SAME: -> (tensor<32x8x16xf32> {nv_tensor_ir.stride = "(1,512,32)"}) +// CHECK: %[[OUT:.*]] = transpose %[[ARG0]] permutation = [2, 0, 1] : tensor<8x16x32xf32> -> tensor<32x8x16xf32> +// CHECK: results %[[OUT]] + +test_output_layout { + %p0 = f32[8,16,32]{2,1,0} parameter(0) + %out = f32[32,8,16]{0,2,1} bitcast(%p0) +} + +// ----- + +// CHECK-LABEL: @test_input_output_layout +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16x32xf32> {nv_tensor_ir.stride = "(1,8,128)"} +// CHECK-SAME: -> (tensor<16x8x32xf32> {nv_tensor_ir.stride = "(8,1,128)"}) +// CHECK: %[[TEMP:.*]] = transpose %[[ARG0]] permutation = [2, 1, 0] : tensor<8x16x32xf32> -> tensor<32x16x8xf32> +// CHECK: %[[OUT:.*]] = transpose %[[TEMP]] permutation = [1, 2, 0] : tensor<32x16x8xf32> -> tensor<16x8x32xf32> +// CHECK: results %[[OUT]] + +test_input_output_layout { + %p0 = f32[8,16,32]{0,1,2} parameter(0) + %out = f32[16,8,32]{1,0,2} bitcast(%p0) +} + +// ----- + +// CHECK-LABEL: @test_transpose_reshape +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16x32xf32> {nv_tensor_ir.stride = "(32,256,1)"} +// CHECK: %[[TEMP:.*]] = transpose %[[ARG0]] permutation = [1, 0, 2] : tensor<8x16x32xf32> -> tensor<16x8x32xf32> +// CHECK: %[[OUT:.*]] = reshape %[[TEMP]] : tensor<16x8x32xf32> -> tensor<128x32xf32> +// CHECK: results %[[OUT]] + +test_transpose_reshape { + %p0 = f32[8,16,32]{2,0,1} parameter(0) + %out = f32[128,32]{1,0} bitcast(%p0) +} + +// ----- + +// CHECK-LABEL: @test_reshape_transpose +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16x32xf32> +// CHECK-SAME: -> (tensor<32x128xf32> {nv_tensor_ir.stride = "(1,32)"}) +// CHECK: %[[TEMP:.*]] = reshape %[[ARG0]] : tensor<8x16x32xf32> -> tensor<128x32xf32> +// CHECK: %[[OUT:.*]] = transpose %[[TEMP]] permutation = [1, 0] : tensor<128x32xf32> -> tensor<32x128xf32> +// CHECK: results %[[OUT]] + +test_reshape_transpose { + %p0 = f32[8,16,32]{2,1,0} parameter(0) + %out = f32[32,128]{0,1} bitcast(%p0) +} + +// ----- + +// CHECK-LABEL: @test_transpose_reshape_transpose +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16x32xf32> {nv_tensor_ir.stride = "(1,8,128)"} +// CHECK-SAME: -> (tensor<32x128xf32> {nv_tensor_ir.stride = "(1,32)"}) +// CHECK: %[[T1:.*]] = transpose %[[ARG0]] permutation = [2, 1, 0] : tensor<8x16x32xf32> -> tensor<32x16x8xf32> +// CHECK: %[[T2:.*]] = reshape %[[T1]] : tensor<32x16x8xf32> -> tensor<128x32xf32> +// CHECK: %[[OUT:.*]] = transpose %[[T2]] permutation = [1, 0] : tensor<128x32xf32> -> tensor<32x128xf32> +// CHECK: results %[[OUT]] + +test_transpose_reshape_transpose { + %p0 = f32[8,16,32]{0,1,2} parameter(0) + %out = f32[32,128]{0,1} bitcast(%p0) +} + +// ----- + +// CHECK-LABEL: @test_reshape_is_bitcast +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = reshape %[[ARG0]] : tensor<8x16xf32> -> tensor<8x4x4xf32> +// CHECK: results %[[OUT]] + +test_reshape_is_bitcast { + %p0 = f32[8,16] parameter(0) + %out = f32[8,4,4] reshape(%p0) +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_broadcast.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_broadcast.hlo new file mode 100644 index 00000000000000..874d349509c258 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_broadcast.hlo @@ -0,0 +1,63 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_brodcast_constant +// CHECK: %[[OUT:.*]] = constant dense<1.000000e+00> : tensor<8x16xf32> +// CHECK: results %[[OUT]] + +test_brodcast_constant { + %c0 = f32[] constant(1) + %out = f32[8,16] broadcast(%c0), dimensions={} +} + +// ----- + +// CHECK-LABEL: @test_broadcast_normal +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[TEMP:.*]] = reshape %[[ARG0]] : tensor<8x16xf32> -> tensor<1x8x16x1xf32> +// CHECK: %[[OUT:.*]] = broadcast %[[TEMP]] : tensor<1x8x16x1xf32> -> tensor<4x8x16x32xf32> +// CHECK: results %[[OUT]] + +test_broadcast_normal { + %p0 = f32[8,16] parameter(0) + %out = f32[4,8,16,32] broadcast(%p0), dimensions={1,2} +} + +// ----- + +// CHECK-LABEL: @test_broadcast_transpose +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[T1:.*]] = transpose %[[ARG0]] permutation = [1, 0] : tensor<8x16xf32> -> tensor<16x8xf32> +// CHECK: %[[T2:.*]] = reshape %[[T1]] : tensor<16x8xf32> -> tensor<16x1x1x8xf32> +// CHECK: %[[OUT:.*]] = broadcast %[[T2]] : tensor<16x1x1x8xf32> -> tensor<16x4x2x8xf32> +// CHECK: results %[[OUT]] + +test_broadcast_transpose { + %p0 = f32[8,16] parameter(0) + %out = f32[16,4,2,8] broadcast(%p0), dimensions={3,0} +} + +// ----- + +// CHECK-LABEL: @test_broadcast_scalar +// CHECK-SAME: %[[ARG0:.*]]: tensor<1xf32> +// CHECK: %[[TEMP:.*]] = reshape %[[ARG0]] : tensor<1xf32> -> tensor<1xf32> +// CHECK: %[[OUT:.*]] = broadcast %[[TEMP]] : tensor<1xf32> -> tensor<32xf32> +// CHECK: results %[[OUT]] + +test_broadcast_scalar { + %p0 = f32[] parameter(0) + %out = f32[32] broadcast(%p0), dimensions={} +} + +// ----- + +// CHECK-LABEL: @test_broadcast_unit +// CHECK-SAME: %[[ARG0:.*]]: tensor<1xf32> +// CHECK: %[[TEMP:.*]] = reshape %[[ARG0]] : tensor<1xf32> -> tensor<1x1xf32> +// CHECK: %[[OUT:.*]] = broadcast %[[TEMP]] : tensor<1x1xf32> -> tensor<1x32xf32> +// CHECK: results %[[OUT]] + +test_broadcast_unit { + %p0 = f32[1] parameter(0) + %out = f32[1,32] broadcast(%p0), dimensions={0} +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_other.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_other.hlo new file mode 100644 index 00000000000000..8dcaaf3e34472b --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_other.hlo @@ -0,0 +1,23 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_slice +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32> +// CHECK: %[[OUT:.*]] = slice %[[ARG0]] starts = [2, 0] limits = [6, 8] strides = [1, 2] : tensor<8x16xf32> -> tensor<4x4xf32> +// CHECK: results %[[OUT]] + +test_slice { + %p0 = f32[8,16] parameter(0) + %out = f32[4,4] slice(%p0), slice={[2:6], [0:8:2]} +} + +// ----- + +// CHECK-LABEL: @test_transpose +// CHECK-SAME: %[[ARG0:.*]]: tensor<2x4x8x16xf32> +// CHECK: %[[OUT:.*]] = transpose %[[ARG0]] permutation = [2, 0, 1, 3] : tensor<2x4x8x16xf32> -> tensor<8x2x4x16xf32> +// CHECK: results %[[OUT]] + +test_transpose { + %p0 = f32[2,4,8,16] parameter(0) + %out = f32[8,2,4,16] transpose(%p0), dimensions={2,0,1,3} +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/other_ops.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/other_ops.hlo new file mode 100644 index 00000000000000..1873f49267e286 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/other_ops.hlo @@ -0,0 +1,45 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_iota_float +// CHECK: %[[OUT:.*]] = iota dimension = 1 : tensor<8x16xf32> +// CHECK: results %[[OUT]] + +test_iota_float { + %out = f32[8,16] iota(), iota_dimension=1 +} + +// ----- + +// CHECK-LABEL: @test_iota_int +// CHECK: %[[OUT:.*]] = iota dimension = 0 : tensor<8x16xui32> +// CHECK: results %[[OUT]] + +test_iota_int { + %out = u32[8,16] iota(), iota_dimension=0 +} + +// ----- + +// CHECK-LABEL: @test_concat_2 +// CHECK-SAME: %[[ARG0:.*]]: tensor<2x16xf32>, %[[ARG1:.*]]: tensor<4x16xf32> +// CHECK: %[[OUT:.*]] = concatenate %[[ARG0]], %[[ARG1]] dimension = 0 : (tensor<2x16xf32>, tensor<4x16xf32>) -> tensor<6x16xf32> +// CHECK: results %[[OUT]] + +test_concat_2 { + %p0 = f32[2,16] parameter(0) + %p1 = f32[4,16] parameter(1) + %out = f32[6,16] concatenate(%p0, %p1), dimensions={0} +} + +// ----- + +// CHECK-LABEL: @test_concat_3 +// CHECK-SAME: %[[ARG0:.*]]: tensor<8x8xf32> +// CHECK: %[[OUT:.*]] = concatenate %[[ARG0]], %[[ARG0]], %[[ARG0]] dimension = 1 : +// CHECK-SAME: (tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>) -> tensor<8x24xf32> +// CHECK: results %[[OUT]] + +test_concat_3 { + %p0 = f32[8,8] parameter(0) + %out = f32[8,24] concatenate(%p0, %p0, %p0), dimensions={1} +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_compare.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_compare.hlo new file mode 100644 index 00000000000000..6512f3a3cf01c9 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_compare.hlo @@ -0,0 +1,101 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// The tests below only verify the lowering of comparison predicates. +// The computations are meaningless and cannot be verified at runtime. + +// CHECK-LABEL: @test_compare_float +// CHECK: arith.cmpf oeq +// CHECK: arith.cmpf une +// CHECK: arith.cmpf olt +// CHECK: arith.cmpf ole +// CHECK: arith.cmpf ogt +// CHECK: arith.cmpf oge + +compare_float { + %lhs = f32[] parameter(0) + %rhs = f32[] parameter(1) + %eq = pred[] compare(%lhs, %rhs), direction=EQ + %ne = pred[] compare(%lhs, %rhs), direction=NE + %lt = pred[] compare(%lhs, %rhs), direction=LT + %le = pred[] compare(%lhs, %rhs), direction=LE + %gt = pred[] compare(%lhs, %rhs), direction=GT + %ge = pred[] compare(%lhs, %rhs), direction=GE + %x1 = pred[] xor(%eq, %ne) + %x2 = pred[] xor(%x1, %lt) + %x3 = pred[] xor(%x2, %le) + %x4 = pred[] xor(%x3, %gt) + %x5 = pred[] xor(%x4, %ge) + ROOT %res = f32[] select(%x5, %lhs, %rhs) +} + +test_compare_float { + %p0 = f32[16,64] parameter(0) + %c0 = f32[] constant(0.0) + ROOT %out = f32[16] reduce(%p0, %c0), dimensions={1}, to_apply=compare_float +} + +// ----- + +// CHECK-LABEL: @test_compare_int_signed +// CHECK: arith.cmpi eq +// CHECK: arith.cmpi ne +// CHECK: arith.cmpi slt +// CHECK: arith.cmpi sle +// CHECK: arith.cmpi sgt +// CHECK: arith.cmpi sge + +compare_signed { + %lhs = s32[] parameter(0) + %rhs = s32[] parameter(1) + %eq = pred[] compare(%lhs, %rhs), direction=EQ + %ne = pred[] compare(%lhs, %rhs), direction=NE + %lt = pred[] compare(%lhs, %rhs), direction=LT + %le = pred[] compare(%lhs, %rhs), direction=LE + %gt = pred[] compare(%lhs, %rhs), direction=GT + %ge = pred[] compare(%lhs, %rhs), direction=GE + %x1 = pred[] xor(%eq, %ne) + %x2 = pred[] xor(%x1, %lt) + %x3 = pred[] xor(%x2, %le) + %x4 = pred[] xor(%x3, %gt) + %x5 = pred[] xor(%x4, %ge) + ROOT %res = s32[] select(%x5, %lhs, %rhs) +} + +test_compare_int_signed { + %p0 = s32[16,64] parameter(0) + %c0 = s32[] constant(0) + ROOT %out = s32[16] reduce(%p0, %c0), dimensions={1}, to_apply=compare_signed +} + +// ----- + +// CHECK-LABEL: @test_compare_int_unsigned +// CHECK: arith.cmpi eq +// CHECK: arith.cmpi ne +// CHECK: arith.cmpi ult +// CHECK: arith.cmpi ule +// CHECK: arith.cmpi ugt +// CHECK: arith.cmpi uge + +compare_unsigned { + %lhs = u32[] parameter(0) + %rhs = u32[] parameter(1) + %eq = pred[] compare(%lhs, %rhs), direction=EQ + %ne = pred[] compare(%lhs, %rhs), direction=NE + %lt = pred[] compare(%lhs, %rhs), direction=LT + %le = pred[] compare(%lhs, %rhs), direction=LE + %gt = pred[] compare(%lhs, %rhs), direction=GT + %ge = pred[] compare(%lhs, %rhs), direction=GE + %x1 = pred[] xor(%eq, %ne) + %x2 = pred[] xor(%x1, %lt) + %x3 = pred[] xor(%x2, %le) + %x4 = pred[] xor(%x3, %gt) + %x5 = pred[] xor(%x4, %ge) + ROOT %res = u32[] select(%x5, %lhs, %rhs) +} + +test_compare_int_unsigned { + %p0 = u32[16,64] parameter(0) + %c0 = u32[] constant(0) + ROOT %out = u32[16] reduce(%p0, %c0), dimensions={1}, to_apply=compare_unsigned +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_float.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_float.hlo new file mode 100644 index 00000000000000..552d91b3d005bf --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_float.hlo @@ -0,0 +1,188 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_reduce_add +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32> +// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) > : tensor<16x64xf32> -> tensor<16x1xf32> +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32> +// CHECK: results %[[OUT]] + +body_add { + %lhs = f32[] parameter(0) + %rhs = f32[] parameter(1) + ROOT %agg = f32[] add(%lhs, %rhs) +} + +test_reduce_add { + %p0 = f32[16,64] parameter(0) + %c0 = f32[] constant(0.0) + ROOT %out = f32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_add +} + +// ----- + +// CHECK-LABEL: @test_reduce_mul +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32> +// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) > : tensor<16x64xf32> -> tensor<16x1xf32> +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32> +// CHECK: results %[[OUT]] + +body_mul { + %lhs = f32[] parameter(0) + %rhs = f32[] parameter(1) + ROOT %agg = f32[] multiply(%lhs, %rhs) +} + +test_reduce_mul { + %p0 = f32[16,64] parameter(0) + %c1 = f32[] constant(1.0) + ROOT %out = f32[16] reduce(%p0, %c1), dimensions={1}, to_apply=body_mul +} + +// ----- + +// CHECK-LABEL: @test_reduce_add_flip +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: f32, %[[RHS:.*]]: f32) { +// CHECK: %[[RESULT:.*]] = arith.addf %[[RHS]], %[[LHS]] : f32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32> +// CHECK: results %[[OUT]] + +body_add_flip { + %lhs = f32[] parameter(0) + %rhs = f32[] parameter(1) + ROOT %agg = f32[] add(%rhs, %lhs) +} + +test_reduce_add_flip { + %p0 = f32[16,64] parameter(0) + %c0 = f32[] constant(0.0) + ROOT %out = f32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_add_flip +} + +// ----- + +// CHECK-LABEL: @test_reduce_mul_flip +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: f32, %[[RHS:.*]]: f32) { +// CHECK: %[[RESULT:.*]] = arith.mulf %[[RHS]], %[[LHS]] : f32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32> +// CHECK: results %[[OUT]] + +body_mul_flip { + %lhs = f32[] parameter(0) + %rhs = f32[] parameter(1) + ROOT %agg = f32[] multiply(%rhs, %lhs) +} + +test_reduce_mul_flip { + %p0 = f32[16,64] parameter(0) + %c1 = f32[] constant(1.0) + ROOT %out = f32[16] reduce(%p0, %c1), dimensions={1}, to_apply=body_mul_flip +} + +// ----- + +// CHECK-LABEL: @test_reduce_min +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: f32, %[[RHS:.*]]: f32) { +// CHECK: %[[RESULT:.*]] = arith.minimumf %[[LHS]], %[[RHS]] : f32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32> +// CHECK: results %[[OUT]] + +body_min { + %lhs = f32[] parameter(0) + %rhs = f32[] parameter(1) + ROOT %agg = f32[] minimum(%lhs, %rhs) +} + +test_reduce_min { + %p0 = f32[16,64] parameter(0) + %c0 = f32[] constant(inf) + ROOT %out = f32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_min +} + +// ----- + +// CHECK-LABEL: @test_reduce_max +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: f32, %[[RHS:.*]]: f32) { +// CHECK: %[[RESULT:.*]] = arith.maximumf %[[LHS]], %[[RHS]] : f32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32> +// CHECK: results %[[OUT]] + +body_max { + %lhs = f32[] parameter(0) + %rhs = f32[] parameter(1) + ROOT %agg = f32[] maximum(%lhs, %rhs) +} + +test_reduce_max { + %p0 = f32[16,64] parameter(0) + %c0 = f32[] constant(-inf) + ROOT %out = f32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_max +} + +// ----- + +// CHECK-LABEL: @test_reduce_clamp +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: f32, %[[RHS:.*]]: f32) { +// CHECK-DAG: %[[LOWER:.*]] = arith.constant -1.000000e+06 : f32 +// CHECK-DAG: %[[UPPER:.*]] = arith.constant 1.000000e+06 : f32 +// CHECK-DAG: %[[ADD:.*]] = arith.addf %[[LHS]], %[[RHS]] : f32 +// CHECK: %[[TEMP:.*]] = arith.maximumf %[[LOWER]], %[[ADD]] : f32 +// CHECK: %[[RESULT:.*]] = arith.minimumf %[[UPPER]], %[[TEMP]] : f32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32> +// CHECK: results %[[OUT]] + +body_clamp { + %lhs = f32[] parameter(0) + %rhs = f32[] parameter(1) + %add = f32[] add(%lhs, %rhs) + %min = f32[] constant(-1e6) + %max = f32[] constant(1e6) + ROOT %agg = f32[] clamp(%min, %add, %max) +} + +test_reduce_clamp { + %p0 = f32[16,64] parameter(0) + %c0 = f32[] constant(0.0) + ROOT %out = f32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_clamp +} + +// ----- + +// CHECK-LABEL: @test_reduce_compare_select +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: f32, %[[RHS:.*]]: f32) { +// CHECK: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 +// CHECK: %[[LHS_IS_ZERO:.*]] = arith.cmpf oeq, %[[LHS]], %[[ZERO]] : f32 +// CHECK: %[[RHS_IS_ZERO:.*]] = arith.cmpf oeq, %[[RHS]], %[[ZERO]] : f32 +// CHECK: %[[ANY_IS_ZERO:.*]] = arith.ori %[[LHS_IS_ZERO]], %[[RHS_IS_ZERO]] : i1 +// CHECK: %[[ONE:.*]] = arith.constant 1.000000e+00 : f32 +// CHECK: %[[RESULT:.*]] = arith.select %[[ANY_IS_ZERO]], %[[ZERO]], %[[ONE]] : f32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32> +// CHECK: results %[[OUT]] + +body_compare_select { + %lhs = f32[] parameter(0) + %rhs = f32[] parameter(1) + %zero = f32[] constant(0.0) + %lhs_is_zero = pred[] compare(%lhs, %zero), direction=EQ + %rhs_is_zero = pred[] compare(%rhs, %zero), direction=EQ + %any_is_zero = pred[] or(%lhs_is_zero, %rhs_is_zero) + %one = f32[] constant(1.0) + ROOT %agg = f32[] select(%any_is_zero, %zero, %one) +} + +test_reduce_compare_select { + %p0 = f32[16,64] parameter(0) + %c1 = f32[] constant(1.0) + ROOT %out = f32[16] reduce(%p0, %c1), dimensions={1}, to_apply=body_compare_select +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_signed.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_signed.hlo new file mode 100644 index 00000000000000..8f47188340e56a --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_signed.hlo @@ -0,0 +1,214 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_reduce_add_flip +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.addi %[[RHS]], %[[LHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32> +// CHECK: results %[[OUT]] + +body_add_flip { + %lhs = s32[] parameter(0) + %rhs = s32[] parameter(1) + ROOT %agg = s32[] add(%rhs, %lhs) +} + +test_reduce_add_flip { + %p0 = s32[16,64] parameter(0) + %c0 = s32[] constant(0) + ROOT %out = s32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_add_flip +} + +// ----- + +// CHECK-LABEL: @test_reduce_mul_flip +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.muli %[[RHS]], %[[LHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32> +// CHECK: results %[[OUT]] + +body_mul_flip { + %lhs = s32[] parameter(0) + %rhs = s32[] parameter(1) + ROOT %agg = s32[] multiply(%rhs, %lhs) +} + +test_reduce_mul_flip { + %p0 = s32[16,64] parameter(0) + %c1 = s32[] constant(1) + ROOT %out = s32[16] reduce(%p0, %c1), dimensions={1}, to_apply=body_mul_flip +} + +// ----- + +// CHECK-LABEL: @test_reduce_min +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.minsi %[[LHS]], %[[RHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32> +// CHECK: results %[[OUT]] + +body_min { + %lhs = s32[] parameter(0) + %rhs = s32[] parameter(1) + ROOT %agg = s32[] minimum(%lhs, %rhs) +} + +test_reduce_min { + %p0 = s32[16,64] parameter(0) + %c0 = s32[] constant(2147483647) + ROOT %out = s32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_min +} + +// ----- + +// CHECK-LABEL: @test_reduce_max +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.maxsi %[[LHS]], %[[RHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32> +// CHECK: results %[[OUT]] + +body_max { + %lhs = s32[] parameter(0) + %rhs = s32[] parameter(1) + ROOT %agg = s32[] maximum(%lhs, %rhs) +} + +test_reduce_max { + %p0 = s32[16,64] parameter(0) + %c0 = s32[] constant(-2147483648) + ROOT %out = s32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_max +} + +// ----- + +// CHECK-LABEL: @test_reduce_clamp +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK-DAG: %[[LOWER:.*]] = arith.constant -1000000 : i32 +// CHECK-DAG: %[[UPPER:.*]] = arith.constant 1000000 : i32 +// CHECK-DAG: %[[ADD:.*]] = arith.addi %[[LHS]], %[[RHS]] : i32 +// CHECK: %[[TEMP:.*]] = arith.maxsi %[[LOWER]], %[[ADD]] : i32 +// CHECK: %[[RESULT:.*]] = arith.minsi %[[UPPER]], %[[TEMP]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32> +// CHECK: results %[[OUT]] + +body_clamp { + %lhs = s32[] parameter(0) + %rhs = s32[] parameter(1) + %add = s32[] add(%lhs, %rhs) + %min = s32[] constant(-1000000) + %max = s32[] constant(1000000) + ROOT %agg = s32[] clamp(%min, %add, %max) +} + +test_reduce_clamp { + %p0 = s32[16,64] parameter(0) + %c0 = s32[] constant(0) + ROOT %out = s32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_clamp +} + +// ----- + +// CHECK-LABEL: @test_reduce_compare_select +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[ZERO:.*]] = arith.constant 0 : i32 +// CHECK: %[[LHS_IS_ZERO:.*]] = arith.cmpi eq, %[[LHS]], %[[ZERO]] : i32 +// CHECK: %[[RHS_IS_ZERO:.*]] = arith.cmpi eq, %[[RHS]], %[[ZERO]] : i32 +// CHECK: %[[ANY_IS_ZERO:.*]] = arith.ori %[[LHS_IS_ZERO]], %[[RHS_IS_ZERO]] : i1 +// CHECK: %[[ONE:.*]] = arith.constant 1 : i32 +// CHECK: %[[RESULT:.*]] = arith.select %[[ANY_IS_ZERO]], %[[ZERO]], %[[ONE]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32> +// CHECK: results %[[OUT]] + +body_compare_select { + %lhs = s32[] parameter(0) + %rhs = s32[] parameter(1) + %zero = s32[] constant(0) + %lhs_is_zero = pred[] compare(%lhs, %zero), direction=EQ + %rhs_is_zero = pred[] compare(%rhs, %zero), direction=EQ + %any_is_zero = pred[] or(%lhs_is_zero, %rhs_is_zero) + %one = s32[] constant(1) + ROOT %agg = s32[] select(%any_is_zero, %zero, %one) +} + +test_reduce_compare_select { + %p0 = s32[16,64] parameter(0) + %c1 = s32[] constant(1) + ROOT %out = s32[16] reduce(%p0, %c1), dimensions={1}, to_apply=body_compare_select +} + +// ----- + +// CHECK-LABEL: @test_reduce_and +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.andi %[[LHS]], %[[RHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32> +// CHECK: results %[[OUT]] + +body_and { + %lhs = s32[] parameter(0) + %rhs = s32[] parameter(1) + ROOT %agg = s32[] and(%lhs, %rhs) +} + +test_reduce_and { + %p0 = s32[16,64] parameter(0) + %c0 = s32[] constant(-1) + ROOT %out = s32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_and +} + +// ----- + +// CHECK-LABEL: @test_reduce_or +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.ori %[[LHS]], %[[RHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32> +// CHECK: results %[[OUT]] + +body_or { + %lhs = s32[] parameter(0) + %rhs = s32[] parameter(1) + ROOT %agg = s32[] or(%lhs, %rhs) +} + +test_reduce_or { + %p0 = s32[16,64] parameter(0) + %c0 = s32[] constant(0) + ROOT %out = s32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_or +} + +// ----- + +// CHECK-LABEL: @test_reduce_xor +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.xori %[[LHS]], %[[RHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32> +// CHECK: results %[[OUT]] + +body_xor { + %lhs = s32[] parameter(0) + %rhs = s32[] parameter(1) + ROOT %agg = s32[] xor(%lhs, %rhs) +} + +test_reduce_xor { + %p0 = s32[16,64] parameter(0) + %c0 = s32[] constant(0) + ROOT %out = s32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_xor +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_unsigned.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_unsigned.hlo new file mode 100644 index 00000000000000..dc97e8f4074f18 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_unsigned.hlo @@ -0,0 +1,216 @@ +// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s + +// CHECK-LABEL: @test_reduce_add_flip +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.addi %[[RHS]], %[[LHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32> +// CHECK: results %[[OUT]] + +body_add_flip { + %lhs = u32[] parameter(0) + %rhs = u32[] parameter(1) + ROOT %agg = u32[] add(%rhs, %lhs) +} + +test_reduce_add_flip { + %p0 = u32[16,64] parameter(0) + %c0 = u32[] constant(0) + ROOT %out = u32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_add_flip +} + +// ----- + +// CHECK-LABEL: @test_reduce_mul_flip +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.muli %[[RHS]], %[[LHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32> +// CHECK: results %[[OUT]] + +body_mul_flip { + %lhs = u32[] parameter(0) + %rhs = u32[] parameter(1) + ROOT %agg = u32[] multiply(%rhs, %lhs) +} + +test_reduce_mul_flip { + %p0 = u32[16,64] parameter(0) + %c1 = u32[] constant(1) + ROOT %out = u32[16] reduce(%p0, %c1), dimensions={1}, to_apply=body_mul_flip +} + +// ----- + +// CHECK-LABEL: @test_reduce_min +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.minui %[[LHS]], %[[RHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32> +// CHECK: results %[[OUT]] + +body_min { + %lhs = u32[] parameter(0) + %rhs = u32[] parameter(1) + ROOT %agg = u32[] minimum(%lhs, %rhs) +} + +test_reduce_min { + %p0 = u32[16,64] parameter(0) + %c0 = u32[] constant(4294967295) + ROOT %out = u32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_min +} + +// ----- + +// CHECK-LABEL: @test_reduce_max +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.maxui %[[LHS]], %[[RHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32> +// CHECK: results %[[OUT]] + +body_max { + %lhs = u32[] parameter(0) + %rhs = u32[] parameter(1) + ROOT %agg = u32[] maximum(%lhs, %rhs) +} + +test_reduce_max { + %p0 = u32[16,64] parameter(0) + %c0 = u32[] constant(0) + ROOT %out = u32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_max +} + +// ----- + +// CHECK-LABEL: @test_reduce_clamp +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui16> +// CHECK: %[[CONV:.*]] = convert %[[ARG0]] +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[CONV]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK-DAG: %[[LOWER:.*]] = arith.constant 0 : i32 +// CHECK-DAG: %[[UPPER:.*]] = arith.constant 1000000 : i32 +// CHECK-DAG: %[[ADD:.*]] = arith.addi %[[LHS]], %[[RHS]] : i32 +// CHECK: %[[TEMP:.*]] = arith.maxui %[[LOWER]], %[[ADD]] : i32 +// CHECK: %[[RESULT:.*]] = arith.minui %[[UPPER]], %[[TEMP]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32> +// CHECK: results %[[OUT]] + +body_clamp { + %lhs = u32[] parameter(0) + %rhs = u32[] parameter(1) + %add = u32[] add(%lhs, %rhs) + %min = u32[] constant(0) + %max = u32[] constant(1000000) + ROOT %agg = u32[] clamp(%min, %add, %max) +} + +test_reduce_clamp { + %p0 = u16[16,64] parameter(0) + %p0_ext = u32[16,64] convert(%p0) + %c0 = u32[] constant(0) + ROOT %out = u32[16] reduce(%p0_ext, %c0), dimensions={1}, to_apply=body_clamp +} + +// ----- + +// CHECK-LABEL: @test_reduce_compare_select +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[ZERO:.*]] = arith.constant 0 : i32 +// CHECK: %[[LHS_IS_ZERO:.*]] = arith.cmpi eq, %[[LHS]], %[[ZERO]] : i32 +// CHECK: %[[RHS_IS_ZERO:.*]] = arith.cmpi eq, %[[RHS]], %[[ZERO]] : i32 +// CHECK: %[[ANY_IS_ZERO:.*]] = arith.ori %[[LHS_IS_ZERO]], %[[RHS_IS_ZERO]] : i1 +// CHECK: %[[ONE:.*]] = arith.constant 1 : i32 +// CHECK: %[[RESULT:.*]] = arith.select %[[ANY_IS_ZERO]], %[[ZERO]], %[[ONE]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32> +// CHECK: results %[[OUT]] + +body_compare_select { + %lhs = u32[] parameter(0) + %rhs = u32[] parameter(1) + %zero = u32[] constant(0) + %lhs_is_zero = pred[] compare(%lhs, %zero), direction=EQ + %rhs_is_zero = pred[] compare(%rhs, %zero), direction=EQ + %any_is_zero = pred[] or(%lhs_is_zero, %rhs_is_zero) + %one = u32[] constant(1) + ROOT %agg = u32[] select(%any_is_zero, %zero, %one) +} + +test_reduce_compare_select { + %p0 = u32[16,64] parameter(0) + %c1 = u32[] constant(1) + ROOT %out = u32[16] reduce(%p0, %c1), dimensions={1}, to_apply=body_compare_select +} + +// ----- + +// CHECK-LABEL: @test_reduce_and +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.andi %[[LHS]], %[[RHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32> +// CHECK: results %[[OUT]] + +body_and { + %lhs = u32[] parameter(0) + %rhs = u32[] parameter(1) + ROOT %agg = u32[] and(%lhs, %rhs) +} + +test_reduce_and { + %p0 = u32[16,64] parameter(0) + %c0 = u32[] constant(4294967295) + ROOT %out = u32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_and +} + +// ----- + +// CHECK-LABEL: @test_reduce_or +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.ori %[[LHS]], %[[RHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32> +// CHECK: results %[[OUT]] + +body_or { + %lhs = u32[] parameter(0) + %rhs = u32[] parameter(1) + ROOT %agg = u32[] or(%lhs, %rhs) +} + +test_reduce_or { + %p0 = u32[16,64] parameter(0) + %c0 = u32[] constant(0) + ROOT %out = u32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_or +} + +// ----- + +// CHECK-LABEL: @test_reduce_xor +// CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32> +// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) { +// CHECK: %[[RESULT:.*]] = arith.xori %[[LHS]], %[[RHS]] : i32 +// CHECK: yield %[[RESULT]] +// CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32> +// CHECK: results %[[OUT]] + +body_xor { + %lhs = u32[] parameter(0) + %rhs = u32[] parameter(1) + ROOT %agg = u32[] xor(%lhs, %rhs) +} + +test_reduce_xor { + %p0 = u32[16,64] parameter(0) + %c0 = u32[] constant(0) + ROOT %out = u32[16] reduce(%p0, %c0), dimensions={1}, to_apply=body_xor +} diff --git a/third_party/xla/xla/backends/gpu/codegen/tools/BUILD b/third_party/xla/xla/backends/gpu/codegen/tools/BUILD index 6f22e415886697..6ad8aafadc13ed 100644 --- a/third_party/xla/xla/backends/gpu/codegen/tools/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/tools/BUILD @@ -85,6 +85,35 @@ xla_cc_binary( ], ) +xla_cc_binary( + name = "fusion_to_tensorir", + testonly = 1, + srcs = ["fusion_to_tensorir.cc"], + linkopts = ["-Wl,-rpath,$$ORIGIN/../lit_lib"], + tags = [ + "cuda-only", + "gpu", + ], + visibility = [":codegen_tests"], + deps = [ + "//xla/backends/gpu/codegen/tensor_ir:conversion", + "//xla/backends/gpu/codegen/tensor_ir:support", + "//xla/hlo/ir:hlo", + "//xla/tools:hlo_module_loader", + "@cuda_tile//:CudaTileDialect", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:MlirTranslateMain", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TranslateLib", + "@tensor_ir//:NVTensorIRCudaTilePipelines", + "@tensor_ir//:NVTensorIRDialect", + "@tensor_ir//:NVTensorIRToCudaTileConversion", + ], +) + xla_cc_binary( name = "gpu_test_correctness", testonly = 1, diff --git a/third_party/xla/xla/backends/gpu/codegen/tools/fusion_to_tensorir.cc b/third_party/xla/xla/backends/gpu/codegen/tools/fusion_to_tensorir.cc new file mode 100644 index 00000000000000..50249e81c5b047 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/tools/fusion_to_tensorir.cc @@ -0,0 +1,108 @@ +/* Copyright 2026 The OpenXLA Authors. + +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 + +#include "cuda_tile/Dialect/CudaTile/IR/Dialect.h" +#include "mlir/Tools/mlir-translate/MlirTranslateMain.h" +#include "mlir/Tools/mlir-translate/Translation.h" +#include "tensor_ir/Compiler/CudaTile/Pipelines.h" +#include "tensor_ir/Conversion/TensorToCudaTile/Options.h" +#include "tensor_ir/Dialect/TensorIR.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/LogicalResult.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Diagnostics.h" +#include "mlir/IR/DialectRegistry.h" +#include "mlir/IR/Location.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Support/LLVM.h" +#include "xla/backends/gpu/codegen/tensor_ir/conversion.h" +#include "xla/backends/gpu/codegen/tensor_ir/support.h" +#include "xla/hlo/ir/hlo_casting_utils.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instructions.h" +#include "xla/tools/hlo_module_loader.h" + +namespace xla::gpu::tensor_ir { +namespace { + +// NOLINTNEXTLINE +llvm::cl::opt compile_flag("compile", + llvm::cl::desc("Compile to CudaTile dialect."), + llvm::cl::init(false)); + +mlir::OwningOpRef HloToTensorIRTranslate( + llvm::StringRef input, mlir::MLIRContext* context) { + context->loadAllAvailableDialects(); + + auto hlo_module = xla::LoadModuleFromData(input, "hlo"); + if (!hlo_module.ok()) { + mlir::emitError(mlir::UnknownLoc::get(context)) + << hlo_module.status().message(); + return nullptr; + } + + const HloComputation* comp = (*hlo_module)->entry_computation(); + if (auto fusion = DynCast(comp->root_instruction()); + fusion != nullptr) { + comp = fusion->fused_instructions_computation(); + } + + if (auto decision = IsSupportedFusionComputation(*comp); + !decision.IsAllowed()) { + mlir::emitError(mlir::UnknownLoc::get(context)) << decision.Explain(); + return nullptr; + } + + auto module_or = ConvertFusionComputation(*comp, context); + if (!module_or.ok()) { + mlir::emitError(mlir::UnknownLoc::get(context)) + << module_or.status().message(); + return nullptr; + } + mlir::OwningOpRef module = *std::move(module_or); + + if (compile_flag) { + mlir::PassManager pass_manager(context); + mlir::nv_tensor_ir::TensorToCudaTilePipelineOptions options; + mlir::nv_tensor_ir::buildTensorToCudaTileConversionPipeline(pass_manager, + options); + if (llvm::failed(pass_manager.run(*module))) { + return nullptr; + } + } + + return module; +} + +static mlir::TranslateToMLIRRegistration hlo_to_tensorir_registration( + "hlo-to-tensorir", "Translate HLO to TensorIR", HloToTensorIRTranslate, + [](mlir::DialectRegistry& registry) { + registry.insert(); + }); + +} // namespace +} // namespace xla::gpu::tensor_ir + +int main(int argc, char** argv) { + return mlir::failed( + mlir::mlirTranslateMain(argc, argv, "HLO Fusion to TensorIR")); +} From 578270a49f17e64e2f4e83f46b52e0c02d042ba8 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sat, 29 Aug 2026 00:43:26 -0700 Subject: [PATCH 2/4] Automated Code Change PiperOrigin-RevId: 973005794 --- .../xla/codegen/emitters/transforms/convert_pure_call_ops.cc | 1 - .../xla/xla/codegen/emitters/transforms/erase_dead_functions.cc | 1 - .../xla/xla/codegen/emitters/transforms/expand_float_ops.cc | 1 - .../xla/xla/codegen/emitters/transforms/expand_integer_power.cc | 1 - .../xla/xla/codegen/emitters/transforms/flatten_tensors.cc | 1 - .../xla/xla/codegen/emitters/transforms/lower_pdl_ops.cc | 1 - .../xla/xla/codegen/emitters/transforms/lower_to_llvm_cpu.cc | 2 -- .../xla/xla/codegen/emitters/transforms/lower_xla_to_scf.cc | 1 - .../codegen/emitters/transforms/merge_pointers_to_same_slice.cc | 1 - .../xla/codegen/emitters/transforms/propagate_slice_indices.cc | 1 - .../xla/codegen/emitters/transforms/safe_integer_arithmetic.cc | 1 - .../xla/xla/codegen/emitters/transforms/simplify_affine.cc | 1 - .../xla/xla/codegen/emitters/transforms/unswitch_loops.cc | 1 - .../xla/codegen/emitters/transforms/vectorize_loads_stores.cc | 1 - 14 files changed, 15 deletions(-) diff --git a/third_party/xla/xla/codegen/emitters/transforms/convert_pure_call_ops.cc b/third_party/xla/xla/codegen/emitters/transforms/convert_pure_call_ops.cc index 612078c2c885a3..0512306d183688 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/convert_pure_call_ops.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/convert_pure_call_ops.cc @@ -9,7 +9,6 @@ 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 #include #include "mlir/Dialect/Func/IR/FuncOps.h" diff --git a/third_party/xla/xla/codegen/emitters/transforms/erase_dead_functions.cc b/third_party/xla/xla/codegen/emitters/transforms/erase_dead_functions.cc index e0d124c3119e83..6fdd88fdc6b194 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/erase_dead_functions.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/erase_dead_functions.cc @@ -12,7 +12,6 @@ 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 #include #include "mlir/Dialect/Func/IR/FuncOps.h" diff --git a/third_party/xla/xla/codegen/emitters/transforms/expand_float_ops.cc b/third_party/xla/xla/codegen/emitters/transforms/expand_float_ops.cc index 5e49c7f87a87eb..41f1e055117117 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/expand_float_ops.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/expand_float_ops.cc @@ -16,7 +16,6 @@ limitations under the License. #include #include #include -#include #include #include "llvm/ADT/APFloat.h" diff --git a/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc b/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc index e23499ddd11eba..23678c37ef094a 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include -#include #include #include "mhlo/IR/hlo_ops.h" diff --git a/third_party/xla/xla/codegen/emitters/transforms/flatten_tensors.cc b/third_party/xla/xla/codegen/emitters/transforms/flatten_tensors.cc index 806cfec29b34c1..b5c7aaa107cbfd 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/flatten_tensors.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/flatten_tensors.cc @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include -#include #include #include diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_pdl_ops.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_pdl_ops.cc index 8092faaf0f851d..02124e52767e65 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/lower_pdl_ops.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/lower_pdl_ops.cc @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include #include #include "mlir/Dialect/Arith/IR/Arith.h" diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_to_llvm_cpu.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_to_llvm_cpu.cc index 4ad8acdb2975d7..c111c939c0c994 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/lower_to_llvm_cpu.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/lower_to_llvm_cpu.cc @@ -15,8 +15,6 @@ limitations under the License. #include "xla/codegen/emitters/transforms/lower_to_llvm_cpu.h" -#include - #include "mlir/Dialect/Func/IR/FuncOps.h" // IWYU pragma: keep, needed by lower_to_llvm_cpu.h.inc. #include "mlir/Pass/Pass.h" #include "mlir/Support/LLVM.h" diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_to_scf.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_to_scf.cc index d77db7068b58fb..baaf8a57fd85c8 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_to_scf.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_to_scf.cc @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include -#include #include #include "llvm/ADT/SmallVector.h" diff --git a/third_party/xla/xla/codegen/emitters/transforms/merge_pointers_to_same_slice.cc b/third_party/xla/xla/codegen/emitters/transforms/merge_pointers_to_same_slice.cc index 2a296bb9f7582d..b8fb409f3704a6 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/merge_pointers_to_same_slice.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/merge_pointers_to_same_slice.cc @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include -#include #include #include diff --git a/third_party/xla/xla/codegen/emitters/transforms/propagate_slice_indices.cc b/third_party/xla/xla/codegen/emitters/transforms/propagate_slice_indices.cc index 6751ef37be9380..a76112a3a4daa2 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/propagate_slice_indices.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/propagate_slice_indices.cc @@ -12,7 +12,6 @@ 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 #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/IR/BuiltinAttributes.h" diff --git a/third_party/xla/xla/codegen/emitters/transforms/safe_integer_arithmetic.cc b/third_party/xla/xla/codegen/emitters/transforms/safe_integer_arithmetic.cc index 7a88eda3d5832b..fedfabf8a83575 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/safe_integer_arithmetic.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/safe_integer_arithmetic.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include -#include #include #include "mlir/Dialect/Arith/IR/Arith.h" diff --git a/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc b/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc index 8ba93f18ad2866..15f442ed6544e5 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include #include -#include #include #include diff --git a/third_party/xla/xla/codegen/emitters/transforms/unswitch_loops.cc b/third_party/xla/xla/codegen/emitters/transforms/unswitch_loops.cc index cac55cc994fec6..ae26b6947a9340 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/unswitch_loops.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/unswitch_loops.cc @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include -#include #include #include "mlir/Dialect/Arith/IR/Arith.h" diff --git a/third_party/xla/xla/codegen/emitters/transforms/vectorize_loads_stores.cc b/third_party/xla/xla/codegen/emitters/transforms/vectorize_loads_stores.cc index 525fd30114df34..04cb4743fb1d17 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/vectorize_loads_stores.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/vectorize_loads_stores.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include #include -#include #include #include "absl/algorithm/container.h" From 4cd4321b844b6d1686aeea09f02087814258134b Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sat, 29 Aug 2026 02:13:38 -0700 Subject: [PATCH 3/4] [XLA] Expose ShapeVerifier::CheckCompositeCall as static method Exposes ShapeVerifier::CheckCompositeCall as a static helper to allow derived shape verifiers to validate composite call frontend attributes. PiperOrigin-RevId: 973030477 --- third_party/xla/xla/service/hlo_verifier.cc | 62 +++++++++++---------- third_party/xla/xla/service/hlo_verifier.h | 1 + 2 files changed, 35 insertions(+), 28 deletions(-) diff --git a/third_party/xla/xla/service/hlo_verifier.cc b/third_party/xla/xla/service/hlo_verifier.cc index b6ab4b9cf7a4b5..62a54a223899fe 100644 --- a/third_party/xla/xla/service/hlo_verifier.cc +++ b/third_party/xla/xla/service/hlo_verifier.cc @@ -1687,40 +1687,46 @@ absl::Status ShapeVerifier::HandleFusion(HloInstruction* fusion) { return absl::OkStatus(); } +absl::Status ShapeVerifier::CheckCompositeCall(const HloInstruction* call) { + if (!call->is_composite()) { + return absl::OkStatus(); + } + TF_RET_CHECK(call->has_frontend_attributes()) + << "A composite call op must have frontend attributes"; + auto map = call->frontend_attributes().map(); + if (auto name = map.find("composite.name"); + name == map.end() || name->second.empty()) { + return InvalidArgument( + "A composite call op must have frontend attributes with key " + "composite.name whose value is non-empty"); + } + if (auto attributes = map.find("composite.attributes"); + attributes != map.end() && attributes->second.empty()) { + return InvalidArgument( + "A composite call op must have frontend attributes with key " + "composite.attributes whose value is default: {} or non-empty"); + } + if (auto version_str = map.find("composite.version"); + version_str != map.end()) { + int64_t version = 0; + if (!absl::SimpleAtoi(version_str->second, &version) || version < 0) { + return InvalidArgument( + "A composite call op must have frontend attributes with a " + "composite.version whose value is a non-negative integer but got: " + "%s", + version_str->second); + } + } + return absl::OkStatus(); +} + absl::Status ShapeVerifier::HandleCall(HloInstruction* call) { ABSL_RETURN_IF_ERROR( CheckParameterCount(call, call->to_apply(), call->operand_count())); for (int64_t i = 0; i < call->to_apply()->num_parameters(); ++i) { ABSL_RETURN_IF_ERROR(CheckOperandAndParameter(call, i, call->to_apply(), i)); } - if (call->is_composite()) { - TF_RET_CHECK(call->has_frontend_attributes()) - << "A composite call op must have frontend attributes"; - auto map = call->frontend_attributes().map(); - if (auto name = map.find("composite.name"); - name == map.end() || name->second.empty()) { - return InvalidArgument( - "A composite call op must have frontend attributes with key " - "composite.name whose value is non-empty"); - } - if (auto attributes = map.find("composite.attributes"); - attributes != map.end() && attributes->second.empty()) { - return InvalidArgument( - "A composite call op must have frontend attributes with key " - "composite.attributes whose value is default: {} or non-empty"); - } - if (auto version_str = map.find("composite.version"); - version_str != map.end()) { - int64_t version = 0; - if (!absl::SimpleAtoi(version_str->second, &version) || version < 0) { - return InvalidArgument( - "A composite call op must have frontend attributes with a " - "composite.version whose value is a non-negative integer but got: " - "%s", - version_str->second); - } - } - } + ABSL_RETURN_IF_ERROR(CheckCompositeCall(call)); // The shape of kCall should match the shape of the computation it calls. return CheckShape(call, call->to_apply()->root_instruction()->shape()); } diff --git a/third_party/xla/xla/service/hlo_verifier.h b/third_party/xla/xla/service/hlo_verifier.h index aeec726842ab7e..61f0d3e9ca8cf8 100644 --- a/third_party/xla/xla/service/hlo_verifier.h +++ b/third_party/xla/xla/service/hlo_verifier.h @@ -324,6 +324,7 @@ class ShapeVerifier : public DfsHloVisitor { static absl::Status CheckParameterCount( const HloInstruction* calling_instruction, const HloComputation* computation, int expected); + static absl::Status CheckCompositeCall(const HloInstruction* call); // Check a unary (binary, etc) instruction's shape against the inferred shape. absl::Status CheckUnaryShape(const HloInstruction* instruction); From 6efd011a63432c2b86d890a4b3a400cf539ec8c3 Mon Sep 17 00:00:00 2001 From: Alexander Belyaev Date: Sat, 29 Aug 2026 02:18:40 -0700 Subject: [PATCH 4/4] [XLA:CPU] Tighten the bounds for the offsets when checking if the mask is needed. Before we had [min(), max()] as the limit, but we can be more precise and that leads to more folding/simplification. PiperOrigin-RevId: 973032009 --- .../xla/xla/backends/cpu/codegen/BUILD | 1 - .../backends/cpu/codegen/fusion_compiler.cc | 22 +++++++++++++------ .../tiled/tests/elementwise/minimum_mask.hlo | 18 +++++---------- .../cpu/codegen/tiled/transforms/BUILD | 1 + .../transforms/tests/vectorize_xtile.mlir | 14 +++++++----- .../tiled/transforms/vectorize_xtile.cc | 19 +++++++++++----- 6 files changed, 43 insertions(+), 32 deletions(-) diff --git a/third_party/xla/xla/backends/cpu/codegen/BUILD b/third_party/xla/xla/backends/cpu/codegen/BUILD index 40f90b5b5e18df..3404d3c57148bc 100644 --- a/third_party/xla/xla/backends/cpu/codegen/BUILD +++ b/third_party/xla/xla/backends/cpu/codegen/BUILD @@ -228,7 +228,6 @@ cc_library( "//xla/tsl/framework/mlir:status_scoped_diagnostic_handler", "@com_google_absl//absl/base:config", "@com_google_absl//absl/functional:any_invocable", - "@com_google_absl//absl/functional:function_ref", "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", diff --git a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc index d16764f7dc8001..0734de13a5f737 100644 --- a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc +++ b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc @@ -22,7 +22,6 @@ limitations under the License. #include #include "absl/base/config.h" // IWYU pragma: keep -#include "absl/functional/function_ref.h" #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" @@ -133,6 +132,14 @@ limitations under the License. namespace xla::cpu { namespace { +emitters::SimplifyArithPassOptions GetSimplifyArithPassOptions( + bool fast_min_max) { + emitters::SimplifyArithPassOptions options; + options.fast_min_max_ = fast_min_max; + options.explicit_nan_propagation_ = false; + return options; +} + absl::Status RunPassPipeline(mlir::ModuleOp module, mlir::PassManager& pm, mlir::interpreter::MlirCompilationTrace* trace, int32_t verification_level) { @@ -231,11 +238,8 @@ void AddScalarOptimizationPasses(mlir::OpPassManager& pm, // These passes are primarily responsible for lowering individual ops to // their LLVM equivalent. void AddGenericLoweringPasses(mlir::OpPassManager& pm, bool fast_min_max) { - emitters::SimplifyArithPassOptions simplify_arith_options; - simplify_arith_options.fast_min_max_ = fast_min_max; - simplify_arith_options.explicit_nan_propagation_ = false; - pm.addNestedPass( - emitters::createSimplifyArithPass(simplify_arith_options)); + pm.addNestedPass(emitters::createSimplifyArithPass( + GetSimplifyArithPassOptions(fast_min_max))); pm.addPass(emitters::createExpandIntegerPowerPass()); pm.addPass(emitters::createSimplifyAffinePass()); pm.addPass(mlir::createCanonicalizerPass()); @@ -405,9 +409,13 @@ void AddNewXtileToVectorPasses(mlir::OpPassManager& pm) { pm.addPass(xtile::createLegalizeUnsignedIntegersAsSignlessPass()); pm.addPass(mlir::createCanonicalizerPass()); pm.addPass(cpu::createVectorizeXTilePass()); - pm.addPass(cpu::createLowerXTileEntryPass()); + pm.addNestedPass(emitters::createSimplifyArithPass( + GetSimplifyArithPassOptions(/*fast_min_max=*/false))); + pm.addPass(mlir::createCanonicalizerPass()); + pm.addPass(mlir::createCSEPass()); + pm.addNestedPass( mlir::stablehlo::createStablehloTargetIndependentOptimizationPass()); diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_mask.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_mask.hlo index dfbe21f909732f..8f0e0a6673a99d 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_mask.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_mask.hlo @@ -67,23 +67,17 @@ fusion { // NEW-VECTOR: %[[CMP0:.*]] = arith.cmpi sge, %[[BOUNDS]]#0, %{{.*}} : index // NEW-VECTOR: %[[CMP1:.*]] = arith.cmpi sge, %[[BOUNDS]]#1, %{{.*}} : index // NEW-VECTOR: %[[COND:.*]] = arith.andi %[[CMP0]], %[[CMP1]] : i1 -// NEW-VECTOR: %[[R0:.*]] = scf.if %[[COND]] -> (vector<16x32xf32>) { +// NEW-VECTOR: %[[R:.*]]:2 = scf.if %[[COND]] -> (vector<16x32xf32>, vector<16x32xf32>) { // NEW-VECTOR: %[[R0_IN:.*]] = vector.transfer_read %[[ARG0]][%{{.*}}, %{{.*}}], %[[PAD]] {in_bounds = [true, true]} : memref<100x100xf32>, vector<16x32xf32> -// NEW-VECTOR: scf.yield %[[R0_IN]] +// NEW-VECTOR: %[[R1_IN:.*]] = vector.transfer_read %[[ARG1]][%{{.*}}, %{{.*}}], %[[PAD]] {in_bounds = [true, true]} : memref<100x100xf32>, vector<16x32xf32> +// NEW-VECTOR: scf.yield %[[R0_IN]], %[[R1_IN]] // NEW-VECTOR: } else { // NEW-VECTOR: %[[MASK0:.*]] = vector.create_mask %{{.*}}, %{{.*}} : vector<16x32xi1> // NEW-VECTOR: %[[R0_MASK:.*]] = vector.transfer_read %[[ARG0]][%{{.*}}, %{{.*}}], %[[PAD]], %[[MASK0]] : memref<100x100xf32>, vector<16x32xf32> -// NEW-VECTOR: scf.yield %[[R0_MASK]] -// NEW-VECTOR: } -// NEW-VECTOR: %[[R1:.*]] = scf.if %[[COND]] -> (vector<16x32xf32>) { -// NEW-VECTOR: %[[R1_IN:.*]] = vector.transfer_read %[[ARG1]][%{{.*}}, %{{.*}}], %[[PAD]] {in_bounds = [true, true]} : memref<100x100xf32>, vector<16x32xf32> -// NEW-VECTOR: scf.yield %[[R1_IN]] -// NEW-VECTOR: } else { -// NEW-VECTOR: %[[MASK1:.*]] = vector.create_mask %{{.*}}, %{{.*}} : vector<16x32xi1> -// NEW-VECTOR: %[[R1_MASK:.*]] = vector.transfer_read %[[ARG1]][%{{.*}}, %{{.*}}], %[[PAD]], %[[MASK1]] : memref<100x100xf32>, vector<16x32xf32> -// NEW-VECTOR: scf.yield %[[R1_MASK]] +// NEW-VECTOR: %[[R1_MASK:.*]] = vector.transfer_read %[[ARG1]][%{{.*}}, %{{.*}}], %[[PAD]], %[[MASK0]] : memref<100x100xf32>, vector<16x32xf32> +// NEW-VECTOR: scf.yield %[[R0_MASK]], %[[R1_MASK]] // NEW-VECTOR: } -// NEW-VECTOR: %[[MIN:.*]] = arith.minimumf %[[R0]], %[[R1]] +// NEW-VECTOR: %[[MIN:.*]] = arith.minimumf %[[R]]#0, %[[R]]#1 // NEW-VECTOR: scf.if %[[COND]] { // NEW-VECTOR: vector.transfer_write %[[MIN]], %[[ARG2]][%{{.*}}, %{{.*}}] {in_bounds = [true, true]} : vector<16x32xf32>, memref<100x100xf32> // NEW-VECTOR: } else { diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD index 593c03155e27cf..6dc1b55ac96e50 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD @@ -89,6 +89,7 @@ cc_library( "//xla/codegen/xtile/ir:xtile", "//xla/codegen/xtile/ir/transforms:passes", "//xla/hlo/analysis:indexing_analysis", + "//xla/hlo/analysis:interval", "//xla/hlo/analysis:symbolic_map", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log:check", diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/vectorize_xtile.mlir b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/vectorize_xtile.mlir index 8afef5b3ae5e61..afdf033e7b769a 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/vectorize_xtile.mlir +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/vectorize_xtile.mlir @@ -626,7 +626,7 @@ func.func @test_extract_aligned(%arg0: memref<128xf32>, %arg1: index) -> tensor< %0 = xtile.extract %arg0[%c0] [8] [1] : memref<128xf32> -> tensor<8xf32> return %0 : tensor<8xf32> } -// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8){{.*}}"> +// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8), domain: d0 in [0, 0]"> // CHECK-LABEL: @test_extract_aligned // CHECK-DAG: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[INDEXING:[^:]+]] = xla.apply_indexing #indexing_map(%{{.*}}) @@ -642,11 +642,12 @@ func.func @test_extract_aligned(%arg0: memref<128xf32>, %arg1: index) -> tensor< // ----- -func.func @test_extract_unaligned(%arg0: memref<128xf32>, %arg1: index) -> tensor<8xf32> { +func.func @test_extract_unaligned(%arg0: memref<128xf32>, + %arg1: index {xla.range = [0 : index, 12 : index]}) -> tensor<8xf32> { %0 = xtile.extract %arg0[%arg1] [8] [1] : memref<128xf32> -> tensor<8xf32> return %0 : tensor<8xf32> } -// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8){{.*}}"> +// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8), domain: d0 in [0, 12]"> // CHECK-LABEL: @test_extract_unaligned // CHECK-DAG: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[INDEXING:[^:]+]] = xla.apply_indexing #indexing_map(%{{.*}}) @@ -667,7 +668,7 @@ func.func @test_insert_aligned(%arg0: tensor<8xf32>, %arg1: memref<128xf32>) { xtile.insert %arg0 into %arg1[%c0] [8] [1] : tensor<8xf32> -> memref<128xf32> return } -// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8){{.*}}"> +// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8), domain: d0 in [0, 0]"> // CHECK-LABEL: @test_insert_aligned // CHECK: %[[INDEXING:[^:]+]] = xla.apply_indexing #indexing_map(%{{.*}}) // CHECK: %[[COND:.*]] = arith.cmpi sge, %[[INDEXING]], %{{.*}} : index @@ -680,11 +681,12 @@ func.func @test_insert_aligned(%arg0: tensor<8xf32>, %arg1: memref<128xf32>) { // ----- -func.func @test_insert_unaligned(%arg0: tensor<8xf32>, %arg1: memref<128xf32>, %arg2: index) { +func.func @test_insert_unaligned(%arg0: tensor<8xf32>, %arg1: memref<128xf32>, + %arg2: index {xla.range = [0 : index, 42 : index]}) { xtile.insert %arg0 into %arg1[%arg2] [8] [1] : tensor<8xf32> -> memref<128xf32> return } -// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8){{.*}}"> +// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8), domain: d0 in [0, 42]"> // CHECK-LABEL: @test_insert_unaligned // CHECK: %[[INDEXING:[^:]+]] = xla.apply_indexing #indexing_map(%{{.*}}) // CHECK: %[[COND:.*]] = arith.cmpi sge, %[[INDEXING]], %{{.*}} : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc index ca3f38efe43f7b..a0c4b6c43315d2 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc @@ -54,6 +54,7 @@ limitations under the License. #include "xla/codegen/xtile/ir/xtile_dialect.h" #include "xla/codegen/xtile/ir/xtile_ops.h" #include "xla/hlo/analysis/indexing_map.h" +#include "xla/hlo/analysis/interval.h" #include "xla/hlo/analysis/symbolic_expr.h" #include "xla/hlo/analysis/symbolic_map.h" #include "xla/util.h" @@ -74,6 +75,17 @@ namespace shlo = ::mlir::stablehlo; using ::mlir::Value; using ::mlir::ValueRange; +std::vector GetVars(ValueRange values) { + std::vector vars; + vars.reserve(values.size()); + for (Value offset : values) { + vars.push_back(xla::IndexingMap::Variable{GetRange(offset).value_or( + Interval{std::numeric_limits::min(), + std::numeric_limits::max()})}); + } + return vars; +} + xla::SymbolicMap GetBoundsCheckSymbolicMap(mlir::MLIRContext* ctx, llvm::ArrayRef memref_shape, llvm::ArrayRef tile_shape) { @@ -92,17 +104,12 @@ Value GetIsInBoundsCondition(mlir::OpBuilder& builder, mlir::Location loc, ValueRange offsets, Value memref, llvm::ArrayRef tile_shape) { auto memref_shape = mlir::cast(memref.getType()).getShape(); - int rank = memref_shape.size(); xla::SymbolicMap symbolic_map = GetBoundsCheckSymbolicMap(builder.getContext(), memref_shape, tile_shape); - std::vector vars( - rank, xla::IndexingMap::Variable{std::numeric_limits::min(), - std::numeric_limits::max()}); - xla::IndexingMap indexing_map(symbolic_map, - /*dimensions=*/std::move(vars), + /*dimensions=*/GetVars(offsets), /*range_vars=*/{}, /*rt_vars=*/{}); auto apply_indexing =