diff --git a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir index a4d4f6a6ced441..195c1fadfad69b 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir +++ b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir @@ -2712,7 +2712,7 @@ func.func @neg_dynamic(%arg0: tensor) -> tensor { func.func @random_shuffle_float_vector(%arg0: tensor<2xf32>) -> tensor<2xf32> { // CHECK: "mhlo.sort" // CHECK: ^bb0([[KEY_LHS:%.*]]: tensor, [[KEY_RHS:%.*]]: tensor, {{.*}}: tensor, {{.*}}: tensor): - // CHECK: mhlo.compare LT, [[KEY_LHS]], [[KEY_RHS]], TOTALORDER : (tensor, tensor) -> tensor + // CHECK: mhlo.compare LT, [[KEY_LHS]], [[KEY_RHS]] : (tensor, tensor) -> tensor %0 = "tf.RandomShuffle"(%arg0) {seed = 1 : i64, seed2 = 2 : i64} : (tensor<2xf32>) -> tensor<2xf32> func.return %0 : tensor<2xf32> } diff --git a/tensorflow/core/kernels/bincount_op_gpu.cu.cc b/tensorflow/core/kernels/bincount_op_gpu.cu.cc index 056d78e555d762..98ed253ba6c46f 100644 --- a/tensorflow/core/kernels/bincount_op_gpu.cu.cc +++ b/tensorflow/core/kernels/bincount_op_gpu.cu.cc @@ -17,6 +17,8 @@ limitations under the License. #define EIGEN_USE_GPU +#include + #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" @@ -118,7 +120,8 @@ __global__ void BincountReduceKernel(const Tidx* in, T* out, const int nthreads, const Tidx num_bins) { GPU_1D_KERNEL_LOOP(index, nthreads) { Tidx bin = ldg(in + index); - if (bin < num_bins) { + if (static_cast>(bin) < + static_cast>(num_bins)) { out[bin] = T(1); } } @@ -149,7 +152,8 @@ __global__ void BincountColReduceKernel(const Tidx* in, const T* weights, const int nthreads = num_rows * num_cols; GPU_1D_KERNEL_LOOP(index, nthreads) { Tidx bin = ldg(in + index); - if (bin < num_bins) { + if (static_cast>(bin) < + static_cast>(num_bins)) { int row = index / num_cols; int offset = row * num_bins + bin; if (binary_count) { @@ -179,7 +183,8 @@ __global__ void BincountColReduceSharedKernel(const Tidx* in, const T* weights, const int nthreads = num_rows * num_cols; GPU_1D_KERNEL_LOOP(index, nthreads) { Tidx bin = ldg(in + index); - if (bin < num_bins) { + if (static_cast>(bin) < + static_cast>(num_bins)) { int row = index / num_cols; int offset = row * num_bins + bin; if (binary_count) { diff --git a/tensorflow/core/kernels/check_numerics_op.cc b/tensorflow/core/kernels/check_numerics_op.cc index fb34c7a4bd79ba..5535fe048589e9 100644 --- a/tensorflow/core/kernels/check_numerics_op.cc +++ b/tensorflow/core/kernels/check_numerics_op.cc @@ -289,8 +289,10 @@ class CheckNumericsOp : public AsyncOpKernel { const int is_nan = abnormality_indicators(0); const int is_inf = abnormality_indicators(1); if (is_nan || is_inf) { - LOG(ERROR) << "abnormal_detected_host @" << abnormality_indicators.data() - << " = {" << is_nan << ", " << is_inf << "} " << message_; + // Not LOG(ERROR): the InvalidArgumentError set below already reports + // this to the caller, so logging per detection would duplicate it. + VLOG(1) << "abnormal_detected_host @" << abnormality_indicators.data() + << " = {" << is_nan << ", " << is_inf << "} " << message_; std::string anomalies; if (is_nan && is_inf) { diff --git a/tensorflow/core/kernels/conv_ops_fused_impl.h b/tensorflow/core/kernels/conv_ops_fused_impl.h index 2a62c27d1f4a6e..b7f9f08eafa7bb 100644 --- a/tensorflow/core/kernels/conv_ops_fused_impl.h +++ b/tensorflow/core/kernels/conv_ops_fused_impl.h @@ -375,12 +375,21 @@ struct LaunchFusedConv2DOp { const int64_t out_cols = GetTensorDim(*output, params.data_format, 'W'); const int64_t out_depths = GetTensorDim(*output, params.data_format, 'C'); - // Bias of the following dimensions: [ output_depth ] + // Bias of the following dimensions: [ output_depth ] or [1, ..., 1, + // output_depth] const Tensor& bias = context->input(2); - OP_REQUIRES(context, bias.dims() == 1, - absl::InvalidArgumentError(absl::StrCat( - "bias must be 1-dimensional", bias.shape().DebugString()))); - OP_REQUIRES(context, bias.dim_size(0) == out_depths, + OP_REQUIRES(context, bias.dims() >= 1, + absl::InvalidArgumentError( + absl::StrCat("bias must be at least 1-dimensional, got: ", + bias.shape().DebugString()))); + for (int i = 0; i < bias.dims() - 1; ++i) { + OP_REQUIRES(context, bias.dim_size(i) == 1, + absl::InvalidArgumentError(absl::StrCat( + "For bias_dims > 1, all except the last dimension " + "must be 1, got: ", + bias.shape().DebugString()))); + } + OP_REQUIRES(context, bias.dim_size(bias.dims() - 1) == out_depths, absl::InvalidArgumentError( absl::StrCat("bias depth must be equal to out depth", bias.shape().DebugString()))); diff --git a/tensorflow/core/kernels/dilation_ops_gpu.cu.cc b/tensorflow/core/kernels/dilation_ops_gpu.cu.cc index 788ad9f5127b6d..d8da9c6f2acc34 100644 --- a/tensorflow/core/kernels/dilation_ops_gpu.cu.cc +++ b/tensorflow/core/kernels/dilation_ops_gpu.cu.cc @@ -36,27 +36,28 @@ namespace { template __global__ void DilationKernel( - const int32_t nthreads, const T* __restrict__ input_ptr, - const T* __restrict__ filter_ptr, int batch, int input_rows, int input_cols, - int depth, int filter_rows, int filter_cols, int output_rows, - int output_cols, int stride_rows, int stride_cols, int rate_rows, - int rate_cols, int pad_top, int pad_left, T* __restrict__ output_ptr) { - GPU_1D_KERNEL_LOOP(out_idx, nthreads) { + const int64_t nthreads, const T* __restrict__ input_ptr, + const T* __restrict__ filter_ptr, int64_t batch, int64_t input_rows, + int64_t input_cols, int64_t depth, int64_t filter_rows, int64_t filter_cols, + int64_t output_rows, int64_t output_cols, int stride_rows, int stride_cols, + int rate_rows, int rate_cols, int pad_top, int pad_left, + T* __restrict__ output_ptr) { + for (int64_t out_idx : GpuGridRangeX(nthreads)) { // out_idx = d + depth * (w_out + output_cols * (h_out + output_rows * b)) - const int d = out_idx % depth; - const int out_idx2 = out_idx / depth; - const int w_out = out_idx2 % output_cols; - const int out_idx3 = out_idx2 / output_cols; - const int h_out = out_idx3 % output_rows; - const int b = out_idx3 / output_rows; - int h_beg = h_out * stride_rows - pad_top; - int w_beg = w_out * stride_cols - pad_left; + const int64_t d = out_idx % depth; + const int64_t out_idx2 = out_idx / depth; + const int64_t w_out = out_idx2 % output_cols; + const int64_t out_idx3 = out_idx2 / output_cols; + const int64_t h_out = out_idx3 % output_rows; + const int64_t b = out_idx3 / output_rows; + int64_t h_beg = h_out * stride_rows - pad_top; + int64_t w_beg = w_out * stride_cols - pad_left; T cur_val = Eigen::NumTraits::lowest(); - for (int h = 0; h < filter_rows; ++h) { - const int h_in = h_beg + h * rate_rows; + for (int64_t h = 0; h < filter_rows; ++h) { + const int64_t h_in = h_beg + h * rate_rows; if (h_in >= 0 && h_in < input_rows) { - for (int w = 0; w < filter_cols; ++w) { - const int w_in = w_beg + w * rate_cols; + for (int64_t w = 0; w < filter_cols; ++w) { + const int64_t w_in = w_beg + w * rate_cols; if (w_in >= 0 && w_in < input_cols) { const T val = input_ptr[d + depth * (w_in + @@ -75,33 +76,33 @@ __global__ void DilationKernel( template __global__ void DilationBackpropInputKernel( - const int32_t nthreads, const T* __restrict__ input_ptr, + const int64_t nthreads, const T* __restrict__ input_ptr, const T* __restrict__ filter_ptr, const T* __restrict__ out_backprop_ptr, - int batch, int input_rows, int input_cols, int depth, int filter_rows, - int filter_cols, int output_rows, int output_cols, int stride_rows, - int stride_cols, int rate_rows, int rate_cols, int pad_top, int pad_left, - T* __restrict__ in_backprop_ptr) { - GPU_1D_KERNEL_LOOP(out_idx, nthreads) { + int64_t batch, int64_t input_rows, int64_t input_cols, int64_t depth, + int64_t filter_rows, int64_t filter_cols, int64_t output_rows, + int64_t output_cols, int stride_rows, int stride_cols, int rate_rows, + int rate_cols, int pad_top, int pad_left, T* __restrict__ in_backprop_ptr) { + for (int64_t out_idx : GpuGridRangeX(nthreads)) { // out_idx = d + depth * (w_out + output_cols * (h_out + output_rows * b)) - const int d = out_idx % depth; - const int out_idx2 = out_idx / depth; - const int w_out = out_idx2 % output_cols; - const int out_idx3 = out_idx2 / output_cols; - const int h_out = out_idx3 % output_rows; - const int b = out_idx3 / output_rows; - int h_beg = h_out * stride_rows - pad_top; - int w_beg = w_out * stride_cols - pad_left; + const int64_t d = out_idx % depth; + const int64_t out_idx2 = out_idx / depth; + const int64_t w_out = out_idx2 % output_cols; + const int64_t out_idx3 = out_idx2 / output_cols; + const int64_t h_out = out_idx3 % output_rows; + const int64_t b = out_idx3 / output_rows; + int64_t h_beg = h_out * stride_rows - pad_top; + int64_t w_beg = w_out * stride_cols - pad_left; T cur_val = Eigen::NumTraits::lowest(); - int h_in_max = (h_beg < 0) ? 0 : h_beg; - int w_in_max = (w_beg < 0) ? 0 : w_beg; + int64_t h_in_max = (h_beg < 0) ? 0 : h_beg; + int64_t w_in_max = (w_beg < 0) ? 0 : w_beg; // In the case of multiple argmax branches, we only back-propagate along the // last branch, i.e., the one with largest value of `h * filter_cols + w`, // similarly to the max-pooling backward routines. - for (int h = 0; h < filter_rows; ++h) { - const int h_in = h_beg + h * rate_rows; + for (int64_t h = 0; h < filter_rows; ++h) { + const int64_t h_in = h_beg + h * rate_rows; if (h_in >= 0 && h_in < input_rows) { - for (int w = 0; w < filter_cols; ++w) { - const int w_in = w_beg + w * rate_cols; + for (int64_t w = 0; w < filter_cols; ++w) { + const int64_t w_in = w_beg + w * rate_cols; if (w_in >= 0 && w_in < input_cols) { const T val = input_ptr[d + depth * (w_in + @@ -125,33 +126,34 @@ __global__ void DilationBackpropInputKernel( template __global__ void DilationBackpropFilterKernel( - const int32_t nthreads, const T* __restrict__ input_ptr, + const int64_t nthreads, const T* __restrict__ input_ptr, const T* __restrict__ filter_ptr, const T* __restrict__ out_backprop_ptr, - int batch, int input_rows, int input_cols, int depth, int filter_rows, - int filter_cols, int output_rows, int output_cols, int stride_rows, - int stride_cols, int rate_rows, int rate_cols, int pad_top, int pad_left, + int64_t batch, int64_t input_rows, int64_t input_cols, int64_t depth, + int64_t filter_rows, int64_t filter_cols, int64_t output_rows, + int64_t output_cols, int stride_rows, int stride_cols, int rate_rows, + int rate_cols, int pad_top, int pad_left, T* __restrict__ filter_backprop_ptr) { - GPU_1D_KERNEL_LOOP(out_idx, nthreads) { + for (int64_t out_idx : GpuGridRangeX(nthreads)) { // out_idx = d + depth * (w_out + output_cols * (h_out + output_rows * b)) - const int d = out_idx % depth; - const int out_idx2 = out_idx / depth; - const int w_out = out_idx2 % output_cols; - const int out_idx3 = out_idx2 / output_cols; - const int h_out = out_idx3 % output_rows; - const int b = out_idx3 / output_rows; - int h_beg = h_out * stride_rows - pad_top; - int w_beg = w_out * stride_cols - pad_left; + const int64_t d = out_idx % depth; + const int64_t out_idx2 = out_idx / depth; + const int64_t w_out = out_idx2 % output_cols; + const int64_t out_idx3 = out_idx2 / output_cols; + const int64_t h_out = out_idx3 % output_rows; + const int64_t b = out_idx3 / output_rows; + int64_t h_beg = h_out * stride_rows - pad_top; + int64_t w_beg = w_out * stride_cols - pad_left; T cur_val = Eigen::NumTraits::lowest(); - int h_max = 0; - int w_max = 0; + int64_t h_max = 0; + int64_t w_max = 0; // In the case of multiple argmax branches, we only back-propagate along the // last branch, i.e., the one with largest value of `h * filter_cols + w`, // similarly to the max-pooling backward routines. - for (int h = 0; h < filter_rows; ++h) { - const int h_in = h_beg + h * rate_rows; + for (int64_t h = 0; h < filter_rows; ++h) { + const int64_t h_in = h_beg + h * rate_rows; if (h_in >= 0 && h_in < input_rows) { - for (int w = 0; w < filter_cols; ++w) { - const int w_in = w_beg + w * rate_cols; + for (int64_t w = 0; w < filter_cols; ++w) { + const int64_t w_in = w_beg + w * rate_cols; if (w_in >= 0 && w_in < input_cols) { const T val = input_ptr[d + depth * (w_in + @@ -182,23 +184,26 @@ struct Dilation { typename TTypes::ConstTensor filter, int stride_rows, int stride_cols, int rate_rows, int rate_cols, int pad_top, int pad_left, typename TTypes::Tensor output) { - const int batch = input.dimension(0); - const int input_rows = input.dimension(1); - const int input_cols = input.dimension(2); - const int depth = input.dimension(3); - - const int filter_rows = filter.dimension(0); - const int filter_cols = filter.dimension(1); - - const int output_rows = output.dimension(1); - const int output_cols = output.dimension(2); - - const int total_count = batch * output_rows * output_cols * depth; - GpuLaunchConfig config = GetGpuLaunchConfig(total_count, d); - + const int64_t batch = input.dimension(0); + const int64_t input_rows = input.dimension(1); + const int64_t input_cols = input.dimension(2); + const int64_t depth = input.dimension(3); + + const int64_t filter_rows = filter.dimension(0); + const int64_t filter_cols = filter.dimension(1); + + const int64_t output_rows = output.dimension(1); + const int64_t output_cols = output.dimension(2); + + const int64_t total_count = batch * output_rows * output_cols * depth; + // Occupancy-aware config: heuristic 1024-thread blocks exceed registers + // on cuda13 for these kernels (too many resources requested for launch). + absl::StatusOr config = + GetGpuLaunchConfig64(total_count, d, DilationKernel, 0, 0); + CHECK_OK(config.status()); // Crash OK TF_CHECK_OK(GpuLaunchKernel( - DilationKernel, config.block_count, config.thread_per_block, 0, - d.stream(), config.virtual_thread_count, input.data(), filter.data(), + DilationKernel, config->block_count, config->thread_per_block, 0, + d.stream(), config->virtual_thread_count, input.data(), filter.data(), batch, input_rows, input_cols, depth, filter_rows, filter_cols, output_rows, output_cols, stride_rows, stride_cols, rate_rows, rate_cols, pad_top, pad_left, output.data())); @@ -213,33 +218,36 @@ struct DilationBackpropInput { int stride_rows, int stride_cols, int rate_rows, int rate_cols, int pad_top, int pad_left, typename TTypes::Tensor in_backprop) { - const int batch = input.dimension(0); - const int input_rows = input.dimension(1); - const int input_cols = input.dimension(2); - const int depth = input.dimension(3); + const int64_t batch = input.dimension(0); + const int64_t input_rows = input.dimension(1); + const int64_t input_cols = input.dimension(2); + const int64_t depth = input.dimension(3); - const int filter_rows = filter.dimension(0); - const int filter_cols = filter.dimension(1); + const int64_t filter_rows = filter.dimension(0); + const int64_t filter_cols = filter.dimension(1); - const int output_rows = out_backprop.dimension(1); - const int output_cols = out_backprop.dimension(2); + const int64_t output_rows = out_backprop.dimension(1); + const int64_t output_cols = out_backprop.dimension(2); - int total_count; - GpuLaunchConfig config; + int64_t total_count; + absl::StatusOr config; // Initialize in_backprop with all zeros. total_count = batch * input_rows * input_cols * depth; - config = GetGpuLaunchConfig(total_count, d); - TF_CHECK_OK(GpuLaunchKernel(SetZero, config.block_count, - config.thread_per_block, 0, d.stream(), - total_count, in_backprop.data())); + config = GetGpuLaunchConfig64(total_count, d); + CHECK_OK(config.status()); // Crash OK + TF_CHECK_OK(GpuLaunchKernel( + SetZero, config->block_count, config->thread_per_block, 0, + d.stream(), config->virtual_thread_count, in_backprop.data())); // Accumulate. total_count = batch * output_rows * output_cols * depth; - config = GetGpuLaunchConfig(total_count, d); + config = GetGpuLaunchConfig64(total_count, d, + DilationBackpropInputKernel, 0, 0); + CHECK_OK(config.status()); // Crash OK TF_CHECK_OK(GpuLaunchKernel( - DilationBackpropInputKernel, config.block_count, - config.thread_per_block, 0, d.stream(), config.virtual_thread_count, + DilationBackpropInputKernel, config->block_count, + config->thread_per_block, 0, d.stream(), config->virtual_thread_count, input.data(), filter.data(), out_backprop.data(), batch, input_rows, input_cols, depth, filter_rows, filter_cols, output_rows, output_cols, stride_rows, stride_cols, rate_rows, rate_cols, pad_top, pad_left, @@ -255,33 +263,36 @@ struct DilationBackpropFilter { int stride_rows, int stride_cols, int rate_rows, int rate_cols, int pad_top, int pad_left, typename TTypes::Tensor filter_backprop) { - const int batch = input.dimension(0); - const int input_rows = input.dimension(1); - const int input_cols = input.dimension(2); - const int depth = input.dimension(3); + const int64_t batch = input.dimension(0); + const int64_t input_rows = input.dimension(1); + const int64_t input_cols = input.dimension(2); + const int64_t depth = input.dimension(3); - const int filter_rows = filter.dimension(0); - const int filter_cols = filter.dimension(1); + const int64_t filter_rows = filter.dimension(0); + const int64_t filter_cols = filter.dimension(1); - const int output_rows = out_backprop.dimension(1); - const int output_cols = out_backprop.dimension(2); + const int64_t output_rows = out_backprop.dimension(1); + const int64_t output_cols = out_backprop.dimension(2); - int total_count; - GpuLaunchConfig config; + int64_t total_count; + absl::StatusOr config; // Initialize filter_backprop with all zeros. total_count = filter_rows * filter_cols * depth; - config = GetGpuLaunchConfig(total_count, d); - TF_CHECK_OK(GpuLaunchKernel(SetZero, config.block_count, - config.thread_per_block, 0, d.stream(), - total_count, filter_backprop.data())); + config = GetGpuLaunchConfig64(total_count, d); + CHECK_OK(config.status()); // Crash OK + TF_CHECK_OK(GpuLaunchKernel( + SetZero, config->block_count, config->thread_per_block, 0, + d.stream(), config->virtual_thread_count, filter_backprop.data())); // Accumulate. total_count = batch * output_rows * output_cols * depth; - config = GetGpuLaunchConfig(total_count, d); + config = GetGpuLaunchConfig64(total_count, d, + DilationBackpropFilterKernel, 0, 0); + CHECK_OK(config.status()); // Crash OK TF_CHECK_OK(GpuLaunchKernel( - DilationBackpropFilterKernel, config.block_count, - config.thread_per_block, 0, d.stream(), config.virtual_thread_count, + DilationBackpropFilterKernel, config->block_count, + config->thread_per_block, 0, d.stream(), config->virtual_thread_count, input.data(), filter.data(), out_backprop.data(), batch, input_rows, input_cols, depth, filter_rows, filter_cols, output_rows, output_cols, stride_rows, stride_cols, rate_rows, rate_cols, pad_top, pad_left, diff --git a/tensorflow/core/kernels/fused_eigen_output_kernels.h b/tensorflow/core/kernels/fused_eigen_output_kernels.h index 7cc19a9a67ab92..570d0095a21d15 100644 --- a/tensorflow/core/kernels/fused_eigen_output_kernels.h +++ b/tensorflow/core/kernels/fused_eigen_output_kernels.h @@ -411,12 +411,23 @@ using WithFusedBatchNormAndLeakyRelu = FusedBatchNormOutputKernel; template absl::Status InitBiasAddArgs(OpKernelContext* context, BiasAddArgs* args, const float* leakyrelu_alpha = nullptr) { - // Bias of the following dimensions: [ output_depth ] + // Bias of the following dimensions: [ output_depth ] or [1, ..., 1, + // output_depth] const Tensor& bias = context->input(2); - if (bias.dims() != 1) + if (bias.dims() < 1) { return absl::InvalidArgumentError( - absl::StrCat("bias must be 1-dimensional", bias.shape().DebugString())); + absl::StrCat("bias must be at least 1-dimensional, got: ", + bias.shape().DebugString())); + } + for (int i = 0; i < bias.dims() - 1; ++i) { + if (bias.dim_size(i) != 1) { + return absl::InvalidArgumentError( + absl::StrCat("For bias_dims > 1, all except the last dimension " + "must be 1, got: ", + bias.shape().DebugString())); + } + } const auto data_ptr = [](const Tensor& tensor) -> const T* { return reinterpret_cast(tensor.tensor_data().data()); diff --git a/tensorflow/core/kernels/gpu_prim.h b/tensorflow/core/kernels/gpu_prim.h index 7f27caa9b825f2..52954116f0218a 100644 --- a/tensorflow/core/kernels/gpu_prim.h +++ b/tensorflow/core/kernels/gpu_prim.h @@ -1,15 +1,16 @@ /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 -To in writing unless required by applicable law or agreed, -distributed on an, software distributed under the license is "AS IS" -BASIS, WITHOUT OF ANY KIND WARRANTIES OR CONDITIONS, either express -or implied. For the specific language governing permissions and -limitations under the license, the license you must see. +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 TENSORFLOW_CORE_KERNELS_GPU_PRIM_H_ #define TENSORFLOW_CORE_KERNELS_GPU_PRIM_H_ diff --git a/tensorflow/core/kernels/image/non_max_suppression_op.cu.cc b/tensorflow/core/kernels/image/non_max_suppression_op.cu.cc index 9b734a2cdfa39f..2b666090b77879 100644 --- a/tensorflow/core/kernels/image/non_max_suppression_op.cu.cc +++ b/tensorflow/core/kernels/image/non_max_suppression_op.cu.cc @@ -484,6 +484,9 @@ absl::Status CheckValidInputs(const Tensor& boxes, const Tensor& scores, max_output_size.shape().DebugString(), " (Shape must be rank 0 but is ", "rank ", max_output_size.dims(), ")")); } + if (max_output_size.scalar()() < 0) { + return absl::InvalidArgumentError("max_output_size must be non-negative"); + } if (!TensorShapeUtils::IsScalar(iou_threshold.shape())) { return absl::InvalidArgumentError(absl::StrCat( "iou_threshold must be 0-D, got shape ", diff --git a/tensorflow/core/kernels/image/non_max_suppression_op_gpu_test.cc b/tensorflow/core/kernels/image/non_max_suppression_op_gpu_test.cc index 76b2a9670d6460..08bd4bea40185f 100644 --- a/tensorflow/core/kernels/image/non_max_suppression_op_gpu_test.cc +++ b/tensorflow/core/kernels/image/non_max_suppression_op_gpu_test.cc @@ -213,6 +213,18 @@ TEST_F(NonMaxSuppressionV2GPUOpTest, TestInvalidIOUThreshold) { << s; } +TEST_F(NonMaxSuppressionV2GPUOpTest, TestInvalidMaxOutputSize) { + MakeOp(); + AddInputFromArray(TensorShape({1, 4}), {0, 0, 1, 1}); + AddInputFromArray(TensorShape({1}), {.9f}); + AddInputFromArray(TensorShape({}), {-1}); + AddInputFromArray(TensorShape({}), {.5f}); + + absl::Status s = RunOpKernel(); + ASSERT_FALSE(s.ok()); + EXPECT_TRUE(absl::StrContains(s.message(), "must be non-negative")) << s; +} + TEST_F(NonMaxSuppressionV2GPUOpTest, TestEmptyInput) { MakeOp(); AddInputFromArray(TensorShape({0, 4}), {}); @@ -446,6 +458,19 @@ TEST_F(NonMaxSuppressionV3GPUOpTest, TestInvalidIOUThreshold) { << s; } +TEST_F(NonMaxSuppressionV3GPUOpTest, TestInvalidMaxOutputSize) { + MakeOp(); + AddInputFromArray(TensorShape({1, 4}), {0, 0, 1, 1}); + AddInputFromArray(TensorShape({1}), {.9f}); + AddInputFromArray(TensorShape({}), {-1}); + AddInputFromArray(TensorShape({}), {.5f}); + AddInputFromArray(TensorShape({}), {0.0f}); + + absl::Status s = RunOpKernel(); + ASSERT_FALSE(s.ok()); + EXPECT_TRUE(absl::StrContains(s.message(), "must be non-negative")) << s; +} + TEST_F(NonMaxSuppressionV3GPUOpTest, TestEmptyInput) { MakeOp(); AddInputFromArray(TensorShape({0, 4}), {}); @@ -539,6 +564,19 @@ TEST_F(NonMaxSuppressionV4GPUOpTest, DifferentInputAndThresholdTypesWorks) { test::ExpectTensorEqual(expected_num_valid, *GetOutput(1)); } +TEST_F(NonMaxSuppressionV4GPUOpTest, TestInvalidMaxOutputSize) { + MakeOp(); + AddInputFromArray(TensorShape({1, 4}), {0, 0, 1, 1}); + AddInputFromArray(TensorShape({1}), {.9f}); + AddInputFromArray(TensorShape({}), {-1}); + AddInputFromArray(TensorShape({}), {.5f}); + AddInputFromArray(TensorShape({}), {0.0f}); + + absl::Status s = RunOpKernel(); + ASSERT_FALSE(s.ok()); + EXPECT_TRUE(absl::StrContains(s.message(), "must be non-negative")) << s; +} + #endif } // namespace diff --git a/tensorflow/core/kernels/matmul_op_fused.cc b/tensorflow/core/kernels/matmul_op_fused.cc index 25e4959aff8136..85ca1a00cf6ca4 100644 --- a/tensorflow/core/kernels/matmul_op_fused.cc +++ b/tensorflow/core/kernels/matmul_op_fused.cc @@ -453,10 +453,19 @@ struct LaunchFusedMatMulOp { // + . Therefore, the bias tensor is required. const Tensor& bias = context->input(2); - if (bias.dims() != 1) { - OP_REQUIRES_OK(context, absl::InvalidArgumentError( - absl::StrCat("bias must be 1-dimensional", - bias.shape().DebugString()))); + if (bias.dims() < 1) { + OP_REQUIRES_OK(context, absl::InvalidArgumentError(absl::StrCat( + "bias must be at least 1-dimensional, got: ", + bias.shape().DebugString()))); + } + for (int i = 0; i < bias.dims() - 1; ++i) { + if (bias.dim_size(i) != 1) { + OP_REQUIRES_OK(context, + absl::InvalidArgumentError(absl::StrCat( + "For bias_dims > 1, all except the last dimension " + "must be 1, got: ", + bias.shape().DebugString()))); + } } auto a_ptr = AsDeviceMemory(a.template flat().data(), diff --git a/tensorflow/core/kernels/matmul_op_test.cc b/tensorflow/core/kernels/matmul_op_test.cc index 51b8fc71e1a1e6..cf43fffeef3c11 100644 --- a/tensorflow/core/kernels/matmul_op_test.cc +++ b/tensorflow/core/kernels/matmul_op_test.cc @@ -288,6 +288,25 @@ class FusedMatMulOpTest : public OpsTestBase { run_fused); } + void VerifyMatMulWithBroadcastBias(int m, int k, int n, bool transpose_a, + bool transpose_b) { + DataType dtype = DataTypeToEnum::v(); + Tensor lhs(dtype, {transpose_a ? k : m, transpose_a ? m : k}); + lhs.flat() = lhs.flat().setRandom(); + Tensor rhs(dtype, {transpose_b ? n : k, transpose_b ? k : n}); + rhs.flat() = rhs.flat().setRandom(); + Tensor bias(dtype, {1, n}); + bias.flat() = bias.flat().setRandom(); + + Tensor fused_matmul; + bool skipped = false; + RunFusedMatMulOp(lhs, rhs, {bias}, {"BiasAdd"}, transpose_a, transpose_b, + &fused_matmul, /*allow_gpu_device=*/true, &skipped); + if (!skipped) { + ASSERT_EQ(fused_matmul.shape(), TensorShape({m, n})); + } + } + // Verifies that computing MatMul+BiasAdd+{Activation} in a graph is identical // to FusedMatMul. void VerifyConv2DWithBiasAndActivation(int m, int k, int n, bool transpose_a, @@ -356,6 +375,10 @@ TYPED_TEST_P(FusedMatMulWithBiasOpTest, MatMul1x256x1) { this->VerifyMatMulWithBias(1, 256, 1, false, false); } +TYPED_TEST_P(FusedMatMulWithBiasOpTest, MatMulWithBroadcastBias) { + this->VerifyMatMulWithBroadcastBias(8, 16, 32, false, false); +} + static auto GetActivations(DataType dtype) { // "GeluExact", "Tanh", "Sigmoid" fusions are only supported for half-float // datatype @@ -407,6 +430,7 @@ REGISTER_TYPED_TEST_SUITE_P(FusedMatMulWithBiasOpTest, // MatMul1x256x256, // MatMul256x256x1, // MatMul1x256x1, // + MatMulWithBroadcastBias, // MatMul256x128x64WithActivation, // MatMul1x256x256WithActivation, // MatMul256x256x1WithActivation, // diff --git a/tensorflow/core/kernels/maxpooling_op.cc b/tensorflow/core/kernels/maxpooling_op.cc index 6f7d68cfdec580..e5f94a0ff9bf8f 100644 --- a/tensorflow/core/kernels/maxpooling_op.cc +++ b/tensorflow/core/kernels/maxpooling_op.cc @@ -1056,11 +1056,11 @@ class MaxPoolingWithArgmaxOp : public OpKernel { bool include_batch_in_index_; }; -template +template struct LaunchMaxPoolingGradWithArgmax; -template -struct LaunchMaxPoolingGradWithArgmax { +template +struct LaunchMaxPoolingGradWithArgmax { typedef Eigen::Map> EigenMatrixMap; @@ -1080,7 +1080,7 @@ struct LaunchMaxPoolingGradWithArgmax { { auto grad_out_flat = grad_out->flat(); - auto argmax_flat = argmax.flat(); + auto argmax_flat = argmax.flat(); auto grad_in_flat = grad_in.flat(); const int64_t output_start = start * output_size_per_batch; @@ -1116,8 +1116,7 @@ struct LaunchMaxPoolingGradWithArgmax { } }; -// TODO(b/175733711): Support int32 argmax type in MaxPoolGradWithArgmax op. -template +template class MaxPoolingGradWithArgmaxOp : public OpKernel { public: explicit MaxPoolingGradWithArgmaxOp(OpKernelConstruction* context) @@ -1195,7 +1194,7 @@ class MaxPoolingGradWithArgmaxOp : public OpKernel { if (out_shape.num_elements() == 0) return; // nothing to be done - LaunchMaxPoolingGradWithArgmax::launch( + LaunchMaxPoolingGradWithArgmax::launch( context, params, grad_in, argmax, grad_out, include_batch_in_index_); } @@ -1529,7 +1528,7 @@ struct LaunchMaxPoolingWithArgmax { }; template -struct LaunchMaxPoolingGradWithArgmax { +struct LaunchMaxPoolingGradWithArgmax { static void launch(OpKernelContext* context, const PoolParameters& params, const Tensor& grad_in, const Tensor& argmax, Tensor* grad_out, const bool include_batch_in_index) { @@ -1605,21 +1604,26 @@ struct LaunchMaxPoolingGradGradWithArgmax { .Device(DEVICE_##D) \ .TypeConstraint("T") \ .TypeConstraint("Targmax"), \ - MaxPoolingGradWithArgmaxOp); + MaxPoolingGradWithArgmaxOp); // Below kernels implemented only for CPU device. -#define REGISTER_CPU_ONLY_POOL_KERNELS(T) \ - REGISTER_KERNEL_BUILDER( \ - Name("MaxPool").Device(DEVICE_CPU).TypeConstraint("T"), \ - MaxPoolingOp); \ - REGISTER_KERNEL_BUILDER( \ - Name("MaxPoolV2").Device(DEVICE_CPU).TypeConstraint("T"), \ - MaxPoolingV2Op); \ - REGISTER_KERNEL_BUILDER(Name("MaxPoolWithArgmax") \ - .Device(DEVICE_CPU) \ - .TypeConstraint("Targmax") \ - .TypeConstraint("T"), \ - MaxPoolingWithArgmaxOp); +#define REGISTER_CPU_ONLY_POOL_KERNELS(T) \ + REGISTER_KERNEL_BUILDER( \ + Name("MaxPool").Device(DEVICE_CPU).TypeConstraint("T"), \ + MaxPoolingOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("MaxPoolV2").Device(DEVICE_CPU).TypeConstraint("T"), \ + MaxPoolingV2Op); \ + REGISTER_KERNEL_BUILDER(Name("MaxPoolWithArgmax") \ + .Device(DEVICE_CPU) \ + .TypeConstraint("Targmax") \ + .TypeConstraint("T"), \ + MaxPoolingWithArgmaxOp); \ + REGISTER_KERNEL_BUILDER(Name("MaxPoolGradWithArgmax") \ + .Device(DEVICE_CPU) \ + .TypeConstraint("T") \ + .TypeConstraint("Targmax"), \ + MaxPoolingGradWithArgmaxOp); TF_CALL_REAL_NUMBER_TYPES(REGISTER_CPU_ONLY_POOL_KERNELS); #undef REGISTER_CPU_ONLY_POOL_KERNELS diff --git a/tensorflow/core/kernels/pack_op.cc b/tensorflow/core/kernels/pack_op.cc index 95331ce42dc481..a43723bd438cf6 100644 --- a/tensorflow/core/kernels/pack_op.cc +++ b/tensorflow/core/kernels/pack_op.cc @@ -28,6 +28,20 @@ limitations under the License. #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM +// Forward declaration of the GPU scalar-pack helper defined in +// pack_op_gpu.cu.cc (compiled with EIGEN_USE_GPU by NVCC/hipcc). +// Calling into that translation unit avoids instantiating Eigen +// GPU-device expressions inside a regular .cc file where +// EIGEN_USE_GPU is not defined, which would trigger Eigen's +// static assertion: +// "Default executor instantiated with non-default device." +namespace tensorflow { +template +void PackScalarsOnGPU(OpKernelContext* c, int num, Tensor* output); +} // namespace tensorflow +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM + namespace tensorflow { typedef Eigen::ThreadPoolDevice CPUDevice; @@ -80,6 +94,26 @@ class PackOp : public OpKernel { Tensor* output; OP_REQUIRES_OK(c, c->allocate_output(0, output_shape, &output)); +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + // Special case: packing 0-D (scalar) inputs on GPU. + // + // The generic path below flattens inputs to {before_dim, after_dim} + // matrices and calls ConcatGPU. For scalars before_dim == after_dim == 1, + // giving {1,1} input matrices and a {1,num} output matrix. Some GPU + // drivers reject the resulting launch configuration, causing a + // "Can't concatenate scalars" error at runtime. Dispatch to + // PackScalarsOnGPU (defined in pack_op_gpu.cu.cc, compiled with + // EIGEN_USE_GPU) which writes each scalar element directly into the + // output flat vector. The indirection through a separate .cu.cc TU is + // required: Eigen device() assignments must be compiled with + // EIGEN_USE_GPU defined (i.e. by NVCC/hipcc), otherwise Eigen fires a + // static assertion even inside discarded if-constexpr branches. + if (std::is_same::value && first_input.dims() == 0) { + PackScalarsOnGPU(c, num, output); + return; + } +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM + int64_t before_dim = 1; for (int i = 0; i < axis; ++i) { before_dim *= output_shape.dim_size(i); diff --git a/tensorflow/core/kernels/pack_op_gpu.cu.cc b/tensorflow/core/kernels/pack_op_gpu.cu.cc new file mode 100644 index 00000000000000..8b05253a20a02f --- /dev/null +++ b/tensorflow/core/kernels/pack_op_gpu.cu.cc @@ -0,0 +1,80 @@ +/* Copyright 2015 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// GPU-side implementation of the PackScalarsOnGPU helper for PackOp. +// +// This file is compiled by NVCC / hipcc (i.e. as a .cu.cc source) which +// defines EIGEN_USE_GPU before including Eigen headers. That is required for +// the Eigen device() assignment used below; doing the same in the plain .cc +// file (where EIGEN_USE_GPU is NOT defined) would trigger: +// +// static assertion failed: "Default executor instantiated with non-default +// device. You must #define EIGEN_USE_THREADS, EIGEN_USE_GPU or +// EIGEN_USE_SYCL before including Eigen headers." +// +// pack_op.cc forward-declares PackScalarsOnGPU and calls it for the GPU +// device path; the linker resolves the call to the explicit instantiations +// at the bottom of this file. + +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + +#define EIGEN_USE_GPU + +#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/register_types.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_types.h" + +namespace tensorflow { + +// Packs `num` 0-D scalar input tensors (c->input(0) … c->input(num-1)) into +// a pre-allocated 1-D output tensor `output` of shape [num]. +// Each scalar element i is written to output[i] via an Eigen GPU chip +// assignment so that the operation is correctly dispatched to the GPU stream. +template +void PackScalarsOnGPU(OpKernelContext* c, int num, Tensor* output) { + const Tensor& first_input = c->input(0); + auto output_vec = output->flat(); + const Eigen::GpuDevice& d = c->eigen_device(); + for (int i = 0; i < num; ++i) { + const Tensor& input = c->input(i); + OP_REQUIRES(c, first_input.shape().IsSameSize(input.shape()), + absl::InvalidArgumentError(absl::StrCat( + "Shapes of all inputs must match: values[0].shape = ", + first_input.shape().DebugString(), " != values[", i, + "].shape = ", input.shape().DebugString()))); + output_vec.template chip<0>(i).device(d) = input.scalar(); + } +} + +// Explicit instantiations for all types registered with REGISTER_GPU in +// pack_op.cc. Must match TF_CALL_GPU_ALL_TYPES + the extra integer types. +#define INSTANTIATE(T) \ + template void PackScalarsOnGPU(OpKernelContext*, int, Tensor*); + +TF_CALL_int64(INSTANTIATE); +TF_CALL_int16(INSTANTIATE); +TF_CALL_uint32(INSTANTIATE); +TF_CALL_uint64(INSTANTIATE); +TF_CALL_GPU_ALL_TYPES(INSTANTIATE); +TF_CALL_float8_e5m2(INSTANTIATE); +TF_CALL_float8_e4m3fn(INSTANTIATE); + +#undef INSTANTIATE + +} // namespace tensorflow + +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/tensorflow/core/kernels/ragged_gather_op.cc b/tensorflow/core/kernels/ragged_gather_op.cc index d280c948303179..72e5c840e86ad5 100644 --- a/tensorflow/core/kernels/ragged_gather_op.cc +++ b/tensorflow/core/kernels/ragged_gather_op.cc @@ -128,6 +128,7 @@ class RaggedGatherOpBase : public OpKernel { std::vector>* value_slices, SPLITS_TYPE* num_values) { *num_values = 0; + int64_t total_values = 0; value_slices->clear(); int num_splits = indices_in.dims() - 1 + params_nested_splits_in.size(); @@ -188,9 +189,17 @@ class RaggedGatherOpBase : public OpKernel { } if (limit != start) { value_slices->emplace_back(start, limit); - *num_values += limit - start; + total_values += limit - start; } } + // The value count is accumulated in a wide type and validated before it is + // narrowed to SPLITS_TYPE, so an int32 Tsplits cannot silently wrap and + // drive an undersized output allocation in WriteValues. + if (total_values > std::numeric_limits::max()) { + return absl::InvalidArgumentError(absl::StrCat( + "Number of values ", total_values, " exceeds limits of Tsplits")); + } + *num_values = static_cast(total_values); return absl::OkStatus(); } diff --git a/tensorflow/core/kernels/random_index_shuffle_ops.cc b/tensorflow/core/kernels/random_index_shuffle_ops.cc index f817f69b7d4bd0..515d08679cc350 100644 --- a/tensorflow/core/kernels/random_index_shuffle_ops.cc +++ b/tensorflow/core/kernels/random_index_shuffle_ops.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "tensorflow/core/framework/bounds_check.h" @@ -68,6 +69,11 @@ class RandomIndexShuffleOp : public OpKernel { explicit RandomIndexShuffleOp(OpKernelConstruction* context) : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr(kRounds, &rounds_)); + OP_REQUIRES( + context, rounds_ >= 4 && rounds_ % 2 == 0, + absl::InvalidArgumentError(absl::StrCat( + "rounds must be an even integer greater than or equal to 4, but ", + "got ", rounds_))); } void Compute(OpKernelContext* context) override { diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc index 538d0b9b0c8d2e..6b8f50774a0cf1 100644 --- a/tensorflow/core/kernels/resource_variable_ops.cc +++ b/tensorflow/core/kernels/resource_variable_ops.cc @@ -773,14 +773,21 @@ class ResourceGatherOp : public OpKernel { " (resource variable dtype)"))); const Tensor& params = *v->tensor(); const Tensor& indices = c->input(1); + OP_REQUIRES( c, TensorShapeUtils::IsVectorOrHigher(params.shape()), absl::InvalidArgumentError("params must be at least 1 dimensional")); - OP_REQUIRES(c, params.shape().dims() >= batch_dims_, + + OP_REQUIRES( + c, batch_dims_ < params.dims(), + absl::InvalidArgumentError(absl::StrCat( + "batch_dims (", batch_dims_, ") must be less than rank(params) (", + params.dims(), ")."))); + + OP_REQUIRES(c, batch_dims_ <= indices.dims(), absl::InvalidArgumentError( - absl::StrCat("params must have at least ", batch_dims_, - " (batch_dims) dimensions but it has shape ", - params.shape().DebugString()))); + absl::StrCat("Expected batch_dims <= rank(indices) (", + indices.dims(), "), but got ", batch_dims_))); // Check that we have enough index space const int64_t N = indices.NumElements(); diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index e358f4521598b1..79c5e531786ee5 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -48,6 +48,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/framework/tensor_util.h" #include "tensorflow/core/framework/types.h" +#include "tensorflow/core/kernels/fill_functor.h" #include "tensorflow/core/kernels/segment_reduction_ops.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/bfloat16.h" @@ -1342,7 +1343,17 @@ class SparseSegmentGradOpBase : public OpKernel { OP_REQUIRES_OK(context, output_shape.SetDimWithStatus(0, M)); Tensor* output = nullptr; OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); - if (M == 0 || N == 0) return; + if (M == 0 || N == 0) { + // `allocate_output` does not initialize the buffer, and the functor below + // only runs when N > 0. Without this, an empty `indices` combined with + // `output_dim0` > 0 returns the freshly allocated output uninitialized. + // Zero is the correct gradient for a row that received no contribution. + if (output->NumElements() > 0) { + functor::SetZeroFunctor()(context->eigen_device(), + output->flat()); + } + return; + } OP_REQUIRES(context, input.dim_size(0) > 0, absl::InvalidArgumentError("Invalid number of segments")); diff --git a/tensorflow/core/kernels/summary_kernels.cc b/tensorflow/core/kernels/summary_kernels.cc index f960a892040345..5f4e913b7bbe11 100644 --- a/tensorflow/core/kernels/summary_kernels.cc +++ b/tensorflow/core/kernels/summary_kernels.cc @@ -241,6 +241,10 @@ class WriteScalarSummaryOp : public OpKernel { const Tensor* t; OP_REQUIRES_OK(ctx, ctx->input("value", &t)); + OP_REQUIRES(ctx, t->NumElements() == 1, + errors::InvalidArgument( + "value must contain exactly one element, got shape ", + t->shape().DebugString())); OP_REQUIRES_OK(ctx, s->WriteScalar(step, *t, tag)); } diff --git a/tensorflow/core/kernels/training_ops.cc b/tensorflow/core/kernels/training_ops.cc index 1def8b120223a6..ef7bd22e57feed 100644 --- a/tensorflow/core/kernels/training_ops.cc +++ b/tensorflow/core/kernels/training_ops.cc @@ -1117,6 +1117,11 @@ class ApplyAdadeltaOp : public OpKernel { absl::InvalidArgumentError(absl::StrCat( "var and grad do not have the same shape", var.shape().DebugString(), " ", grad.shape().DebugString()))); + OP_REQUIRES(ctx, var.shape().IsSameSize(accum_update.shape()), + absl::InvalidArgumentError(absl::StrCat( + "var and accum_update do not have the same shape", + var.shape().DebugString(), " ", + accum_update.shape().DebugString()))); } void DoCompute(OpKernelContext* ctx) { diff --git a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py index 70d4a758cb5059..293d0c67f4b4e4 100644 --- a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py @@ -273,6 +273,80 @@ def test_bincount_all_binary_with_weights(self, dtype): gen_math_ops.dense_bincount( input=inp, weights=np_weight, size=size, binary_output=True))) + @test_util.run_gpu_only + def test_bincount_all_binary_with_negative_input_on_gpu(self): + # Regression test for https://github.com/tensorflow/tensorflow/issues/103995 + # BincountReduceKernel only checked `bin < num_bins`, so a negative + # input value indexed before the start of the output buffer on GPU + # (an out-of-bounds write) instead of being skipped like the CPU + # kernel does. + size = 7 + inp = np.array( + [-1737321568, 32584, 7, 4, 3, 7, 7, 2, 5, 4, 1, 7, 5, 1], dtype=np.int32 + ) + expected_out = np.zeros((size,)) + for v in inp: + if 0 <= v < size: + expected_out[v] = 1 + with ops.device("/GPU:0"): + out = self.evaluate( + gen_math_ops.dense_bincount( + input=inp, weights=[], size=size, binary_output=True + ) + ) + self.assertAllEqual(expected_out, out) + + def _test_bincount_col_binary_with_negative_input( + self, num_rows, num_cols, size, dtype + ): + # Regression test for the same out-of-bounds write as above, but in + # BincountColReduceKernel / BincountColReduceSharedKernel, which had + # the identical missing lower-bound check. + np.random.seed(42) + inp = np.random.randint(0, size, (num_rows, num_cols), dtype=dtype) + inp[0, 0] = -1737321568 + inp[num_rows - 1, num_cols - 1] = -1 + expected_out = np.zeros((num_rows, size)) + for row in range(num_rows): + for v in inp[row]: + if 0 <= v < size: + expected_out[row, v] = 1 + with ops.device("/GPU:0"): + out = self.evaluate( + gen_math_ops.dense_bincount( + input=inp, weights=[], size=size, binary_output=True + ) + ) + self.assertAllEqual(expected_out, out) + + @parameterized.parameters([ + { + "dtype": np.int32, + }, + { + "dtype": np.int64, + }, + ]) + @test_util.run_gpu_only + def test_col_reduce_shared_memory_with_negative_input(self, dtype): + # num_rows * size small enough to select BincountColReduceSharedKernel, + # matching test_col_reduce_shared_memory below. + self._test_bincount_col_binary_with_negative_input(128, 27, 10, dtype) + + @parameterized.parameters([ + { + "dtype": np.int32, + }, + { + "dtype": np.int64, + }, + ]) + @test_util.run_gpu_only + def test_col_reduce_global_memory_with_negative_input(self, dtype): + # num_rows * size large enough to select BincountColReduceKernel, + # matching test_col_reduce_global_memory below. + self._test_bincount_col_binary_with_negative_input(128, 27, 1024, dtype) + def _test_bincount_col_count(self, num_rows, num_cols, size, dtype): np.random.seed(42) inp = np.random.randint(0, size, (num_rows, num_cols), dtype=dtype) diff --git a/tensorflow/python/kernel_tests/math_ops/matmul_op_test.py b/tensorflow/python/kernel_tests/math_ops/matmul_op_test.py index ae764c8ff126b9..bb7d462e141845 100644 --- a/tensorflow/python/kernel_tests/math_ops/matmul_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/matmul_op_test.py @@ -19,6 +19,7 @@ import numpy as np from tensorflow.python import tf2 +from tensorflow.python.eager import def_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -55,6 +56,23 @@ def testBatchMatMulV3MixedPrec(self): self.assertAllEqual([[5, 11], [11, 25]], c) +@test_util.with_eager_op_as_function +class MatMulAddFusionTest(test_lib.TestCase): + """Test broadcasting behavior when MatMul and Add are fused in graph execution.""" + + def testMatMulAddBroadcastInTfFunction(self): + @def_function.function + def f(x): + return math_ops.add( + math_ops.matmul(x, x, transpose_a=True), + math_ops.matmul(array_ops.ones([1, 1], dtype=x.dtype), x), + ) + + x = constant_op.constant(np.ones([1, 7], dtype=np.float32)) + res = f(x) + self.assertEqual(tuple(res.shape), (7, 7)) + + @test_util.with_eager_op_as_function class MatVecTest(test_lib.TestCase): """Simple test for matvec, which is sugar on top of matmul.""" diff --git a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py index 2f7adecacfd6bf..fcf5ae14da605d 100644 --- a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py @@ -1182,6 +1182,27 @@ def testGradientExplicitSingleOutput(self): tf_xgrad = tf_op(tf_ygrad, indices, segment_ids, output_dim0) self.assertAllClose(tf_xgrad, np_xgrad) + def testGradientEmptyIndicesWithNonZeroOutputDim(self): + ops_list = [ + math_ops.sparse_segment_sum_grad, + math_ops.sparse_segment_mean_grad, + math_ops.sparse_segment_sqrt_n_grad, + ] + indices = [] + segment_ids = [] + output_dim0 = 4 + for dtype in [ + dtypes_lib.float16, + dtypes_lib.bfloat16, + dtypes_lib.float32, + dtypes_lib.float64, + ]: + grad = constant_op.constant([], shape=[0, 3], dtype=dtype) + for tf_op in ops_list: + res = self.evaluate(tf_op(grad, indices, segment_ids, output_dim0)) + self.assertEqual(res.shape, (4, 3)) + self.assertAllEqual(res, np.zeros((4, 3), dtype=dtype.as_numpy_dtype)) + def testGradientV2Explicit(self): # Note that the GPU implem has different paths for different inner sizes. for inner_size in (1, 2, 3, 32): diff --git a/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py b/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py index 1ff2a932836371..91c180dd613422 100644 --- a/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py @@ -27,6 +27,7 @@ from tensorflow.python.framework import tensor_spec from tensorflow.python.ops import array_ops from tensorflow.python.ops import gradient_checker_v2 +from tensorflow.python.ops import math_ops from tensorflow.python.platform import test @@ -519,6 +520,66 @@ def testNullTensor(self): xt = array_ops.transpose(x, [0, 2, 1]) self.assertAllEqual(xt.shape, (1, 0, 4)) + def testConcatScalarRankDynamic(self): + """Regression test for GitHub issue #124017. + + When tf.rank() returns a 0-D scalar tensor at graph/function build time + (i.e. the static rank is unknown) and the user builds a permutation via + tf.concat([[num_dims - 1], tf.range(num_dims - 1)], axis=0), the + autopacking of [num_dims - 1] previously failed on GPU with: + "ConcatOp: Can't concatenate scalars (use tf.stack instead)" + because the GPU Pack kernel did not support 0-D scalar inputs. + + The fix adds an explicit scalar-input path in the C++ GPU PackOp kernel + (tensorflow/core/kernels/pack_op.cc) that writes each scalar element + directly into the output vector via Eigen chip operations, bypassing the + ConcatGPU helper that was incompatible with 0-D inputs. + """ + + @def_function.function( + input_signature=[ + tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32), + ] + ) + def move_channel_first(image): + """Move the last axis to the front (channels-first transform).""" + # tf.rank returns a 0-D scalar when static rank is unknown. + num_dims = array_ops.rank(image) + # Building perm with a scalar in a Python list triggers autopacking. + perm = array_ops.concat( + [[num_dims - 1], math_ops.range(num_dims - 1)], + axis=0, + ) + return array_ops.transpose(image, perm) + + # 4-D input: [batch, H, W, C] -> [C, batch, H, W] + image = constant_op.constant( + [[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]] + ) # shape (1, 2, 2, 2) + result = move_channel_first(image) + # channel dim (size 2) should now be first + self.assertAllEqual(result.shape[-3:], (1, 2, 2)) + self.assertEqual(result.shape[0], 2) + + @def_function.function( + input_signature=[ + tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32), + ] + ) + def move_channel_last(image): + """Move the first axis to the end (channels-last transform).""" + num_dims = array_ops.rank(image) + inverse_perm = array_ops.concat( + [math_ops.range(1, num_dims), [0]], + axis=0, + ) + return array_ops.transpose(image, inverse_perm) + + # Verify round-trip: channels-first -> channels-last gives original shape. + channels_first = move_channel_first(image) # (2, 1, 2, 2) + restored = move_channel_last(channels_first) # (1, 2, 2, 2) + self.assertAllEqual(restored, image) + def testScalar(self): with self.cached_session(): x = constant_op.constant(42, dtype=dtypes.float32, shape=[]) diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py index c7f707285e1e23..78a7a851619332 100644 --- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py @@ -1091,12 +1091,15 @@ def testMaxPoolingGradWithArgmax(self): tensor_input = [11.0, 12.0, 13.0, 14.0, 21.0, 22.0, 23.0, 24.0] Config = collections.namedtuple( - "Config", ["use_gpu", "include_batch_in_index", "argmax"]) + "Config", ["use_gpu", "include_batch_in_index", "argmax", "Targmax"] + ) configs = [ - Config(False, False, [0, 1, 3, 5, 0, 2, 6, 8]), - Config(False, True, [0, 1, 3, 5, 9, 11, 15, 17]), - Config(True, False, [0, 1, 3, 5, 0, 2, 6, 8]), - Config(True, True, [0, 1, 3, 5, 9, 11, 15, 17]) + Config(False, False, [0, 1, 3, 5, 0, 2, 6, 8], dtypes.int64), + Config(False, True, [0, 1, 3, 5, 9, 11, 15, 17], dtypes.int64), + Config(False, False, [0, 1, 3, 5, 0, 2, 6, 8], dtypes.int32), + Config(False, True, [0, 1, 3, 5, 9, 11, 15, 17], dtypes.int32), + Config(True, False, [0, 1, 3, 5, 0, 2, 6, 8], dtypes.int64), + Config(True, True, [0, 1, 3, 5, 9, 11, 15, 17], dtypes.int64), ] for config in configs: @@ -1104,7 +1107,8 @@ def testMaxPoolingGradWithArgmax(self): orig_in = constant_op.constant(orig_input, shape=[2, 3, 3, 1]) t = constant_op.constant(tensor_input, shape=[2, 2, 2, 1]) argmax_t = constant_op.constant( - config.argmax, shape=[2, 2, 2, 1], dtype=dtypes.int64) + config.argmax, shape=[2, 2, 2, 1], dtype=config.Targmax + ) out_op = gen_nn_ops.max_pool_grad_with_argmax( orig_in, t, diff --git a/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py b/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py index c7ab54a722b506..251e01e12b9e5b 100644 --- a/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py +++ b/tensorflow/python/kernel_tests/random/random_index_shuffle_test.py @@ -118,6 +118,18 @@ def test_index_greater_than_max_index(self): ): self.evaluate(stateless.index_shuffle(5, seed=(1, 2), max_index=4)) + @parameterized.parameters(-1, 0, 2, 3, 5) + def test_invalid_rounds(self, rounds): + with self.assertRaisesRegex( + errors.InvalidArgumentError, + 'rounds must be an even integer', + ): + self.evaluate( + gen_random_index_shuffle_ops.random_index_shuffle( + index=0, seed=[1, 2, 3], max_index=10, rounds=rounds + ) + ) + if __name__ == '__main__': test.main() diff --git a/tensorflow/python/kernel_tests/summary_ops/summary_ops_test.py b/tensorflow/python/kernel_tests/summary_ops/summary_ops_test.py index 553193f6d15c8d..c34eba1777d3ee 100644 --- a/tensorflow/python/kernel_tests/summary_ops/summary_ops_test.py +++ b/tensorflow/python/kernel_tests/summary_ops/summary_ops_test.py @@ -34,6 +34,7 @@ from tensorflow.python.framework import test_util from tensorflow.python.lib.io import tf_record from tensorflow.python.module import module +from tensorflow.python.ops import gen_summary_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import summary_ops_v2 as summary_ops from tensorflow.python.ops import variables @@ -1590,6 +1591,50 @@ def summary_op_fn(): self.exec_summary_op(summary_op_fn) +class WriteScalarSummaryRawOpsTest(test_util.TensorFlowTestCase): + + @test_util.run_in_graph_and_eager_modes + def testValidScalarAndSingleElementInputs(self): + logdir = self.get_temp_dir() + writer = summary_ops.create_file_writer_v2(logdir) + self.evaluate(writer.init()) + step = constant_op.constant(1, dtype=dtypes.int64) + tag = constant_op.constant('test_tag', dtype=dtypes.string) + + for val in [2.5, [2.5], [[2.5]]]: + t = constant_op.constant(val, dtype=dtypes.float32) + self.evaluate( + gen_summary_ops.write_scalar_summary( + writer=writer._resource, # pylint: disable=protected-access + step=step, + tag=tag, + value=t, + ) + ) + + @test_util.run_in_graph_and_eager_modes + def testInvalidEmptyAndMultiElementInputs(self): + logdir = self.get_temp_dir() + writer = summary_ops.create_file_writer_v2(logdir) + self.evaluate(writer.init()) + step = constant_op.constant(1, dtype=dtypes.int64) + tag = constant_op.constant('test_tag', dtype=dtypes.string) + + for invalid_val in [[], [1.0, 2.0], [[1.0, 2.0]]]: + t = constant_op.constant(invalid_val, dtype=dtypes.float32) + with self.assertRaisesRegex( + errors.InvalidArgumentError, 'must contain exactly one element' + ): + self.evaluate( + gen_summary_ops.write_scalar_summary( + writer=writer._resource, # pylint: disable=protected-access + step=step, + tag=tag, + value=t, + ) + ) + + def events_from_file(filepath): """Returns all events in a single event file. diff --git a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py index f6f0b9e1762e1a..93640351eb1fe7 100644 --- a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py +++ b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py @@ -1750,6 +1750,42 @@ def testGatherWithBatchDimsMatchesTensor(self, params_shape, indices_shape, self.assertAllEqual(output_shape, result.shape.as_list()) self.assertAllEqual(expected, result) + @test_util.run_in_graph_and_eager_modes + def testGatherWithBatchDimsEqualTensorRank(self): + var = resource_variable_ops.ResourceVariable( + [10, 20, 30, 40, 50], name="var0", dtype=dtypes.float32 + ) + indices = constant_op.constant([1, 3, 4], dtype=dtypes.int32) + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), "rank" + ): + _ = resource_variable_ops.resource_gather( + resource=var.handle, + indices=indices, + dtype=dtypes.float32, + batch_dims=1, + validate_indices=True, + ) + + @test_util.run_in_graph_and_eager_modes + def testGatherWithBatchDimsGreaterThanIndicesRank(self): + var = resource_variable_ops.ResourceVariable( + [[10, 20], [30, 40]], name="var0", dtype=dtypes.float32 + ) + # indices has rank 0 (scalar) + indices = constant_op.constant(1, dtype=dtypes.int32) + # batch_dims = 1 > indices.dims() = 0 + with self.assertRaisesRegex( + (errors.InvalidArgumentError, ValueError), "rank" + ): + _ = resource_variable_ops.resource_gather( + resource=var.handle, + indices=indices, + dtype=dtypes.float32, + batch_dims=1, + validate_indices=True, + ) + @parameterized.parameters([ dict(dtype=dtypes.bool), dict(dtype=dtypes.int64), diff --git a/tensorflow/python/ops/ragged/BUILD b/tensorflow/python/ops/ragged/BUILD index 2f1bb55b55d132..4c8416eb018483 100644 --- a/tensorflow/python/ops/ragged/BUILD +++ b/tensorflow/python/ops/ragged/BUILD @@ -965,6 +965,7 @@ py_test( "//tensorflow/python/framework:test_lib", "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:gradients_impl", + "//tensorflow/python/ops:ragged_array_ops_gen", "//tensorflow/python/platform:test", "//third_party/py/numpy", "@absl_py//absl/testing:parameterized", diff --git a/tensorflow/python/ops/ragged/ragged_gather_op_test.py b/tensorflow/python/ops/ragged/ragged_gather_op_test.py index b7bd1f1a9aa150..d20cc3744060e1 100644 --- a/tensorflow/python/ops/ragged/ragged_gather_op_test.py +++ b/tensorflow/python/ops/ragged/ragged_gather_op_test.py @@ -15,7 +15,6 @@ """Tests for ragged_array_ops.gather.""" from absl.testing import parameterized - import numpy as np from tensorflow.python.eager import context @@ -26,6 +25,7 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_ragged_array_ops from tensorflow.python.ops import gradients_impl from tensorflow.python.ops.ragged import ragged_factory_ops from tensorflow.python.ops.ragged import ragged_gather_ops @@ -222,6 +222,26 @@ def testOutOfBoundsError(self): r'indices\[1\] = 3 is not in \[0, 2\)'): self.evaluate(ragged_gather_ops.gather(ragged_params, ragged_indices)) + def testValueCountOverflowError(self): + # The gathered value count (num_indices * gathered_row_length) overflows an + # int32 Tsplits: 65537 * 65536 = 2**32 + 65536. Before the fix this wrapped + # to 65536, drove an undersized output allocation, and the copy loop then + # wrote past it (heap-buffer-overflow). The raw op pins Tsplits=int32; the + # high-level gather defaults to int64 splits and cannot reproduce. + with self.assertRaisesRegex( + errors.InvalidArgumentError, 'exceeds limits of Tsplits' + ): + self.evaluate( + gen_ragged_array_ops.ragged_gather( + params_nested_splits=[ + constant_op.constant([0, 65536], dtype=dtypes.int32) + ], + params_dense_values=array_ops.zeros([65536], dtype=dtypes.uint8), + indices=array_ops.zeros([65537], dtype=dtypes.int32), + OUTPUT_RAGGED_RANK=1, + ) + ) + def testUnknownIndicesRankError(self): if context.executing_eagerly(): return diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py index fe06bd08e8faa2..e45edfe9fae226 100644 --- a/tensorflow/python/training/training_ops_test.py +++ b/tensorflow/python/training/training_ops_test.py @@ -198,6 +198,63 @@ def testApplyAdagrad(self): grad = np.arange(100).astype(dtype) self._testTypesForAdagrad(x, y, lr, grad, use_gpu) + @test_util.run_v1_only( + "ApplyAdadelta op returns a ref, so it is not supported in eager mode." + ) + def testApplyAdadeltaMismatchedShape(self): + for dtype in [dtypes.float32]: + var = variable_v1.VariableV1([1.0, 2.0], dtype=dtype) + accum = variable_v1.VariableV1([1.0, 2.0], dtype=dtype) + accum_update = variable_v1.VariableV1([1.0], dtype=dtype) + lr = constant_op.constant(1.0, dtype=dtype) + rho = constant_op.constant(0.9, dtype=dtype) + epsilon = constant_op.constant(1e-6, dtype=dtype) + grad = constant_op.constant([0.1, 0.2], dtype=dtype) + with self.cached_session() as sess: + self.evaluate(variables.global_variables_initializer()) + with self.assertRaises((ValueError, errors.InvalidArgumentError)): + self.evaluate( + gen_training_ops.apply_adadelta( + var, + accum, + accum_update, + lr, + rho, + epsilon, + grad, + use_locking=False, + ) + ) + + @test_util.run_v1_only( + "ResourceApplyAdadelta op returns a ref, so it is not " + "supported in eager mode." + ) + def testResourceApplyAdadeltaMismatchedShape(self): + for dtype in [dtypes.float32]: + var = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype) + accum = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype) + accum_update = resource_variable_ops.ResourceVariable([1.0], dtype=dtype) + lr = constant_op.constant(1.0, dtype=dtype) + rho = constant_op.constant(0.9, dtype=dtype) + epsilon = constant_op.constant(1e-6, dtype=dtype) + grad = constant_op.constant([0.1, 0.2], dtype=dtype) + with self.cached_session() as sess: + self.evaluate(variables.global_variables_initializer()) + with self.assertRaises((ValueError, errors.InvalidArgumentError)): + self.evaluate( + gen_training_ops.resource_apply_adadelta( + var.handle, + accum.handle, + accum_update.handle, + lr, + rho, + epsilon, + grad, + use_locking=False, + ) + ) + @test_util.run_v1_only("ApplyFtrl op returns a ref, so it is not " "supported in eager mode.") def testApplyFtrl(self): diff --git a/third_party/xla/build_tools/ci/build.py b/third_party/xla/build_tools/ci/build.py index c4b7150b1a2981..20ea7038739ac5 100755 --- a/third_party/xla/build_tools/ci/build.py +++ b/third_party/xla/build_tools/ci/build.py @@ -246,7 +246,6 @@ def commands(self) -> List[List[str]]: ) ) cmds.append(self.bazel_command(subcommand=self.subcommand)) - cmds.append(["bazel", "analyze-profile", "profile.json.gz"]) return cmds diff --git a/third_party/xla/build_tools/ci/golden_commands.txt b/third_party/xla/build_tools/ci/golden_commands.txt index e65e3f069635d5..bffa12b8d1b742 100644 --- a/third_party/xla/build_tools/ci/golden_commands.txt +++ b/third_party/xla/build_tools/ci/golden_commands.txt @@ -1,105 +1,86 @@ # BEGIN BuildType.JAX_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters= --test_tag_filters= --config=rbe_linux_x86_64 --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test bazel test --build_tag_filters= --test_tag_filters= --config=rbe_linux_x86_64 --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test -bazel analyze-profile profile.json.gz # END BuildType.JAX_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS # BEGIN BuildType.JAX_LINUX_X86_GPU_T4_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-multiaccelerator --test_tag_filters=-multiaccelerator --config=rbe_linux_x86_64_cuda13 --test_env=JAX_SKIP_SLOW_TESTS=1 --test_env=TF_CPP_MIN_LOG_LEVEL=0 --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --@local_config_cuda//cuda:override_include_cuda_libs --nobuild -- //tests:gpu_tests //tests:backend_independent_tests //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests //jaxlib/tools:check_gpu_wheel_sources_test //jaxlib/tools:jax_cuda_plugin_wheel_size_test //jaxlib/tools:jax_cuda_pjrt_wheel_size_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test bazel test --build_tag_filters=-multiaccelerator --test_tag_filters=-multiaccelerator --config=rbe_linux_x86_64_cuda13 --test_env=JAX_SKIP_SLOW_TESTS=1 --test_env=TF_CPP_MIN_LOG_LEVEL=0 --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --@local_config_cuda//cuda:override_include_cuda_libs -- //tests:gpu_tests //tests:backend_independent_tests //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests //jaxlib/tools:check_gpu_wheel_sources_test //jaxlib/tools:jax_cuda_plugin_wheel_size_test //jaxlib/tools:jax_cuda_pjrt_wheel_size_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test -bazel analyze-profile profile.json.gz # END BuildType.JAX_LINUX_X86_GPU_T4_GITHUB_ACTIONS # BEGIN BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS bazel --output_base=$GITHUB_WORKSPACE\bazel_output_base build --build_tag_filters= --test_tag_filters= --config=rbe_windows_amd64 --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE\openxla\xla --override_module=xla=$GITHUB_WORKSPACE\openxla\xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_runfile_links=False -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test -bazel analyze-profile profile.json.gz # END BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only=False --//xla/tsl:ci_build --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only=False --//xla/tsl:ci_build -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_ARM64_CPU_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_ARM64_CPU_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_CPU_WORKSPACE_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --config=workspace --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --config=workspace --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_CPU_WORKSPACE_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any --test_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --test_sharding_strategy=disabled --test_strategy=exclusive --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any --test_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --test_sharding_strategy=disabled --test_strategy=exclusive -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --test_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --config=warnings --config=rbe_linux_cpu --config=rocm_clang_hermetic --config=rocm_ci_hermetic --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel build --build_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --test_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --config=warnings --config=rbe_linux_cpu --config=rocm_clang_hermetic --config=rocm_ci_hermetic --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_ONEAPI_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=oneapi-only,requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --test_tag_filters=oneapi-only,-requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --config=nonccl --config=rbe_linux_cpu --config=sycl --config=sycl_hermetic --config=icpx_clang --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel build --build_tag_filters=oneapi-only,requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --test_tag_filters=oneapi-only,-requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --config=nonccl --config=rbe_linux_cpu --config=sycl --config=sycl_hermetic --config=icpx_clang --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_ONEAPI_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_T4_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/... //build_tools/... @tsl//tsl/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_LINUX_X86_GPU_T4_GITHUB_ACTIONS # BEGIN BuildType.XLA_MACOS_ARM64_CPU_KOKORO df -h bazel --version mkdir -p /tmpfs/bazel_output bazel test --build_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --macos_minimum_os=10.15 --test_tmpdir=/tmpfs/bazel_output --test_size_filters=small,medium --//xla/tsl:ci_build -- //xla/... -//xla/python_api/... -//xla/python/... -//xla/service/gpu/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_MACOS_ARM64_CPU_KOKORO # BEGIN BuildType.XLA_MACOS_X86_CPU_KOKORO sudo wget --no-verbose -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-darwin-amd64 @@ -107,9 +88,7 @@ chmod +x /usr/local/bin/bazel bazel --version mkdir -p /Volumes/BuildData/bazel_output bazel test --build_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --macos_minimum_os=10.15 --test_tmpdir=/Volumes/BuildData/bazel_output --//xla/tsl:ci_build -- //xla/... -//xla/python_api/... -//xla/python/... -//xla/service/gpu/... -bazel analyze-profile profile.json.gz # END BuildType.XLA_MACOS_X86_CPU_KOKORO # BEGIN BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS bazel --output_user_root=C:/x build --build_tag_filters=-no_windows,-no_oss,-gpu,-tpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_windows,-no_oss,-gpu,-tpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_xla_windows_x86_cpu_2022 --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -//xla/tpu/... -//xla/backends/cpu/collectives:gloo_collectives_test -//xla/backends/cpu/collectives:mpi_collectives -//xla/backends/cpu/collectives:mpi_communicator -//xla/python/transfer/... -//xla/backends/profiler/subprocess:subprocess_profiling_session -//xla/backends/profiler/subprocess:subprocess_profiling_session_test -//xla/backends/profiler/subprocess:subprocess_registry -//xla/backends/profiler/subprocess:subprocess_registry_test -//xla/backends/cpu/runtime/ynnpack:ynn_fusion_thunk -//xla/backends/cpu/runtime/ynnpack:ynn_interop -//xla/backends/cpu/runtime/ynnpack:ynn_threadpool -//xla/backends/gpu/... -//xla/codegen/emitters/tests/... -//xla/service/gpu/... -//xla/codegen/intrinsic/cpp:eigen_unary_test -bazel analyze-profile profile.json.gz # END BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index ec90dce6dc5dc6..fb55a869629448 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -18,1117 +18,11 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel --- stablehlo/BUILD.bazel +++ stablehlo/BUILD.bazel -@@ -11,24 +11,42 @@ - # 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. -+ - load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "gentbl_filegroup", "td_library") --load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library", "cc_test") -+load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library") -+load("//build_tools:utils.bzl", "get_compatible_with_portable") -+# copybara:uncomment load("@rules_license//rules:license.bzl", "license") - - package( -- default_visibility = ["//visibility:public"], -+ # copybara:uncomment default_applicable_licenses = [":license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], - licenses = ["notice"], +@@ -1960,3 +1960,17 @@ + name = "all_tests", + tests = [":stablehlo_ci_tests"], ) - -+# copybara:uncomment_begin -+# license( -+# name = "license", -+# package_name = "stablehlo", -+# ) -+# copybara:uncomment_end + - exports_files([ - "LICENSE", - # Export to be used by MLIR-free python binding impls - "stablehlo/integrations/python/StablehloApi.cpp", - "stablehlo/integrations/python/StablehloApi.h", -+ # Export TD files for builder code generation. -+ "stablehlo/dialect/ChloOps.td", -+ "stablehlo/dialect/StablehloOps.td", - ]) - - filegroup( - name = "stablehlo_ops_td_filegroup", - srcs = glob(["stablehlo/dialect/*.td"]), -+ compatible_with = get_compatible_with_portable(), - ) - - cc_library( -@@ -39,7 +57,8 @@ - hdrs = [ - "stablehlo/dialect/Base.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base_attr_interfaces_inc_gen", - "@llvm-project//llvm:Support", -@@ -57,7 +76,7 @@ - - gentbl_cc_library( - name = "base_attr_interfaces_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/BaseAttrInterfaces.cpp.inc": ["-gen-attr-interface-defs"], - "stablehlo/dialect/BaseAttrInterfaces.h.inc": ["-gen-attr-interface-decls"], -@@ -72,6 +91,7 @@ - srcs = [ - "stablehlo/dialect/Base.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -@@ -88,7 +108,8 @@ - hdrs = [ - "stablehlo/dialect/BroadcastUtils.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - "@llvm-project//llvm:Support", -@@ -100,7 +121,7 @@ - - gentbl_cc_library( - name = "chlo_attrs_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/ChloAttrs.cpp.inc": ["-gen-attrdef-defs"], - "stablehlo/dialect/ChloAttrs.h.inc": ["-gen-attrdef-decls"], -@@ -126,7 +147,6 @@ - name = "chlo_capi", - srcs = CHLO_CAPI_SOURCES, - hdrs = CHLO_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":chlo_ops", - "@llvm-project//llvm:Support", -@@ -138,7 +158,7 @@ - cc_library( - name = "chlo_capi_headers", - hdrs = CHLO_CAPI_HEADERS, -- strip_include_prefix = ".", -+ includes = ["."], - deps = [ - "@llvm-project//mlir:CAPIIRHeaders", - ], -@@ -149,7 +169,6 @@ - name = "chlo_capi_objects", - srcs = CHLO_CAPI_SOURCES, - hdrs = CHLO_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":chlo_ops", - "@llvm-project//llvm:Support", -@@ -160,7 +179,7 @@ - - gentbl_cc_library( - name = "chlo_enums_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/ChloEnums.cpp.inc": ["-gen-enum-defs"], - "stablehlo/dialect/ChloEnums.h.inc": ["-gen-enum-decls"], -@@ -174,7 +193,7 @@ - - gentbl_cc_library( - name = "chlo_ops_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/ChloOps.cpp.inc": ["-gen-op-defs"], - "stablehlo/dialect/ChloOps.h.inc": ["-gen-op-decls"], -@@ -200,6 +219,7 @@ - - gentbl_filegroup( - name = "chlo_ops_py_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = [ - ( - [ -@@ -223,6 +243,7 @@ - "stablehlo/dialect/ChloEnums.td", - "stablehlo/dialect/ChloOps.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - ":base_td_files", -@@ -249,7 +270,8 @@ - "stablehlo/dialect/ChloBytecode.h", - "stablehlo/dialect/ChloOps.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":broadcast_utils", -@@ -278,6 +300,7 @@ - - gentbl_cc_library( - name = "chlo_rewriters_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/ChloDecompositionPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/ChloDecompositionPatterns.td", -@@ -289,6 +312,7 @@ - - gentbl_cc_library( - name = "stablehlo_aggressive_simplification_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/optimization/StablehloAggressiveSimplificationPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/optimization/StablehloAggressiveSimplificationPatterns.td", -@@ -299,6 +323,7 @@ - - gentbl_cc_library( - name = "stablehlo_legalize_deprecated_ops_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td", -@@ -309,6 +334,7 @@ - - gentbl_cc_library( - name = "vhlo_rewriters_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/VhloToVersionPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/VhloToVersionPatterns.td", -@@ -319,6 +345,7 @@ - - gentbl_cc_library( - name = "stablehlo_create_compatibility_expander_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/StablehloCompatibilityExpanderPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/StablehloCompatibilityExpanderPatterns.td", -@@ -329,6 +356,7 @@ - - gentbl_cc_library( - name = "stablehlo_create_complex_math_expander_inc_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = {"stablehlo/transforms/StablehloComplexMathExpanderPatterns.h.inc": ["--gen-rewriters"]}, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/transforms/StablehloComplexMathExpanderPatterns.td", -@@ -345,7 +373,8 @@ - hdrs = [ - "stablehlo/reference/InterpreterOps.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":interpreter_ops_inc_gen", -@@ -364,7 +393,7 @@ - - gentbl_cc_library( - name = "interpreter_ops_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/reference/InterpreterOps.cpp.inc": ["-gen-op-defs"], - "stablehlo/reference/InterpreterOps.h.inc": ["-gen-op-decls"], -@@ -381,6 +410,7 @@ - srcs = [ - "stablehlo/reference/InterpreterOps.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - ":base_td_files", -@@ -389,6 +419,7 @@ - - gentbl_cc_library( - name = "interpreter_pass_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "stablehlo/reference/InterpreterPasses.h.inc": [ -@@ -412,7 +443,8 @@ - hdrs = [ - "stablehlo/reference/InterpreterPasses.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":interpreter_ops", - ":interpreter_pass_inc_gen", -@@ -426,6 +458,7 @@ - - gentbl_cc_library( - name = "linalg_pass_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "stablehlo/conversions/linalg/transforms/Passes.h.inc": [ -@@ -461,6 +494,7 @@ - "stablehlo/conversions/linalg/transforms/Rewriters.h", - "stablehlo/conversions/linalg/transforms/TypeConversion.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - "stablehlo_type_conversions", -@@ -503,6 +537,7 @@ - hdrs = [ - "stablehlo/reference/Api.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - ":interpreter_ops", -@@ -534,7 +569,8 @@ - hdrs = [ - "stablehlo/reference/Axes.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", -@@ -550,6 +586,7 @@ - hdrs = [ - "stablehlo/reference/Configuration.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - ":reference_errors", -@@ -569,7 +606,8 @@ - hdrs = [ - "stablehlo/reference/Element.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_errors", - ":reference_types", -@@ -588,7 +626,8 @@ - hdrs = [ - "stablehlo/reference/Errors.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - ], -@@ -602,7 +641,8 @@ - hdrs = [ - "stablehlo/reference/Index.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", -@@ -618,6 +658,7 @@ - hdrs = [ - "stablehlo/reference/NumPy.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - ":reference_tensor", -@@ -635,7 +676,8 @@ - hdrs = [ - "stablehlo/reference/Ops.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":reference_axes", -@@ -669,7 +711,8 @@ - hdrs = [ - "stablehlo/reference/Process.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_process_grid", - ":reference_tensor", -@@ -686,7 +729,8 @@ - hdrs = [ - "stablehlo/reference/ProcessGrid.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_tensor", - "@llvm-project//llvm:Support", -@@ -703,7 +747,8 @@ - hdrs = [ - "stablehlo/reference/Scope.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_tensor", - ":reference_token", -@@ -722,7 +767,8 @@ - hdrs = [ - "stablehlo/reference/Tensor.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_axes", - ":reference_element", -@@ -743,7 +789,8 @@ - hdrs = [ - "stablehlo/reference/Token.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_errors", - ":stablehlo_ops", -@@ -761,7 +808,8 @@ - hdrs = [ - "stablehlo/reference/Types.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Support", -@@ -776,7 +824,8 @@ - hdrs = [ - "stablehlo/reference/Value.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":reference_errors", - ":reference_tensor", -@@ -795,7 +844,7 @@ - hdrs = [ - "stablehlo/dialect/Register.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - deps = [ - ":chlo_ops", - ":stablehlo_ops", -@@ -817,7 +866,7 @@ - hdrs = [ - "stablehlo/dialect/AssemblyFormat.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - deps = [ - ":base", - "@llvm-project//llvm:Support", -@@ -828,7 +877,7 @@ - - gentbl_cc_library( - name = "stablehlo_attrs_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/StablehloAttrs.cpp.inc": ["-gen-attrdef-defs"], - "stablehlo/dialect/StablehloAttrs.h.inc": ["-gen-attrdef-decls"], -@@ -860,7 +909,6 @@ - name = "stablehlo_dialect_capi", - srcs = STABLEHLO_DIALECT_CAPI_SOURCES, - hdrs = STABLEHLO_DIALECT_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":interpreter_ops", - ":stablehlo_ops", -@@ -878,7 +926,7 @@ - cc_library( - name = "stablehlo_dialect_capi_headers", - hdrs = STABLEHLO_DIALECT_CAPI_HEADERS, -- strip_include_prefix = ".", -+ includes = ["."], - deps = [ - "@llvm-project//mlir:CAPIIRHeaders", - ], -@@ -889,7 +937,6 @@ - name = "stablehlo_dialect_capi_objects", - srcs = STABLEHLO_DIALECT_CAPI_SOURCES, - hdrs = STABLEHLO_DIALECT_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":interpreter_ops", - ":stablehlo_ops", -@@ -923,7 +970,6 @@ - name = "stablehlo_unified_capi", - srcs = STABLEHLO_UNIFIED_CAPI_SOURCES, - hdrs = STABLEHLO_UNIFIED_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":linalg_passes", - ":reference_api", -@@ -951,7 +997,7 @@ - cc_library( - name = "stablehlo_unified_capi_headers", - hdrs = STABLEHLO_UNIFIED_CAPI_HEADERS, -- strip_include_prefix = ".", -+ includes = ["."], - deps = [ - "@llvm-project//mlir:CAPIIRHeaders", - ], -@@ -967,7 +1013,6 @@ - name = "stablehlo_unified_capi_objects", - srcs = STABLEHLO_UNIFIED_CAPI_SOURCES, - hdrs = STABLEHLO_UNIFIED_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":linalg_passes", - ":reference_api", -@@ -993,7 +1038,7 @@ - - gentbl_cc_library( - name = "stablehlo_enums_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/StablehloEnums.cpp.inc": ["-gen-enum-defs"], - "stablehlo/dialect/StablehloEnums.h.inc": ["-gen-enum-decls"], -@@ -1007,7 +1052,7 @@ - - gentbl_cc_library( - name = "stablehlo_types_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/StablehloTypeDefs.cpp.inc": [ - "-gen-typedef-defs", -@@ -1027,7 +1072,7 @@ - - gentbl_cc_library( - name = "stablehlo_ops_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/StablehloOps.cpp.inc": ["-gen-op-defs"], - "stablehlo/dialect/StablehloOps.h.inc": ["-gen-op-decls"], -@@ -1049,6 +1094,7 @@ - - gentbl_filegroup( - name = "stablehlo_ops_py_gen", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = [ - ( - [ -@@ -1085,6 +1131,7 @@ - "stablehlo/dialect/StablehloTypes.td", - "stablehlo/transforms/ChloDecompositionPatternsMath.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - ":base_td_files", -@@ -1096,6 +1143,7 @@ - - gentbl_cc_library( - name = "stablehlo_pass_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "stablehlo/transforms/Passes.h.inc": [ -@@ -1118,7 +1166,8 @@ - hdrs = [ - "stablehlo/transforms/StablehloBroadcastLowering.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":stablehlo_ops", - "@llvm-project//llvm:Support", -@@ -1136,7 +1185,8 @@ - hdrs = [ - "stablehlo/transforms/PassUtils.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":chlo_ops", -@@ -1145,6 +1195,7 @@ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:ComplexDialect", - "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:QuantOps", - "@llvm-project//mlir:Support", - ], - ) -@@ -1179,7 +1230,8 @@ - "stablehlo/transforms/Passes.h", - "stablehlo/transforms/StablehloRefineShapes.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":broadcast_utils", -@@ -1229,6 +1281,7 @@ - - gentbl_cc_library( - name = "stablehlo_passes_optimization_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "stablehlo/transforms/optimization/Passes.h.inc": [ -@@ -1254,7 +1307,8 @@ - hdrs = [ - "stablehlo/transforms/optimization/Passes.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":chlo_ops", -@@ -1286,7 +1340,8 @@ - hdrs = [ - "stablehlo/api/PortableApi.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":register", - ":stablehlo_ops", -@@ -1310,6 +1365,7 @@ - hdrs = [ - "stablehlo/dialect/Serialization.h", - ], -+ compatible_with = get_compatible_with_portable(), - deps = [ - ":stablehlo_ops", - ":stablehlo_passes", -@@ -1333,7 +1389,8 @@ - hdrs = [ - "stablehlo/transforms/conversions/TypeConversion.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", -@@ -1350,7 +1407,8 @@ - hdrs = [ - "stablehlo/dialect/TypeInference.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":stablehlo_assembly_format", -@@ -1371,7 +1429,8 @@ - hdrs = [ - "stablehlo/dialect/ReplicaGroupUtils.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":stablehlo_ops", - "@llvm-project//llvm:Support", -@@ -1390,7 +1449,8 @@ - "stablehlo/dialect/StablehloBytecode.h", - "stablehlo/dialect/StablehloOps.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":stablehlo_assembly_format", -@@ -1439,11 +1499,13 @@ - srcs = [ - "stablehlo/tools/StablehloTranslateMain.cpp", - ], -+ compatible_with = get_compatible_with_portable(), - deps = [ - ":interpreter_ops", - ":reference_api", - ":reference_configuration", - ":reference_errors", -+ ":reference_numpy", - ":reference_ops", - ":reference_process", - ":reference_process_grid", -@@ -1478,6 +1540,7 @@ - srcs = [ - "stablehlo/tools/StablehloOptMain.cpp", - ], -+ compatible_with = get_compatible_with_portable(), - deps = [ - ":interpreter_ops", - ":interpreter_passes", -@@ -1486,8 +1549,6 @@ - ":stablehlo_passes", - ":stablehlo_passes_optimization", - ":tosa_passes", -- "//stablehlo/tests:check_ops", -- "//stablehlo/tests:test_utils", - "@llvm-project//llvm:Support", - "@llvm-project//mlir:AllExtensions", - "@llvm-project//mlir:AllPassesAndDialects", -@@ -1497,18 +1558,24 @@ - "@llvm-project//mlir:RegisterAllExtensions", - "@llvm-project//mlir:RegisterAllPasses", - "@llvm-project//mlir:TosaDialect", -+ # copybara:uncomment "@shardy//shardy/dialect/sdy/ir:register", -+ "//stablehlo/tests:check_ops", -+ "//stablehlo/tests:test_utils", - ], - ) - - gentbl_cc_library( - name = "tosa_pass_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "stablehlo/conversions/tosa/transforms/Passes.h.inc": [ - "-gen-pass-decls", - "-name=StablehloTOSATransforms", - ], -- "stablehlo/conversions/tosa/transforms/stablehlo_tosa_passes.md": ["-gen-pass-doc"], -+ "stablehlo/conversions/tosa/transforms/stablehlo_tosa_passes.md": [ -+ "-gen-pass-doc", -+ ], - }, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/conversions/tosa/transforms/Passes.td", -@@ -1526,7 +1593,8 @@ - hdrs = [ - "stablehlo/conversions/tosa/transforms/Passes.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":stablehlo_ops", - ":tosa_pass_inc_gen", -@@ -1549,6 +1617,8 @@ - - gentbl_cc_library( - name = "tosa_pdll_inc_gen", -+ compatible_with = get_compatible_with_portable(), -+ strip_include_prefix = ".", - tbl_outs = {"stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll.h.inc": ["-x=cpp"]}, - tblgen = "@llvm-project//mlir:mlir-pdll", - td_file = "stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.pdll", -@@ -1567,7 +1637,8 @@ - hdrs = [ - "stablehlo/dialect/Version.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", -@@ -1587,7 +1658,6 @@ - name = "vhlo_capi", - srcs = VHLO_CAPI_SOURCES, - hdrs = VHLO_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":vhlo_ops", - "@llvm-project//mlir:CAPIIR", -@@ -1598,7 +1668,7 @@ - cc_library( - name = "vhlo_capi_headers", - hdrs = VHLO_CAPI_HEADERS, -- strip_include_prefix = ".", -+ includes = ["."], - deps = [ - "@llvm-project//mlir:CAPIIRHeaders", - ], -@@ -1609,7 +1679,6 @@ - name = "vhlo_capi_objects", - srcs = VHLO_CAPI_SOURCES, - hdrs = VHLO_CAPI_HEADERS, -- strip_include_prefix = ".", - deps = [ - ":vhlo_ops", - "@llvm-project//mlir:CAPIIRObjects", -@@ -1627,15 +1696,13 @@ - - gentbl_filegroup( - name = "vhlo_ops_py_gen", -- tbl_outs = [ -- ( -- [ -- "-gen-python-op-bindings", -- "-bind-dialect=vhlo", -- ], -- "stablehlo/integrations/python/mlir/dialects/_vhlo_ops_gen.py", -- ), -- ], -+ compatible_with = get_compatible_with_portable(), -+ tbl_outs = { -+ "stablehlo/integrations/python/mlir/dialects/_vhlo_ops_gen.py": [ -+ "-gen-python-op-bindings", -+ "-bind-dialect=vhlo", -+ ], -+ }, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "stablehlo/integrations/python/mlir/dialects/VhloOps.td", - deps = [ -@@ -1646,7 +1713,7 @@ - - gentbl_cc_library( - name = "vhlo_attr_interfaces_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloAttrInterfaces.cpp.inc": ["-gen-attr-interface-defs"], - "stablehlo/dialect/VhloAttrInterfaces.h.inc": ["-gen-attr-interface-decls"], -@@ -1660,7 +1727,7 @@ - - gentbl_cc_library( - name = "vhlo_attrs_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloAttrs.cpp.inc": ["-gen-attrdef-defs"], - "stablehlo/dialect/VhloAttrs.h.inc": ["-gen-attrdef-decls"], -@@ -1674,7 +1741,7 @@ - - gentbl_cc_library( - name = "vhlo_enums_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloEnums.cpp.inc": ["-gen-enum-defs"], - "stablehlo/dialect/VhloEnums.h.inc": ["-gen-enum-decls"], -@@ -1688,7 +1755,7 @@ - - gentbl_cc_library( - name = "vhlo_op_interfaces_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloOpInterfaces.cpp.inc": ["-gen-op-interface-defs"], - "stablehlo/dialect/VhloOpInterfaces.h.inc": ["-gen-op-interface-decls"], -@@ -1710,7 +1777,8 @@ - "stablehlo/dialect/VhloBytecode.h", - "stablehlo/dialect/VhloOps.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":base", - ":stablehlo_assembly_format", -@@ -1734,7 +1802,7 @@ - - gentbl_cc_library( - name = "vhlo_ops_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloOps.cpp.inc": ["-gen-op-defs"], - "stablehlo/dialect/VhloOps.h.inc": ["-gen-op-decls"], -@@ -1756,6 +1824,7 @@ - "stablehlo/dialect/VhloOps.td", - "stablehlo/dialect/VhloTypes.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - "@llvm-project//mlir:BuiltinDialectTdFiles", -@@ -1772,7 +1841,8 @@ - hdrs = [ - "stablehlo/dialect/VhloTypes.h", - ], -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), -+ includes = ["."], - deps = [ - ":stablehlo_assembly_format", - ":version", -@@ -1790,7 +1860,7 @@ - - gentbl_cc_library( - name = "vhlo_type_interfaces_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloTypeInterfaces.cpp.inc": ["-gen-type-interface-defs"], - "stablehlo/dialect/VhloTypeInterfaces.h.inc": ["-gen-type-interface-decls"], -@@ -1804,7 +1874,7 @@ - - gentbl_cc_library( - name = "vhlo_types_inc_gen", -- strip_include_prefix = ".", -+ compatible_with = get_compatible_with_portable(), - tbl_outs = { - "stablehlo/dialect/VhloTypeDefs.cpp.inc": ["-gen-typedef-defs"], - "stablehlo/dialect/VhloTypeDefs.h.inc": ["-gen-typedef-decls"], -@@ -1826,227 +1896,21 @@ - ], - ) - --### Builder APIs -- --##### --# TableGen --##### -- --cc_binary( -- name = "mlir_builder_tblgen", -- srcs = ["stablehlo/integrations/cpp/builder/MlirBuilderTblgen.cpp"], -- deps = [ -- "@llvm-project//llvm:Support", -- "@llvm-project//llvm:TableGen", -- "@llvm-project//mlir:Support", -- "@llvm-project//mlir:TableGen", -- ], --) -- --#### --## MlirBuilder base class --#### -- --cc_library( -- name = "mlir_builder", -- srcs = ["stablehlo/integrations/cpp/builder/MlirBuilder.cpp"], -- hdrs = ["stablehlo/integrations/cpp/builder/MlirBuilder.h"], -- strip_include_prefix = ".", -- deps = [ -- ":attr_type_builder_util", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- ], --) -- --cc_test( -- name = "mlir_builder_test", -- srcs = ["stablehlo/integrations/cpp/builder/MlirBuilderTest.cpp"], -- deps = [ -- ":attr_type_builder_util", -- ":chlo_builder", -- ":func_builder", -- ":mlir_builder", -- ":register", -- ":stablehlo_builder", -- ":stablehlo_ops", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:FuncDialect", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- "@llvm-project//mlir:TosaDialect", -- "@llvm-project//third-party/unittest:gmock", -- "@llvm-project//third-party/unittest:gtest", -- ], --) -- --#### --## Attr / Type Builder Helpers --#### -- --cc_library( -- name = "attr_type_builder_util", -- srcs = ["stablehlo/integrations/cpp/builder/AttrTypeBuilderUtil.cpp"], -- hdrs = ["stablehlo/integrations/cpp/builder/AttrTypeBuilderUtil.h"], -- strip_include_prefix = ".", -- deps = [ -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- ], --) -- --cc_test( -- name = "attr_type_builder_util_test", -- srcs = ["stablehlo/integrations/cpp/builder/AttrTypeBuilderUtilTest.cpp"], -- deps = [ -- ":attr_type_builder_util", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- "@llvm-project//third-party/unittest:gmock", -- "@llvm-project//third-party/unittest:gtest", -- ], --) -- --##### --## Dialect-specific builders --#### -- --gentbl_cc_library( -- name = "chlo_builder_inc", -- tbl_outs = { -- "stablehlo/integrations/cpp/builder/ChloBuilder.cpp.inc": ["-gen-builder-defs"], -- "stablehlo/integrations/cpp/builder/ChloBuilder.h.inc": ["-gen-builder-decls"], -- "stablehlo/integrations/cpp/builder/ChloBuilder.md": ["-gen-builder-docs"], -- }, -- tblgen = ":mlir_builder_tblgen", -- td_file = ":stablehlo/dialect/ChloOps.td", -- deps = [ -- ":chlo_ops_td_files", -- "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -- "@llvm-project//mlir:OpBaseTdFiles", -- "@llvm-project//mlir:SideEffectInterfacesTdFiles", -- ], --) -- --cc_library( -- name = "chlo_builder", -- srcs = ["stablehlo/integrations/cpp/builder/ChloBuilder.cpp"], -- hdrs = ["stablehlo/integrations/cpp/builder/ChloBuilder.h"], -- strip_include_prefix = ".", -- deps = [ -- ":chlo_builder_inc", -- ":chlo_ops", -- ":mlir_builder", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- ], --) -- --cc_test( -- name = "chlo_builder_test", -- srcs = ["stablehlo/integrations/cpp/builder/ChloBuilderTest.cpp"], -- deps = [ -- ":attr_type_builder_util", -- ":chlo_builder", -- ":func_builder", -- ":mlir_builder", -- ":register", -- ":stablehlo_builder", -- ":stablehlo_ops", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- "@llvm-project//third-party/unittest:gmock", -- "@llvm-project//third-party/unittest:gtest", -- ], --) -- --gentbl_cc_library( -- name = "func_builder_inc", -- tbl_outs = { -- "stablehlo/integrations/cpp/builder/FuncBuilder.cpp.inc": ["-gen-builder-defs"], -- "stablehlo/integrations/cpp/builder/FuncBuilder.h.inc": ["-gen-builder-decls"], -- "stablehlo/integrations/cpp/builder/FuncBuilder.md": ["-gen-builder-docs"], -- }, -- tblgen = ":mlir_builder_tblgen", -- td_file = "@llvm-project//mlir:include/mlir/Dialect/Func/IR/FuncOps.td", -- deps = [ -- "@llvm-project//mlir:FuncTdFiles", -- "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -- "@llvm-project//mlir:OpBaseTdFiles", -- "@llvm-project//mlir:SideEffectInterfacesTdFiles", -- ], --) -- --cc_library( -- name = "func_builder", -- srcs = ["stablehlo/integrations/cpp/builder/FuncBuilder.cpp"], -- hdrs = ["stablehlo/integrations/cpp/builder/FuncBuilder.h"], -- strip_include_prefix = ".", -- deps = [ -- ":func_builder_inc", -- ":mlir_builder", -- "@llvm-project//mlir:FuncDialect", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- ], --) -- --gentbl_cc_library( -- name = "stablehlo_builder_inc", -- tbl_outs = { -- "stablehlo/integrations/cpp/builder/StablehloBuilder.cpp.inc": ["-gen-builder-defs"], -- "stablehlo/integrations/cpp/builder/StablehloBuilder.h.inc": ["-gen-builder-decls"], -- "stablehlo/integrations/cpp/builder/StablehloBuilder.md": ["-gen-builder-docs"], -- }, -- tblgen = ":mlir_builder_tblgen", -- td_file = ":stablehlo/dialect/StablehloOps.td", -- deps = [ -- ":base_td_files", -- ":stablehlo_ops_td_filegroup", -- "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -- "@llvm-project//mlir:OpBaseTdFiles", -- "@llvm-project//mlir:ShapeOpsTdFiles", -- "@llvm-project//mlir:SideEffectInterfacesTdFiles", -- ], --) -- --cc_library( -- name = "stablehlo_builder", -- srcs = ["stablehlo/integrations/cpp/builder/StablehloBuilder.cpp"], -- hdrs = ["stablehlo/integrations/cpp/builder/StablehloBuilder.h"], -- strip_include_prefix = ".", -- deps = [ -- ":attr_type_builder_util", -- ":mlir_builder", -- ":stablehlo_broadcast_lowering", -- ":stablehlo_builder_inc", -- ":stablehlo_ops", -- ":stablehlo_type_inference", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:FuncDialect", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:InferTypeOpInterface", -- "@llvm-project//mlir:Support", -- ], -+test_suite( -+ name = "all_tests", -+ tests = [":stablehlo_ci_tests"], - ) - --cc_test( -- name = "stablehlo_builder_test", -- srcs = ["stablehlo/integrations/cpp/builder/StablehloBuilderTest.cpp"], -- deps = [ -- ":attr_type_builder_util", -- ":func_builder", -- ":mlir_builder", -- ":register", -- ":stablehlo_builder", -- ":stablehlo_ops", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:Support", -- "@llvm-project//third-party/unittest:gmock", -- "@llvm-project//third-party/unittest:gtest", -- ], --) +# copybara:uncomment_begin +# alias( +# name = "check_ops", @@ -1142,744 +36,34 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel +# deprecation = "Use //stablehlo/tests:test_utils instead.", +# ) +# copybara:uncomment_end -diff --ruN a/stablehlo/build_tools/BUILD.bazel b/stablehlo/build_tools/BUILD.bazel ---- stablehlo/build_tools/BUILD.bazel -+++ stablehlo/build_tools/BUILD.bazel -@@ -0,0 +1,38 @@ -+# Copyright 2025 The StableHLO Authors. All Rights Reserved. -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+# You may obtain a copy of the License at -+# -+# https://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. -+ -+package( -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], -+ licenses = ["notice"], -+) -+ -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ -+exports_files( -+ [ -+ "lit_test_suite.bzl", -+ "utils.bzl", -+ ], -+ visibility = ["//:__subpackages__"], -+) diff --ruN a/stablehlo/build_tools/lit_test_suite.bzl b/stablehlo/build_tools/lit_test_suite.bzl --- stablehlo/build_tools/lit_test_suite.bzl +++ stablehlo/build_tools/lit_test_suite.bzl -@@ -0,0 +1,71 @@ -+# Copyright 2026 The StableHLO Authors. All Rights Reserved. -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+# You may obtain a copy of the License at -+# -+# https://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. -+ -+"""Defines the `lit_test_suite` macro.""" -+ -+load("@llvm-project//llvm:lit_test.bzl", "lit_test") -+ -+def lit_test_suite( -+ name, -+ srcs, -+ size_per_test = "small", -+ data = None, -+ tags = None, -+ args = None): -+ """Defines a `lit_test` for each file in `srcs`; wraps them in a `test_suite`. -+ -+ Args: -+ name: The name of the test suite. -+ srcs: The list/glob of `.mlir` LIT test files included in the test suite. -+ size_per_test: The size of the individual tests comprising the suite. -+ data: List of labels for the test's data dependencies (i.e. runtime deps). -+ tags: List of tags to attach to the tests. -+ args: List of extra arguments to pass to the test driver. -+ """ -+ -+ if data == None: -+ data = [] -+ if tags == None: -+ tags = [] -+ if args == None: -+ args = [] -+ -+ tests_in_suite = [] -+ -+ for src in srcs: +@@ -44,11 +44,13 @@ + tests_in_suite = [] + + for src in srcs: +- test_name = ("%s.test" % src).replace("/", "_") + test_name = ("%s.test" % src).replace("/", "_") # copybara:comment + + # copybara:uncomment test_name = "%s.test" % src -+ lit_test( -+ name = test_name, # copybara:comment -+ # copybara:uncomment name = src, + lit_test( + name = test_name, # copybara:comment + # copybara:uncomment name = src, +- srcs = [src], + srcs = [src], # copybara:comment -+ # copybara:uncomment src = src, -+ args = args, -+ data = data + native.glob( -+ ["%s.bc" % src], -+ # Allow empty since not all `.mlir` tests have `.mlir.bc` files. -+ allow_empty = True, -+ ), -+ deps = ["@rules_python//python/runfiles"], # copybara:comment + # copybara:uncomment src = src, + args = args, + data = data + native.glob( +@@ -57,7 +59,7 @@ + allow_empty = True, + ), + deps = ["@rules_python//python/runfiles"], # copybara:comment +- # copybara:uncomment driver = "//third_party/llvm/llvm-project/mlir:run_lit.sh", + # copybara:uncomment driver = "@llvm-project//mlir:run_lit.sh", -+ tags = tags, -+ size = size_per_test, -+ ) -+ tests_in_suite.append(test_name) -+ -+ native.test_suite( -+ name = name, -+ tests = tests_in_suite, -+ ) -diff --ruN a/stablehlo/build_tools/utils.bzl b/stablehlo/build_tools/utils.bzl ---- stablehlo/build_tools/utils.bzl -+++ stablehlo/build_tools/utils.bzl -@@ -24,3 +24,6 @@ - def workspace_name(): - """Return the name of the workspace.""" - return "_main" if is_bzlmod_enabled() else "stablehlo" -+ -+def get_compatible_with_portable(): -+ return [] -diff --ruN a/stablehlo/examples/CMakeLists.txt b/stablehlo/examples/CMakeLists.txt ---- stablehlo/examples/CMakeLists.txt -+++ stablehlo/examples/CMakeLists.txt -@@ -12,7 +12,7 @@ - # See the License for the specific language governing permissions and - # limitations under the License. - --add_subdirectory(c++) -+add_subdirectory(cpp) - - if(STABLEHLO_ENABLE_BINDINGS_PYTHON) - add_subdirectory(python) -diff --ruN a/stablehlo/examples/c++/BUILD.bazel b/stablehlo/examples/c++/BUILD.bazel ---- stablehlo/examples/c++/BUILD.bazel -+++ stablehlo/examples/c++/BUILD.bazel -@@ -1,53 +0,0 @@ --# Copyright 2023 The StableHLO Authors. All Rights Reserved. --# --# Licensed under the Apache License, Version 2.0 (the "License"); --# you may not use this file except in compliance with the License. --# You may obtain a copy of the License at --# --# https://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. --load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_test") -- --package( -- default_visibility = ["//visibility:public"], -- licenses = ["notice"], --) -- --cc_binary( -- name = "example_add", -- srcs = [ -- "ExampleAdd.cpp", -- ], -- deps = [ -- "//:reference_api", -- "//:reference_configuration", -- "//:stablehlo_ops", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:FuncDialect", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:QuantOps", -- "@llvm-project//mlir:Support", -- ], --) -- --cc_test( -- name = "example_add_test", -- srcs = [ -- "ExampleAdd.cpp", -- ], -- deps = [ -- "//:reference_api", -- "//:reference_configuration", -- "//:stablehlo_ops", -- "@llvm-project//llvm:Support", -- "@llvm-project//mlir:FuncDialect", -- "@llvm-project//mlir:IR", -- "@llvm-project//mlir:QuantOps", -- "@llvm-project//mlir:Support", -- ], --) -diff --ruN a/stablehlo/examples/c++/CMakeLists.txt b/stablehlo/examples/c++/CMakeLists.txt ---- stablehlo/examples/c++/CMakeLists.txt -+++ stablehlo/examples/c++/CMakeLists.txt -@@ -1,29 +0,0 @@ --# Copyright 2024 The StableHLO 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 --# --# https://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. -- --# TODO(fzakaria): As we think about adding more tests we should consider --# making this a function to be DRY. --add_executable(example-add ExampleAdd.cpp) --llvm_update_compile_flags(example-add) --target_link_libraries(example-add PRIVATE StablehloOps StablehloReferenceApi) -- --mlir_check_all_link_libraries(example-add) --add_custom_target(example-add-test -- COMMAND example-add -- WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} -- COMMENT "Executing example-add to validate it works." -- DEPENDS -- example-add --) --add_dependencies(check-stablehlo-quick example-add-test) -diff --ruN a/stablehlo/examples/c++/ExampleAdd.cpp b/stablehlo/examples/c++/ExampleAdd.cpp ---- stablehlo/examples/c++/ExampleAdd.cpp -+++ stablehlo/examples/c++/ExampleAdd.cpp -@@ -1,93 +0,0 @@ --/* Copyright 2022 The StableHLO 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 "llvm/ADT/ArrayRef.h" --#include "llvm/ADT/SmallVector.h" --#include "llvm/Support/LogicalResult.h" --#include "mlir/Dialect/Func/IR/FuncOps.h" --#include "mlir/Dialect/Quant/IR/Quant.h" --#include "mlir/IR/Attributes.h" --#include "mlir/IR/Block.h" --#include "mlir/IR/Builders.h" --#include "mlir/IR/BuiltinAttributes.h" --#include "mlir/IR/BuiltinOps.h" --#include "mlir/IR/BuiltinTypes.h" --#include "mlir/IR/Location.h" --#include "mlir/IR/MLIRContext.h" --#include "mlir/IR/OwningOpRef.h" --#include "mlir/IR/Value.h" --#include "mlir/IR/Verifier.h" --#include "mlir/Support/LLVM.h" --#include "stablehlo/dialect/StablehloOps.h" --#include "stablehlo/reference/Api.h" --#include "stablehlo/reference/Configuration.h" -- --int main() { -- mlir::MLIRContext context; -- -- /** create module **/ -- mlir::OwningOpRef module = -- mlir::ModuleOp::create(mlir::UnknownLoc::get(&context)); -- module->getContext()->loadDialect(); -- module->getContext()->loadDialect(); -- module->getContext()->loadDialect(); -- module->setName("test_module"); -- -- /** create function **/ -- // create function argument and result types. -- auto tensorType = -- mlir::RankedTensorType::get({3, 4}, mlir::Float32Type::get(&context)); -- auto func_type = -- mlir::FunctionType::get(&context, {tensorType, tensorType}, {tensorType}); -- -- // create the function and map arguments. -- llvm::ArrayRef attrs; -- auto function = mlir::func::FuncOp::create(mlir::UnknownLoc::get(&context), -- "main", func_type, attrs); -- function.setVisibility(mlir::func::FuncOp::Visibility::Public); -- module->push_back(function); -- -- // create function block with add operations. -- mlir::Block* block = function.addEntryBlock(); -- llvm::SmallVector arguments(block->args_begin(), -- block->args_end()); -- mlir::OpBuilder block_builder = mlir::OpBuilder::atBlockEnd(block); -- mlir::Location loc = block_builder.getUnknownLoc(); -- -- llvm::SmallVector attributes; -- mlir::Operation* op = -- mlir::stablehlo::AddOp::create(block_builder, loc, arguments, attributes) -- .getOperation(); -- mlir::func::ReturnOp::create(block_builder, loc, op->getResult(0)); -- -- /** verify and dump the module **/ -- assert(mlir::succeeded(mlir::verify(module.get()))); -- -- /* interpret the function "main" with concrete inputs **/ -- auto getConstValue = [&](double val) { -- return mlir::DenseElementsAttr::get( -- tensorType, -- block_builder.getFloatAttr(tensorType.getElementType(), val)); -- }; -- -- auto inputValue1 = getConstValue(10.0); -- auto inputValue2 = getConstValue(20.0); -- auto expectedValue = getConstValue(30.0); -- -- mlir::stablehlo::InterpreterConfiguration config; -- auto results = evalModule(*module, {inputValue1, inputValue2}, config); -- return failed(results) || (*results)[0] != expectedValue; --} -diff --ruN a/stablehlo/examples/c++/README.md b/stablehlo/examples/c++/README.md ---- stablehlo/examples/c++/README.md -+++ stablehlo/examples/c++/README.md -@@ -1,45 +0,0 @@ --# Examples -- --This directory should contain a series of examples as a starting point --for how to use StableHLO. -- --Note: If you have a great example to highlight, we welcome contributions! -- --* [example-add](./ExampleAdd.cpp): A simple example that demonstrates how to -- * Use the StableHLO library to add two numbers. -- * Interpret a StableHLO program using concrete inputs. -- --```c++ -- // Assume 'module' is an MLIR module with function "main" containing StableHLO -- // operations. -- llvm::outs() << "Program:\n " << module << "\n"; -- -- // Create concrete inputs to be used for interpreting "main". -- auto inputValue1 = mlir::DenseElementsAttr::get( -- tensorType, block_builder.getFloatAttr(tensorType.getElementType(), -- static_cast(10))); -- auto inputValue2 = mlir::DenseElementsAttr::get( -- tensorType, block_builder.getFloatAttr(tensorType.getElementType(), -- static_cast(20))); -- llvm::outs() << "Inputs: " << inputValue1 << ", " << inputValue2 << "\n"; -- -- -- mlir::stablehlo::InterpreterConfiguration config; -- auto results = evalModule(module, {inputValue1, inputValue2}, config); -- llvm::outs() << "Output: " << (*results)[0]; --``` -- --Output: -- --```mlir --Program: --module @test_module { -- func.func @main(%arg0: tensor<3x4xf32>, %arg1: tensor<3x4xf32>) -> tensor<3x4xf32> { -- %0 = stablehlo.add %arg0, %arg1 : tensor<3x4xf32> -- %1 = stablehlo.add %arg0, %arg1 : tensor<3x4xf32> -- return %1 : tensor<3x4xf32> -- } --} --Inputs: dense<1.000000e+01> : tensor<3x4xf32>, dense<2.000000e+01> : tensor<3x4xf32> --Output: dense<3.000000e+01> : tensor<3x4xf32> --``` -diff --ruN a/stablehlo/examples/cpp/BUILD.bazel b/stablehlo/examples/cpp/BUILD.bazel ---- stablehlo/examples/cpp/BUILD.bazel -+++ stablehlo/examples/cpp/BUILD.bazel -@@ -0,0 +1,66 @@ -+# Copyright 2023 The StableHLO Authors. All Rights Reserved. -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+# You may obtain a copy of the License at -+# -+# https://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. -+ -+load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_test") -+ -+package( -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], -+ licenses = ["notice"], -+) -+ -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ -+cc_binary( -+ name = "example_add", -+ srcs = [ -+ "ExampleAdd.cpp", -+ ], -+ deps = [ -+ "//:reference_api", -+ "//:reference_configuration", -+ "//:stablehlo_ops", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:FuncDialect", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:QuantOps", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+cc_test( -+ name = "example_add_test", -+ srcs = [ -+ "ExampleAdd.cpp", -+ ], -+ deps = [ -+ "//:reference_api", -+ "//:reference_configuration", -+ "//:stablehlo_ops", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:FuncDialect", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:QuantOps", -+ "@llvm-project//mlir:Support", -+ ], -+) -diff --ruN a/stablehlo/examples/cpp/CMakeLists.txt b/stablehlo/examples/cpp/CMakeLists.txt ---- stablehlo/examples/cpp/CMakeLists.txt -+++ stablehlo/examples/cpp/CMakeLists.txt -@@ -0,0 +1,29 @@ -+# Copyright 2024 The StableHLO 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 -+# -+# https://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. -+ -+# TODO(fzakaria): As we think about adding more tests we should consider -+# making this a function to be DRY. -+add_executable(example-add ExampleAdd.cpp) -+llvm_update_compile_flags(example-add) -+target_link_libraries(example-add PRIVATE StablehloOps StablehloReferenceApi) -+ -+mlir_check_all_link_libraries(example-add) -+add_custom_target(example-add-test -+ COMMAND example-add -+ WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} -+ COMMENT "Executing example-add to validate it works." -+ DEPENDS -+ example-add -+) -+add_dependencies(check-stablehlo-quick example-add-test) -diff --ruN a/stablehlo/examples/cpp/ExampleAdd.cpp b/stablehlo/examples/cpp/ExampleAdd.cpp ---- stablehlo/examples/cpp/ExampleAdd.cpp -+++ stablehlo/examples/cpp/ExampleAdd.cpp -@@ -0,0 +1,93 @@ -+/* Copyright 2022 The StableHLO 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 "llvm/ADT/ArrayRef.h" -+#include "llvm/ADT/SmallVector.h" -+#include "llvm/Support/LogicalResult.h" -+#include "mlir/Dialect/Func/IR/FuncOps.h" -+#include "mlir/Dialect/Quant/IR/Quant.h" -+#include "mlir/IR/Attributes.h" -+#include "mlir/IR/Block.h" -+#include "mlir/IR/Builders.h" -+#include "mlir/IR/BuiltinAttributes.h" -+#include "mlir/IR/BuiltinOps.h" -+#include "mlir/IR/BuiltinTypes.h" -+#include "mlir/IR/Location.h" -+#include "mlir/IR/MLIRContext.h" -+#include "mlir/IR/OwningOpRef.h" -+#include "mlir/IR/Value.h" -+#include "mlir/IR/Verifier.h" -+#include "mlir/Support/LLVM.h" -+#include "stablehlo/dialect/StablehloOps.h" -+#include "stablehlo/reference/Api.h" -+#include "stablehlo/reference/Configuration.h" -+ -+int main() { -+ mlir::MLIRContext context; -+ -+ /** create module **/ -+ mlir::OwningOpRef module = -+ mlir::ModuleOp::create(mlir::UnknownLoc::get(&context)); -+ module->getContext()->loadDialect(); -+ module->getContext()->loadDialect(); -+ module->getContext()->loadDialect(); -+ module->setName("test_module"); -+ -+ /** create function **/ -+ // create function argument and result types. -+ auto tensorType = -+ mlir::RankedTensorType::get({3, 4}, mlir::Float32Type::get(&context)); -+ auto func_type = -+ mlir::FunctionType::get(&context, {tensorType, tensorType}, {tensorType}); -+ -+ // create the function and map arguments. -+ llvm::ArrayRef attrs; -+ auto function = mlir::func::FuncOp::create(mlir::UnknownLoc::get(&context), -+ "main", func_type, attrs); -+ function.setVisibility(mlir::func::FuncOp::Visibility::Public); -+ module->push_back(function); -+ -+ // create function block with add operations. -+ mlir::Block* block = function.addEntryBlock(); -+ llvm::SmallVector arguments(block->args_begin(), -+ block->args_end()); -+ mlir::OpBuilder block_builder = mlir::OpBuilder::atBlockEnd(block); -+ mlir::Location loc = block_builder.getUnknownLoc(); -+ -+ llvm::SmallVector attributes; -+ mlir::Operation* op = -+ mlir::stablehlo::AddOp::create(block_builder, loc, arguments, attributes) -+ .getOperation(); -+ mlir::func::ReturnOp::create(block_builder, loc, op->getResult(0)); -+ -+ /** verify and dump the module **/ -+ assert(mlir::succeeded(mlir::verify(module.get()))); -+ -+ /* interpret the function "main" with concrete inputs **/ -+ auto getConstValue = [&](double val) { -+ return mlir::DenseElementsAttr::get( -+ tensorType, -+ block_builder.getFloatAttr(tensorType.getElementType(), val)); -+ }; -+ -+ auto inputValue1 = getConstValue(10.0); -+ auto inputValue2 = getConstValue(20.0); -+ auto expectedValue = getConstValue(30.0); -+ -+ mlir::stablehlo::InterpreterConfiguration config; -+ auto results = evalModule(*module, {inputValue1, inputValue2}, config); -+ return failed(results) || (*results)[0] != expectedValue; -+} -diff --ruN a/stablehlo/examples/cpp/README.md b/stablehlo/examples/cpp/README.md ---- stablehlo/examples/cpp/README.md -+++ stablehlo/examples/cpp/README.md -@@ -0,0 +1,45 @@ -+# Examples -+ -+This directory should contain a series of examples as a starting point -+for how to use StableHLO. -+ -+Note: If you have a great example to highlight, we welcome contributions! -+ -+* [example-add](./ExampleAdd.cpp): A simple example that demonstrates how to -+ * Use the StableHLO library to add two numbers. -+ * Interpret a StableHLO program using concrete inputs. -+ -+```c++ -+ // Assume 'module' is an MLIR module with function "main" containing StableHLO -+ // operations. -+ llvm::outs() << "Program:\n " << module << "\n"; -+ -+ // Create concrete inputs to be used for interpreting "main". -+ auto inputValue1 = mlir::DenseElementsAttr::get( -+ tensorType, block_builder.getFloatAttr(tensorType.getElementType(), -+ static_cast(10))); -+ auto inputValue2 = mlir::DenseElementsAttr::get( -+ tensorType, block_builder.getFloatAttr(tensorType.getElementType(), -+ static_cast(20))); -+ llvm::outs() << "Inputs: " << inputValue1 << ", " << inputValue2 << "\n"; -+ -+ -+ mlir::stablehlo::InterpreterConfiguration config; -+ auto results = evalModule(module, {inputValue1, inputValue2}, config); -+ llvm::outs() << "Output: " << (*results)[0]; -+``` -+ -+Output: -+ -+```mlir -+Program: -+module @test_module { -+ func.func @main(%arg0: tensor<3x4xf32>, %arg1: tensor<3x4xf32>) -> tensor<3x4xf32> { -+ %0 = stablehlo.add %arg0, %arg1 : tensor<3x4xf32> -+ %1 = stablehlo.add %arg0, %arg1 : tensor<3x4xf32> -+ return %1 : tensor<3x4xf32> -+ } -+} -+Inputs: dense<1.000000e+01> : tensor<3x4xf32>, dense<2.000000e+01> : tensor<3x4xf32> -+Output: dense<3.000000e+01> : tensor<3x4xf32> -+``` -diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel b/stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel ---- stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel -+++ stablehlo/stablehlo/conversions/linalg/tests/BUILD.bazel -@@ -11,14 +11,27 @@ - # 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. -+ - load("@bazel_skylib//rules:expand_template.bzl", "expand_template") --load("@llvm-project//llvm:lit_test.bzl", "lit_test") -+load("//build_tools:lit_test_suite.bzl", "lit_test_suite") - - package( -- default_visibility = ["//visibility:public"], -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], - licenses = ["notice"], - ) - -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ - # Equivalent of configure_lit_site_cfg from CMakeLists.txt. - expand_template( - name = "lit_site_cfg_py_gen", -@@ -39,25 +52,15 @@ - template = "lit.site.cfg.py.in", - ) - --# Equivalent of add_lit_testsuite from CMakeLists.txt. --[ -- lit_test( -- name = "%s.test" % src, -- size = "small", -- srcs = [src], -- data = [ -- "lit.cfg.py", -- "lit.site.cfg.py", -- "//:stablehlo-opt", -- "@llvm-project//llvm:FileCheck", -- ], -- tags = ["stablehlo_linalg_tests"], -- deps = ["@rules_python//python/runfiles"], -- ) -- for src in glob(["**/*.mlir"]) --] -- --test_suite( -+lit_test_suite( - name = "stablehlo_linalg_tests", -+ srcs = glob(["**/*.mlir"]), -+ data = [ -+ "lit.cfg.py", -+ "lit.site.cfg.py", -+ "//:stablehlo-opt", -+ "@llvm-project//llvm:FileCheck", -+ ], -+ size_per_test = "small", - tags = ["stablehlo_linalg_tests"], - ) -diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir b/stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir ---- stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir -+++ stablehlo/stablehlo/conversions/linalg/tests/convolution.mlir -@@ -52,8 +52,8 @@ - // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM0]], %[[DIM2]]) - // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ZERO]]{{.*}}outs(%[[INIT]] - // CHECK: linalg.conv_1d_nwc_wcf --// CHECK-SAME: {dilations = dense<1> : tensor<1xi64> --// CHECK-SAME: someattr -+// CHECK-SAME: {someattr -+// CHECK-SAME: dilations = dense<1> : tensor<1xi64> - // CHECK-SAME: strides = dense<1> : tensor<1xi64>} - // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor, tensor<2x?x?xf32>) - // CHECK-SAME: outs(%[[FILL]] : tensor) -> tensor -@@ -360,7 +360,7 @@ - // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x3x4x2x3xf32> - // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<2x3x4x2x3xf32>) -> tensor<2x3x4x2x3xf32> - // CHECK: %[[OUT:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm --// CHECK-SAME: {dilations = dense<1> : tensor<2xi64>, someattr, strides = dense<1> : tensor<2xi64>} -+// CHECK-SAME: {someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} - // CHECK-SAME: ins(%[[IN]], %[[EXPAND]] : tensor<2x4x5x2xf32>, tensor<2x2x2x3xf32>) - // CHECK-SAME: outs(%[[FILL]] : tensor<2x3x4x2x3xf32>) -> tensor<2x3x4x2x3xf32> - // CHECK: %{{.+}} = tensor.collapse_shape %[[OUT]] -@@ -409,7 +409,7 @@ - // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x3x6x2x2xf32> - // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ZERO]] : f32) outs(%[[INIT]] : tensor<2x3x6x2x2xf32>) -> tensor<2x3x6x2x2xf32> - // CHECK: %[[OUT:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm --// CHECK-SAME: {dilations = dense<1> : tensor<2xi64>, someattr, strides = dense<1> : tensor<2xi64>} -+// CHECK-SAME: {someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} - // CHECK-SAME: ins(%[[PAD]], %[[EXPAND]] : tensor<2x4x7x2xf32>, tensor<2x2x2x2xf32>) - // CHECK-SAME: outs(%[[FILL]] : tensor<2x3x6x2x2xf32>) -> tensor<2x3x6x2x2xf32> - // CHECK: %{{.+}} = tensor.collapse_shape %[[OUT]] -diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/reduce.mlir b/stablehlo/stablehlo/conversions/linalg/tests/reduce.mlir ---- stablehlo/stablehlo/conversions/linalg/tests/reduce.mlir -+++ stablehlo/stablehlo/conversions/linalg/tests/reduce.mlir -@@ -394,8 +394,8 @@ - // CHECK: %[[INIT_VAL:.+]] = tensor.extract %[[ARG1]][] : tensor - // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[INIT_VAL]] : f32) outs(%[[INIT]] : tensor<1x8x8x64xf32>) -> tensor<1x8x8x64xf32> - // CHECK: %[[RES:.+]] = linalg.pooling_nhwc_min --// CHECK-SAME: {dilations = dense<1> : vector<2xi64> --// CHECK-SAME: someattr, -+// CHECK-SAME: {someattr, -+// CHECK-SAME: dilations = dense<1> : vector<2xi64> - // CHECK-SAME: strides = dense<2> : vector<2xi64>} - // CHECK-SAME: ins(%[[ARG0]], %[[WINDOW]] : tensor<1x17x17x64xf32>, tensor<3x3xf32>) - // CHECK-SAME: outs(%[[FILL]] : tensor<1x8x8x64xf32>) -> tensor<1x8x8x64xf32> + tags = tags, + size = size_per_test, + ) diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h b/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h --- stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h +++ stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h @@ -1939,76 +123,6 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgU /// Extracts integer values from the attribute |elements|. SmallVector extract1DVector(DenseIntElementsAttr elements); -diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel ---- stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel -+++ stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel -@@ -11,14 +11,27 @@ - # 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. -+ - load("@bazel_skylib//rules:expand_template.bzl", "expand_template") --load("@llvm-project//llvm:lit_test.bzl", "lit_test") -+load("//build_tools:lit_test_suite.bzl", "lit_test_suite") - - package( -- default_visibility = ["//visibility:public"], -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], - licenses = ["notice"], - ) - -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ - # Equivalent of configure_lit_site_cfg from CMakeLists.txt. - expand_template( - name = "lit_site_cfg_py_gen", -@@ -39,26 +52,16 @@ - template = "lit.site.cfg.py.in", - ) - --# Equivalent of add_lit_testsuite from CMakeLists.txt. --[ -- lit_test( -- name = "%s.test" % src, -- size = "small", -- srcs = [src], -- data = [ -- "lit.cfg.py", -- "lit.site.cfg.py", -- "//:stablehlo-opt", -- "//:stablehlo-translate", -- "@llvm-project//llvm:FileCheck", -- ], -- tags = ["stablehlo_tosa_tests"], -- deps = ["@rules_python//python/runfiles"], -- ) -- for src in glob(["**/*.mlir"]) --] -- --test_suite( -+lit_test_suite( - name = "stablehlo_tosa_tests", -+ srcs = glob(["**/*.mlir"]), -+ data = [ -+ "lit.cfg.py", -+ "lit.site.cfg.py", -+ "//:stablehlo-opt", -+ "//:stablehlo-translate", -+ "@llvm-project//llvm:FileCheck", -+ ], -+ size_per_test = "small", - tags = ["stablehlo_tosa_tests"], - ) diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.cpp b/stablehlo/stablehlo/dialect/ChloBytecode.cpp --- stablehlo/stablehlo/dialect/ChloBytecode.cpp +++ stablehlo/stablehlo/dialect/ChloBytecode.cpp @@ -2175,7 +289,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dial //===----------------------------------------------------------------------===// // Ragged dot op definitions. -@@ -90,9 +94,12 @@ +@@ -90,7 +94,9 @@ let cppNamespace = "::mlir::chlo"; } @@ -2186,9 +300,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dial def CHLO_PrecisionConfigAttr: TypedArrayAttrBase; - -+ - #endif // STABLEHLO_DIALECT_CHLO_ENUMS diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/dialect/ChloOps.cpp --- stablehlo/stablehlo/dialect/ChloOps.cpp +++ stablehlo/stablehlo/dialect/ChloOps.cpp @@ -2310,301 +421,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/Register.h b/stablehlo/stablehlo/dialec } // namespace stablehlo } // namespace mlir -diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp ---- stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp -+++ stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp -@@ -20,6 +20,7 @@ - #include - - #include "llvm/ADT/DenseSet.h" -+#include "llvm/ADT/STLExtras.h" - #include "llvm/ADT/SmallVector.h" - #include "llvm/ADT/StringRef.h" - #include "mlir/IR/Attributes.h" -@@ -31,57 +32,53 @@ - namespace mlir { - namespace stablehlo { - --static SmallVector> -+namespace { -+ -+struct ReindexedAxes { -+ SmallVector splitAxisSizes; -+ SmallVector groupedAxisIndices; -+}; -+ -+// Generates replica groups from the reshaped mesh axis sizes and the indices of -+// the communication axes using a Reshape-Transpose permutation. -+SmallVector> - flattenedReplicaGroupsFromTransposePermutation( -- const SmallVector& meshAxisNames, -- const SmallVector& commAxisNames, -- const llvm::DenseSet& commAxisSet, -- const SmallVector& axisSizes, -- const SmallVector& deviceIds, int64_t totalDevices) { -- // Reshape and Transpose equivalence bridging XLA TileAssignment behavior. -+ ArrayRef axisSizes, ArrayRef groupedAxisIndices, -+ ArrayRef deviceIds, int64_t totalDevices) { -+ llvm::DenseSet groupedAxisSet(groupedAxisIndices.begin(), -+ groupedAxisIndices.end()); - SmallVector transposeAxes; - // Non-grouped axes first -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -- if (!commAxisSet.count(meshAxisNames[i])) { -+ for (size_t i = 0; i < axisSizes.size(); ++i) { -+ if (!groupedAxisSet.count(i)) { - transposeAxes.push_back(i); - } - } -- // Grouped axes -- for (const auto& name : commAxisNames) { -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -- if (meshAxisNames[i] == name) { -- transposeAxes.push_back(i); -- break; -- } -- } -+ // Grouped axes in the specified order -+ for (int64_t idx : groupedAxisIndices) { -+ transposeAxes.push_back(idx); - } - -- SmallVector transposedSizes(meshAxisNames.size()); -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -+ SmallVector transposedSizes(axisSizes.size()); -+ for (size_t i = 0; i < axisSizes.size(); ++i) { - transposedSizes[i] = axisSizes[transposeAxes[i]]; - } - -- // Compute strides for original shape -- SmallVector originalStrides(meshAxisNames.size(), 1); -- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { -+ // Compute strides for reshaped shape -+ SmallVector originalStrides(axisSizes.size(), 1); -+ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { - originalStrides[i] = originalStrides[i + 1] * axisSizes[i + 1]; - } - - // Compute strides for transposed shape -- SmallVector transposedStrides(meshAxisNames.size(), 1); -- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { -+ SmallVector transposedStrides(axisSizes.size(), 1); -+ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { - transposedStrides[i] = transposedStrides[i + 1] * transposedSizes[i + 1]; - } - -- // Generate chunks - int64_t numDevicesPerGroup = 1; -- for (auto name : commAxisNames) { -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -- if (meshAxisNames[i] == name) { -- numDevicesPerGroup *= axisSizes[i]; -- break; -- } -- } -+ for (int64_t idx : groupedAxisIndices) { -+ numDevicesPerGroup *= axisSizes[idx]; - } - int64_t numGroups = totalDevices / numDevicesPerGroup; - -@@ -93,7 +90,7 @@ - for (int64_t j = 0; j < numDevicesPerGroup; ++j) { - int64_t linearTransposeIdx = i * numDevicesPerGroup + j; - int64_t originalIndex = 0; -- for (size_t k = 0; k < meshAxisNames.size(); ++k) { -+ for (size_t k = 0; k < axisSizes.size(); ++k) { - int64_t coord = - (linearTransposeIdx / transposedStrides[k]) % transposedSizes[k]; - originalIndex += coord * originalStrides[transposeAxes[k]]; -@@ -102,10 +99,136 @@ - } - groups.push_back(std::move(group)); - } -- - return groups; - } - -+// Splits mesh axes based on sub-axis references and computes the corresponding -+// indices for the communication axes. -+FailureOr computeReindexedAxes(ArrayRef axesInMesh, -+ ArrayAttr commAxes, -+ Location loc) { -+ ReindexedAxes result; -+ -+ // Validate commAxes and verify that all mesh axes exist and have valid sizes. -+ for (auto attr : commAxes) { -+ auto shloAxisRef = llvm::dyn_cast(attr); -+ if (!shloAxisRef) { -+ return emitError(loc) << "expected AxisRefAttr in comm_axes"; -+ } -+ StringRef axisName = shloAxisRef.getName(); -+ bool found = false; -+ for (auto meshAxis : axesInMesh) { -+ if (meshAxis.getName() == axisName) { -+ found = true; -+ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { -+ int64_t preSize = subAxisInfo.getPreSize(); -+ int64_t size = subAxisInfo.getSize(); -+ if (preSize < 1 || size < 1) { -+ return emitError(loc) -+ << "sub-axis pre_size and size must be at least 1"; -+ } -+ int64_t nextPreSize = preSize * size; -+ if (nextPreSize > meshAxis.getSize() || -+ meshAxis.getSize() % nextPreSize != 0) { -+ return emitError(loc) -+ << "sub-axis (pre_size * size) must divide mesh axis size"; -+ } -+ } -+ break; -+ } -+ } -+ if (!found) { -+ return emitError(loc) -+ << "axis '" << axisName << "' not found in mesh definition"; -+ } -+ } -+ -+ // Split each mesh axis according to the referenced subaxes. -+ struct SplitDim { -+ StringRef axisName; -+ int64_t preSize; -+ int64_t size; -+ int64_t dimIndex; -+ }; -+ SmallVector splitDims; -+ -+ for (auto meshAxis : axesInMesh) { -+ StringRef axisName = meshAxis.getName(); -+ int64_t axisSize = meshAxis.getSize(); -+ -+ if (axisSize == 1) { -+ int64_t dimIdx = result.splitAxisSizes.size(); -+ result.splitAxisSizes.push_back(1); -+ splitDims.push_back({axisName, /*preSize=*/1, /*size=*/1, dimIdx}); -+ continue; -+ } -+ -+ SmallVector preSizes = {1, axisSize}; -+ for (auto attr : commAxes) { -+ auto shloAxisRef = llvm::cast(attr); -+ if (shloAxisRef.getName() == axisName) { -+ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { -+ preSizes.push_back(subAxisInfo.getPreSize()); -+ preSizes.push_back(subAxisInfo.getPreSize() * subAxisInfo.getSize()); -+ } -+ } -+ } -+ -+ llvm::sort(preSizes); -+ preSizes.erase(llvm::unique(preSizes), preSizes.end()); -+ -+ for (size_t j = 0; j < preSizes.size() - 1; ++j) { -+ int64_t segPreSize = preSizes[j]; -+ int64_t segSize = preSizes[j + 1] / segPreSize; -+ int64_t dimIdx = result.splitAxisSizes.size(); -+ result.splitAxisSizes.push_back(segSize); -+ splitDims.push_back({axisName, segPreSize, segSize, dimIdx}); -+ } -+ } -+ -+ // Map each communication axis to its corresponding split dimension. -+ llvm::DenseSet groupedSet; -+ for (auto attr : commAxes) { -+ auto shloAxisRef = llvm::cast(attr); -+ StringRef axisName = shloAxisRef.getName(); -+ int64_t reqPreSize = 1; -+ int64_t reqSize = 0; -+ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { -+ reqPreSize = subAxisInfo.getPreSize(); -+ reqSize = subAxisInfo.getSize(); -+ } else { -+ for (auto meshAxis : axesInMesh) { -+ if (meshAxis.getName() == axisName) { -+ reqSize = meshAxis.getSize(); -+ break; -+ } -+ } -+ } -+ -+ bool matched = false; -+ for (const auto& splitDim : splitDims) { -+ if (splitDim.axisName == axisName && splitDim.preSize == reqPreSize && -+ splitDim.size == reqSize) { -+ if (!groupedSet.insert(splitDim.dimIndex).second) { -+ return emitError(loc) -+ << "Duplicate or overlapping communication axis: " << axisName; -+ } -+ result.groupedAxisIndices.push_back(splitDim.dimIndex); -+ matched = true; -+ break; -+ } -+ } -+ if (!matched) { -+ return emitError(loc) << "Invalid or overlapping communication axis on '" -+ << axisName << "'"; -+ } -+ } -+ -+ return result; -+} -+ -+} // namespace -+ - FailureOr>> flattenReplicaGroupMeshAxes( - Attribute meshAttr, ArrayAttr commAxes, Location loc) { - // Reference implementation for flattening a mesh-axes based replica group -@@ -120,34 +243,13 @@ - if (!mesh) - return emitOptionalError(loc, "expected stablehlo.mesh for mesh attribute"); - -- auto axesInMesh = mesh.getAxes(); -- -- // Identify which axes are communication axes. -- llvm::SmallVector commAxisNames; -- llvm::DenseSet commAxisSet; -- for (auto attr : commAxes) { -- auto shloAxisRef = llvm::dyn_cast(attr); -- if (!shloAxisRef) { -- return emitError(loc) << "expected AxisRefAttr in comm_axes"; -- } -- if (shloAxisRef.getSubAxisInfo()) { -- return emitError(loc) << "Subaxes are not supported in " -- "flattenReplicaGroupMeshAxes"; -- } -- commAxisNames.push_back(shloAxisRef.getName()); -- commAxisSet.insert(shloAxisRef.getName()); -- } -- -- // Calculate total devices and axis sizes -+ FailureOr reindexedAxes = -+ computeReindexedAxes(mesh.getAxes(), commAxes, loc); -+ if (failed(reindexedAxes)) return failure(); - - int64_t totalDevices = 1; -- SmallVector axisSizes; -- SmallVector meshAxisNames; -- for (auto meshAxis : axesInMesh) { -- auto typedMeshAxis = llvm::cast(meshAxis); -- axisSizes.push_back(typedMeshAxis.getSize()); -- meshAxisNames.push_back(typedMeshAxis.getName()); -- totalDevices *= typedMeshAxis.getSize(); -+ for (auto meshAxis : mesh.getAxes()) { -+ totalDevices *= llvm::cast(meshAxis).getSize(); - } - - SmallVector deviceIds; -@@ -160,8 +262,8 @@ - } - - return flattenedReplicaGroupsFromTransposePermutation( -- meshAxisNames, commAxisNames, commAxisSet, axisSizes, deviceIds, -- totalDevices); -+ reindexedAxes->splitAxisSizes, reindexedAxes->groupedAxisIndices, -+ deviceIds, totalDevices); - } - - } // namespace stablehlo diff --ruN a/stablehlo/stablehlo/dialect/StablehloBytecode.h b/stablehlo/stablehlo/dialect/StablehloBytecode.h --- stablehlo/stablehlo/dialect/StablehloBytecode.h +++ stablehlo/stablehlo/dialect/StablehloBytecode.h @@ -2686,7 +502,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo def STABLEHLO_RNG_ALGORITHM_DEFAULT : I32EnumAttrCase<"DEFAULT", 0>; def STABLEHLO_RNG_ALGORITHM_THREE_FRY : I32EnumAttrCase<"THREE_FRY", 1>; -@@ -220,6 +232,9 @@ +@@ -220,6 +232,8 @@ let cppNamespace = "::mlir::stablehlo"; } @@ -2696,7 +512,6 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo +} #endif // STABLEHLO_DIALECT_STABLEHLO_ENUMS -+ diff --ruN a/stablehlo/stablehlo/dialect/VhloBytecode.h b/stablehlo/stablehlo/dialect/VhloBytecode.h --- stablehlo/stablehlo/dialect/VhloBytecode.h +++ stablehlo/stablehlo/dialect/VhloBytecode.h @@ -2883,315 +698,6 @@ diff --ruN a/stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp b/stable resultsVec.push_back(wrap(result)); } return mlirArrayAttrGet(mlirModuleGetContext(module), resultsVec.size(), -diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel b/stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel ---- stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel -+++ stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel -@@ -0,0 +1,281 @@ -+# Copyright 2025 The StableHLO Authors. All Rights Reserved. -+# -+# Licensed under the Apache License, Version 2.0 (the "License"); -+# you may not use this file except in compliance with the License. -+# You may obtain a copy of the License at -+# -+# https://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. -+ -+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library") -+load("@rules_cc//cc:cc_binary.bzl", "cc_binary") -+load("@rules_cc//cc:cc_library.bzl", "cc_library") -+load("@rules_cc//cc:cc_test.bzl", "cc_test") -+load("//build_tools:utils.bzl", "get_compatible_with_portable") -+ -+package( -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], -+ licenses = ["notice"], -+) -+ -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ -+filegroup( -+ name = "stablehlo_builder_filegroup", -+ srcs = glob(["**"]), -+) -+ -+##### -+# TableGen -+##### -+ -+cc_binary( -+ name = "mlir_builder_tblgen", -+ srcs = ["MlirBuilderTblgen.cpp"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ "@llvm-project//llvm:Support", -+ "@llvm-project//llvm:TableGen", -+ "@llvm-project//mlir:Support", -+ "@llvm-project//mlir:TableGen", -+ ], -+) -+ -+#### -+## MlirBuilder base class -+#### -+ -+cc_library( -+ name = "mlir_builder", -+ srcs = ["MlirBuilder.cpp"], -+ hdrs = ["MlirBuilder.h"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+#### -+## Attr / Type Builder Helpers -+#### -+ -+cc_library( -+ name = "attr_type_builder_util", -+ srcs = ["AttrTypeBuilderUtil.cpp"], -+ hdrs = ["AttrTypeBuilderUtil.h"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+cc_test( -+ name = "attr_type_builder_util_test", -+ srcs = ["AttrTypeBuilderUtilTest.cpp"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ "@llvm-project//third-party/unittest:gmock", -+ "@llvm-project//third-party/unittest:gtest", -+ ], -+) -+ -+##### -+## Dialect-specific builders -+#### -+ -+gentbl_cc_library( -+ name = "chlo_builder_inc", -+ compatible_with = get_compatible_with_portable(), -+ tbl_outs = { -+ "ChloBuilder.h.inc": ["-gen-builder-decls"], -+ "ChloBuilder.cpp.inc": ["-gen-builder-defs"], -+ "ChloBuilder.md": ["-gen-builder-docs"], -+ }, -+ tblgen = ":mlir_builder_tblgen", -+ td_file = "//:stablehlo/dialect/ChloOps.td", -+ deps = [ -+ "//:chlo_ops_td_files", -+ "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -+ "@llvm-project//mlir:OpBaseTdFiles", -+ "@llvm-project//mlir:SideEffectInterfacesTdFiles", -+ ], -+) -+ -+cc_library( -+ name = "chlo_builder", -+ srcs = ["ChloBuilder.cpp"], -+ hdrs = ["ChloBuilder.h"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":chlo_builder_inc", -+ ":mlir_builder", -+ "//:chlo_ops", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+gentbl_cc_library( -+ name = "func_builder_inc", -+ compatible_with = get_compatible_with_portable(), -+ tbl_outs = { -+ "FuncBuilder.h.inc": ["-gen-builder-decls"], -+ "FuncBuilder.cpp.inc": ["-gen-builder-defs"], -+ "FuncBuilder.md": ["-gen-builder-docs"], -+ }, -+ tblgen = ":mlir_builder_tblgen", -+ td_file = "@llvm-project//mlir:include/mlir/Dialect/Func/IR/FuncOps.td", -+ deps = [ -+ "@llvm-project//mlir:FuncTdFiles", -+ "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -+ "@llvm-project//mlir:OpBaseTdFiles", -+ "@llvm-project//mlir:SideEffectInterfacesTdFiles", -+ ], -+) -+ -+cc_library( -+ name = "func_builder", -+ srcs = ["FuncBuilder.cpp"], -+ hdrs = ["FuncBuilder.h"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":func_builder_inc", -+ ":mlir_builder", -+ "@llvm-project//mlir:FuncDialect", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+gentbl_cc_library( -+ name = "stablehlo_builder_inc", -+ compatible_with = get_compatible_with_portable(), -+ tbl_outs = { -+ "StablehloBuilder.h.inc": ["-gen-builder-decls"], -+ "StablehloBuilder.cpp.inc": ["-gen-builder-defs"], -+ "StablehloBuilder.md": ["-gen-builder-docs"], -+ }, -+ tblgen = ":mlir_builder_tblgen", -+ td_file = "//:stablehlo/dialect/StablehloOps.td", -+ deps = [ -+ "//:base_td_files", -+ "//:stablehlo_ops_td_filegroup", -+ "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", -+ "@llvm-project//mlir:OpBaseTdFiles", -+ "@llvm-project//mlir:ShapeOpsTdFiles", -+ "@llvm-project//mlir:SideEffectInterfacesTdFiles", -+ ], -+) -+ -+cc_library( -+ name = "stablehlo_builder", -+ srcs = ["StablehloBuilder.cpp"], -+ hdrs = ["StablehloBuilder.h"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ ":mlir_builder", -+ ":stablehlo_builder_inc", -+ "//:stablehlo_broadcast_lowering", -+ "//:stablehlo_ops", -+ "//:stablehlo_type_inference", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:InferTypeOpInterface", -+ "@llvm-project//mlir:Support", -+ ], -+) -+ -+cc_test( -+ name = "stablehlo_builder_test", -+ srcs = ["StablehloBuilderTest.cpp"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ ":func_builder", -+ ":mlir_builder", -+ ":stablehlo_builder", -+ "//:register", -+ "//:stablehlo_ops", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ "@llvm-project//third-party/unittest:gmock", -+ "@llvm-project//third-party/unittest:gtest", -+ ], -+) -+ -+cc_test( -+ name = "chlo_builder_test", -+ srcs = ["ChloBuilderTest.cpp"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ ":chlo_builder", -+ ":func_builder", -+ ":mlir_builder", -+ ":stablehlo_builder", -+ "//:register", -+ "//:stablehlo_broadcast_lowering", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ "@llvm-project//third-party/unittest:gmock", -+ "@llvm-project//third-party/unittest:gtest", -+ ], -+) -+ -+cc_test( -+ name = "mlir_builder_test", -+ srcs = ["MlirBuilderTest.cpp"], -+ compatible_with = get_compatible_with_portable(), -+ includes = ["../../../.."], -+ deps = [ -+ ":attr_type_builder_util", -+ ":chlo_builder", -+ ":func_builder", -+ ":mlir_builder", -+ ":stablehlo_builder", -+ "//:register", -+ "//:stablehlo_ops", -+ "@llvm-project//llvm:Support", -+ "@llvm-project//mlir:FuncDialect", -+ "@llvm-project//mlir:IR", -+ "@llvm-project//mlir:Support", -+ "@llvm-project//mlir:TosaDialect", -+ "@llvm-project//third-party/unittest:gmock", -+ "@llvm-project//third-party/unittest:gtest", -+ ], -+) -diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/ChloBuilder.cpp b/stablehlo/stablehlo/integrations/cpp/builder/ChloBuilder.cpp ---- stablehlo/stablehlo/integrations/cpp/builder/ChloBuilder.cpp -+++ stablehlo/stablehlo/integrations/cpp/builder/ChloBuilder.cpp -@@ -13,7 +13,7 @@ - limitations under the License. - ==============================================================================*/ - --#include "ChloBuilder.h" -+#include "stablehlo/integrations/cpp/builder/ChloBuilder.h" - - #include - -diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/MlirBuilder.cpp b/stablehlo/stablehlo/integrations/cpp/builder/MlirBuilder.cpp ---- stablehlo/stablehlo/integrations/cpp/builder/MlirBuilder.cpp -+++ stablehlo/stablehlo/integrations/cpp/builder/MlirBuilder.cpp -@@ -13,7 +13,7 @@ - limitations under the License. - ==============================================================================*/ - --#include "MlirBuilder.h" -+#include "stablehlo/integrations/cpp/builder/MlirBuilder.h" - - #include "mlir/IR/Location.h" - #include "mlir/IR/Value.h" diff --ruN a/stablehlo/stablehlo/integrations/python/StablehloApi.cpp b/stablehlo/stablehlo/integrations/python/StablehloApi.cpp --- stablehlo/stablehlo/integrations/python/StablehloApi.cpp +++ stablehlo/stablehlo/integrations/python/StablehloApi.cpp @@ -3927,267 +1433,6 @@ diff --ruN a/stablehlo/stablehlo/reference/Value.cpp b/stablehlo/stablehlo/refer if (isTensor()) getTensor().print(os); else if (isToken()) -diff --ruN a/stablehlo/stablehlo/testdata/BUILD.bazel b/stablehlo/stablehlo/testdata/BUILD.bazel ---- stablehlo/stablehlo/testdata/BUILD.bazel -+++ stablehlo/stablehlo/testdata/BUILD.bazel -@@ -11,14 +11,27 @@ - # 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. -+ - load("@bazel_skylib//rules:expand_template.bzl", "expand_template") --load("@llvm-project//llvm:lit_test.bzl", "lit_test") -+load("//build_tools:lit_test_suite.bzl", "lit_test_suite") - - package( -- default_visibility = ["//visibility:public"], -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], - licenses = ["notice"], - ) - -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ - # Equivalent of configure_lit_site_cfg from CMakeLists.txt. - expand_template( - name = "lit_site_cfg_py_gen", -@@ -39,46 +52,42 @@ - template = "lit.site.cfg.py.in", - ) - --# Equivalent of add_lit_testsuite from CMakeLists.txt. --[ -- lit_test( -- name = "%s.test" % src, -- size = "small", -- srcs = [src], -- data = [ -- "lit.cfg.py", -- "lit.site.cfg.py", -- "//:stablehlo-opt", -- "//:stablehlo-translate", -- "@llvm-project//llvm:FileCheck", -- ], -- tags = ["stablehlo_testdata_tests"], -- deps = ["@rules_python//python/runfiles"], -- ) -- for src in glob( -+lit_test_suite( -+ name = "stablehlo_testdata_small_tests", -+ srcs = glob( - ["**/*.mlir"], - exclude = ["**/*.large.mlir"], -- ) --] -+ ), -+ data = [ -+ "lit.cfg.py", -+ "lit.site.cfg.py", -+ "//:stablehlo-opt", -+ "//:stablehlo-translate", -+ "@llvm-project//llvm:FileCheck", -+ ], -+ size_per_test = "small", -+ tags = [ -+ "stablehlo_testdata_small_tests", -+ "stablehlo_testdata_tests", -+ ], -+) - --# Give large tests more RAM / timeout --[ -- lit_test( -- name = "%s.test" % src, -- size = "large", -- srcs = [src], -- data = [ -- "lit.cfg.py", -- "lit.site.cfg.py", -- "//:stablehlo-opt", -- "//:stablehlo-translate", -- "@llvm-project//llvm:FileCheck", -- ], -- tags = ["stablehlo_testdata_tests"], -- deps = ["@rules_python//python/runfiles"], -- ) -- for src in glob(["**/*.large.mlir"]) --] -+lit_test_suite( -+ name = "stablehlo_testdata_large_tests", -+ srcs = glob(["**/*.large.mlir"]), -+ data = [ -+ "lit.cfg.py", -+ "lit.site.cfg.py", -+ "//:stablehlo-opt", -+ "//:stablehlo-translate", -+ "@llvm-project//llvm:FileCheck", -+ ], -+ size_per_test = "large", -+ tags = [ -+ "stablehlo_testdata_large_tests", -+ "stablehlo_testdata_tests", -+ ], -+) - - test_suite( - name = "stablehlo_testdata_tests", -diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/BUILD.bazel ---- stablehlo/stablehlo/tests/BUILD.bazel -+++ stablehlo/stablehlo/tests/BUILD.bazel -@@ -11,16 +11,30 @@ - # 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. -+ - load("@bazel_skylib//rules:expand_template.bzl", "expand_template") --load("@llvm-project//llvm:lit_test.bzl", "lit_test") - load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library") - load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") -+load("//build_tools:lit_test_suite.bzl", "lit_test_suite") -+load("//build_tools:utils.bzl", "get_compatible_with_portable") - - package( -- default_visibility = ["//visibility:public"], -+ # copybara:uncomment default_applicable_licenses = ["//:license"], -+ default_visibility = [ -+ "//visibility:public", # copybara:comment -+ # copybara:uncomment "//:stablehlo_friends", -+ ], - licenses = ["notice"], - ) - -+# copybara:uncomment_begin -+# filegroup( -+# name = "all_files", -+# srcs = glob(["**"]), -+# visibility = ["//:__subpackages__"], -+# ) -+# copybara:uncomment_end -+ - cc_library( - name = "check_ops", - srcs = [ -@@ -29,6 +43,7 @@ - hdrs = [ - "CheckOps.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - ":check_ops_inc_gen", -@@ -67,6 +82,7 @@ - - gentbl_cc_library( - name = "check_ops_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - tbl_outs = { - "CheckOps.cpp.inc": ["-gen-op-defs"], -@@ -84,6 +100,7 @@ - srcs = [ - "CheckOps.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - "//:base_td_files", -@@ -98,6 +115,7 @@ - hdrs = [ - "TestUtils.h", - ], -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", - deps = [ - ":test_utils_inc_gen", -@@ -120,11 +138,14 @@ - - gentbl_cc_library( - name = "test_utils_inc_gen", -+ compatible_with = get_compatible_with_portable(), - strip_include_prefix = ".", -- tbl_outs = {"TestUtils.h.inc": [ -- "-gen-pass-decls", -- "-name=HloTest", -- ]}, -+ tbl_outs = { -+ "TestUtils.h.inc": [ -+ "-gen-pass-decls", -+ "-name=HloTest", -+ ], -+ }, - tblgen = "@llvm-project//mlir:mlir-tblgen", - td_file = "TestUtils.td", - deps = [ -@@ -137,6 +158,7 @@ - srcs = [ - "TestUtils.td", - ], -+ compatible_with = get_compatible_with_portable(), - includes = ["."], - deps = [ - "@llvm-project//mlir:PassBaseTdFiles", -@@ -163,31 +185,27 @@ - template = "lit.site.cfg.py.in", - ) - --# Equivalent of add_lit_testsuite from CMakeLists.txt. --[ -- lit_test( -- name = "%s.test" % src, -- size = "small", -- srcs = [src], -- data = [ -- "lit.cfg.py", -- "lit.site.cfg.py", -- "//:stablehlo-opt", -- "//:stablehlo-translate", -- "@llvm-project//llvm:FileCheck", -- "@llvm-project//llvm:not", -- # Allow empty since not all `.mlir` tests have `.mlir.bc` test files. -- ] + glob( -- ["%s.bc" % src], -- allow_empty = True, -- ), -- tags = ["stablehlo_tests"], -- deps = ["@rules_python//python/runfiles"], -- ) -- for src in glob(["**/*.mlir"]) --] -- --test_suite( -+lit_test_suite( - name = "stablehlo_tests", -+ srcs = glob([ -+ "*.mlir", -+ "**/*.mlir", -+ ]), -+ data = [ -+ "lit.cfg.py", -+ "lit.site.cfg.py", -+ "//:stablehlo-opt", -+ "//:stablehlo-translate", -+ "@llvm-project//llvm:FileCheck", -+ "@llvm-project//llvm:not", -+ ], -+ size_per_test = "small", - tags = ["stablehlo_tests"], - ) -+ -+test_suite( -+ name = "all_tests", # alias for `:stablehlo_tests` -+ tests = [ -+ "//stablehlo/tests:stablehlo_tests", -+ ], -+) diff --ruN a/stablehlo/stablehlo/tests/CheckOps.cpp b/stablehlo/stablehlo/tests/CheckOps.cpp --- stablehlo/stablehlo/tests/CheckOps.cpp +++ stablehlo/stablehlo/tests/CheckOps.cpp @@ -4348,196 +1593,6 @@ diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/st // CHECK-LABEL: func.func @mulhi_i32( // CHECK-SAME: %[[ARG0:.*]]: tensor<4xi32>, // CHECK-SAME: %[[ARG1:.*]]: tensor<4xi32>) -> tensor<4xi32> { -diff --ruN a/stablehlo/stablehlo/tests/interpret/all_gather.mlir b/stablehlo/stablehlo/tests/interpret/all_gather.mlir ---- stablehlo/stablehlo/tests/interpret/all_gather.mlir -+++ stablehlo/stablehlo/tests/interpret/all_gather.mlir -@@ -133,3 +133,41 @@ - func.return - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @all_gather(%operand : tensor<1xi64>) -> tensor<2xi64> { -+ %result = "stablehlo.all_gather"(%operand) { -+ all_gather_dim = 0 : i64, -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<1xi64>) -> tensor<2xi64> -+ return %result : tensor<2xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[0]> : tensor<1xi64> -+ %p1 = stablehlo.constant dense<[1]> : tensor<1xi64> -+ %p2 = stablehlo.constant dense<[2]> : tensor<1xi64> -+ %p3 = stablehlo.constant dense<[3]> : tensor<1xi64> -+ %p4 = stablehlo.constant dense<[4]> : tensor<1xi64> -+ %p5 = stablehlo.constant dense<[5]> : tensor<1xi64> -+ %p6 = stablehlo.constant dense<[6]> : tensor<1xi64> -+ %p7 = stablehlo.constant dense<[7]> : tensor<1xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@all_gather], [@all_gather], [@all_gather], [@all_gather], -+ [@all_gather], [@all_gather], [@all_gather], [@all_gather]] -+ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> -+ (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, -+ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -+ check.expect_eq_const %results#0, dense<[0, 2]> : tensor<2xi64> -+ check.expect_eq_const %results#1, dense<[1, 3]> : tensor<2xi64> -+ check.expect_eq_const %results#2, dense<[0, 2]> : tensor<2xi64> -+ check.expect_eq_const %results#3, dense<[1, 3]> : tensor<2xi64> -+ check.expect_eq_const %results#4, dense<[4, 6]> : tensor<2xi64> -+ check.expect_eq_const %results#5, dense<[5, 7]> : tensor<2xi64> -+ check.expect_eq_const %results#6, dense<[4, 6]> : tensor<2xi64> -+ check.expect_eq_const %results#7, dense<[5, 7]> : tensor<2xi64> -+ func.return -+ } -+} -diff --ruN a/stablehlo/stablehlo/tests/interpret/all_reduce.mlir b/stablehlo/stablehlo/tests/interpret/all_reduce.mlir ---- stablehlo/stablehlo/tests/interpret/all_reduce.mlir -+++ stablehlo/stablehlo/tests/interpret/all_reduce.mlir -@@ -135,3 +135,45 @@ - func.return - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @all_reduce(%operand : tensor<1xi64>) -> tensor<1xi64> { -+ %result = "stablehlo.all_reduce"(%operand) ({ -+ ^bb0(%arg0: tensor, %arg1: tensor): -+ %0 = stablehlo.add %arg0, %arg1 : tensor -+ stablehlo.return %0 : tensor -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]>, -+ channel_handle = #stablehlo.channel_handle -+ } : (tensor<1xi64>) -> tensor<1xi64> -+ return %result : tensor<1xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[10]> : tensor<1xi64> -+ %p1 = stablehlo.constant dense<[20]> : tensor<1xi64> -+ %p2 = stablehlo.constant dense<[30]> : tensor<1xi64> -+ %p3 = stablehlo.constant dense<[40]> : tensor<1xi64> -+ %p4 = stablehlo.constant dense<[50]> : tensor<1xi64> -+ %p5 = stablehlo.constant dense<[60]> : tensor<1xi64> -+ %p6 = stablehlo.constant dense<[70]> : tensor<1xi64> -+ %p7 = stablehlo.constant dense<[80]> : tensor<1xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce], -+ [@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce]] -+ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> -+ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -+ check.expect_eq_const %results#0, dense<[40]> : tensor<1xi64> -+ check.expect_eq_const %results#1, dense<[60]> : tensor<1xi64> -+ check.expect_eq_const %results#2, dense<[40]> : tensor<1xi64> -+ check.expect_eq_const %results#3, dense<[60]> : tensor<1xi64> -+ check.expect_eq_const %results#4, dense<[120]> : tensor<1xi64> -+ check.expect_eq_const %results#5, dense<[140]> : tensor<1xi64> -+ check.expect_eq_const %results#6, dense<[120]> : tensor<1xi64> -+ check.expect_eq_const %results#7, dense<[140]> : tensor<1xi64> -+ func.return -+ } -+} -diff --ruN a/stablehlo/stablehlo/tests/interpret/all_to_all.mlir b/stablehlo/stablehlo/tests/interpret/all_to_all.mlir ---- stablehlo/stablehlo/tests/interpret/all_to_all.mlir -+++ stablehlo/stablehlo/tests/interpret/all_to_all.mlir -@@ -172,3 +172,43 @@ - func.return %results#0, %results#1, %results#2, %results#3 : tensor<4x2xi64>, tensor<6x2xi32>, tensor<4x2xi64>, tensor<6x2xi32> - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @all_to_all(%operand : tensor<2x1xi64>) -> tensor<1x2xi64> { -+ %result = "stablehlo.all_to_all"(%operand) { -+ split_dimension = 0 : i64, -+ concat_dimension = 1 : i64, -+ split_count = 2 : i64, -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<2x1xi64>) -> tensor<1x2xi64> -+ return %result : tensor<1x2xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[[1], [2]]> : tensor<2x1xi64> -+ %p1 = stablehlo.constant dense<[[3], [4]]> : tensor<2x1xi64> -+ %p2 = stablehlo.constant dense<[[10], [20]]> : tensor<2x1xi64> -+ %p3 = stablehlo.constant dense<[[30], [40]]> : tensor<2x1xi64> -+ %p4 = stablehlo.constant dense<[[5], [6]]> : tensor<2x1xi64> -+ %p5 = stablehlo.constant dense<[[7], [8]]> : tensor<2x1xi64> -+ %p6 = stablehlo.constant dense<[[50], [60]]> : tensor<2x1xi64> -+ %p7 = stablehlo.constant dense<[[70], [80]]> : tensor<2x1xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all], -+ [@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all]] -+ } : (tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, -+ tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>) -> -+ (tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, -+ tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>) -+ check.expect_eq_const %results#0, dense<[[1, 10]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#1, dense<[[3, 30]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#2, dense<[[2, 20]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#3, dense<[[4, 40]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#4, dense<[[5, 50]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#5, dense<[[7, 70]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#6, dense<[[6, 60]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#7, dense<[[8, 80]]> : tensor<1x2xi64> -+ func.return -+ } -+} -diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir ---- stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir -+++ stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir -@@ -90,3 +90,45 @@ - func.return - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @reduce_scatter(%operand : tensor<2xi64>) -> tensor<1xi64> { -+ %result = "stablehlo.reduce_scatter"(%operand) ({ -+ ^bb0(%arg0: tensor, %arg1: tensor): -+ %0 = stablehlo.add %arg0, %arg1 : tensor -+ stablehlo.return %0 : tensor -+ }) { -+ scatter_dimension = 0 : i64, -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<2xi64>) -> tensor<1xi64> -+ return %result : tensor<1xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[1, 2]> : tensor<2xi64> -+ %p1 = stablehlo.constant dense<[3, 4]> : tensor<2xi64> -+ %p2 = stablehlo.constant dense<[10, 20]> : tensor<2xi64> -+ %p3 = stablehlo.constant dense<[30, 40]> : tensor<2xi64> -+ %p4 = stablehlo.constant dense<[5, 6]> : tensor<2xi64> -+ %p5 = stablehlo.constant dense<[7, 8]> : tensor<2xi64> -+ %p6 = stablehlo.constant dense<[50, 60]> : tensor<2xi64> -+ %p7 = stablehlo.constant dense<[70, 80]> : tensor<2xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], -+ [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter]] -+ } : (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, -+ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> -+ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -+ check.expect_eq_const %results#0, dense<[11]> : tensor<1xi64> -+ check.expect_eq_const %results#1, dense<[33]> : tensor<1xi64> -+ check.expect_eq_const %results#2, dense<[22]> : tensor<1xi64> -+ check.expect_eq_const %results#3, dense<[44]> : tensor<1xi64> -+ check.expect_eq_const %results#4, dense<[55]> : tensor<1xi64> -+ check.expect_eq_const %results#5, dense<[77]> : tensor<1xi64> -+ check.expect_eq_const %results#6, dense<[66]> : tensor<1xi64> -+ check.expect_eq_const %results#7, dense<[88]> : tensor<1xi64> -+ func.return -+ } -+} diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir --- stablehlo/stablehlo/tests/ops_chlo.mlir +++ stablehlo/stablehlo/tests/ops_chlo.mlir @@ -4626,39 +1681,13 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir b/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir --- stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir +++ stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir -@@ -1,11 +1,11 @@ +@@ -1,4 +1,4 @@ -// RUN: stablehlo-opt %s -stablehlo-compatibility-expander="target=1.14.0" -allow-unregistered-dialect | FileCheck %s +// RUN: stablehlo-opt %s -stablehlo-compatibility-expander="target=1.14.0" | FileCheck %s module { - - // CHECK-LABEL: @all_reduce_rgv3 - func.func @all_reduce_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_reduce"(%arg0) ({ - ^bb0(%arg1: tensor, %arg2: tensor): - %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -@@ -19,7 +19,7 @@ - - // CHECK-LABEL: @all_gather_rgv3 - func.func @all_gather_rgv3(%arg0: tensor<4xf32>) -> tensor<8xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_gather"(%arg0) { - all_gather_dim = 0 : i64, - replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -@@ -30,7 +30,7 @@ - - // CHECK-LABEL: @all_to_all_rgv3 - func.func @all_to_all_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_to_all"(%arg0) { - concat_dimension = 0 : i64, - split_dimension = 0 : i64, -@@ -40,11 +40,11 @@ +@@ -40,7 +40,7 @@ return %0 : tensor<4xf32> } @@ -4667,12 +1696,7 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compat // CHECK-LABEL: @all_reduce_sdy_mesh func.func @all_reduce_sdy_mesh(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_reduce"(%arg0) ({ - ^bb0(%arg1: tensor, %arg2: tensor): - %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -@@ -55,11 +55,11 @@ +@@ -55,7 +55,7 @@ return %0 : tensor<4xf32> } @@ -4681,11 +1705,6 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compat // CHECK-LABEL: @all_reduce_sdy_mesh_dev func.func @all_reduce_sdy_mesh_dev(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_reduce"(%arg0) ({ - ^bb0(%arg1: tensor, %arg2: tensor): - %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor @@ -70,7 +70,7 @@ return %0 : tensor<4xf32> } @@ -4695,63 +1714,6 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compat // CHECK-LABEL: @all_reduce_sdy_mesh_max func.func @all_reduce_sdy_mesh_max(%arg0: tensor<4xf32>) -> tensor<4xf32> { -@@ -84,4 +84,56 @@ - } : (tensor<4xf32>) -> tensor<4xf32> - return %0 : tensor<4xf32> - } -+ -+ // CHECK-LABEL: @all_reduce_subaxis -+ func.func @all_reduce_subaxis(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3], [4, 6], [5, 7]]> : tensor<4x2xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } -+ -+ // CHECK-LABEL: @all_reduce_subaxis_order_1 -+ func.func @all_reduce_subaxis_order_1(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]]> : tensor<3x10xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } -+ -+ // CHECK-LABEL: @all_reduce_subaxis_order_2 -+ func.func @all_reduce_subaxis_order_2(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 5, 1, 6, 2, 7, 3, 8, 4, 9], [10, 15, 11, 16, 12, 17, 13, 18, 14, 19], [20, 25, 21, 26, 22, 27, 23, 28, 24, 29]]> : tensor<3x10xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } -+ -+ // CHECK-LABEL: @all_reduce_size_one_axis -+ func.func @all_reduce_size_one_axis(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } - } diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir --- stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir +++ stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir diff --git a/third_party/xla/third_party/stablehlo/workspace.bzl b/third_party/xla/third_party/stablehlo/workspace.bzl index c8508914e189ec..8764fbd0a3f908 100644 --- a/third_party/xla/third_party/stablehlo/workspace.bzl +++ b/third_party/xla/third_party/stablehlo/workspace.bzl @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def repo(): # LINT.IfChange - STABLEHLO_COMMIT = "639335932274930617d9724fb25ec9f2f52c5e40" - STABLEHLO_SHA256 = "dd590949798e26838d9ac3d6016def9b508caa60d0d7a69535ed2f5454ad30b8" + STABLEHLO_COMMIT = "b9029ac2228a12e176ea20e56f7d3dcb6bf8a019" + STABLEHLO_SHA256 = "16af6c8f0baacda15dd72d98dc295a213be455b8fa1899af46e144d6dae670bd" # LINT.ThenChange(Google-internal path) tf_http_archive( diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index bfe3b67572bec5..d96b6dfca36538 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -3500,6 +3500,7 @@ cc_library( "//xla/runtime:buffer_use", "//xla/runtime:device_id", "//xla/service:buffer_assignment", + "//xla/service:shaped_slice", "//xla/stream_executor:device_address", "//xla/stream_executor:event", "//xla/stream_executor:stream", diff --git a/third_party/xla/xla/backends/gpu/runtime/async_thunk.h b/third_party/xla/xla/backends/gpu/runtime/async_thunk.h index df5c119602697a..419167ac764b5c 100644 --- a/third_party/xla/xla/backends/gpu/runtime/async_thunk.h +++ b/third_party/xla/xla/backends/gpu/runtime/async_thunk.h @@ -62,8 +62,8 @@ class AsyncStartThunk : public Thunk { absl::Status Initialize(const InitializeParams& params) override; absl::Status ExecuteOnStream(const ExecuteParams& params) override; - // TODO(b/527907619): Implement this properly once we have figured out how - // buffer uses should look like for async thunks. + // AsyncStartThunk touches no buffers; this is explicit rather than relying on + // the base class default. BufferUses buffer_uses() const override { return {}; } ExecutionStreamId execution_stream_id() const { return execution_stream_id_; } diff --git a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc index 1ac64f4c193a20..d1c0ab247b1ce5 100644 --- a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc @@ -36,6 +36,7 @@ limitations under the License. #include "xla/backends/gpu/runtime/thunk.pb.h" #include "xla/runtime/device_id.h" #include "xla/service/buffer_assignment.h" +#include "xla/service/shaped_slice.h" #include "xla/shape.h" #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/event.h" @@ -107,13 +108,12 @@ HostSendRecvAsyncEvents::Extract(se::StreamExecutor* executor, //===----------------------------------------------------------------------===// HostSendThunk::HostSendThunk( - ThunkInfo thunk_info, Shape shape, BufferAllocation::Slice buffer, - int64_t channel_id, std::shared_ptr events, + ThunkInfo thunk_info, ShapedSlice buffer, int64_t channel_id, + std::shared_ptr events, absl::flat_hash_map frontend_attrs, std::optional device_constraint) : HostAsyncThunk(Thunk::kHostSend, thunk_info), - shape_(shape), - buffer_(buffer), + slice_(buffer), channel_id_(channel_id), events_(std::move(events)), frontend_attrs_(std::move(frontend_attrs)), @@ -123,8 +123,9 @@ absl::StatusOr HostSendThunk::ToProto() const { ThunkProto proto; *proto.mutable_thunk_info() = thunk_info().ToProto(); HostSendThunkProto& host_send_thunk_proto = *proto.mutable_host_send_thunk(); - *host_send_thunk_proto.mutable_shape() = shape_.ToProto(); - ABSL_ASSIGN_OR_RETURN(*host_send_thunk_proto.mutable_buffer(), buffer_.ToProto()); + *host_send_thunk_proto.mutable_shape() = slice_.shape.ToProto(); + ABSL_ASSIGN_OR_RETURN(*host_send_thunk_proto.mutable_buffer(), + slice_.slice.ToProto()); host_send_thunk_proto.set_channel_id(channel_id_); host_send_thunk_proto.mutable_frontend_attrs()->insert( frontend_attrs_.begin(), frontend_attrs_.end()); @@ -160,13 +161,13 @@ absl::StatusOr> HostSendThunk::FromProto( AsyncEventsUniqueId(proto.async_events_unique_id()), std::make_shared()); return std::make_unique( - thunk_info, std::move(shape), buffer, proto.channel_id(), + thunk_info, ShapedSlice{buffer, shape}, proto.channel_id(), async_event_it->second, std::move(frontend_attrs), device_constraint); } absl::Status HostSendThunk::ExecuteOnStream(const ExecuteParams& params) { VLOG(3) << "Send buffer: channel_id=" << channel_id_ - << "; shape=" << shape_.ToString(); + << "; shape=" << slice_.shape.ToString(); ABSL_ASSIGN_OR_RETURN(bool skip, ShouldSkip("sending buffer", params, device_constraint_)); @@ -186,13 +187,13 @@ absl::Status HostSendThunk::ExecuteOnStream(const ExecuteParams& params) { } se::DeviceAddressBase src = - params.buffer_allocations->GetDeviceAddress(buffer_); + params.buffer_allocations->GetDeviceAddress(slice_.slice); // Send buffer to a handler registered with the executable. if (auto* send = params.send_device_memory_function) { ABSL_ASSIGN_OR_RETURN( AsyncValueRef> done, - (*send)(channel_id_, stream, shape_, src, frontend_attrs_)); + (*send)(channel_id_, stream, slice_.shape, src, frontend_attrs_)); return events_->Emplace(stream->parent(), channel_id_, std::move(done)); } @@ -212,15 +213,15 @@ std::optional HostSendThunk::GetAsyncEventsUniqueId() //===----------------------------------------------------------------------===// // HostSendDoneThunk //===----------------------------------------------------------------------===// - HostSendDoneThunk::HostSendDoneThunk( - ThunkInfo thunk_info, int64_t channel_id, + ThunkInfo thunk_info, std::optional slice, int64_t channel_id, std::shared_ptr events, std::optional device_constraint) : HostAsyncThunk(Thunk::kHostSendDone, thunk_info), channel_id_(channel_id), events_(std::move(events)), - device_constraint_(device_constraint) {} + device_constraint_(device_constraint), + slice_(slice) {} absl::StatusOr HostSendDoneThunk::ToProto() const { ThunkProto proto; @@ -239,6 +240,10 @@ absl::StatusOr HostSendDoneThunk::ToProto() const { } host_send_done_thunk_proto.set_async_events_unique_id( async_events_unique_id.value().value()); + if (slice_.has_value()) { + ABSL_ASSIGN_OR_RETURN(*host_send_done_thunk_proto.mutable_buffer(), + slice_->ToProto()); + } return proto; } @@ -255,9 +260,15 @@ absl::StatusOr> HostSendDoneThunk::FromProto( AsyncEventsUniqueId(proto.async_events_unique_id()), std::make_shared()); - return std::make_unique(thunk_info, proto.channel_id(), - async_event_it->second, - device_constraint); + std::optional slice; + if (proto.has_buffer()) { + ABSL_ASSIGN_OR_RETURN(slice, + ShapedSlice::FromProto(proto.buffer(), allocations)); + } + + return std::make_unique( + thunk_info, slice, proto.channel_id(), async_event_it->second, + device_constraint); } absl::Status HostSendDoneThunk::ExecuteOnStream(const ExecuteParams& params) { @@ -408,13 +419,14 @@ std::optional HostRecvThunk::GetAsyncEventsUniqueId() //===----------------------------------------------------------------------===// HostRecvDoneThunk::HostRecvDoneThunk( - ThunkInfo thunk_info, int64_t channel_id, + ThunkInfo thunk_info, std::optional slice, int64_t channel_id, std::shared_ptr events, std::optional device_constraint) : HostAsyncThunk(Thunk::kHostRecvDone, thunk_info), channel_id_(channel_id), events_(std::move(events)), - device_constraint_(device_constraint) {} + device_constraint_(device_constraint), + slice_(slice) {} absl::StatusOr HostRecvDoneThunk::ToProto() const { ThunkProto proto; @@ -433,6 +445,11 @@ absl::StatusOr HostRecvDoneThunk::ToProto() const { } host_recv_done_thunk_proto.set_async_events_unique_id( async_events_unique_id.value().value()); + + if (slice_.has_value()) { + ABSL_ASSIGN_OR_RETURN(*host_recv_done_thunk_proto.mutable_buffer(), + slice_->ToProto()); + } return proto; } @@ -449,9 +466,15 @@ absl::StatusOr> HostRecvDoneThunk::FromProto( AsyncEventsUniqueId(proto.async_events_unique_id()), std::make_shared()); - return std::make_unique(thunk_info, proto.channel_id(), - async_event_it->second, - device_constraint); + std::optional slice; + if (proto.has_buffer()) { + ABSL_ASSIGN_OR_RETURN(slice, + ShapedSlice::FromProto(proto.buffer(), allocations)); + } + + return std::make_unique( + thunk_info, slice, proto.channel_id(), async_event_it->second, + device_constraint); } absl::Status HostRecvDoneThunk::ExecuteOnStream(const ExecuteParams& params) { diff --git a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h index a3517382828389..9a83f160943c35 100644 --- a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h +++ b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h @@ -34,6 +34,7 @@ limitations under the License. #include "xla/runtime/buffer_use.h" #include "xla/runtime/device_id.h" #include "xla/service/buffer_assignment.h" +#include "xla/service/shaped_slice.h" #include "xla/shape.h" #include "xla/stream_executor/event.h" #include "xla/stream_executor/stream_executor.h" @@ -100,8 +101,7 @@ class HostSendThunk : public HostAsyncThunk { absl::Span allocations, HostSendRecvAsyncEventsMap& async_events_map); - HostSendThunk(ThunkInfo thunk_info, Shape shape, - BufferAllocation::Slice buffer, int64_t channel_id, + HostSendThunk(ThunkInfo thunk_info, ShapedSlice buffer, int64_t channel_id, std::shared_ptr events, absl::flat_hash_map frontend_attrs, std::optional device_constraint); @@ -110,7 +110,7 @@ class HostSendThunk : public HostAsyncThunk { BufferUses buffer_uses() const override { return { - BufferUse::Read(buffer_, shape_), + BufferUse::Read(slice_.slice, slice_.shape), }; } @@ -121,8 +121,7 @@ class HostSendThunk : public HostAsyncThunk { bool IsAsyncStart() const override { return events_ != nullptr; } private: - Shape shape_; - BufferAllocation::Slice buffer_; + ShapedSlice slice_; int64_t channel_id_; @@ -142,15 +141,22 @@ class HostSendDoneThunk : public HostAsyncThunk { absl::Span allocations, HostSendRecvAsyncEventsMap& async_events_map); - HostSendDoneThunk(ThunkInfo thunk_info, int64_t channel_id, + HostSendDoneThunk(ThunkInfo thunk_info, std::optional slice, + int64_t channel_id, std::shared_ptr events, std::optional device_constraint); absl::Status ExecuteOnStream(const ExecuteParams& params) override; - // TODO(b/527907619): Implement this properly once we have figured out how - // buffer uses should look like for async thunks. - BufferUses buffer_uses() const override { return {}; } + BufferUses buffer_uses() const override { + // TODO(b/527907619): make mandatory after 2027-1-1 + if (!slice_.has_value()) { + return {}; + } + return { + BufferUse::Read(slice_->slice, slice_->shape), + }; + } absl::StatusOr ToProto() const override; @@ -163,12 +169,12 @@ class HostSendDoneThunk : public HostAsyncThunk { std::shared_ptr events_; std::optional device_constraint_; + std::optional slice_; }; //===----------------------------------------------------------------------===// // HostRecvThunk //===----------------------------------------------------------------------===// - class HostRecvThunk : public HostAsyncThunk { public: static absl::StatusOr> FromProto( @@ -218,15 +224,22 @@ class HostRecvDoneThunk : public HostAsyncThunk { absl::Span allocations, HostSendRecvAsyncEventsMap& async_events_map); - HostRecvDoneThunk(ThunkInfo thunk_info, int64_t channel_id, + HostRecvDoneThunk(ThunkInfo thunk_info, std::optional slice, + int64_t channel_id, std::shared_ptr events, std::optional device_constraint); absl::Status ExecuteOnStream(const ExecuteParams& params) override; - // TODO(b/527907619): Implement this properly once we have figured out how - // buffer uses should look like for async thunks. - BufferUses buffer_uses() const override { return {}; } + BufferUses buffer_uses() const override { + // TODO(b/527907619): make mandatory after 2027-1-1 + if (!slice_.has_value()) { + return {}; + } + return { + BufferUse::Write(slice_->slice, slice_->shape), + }; + } absl::StatusOr ToProto() const override; @@ -239,6 +252,7 @@ class HostRecvDoneThunk : public HostAsyncThunk { std::shared_ptr events_; std::optional device_constraint_; + std::optional slice_; }; } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk.proto b/third_party/xla/xla/backends/gpu/runtime/thunk.proto index 24bbab38e2aed7..b32f8d74310d25 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk.proto +++ b/third_party/xla/xla/backends/gpu/runtime/thunk.proto @@ -149,6 +149,7 @@ message HostSendDoneThunkProto { int64 channel_id = 1; optional int64 device_constraint = 2; uint64 async_events_unique_id = 3; + ShapedSliceProto buffer = 4; } message HostRecvThunkProto { @@ -164,6 +165,7 @@ message HostRecvDoneThunkProto { int64 channel_id = 1; optional int64 device_constraint = 2; uint64 async_events_unique_id = 3; + ShapedSliceProto buffer = 4; } message HostExecuteStartThunkProto { diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc index bb6f8587dd4448..7e0e395940053f 100644 --- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc @@ -1698,9 +1698,7 @@ absl::StatusOr CudnnFusedConvRewriter::RunImpl( auto* cc = compute_capability_.cuda_compute_capability(); ABSL_ASSIGN_OR_RETURN(changed, F8GraphConv(comp, *cc, dnn_version_, toolkit_version_)); - if (changed) { - return changed; - } + any_changed |= changed; } // Fuse "inside out" starting with the operations closest to the conv. ABSL_ASSIGN_OR_RETURN(changed, FuseRemoveConvertInConv(comp)); diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter_test.cc index c5205ee226492b..d55df350ed27a2 100644 --- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fused_conv_rewriter_test.cc @@ -37,6 +37,7 @@ limitations under the License. #include "xla/comparison_util.h" #include "xla/debug_options_flags.h" #include "xla/error_spec.h" +#include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_print_options.h" @@ -907,6 +908,84 @@ TEST_F(CudnnFusedConvRewriterTest, TestConvF8) { )"); } +// The pass used to stop after the first computation in which an FP8 +// convolution was rewritten, leaving the FP8 convolutions of the remaining +// computations of the module unrewritten. +TEST_F(CudnnFusedConvRewriterTest, TestConvF8InMultipleComputations) { + MAYBE_SKIP_TEST("F8"); + const std::string kHloString = R"( + HloModule Test + + branch_a { + p = (f8e4m3fn[1,128,6,6], f8e4m3fn[3,3,128,16]) parameter(0) + input_a = f8e4m3fn[1,128,6,6] get-tuple-element(p), index=0 + filter_a = f8e4m3fn[3,3,128,16] get-tuple-element(p), index=1 + ROOT conv_a = f8e4m3fn[1,16,6,6] convolution(input_a, filter_a), window={size=3x3 pad=1_1x1_1}, dim_labels=bf01_01io->bf01, feature_group_count=1 + } + + branch_b { + p = (f8e4m3fn[1,128,6,6], f8e4m3fn[3,3,128,16]) parameter(0) + input_b = f8e4m3fn[1,128,6,6] get-tuple-element(p), index=0 + filter_b = f8e4m3fn[3,3,128,16] get-tuple-element(p), index=1 + ROOT conv_b = f8e4m3fn[1,16,6,6] convolution(input_b, filter_b), window={size=3x3 pad=1_1x1_1}, dim_labels=bf01_01io->bf01, feature_group_count=1 + } + + ENTRY Test { + predicate = pred[] parameter(0) + input = f8e4m3fn[1,128,6,6] parameter(1) + filter_a = f8e4m3fn[3,3,128,16] parameter(2) + filter_b = f8e4m3fn[3,3,128,16] parameter(3) + operands_a = (f8e4m3fn[1,128,6,6], f8e4m3fn[3,3,128,16]) tuple(input, filter_a) + operands_b = (f8e4m3fn[1,128,6,6], f8e4m3fn[3,3,128,16]) tuple(input, filter_b) + ROOT result = f8e4m3fn[1,16,6,6] conditional(predicate, operands_a, operands_b), true_computation=branch_a, false_computation=branch_b + })"; + + // Both branch computations must get a ForwardGraph Custom Call. + const se::CudaComputeCapability hopper{se::CudaComputeCapability::kHopper, 0}; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHloString)); + ASSERT_OK_AND_ASSIGN(bool changed, + RunHloPass(ConvRewriter(hopper), module.get())); + EXPECT_TRUE(changed); + ASSERT_OK_AND_ASSIGN( + changed, RunHloPass(CudnnFusedConvRewriter(hopper, GetDnnVersion(), + GetToolkitVersion()), + module.get())); + EXPECT_TRUE(changed); + for (absl::string_view name : {"branch_a", "branch_b"}) { + HloComputation* branch = module->GetComputationWithName(name); + ASSERT_THAT(branch, NotNull()) << name; + const HloInstruction* conv = nullptr; + for (const HloInstruction* instr : branch->instructions()) { + if (instr->opcode() == HloOpcode::kCustomCall) { + conv = instr; + } + } + ASSERT_THAT(conv, NotNull()) << name; + EXPECT_EQ(conv->custom_call_target(), kCudnnConvForwardGraphCallTarget) + << name; + } + + bool fp8_supported = GetDnnVersion() >= se::dnn::VersionInfo{9, 8, 0} + ? GetCudaComputeCapability().IsAtLeastAda() + : GetCudaComputeCapability().IsAtLeastHopper(); + if (fp8_supported) { + // Through the full pipeline, neither convolution may fall back to the + // legacy (non-FP8) Custom Call. + absl::StatusOr filecheck_result = + RunFileCheck(GetOptimizedHlo(kHloString), R"( +// CHECK-NOT: custom_call_target="__cudnn$convForward" +// CHECK: custom_call_target="__cudnn$convForwardGraph" +// CHECK-NOT: custom_call_target="__cudnn$convForward" +// CHECK: custom_call_target="__cudnn$convForwardGraph" +// CHECK-NOT: custom_call_target="__cudnn$convForward" + )"); + ASSERT_TRUE(filecheck_result.ok()) << filecheck_result.status(); + EXPECT_TRUE(*filecheck_result); + EXPECT_TRUE(RunAndCompare(kHloString, ErrorSpec{0.15, 0.15})) << kHloString; + } +} + TEST_F(CudnnFusedConvRewriterTest, TestConvScaledOutputF8) { MAYBE_SKIP_TEST("F8"); TestF8( diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h index acc6b37b43dfeb..42ba52fd318d12 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h @@ -18,6 +18,7 @@ limitations under the License. #define _USE_MATH_DEFINES +#include #include #include #include @@ -812,7 +813,6 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, // Optional handler exercised when evaluating literals. EvalLiteralHandler eval_literal_handler_; - // Set by EvaluateInternal and opportunistically used by the HandleXXX // functions. When non-empty, the HandleXXX function may evaluate the // instruction at only the given shape index. diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD index 15b4cc8d75fd10..cdf5463c660cee 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD @@ -63,7 +63,6 @@ cc_library( "//xla/hlo/ir:hlo", "//xla/mlir_hlo", "//xla/mlir_hlo:unregistered_attributes", - "//xla/tsl/platform:errors", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/async_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/async_importer.cc index c6e2739e5c4012..8349a5716dc71a 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/async_importer.cc +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/async_importer.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -49,7 +50,6 @@ limitations under the License. #include "xla/hlo/translate/hlo_to_mhlo/hlo_utils.h" #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h" #include "xla/mlir_hlo/utils/unregistered_attributes.h" -#include "xla/tsl/platform/errors.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -201,6 +201,79 @@ absl::StatusOr ImportStablehloAsyncDone( return op.getOperation(); } +// Finds the defining MHLO AsyncStartOp by traversing backwards through +// transparent dataflow intermediaries (tuples, optimization barriers, and +// get-tuple-element). Maintains a tuple index stack to correctly navigate +// through multi-level nested tuples. +mlir::mhlo::AsyncStartOp FindMhloAsyncStart(mlir::Value value) { + llvm::SmallVector tuple_indices_stack; + while (value) { + mlir::Operation* defining_op = value.getDefiningOp(); + if (!defining_op) { + return nullptr; + } + + // Base case: Reached the target AsyncStartOp. + if (auto async_start = + llvm::dyn_cast(defining_op)) { + return async_start; + } + + // GetTupleElement (StableHLO / MHLO): Push the tuple index onto the + // stack so that when we encounter the matching TupleOp upstream, we select + // the correct operand branch. + if (auto gte = + llvm::dyn_cast(defining_op)) { + tuple_indices_stack.push_back(gte.getIndex()); + value = gte.getOperand(); + continue; + } + if (auto gte = llvm::dyn_cast(defining_op)) { + tuple_indices_stack.push_back(gte.getIndex()); + value = gte.getOperand(); + continue; + } + + // TupleOp (StableHLO / MHLO): Pop the expected tuple index from the + // stack to step into the corresponding tuple element operand. If no index + // was tracked, the tuple was consumed directly rather than by element, so + // it cannot be an async-start op. + if (llvm::isa(defining_op)) { + if (tuple_indices_stack.empty()) { + return nullptr; + } + int64_t idx = tuple_indices_stack.pop_back_val(); + if (idx >= defining_op->getNumOperands()) { + return nullptr; + } + value = defining_op->getOperand(idx); + continue; + } + + // OptimizationBarrierOp (StableHLO / MHLO): The barrier is transparent; + // trace the corresponding input operand matching this OpResult's result + // number. + if (llvm::isa(defining_op)) { + auto op_result = mlir::dyn_cast(value); + if (!op_result) { + return nullptr; + } + unsigned res_num = op_result.getResultNumber(); + if (res_num >= defining_op->getNumOperands()) { + return nullptr; + } + value = defining_op->getOperand(res_num); + continue; + } + + // Non-transparent defining op; async-start could not be found along this + // path. + return nullptr; + } + return nullptr; +} + absl::StatusOr ImportOldStyleAsyncDone( llvm::SmallVectorImpl& attributes, const llvm::SmallVectorImpl& operands, mlir::Location loc, @@ -208,7 +281,7 @@ absl::StatusOr ImportOldStyleAsyncDone( bool useBundleResult = false) { assert(operands.size() == 1 && "*-done ops must take only a single async_bundle operand"); - auto async_start = operands[0].getDefiningOp(); + auto async_start = FindMhloAsyncStart(operands[0]); if (!async_start) { return InvalidArgument("*-start requires *-done as input"); } @@ -229,7 +302,9 @@ absl::StatusOr ImportOldStyleAsyncDone( operands, attributes); return {op}; } - if (useBundleResult) result_type = async_bundle.getTypes()[1]; + if (useBundleResult) { + result_type = async_bundle.getTypes()[1]; + } auto op = mlir::mhlo::AsyncDoneOp::create(*builder, loc, Untuple(result_type), operands, attributes); return CreateTupleFromOpResults(builder, loc, op.getOperation(), result_type); diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import_async.hlo b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import_async.hlo index 14672943c8e049..d3b603c33adf2a 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import_async.hlo +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import_async.hlo @@ -295,6 +295,33 @@ ENTRY %test_all_gather_start_mesh { ROOT ag-done = f32[128,128] all-gather-done(ag-start) } +// ----- + +HloModule main + +// CHECK-LABEL: func.func private @copy_(%arg0: tensor<128x32xf32>) +// CHECK-NEXT: mhlo.copy + +// CHECK-LABEL: func.func @main +// CHECK-SAME: (%[[ARG0:.*]]: tensor<1024x1024xf32>, %[[ARG1:.*]]: tensor<128x32xf32>) +// CHECK: %[[BARRIER:.*]] = stablehlo.optimization_barrier %[[ARG0]] +// CHECK: %[[START:.*]] = "mhlo.async_start"(%[[ARG1]]) <{called_computation = @copy_, execution_thread = "main"}> +// CHECK: %[[DONE:.*]] = "mhlo.async_done"(%[[START]]) +// CHECK-SAME: called_computation = @copy_ +// CHECK-SAME: execution_thread = "main" +// CHECK: return %[[BARRIER]], %[[DONE]] +ENTRY %test_async_copy_opt_barrier (dot_general.1: f32[1024,1024], shard_map.5: f32[128,32]) -> (f32[1024,1024], f32[128,32]) { + %dot_general.1 = f32[1024,1024]{1,0} parameter(0) + %shard_map.5 = f32[128,32]{1,0} parameter(1) + %opt_bar = f32[1024,1024]{1,0} opt-barrier(%dot_general.1) + %copy-start = (f32[128,32]{1,0}, f32[128,32]{1,0}, u32[]) copy-start(%shard_map.5) + %tuple.4 = (f32[1024,1024]{1,0}, (f32[128,32]{1,0}, f32[128,32]{1,0}, u32[])) tuple(%opt_bar, %copy-start) + %get-tuple-element.0 = f32[1024,1024]{1,0} get-tuple-element(%tuple.4), index=0 + %get-tuple-element.1 = (f32[128,32]{1,0}, f32[128,32]{1,0}, u32[]) get-tuple-element(%tuple.4), index=1 + %copy-done = f32[128,32]{1,0} copy-done(%get-tuple-element.1) + ROOT %root = (f32[1024,1024]{1,0}, f32[128,32]{1,0}) tuple(%get-tuple-element.0, %copy-done) +} + // BROKEN: b/TODO: support roundtrip of async custom calls? // HloModule main, entry_computation_layout={(f32[10]{0})->(f32[20]{0})} diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc index 89a0171789b02c..09432b11e5d080 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -955,6 +955,28 @@ LogicalResult TransposeOp::verify() { return success(); } +OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) { + if (llvm::is_sorted(getPermutation())) { + return getVector(); + } + if (auto cst = dyn_cast_if_present(adaptor.getVector())) { + return cst.reshape(getType()); + } + if (auto prev_op = getVector().getDefiningOp()) { + ArrayRef prev_perm = prev_op.getPermutation(); + ArrayRef curr_perm = getPermutation(); + CHECK_EQ(prev_perm.size(), curr_perm.size()); + SmallVector composed_perm(curr_perm.size()); + for (size_t i = 0; i < curr_perm.size(); ++i) { + composed_perm[i] = prev_perm[curr_perm[i]]; + } + if (llvm::is_sorted(composed_perm)) { + return prev_op.getVector(); + } + } + return nullptr; +} + LogicalResult MemRefBitcastOp::verify() { auto src_ty = getInput().getType(); auto tgt_ty = getType(); @@ -1147,49 +1169,73 @@ LogicalResult CompressStoreVregOp::verify() { return success(); } -LogicalResult VectorCompressStoreOp::verify() { - MemRefType ref_ty = getBase().getType(); +template +LogicalResult verifyCompressStoreExpandLoad( + Op op, VectorType ty, int32_t compress_expand_dim, StringRef vector_name, + StringRef compress_expand_dim_name) { + MemRefType ref_ty = op.getBase().getType(); const int64_t rank = ref_ty.getRank(); - VectorType value_ty = getValueToStore().getType(); - if (value_ty.getRank() != rank) { - return emitOpError("Expected valueToStore to have the same rank as base (") - << rank << "). Got: " << value_ty.getRank() << "."; - } - if (llvm::size(getIndices()) != rank) { - return emitOpError("Expected ") << rank << " indices."; - } - const int32_t compress_dim = getCompressDim(); - if (compress_dim < 0 || compress_dim >= rank) { - return emitOpError("Expected compress_dim to be in [0, ") - << rank << "). Got: " << compress_dim << "."; - } + if (ty.getRank() != rank) { + return op.emitOpError("Expected ") + << vector_name << " to have the same rank as base (" << rank + << "). Got: " << ty.getRank() << "."; + } + if (llvm::size(op.getIndices()) != rank) { + return op.emitOpError("Expected ") << rank << " indices."; + } + if (compress_expand_dim < 0 || compress_expand_dim >= rank) { + return op.emitOpError("Expected ") + << compress_expand_dim_name << " to be in [0, " << rank + << "). Got: " << compress_expand_dim << "."; + } + const StringRef non_target_dim_adjective = + compress_expand_dim_name == "compress_dim" ? "Non-compressed" + : "Non-expanded"; for (int64_t i = 0; i < rank; ++i) { - if (i == compress_dim || ref_ty.isDynamicDim(i)) { + if (i == compress_expand_dim || ref_ty.isDynamicDim(i)) { continue; } - if (value_ty.getDimSize(i) > ref_ty.getDimSize(i)) { - return emitOpError("Non-compressed dimension ") - << i << " of valueToStore goes out of bounds of base (" - << value_ty.getDimSize(i) << " > " << ref_ty.getDimSize(i) << ")."; + if (ty.getDimSize(i) > ref_ty.getDimSize(i)) { + return op.emitOpError(non_target_dim_adjective) + << " dimension " << i << " of " << vector_name + << " goes out of bounds of base (" << ty.getDimSize(i) << " > " + << ref_ty.getDimSize(i) << ")."; } } if (ref_ty.getMemorySpace() && !HasMemorySpace(ref_ty, MemorySpace::kVmem)) { - return emitOpError("Expected base memref to be in VMEM."); + return op.emitOpError("Expected base memref to be in VMEM."); } - if (value_ty.getElementType() != ref_ty.getElementType()) { - return emitOpError("Expected base and valueToStore element type to match"); + if (ty.getElementType() != ref_ty.getElementType()) { + return op.emitOpError("Expected base and ") + << vector_name << " element type to match"; } - // Note: We deliberately do not go through verifyStoreOp, which rejects masked - // stores of non-32-bit element types. The mask here is mandatory and narrow - // element types are allowed. - if (value_ty.getShape() != getMask().getType().getShape()) { - return emitOpError("Expected mask shape to match value shape: (") - << value_ty.getShape() << "). Got: (" - << getMask().getType().getShape() << ")."; + // Note: We deliberately do not go through verifyStoreOp / verifyLoadOp, which + // reject masked stores/loads of non-32-bit element types. The mask here is + // mandatory and narrow element types are allowed. They are rejected later, in + // the vector layout passes. + if (ty.getShape() != op.getMask().getType().getShape()) { + return op.emitOpError("Expected mask shape to match ") + << (vector_name == "valueToStore" ? "value" : "result") + << " shape: (" << ty.getShape() << "). Got: (" + << op.getMask().getType().getShape() << ")."; } return success(); } +LogicalResult VectorCompressStoreOp::verify() { + return verifyCompressStoreExpandLoad( + *this, getValueToStore().getType(), getCompressDim(), + /*vector_name=*/"valueToStore", + /*compress_expand_dim_name=*/"compress_dim"); +} + +LogicalResult VectorExpandLoadOp::verify() { + return verifyCompressStoreExpandLoad( + *this, getResult().getType(), getExpandDim(), + /*vector_name=*/"result", + /*compress_expand_dim_name=*/"expand_dim"); +} + template LogicalResult verifyLoadOp(Op op) { MemRefType ref_ty = op.getBase().getType(); diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td index 084deae550db80..1130919108db8c 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td @@ -413,6 +413,49 @@ def TPU_ExpandLoadVregOp : TPU_Op<"expand_load_vreg", [DefaultMemRead]> { let hasVerifier = 1; } +def TPU_VectorExpandLoadOp : TPU_Op<"vector_expand_load", [DefaultMemRead]> { + let summary = "Performs a masked, expanded load of a vector."; + let description = [{ + In contrast to `tpu.vector_load`, which has non-expanded semantics + (masked-out elements leave gaps in the addresses that are read), masked-out + elements are skipped entirely when striding along `expand_dim`: + `result[i_0, i_1, ...]`, if not masked, is loaded from + `base[j_0, j_1, ...]` where: + - For `n != expand_dim`, `j_n = indices_n + i_n` (same as a regular load) + - For `n == expand_dim`, `j_n = indices_n + c`, where `c` is the count of + enabled (true) values in `mask[i_0, i_1, ..., 0:i_n, ...]` + + Elements at masked-off positions are undefined. + + For example, for a rank-1 `result` (and hence `expand_dim` = 0), loading + with mask [true, false, true, false] yields + [base[indices], , base[indices + 1], ]. + + base : Memref to load from. + indices : Scalar indices into `base`, identifying the position of the + first loaded element. Must be of the same rank as the base + memref shape. + mask : Mask of the same shape as `result`. Masking has a granularity + of 32 bits: elements that share a 32-bit word must have the + same mask value. + expand_dim: The dimension along which the enabled elements are read + consecutively. + result : The loaded vector. Must be of the same rank as the base memref + shape. + }]; + let arguments = (ins + AnyMemRef:$base, + Variadic:$indices, + VectorOfNonZeroRankOf<[I1]>:$mask, + I32Attr:$expand_dim + ); + let results = (outs AnyVectorOfNonZeroRank:$result); + let assemblyFormat = [{ + $base `[` $indices `]` `,` $mask attr-dict `:` type($base) `,` type($result) `,` type($mask) + }]; + let hasVerifier = 1; +} + // tpu.vector_load loads a vector from memory into a register. // // base : Memref to load from. @@ -1878,6 +1921,7 @@ def TPU_TransposeOp : TPU_Op<"transpose", [Pure]> { } }]; let hasVerifier = 1; + let hasFolder = 1; } def TPU_LogOp : TPU_Op<"log"> { diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc index 2fa8f8df4ecae9..0540083953a131 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc @@ -765,6 +765,188 @@ TEST_F(TpuOpsVerificationTest, VectorCompressStoreDynamicBaseDimAllowed) { ASSERT_OK(VerifyOp(cs)); } +TEST_F(TpuOpsVerificationTest, VectorExpandLoadValid) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_OK(VerifyOp(el)); +} + +// Packed (sub-32-bit) types must pass verification. +TEST_F(TpuOpsVerificationTest, VectorExpandLoadValidPacked) { + auto c0 = Create(0); + Value memref = + Create( + GetMemRefType({4, 16}, builder().getBF16Type(), MemorySpace::kVmem)) + .getMemref(); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, builder().getBF16Type()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_OK(VerifyOp(el)); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadInvalidMemorySpace) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kHbm); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_THAT(VerifyOp(el), + StatusIs(_, HasSubstr("Expected base memref to be in VMEM."))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadMismatchedRank) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{8}, + /*values=*/SmallVector(8, true)); + auto el = Create( + /*result=*/VectorType::get({8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_THAT(VerifyOp(el), + StatusIs(_, HasSubstr("Expected result to have the same rank as " + "base (2). Got: 1."))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadMismatchedIndicesCount) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_THAT(VerifyOp(el), StatusIs(_, HasSubstr("Expected 2 indices."))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadExpandDimOutOfRange) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(2)); + + ASSERT_THAT( + VerifyOp(el), + StatusIs(_, HasSubstr("Expected expand_dim to be in [0, 2). Got: 2."))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadNonMinormostExpandDim) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(0)); + + ASSERT_OK(VerifyOp(el)); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadMismatchedMaskShape) { + auto c0 = Create(0); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 4}, + /*values=*/SmallVector(8, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_THAT( + VerifyOp(el), + StatusIs(_, HasSubstr("Expected mask shape to match result shape"))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadNonExpandedDimOutOfBounds) { + auto c0 = Create(0); + Value memref = AllocaI32({2, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{4, 8}, + /*values=*/SmallVector(32, true)); + auto el = Create( + /*result=*/VectorType::get({4, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_THAT(VerifyOp(el), + StatusIs(_, HasSubstr("Non-expanded dimension 0 of result goes " + "out of bounds of base"))); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadDynamicIndicesAllowed) { + // Non-expanded dimension fits (2 <= 4), but constant index offset plus size + // exceeds base dimension (3 + 2 = 5 > 4). All indices are treated as dynamic, + // so verification passes. + auto c0 = Create(0); + auto c3 = Create(3); + Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{2, 8}, + /*values=*/SmallVector(16, true)); + auto el = Create( + /*result=*/VectorType::get({2, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c3, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_OK(VerifyOp(el)); +} + +TEST_F(TpuOpsVerificationTest, VectorExpandLoadDynamicBaseDimAllowed) { + auto c0 = Create(0); + Value memref = AllocaI32({ShapedType::kDynamic, 16}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{8, 8}, + /*values=*/SmallVector(64, true)); + auto el = Create( + /*result=*/VectorType::get({8, 8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask, + /*expand_dim=*/builder().getI32IntegerAttr(1)); + + ASSERT_OK(VerifyOp(el)); +} + TEST_F(TpuOpsVerificationTest, UnpackSubelementsValidIndex) { Value source = ConstantI8Vector(/*shape=*/{4, 8}, /*values=*/{1}); auto unpack = Create( diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc index e39b70dfd60a3c..3713b9c37535a1 100644 --- a/third_party/xla/xla/service/gpu/thunk_emitter.cc +++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc @@ -2516,7 +2516,7 @@ absl::StatusOr ThunkEmitter::EmitHostSend( return ThunkSequence::Of( Thunk::ThunkInfo::WithProfileAnnotation( instr, ir_emitter_context_->GetNextThunkId()), - src->shape(), slice.slice, *instr->channel_id(), send_recv_events_, + slice, *instr->channel_id(), send_recv_events_, ConvertFrontendAttributes(instr->frontend_attributes()), DeviceConstraint(instr)); } @@ -2550,10 +2550,13 @@ absl::StatusOr ThunkEmitter::EmitHostSendDone( "Unknown channel id in host transfer send done instruction"); } + const HloInstruction* src = host_transfer->operand(0); + ABSL_ASSIGN_OR_RETURN(ShapedSlice slice, GetShapedSliceForHlo(src, {})); + return ThunkSequence::Of( Thunk::ThunkInfo::WithProfileAnnotation( done, ir_emitter_context_->GetNextThunkId()), - *host_transfer->channel_id(), send_recv_events_, + slice, *host_transfer->channel_id(), send_recv_events_, DeviceConstraint(host_transfer)); } @@ -2565,10 +2568,13 @@ absl::StatusOr ThunkEmitter::EmitHostRecvDone( "Unknown channel id in host transfer recv done instruction"); } + const HloInstruction* src = host_transfer->operand(0); + ABSL_ASSIGN_OR_RETURN(ShapedSlice slice, GetShapedSliceForHlo(src, {})); + return ThunkSequence::Of( Thunk::ThunkInfo::WithProfileAnnotation( done, ir_emitter_context_->GetNextThunkId()), - *host_transfer->channel_id(), send_recv_events_, + slice, *host_transfer->channel_id(), send_recv_events_, DeviceConstraint(host_transfer)); } diff --git a/third_party/xla/xla/service/hlo_cse.cc b/third_party/xla/xla/service/hlo_cse.cc index 440eb16a92ad35..b67f565edf1a6f 100644 --- a/third_party/xla/xla/service/hlo_cse.cc +++ b/third_party/xla/xla/service/hlo_cse.cc @@ -112,12 +112,14 @@ absl::StatusOr CombineConstants( return combined > 0; } -// An instruction is considered to be equivalent to another only if they -// share the exact same set of operands. +// This differs from "normal" HLO Instruction hashing because it takes the +// operands of an instruction into account. For the purposes of CSEs, two +// instructions that have different operands are never equivalent, so +// disambiguating by operands is always useful - but not always sufficient. struct CseKey { template - friend H AbslHashValue(H h, const CseKey& key) { - auto instruction = key.hlo; + static H HashInstruction(H h, const HloInstruction* instruction, + bool hash_operands) { h = instruction->shape().IsArray() ? H::combine(std::move(h), instruction->opcode(), instruction->shape().dimensions()) @@ -144,26 +146,28 @@ struct CseKey { h = result_accuracy_hash(std::move(h), instruction->result_accuracy()); // Hash operands, ignoring operand order on commutative ops. - if (HloOpcodeIsBinaryCommutative(instruction->opcode())) { - CHECK_EQ(instruction->operand_count(), 2); - auto id0 = instruction->operand(0)->unique_id(); - if (instruction->operand(0)->opcode() == HloOpcode::kIota) { - id0 = 0; - } - auto id1 = instruction->operand(1)->unique_id(); - if (instruction->operand(1)->opcode() == HloOpcode::kIota) { - id1 = 0; - } - if (id0 > id1) { - std::swap(id0, id1); - } - h = H::combine(std::move(h), id0, id1); - } else { - for (auto operand : instruction->operands()) { - if (operand->opcode() == HloOpcode::kIota) { - continue; + if (hash_operands) { + if (HloOpcodeIsBinaryCommutative(instruction->opcode())) { + CHECK_EQ(instruction->operand_count(), 2); + auto id0 = instruction->operand(0)->unique_id(); + if (instruction->operand(0)->opcode() == HloOpcode::kIota) { + id0 = 0; + } + auto id1 = instruction->operand(1)->unique_id(); + if (instruction->operand(1)->opcode() == HloOpcode::kIota) { + id1 = 0; + } + if (id0 > id1) { + std::swap(id0, id1); + } + h = H::combine(std::move(h), id0, id1); + } else { + for (auto operand : instruction->operands()) { + if (operand->opcode() == HloOpcode::kIota) { + continue; + } + h = H::combine(std::move(h), operand->unique_id()); } - h = H::combine(std::move(h), operand->unique_id()); } } @@ -171,6 +175,31 @@ struct CseKey { h = H::combine(std::move(h), c->root_instruction()->opcode()); } switch (instruction->opcode()) { + case HloOpcode::kFusion: + // The return post-order for two equal computations will always be the + // same, since it's determined purely by operand order, so we should + // never expect to hash two equal computations to different values. Note + // that the reasons this works are somewhat subtle: (a) We never + // consider kFusion commutative, so a dependency on operand order is + // safe. (b) Control dependencies may make the post-order traversal + // order depend on insertion order, but that's not an issue for fusion + // computations. (c) The post-order returned by + // HloComputation::MakeInstructionPostOrder() can depend on insertion + // order if dead instructions are present, which is why we use + // MakeInstructionPostOrderFrom() on the root. + for (const HloInstruction* fused_instruction : + instruction->fused_instructions_computation() + ->MakeInstructionPostOrderFrom( + *instruction->fused_instructions_computation() + ->root_instruction())) { + // Internal instructions of equivalent fusions belong to different + // HloComputations (whose IDs are encoded into + // HloInstruction::unique_id()), so we can only hash operand + // unique_ids for top-level instructions. + h = HashInstruction(std::move(h), fused_instruction, + /*hash_operands=*/false); + } + return H::combine(std::move(h), instruction->fusion_kind()); case HloOpcode::kSlice: return H::combine(std::move(h), instruction->slice_starts(), instruction->slice_strides()); @@ -227,6 +256,11 @@ struct CseKey { return std::move(h); } } + + template + friend H AbslHashValue(H h, const CseKey& key) { + return HashInstruction(std::move(h), key.hlo, /*hash_operands=*/true); + } HloInstruction* hlo; }; diff --git a/third_party/xla/xla/service/hlo_cse_test.cc b/third_party/xla/xla/service/hlo_cse_test.cc index 68dad767e5f12c..9336949ec3e0cb 100644 --- a/third_party/xla/xla/service/hlo_cse_test.cc +++ b/third_party/xla/xla/service/hlo_cse_test.cc @@ -1089,6 +1089,60 @@ TEST_F(HloCseTest, ResultAccuracyCseKey) { EXPECT_EQ(root->operand(2), root->operand(3)); } +TEST_F(HloCseTest, FusionCseKey) { + constexpr absl::string_view kHlo = R"( + HloModule m + + fused_comp.0 { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + mul = f32[4] multiply(p0, p1) + ROOT add = f32[4] add(p0, mul) + } + + fused_comp.1 { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + mul = f32[4] multiply(p0, p1) + ROOT add = f32[4] add(p0, mul) + } + + fused_comp.2 { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + sub = f32[4] subtract(p0, p1) + ROOT add = f32[4] add(p0, sub) + } + + fused_comp.3 { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + mul = f32[4] multiply(p1, p1) + ROOT add = f32[4] add(p0, mul) + } + + ENTRY main { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + f0 = f32[4] fusion(p0, p1), kind=kLoop, calls=fused_comp.0 + f1 = f32[4] fusion(p0, p1), kind=kLoop, calls=fused_comp.1 + f2 = f32[4] fusion(p0, p1), kind=kLoop, calls=fused_comp.2 + f3 = f32[4] fusion(p0, p1), kind=kLoop, calls=fused_comp.3 + ROOT t = tuple(f0, f1, f2, f3) + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(kHlo)); + HloCSE cse(/*is_layout_sensitive=*/false); + EXPECT_THAT(cse.Run(m.get()), absl_testing::IsOkAndHolds(true)); + const HloInstruction* root = m->entry_computation()->root_instruction(); + ASSERT_EQ(root->operand_count(), 4); + EXPECT_EQ(root->operand(0), root->operand(1)); + EXPECT_NE(root->operand(0), root->operand(2)); + EXPECT_NE(root->operand(0), root->operand(3)); + EXPECT_NE(root->operand(2), root->operand(3)); +} + TEST_F(HloCseTest, ScalarCustomCallNoOperands) { constexpr absl::string_view kHlo = R"( HloModule main