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;