Use mlir::OpState::operator->() to get to methods of mlir::Operation.
This is a preparation step to remove those methods from OpState.
PiperOrigin-RevId: 348434761
Change-Id: I975647b30b42722b9af870f024efe8d05199dd1e
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/mlir_compiler_impl.cc b/tensorflow/compiler/xla/service/mlir_gpu/mlir_compiler_impl.cc
index 9547f43..757d496 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/mlir_compiler_impl.cc
+++ b/tensorflow/compiler/xla/service/mlir_gpu/mlir_compiler_impl.cc
@@ -300,9 +300,9 @@
auto num_original_args = kernel.getNumArguments();
std::vector<LLVMType> new_arg_types(buffers.size(), ptr_type);
- kernel.setAttr(kernel.getTypeAttrName(),
- mlir::TypeAttr::get(LLVMType::getFunctionTy(
- void_type, new_arg_types, /*isVarArg=*/false)));
+ kernel->setAttr(kernel.getTypeAttrName(),
+ mlir::TypeAttr::get(LLVMType::getFunctionTy(
+ void_type, new_arg_types, /*isVarArg=*/false)));
std::vector<Value> original_args(kernel.args_begin(), kernel.args_end());
std::vector<mlir::Type> as_mlir_types(new_arg_types.begin(),
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/passes.cc b/tensorflow/compiler/xla/service/mlir_gpu/passes.cc
index 84751bc..2cabd92 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/passes.cc
+++ b/tensorflow/compiler/xla/service/mlir_gpu/passes.cc
@@ -257,8 +257,8 @@
auto new_kernel = kernel_builder.create<mlir::gpu::GPUFuncOp>(
kernel.getLoc(), kernel.getName(),
kernel_builder.getFunctionType(operand_types, {}));
- new_kernel.setAttr(mlir::gpu::GPUDialect::getKernelFuncAttrName(),
- kernel_builder.getUnitAttr());
+ new_kernel->setAttr(mlir::gpu::GPUDialect::getKernelFuncAttrName(),
+ kernel_builder.getUnitAttr());
// Create a map from old kernel argument to new one.
mlir::BlockAndValueMapping old_kernel_to_new;