Repository: alibaba/BladeDISC (language: python)

Issue:
Revert changes on XLA under tf_community.
As is done in https://github.com/alibaba/BladeDISC/pull/645, we won't build internal decoupled xla from BladeDISC, but there still exists on related changes under xla directries of tf_community. To avoid further rebase conflicts, it better to revert them, too.


Discussion / hints:
yyds

Fix patch:
diff --git a/docs/developers/pass_pipeline.md b/docs/developers/pass_pipeline.md
index 99117fbd719..9215675a191 100644
--- a/docs/developers/pass_pipeline.md
+++ b/docs/developers/pass_pipeline.md
@@ -46,7 +46,7 @@ tutorials in this document. Please refer to
 [build_from_source.md](../build_from_source.md)
 to build the binaries if you are in other CUDA environments.
 
-![passpipeline](./pics/pass_pipeline.png) 
+![passpipeline](./pics/pass_pipeline.png)
 <center>PassPipeline of DISC</center>
 
 Before we start, let's recap the major problems we need to solve for the
@@ -64,7 +64,7 @@ shape in some cases. As a solution, we extend HLO with a set of IR
 supplementation. In general, these extended operations are prefixed with
 "Dynamic" (one exception is mhlo.RealDynamicSliceOp since there is already one
 named with DynamicSliceOp). These extended part has already been upstreamed into
-[Mlir-HLO](https://github.com/tensorflow/mlir-hlo). 
+[Mlir-HLO](https://github.com/tensorflow/mlir-hlo).
 
 
 Here we demo the difference in op definition with SliceOp as an example:
@@ -74,7 +74,7 @@ mhlo.SliceOp:
 ```
 def HLO_SliceOp: HLO_Op<
       "slice",
-      [NoSideEffect, SameOperandsAndResultElementType,
+      [Pure, SameOperandsAndResultElementType,
        AllTypesMatch<["start_indices", "limit_indices", "strides"]>,
        DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let arguments = (ins
@@ -96,7 +96,7 @@ mhlo.RealDynamicSliceOp:
 ```
 def HLO_RealDynamicSliceOp: HLO_ShapedInterfaceOp<
       "real_dynamic_slice",
-      [NoSideEffect, AllElementTypesMatch<["operand", "result"]>,
+      [Pure, AllElementTypesMatch<["operand", "result"]>,
        AllTypesMatch<["start_indices", "limit_indices", "strides"]>]> {
   let summary = "Real Dynamic Slice operator";
   let description = [{
@@ -134,7 +134,7 @@ calculated shapes of each buffer will guide the buffer allocation and launch
 dimension at runtime. BladeDISC does bufferization during the conversion from
 mhlo to lmhlo. Rather than using an interpreter or virtual machine, BladeDISC
 compiles and generates the code of computations on both host and device side,
-and also the runtime flows (buffer management, kernel launch, et.al.). 
+and also the runtime flows (buffer management, kernel launch, et.al.).
 
 Runtime Abstraction Layer (RAL) is used to isolate the complexity between the
 compiler and different runtime environments, so that the compiler does not see
@@ -187,9 +187,9 @@ tutorial:
 !wget http://pai-blade.cn-hangzhou.oss.aliyun-inc.com/bladedisc_notebook_binaries/disc_compiler_main -O disc_compiler_main
 
 !wget http://pai-blade.cn-hangzhou.oss.aliyun-inc.com/bladedisc_notebook_binaries/tf-opt -O tf-opt
-    
+
 !wget http://pai-blade.cn-hangzhou.oss.aliyun-inc.com/bladedisc_notebook_binaries/disc-opt -O disc-opt
-    
+
 !wget http://pai-blade.cn-hangzhou.oss.aliyun-inc.com/bladedisc_notebook_binaries/tutorial.mlir -O tutorial.mlir
 
 !chmod +x disc_compiler_main
@@ -200,7 +200,7 @@ tutorial:
 ```
 
 We choose an input subgraph that contains as mush essential elements as
-possible to demo the pass pipeline. 
+possible to demo the pass pipeline.
 
 
 ```python
@@ -271,7 +271,7 @@ is responsible to apply the logged attributes onto the IR.
 
 After this pass, we should get the IR mainly made of Mhlo Dialect and Standard
 Dialect, in which the Standard Dialect ops are for shape calculations. This is
-part of the shape calculation and more will be lowered in the following phases. 
+part of the shape calculation and more will be lowered in the following phases.
 
 
 ```python
@@ -617,7 +617,7 @@ RAL uses the `DispatchOp` in `disc_ral` dialect to model external Library call
 ops (e.g. gemm/conv) to `disc_ral.dispatch` ops.  At the end of the pass
 pipeline, the `disc_ral.dispatch` ops will be eventually lowered into  a
 uniformed type-erased style to make the compiled binary having stable and clean
-ABI. 
+ABI.
 
 2. Lower const ops. Similar to other stateful operations, we use a library call
    to hide the const initialization process and use the RAL context object to
@@ -839,7 +839,7 @@ Host side compilation can be roughly divided into two parts.
   multi-threading execution. Readers may refer to `disc-outline-cpu-kernel` and
   `disc-cpu-map-parallel-loop` for more information. We won't expand the detail
   in this doc.
- 
+
 *  Scheduling logic codegen. The scheduling logic includes kernel launching
    setting, data movement, synchronization, and buffer management. The
    scheduling logic will be lowered to LLVM IR and then be compiled into binary
@@ -919,5 +919,5 @@ form. The basic idea is:
 
 After this phase, the LLVM Dialect will then be converted to LLVM IR, and sent
 to LLVM backend. By this step, we've got a binary, which is executable and will
-link different runtime libraries for different scenarios. 
+link different runtime libraries for different scenarios.
 
diff --git a/pytorch_blade/WORKSPACE b/pytorch_blade/WORKSPACE
index a93575a7b5f..77bc44a2e9d 100644
--- a/pytorch_blade/WORKSPACE
+++ b/pytorch_blade/WORKSPACE
@@ -2,6 +2,7 @@ workspace(name = "org_torch_blade")
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
+load("//bazel/torch_mlir:repo.bzl", "blade_http_archive")
 
 http_archive(
     name = "com_github_bazelbuild_buildtools",
@@ -12,6 +13,15 @@ http_archive(
     ],
 )
 
+http_archive(
+    name = "bazel_skylib",
+    sha256 = "74d544d96f4a5bb630d465ca8bbcfe231e3594e5aae57e1edbf17a6eb3ca2506",
+    urls = [
+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
+        "https://github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
+    ],
+)
+
 http_archive(
     name = "rules_python",
     urls = [
@@ -118,10 +128,16 @@ new_local_repository(
     path = "../../platform_alibaba/blade_gemm"
 )
 
-new_local_repository(
+blade_http_archive(
   name = "mlir-hlo",
-  build_file = "third_party/torch-mlir/externals/mlir-hlo/BUILD",
-  path = "third_party/torch-mlir/externals/mlir-hlo"
+  sha256 = "c7d9b8667bdcc2865701571e7cdef2ba21762b619bb5505cdec3af6f8b36fbb6",
+  strip_prefix = "mlir-hlo-36238f16441cd1a884af988d4400d2ebb0c75bbc",
+  urls = [
+    "https://github.com/tensorflow/mlir-hlo/archive/36238f16441cd1a884af988d4400d2ebb0c75bbc.zip",
+  ],
+  patch_file = [
+    "//bazel/torch_mlir:disable-simplify-dynamic-gather-to-gather.patch",
+  ]
 )
 
 new_local_repository(
diff --git a/pytorch_blade/bazel/torch_mlir/disable-simplify-dynamic-gather-to-gather.patch b/pytorch_blade/bazel/torch_mlir/disable-simplify-dynamic-gather-to-gather.patch
new file mode 100644
index 00000000000..028565e05ae
--- /dev/null
+++ b/pytorch_blade/bazel/torch_mlir/disable-simplify-dynamic-gather-to-gather.patch
@@ -0,0 +1,13 @@
+diff --git a/lib/Dialect/mhlo/IR/hlo_ops.cc b/lib/Dialect/mhlo/IR/hlo_ops.cc
+index feb575f0..923b53e1 100644
+--- a/lib/Dialect/mhlo/IR/hlo_ops.cc
++++ b/lib/Dialect/mhlo/IR/hlo_ops.cc
+@@ -1784,7 +1784,7 @@ LogicalResult simplifyDynamicGatherToGather(DynamicGatherOp op,
+ 
+ void DynamicGatherOp::getCanonicalizationPatterns(RewritePatternSet& result,
+                                                   MLIRContext* context) {
+-  result.add(simplifyDynamicGatherToGather);
++  // result.add(simplifyDynamicGatherToGather);
+ }
+ 
+ LogicalResult DynamicGatherOp::reifyReturnTypeShapes(
diff --git a/pytorch_blade/bazel/torch_mlir/repo.bzl b/pytorch_blade/bazel/torch_mlir/repo.bzl
new file mode 100644
index 00000000000..95dfe577348
--- /dev/null
+++ b/pytorch_blade/bazel/torch_mlir/repo.bzl
@@ -0,0 +1,69 @@
+"""Utilities for defining Blade Bazel dependencies."""
+
+def _get_link_dict(ctx, link_files, build_file):
+    link_dict = {ctx.path(v): ctx.path(Label(k)) for k, v in link_files.items()}
+    if build_file:
+        # Use BUILD.bazel because it takes precedence over BUILD.
+        link_dict[ctx.path("BUILD.bazel")] = ctx.path(Label(build_file))
+    return link_dict
+
+def _blade_http_archive_impl(ctx):
+    # Construct all paths early on to prevent rule restart. We want the
+    # attributes to be strings instead of labels because they refer to files
+    # in the TensorFlow repository, not files in repos depending on TensorFlow.
+    # See also https://github.com/bazelbuild/bazel/issues/10515.
+    link_dict = _get_link_dict(ctx, ctx.attr.link_files, ctx.attr.build_file)
+
+    # For some reason, we need to "resolve" labels once before the
+    # download_and_extract otherwise it'll invalidate and re-download the
+    # archive each time.
+    # https://github.com/bazelbuild/bazel/issues/10515
+    patch_files = ctx.attr.patch_file
+    for patch_file in patch_files:
+        if patch_file:
+            ctx.path(Label(patch_file))
+
+    ctx.download_and_extract(
+        url = ctx.attr.urls,
+        sha256 = ctx.attr.sha256,
+        type = ctx.attr.type,
+        stripPrefix = ctx.attr.strip_prefix,
+    )
+    if patch_files:
+        for patch_file in patch_files:
+            patch_file = ctx.path(Label(patch_file)) if patch_file else None
+            if patch_file:
+                ctx.patch(patch_file, strip = 1)
+
+    for dst, src in link_dict.items():
+        ctx.delete(dst)
+        ctx.symlink(src, dst)
+
+_blade_http_archive = repository_rule(
+    implementation = _blade_http_archive_impl,
+    attrs = {
+        "sha256": attr.string(mandatory = True),
+        "urls": attr.string_list(mandatory = True),
+        "strip_prefix": attr.string(),
+        "type": attr.string(),
+        "patch_file": attr.string_list(),
+        "build_file": attr.string(),
+        "system_build_file": attr.string(),
+        "link_files": attr.string_dict(),
+        "system_link_files": attr.string_dict(),
+    },
+    environ = ["TF_SYSTEM_LIBS"],
+)
+
+def blade_http_archive(name, sha256, urls, **kwargs):
+    if native.existing_rule(name):
+        print("\n\033[1;33mWarning:\033[0m skipping import of repository '" +
+              name + "' because it already exists.\n")
+        return
+
+    _blade_http_archive(
+        name = name,
+        sha256 = sha256,
+        urls = urls,
+        **kwargs
+    )
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp
index ace58a0a988..a9dad50a61f 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp
@@ -16,10 +16,10 @@
 #include "torch-mlir/Dialect/Torch/Transforms/Passes.h"
 #include "torch-mlir/Dialect/TorchConversion/Transforms/Passes.h"
 
-namespace {
+namespace impl {
 #define GEN_PASS_REGISTRATION
 #include "torch-mlir/Conversion/MhloPasses.h.inc"
-} // end namespace
+} // end namespace impl
 
 using namespace mlir;
 using namespace mlir::torch;
@@ -27,7 +27,7 @@ using namespace mlir::torch::Torch;
 using namespace mlir::torch::TorchConversion;
 
 void ::mlir::torch::registerTorchToMhloPasses() {
-  ::registerPasses();
+  ::impl::registerPasses();
   ::mlir::PassPipelineRegistration<Torch::TorchLoweringPipelineOptions>(
       "torch-backend-to-mhlo-backend-pipeline",
       "Pipeline lowering torch backend contract to mhlo backend "
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp
index 88849fa1ea0..f122d99796a 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp
@@ -23,7 +23,7 @@
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
 #include "mlir-hlo/utils/hlo_utils.h"
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Matchers.h"
@@ -555,10 +555,10 @@ LogicalResult ConvertAtenOp<AtenFlipOp>::matchAndRewrite(
   return success();
 }
 
-// ValsemVariantAtenUniformOp
+// AtenUniformOp
 template <>
-LogicalResult ConvertAtenOp<ValsemVariantAtenUniformOp>::matchAndRewrite(
-    ValsemVariantAtenUniformOp op,
+LogicalResult ConvertAtenOp<AtenUniformOp>::matchAndRewrite(
+    AtenUniformOp op,
     OpAdaptor adaptor,
     ConversionPatternRewriter& rewriter) const {
   auto inputTy = adaptor.self().getType().template cast<RankedTensorType>();
@@ -724,7 +724,7 @@ static llvm::Optional<ValueRange> getMaxValueInDim(
   auto valueReduceOp = rewriter.create<mhlo::ReduceOp>(
       op->getLoc(), input, initValue, dimensions);
   {
-    Block& block = valueReduceOp.body().emplaceBlock();
+    Block& block = valueReduceOp.getBody().emplaceBlock();
     auto argumentType = RankedTensorType::get({}, inputTy.getElementType());
     block.addArgument(argumentType, op->getLoc());
     block.addArgument(argumentType, op->getLoc());
@@ -785,7 +785,7 @@ static llvm::Optional<ValueRange> getMaxIndicesInDim(
       },
       dimensions);
   {
-    Block& block = indicesReduceOp.body().emplaceBlock();
+    Block& block = indicesReduceOp.getBody().emplaceBlock();
 
     // Add block arguments
     auto blockValArgumentType =
@@ -1067,7 +1067,7 @@ class DiscConvertTorchToMhlo
     registry.insert<chlo::ChloDialect>();
     registry.insert<mhlo::MhloDialect>();
     registry.insert<tensor::TensorDialect>();
-    registry.insert<arith::ArithmeticDialect>();
+    registry.insert<arith::ArithDialect>();
     registry.insert<Torch::TorchDialect>();
     torch::TorchConversion::getBackendTypeConversionDependentDialects(registry);
   }
@@ -1079,7 +1079,7 @@ class DiscConvertTorchToMhlo
         chlo::ChloDialect,
         mhlo::MhloDialect,
         tensor::TensorDialect,
-        arith::ArithmeticDialect,
+        arith::ArithDialect,
         Torch::TorchDialect>();
 
     TypeConverter typeConverter;
@@ -1125,7 +1125,7 @@ class DiscConvertTorchToMhlo
     INSERT_ATENOP_PATTERN(AtenDropoutOp);
     INSERT_ATENOP_PATTERN(TensorStaticInfoCastOp);
     INSERT_ATENOP_PATTERN(AtenFlipOp);
-    INSERT_ATENOP_PATTERN(ValsemVariantAtenUniformOp);
+    INSERT_ATENOP_PATTERN(AtenUniformOp);
     INSERT_ATENOP_PATTERN(AtenTensorOp);
     INSERT_ATENOP_PATTERN(AtenTensorBoolOp);
     INSERT_ATENOP_PATTERN(AtenTensorFloatOp);
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/ApplyValueSemantics.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/ApplyValueSemantics.cpp
index b130ec5fa22..24392141b03 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/ApplyValueSemantics.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/ApplyValueSemantics.cpp
@@ -16,7 +16,7 @@
 #include "torch-mlir/Dialect/Torch/IR/TorchTypes.h"
 #include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.h"
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Math/IR/Math.h"
@@ -106,7 +106,7 @@ class ApplyValueSemanticsPass
     target.addLegalDialect<
         Torch::TorchDialect,
         func::FuncDialect,
-        arith::ArithmeticDialect,
+        arith::ArithDialect,
         tensor::TensorDialect,
         cf::ControlFlowDialect,
         math::MathDialect>();
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp
index c10f8188a4a..2b3e603a34f 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp
@@ -9,7 +9,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Math/IR/Math.h"
@@ -139,7 +139,7 @@ class DiscConvertTorchToDiscMhlo
 
     ConversionTarget target(*context);
     target.addLegalDialect<
-        arith::ArithmeticDialect,
+        arith::ArithDialect,
         chlo::ChloDialect,
         mhlo::MhloDialect,
         mhlo_disc::MhloDiscDialect,
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp
index 66eba6e158d..907cc921243 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp
@@ -19,7 +19,7 @@
 #include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.h"
 #include "torch-mlir/Dialect/TorchConversion/Transforms/BackendTypeConversion.h"
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Math/IR/Math.h"
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp
index f98ef45cd60..6323c620863 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp
@@ -20,7 +20,7 @@
 
 #include <mlir-hlo/Dialect/mhlo/IR/hlo_ops.h> // from tf repo
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/FuncConversions.h"
@@ -153,7 +153,7 @@ class VerifyMhloBackendContractPass
     target.addDynamicallyLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>(
         opHasLegalTypes);
     target.addLegalDialect<
-        arith::ArithmeticDialect,
+        arith::ArithDialect,
         chlo::ChloDialect,
         mhlo::MhloDialect,
         mhlo_disc::MhloDiscDialect,
diff --git a/pytorch_blade/third_party/torch-mlir b/pytorch_blade/third_party/torch-mlir
index 9e0a37fe290..785c103ea69 160000
--- a/pytorch_blade/third_party/torch-mlir
+++ b/pytorch_blade/third_party/torch-mlir
@@ -1,1 +1,1 @@
-Subproject commit 9e0a37fe2902bc60c891aa5e37749f0e01636bdf
+Subproject commit 785c103ea692c2e2a1d062567970645e175f01ef
diff --git a/tao_compiler/decoupling/tao_compiler_main.cc b/tao_compiler/decoupling/tao_compiler_main.cc
index dead1f94066..dc233438d43 100644
--- a/tao_compiler/decoupling/tao_compiler_main.cc
+++ b/tao_compiler/decoupling/tao_compiler_main.cc
@@ -138,8 +138,9 @@ Status RealMain(int argc, char** argv) {
   }
 
   DeviceType device_type(input.options().device_type());
-  auto* compiler_wrapper =
-      CompilerBase::GetCompilerForDevice(device_type).ValueOrDie();
+  auto status_or = CompilerBase::GetCompilerForDevice(device_type);
+  if (!status_or.ok()) return status_or.status();
+  auto* compiler_wrapper = status_or.value();
   TF_RETURN_IF_ERROR(compiler_wrapper->Compile(input, output_fn));
   tao::TaoCompilerTrace::Instance()->Shutdown();
   return Status::OK();
@@ -148,10 +149,10 @@ Status RealMain(int argc, char** argv) {
 }  // namespace tensorflow
 
 int main(int argc, char** argv) {
-  std::vector<tensorflow::Flag> flag_list;
+  std::vector<tsl::Flag> flag_list;
   xla::AppendDebugOptionsFlags(&flag_list);
-  std::string usage = tensorflow::Flags::Usage(argv[0], flag_list);
-  if (!tensorflow::Flags::Parse(&argc, argv, flag_list)) {
+  std::string usage = tsl::Flags::Usage(argv[0], flag_list);
+  if (!tsl::Flags::Parse(&argc, argv, flag_list)) {
     LOG(ERROR) << "\n" << usage;
     return -1;
   }
diff --git a/tao_compiler/decoupling/tao_compiler_trace.h b/tao_compiler/decoupling/tao_compiler_trace.h
index a1ec9b4eaf8..31f4f3ec18c 100644
--- a/tao_compiler/decoupling/tao_compiler_trace.h
+++ b/tao_compiler/decoupling/tao_compiler_trace.h
@@ -37,7 +37,7 @@ class TaoCompilerTrace {
 
  private:
   struct Event {
-    tensorflow::uint64 timestamp_us;
+    uint64_t timestamp_us;
     std::string key;
     std::string value;
   };
diff --git a/tao_compiler/mlir/disc/BUILD b/tao_compiler/mlir/disc/BUILD
index 648384b1b9e..39cdc071c59 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -352,7 +352,7 @@ cc_library(
         ":input_inline_fusion_pattern",
         ":pass_details",
         ":placement_utils",
-        "//tensorflow/compiler/xla/mlir_hlo:LmhloPassIncGen",
+        "//tensorflow/compiler/xla/mlir_hlo:lmhlo_pass_inc_gen",
         "//tensorflow/compiler/xla/mlir_hlo:mlir_hlo",
         "//tensorflow/compiler/xla/mlir_hlo:lhlo",
         "//tensorflow/compiler/xla/mlir_hlo:map_lmhlo_to_scalar_op",
@@ -376,7 +376,7 @@ cc_library(
     srcs = ["transforms/disc_fuse_splat_const.cc"],
     deps = [
         ":pass_details",
-        "//tensorflow/compiler/xla/mlir_hlo:LmhloPassIncGen",
+        "//tensorflow/compiler/xla/mlir_hlo:lmhlo_pass_inc_gen",
         "//tensorflow/compiler/xla/mlir_hlo:lhlo",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
@@ -974,7 +974,7 @@ disc_cc_library(
         "//tensorflow/compiler/xla:status",
         "//tensorflow/compiler/xla:statusor",
         "//tensorflow/compiler/xla/stream_executor/gpu:asm_compiler",
-        "//tensorflow/core/platform:cuda_libdevice_path",
+        "//tensorflow/tsl/platform:cuda_libdevice_path",
         "//tensorflow/core:lib",
     ] + if_cuda_is_configured([
         "//tensorflow/compiler/xla/stream_executor/cuda:cuda_asm_compiler",
@@ -1016,6 +1016,7 @@ cc_library(
     deps = [
         ":pass_details",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
@@ -1112,6 +1113,7 @@ cc_library(
         ":pass_details",
         "//tensorflow/compiler/xla/mlir_hlo:mlir_hlo",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:MemRefDialect",
@@ -1635,7 +1637,7 @@ cc_library(
         ":fusion_utils",
         ":pass_details",
         ":placement_utils",
-        "//tensorflow/compiler/xla/mlir_hlo:LmhloPassIncGen",
+        "//tensorflow/compiler/xla/mlir_hlo:lmhlo_pass_inc_gen",
         "//tensorflow/compiler/xla/mlir_hlo:mlir_hlo",
         "//tensorflow/compiler/xla/mlir_hlo:lhlo",
         "@llvm-project//llvm:Support",
@@ -2393,7 +2395,7 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:ArithmeticDialect",
+        "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:ControlFlowInterfaces",
         "@llvm-project//mlir:Dialect",
         "@llvm-project//mlir:FuncDialect",
@@ -2446,7 +2448,7 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:ArithmeticDialect",
+        "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:ControlFlowInterfaces",
         "@llvm-project//mlir:Dialect",
         "@llvm-project//mlir:FuncDialect",
diff --git a/tao_compiler/mlir/disc/IR/disc_shape_ops.cc b/tao_compiler/mlir/disc/IR/disc_shape_ops.cc
index 44e675bc483..75813e1387b 100644
--- a/tao_compiler/mlir/disc/IR/disc_shape_ops.cc
+++ b/tao_compiler/mlir/disc/IR/disc_shape_ops.cc
@@ -15,7 +15,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/disc/IR/disc_shape_ops.h"
 
 #include "llvm/Support/Debug.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/PatternMatch.h"
 
@@ -49,28 +49,28 @@ struct LinearizeOfDelinearizeOp : public OpRewritePattern<LinearizeOp> {
 
   LogicalResult matchAndRewrite(LinearizeOp op,
                                 PatternRewriter& rewriter) const override {
-    if (!op.multiDimIndexes().size()) {
+    if (!op.getMultiDimIndexes().size()) {
       return failure();
     }
 
     auto delinearizeOp = dyn_cast_or_null<DelinearizeOp>(
-        op.multiDimIndexes().front().getDefiningOp());
+        op.getMultiDimIndexes().front().getDefiningOp());
     if (!delinearizeOp) return failure();
 
-    if (op.multiDimIndexes().size() != delinearizeOp->getResults().size())
+    if (op.getMultiDimIndexes().size() != delinearizeOp->getResults().size())
       return failure();
 
     for (auto&& z :
-         llvm::zip(op.multiDimIndexes(), delinearizeOp->getResults())) {
+         llvm::zip(op.getMultiDimIndexes(), delinearizeOp->getResults())) {
       if (std::get<0>(z) != std::get<1>(z)) return failure();
     }
 
-    for (auto&& z :
-         llvm::zip(delinearizeOp.shapeDimIndexes(), op.shapeDimIndexes())) {
+    for (auto&& z : llvm::zip(delinearizeOp.getShapeDimIndexes(),
+                              op.getShapeDimIndexes())) {
       if (std::get<0>(z) != std::get<1>(z)) return failure();
     }
 
-    rewriter.replaceOp(op, {delinearizeOp.linearIndex()});
+    rewriter.replaceOp(op, {delinearizeOp.getLinearIndex()});
     return success();
   }
 };
@@ -105,13 +105,14 @@ struct RemoveSizeOneDimOfLinearizeOp : public OpRewritePattern<LinearizeOp> {
 
   LogicalResult matchAndRewrite(LinearizeOp op,
                                 PatternRewriter& rewriter) const override {
-    if (!op.multiDimIndexes().size()) {
+    if (!op.getMultiDimIndexes().size()) {
       return failure();
     }
 
     SmallVector<Value> newMultiDimIndexes;
     SmallVector<Value> newShapeDimIndexes;
-    for (auto&& z : llvm::zip(op.multiDimIndexes(), op.shapeDimIndexes())) {
+    for (auto&& z :
+         llvm::zip(op.getMultiDimIndexes(), op.getShapeDimIndexes())) {
       Value idx = std::get<0>(z);
       Value dimSize = std::get<1>(z);
       auto constOp =
@@ -122,7 +123,7 @@ struct RemoveSizeOneDimOfLinearizeOp : public OpRewritePattern<LinearizeOp> {
       newShapeDimIndexes.push_back(dimSize);
     }
 
-    if (newMultiDimIndexes.size() == op.multiDimIndexes().size()) {
+    if (newMultiDimIndexes.size() == op.getMultiDimIndexes().size()) {
       return failure();
     }
 
@@ -137,23 +138,24 @@ struct DelinearizeOfLinearizeOp : public OpRewritePattern<DelinearizeOp> {
 
   LogicalResult matchAndRewrite(DelinearizeOp op,
                                 PatternRewriter& rewriter) const override {
-    if (!op.shapeDimIndexes().size()) {
+    if (!op.getShapeDimIndexes().size()) {
       return failure();
     }
 
     auto linearizeOp =
-        dyn_cast_or_null<LinearizeOp>(op.linearIndex().getDefiningOp());
+        dyn_cast_or_null<LinearizeOp>(op.getLinearIndex().getDefiningOp());
     if (!linearizeOp) return failure();
 
-    if (linearizeOp.shapeDimIndexes().size() != op.shapeDimIndexes().size())
+    if (linearizeOp.getShapeDimIndexes().size() !=
+        op.getShapeDimIndexes().size())
       return failure();
 
     for (auto&& z :
-         llvm::zip(linearizeOp.shapeDimIndexes(), op.shapeDimIndexes())) {
+         llvm::zip(linearizeOp.getShapeDimIndexes(), op.getShapeDimIndexes())) {
       if (std::get<0>(z) != std::get<1>(z)) return failure();
     }
 
-    rewriter.replaceOp(op, linearizeOp.multiDimIndexes());
+    rewriter.replaceOp(op, linearizeOp.getMultiDimIndexes());
     return success();
   }
 };
@@ -170,14 +172,14 @@ struct RemoveSizeOneDimOfDelinearizeOp
 
   LogicalResult matchAndRewrite(DelinearizeOp op,
                                 PatternRewriter& rewriter) const override {
-    if (!op.shapeDimIndexes().size()) {
+    if (!op.getShapeDimIndexes().size()) {
       return failure();
     }
 
     SmallVector<bool> sizeOneDimVec;
     SmallVector<Value> newShapeDimIndexes;
     SmallVector<Type> newResultTypes;
-    for (Value dimSize : op.shapeDimIndexes()) {
+    for (Value dimSize : op.getShapeDimIndexes()) {
       auto constOp =
           dyn_cast_or_null<arith::ConstantIndexOp>(dimSize.getDefiningOp());
       bool isSizeOne =
@@ -189,12 +191,12 @@ struct RemoveSizeOneDimOfDelinearizeOp
       }
     }
 
-    if (newShapeDimIndexes.size() == op.shapeDimIndexes().size()) {
+    if (newShapeDimIndexes.size() == op.getShapeDimIndexes().size()) {
       return failure();
     }
 
     auto newOp = rewriter.create<DelinearizeOp>(
-        op.getLoc(), newResultTypes, op.linearIndex(), newShapeDimIndexes);
+        op.getLoc(), newResultTypes, op.getLinearIndex(), newShapeDimIndexes);
     Value zero = rewriter.create<arith::ConstantIndexOp>(op.getLoc(), 0);
 
     int nonSizeOneDimIdx = 0;
@@ -219,7 +221,7 @@ struct IdentityTieShapeOp : public OpRewritePattern<TieShapeOp> {
                                 PatternRewriter& rewriter) const override {
     // Do not touch tie_shape op with symbolic dim ref attrs.
     if (op->hasAttr(SymbolicDimOp::getSymbolicDimAttrName())) return success();
-    Value operand = op.value();
+    Value operand = op.getValue();
     auto operandTy = operand.getType().dyn_cast<RankedTensorType>();
     if (!operandTy) return failure();
 
@@ -228,7 +230,7 @@ struct IdentityTieShapeOp : public OpRewritePattern<TieShapeOp> {
     //   %d1 = tensor.dim %0, %c1, or %d1 is a constant
     bool allDimMatch = true;
     for (auto& en : llvm::enumerate(
-             llvm::zip(operandTy.getShape(), op.shapeDimIndexes()))) {
+             llvm::zip(operandTy.getShape(), op.getShapeDimIndexes()))) {
       int64_t idx = en.index();
       int64_t staticDim = std::get<0>(en.value());
       Value dynamicDim = std::get<1>(en.value());
@@ -318,11 +320,11 @@ void SymbolicDimOp::setDimSize(int64_t val) {
   OpBuilder b(*this);
   (*this)->setAttr("value", b.getI64IntegerAttr(val));
   if (val == -1) {
-    setKnownNegativeOne(true);
+    updateKnownNegativeOne(true);
   } else if (val >= 0) {
-    setKnownNonNegative(true);
-    if (val != 0) setKnownNonSizeZero(true);
-    if (val != 1) setKnownNonSizeOne(true);
+    updateKnownNonNegative(true);
+    if (val != 0) updateKnownNonSizeZero(true);
+    if (val != 1) updateKnownNonSizeOne(true);
   }
 }
 
@@ -330,26 +332,26 @@ bool SymbolicDimOp::isDynamic() {
   return getDimSize() == ShapedType::kDynamicSize;
 }
 
-void SymbolicDimOp::setKnownNonNegative(bool flag) {
+void SymbolicDimOp::updateKnownNonNegative(bool flag) {
   OpBuilder b(*this);
   (*this)->setAttr("knownNonNegative", b.getBoolAttr(flag));
 }
 
-void SymbolicDimOp::setKnownNegativeOne(bool flag) {
+void SymbolicDimOp::updateKnownNegativeOne(bool flag) {
   OpBuilder b(*this);
   (*this)->setAttr("knownNegativeOne", b.getBoolAttr(flag));
   if (flag) {
-    setKnownNonSizeOne(true);
-    setKnownNonSizeZero(true);
+    updateKnownNonSizeOne(true);
+    updateKnownNonSizeZero(true);
   }
 }
 
-void SymbolicDimOp::setKnownNonSizeOne(bool flag) {
+void SymbolicDimOp::updateKnownNonSizeOne(bool flag) {
   OpBuilder b(*this);
   (*this)->setAttr("knownNonSizeOne", b.getBoolAttr(flag));
 }
 
-void SymbolicDimOp::setKnownNonSizeZero(bool flag) {
+void SymbolicDimOp::updateKnownNonSizeZero(bool flag) {
   OpBuilder b(*this);
   (*this)->setAttr("knownNonSizeZero", b.getBoolAttr(flag));
 }
@@ -359,19 +361,22 @@ LogicalResult SymbolicDimOp::Merge(SymbolicDimOp other) {
     return failure();
   if (isDynamic() && !other.isDynamic()) setDimSize(other.getDimSize());
 
-  bool knownNonNegativeFlag = knownNonNegative() || other.knownNonNegative();
-  bool knownNegativeOneFlag = knownNegativeOne() || other.knownNegativeOne();
-  bool knownNonSizeOneFlag =
-      knownNonSizeOne() || other.knownNonSizeOne() || knownNegativeOneFlag;
-  bool knownNonSizeZeroFlag =
-      knownNonSizeZero() || other.knownNonSizeZero() || knownNegativeOneFlag;
+  bool knownNonNegativeFlag =
+      getKnownNonNegative() || other.getKnownNonNegative();
+  bool knownNegativeOneFlag =
+      getKnownNegativeOne() || other.getKnownNegativeOne();
+  bool knownNonSizeOneFlag = getKnownNonSizeOne() ||
+                             other.getKnownNonSizeOne() || knownNegativeOneFlag;
+  bool knownNonSizeZeroFlag = getKnownNonSizeZero() ||
+                              other.getKnownNonSizeZero() ||
+                              knownNegativeOneFlag;
 
   if (knownNonNegativeFlag && knownNegativeOneFlag) return failure();
 
-  setKnownNonSizeZero(knownNonSizeZeroFlag);
-  setKnownNonSizeOne(knownNonSizeOneFlag);
-  setKnownNegativeOne(knownNegativeOneFlag);
-  setKnownNonNegative(knownNonNegativeFlag);
+  updateKnownNonSizeZero(knownNonSizeZeroFlag);
+  updateKnownNonSizeOne(knownNonSizeOneFlag);
+  updateKnownNegativeOne(knownNegativeOneFlag);
+  updateKnownNonNegative(knownNonNegativeFlag);
 
   return success();
 }
diff --git a/tao_compiler/mlir/disc/IR/disc_shape_ops.td b/tao_compiler/mlir/disc/IR/disc_shape_ops.td
index 82c127acc32..e99b90f472c 100644
--- a/tao_compiler/mlir/disc/IR/disc_shape_ops.td
+++ b/tao_compiler/mlir/disc/IR/disc_shape_ops.td
@@ -32,7 +32,7 @@ class DISCShape_Op<string mnemonic, list<Trait> traits> :
   let hasVerifier = 1;
 }
 
-def DISCShape_LinearizeOp: DISCShape_Op<"linearize", [NoSideEffect, AttrSizedOperandSegments]> {
+def DISCShape_LinearizeOp: DISCShape_Op<"linearize", [Pure, AttrSizedOperandSegments]> {
   let summary = "calculate the linear index from multidim index";
   let description = [{
     Calculate the linear index using provided multidim indexes and shape.
@@ -46,7 +46,7 @@ def DISCShape_LinearizeOp: DISCShape_Op<"linearize", [NoSideEffect, AttrSizedOpe
   let hasCanonicalizer = 1;
 }
 
-def DISCShape_DelinearizeOp: DISCShape_Op<"delinearize", [NoSideEffect]> {
+def DISCShape_DelinearizeOp: DISCShape_Op<"delinearize", [Pure]> {
   let summary = "calculate the multidim index from linear index";
   let description = [{
     Calculate the multidim index using provided linear index and shape.
@@ -60,7 +60,7 @@ def DISCShape_DelinearizeOp: DISCShape_Op<"delinearize", [NoSideEffect]> {
   let hasCanonicalizer = 1;
 }
 
-def DISCShape_TieShapeOp: DISCShape_Op<"tie_shape", [NoSideEffect]> {
+def DISCShape_TieShapeOp: DISCShape_Op<"tie_shape", [Pure]> {
   let summary = "binding a set of multidim index with a tensor value";
   let description = [{
     Binding a set of multidim index with a tensor value.
@@ -104,16 +104,16 @@ def DISCShape_SymbolicDimOp:
     bool knownSizeOne() { return getDimSize() == 1; }
 
     // Sets `knownNonNegative` to the value of `flag`
-    void setKnownNonNegative(bool flag);
+    void updateKnownNonNegative(bool flag);
 
     // Sets `knownNegativeOne` to the value of `flag`
-    void setKnownNegativeOne(bool flag);
+    void updateKnownNegativeOne(bool flag);
 
     // Sets `knownNonSizeOne` to the value of `flag`
-    void setKnownNonSizeOne(bool flag);
+    void updateKnownNonSizeOne(bool flag);
 
     // Sets `knownNonSizeZero` to the value of `flag`
-    void setKnownNonSizeZero(bool flag);
+    void updateKnownNonSizeZero(bool flag);
 
     // Try to merge two SymbolicDim ops.
     // Returns failure if fail to merge.
@@ -127,7 +127,7 @@ def DISCShape_SymbolicDimOp:
 }
 
 def DISCShape_DimOp:
-      DISCShape_Op<"dim", [NoSideEffect]> {
+      DISCShape_Op<"dim", [Pure]> {
   let summary = "binding to a symbolic dimension size";
   let description = [{
     Binding to a symbolic dimension size.
diff --git a/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.td b/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.td
index 3daa1e730df..4d81d75c8e3 100644
--- a/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.td
+++ b/tao_compiler/mlir/disc/IR/disc_tf_additional_ops.td
@@ -27,7 +27,7 @@ include "mlir/IR/SymbolInterfaces.td"
 
 
 
-def TF_DiscFakeQuantOp : TF_Op<"DiscFakeQuant", [NoSideEffect, TF_NoConstantFold]> {
+def TF_DiscFakeQuantOp : TF_Op<"DiscFakeQuant", [Pure, TF_NoConstantFold]> {
 
   let summary = "FakeQuantOp defined by DISC to carry quantization info.";
 
diff --git a/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc b/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
index 9f043e3f6ef..e34a4999f10 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.cc
@@ -18,7 +18,7 @@ limitations under the License.
 #include <unordered_map>
 
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "tensorflow/compiler/mlir/disc/IR/custom_call_base.h"
 #include "tensorflow/compiler/mlir/disc/IR/hlo_disc_enums.cc.inc"
@@ -81,8 +81,8 @@ LogicalResult CustomCallOp::reifyReturnTypeShapes(
     OpBuilder& builder, ValueRange operands,
     SmallVectorImpl<Value>& reifiedReturnShapes) {
   CustomCallOp::Adaptor adaptor(operands);
-  ValueRange args = adaptor.args();
-  StringRef target = call_target_name();
+  ValueRange args = adaptor.getArgs();
+  StringRef target = getCallTargetName();
   auto reify_shapes_func =
       CustomCallRegistry::Global().FindReifyShapesFunc(target.str());
   if (!reify_shapes_func) {
@@ -99,18 +99,19 @@ LogicalResult CustomCallOp::verify() { return Verify(*this); }
 
 template <typename T>
 LogicalResult QuantVerify(T* op) {
-  auto inputTy = op->input().getType().template dyn_cast<RankedTensorType>();
-  auto scaleTy = op->scale().getType().template dyn_cast<RankedTensorType>();
+  auto inputTy = op->getInput().getType().template dyn_cast<RankedTensorType>();
+  auto scaleTy = op->getScale().getType().template dyn_cast<RankedTensorType>();
   auto zeroPointTy =
-      op->zero_point().getType().template dyn_cast<RankedTensorType>();
-  auto resultTy = op->result().getType().template dyn_cast<RankedTensorType>();
+      op->getZeroPoint().getType().template dyn_cast<RankedTensorType>();
+  auto resultTy =
+      op->getResult().getType().template dyn_cast<RankedTensorType>();
   if (!inputTy || !scaleTy || !zeroPointTy || !resultTy)
     return op->emitOpError() << "only support ranked input.\n";
   if (inputTy.getShape() != resultTy.getShape())
     return op->emitOpError() << "input and result have mismatch shape.\n";
   if (scaleTy.getRank() != zeroPointTy.getRank())
     return op->emitOpError() << "scale and zero_point have mismatch rank.\n";
-  auto axis = op->axis().template getValues<int64_t>();
+  auto axis = op->getAxis().template getValues<int64_t>();
   if (axis.size() != scaleTy.getRank())
     return op->emitOpError() << "num of quantized axes (len(axis)) is not "
                                 "equal to the rank of scale tensor\n";
@@ -137,9 +138,11 @@ LogicalResult DequantizeOp::verify() { return QuantVerify(this); }
 
 template <typename T>
 LogicalResult CommonVerifyForQuantizedComputeIntensiveOp(T* op) {
-  auto inputTy = op->input().getType().template dyn_cast<RankedTensorType>();
-  auto weightTy = op->weight().getType().template dyn_cast<RankedTensorType>();
-  auto resultTy = op->result().getType().template dyn_cast<RankedTensorType>();
+  auto inputTy = op->getInput().getType().template dyn_cast<RankedTensorType>();
+  auto weightTy =
+      op->getWeight().getType().template dyn_cast<RankedTensorType>();
+  auto resultTy =
+      op->getResult().getType().template dyn_cast<RankedTensorType>();
 
   if (!inputTy || !weightTy || !resultTy ||
       inputTy.getRank() != weightTy.getRank() ||
@@ -149,9 +152,9 @@ LogicalResult CommonVerifyForQuantizedComputeIntensiveOp(T* op) {
   }
 
   auto inputScaleTy =
-      op->input_scale().getType().template dyn_cast<RankedTensorType>();
+      op->getInputScale().getType().template dyn_cast<RankedTensorType>();
   auto inputZeroPointTy =
-      op->input_zero_point().getType().template dyn_cast<RankedTensorType>();
+      op->getInputZeroPoint().getType().template dyn_cast<RankedTensorType>();
   if (!inputScaleTy || !inputZeroPointTy || inputScaleTy.getRank() != 0 ||
       inputZeroPointTy.getRank() != 0) {
     return op->emitOpError() << "input_scale and input_zero_point only support "
@@ -159,9 +162,9 @@ LogicalResult CommonVerifyForQuantizedComputeIntensiveOp(T* op) {
   }
 
   auto resultScaleTy =
-      op->result_scale().getType().template dyn_cast<RankedTensorType>();
+      op->getResultScale().getType().template dyn_cast<RankedTensorType>();
   auto resultZeroPointTy =
-      op->result_zero_point().getType().template dyn_cast<RankedTensorType>();
+      op->getResultZeroPoint().getType().template dyn_cast<RankedTensorType>();
   if (!resultScaleTy || !resultZeroPointTy || resultScaleTy.getRank() != 0 ||
       resultZeroPointTy.getRank() != 0) {
     return op->emitOpError() << "result_scale and result_zero_point only "
@@ -169,15 +172,15 @@ LogicalResult CommonVerifyForQuantizedComputeIntensiveOp(T* op) {
   }
 
   auto weightScaleTy =
-      op->weight_scale().getType().template dyn_cast<RankedTensorType>();
+      op->getWeightScale().getType().template dyn_cast<RankedTensorType>();
   auto weightZeroPointTy =
-      op->weight_zero_point().getType().template dyn_cast<RankedTensorType>();
+      op->getWeightZeroPoint().getType().template dyn_cast<RankedTensorType>();
   if (!weightScaleTy || !weightZeroPointTy ||
       weightScaleTy.getShape() != weightZeroPointTy.getShape()) {
     return op->emitOpError()
            << "weight_scale and weight_zero_point have mismatch shape\n";
   }
-  auto axis = op->axis().template getValues<int64_t>();
+  auto axis = op->getAxis().template getValues<int64_t>();
   if (axis.size() != weightScaleTy.getRank())
     return op->emitOpError() << "num of quantized axes (len(axis)) is not "
                                 "equal to the rank of weight_scale tensor\n";
@@ -191,32 +194,32 @@ LogicalResult QuantizedDotGeneralOp::verify() {
 LogicalResult QuantizedDotGeneralOp::reifyReturnTypeShapes(
     OpBuilder& builder, ValueRange operands,
     SmallVectorImpl<Value>& reifiedReturnShapes) {
-  auto lhsType = input().getType().dyn_cast<ShapedType>();
-  auto rhsType = weight().getType().dyn_cast<ShapedType>();
+  auto lhsType = getInput().getType().dyn_cast<ShapedType>();
+  auto rhsType = getWeight().getType().dyn_cast<ShapedType>();
   if (!lhsType || !rhsType) {
     return failure();
   }
 
   Adaptor adaptor(operands);
-  auto dimNumbers = dot_dimension_numbers();
+  auto dimNumbers = getDotDimensionNumbers();
   SmallVector<Value> dimensions;
   for (const int64_t lhsDim : dimNumbers.getLhsBatchingDimensions()) {
     dimensions.push_back(
-        builder.create<tensor::DimOp>(getLoc(), adaptor.input(), lhsDim));
+        builder.create<tensor::DimOp>(getLoc(), adaptor.getInput(), lhsDim));
   }
 
   for (int64_t i = 0; i < lhsType.getRank(); i++) {
     if (!llvm::is_contained(dimNumbers.getLhsContractingDimensions(), i) &&
         !llvm::is_contained(dimNumbers.getLhsBatchingDimensions(), i)) {
       dimensions.push_back(
-          builder.create<tensor::DimOp>(getLoc(), adaptor.input(), i));
+          builder.create<tensor::DimOp>(getLoc(), adaptor.getInput(), i));
     }
   }
   for (int64_t i = 0; i < rhsType.getRank(); i++) {
     if (!llvm::is_contained(dimNumbers.getRhsContractingDimensions(), i) &&
         !llvm::is_contained(dimNumbers.getRhsBatchingDimensions(), i)) {
       dimensions.push_back(
-          builder.create<tensor::DimOp>(getLoc(), adaptor.weight(), i));
+          builder.create<tensor::DimOp>(getLoc(), adaptor.getWeight(), i));
     }
   }
   reifiedReturnShapes.push_back(
@@ -254,8 +257,8 @@ LogicalResult ConvReifyReturnTypeImpl(
     SmallVectorImpl<Value>& reifiedReturnShapes,
     const SmallVector<Value>& spatial_padding_values, Type shape_scalar_type) {
   typename Op::Adaptor adaptor(operands);
-  Value lhs = adaptor.input();
-  Value rhs = adaptor.weight();
+  Value lhs = adaptor.getInput();
+  Value rhs = adaptor.getWeight();
 
   RankedTensorType lhs_type = lhs.getType().dyn_cast<RankedTensorType>();
   RankedTensorType rhs_type = rhs.getType().dyn_cast<RankedTensorType>();
@@ -268,7 +271,7 @@ LogicalResult ConvReifyReturnTypeImpl(
     return maybeCastTo(builder, loc, v, shape_scalar_type);
   };
 
-  auto dimension_numbers = op->dimension_numbers();
+  auto dimension_numbers = op->getDimensionNumbers();
   int64_t input_batch_dimension = dimension_numbers.getInputBatchDimension();
   int64_t kernel_output_feature_dimension =
       dimension_numbers.getKernelOutputFeatureDimension();
@@ -285,7 +288,7 @@ LogicalResult ConvReifyReturnTypeImpl(
   Value lhs_batch_dim = to_shape_scalar_type(
       builder.create<tensor::DimOp>(loc, lhs, input_batch_dimension));
   Value batch_group_count = to_shape_scalar_type(
-      builder.create<arith::ConstantIndexOp>(loc, op->batch_group_count()));
+      builder.create<arith::ConstantIndexOp>(loc, op->getBatchGroupCount()));
   Value batch_dim = to_shape_scalar_type(
       builder.create<arith::DivSIOp>(loc, lhs_batch_dim, batch_group_count));
   int64_t output_batch_dimension = dimension_numbers.getOutputBatchDimension();
@@ -298,9 +301,9 @@ LogicalResult ConvReifyReturnTypeImpl(
       dimension_numbers.getOutputFeatureDimension();
   shape_values[output_feature_dimension] = feature_dim;
 
-  Optional<DenseIntElementsAttr> window_strides_attr = op->window_strides();
-  Optional<DenseIntElementsAttr> lhs_dilation_attr = op->lhs_dilation();
-  Optional<DenseIntElementsAttr> rhs_dilation_attr = op->rhs_dilation();
+  Optional<DenseIntElementsAttr> window_strides_attr = op->getWindowStrides();
+  Optional<DenseIntElementsAttr> lhs_dilation_attr = op->getLhsDilation();
+  Optional<DenseIntElementsAttr> rhs_dilation_attr = op->getRhsDilation();
 
   Value one =
       to_shape_scalar_type(builder.create<arith::ConstantIndexOp>(loc, 1));
@@ -374,7 +377,7 @@ LogicalResult QuantizedDynamicConvOp::reifyReturnTypeShapes(
     OpBuilder& builder, ValueRange operands,
     SmallVectorImpl<Value>& reifiedReturnShapes) {
   QuantizedDynamicConvOp::Adaptor adaptor(operands);
-  Value d_padding = adaptor.d_padding();
+  Value d_padding = adaptor.getDPadding();
 
   RankedTensorType padding_type =
       d_padding.getType().dyn_cast<RankedTensorType>();
@@ -388,7 +391,7 @@ LogicalResult QuantizedDynamicConvOp::reifyReturnTypeShapes(
   };
 
   SmallVector<Value> spatial_padding_values;
-  auto dimension_numbers = this->dimension_numbers();
+  auto dimension_numbers = this->getDimensionNumbers();
   auto input_spatial_dimensions_attr =
       dimension_numbers.getInputSpatialDimensions();
   int64_t padding_num = input_spatial_dimensions_attr.size() * 2;
@@ -412,11 +415,11 @@ LogicalResult SparseReshapeOp::reifyReturnTypeShapes(
     SmallVectorImpl<Value>& reifiedReturnShapes) {
   SparseReshapeOp::Adaptor adaptor(operands);
   auto input_indices_type =
-      adaptor.input_indices().getType().dyn_cast<RankedTensorType>();
+      adaptor.getInputIndices().getType().dyn_cast<RankedTensorType>();
   auto input_shape_type =
-      adaptor.input_shape().getType().dyn_cast<RankedTensorType>();
+      adaptor.getInputShape().getType().dyn_cast<RankedTensorType>();
   auto new_shape_type =
-      adaptor.new_shape().getType().dyn_cast<RankedTensorType>();
+      adaptor.getNewShape().getType().dyn_cast<RankedTensorType>();
   if (!input_indices_type || !input_shape_type || !new_shape_type) {
     return failure();
   }
@@ -449,15 +452,15 @@ LogicalResult SparseReshapeOp::reifyReturnTypeShapes(
 
 LogicalResult SparseReshapeOp::verify() {
   auto input_indices_type =
-      this->input_indices().getType().template dyn_cast<RankedTensorType>();
+      this->getInputIndices().getType().template dyn_cast<RankedTensorType>();
   auto input_shape_type =
-      this->input_shape().getType().template dyn_cast<RankedTensorType>();
+      this->getInputShape().getType().template dyn_cast<RankedTensorType>();
   auto new_shape_type =
-      this->new_shape().getType().template dyn_cast<RankedTensorType>();
+      this->getNewShape().getType().template dyn_cast<RankedTensorType>();
   auto output_indices_type =
-      this->output_indices().getType().template dyn_cast<RankedTensorType>();
+      this->getOutputIndices().getType().template dyn_cast<RankedTensorType>();
   auto output_shape_type =
-      this->output_shape().getType().template dyn_cast<RankedTensorType>();
+      this->getOutputShape().getType().template dyn_cast<RankedTensorType>();
 
   if (!input_indices_type || !input_shape_type || !new_shape_type) {
     return this->emitOpError() << "only support ranked input.\n";
@@ -488,15 +491,15 @@ LogicalResult SparseFillEmptyRowsOp::reifyReturnTypeShapes(
     SmallVectorImpl<Value>& reifiedReturnShapes) {
   SparseFillEmptyRowsOp::Adaptor adaptor(operands);
   // index 0
-  auto indices_type = adaptor.indices().getType().cast<RankedTensorType>();
+  auto indices_type = adaptor.getIndices().getType().cast<RankedTensorType>();
   // index 1
-  auto values_type = adaptor.values().getType().cast<RankedTensorType>();
+  auto values_type = adaptor.getValues().getType().cast<RankedTensorType>();
   // index 2
   auto dense_shape_type =
-      adaptor.dense_shape().getType().cast<RankedTensorType>();
+      adaptor.getDenseShape().getType().cast<RankedTensorType>();
   // index 3
   auto default_value_type =
-      adaptor.default_value().getType().cast<RankedTensorType>();
+      adaptor.getDefaultValue().getType().cast<RankedTensorType>();
 
   Location loc = this->getLoc();
 
@@ -560,12 +563,12 @@ LogicalResult SparseFillEmptyRowsOp::reifyReturnTypeShapes(
 }
 
 LogicalResult SparseFillEmptyRowsOp::verify() {
-  auto indices_type = this->indices().getType().dyn_cast<RankedTensorType>();
-  auto values_type = this->values().getType().dyn_cast<RankedTensorType>();
+  auto indices_type = this->getIndices().getType().dyn_cast<RankedTensorType>();
+  auto values_type = this->getValues().getType().dyn_cast<RankedTensorType>();
   auto dense_shape_type =
-      this->dense_shape().getType().dyn_cast<RankedTensorType>();
+      this->getDenseShape().getType().dyn_cast<RankedTensorType>();
   auto default_value_type =
-      this->default_value().getType().dyn_cast<RankedTensorType>();
+      this->getDefaultValue().getType().dyn_cast<RankedTensorType>();
 
   if (!indices_type || !values_type || !default_value_type) {
     return failure();
@@ -590,15 +593,15 @@ LogicalResult SparseFillEmptyRowsOp::verify() {
   }
 
   auto output_indices_type =
-      this->output_indices().getType().dyn_cast<RankedTensorType>();
+      this->getOutputIndices().getType().dyn_cast<RankedTensorType>();
   auto output_values_type =
-      this->output_values().getType().dyn_cast<RankedTensorType>();
+      this->getOutputValues().getType().dyn_cast<RankedTensorType>();
   auto empty_row_indicator_type =
-      this->empty_row_indicator().getType().dyn_cast<RankedTensorType>();
+      this->getEmptyRowIndicator().getType().dyn_cast<RankedTensorType>();
   auto reverse_index_map_type =
-      this->reverse_index_map().getType().dyn_cast<RankedTensorType>();
+      this->getReverseIndexMap().getType().dyn_cast<RankedTensorType>();
   auto output_elements_type =
-      this->output_elements().getType().dyn_cast<RankedTensorType>();
+      this->getOutputElements().getType().dyn_cast<RankedTensorType>();
   if (!output_indices_type || !output_values_type ||
       !empty_row_indicator_type || !reverse_index_map_type ||
       !output_elements_type) {
@@ -630,10 +633,10 @@ LogicalResult SparseSegmentMeanOp::reifyReturnTypeShapes(
     SmallVectorImpl<Value>& reifiedReturnShapes) {
   SparseSegmentMeanOp::Adaptor adaptor(operands);
   Location loc = this->getLoc();
-  auto data_type = adaptor.data().getType().cast<RankedTensorType>();
-  auto indices_type = adaptor.indices().getType().cast<RankedTensorType>();
+  auto data_type = adaptor.getData().getType().cast<RankedTensorType>();
+  auto indices_type = adaptor.getIndices().getType().cast<RankedTensorType>();
   auto segment_ids_type =
-      adaptor.segment_ids().getType().cast<RankedTensorType>();
+      adaptor.getSegmentIds().getType().cast<RankedTensorType>();
   auto input_rank = data_type.getRank();
   SmallVector<Value, 2> output_shape_values;
   Value segment_size = builder.create<tensor::DimOp>(loc, operands[2], 0);
@@ -659,10 +662,10 @@ LogicalResult SparseSegmentMeanOp::reifyReturnTypeShapes(
 }
 
 LogicalResult SparseSegmentMeanOp::verify() {
-  auto data_type = this->data().getType().dyn_cast<RankedTensorType>();
-  auto indices_type = this->indices().getType().dyn_cast<RankedTensorType>();
+  auto data_type = this->getData().getType().dyn_cast<RankedTensorType>();
+  auto indices_type = this->getIndices().getType().dyn_cast<RankedTensorType>();
   auto segment_ids_type =
-      this->segment_ids().getType().dyn_cast<RankedTensorType>();
+      this->getSegmentIds().getType().dyn_cast<RankedTensorType>();
 
   if (!data_type || !indices_type || !segment_ids_type) {
     return failure();
@@ -688,7 +691,7 @@ LogicalResult SparseSegmentMeanOp::verify() {
     return this->emitOpError() << "input must be at least rank 1";
   }
 
-  auto output_type = this->output().getType().dyn_cast<RankedTensorType>();
+  auto output_type = this->getOutput().getType().dyn_cast<RankedTensorType>();
   if (!output_type) {
     return failure();
   }
@@ -713,7 +716,7 @@ LogicalResult WhereOp::reifyReturnTypeShapes(
     SmallVectorImpl<Value>& reifiedReturnShapes) {
   WhereOp::Adaptor adaptor(operands);
   Location loc = this->getLoc();
-  auto input_type = adaptor.input().getType().cast<RankedTensorType>();
+  auto input_type = adaptor.getInput().getType().cast<RankedTensorType>();
 
   SmallVector<Value, 2> index_shape_values, num_output_elements_shape_values;
 
@@ -742,13 +745,13 @@ LogicalResult WhereOp::reifyReturnTypeShapes(
 }
 
 LogicalResult WhereOp::verify() {
-  auto input_type = this->input().getType().dyn_cast<RankedTensorType>();
+  auto input_type = this->getInput().getType().dyn_cast<RankedTensorType>();
 
   if (!input_type) {
     return failure();
   }
 
-  auto index_type = this->index().getType().dyn_cast<RankedTensorType>();
+  auto index_type = this->getIndex().getType().dyn_cast<RankedTensorType>();
   if (!index_type) {
     return failure();
   }
diff --git a/tao_compiler/mlir/disc/IR/hlo_disc_ops.td b/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
index ea126b14524..13ad0233369 100644
--- a/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
+++ b/tao_compiler/mlir/disc/IR/hlo_disc_ops.td
@@ -40,7 +40,7 @@ class HLODISC_ShapedInterfaceOp<string mnemonic, list<Trait> traits> :
     ["reifyReturnTypeShapes"]>]> {
 }
 
-def HLODISC_H2DOp: HLODISC_ShapedInterfaceOp<"h2d", [NoSideEffect, SameOperandsAndResultType]> {
+def HLODISC_H2DOp: HLODISC_ShapedInterfaceOp<"h2d", [Pure, SameOperandsAndResultType]> {
   let summary = "H2D operator";
   let description = [{
     Copy `operand` from host to device.
@@ -49,7 +49,7 @@ def HLODISC_H2DOp: HLODISC_ShapedInterfaceOp<"h2d", [NoSideEffect, SameOperandsA
   let results = (outs HLO_Tensor);
 }
 
-def HLODISC_D2HOp: HLODISC_ShapedInterfaceOp<"d2h", [NoSideEffect, SameOperandsAndResultType]> {
+def HLODISC_D2HOp: HLODISC_ShapedInterfaceOp<"d2h", [Pure, SameOperandsAndResultType]> {
   let summary = "D2H operator";
   let description = [{
     Copy `operand` from device to host.
@@ -118,7 +118,7 @@ def HLODISC_CustomCallV2Op: HLODISC_Op<"custom_call_v2", []> {
   let results = (outs Variadic<HLO_TensorOrIndexTensor>);
 }
 
-def HLO_FakeQuantOp: HLODISC_Op<"fake_quant", [NoSideEffect]> {
+def HLO_FakeQuantOp: HLODISC_Op<"fake_quant", [Pure]> {
   let summary = "FakeQuant operator";
   let description = [{
     An fake quant op is logically equal to a quant op + a dequant op.
@@ -139,7 +139,7 @@ def HLO_FakeQuantOp: HLODISC_Op<"fake_quant", [NoSideEffect]> {
   let results = (outs HLO_FpTensor:$result);
 }
 
-def HLO_QuantizeOp: HLODISC_Op<"quantize", [NoSideEffect]> {
+def HLO_QuantizeOp: HLODISC_Op<"quantize", [Pure]> {
   let summary = "Quantize operator";
   let description = [{
     Quantize a float tensor to int tensor.
@@ -158,7 +158,7 @@ def HLO_QuantizeOp: HLODISC_Op<"quantize", [NoSideEffect]> {
   let results = (outs HLO_IntTensor:$result);
 }
 
-def HLO_DequantizeOp: HLODISC_Op<"dequantize", [NoSideEffect]> {
+def HLO_DequantizeOp: HLODISC_Op<"dequantize", [Pure]> {
   let summary = "Dequantize operator";
   let description = [{
     Dequantize a int tensor to float tensor.
@@ -175,7 +175,7 @@ def HLO_DequantizeOp: HLODISC_Op<"dequantize", [NoSideEffect]> {
   let results = (outs HLO_FpTensor:$result);
 }
 
-def HLO_QuantizedDotGeneralOp: HLODISC_ShapedInterfaceOp<"quantized_dot_general", [NoSideEffect]> {
+def HLO_QuantizedDotGeneralOp: HLODISC_ShapedInterfaceOp<"quantized_dot_general", [Pure]> {
   let summary = "quantized version of dot gerneal operator";
   let description = [{
     Compute the dot product using result quantized inputs.
@@ -197,7 +197,7 @@ def HLO_QuantizedDotGeneralOp: HLODISC_ShapedInterfaceOp<"quantized_dot_general"
   let results = (outs HLO_Tensor:$result);
 }
 
-def HLO_QuantizedDynamicConvOp: HLODISC_ShapedInterfaceOp<"quantized_dynamic_conv", [NoSideEffect]> {
+def HLO_QuantizedDynamicConvOp: HLODISC_ShapedInterfaceOp<"quantized_dynamic_conv", [Pure]> {
   let summary = "quantized version of dynamic_conv operator";
   let description = [{
     Compute the convolution result using quantized inputs.
@@ -232,7 +232,7 @@ def HLO_QuantizedDynamicConvOp: HLODISC_ShapedInterfaceOp<"quantized_dynamic_con
   let results = (outs HLO_Tensor:$result);
 }
 
-def HLO_SparseReshapeOp : HLODISC_ShapedInterfaceOp<"sparse_reshape", [NoSideEffect]> {
+def HLO_SparseReshapeOp : HLODISC_ShapedInterfaceOp<"sparse_reshape", [Pure]> {
   let summary = "One-to-one mapping of TF_SparseReshapeOp";
   let description = [{
     See tf's sparse_reshape operator.
@@ -252,7 +252,7 @@ def HLO_SparseReshapeOp : HLODISC_ShapedInterfaceOp<"sparse_reshape", [NoSideEff
 // possible types for SparseFillEmptyRows:
 // Supported integral types: int64/int32/uint16/int16/uint8/int8
 // Supported float types: half/bfloat16/float/double
-def HLO_SparseFillEmptyRowsOp : HLODISC_ShapedInterfaceOp<"sparse_fill_empty_rows", [NoSideEffect]> {
+def HLO_SparseFillEmptyRowsOp : HLODISC_ShapedInterfaceOp<"sparse_fill_empty_rows", [Pure]> {
   let summary = "One-to-one mapping of TF_SparseFillEmptyRowsOp";
   let description = [{
     See tf's sparse_fill_empty_rows operator.
@@ -273,7 +273,7 @@ def HLO_SparseFillEmptyRowsOp : HLODISC_ShapedInterfaceOp<"sparse_fill_empty_row
   );
 }
 
-def HLO_SparseSegmentMeanOp : HLODISC_ShapedInterfaceOp<"sparse_segment_mean", [NoSideEffect]> {
+def HLO_SparseSegmentMeanOp : HLODISC_ShapedInterfaceOp<"sparse_segment_mean", [Pure]> {
   let summary = "One-to-one mapping of TF_SparseSegmentMeanOp";
   let description = [{
     See tf's sparse_segment_mean operator.
@@ -289,7 +289,7 @@ def HLO_SparseSegmentMeanOp : HLODISC_ShapedInterfaceOp<"sparse_segment_mean", [
   );
 }
 
-def HLO_WhereOp : HLODISC_ShapedInterfaceOp<"where", [NoSideEffect]> {
+def HLO_WhereOp : HLODISC_ShapedInterfaceOp<"where", [Pure]> {
   let summary = "One-to-one mapping of TF_WhereOp";
   let description = [{
     See tf's where operator.
diff --git a/tao_compiler/mlir/disc/IR/qconv_custom_call_op.cc b/tao_compiler/mlir/disc/IR/qconv_custom_call_op.cc
index e60f970bcbb..1254121cc7a 100644
--- a/tao_compiler/mlir/disc/IR/qconv_custom_call_op.cc
+++ b/tao_compiler/mlir/disc/IR/qconv_custom_call_op.cc
@@ -14,7 +14,7 @@ limitations under the License.
 
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
@@ -65,7 +65,7 @@ LogicalResult reifyReturnTypeShapesDconvI8I8I8Impl(
   };
 
   SmallVector<Value> spatial_padding_values;
-  auto config = op.backend_config().cast<DictionaryAttr>();
+  auto config = op.getBackendConfig().cast<DictionaryAttr>();
   auto dimension_numbers =
       config.getAs<mhlo::ConvDimensionNumbersAttr>("dimension_numbers");
   auto input_spatial_dimensions_attr =
diff --git a/tao_compiler/mlir/disc/IR/rng_uniform_custom_call_op.cc b/tao_compiler/mlir/disc/IR/rng_uniform_custom_call_op.cc
index 450f6d7f6ed..d9a2712e299 100644
--- a/tao_compiler/mlir/disc/IR/rng_uniform_custom_call_op.cc
+++ b/tao_compiler/mlir/disc/IR/rng_uniform_custom_call_op.cc
@@ -17,7 +17,7 @@ limitations under the License.
 #include <atomic>
 
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
diff --git a/tao_compiler/mlir/disc/IR/sparse_custom_call_op.cc b/tao_compiler/mlir/disc/IR/sparse_custom_call_op.cc
index 9e302815cd9..6f48f98d3a6 100644
--- a/tao_compiler/mlir/disc/IR/sparse_custom_call_op.cc
+++ b/tao_compiler/mlir/disc/IR/sparse_custom_call_op.cc
@@ -13,7 +13,7 @@ limitations under the License.
 // This file defines sparse gemm related custom calls.
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
@@ -40,7 +40,7 @@ LogicalResult reifyReturnTypeShapesSparseGemmImpl(
   if (!lhsTy || !rhsTy) return op->emitError() << "not support unranked type\n";
 
   Location loc = op.getLoc();
-  auto config = op.backend_config().cast<DictionaryAttr>();
+  auto config = op.getBackendConfig().cast<DictionaryAttr>();
   int64_t lhs_contracting_dimensions =
       config.getAs<IntegerAttr>("lhs_contracting_dimensions").getInt();
   int64_t rhs_contracting_dimensions =
diff --git a/tao_compiler/mlir/disc/IR/topk_custom_call_op.cc b/tao_compiler/mlir/disc/IR/topk_custom_call_op.cc
index 445df2532d0..1ff405c06da 100644
--- a/tao_compiler/mlir/disc/IR/topk_custom_call_op.cc
+++ b/tao_compiler/mlir/disc/IR/topk_custom_call_op.cc
@@ -53,7 +53,7 @@ LogicalResult reifyReturnTypeShapesImpl<TopKBackendConfig>(
     SmallVectorImpl<Value>& reifiedReturnShapes) {
   llvm::Expected<TopKBackendConfig> backend_config =
       llvm::json::parse<TopKBackendConfig>(
-          op.backend_config().cast<StringAttr>());
+          op.getBackendConfig().cast<StringAttr>());
   int64_t dimension = backend_config->dimension;
 
   Value keys_operand = operands[0];
diff --git a/tao_compiler/mlir/disc/IR/topk_custom_call_op.h b/tao_compiler/mlir/disc/IR/topk_custom_call_op.h
index df37407e7c0..ca7631eafe0 100644
--- a/tao_compiler/mlir/disc/IR/topk_custom_call_op.h
+++ b/tao_compiler/mlir/disc/IR/topk_custom_call_op.h
@@ -31,7 +31,7 @@ using llvm::json::ObjectMapper;
 class CustomCallOp;
 
 // Methods to do deserialization:
-// llvm::json::parse<TopKBackendConfig>(custom_call_op.backend_config());
+// llvm::json::parse<TopKBackendConfig>(custom_call_op.getBackendConfig());
 //
 // Methods to do serialization:
 // raw_ostream << toJSON(topk_backend_config);
diff --git a/tao_compiler/mlir/disc/disc_compiler.cc b/tao_compiler/mlir/disc/disc_compiler.cc
index 440f2476fbf..dbf34dda2cd 100644
--- a/tao_compiler/mlir/disc/disc_compiler.cc
+++ b/tao_compiler/mlir/disc/disc_compiler.cc
@@ -37,7 +37,7 @@ limitations under the License.
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
 #include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"  // from @llvm-project
-#include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/Passes.h"
@@ -382,7 +382,7 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
 
   // bufferize constant ops & index_cast that have tensor types.
   pm.addNestedPass<FuncOp>(disc_ral::createDiscStdBufferizePass());
-  pm.addPass(arith::createArithmeticBufferizePass());
+  pm.addPass(arith::createArithBufferizePass());
   pm.addNestedPass<FuncOp>(createTensorBufferizePass());
   pm.addNestedPass<FuncOp>(bufferization::createFinalizingBufferizePass());
   pm.addNestedPass<FuncOp>(createCanonicalizerPass());
@@ -472,7 +472,7 @@ LogicalResult LowerHLOToLLVM(ModuleOp m, const DISCLoweringOptions& options) {
     // optimization. Then this pass will be enabled by default.
     pm.addNestedPass<FuncOp>(disc_ral::createForLoopUnrollInterleavePass());
   }
-  pm.addNestedPass<FuncOp>(arith::createArithmeticExpandOpsPass());
+  pm.addNestedPass<FuncOp>(arith::createArithExpandOpsPass());
   pm.addNestedPass<FuncOp>(mlir::memref::createFoldMemRefAliasOpsPass());
 
   // Flatten multi dim memref accesses to its 1D format to enable more
@@ -687,7 +687,8 @@ LogicalResult ApplyCpuOptionsBeforeTranslatingToLLVM(
           llvm::errs() << "[[DISC WARNING]] unknown fast_math_level value\n";
           break;
       }
-      op->setAttr("fastmathFlags", LLVM::FMFAttr::get(op->getContext(), fmf));
+      op->setAttr("fastmathFlags",
+                  LLVM::FastmathFlagsAttr::get(op->getContext(), fmf));
     });
     if (VLOG_IS_ON(1)) {
       llvm::dbgs() << "[[DISC DEBUG]] fastmath dump begin: \n"
diff --git a/tao_compiler/mlir/disc/disc_compiler_main.cc b/tao_compiler/mlir/disc/disc_compiler_main.cc
index a4f21987721..c6f737043f7 100644
--- a/tao_compiler/mlir/disc/disc_compiler_main.cc
+++ b/tao_compiler/mlir/disc/disc_compiler_main.cc
@@ -30,10 +30,10 @@ limitations under the License.
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Target/TargetMachine.h"
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
-#include "mlir-hlo/Dialect/lhlo/transforms/register_passes.h"
+#include "mlir-hlo/Dialect/lhlo/transforms/passes.h"
 #include "mlir-hlo/Dialect/lhlo_gpu/IR/lhlo_gpu_ops.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir-hlo/Dialect/mhlo/transforms/register_passes.h"
+#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
 #include "mlir/ExecutionEngine/OptUtils.h"  // from @llvm-project
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
diff --git a/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc b/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc
index 9d60d05f88e..ea5603edd24 100644
--- a/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc
+++ b/tao_compiler/mlir/disc/tools/disc-opt/disc-opt.cc
@@ -14,10 +14,10 @@ limitations under the License.
 ==============================================================================*/
 
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
-#include "mlir-hlo/Dialect/lhlo/transforms/register_passes.h"
+#include "mlir-hlo/Dialect/lhlo/transforms/passes.h"
 #include "mlir-hlo/Dialect/lhlo_gpu/IR/lhlo_gpu_ops.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir-hlo/Dialect/mhlo/transforms/register_passes.h"
+#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
diff --git a/tao_compiler/mlir/disc/tools/disc-replay/disc_interpreter.cc b/tao_compiler/mlir/disc/tools/disc-replay/disc_interpreter.cc
index 864421e7890..fc98cd0d87d 100644
--- a/tao_compiler/mlir/disc/tools/disc-replay/disc_interpreter.cc
+++ b/tao_compiler/mlir/disc/tools/disc-replay/disc_interpreter.cc
@@ -31,9 +31,10 @@ tensorflow::Status DiscInterpreter::Compile(
 
   // compile input proto to executable file
   tensorflow::DeviceType device_type(input.options().device_type());
-  auto* compiler_wrapper =
-      tensorflow::tao::CompilerBase::GetCompilerForDevice(device_type)
-          .ValueOrDie();
+  auto status_or =
+      tensorflow::tao::CompilerBase::GetCompilerForDevice(device_type);
+  if (!status_or.ok()) return status_or.status();
+  auto* compiler_wrapper = status_or.value();
   TF_RETURN_IF_ERROR(compiler_wrapper->Compile(input, tmp_file));
   result.output_fname = tmp_file + ".so";
   result.meta_fname = tmp_file + ".so.pbtxt";
diff --git a/tao_compiler/mlir/disc/transforms/PassDetail.h b/tao_compiler/mlir/disc/transforms/PassDetail.h
index f8d298dc9ec..f3759a530c4 100644
--- a/tao_compiler/mlir/disc/transforms/PassDetail.h
+++ b/tao_compiler/mlir/disc/transforms/PassDetail.h
@@ -50,6 +50,10 @@ namespace mhlo {
 class MhloDialect;
 }
 
+namespace arith {
+class ArithDialect;
+}
+
 namespace disc_ral {
 
 #define GEN_PASS_CLASSES
diff --git a/tao_compiler/mlir/disc/transforms/codegen_utils.cc b/tao_compiler/mlir/disc/transforms/codegen_utils.cc
index 934f6acfad5..345564285a5 100644
--- a/tao_compiler/mlir/disc/transforms/codegen_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/codegen_utils.cc
@@ -16,7 +16,7 @@ limitations under the License.
 
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Transforms/Passes.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
@@ -382,8 +382,7 @@ std::pair<ParallelOp, ParallelOp> tileParallelLoop(ParallelOp op,
       thenBlock.getArgument(ivs.index())
           .replaceAllUsesExcept(newIndex, newIndex);
     }
-    thenBlock.eraseArguments(llvm::to_vector<4>(
-        llvm::seq((unsigned)0, thenBlock.getNumArguments())));
+    thenBlock.eraseArguments(0, thenBlock.getNumArguments());
   } else {
     innerLoop.getLoopBody().takeBody(op.getLoopBody());
     b.setInsertionPointToStart(innerLoop.getBody());
diff --git a/tao_compiler/mlir/disc/transforms/conv_rewriter.cc b/tao_compiler/mlir/disc/transforms/conv_rewriter.cc
index e3625d2759c..7bec61ecdf1 100644
--- a/tao_compiler/mlir/disc/transforms/conv_rewriter.cc
+++ b/tao_compiler/mlir/disc/transforms/conv_rewriter.cc
@@ -62,7 +62,7 @@ LogicalResult fillConvParams(T op, ConvParams<T>& params) {
   auto inputTy = params.input.getType().template dyn_cast<RankedTensorType>();
   auto filterTy = params.filter.getType().template dyn_cast<RankedTensorType>();
   auto paddingTy =
-      op.d_padding().getType().template dyn_cast<RankedTensorType>();
+      op.getDPadding().getType().template dyn_cast<RankedTensorType>();
   auto outputTy = params.output.getType().template dyn_cast<RankedTensorType>();
 
   if (!inputTy || !filterTy || !paddingTy || !outputTy) {
@@ -76,7 +76,7 @@ LogicalResult fillConvParams(T op, ConvParams<T>& params) {
     return op.emitOpError() << "conv op's input rank is less than 3";
   }
 
-  auto dimension_numbers = op.dimension_numbers();
+  auto dimension_numbers = op.getDimensionNumbers();
   params.inputLayout.push_back(dimension_numbers.getInputBatchDimension());
   params.inputLayout.push_back(dimension_numbers.getInputFeatureDimension());
   auto input_spatial_dimensions = dimension_numbers.getInputSpatialDimensions();
@@ -107,8 +107,8 @@ LogicalResult fillConvParams(T op, ConvParams<T>& params) {
 LogicalResult extractConvParams(mhlo::DynamicConvOp op,
                                 ConvParams<mhlo::DynamicConvOp>& params) {
   params.conv = op;
-  params.input = op.lhs();
-  params.filter = op.rhs();
+  params.input = op.getLhs();
+  params.filter = op.getRhs();
   params.output = op.getResult();
 
   return fillConvParams<mhlo::DynamicConvOp>(op, params);
@@ -118,8 +118,8 @@ LogicalResult extractConvParams(
     mlir::mhlo_disc::QuantizedDynamicConvOp op,
     ConvParams<mlir::mhlo_disc::QuantizedDynamicConvOp>& params) {
   params.conv = op;
-  params.input = op.input();
-  params.filter = op.weight();
+  params.input = op.getInput();
+  params.filter = op.getWeight();
   params.output = op.getResult();
 
   return fillConvParams<mlir::mhlo_disc::QuantizedDynamicConvOp>(op, params);
@@ -188,8 +188,8 @@ struct ConvToDynamicConvConvert : public OpRewritePattern<mhlo::ConvolutionOp> {
   LogicalResult matchAndRewrite(mhlo::ConvolutionOp op,
                                 PatternRewriter& rewriter) const override {
     Location loc = op.getLoc();
-    Value input = op.lhs();
-    Value filter = op.rhs();
+    Value input = op.getLhs();
+    Value filter = op.getRhs();
 
     auto input_tp = input.getType().dyn_cast<RankedTensorType>();
     auto filter_tp = filter.getType().dyn_cast<RankedTensorType>();
@@ -205,7 +205,7 @@ struct ConvToDynamicConvConvert : public OpRewritePattern<mhlo::ConvolutionOp> {
       return failure();
     }
 
-    auto i64Vec = ConvertDenseIntAttr(op.padding());
+    auto i64Vec = ConvertDenseIntAttr(op.getPadding());
     std::vector<int32_t> paddingValues(i64Vec.size());
     std::transform(
         i64Vec.begin(), i64Vec.end(), paddingValues.begin(),
diff --git a/tao_compiler/mlir/disc/transforms/disc_algebra_simplifier.cc b/tao_compiler/mlir/disc/transforms/disc_algebra_simplifier.cc
index 38ea25356fd..0934785e4b0 100644
--- a/tao_compiler/mlir/disc/transforms/disc_algebra_simplifier.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_algebra_simplifier.cc
@@ -60,10 +60,10 @@ struct AddZeroTensorOp : public OpRewritePattern<mhlo::AddOp> {
 
   LogicalResult matchAndRewrite(mhlo::AddOp op,
                                 PatternRewriter& rewriter) const override {
-    if (allElementsAreSameValue(op.lhs(), 0)) {
-      rewriter.replaceOp(op, op.rhs());
-    } else if (allElementsAreSameValue(op.rhs(), 0)) {
-      rewriter.replaceOp(op, op.lhs());
+    if (allElementsAreSameValue(op.getLhs(), 0)) {
+      rewriter.replaceOp(op, op.getRhs());
+    } else if (allElementsAreSameValue(op.getRhs(), 0)) {
+      rewriter.replaceOp(op, op.getLhs());
     } else {
       return failure();
     }
@@ -77,10 +77,10 @@ struct MulOneTensorOp : public OpRewritePattern<mhlo::MulOp> {
 
   LogicalResult matchAndRewrite(mhlo::MulOp op,
                                 PatternRewriter& rewriter) const override {
-    if (allElementsAreSameValue(op.lhs(), 1)) {
-      rewriter.replaceOp(op, op.rhs());
-    } else if (allElementsAreSameValue(op.rhs(), 1)) {
-      rewriter.replaceOp(op, op.lhs());
+    if (allElementsAreSameValue(op.getLhs(), 1)) {
+      rewriter.replaceOp(op, op.getRhs());
+    } else if (allElementsAreSameValue(op.getRhs(), 1)) {
+      rewriter.replaceOp(op, op.getLhs());
     } else {
       return failure();
     }
@@ -122,20 +122,20 @@ struct ExpandPowOp : public OpRewritePattern<mhlo::PowOp> {
     auto constTy = constOp.getResult().getType().dyn_cast<RankedTensorType>();
     if (!constTy || !constTy.hasStaticShape()) return failure();
 
-    if (!constOp.value().isSplat() && constTy.getNumElements() != 1)
+    if (!constOp.getValue().isSplat() && constTy.getNumElements() != 1)
       return failure();
 
     if (constTy.getElementType().isIntOrIndex()) {
       exponential =
-          (*constOp.value().getValues<APInt>().begin()).getSExtValue();
+          (*constOp.getValue().getValues<APInt>().begin()).getSExtValue();
       return success();
     }
 
     double fpExponential;
     if (constTy.getElementType().isF32()) {
-      fpExponential = *constOp.value().getValues<float>().begin();
+      fpExponential = *constOp.getValue().getValues<float>().begin();
     } else if (constTy.getElementType().isF64()) {
-      fpExponential = *constOp.value().getValues<double>().begin();
+      fpExponential = *constOp.getValue().getValues<double>().begin();
     } else {
       // unsupported float types.
       return failure();
@@ -149,7 +149,7 @@ struct ExpandPowOp : public OpRewritePattern<mhlo::PowOp> {
 
   LogicalResult tryToExtractMultipler(mhlo::PowOp op,
                                       int64_t& exponential) const {
-    Operation* rhsDefiningOp = op.rhs().getDefiningOp();
+    Operation* rhsDefiningOp = op.getRhs().getDefiningOp();
     if (!rhsDefiningOp) return failure();
 
     if (auto constOp = dyn_cast<mhlo::ConstantOp>(rhsDefiningOp)) {
@@ -328,7 +328,8 @@ struct BroadCastInDimOfReshapeOpCanonicalizationPattern
     if (!matched) return failure();
 
     SmallVector<int64_t> newBcastDims;
-    auto oldBcastDims = op.broadcast_dimensions().template getValues<int64_t>();
+    auto oldBcastDims =
+        op.getBroadcastDimensions().template getValues<int64_t>();
     for (size_t d = 0, rank = dimMap.size(); d < rank; ++d)
       newBcastDims.push_back(oldBcastDims[dimMap[d]]);
 
diff --git a/tao_compiler/mlir/disc/transforms/disc_assign_kernel_name.cc b/tao_compiler/mlir/disc/transforms/disc_assign_kernel_name.cc
index d6d523dc800..e0ec6740839 100644
--- a/tao_compiler/mlir/disc/transforms/disc_assign_kernel_name.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_assign_kernel_name.cc
@@ -85,10 +85,10 @@ LogicalResult AssignKernelNamePass::processLaunchFuncOp(LaunchFuncOp op,
   auto func = m.lookupSymbol<GPUFuncOp>(op.getKernelName());
   func.setName(name);
   OpBuilder b(op);
-  auto newOp =
-      b.create<LaunchFuncOp>(op.getLoc(), func, op.getGridSizeOperandValues(),
-                             op.getBlockSizeOperandValues(),
-                             op.dynamicSharedMemorySize(), op.operands());
+  auto newOp = b.create<LaunchFuncOp>(
+      op.getLoc(), func, op.getGridSizeOperandValues(),
+      op.getBlockSizeOperandValues(), op.getDynamicSharedMemorySize(),
+      op.getKernelOperands());
   op.erase();
   return success();
 }
diff --git a/tao_compiler/mlir/disc/transforms/disc_convert_fake_quant_op.cc b/tao_compiler/mlir/disc/transforms/disc_convert_fake_quant_op.cc
index 2565f0ba9b5..f52d4f64f3e 100644
--- a/tao_compiler/mlir/disc/transforms/disc_convert_fake_quant_op.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_convert_fake_quant_op.cc
@@ -51,7 +51,7 @@ template <>
 struct ConvLikeOpPadding<mhlo::DynamicConvOp> {
   static Value get(mhlo::DynamicConvOp convOp, OpBuilder& builder) {
     if (!convOp) return nullptr;
-    return convOp.d_padding();
+    return convOp.getDPadding();
   }
 };
 
@@ -60,7 +60,7 @@ struct ConvLikeOpPadding<mhlo::ConvolutionOp> {
   static Value get(mhlo::ConvolutionOp convOp, OpBuilder& builder) {
     if (!convOp) return nullptr;
     SmallVector<int32_t> paddingValues;
-    for (const auto& val : (*convOp.padding()).getValues<int64_t>()) {
+    for (const auto& val : (*convOp.getPadding()).getValues<int64_t>()) {
       paddingValues.push_back(static_cast<int32_t>(val));
     }
     RankedTensorType ty = RankedTensorType::get({paddingValues.size()},
@@ -77,34 +77,35 @@ struct QuantizedConvLikeOpConverter
 
   LogicalResult matchAndRewrite(mhlo_disc::FakeQuantOp op,
                                 PatternRewriter& rewriter) const override {
-    auto convOp = op.input().template getDefiningOp<OpTy>();
+    auto convOp = op.getInput().template getDefiningOp<OpTy>();
     if (!convOp) return failure();
 
     Value padding = ConvLikeOpPadding<OpTy>::get(convOp, rewriter);
     if (!padding) return failure();
 
     auto inputFakeQuantOp =
-        convOp.lhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
+        convOp.getLhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
     auto weightFakeQuantOp =
-        convOp.rhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
+        convOp.getRhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
     if (!inputFakeQuantOp || !weightFakeQuantOp) return failure();
-    if (inputFakeQuantOp.use_signed() != weightFakeQuantOp.use_signed() ||
-        inputFakeQuantOp.num_bits() != weightFakeQuantOp.num_bits()) {
+    if (inputFakeQuantOp.getUseSigned() != weightFakeQuantOp.getUseSigned() ||
+        inputFakeQuantOp.getNumBits() != weightFakeQuantOp.getNumBits()) {
       return failure();
     }
-    if (inputFakeQuantOp.use_dynamic() != op.use_dynamic() ||
-        weightFakeQuantOp.use_dynamic() != op.use_dynamic())
+    if (inputFakeQuantOp.getUseDynamic() != op.getUseDynamic() ||
+        weightFakeQuantOp.getUseDynamic() != op.getUseDynamic())
       return failure();
 
     Location loc = op.getLoc();
     auto buildQuantizedTensorType = [&](mhlo_disc::FakeQuantOp fakeQuantOp) {
       Type quantizedElemType =
-          fakeQuantOp.use_signed()
-              ? IntegerType::get(rewriter.getContext(), fakeQuantOp.num_bits())
+          fakeQuantOp.getUseSigned()
+              ? IntegerType::get(rewriter.getContext(),
+                                 fakeQuantOp.getNumBits())
               : IntegerType::get(
-                    rewriter.getContext(), fakeQuantOp.num_bits(),
+                    rewriter.getContext(), fakeQuantOp.getNumBits(),
                     mlir::IntegerType::SignednessSemantics::Unsigned);
-      auto inputTy = fakeQuantOp.input().getType().cast<RankedTensorType>();
+      auto inputTy = fakeQuantOp.getInput().getType().cast<RankedTensorType>();
       return RankedTensorType::get(inputTy.getShape(), quantizedElemType,
                                    inputTy.getEncoding());
     };
@@ -115,10 +116,10 @@ struct QuantizedConvLikeOpConverter
         [&](mhlo_disc::FakeQuantOp fakeQuantOp) {
           auto outTy = buildQuantizedTensorType(fakeQuantOp);
           Value quantizedInput = rewriter.create<mhlo_disc::QuantizeOp>(
-              loc, outTy, fakeQuantOp.input(), fakeQuantOp.scale(),
-              fakeQuantOp.zero_point(), fakeQuantOp.use_symmetric(),
-              fakeQuantOp.axis(), fakeQuantOp.quant_min(),
-              fakeQuantOp.quant_max(), fakeQuantOp.use_dynamic());
+              loc, outTy, fakeQuantOp.getInput(), fakeQuantOp.getScale(),
+              fakeQuantOp.getZeroPoint(), fakeQuantOp.getUseSymmetric(),
+              fakeQuantOp.getAxis(), fakeQuantOp.getQuantMin(),
+              fakeQuantOp.getQuantMax(), fakeQuantOp.getUseDynamic());
           return quantizedInput;
         };
 
@@ -126,17 +127,18 @@ struct QuantizedConvLikeOpConverter
     Value quantizedWeight = buildQuantizedOpFromFakeQuantOp(weightFakeQuantOp);
     Value quantizedConv = rewriter.create<mhlo_disc::QuantizedDynamicConvOp>(
         loc, buildQuantizedTensorType(op), quantizedInput, quantizedWeight,
-        padding, inputFakeQuantOp.scale(), inputFakeQuantOp.zero_point(),
-        weightFakeQuantOp.scale(), weightFakeQuantOp.zero_point(), op.scale(),
-        op.zero_point(), *convOp.window_strides(), *convOp.padding(),
-        *convOp.lhs_dilation(), *convOp.rhs_dilation(),
-        *convOp.window_reversal(), convOp.dimension_numbers(),
-        convOp.feature_group_count(), convOp.batch_group_count(),
-        op.use_symmetric(), weightFakeQuantOp.axis(), op.use_dynamic());
+        padding, inputFakeQuantOp.getScale(), inputFakeQuantOp.getZeroPoint(),
+        weightFakeQuantOp.getScale(), weightFakeQuantOp.getZeroPoint(),
+        op.getScale(), op.getZeroPoint(), *convOp.getWindowStrides(),
+        *convOp.getPadding(), *convOp.getLhsDilation(),
+        *convOp.getRhsDilation(), *convOp.getWindowReversal(),
+        convOp.getDimensionNumbers(), convOp.getFeatureGroupCount(),
+        convOp.getBatchGroupCount(), op.getUseSymmetric(),
+        weightFakeQuantOp.getAxis(), op.getUseDynamic());
 
     Value dequantizedOut = rewriter.create<mhlo_disc::DequantizeOp>(
-        loc, op.getType(), quantizedConv, op.scale(), op.zero_point(),
-        op.use_symmetric(), op.axis(), op.use_dynamic());
+        loc, op.getType(), quantizedConv, op.getScale(), op.getZeroPoint(),
+        op.getUseSymmetric(), op.getAxis(), op.getUseDynamic());
 
     rewriter.replaceOp(op, dequantizedOut);
     return success();
@@ -154,12 +156,12 @@ Value buildQuantizedDotGeneralOp(mhlo::DotOp dotOp, PatternRewriter& rewriter,
   auto dotDimensionAttr = mhlo::DotDimensionNumbersAttr::get(
       rewriter.getContext(), {}, {}, lhsContractingDims, rhsContractingDims);
   return rewriter.create<mhlo_disc::QuantizedDotGeneralOp>(
-      loc, outTy, input, weight, inputFakeQuantOp.scale(),
-      inputFakeQuantOp.zero_point(), weightFakeQuantOp.scale(),
-      weightFakeQuantOp.zero_point(), resultFakeQuantOp.scale(),
-      resultFakeQuantOp.zero_point(), dotDimensionAttr,
-      resultFakeQuantOp.use_symmetric(), weightFakeQuantOp.axis(),
-      resultFakeQuantOp.use_dynamic());
+      loc, outTy, input, weight, inputFakeQuantOp.getScale(),
+      inputFakeQuantOp.getZeroPoint(), weightFakeQuantOp.getScale(),
+      weightFakeQuantOp.getZeroPoint(), resultFakeQuantOp.getScale(),
+      resultFakeQuantOp.getZeroPoint(), dotDimensionAttr,
+      resultFakeQuantOp.getUseSymmetric(), weightFakeQuantOp.getAxis(),
+      resultFakeQuantOp.getUseDynamic());
 }
 
 Value buildQuantizedDotGeneralOp(mhlo::DotGeneralOp dotOp,
@@ -169,12 +171,12 @@ Value buildQuantizedDotGeneralOp(mhlo::DotGeneralOp dotOp,
                                  mhlo_disc::FakeQuantOp weightFakeQuantOp,
                                  mhlo_disc::FakeQuantOp resultFakeQuantOp) {
   return rewriter.create<mhlo_disc::QuantizedDotGeneralOp>(
-      loc, outTy, input, weight, inputFakeQuantOp.scale(),
-      inputFakeQuantOp.zero_point(), weightFakeQuantOp.scale(),
-      weightFakeQuantOp.zero_point(), resultFakeQuantOp.scale(),
-      resultFakeQuantOp.zero_point(), dotOp.dot_dimension_numbers(),
-      resultFakeQuantOp.use_symmetric(), weightFakeQuantOp.axis(),
-      resultFakeQuantOp.use_dynamic());
+      loc, outTy, input, weight, inputFakeQuantOp.getScale(),
+      inputFakeQuantOp.getZeroPoint(), weightFakeQuantOp.getScale(),
+      weightFakeQuantOp.getZeroPoint(), resultFakeQuantOp.getScale(),
+      resultFakeQuantOp.getZeroPoint(), dotOp.getDotDimensionNumbers(),
+      resultFakeQuantOp.getUseSymmetric(), weightFakeQuantOp.getAxis(),
+      resultFakeQuantOp.getUseDynamic());
 }
 
 template <typename OpTy>
@@ -184,14 +186,14 @@ struct QuantizedDotLikeOpConverter
 
   LogicalResult matchAndRewrite(mhlo_disc::FakeQuantOp op,
                                 PatternRewriter& rewriter) const override {
-    auto dotOp = op.input().template getDefiningOp<OpTy>();
+    auto dotOp = op.getInput().template getDefiningOp<OpTy>();
     if (!dotOp) return failure();
 
     if (isa<mhlo::DotOp>(dotOp.getOperation())) {
       auto inputTy =
-          dotOp.lhs().getType().template dyn_cast<RankedTensorType>();
+          dotOp.getLhs().getType().template dyn_cast<RankedTensorType>();
       auto weightTy =
-          dotOp.rhs().getType().template dyn_cast<RankedTensorType>();
+          dotOp.getRhs().getType().template dyn_cast<RankedTensorType>();
       if (!inputTy || !weightTy || inputTy.getRank() != weightTy.getRank() ||
           inputTy.getRank() != 2) {
         return rewriter.notifyMatchFailure(
@@ -200,27 +202,28 @@ struct QuantizedDotLikeOpConverter
     }
 
     auto inputFakeQuantOp =
-        dotOp.lhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
+        dotOp.getLhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
     auto weightFakeQuantOp =
-        dotOp.rhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
+        dotOp.getRhs().template getDefiningOp<mhlo_disc::FakeQuantOp>();
     if (!inputFakeQuantOp || !weightFakeQuantOp) return failure();
-    if (inputFakeQuantOp.use_signed() != weightFakeQuantOp.use_signed() ||
-        inputFakeQuantOp.num_bits() != weightFakeQuantOp.num_bits()) {
+    if (inputFakeQuantOp.getUseSigned() != weightFakeQuantOp.getUseSigned() ||
+        inputFakeQuantOp.getNumBits() != weightFakeQuantOp.getNumBits()) {
       return failure();
     }
-    if (inputFakeQuantOp.use_dynamic() != op.use_dynamic() ||
-        weightFakeQuantOp.use_dynamic() != op.use_dynamic())
+    if (inputFakeQuantOp.getUseDynamic() != op.getUseDynamic() ||
+        weightFakeQuantOp.getUseDynamic() != op.getUseDynamic())
       return failure();
 
     Location loc = op.getLoc();
     auto buildQuantizedTensorType = [&](mhlo_disc::FakeQuantOp fakeQuantOp) {
       Type quantizedElemType =
-          fakeQuantOp.use_signed()
-              ? IntegerType::get(rewriter.getContext(), fakeQuantOp.num_bits())
+          fakeQuantOp.getUseSigned()
+              ? IntegerType::get(rewriter.getContext(),
+                                 fakeQuantOp.getNumBits())
               : IntegerType::get(
-                    rewriter.getContext(), fakeQuantOp.num_bits(),
+                    rewriter.getContext(), fakeQuantOp.getNumBits(),
                     mlir::IntegerType::SignednessSemantics::Unsigned);
-      auto inputTy = fakeQuantOp.input().getType().cast<RankedTensorType>();
+      auto inputTy = fakeQuantOp.getInput().getType().cast<RankedTensorType>();
       return RankedTensorType::get(inputTy.getShape(), quantizedElemType,
                                    inputTy.getEncoding());
     };
@@ -231,11 +234,11 @@ struct QuantizedDotLikeOpConverter
         [&](mhlo_disc::FakeQuantOp fakeQuantOp) {
           auto outTy = buildQuantizedTensorType(fakeQuantOp);
           Value quantizedInput = rewriter.create<mhlo_disc::QuantizeOp>(
-              loc, outTy, fakeQuantOp.input(), fakeQuantOp.scale(),
-              fakeQuantOp.zero_point(), fakeQuantOp.use_symmetric(),
-              fakeQuantOp.axis(), fakeQuantOp.quant_min(),
-              fakeQuantOp.quant_max(), fakeQuantOp.use_dynamic(),
-              fakeQuantOp.round_mode());
+              loc, outTy, fakeQuantOp.getInput(), fakeQuantOp.getScale(),
+              fakeQuantOp.getZeroPoint(), fakeQuantOp.getUseSymmetric(),
+              fakeQuantOp.getAxis(), fakeQuantOp.getQuantMin(),
+              fakeQuantOp.getQuantMax(), fakeQuantOp.getUseDynamic(),
+              fakeQuantOp.getRoundMode());
           return quantizedInput;
         };
 
@@ -245,8 +248,9 @@ struct QuantizedDotLikeOpConverter
         dotOp, rewriter, loc, buildQuantizedTensorType(op), quantizedInput,
         quantizedWeight, inputFakeQuantOp, weightFakeQuantOp, op);
     Value dequantizedOut = rewriter.create<mhlo_disc::DequantizeOp>(
-        loc, op.getType(), quantizedDot, op.scale(), op.zero_point(),
-        op.use_symmetric(), op.axis(), op.use_dynamic(), op.round_mode());
+        loc, op.getType(), quantizedDot, op.getScale(), op.getZeroPoint(),
+        op.getUseSymmetric(), op.getAxis(), op.getUseDynamic(),
+        op.getRoundMode());
 
     rewriter.replaceOp(op, dequantizedOut);
     return success();
@@ -270,7 +274,7 @@ struct FakeQuantOpToIdentityConverter
 
   LogicalResult matchAndRewrite(mhlo_disc::FakeQuantOp op,
                                 PatternRewriter& rewriter) const override {
-    rewriter.replaceOp(op, op.input());
+    rewriter.replaceOp(op, op.getInput());
     return success();
   }
 };
diff --git a/tao_compiler/mlir/disc/transforms/disc_cpu_map_parallel_loop.cc b/tao_compiler/mlir/disc/transforms/disc_cpu_map_parallel_loop.cc
index e0e32c814b0..8d93921a80f 100644
--- a/tao_compiler/mlir/disc/transforms/disc_cpu_map_parallel_loop.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_cpu_map_parallel_loop.cc
@@ -18,7 +18,7 @@ limitations under the License.
 
 #include "llvm/Support/Debug.h"
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BlockAndValueMapping.h"
diff --git a/tao_compiler/mlir/disc/transforms/disc_dense_to_sparse.cc b/tao_compiler/mlir/disc/transforms/disc_dense_to_sparse.cc
index 800104db313..7dc88e0eb97 100644
--- a/tao_compiler/mlir/disc/transforms/disc_dense_to_sparse.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_dense_to_sparse.cc
@@ -208,15 +208,16 @@ struct ExpandDotGeneralOp : public OpRewritePattern<mhlo::DotGeneralOp> {
     if (enable_sparse_convert_) {
       Operation* weightOp;
       if (sparse_weight_pos == 0) {
-        weightOp = op.lhs().getDefiningOp();
+        weightOp = op.getLhs().getDefiningOp();
       } else {
-        weightOp = op.rhs().getDefiningOp();
+        weightOp = op.getRhs().getDefiningOp();
       }
       auto weightConstOp = dyn_cast<mhlo::ConstantOp>(weightOp);
       ArrayRef<int64_t> weightShape =
           weightConstOp.getType().cast<ShapedType>().getShape();
 
-      auto weights_value = weightConstOp.value().getValues<APFloat>().begin();
+      auto weights_value =
+          weightConstOp.getValue().getValues<APFloat>().begin();
       SmallVector<APFloat> compressed_weight(
           weightShape[0] * weightShape[1] / 2, APFloat(0.f));
       std::vector<uint16_t> compressed_index(
@@ -277,7 +278,7 @@ struct ExpandDotGeneralOp : public OpRewritePattern<mhlo::DotGeneralOp> {
         RankedTensorType::get(spgemm_shape, outputTy.getElementType());
 
     // insert spgemm op
-    auto dot_dimension_numbers = op.dot_dimension_numbers();
+    auto dot_dimension_numbers = op.getDotDimensionNumbers();
     auto lhs_contracting_dimensions =
         dot_dimension_numbers.getLhsContractingDimensions();
     auto rhs_contracting_dimensions =
@@ -329,7 +330,7 @@ struct ExpandDotGeneralOp : public OpRewritePattern<mhlo::DotGeneralOp> {
 
     if (weight_shape.size() != 2) return failure();
 
-    auto weight_value = constOp.value().getValues<APFloat>().begin();
+    auto weight_value = constOp.getValue().getValues<APFloat>().begin();
 
     int64_t row = is_transpose ? weight_shape[1] : weight_shape[0];
     int64_t col = is_transpose ? weight_shape[0] : weight_shape[1];
@@ -361,7 +362,7 @@ struct ExpandDotGeneralOp : public OpRewritePattern<mhlo::DotGeneralOp> {
   }
 
   int tryToFindSparseWeight(mhlo::DotGeneralOp op) const {
-    Operation* rhsDefiningOp = op.rhs().getDefiningOp();
+    Operation* rhsDefiningOp = op.getRhs().getDefiningOp();
     if (!rhsDefiningOp) return -1;
 
     if (auto constOp = dyn_cast<mhlo::ConstantOp>(rhsDefiningOp)) {
@@ -370,7 +371,7 @@ struct ExpandDotGeneralOp : public OpRewritePattern<mhlo::DotGeneralOp> {
       }
     }
 
-    Operation* lhsDefiningOp = op.lhs().getDefiningOp();
+    Operation* lhsDefiningOp = op.getLhs().getDefiningOp();
     if (!lhsDefiningOp) return -1;
 
     if (auto constOp = dyn_cast<mhlo::ConstantOp>(lhsDefiningOp)) {
@@ -444,7 +445,7 @@ struct ExpandSparseGemmOp : public OpRewritePattern<mhlo_disc::CustomCallOp> {
 
     if (!inputTy || !outputTy) return success();
 
-    auto config = op.backend_config().cast<DictionaryAttr>();
+    auto config = op.getBackendConfig().cast<DictionaryAttr>();
     int64_t lhs_contracting_dimensions =
         config.getAs<IntegerAttr>("lhs_contracting_dimensions").getInt();
     int64_t rhs_contracting_dimensions =
@@ -485,7 +486,7 @@ struct ExpandSparseGemmOp : public OpRewritePattern<mhlo_disc::CustomCallOp> {
   }
 
   LogicalResult tryToFindSparseGemm(mhlo_disc::CustomCallOp op) const {
-    auto call_target_name = op.call_target_name();
+    auto call_target_name = op.getCallTargetName();
     if (call_target_name == "sparse_gemm") {
       Value lhs = op->getOperand(0);
       Value rhs = op->getOperand(1);
diff --git a/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc b/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc
index 7909a43df3b..1d33c51f664 100644
--- a/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_dot_merge.cc
@@ -13,7 +13,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -240,8 +240,8 @@ bool DotShareOperandMergeConverter::buildShareOperandMap(
   block->walk([&](mhlo::DotGeneralOp op) {
     // get one-side operand shareinfo according to the share_type
     DotShareInfo share_info;
-    share_info.share_operand = (share_type == LEFT) ? op.lhs() : op.rhs();
-    share_info.dimension_numbers = op.dot_dimension_numbers();
+    share_info.share_operand = (share_type == LEFT) ? op.getLhs() : op.getRhs();
+    share_info.dimension_numbers = op.getDotDimensionNumbers();
     auto& shared_op_list = share_operand_map[std::move(share_info)];
     shared_op_list.push_back(op);
   });
@@ -274,11 +274,11 @@ bool DotShareOperandMergeConverter::applyMerging(DotCluster& cluster,
 
   OpBuilder builder(foremost_dot);
   auto non_concat_op =
-      (share_type == LEFT) ? foremost_dot.lhs() : foremost_dot.rhs();
+      (share_type == LEFT) ? foremost_dot.getLhs() : foremost_dot.getRhs();
   auto orig_lhs_type =
-      foremost_dot.lhs().getType().dyn_cast<RankedTensorType>();
+      foremost_dot.getLhs().getType().dyn_cast<RankedTensorType>();
   auto orig_rhs_type =
-      foremost_dot.rhs().getType().dyn_cast<RankedTensorType>();
+      foremost_dot.getRhs().getType().dyn_cast<RankedTensorType>();
   auto orig_concat_op_type =
       (share_type == LEFT) ? orig_rhs_type : orig_lhs_type;
   auto orig_non_concat_op_type =
@@ -292,7 +292,7 @@ bool DotShareOperandMergeConverter::applyMerging(DotCluster& cluster,
   // Find concat dim.
   int64_t concat_dim = -1;
   DenseSet<int64_t> non_concat_dims;
-  auto dim_numbers = foremost_dot.dot_dimension_numbers();
+  auto dim_numbers = foremost_dot.getDotDimensionNumbers();
   // Batching and contracting dims of the to-concat operands.
   const auto& batch_dims = (share_type == LEFT)
                                ? dim_numbers.getRhsBatchingDimensions()
@@ -318,7 +318,7 @@ bool DotShareOperandMergeConverter::applyMerging(DotCluster& cluster,
   int64_t concat_dim_sum = 0;
   for (auto op : ops) {
     mhlo::DotGeneralOp dot = dyn_cast<mhlo::DotGeneralOp>(op);
-    auto to_concat = (share_type == LEFT) ? dot.rhs() : dot.lhs();
+    auto to_concat = (share_type == LEFT) ? dot.getRhs() : dot.getLhs();
     auto concat_op_type = to_concat.getType().dyn_cast<RankedTensorType>();
     auto concat_dim_size = concat_op_type.getDimSize(concat_dim);
     if (concat_dim_size == ShapedType::kDynamicSize) {
@@ -411,8 +411,9 @@ bool DotShareOperandMergeConverter::applyMerging(DotCluster& cluster,
       // Note that we do not use op but op's rhs and lhs to build DimOp.
       // Because op will be replaced and if use op to build DimOp, it will
       // form circles.
-      auto concatenate_op = (share_type == LEFT) ? op.rhs() : op.lhs();
-      auto non_concatenate_op = (share_type == LEFT) ? op.lhs() : op.rhs();
+      auto concatenate_op = (share_type == LEFT) ? op.getRhs() : op.getLhs();
+      auto non_concatenate_op =
+          (share_type == LEFT) ? op.getLhs() : op.getRhs();
       auto orig_dot_type = op.getType().dyn_cast<RankedTensorType>();
 
       SmallVector<Value, 4> start_values(rank, zero);
@@ -581,10 +582,10 @@ bool DotBatchMergeConverter::buildMergingShapeMap(
     MergingShapeEqualMap& equal_merge_shape_map) {
   block->walk([&](mhlo::DotGeneralOp op) {
     MergingShape dot_shape;
-    Value lhs = op.lhs();
-    Value rhs = op.rhs();
+    Value lhs = op.getLhs();
+    Value rhs = op.getRhs();
     // Initialize `dimension_numbers`.
-    dot_shape.dimension_numbers = op.dot_dimension_numbers();
+    dot_shape.dimension_numbers = op.getDotDimensionNumbers();
     auto lhs_batch_dims =
         dot_shape.dimension_numbers.getLhsBatchingDimensions();
     auto rhs_batch_dims =
@@ -663,15 +664,15 @@ bool DotBatchMergeConverter::applyMerging(DotCluster& cluster) {
   // We use the foremost dot to create the builder. Thus we only need to reorder
   // the operands of some newly created ops, rather users of them.
   OpBuilder builder(last_dot);
-  auto orig_lhs_type = last_dot.lhs().getType().dyn_cast<RankedTensorType>();
-  auto orig_rhs_type = last_dot.rhs().getType().dyn_cast<RankedTensorType>();
+  auto orig_lhs_type = last_dot.getLhs().getType().dyn_cast<RankedTensorType>();
+  auto orig_rhs_type = last_dot.getRhs().getType().dyn_cast<RankedTensorType>();
   auto orig_dot_type = last_dot.getType().dyn_cast<RankedTensorType>();
   SmallVector<Value, 4> lhs_operands;
   SmallVector<Value, 4> rhs_operands;
   for (auto op : ops) {
     mhlo::DotGeneralOp dot = dyn_cast<mhlo::DotGeneralOp>(op);
-    auto lhs_expand = expandDim0(builder, loc, dot.lhs());
-    auto rhs_expand = expandDim0(builder, loc, dot.rhs());
+    auto lhs_expand = expandDim0(builder, loc, dot.getLhs());
+    auto rhs_expand = expandDim0(builder, loc, dot.getRhs());
     if (!lhs_expand || !rhs_expand) {
       LLVM_DEBUG(llvm::dbgs() << "Failed to expand dim for dot merge.\n");
       return false;
@@ -712,7 +713,7 @@ bool DotBatchMergeConverter::applyMerging(DotCluster& cluster) {
   auto result_type =
       RankedTensorType::get(result_shapes, orig_dot_type.getElementType());
   // Build dot dimension numbers.
-  auto dim_numbers = last_dot.dot_dimension_numbers();
+  auto dim_numbers = last_dot.getDotDimensionNumbers();
 
   SmallVector<int64_t> lhs_batching_dims;
   auto lhs_batch = dim_numbers.getLhsBatchingDimensions();
diff --git a/tao_compiler/mlir/disc/transforms/disc_flatten_memref_access.cc b/tao_compiler/mlir/disc/transforms/disc_flatten_memref_access.cc
index becd58ff974..578fab740c3 100644
--- a/tao_compiler/mlir/disc/transforms/disc_flatten_memref_access.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_flatten_memref_access.cc
@@ -91,7 +91,7 @@ LogicalResult DiscFlattenMemrefAccessPass::processStoreOp(memref::StoreOp op) {
   SmallVector<Value> dimSizes = disc_ral::getShapeValues(&b, memref);
   Value linear = b.create<disc_shape::LinearizeOp>(loc, b.getIndexType(),
                                                    op.getIndices(), dimSizes);
-  disc_ral::createOffsetStore(b, loc, op.value(), memref, linear);
+  disc_ral::createOffsetStore(b, loc, op.getValue(), memref, linear);
   op->erase();
   return success();
 }
diff --git a/tao_compiler/mlir/disc/transforms/disc_gpu_kernel_to_blob.cc b/tao_compiler/mlir/disc/transforms/disc_gpu_kernel_to_blob.cc
index c046507044a..085bda8200e 100644
--- a/tao_compiler/mlir/disc/transforms/disc_gpu_kernel_to_blob.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_gpu_kernel_to_blob.cc
@@ -36,11 +36,11 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_type_conversion_util.h"
 #include "tensorflow/compiler/xla/status.h"
 #include "tensorflow/compiler/xla/statusor.h"
-#include "tensorflow/core/platform/cuda_libdevice_path.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/path.h"
 #include "tensorflow/core/platform/random.h"
 #include "tensorflow/core/util/env_var.h"
+#include "tensorflow/tsl/platform/cuda_libdevice_path.h"
 
 #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
 #include "tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.h"
@@ -122,7 +122,7 @@ class GpuKernelToBlobPass
           gpu_module.emitError(blob_or.status().error_message());
           return signalPassFailure();
         }
-        const auto& blob = blob_or.ValueOrDie();
+        const auto& blob = blob_or.value();
         std::string blob_string(blob.begin(), blob.end());
         std::string attr_str = std::string(kGpuBinaryAttrName) + "_" + name;
         gpu_module->setAttr(attr_str,
@@ -135,7 +135,7 @@ class GpuKernelToBlobPass
         gpu_module.emitError(blob_or.status().error_message());
         return signalPassFailure();
       }
-      const auto& blob = blob_or.ValueOrDie();
+      const auto& blob = blob_or.value();
       std::string blob_string(blob.begin(), blob.end());
       gpu_module->setAttr(blob_annotation_,
                           mlir::StringAttr::get(&getContext(), blob_string));
@@ -435,12 +435,12 @@ class GpuKernelToBlobPass
  private:
   xla::StatusOr<std::string> GetLibdeviceDir(
       const xla::HloModuleConfig& hlo_module_config) {
-    for (const std::string& cuda_root : tensorflow::CandidateCudaRoots(
+    for (const std::string& cuda_root : tsl::CandidateCudaRoots(
              hlo_module_config.debug_options().xla_gpu_cuda_data_dir())) {
       std::string libdevice_dir =
           tensorflow::io::JoinPath(cuda_root, "nvvm", "libdevice");
       VLOG(2) << "Looking for libdevice at " << libdevice_dir;
-      if (tensorflow::Env::Default()->IsDirectory(libdevice_dir).ok()) {
+      if (tsl::Env::Default()->IsDirectory(libdevice_dir).ok()) {
         VLOG(2) << "Found libdevice dir " << libdevice_dir;
         return libdevice_dir;
       }
diff --git a/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc b/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
index b83f023b101..b15ac73c09b 100644
--- a/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_hlo_legalize_to_lhlo.cc
@@ -19,7 +19,7 @@ limitations under the License.
 #include <algorithm>
 #include <utility>
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Func/Transforms/FuncConversions.h"
@@ -226,10 +226,10 @@ struct DiscHloLegalizeToLhlo
     auto& context = getContext();
     RewritePatternSet patterns(&context);
     ConversionTarget target(context);
-    target.addLegalDialect<
-        arith::ArithmeticDialect, lmhlo_disc::LmhloDiscDialect,
-        bufferization::BufferizationDialect, memref::MemRefDialect,
-        shape::ShapeDialect, tensor::TensorDialect>();
+    target.addLegalDialect<arith::ArithDialect, lmhlo_disc::LmhloDiscDialect,
+                           bufferization::BufferizationDialect,
+                           memref::MemRefDialect, shape::ShapeDialect,
+                           tensor::TensorDialect>();
     target.addIllegalDialect<mhlo_disc::MhloDiscDialect>();
     target.addIllegalOp<disc_shape::TieShapeOp>();
 
diff --git a/tao_compiler/mlir/disc/transforms/disc_hlo_to_std.cc b/tao_compiler/mlir/disc/transforms/disc_hlo_to_std.cc
index 3e1b6558385..a07a4fe9c59 100644
--- a/tao_compiler/mlir/disc/transforms/disc_hlo_to_std.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_hlo_to_std.cc
@@ -14,7 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Attributes.h"
@@ -134,7 +134,7 @@ void ConvertHloToStandardPass::runOnOperation() {
   // Setup target legality.
   MLIRContext& ctx = getContext();
   ConversionTarget target(ctx);
-  target.addLegalDialect<tensor::TensorDialect, arith::ArithmeticDialect>();
+  target.addLegalDialect<tensor::TensorDialect, arith::ArithDialect>();
   target.addIllegalOp<ComputeReshapeShapeOp>();
 
   // Setup conversion patterns.
diff --git a/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_common.cc b/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_common.cc
index e43a985870d..12325c6c4b9 100644
--- a/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_common.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_common.cc
@@ -70,9 +70,9 @@ LogicalResult GenericAtomicRMWOpLoweringWithBitcast::matchAndRewrite(
 
   // Compute the loaded value and branch to the loop block.
   rewriter.setInsertionPointToEnd(initBlock);
-  auto memRefType = atomicOp.memref().getType().cast<MemRefType>();
-  auto dataPtr = getStridedElementPtr(loc, memRefType, adaptor.memref(),
-                                      adaptor.indices(), rewriter);
+  auto memRefType = atomicOp.getMemref().getType().cast<MemRefType>();
+  auto dataPtr = getStridedElementPtr(loc, memRefType, adaptor.getMemref(),
+                                      adaptor.getIndices(), rewriter);
   Value init = rewriter.create<LLVM::LoadOp>(loc, dataPtr);
   rewriter.create<LLVM::BrOp>(loc, init, loopBlock);
 
@@ -83,7 +83,7 @@ LogicalResult GenericAtomicRMWOpLoweringWithBitcast::matchAndRewrite(
   auto loopArgument = loopBlock->getArgument(0);
   BlockAndValueMapping mapping;
   mapping.map(atomicOp.getCurrentValue(), loopArgument);
-  Block& entryBlock = atomicOp.atomic_body().front();
+  Block& entryBlock = atomicOp.getAtomicBody().front();
   for (auto& nestedOp : entryBlock.without_terminator()) {
     Operation* clone = rewriter.clone(nestedOp, mapping);
     mapping.map(nestedOp.getResults(), clone->getResults());
diff --git a/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_to_nvvm_ops.cc b/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_to_nvvm_ops.cc
index 203d4384bc9..3a470597612 100644
--- a/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_to_nvvm_ops.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_to_nvvm_ops.cc
@@ -23,7 +23,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/FormatVariadic.h"
-#include "mlir/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.h"
+#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
@@ -32,7 +32,7 @@
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
@@ -46,7 +46,6 @@
 #include "mlir/lib/Conversion/GPUCommon/GPUOpsLowering.h"
 #include "mlir/lib/Conversion/GPUCommon/IndexIntrinsicsOpLowering.h"
 #include "mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h"
-#include "mlir/lib/Conversion/PassDetail.h"
 #include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
 #include "tensorflow/compiler/mlir/disc/transforms/disc_lower_gpu_ops_common.h"
 
@@ -168,8 +167,7 @@ struct DiscLowerGpuOpsToNVVMOpsPass
     llvmPatterns.add<GenericAtomicRMWOpLoweringWithBitcast>(
         converter, /* PatternBenefit */ 3);
     llvmPatterns.add<RemoveUselessUnrealizedConversionCastOp>(converter);
-    mlir::arith::populateArithmeticToLLVMConversionPatterns(converter,
-                                                            llvmPatterns);
+    mlir::arith::populateArithToLLVMConversionPatterns(converter, llvmPatterns);
     populateMathToLLVMConversionPatterns(converter, patterns);
     populateMemRefToLLVMConversionPatterns(converter, llvmPatterns);
     populateFuncToLLVMConversionPatterns(converter, patterns);
@@ -181,7 +179,7 @@ struct DiscLowerGpuOpsToNVVMOpsPass
     LLVMConversionTarget target(getContext());
     configureGpuToNVVMConversionLegality(target);
     target.addLegalDialect<LLVM::LLVMDialect>();
-    target.addIllegalDialect<arith::ArithmeticDialect, math::MathDialect,
+    target.addIllegalDialect<arith::ArithDialect, math::MathDialect,
                              cf::ControlFlowDialect>();
     target.addIllegalOp<UnrealizedConversionCastOp>();
     if (failed(applyPartialConversion(m, target, llvmFrozenPatterns)))
diff --git a/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_to_rocdl_ops.cc b/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_to_rocdl_ops.cc
index e32ae026d06..e98cf74cea0 100644
--- a/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_to_rocdl_ops.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_lower_gpu_ops_to_rocdl_ops.cc
@@ -23,7 +23,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/FormatVariadic.h"
-#include "mlir/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.h"
+#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
@@ -46,7 +46,6 @@
 #include "mlir/lib/Conversion/GPUCommon/GPUOpsLowering.h"
 #include "mlir/lib/Conversion/GPUCommon/IndexIntrinsicsOpLowering.h"
 #include "mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h"
-#include "mlir/lib/Conversion/PassDetail.h"
 #include "tensorflow/compiler/mlir/disc/transforms/PassDetail.h"
 #include "tensorflow/compiler/mlir/disc/transforms/disc_lower_gpu_ops_common.h"
 
@@ -74,7 +73,7 @@ struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
       gpu::ShuffleOp op, OpAdaptor adaptor,
       ConversionPatternRewriter& rewriter) const override {
     Location loc = op->getLoc();
-    auto valueTy = adaptor.value().getType();
+    auto valueTy = adaptor.getValue().getType();
     auto int32Type = IntegerType::get(rewriter.getContext(), 32);
     auto predTy = IntegerType::get(rewriter.getContext(), 1);
     Value yes = rewriter.create<LLVM::ConstantOp>(
@@ -84,14 +83,14 @@ struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
     // Bit mask of active lanes: `(1 << activeWidth) - 1`.
     Value activeMask = rewriter.create<LLVM::SubOp>(
         loc, int32Type,
-        rewriter.create<LLVM::ShlOp>(loc, int32Type, one, adaptor.width()),
+        rewriter.create<LLVM::ShlOp>(loc, int32Type, one, adaptor.getWidth()),
         one);
     // Clamp lane: `activeWidth - 1`
     Value maskAndClamp =
-        rewriter.create<LLVM::SubOp>(loc, int32Type, adaptor.width(), one);
+        rewriter.create<LLVM::SubOp>(loc, int32Type, adaptor.getWidth(), one);
 
     Value shfl = rewriter.create<ROCDL::ShflBflyOp>(
-        loc, valueTy, activeMask, adaptor.value(), adaptor.offset(),
+        loc, valueTy, activeMask, adaptor.getValue(), adaptor.getOffset(),
         maskAndClamp);
 
     rewriter.replaceOp(op, {shfl, yes});
@@ -232,8 +231,7 @@ struct DiscLowerGpuOpsToROCDLOpsPass
     llvmPatterns.add<GenericAtomicRMWOpLoweringWithBitcast>(
         converter, /* PatternBenefit */ 3);
     llvmPatterns.add<RemoveUselessUnrealizedConversionCastOp>(converter);
-    mlir::arith::populateArithmeticToLLVMConversionPatterns(converter,
-                                                            llvmPatterns);
+    mlir::arith::populateArithToLLVMConversionPatterns(converter, llvmPatterns);
     populateVectorToLLVMConversionPatterns(converter, llvmPatterns);
     // Remove the following since it disappears from LLVM.
     // populateVectorToROCDLConversionPatterns(converter, llvmPatterns);
diff --git a/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc b/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc
index ab90380ffa6..af7dfd4044c 100644
--- a/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_lower_quantize_and_dequantize.cc
@@ -38,36 +38,36 @@ struct QuantizeOpConverter : public OpRewritePattern<mhlo_disc::QuantizeOp> {
 
   LogicalResult matchAndRewrite(mhlo_disc::QuantizeOp op,
                                 PatternRewriter& rewriter) const override {
-    if (op.use_dynamic())
+    if (op.getUseDynamic())
       return rewriter.notifyMatchFailure(op,
                                          "Not support dynamic quantize a.t.m.");
 
-    auto inputTy = op.input().getType().cast<RankedTensorType>();
+    auto inputTy = op.getInput().getType().cast<RankedTensorType>();
     if (!inputTy.getElementType().isF32())
       return rewriter.notifyMatchFailure(
           op, "Only support quantize f32 input a.t.m.");
 
     Location loc = op.getLoc();
-    Value inputShape = rewriter.create<shape::ShapeOfOp>(loc, op.input());
+    Value inputShape = rewriter.create<shape::ShapeOfOp>(loc, op.getInput());
     Value bcastedScale = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
-        loc, inputTy, op.scale(), inputShape, op.axis());
-    auto zeroPointTy = op.zero_point().getType().cast<RankedTensorType>();
+        loc, inputTy, op.getScale(), inputShape, op.getAxis());
+    auto zeroPointTy = op.getZeroPoint().getType().cast<RankedTensorType>();
     auto castedZeroPointTy =
         RankedTensorType::get(zeroPointTy.getShape(), inputTy.getElementType(),
                               zeroPointTy.getEncoding());
     Value castedZeroPoint = rewriter.create<mhlo::ConvertOp>(
-        loc, castedZeroPointTy, op.zero_point());
+        loc, castedZeroPointTy, op.getZeroPoint());
     Value bcastedZeroPoint = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
-        loc, inputTy, castedZeroPoint, inputShape, op.axis());
+        loc, inputTy, castedZeroPoint, inputShape, op.getAxis());
     auto quantMinMaxTy = RankedTensorType::get({}, inputTy.getElementType());
     Value quantMin = rewriter.create<mhlo::ConstantOp>(
         loc, DenseElementsAttr::get(
                  quantMinMaxTy,
-                 static_cast<float>(static_cast<int64_t>(op.quant_min()))));
+                 static_cast<float>(static_cast<int64_t>(op.getQuantMin()))));
     Value quantMax = rewriter.create<mhlo::ConstantOp>(
         loc, DenseElementsAttr::get(
                  quantMinMaxTy,
-                 static_cast<float>(static_cast<int64_t>(op.quant_max()))));
+                 static_cast<float>(static_cast<int64_t>(op.getQuantMax()))));
     // quantMin/Max should always be scalar.
     Value bcastedQuantMin = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
         loc, inputTy, quantMin, inputShape, rewriter.getI64TensorAttr({}));
@@ -76,12 +76,12 @@ struct QuantizeOpConverter : public OpRewritePattern<mhlo_disc::QuantizeOp> {
 
     // output = clip(round(\frac{input}{scale} + zero\_point), quant\_min,
     // quant\_max)
-    Value t0 = rewriter.create<mhlo::DivOp>(loc, op.input(), bcastedScale);
+    Value t0 = rewriter.create<mhlo::DivOp>(loc, op.getInput(), bcastedScale);
     Value t1 = rewriter.create<mhlo::AddOp>(loc, t0, bcastedZeroPoint);
     Value t2;
-    if (op.round_mode() == mlir::mhlo_disc::RoundModeEnum::RoundHalfToEven) {
+    if (op.getRoundMode() == mlir::mhlo_disc::RoundModeEnum::RoundHalfToEven) {
       t2 = rewriter.create<mhlo::RoundNearestEvenOp>(loc, t1);
-    } else if (op.round_mode() ==
+    } else if (op.getRoundMode() ==
                mlir::mhlo_disc::RoundModeEnum::RoundHalfAwayFromZero) {
       t2 = rewriter.create<mhlo::RoundOp>(loc, t1);
     } else {
@@ -102,30 +102,31 @@ struct DequantizeOpConverter
 
   LogicalResult matchAndRewrite(mhlo_disc::DequantizeOp op,
                                 PatternRewriter& rewriter) const override {
-    if (op.use_dynamic())
+    if (op.getUseDynamic())
       return rewriter.notifyMatchFailure(
           op, "Not support dynamic dequantize a.t.m.");
 
     auto outTy = op.getType().cast<RankedTensorType>();
-    auto scaleTy = op.scale().getType().cast<RankedTensorType>();
+    auto scaleTy = op.getScale().getType().cast<RankedTensorType>();
     if (outTy.getElementType() != scaleTy.getElementType() ||
         !outTy.getElementType().isF32())
       return rewriter.notifyMatchFailure(
           op, "Only support dequantize to f32 a.t.m.");
 
     Location loc = op.getLoc();
-    auto inputTy = op.input().getType().cast<RankedTensorType>();
-    Value inputShape = rewriter.create<shape::ShapeOfOp>(loc, op.input());
+    auto inputTy = op.getInput().getType().cast<RankedTensorType>();
+    Value inputShape = rewriter.create<shape::ShapeOfOp>(loc, op.getInput());
     Value bcastedScale = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
-        loc, outTy, op.scale(), inputShape, op.axis());
-    auto zeroPointTy = op.zero_point().getType().cast<RankedTensorType>();
+        loc, outTy, op.getScale(), inputShape, op.getAxis());
+    auto zeroPointTy = op.getZeroPoint().getType().cast<RankedTensorType>();
     auto bcastedInputOrZeroPointTy =
         RankedTensorType::get(inputTy.getShape(), zeroPointTy.getElementType(),
                               zeroPointTy.getEncoding());
     Value bcastedZeroPoint = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
-        loc, bcastedInputOrZeroPointTy, op.zero_point(), inputShape, op.axis());
+        loc, bcastedInputOrZeroPointTy, op.getZeroPoint(), inputShape,
+        op.getAxis());
     Value castedInput = rewriter.create<mhlo::ConvertOp>(
-        loc, bcastedInputOrZeroPointTy, op.input());
+        loc, bcastedInputOrZeroPointTy, op.getInput());
 
     // output = (input - zero\_point) \times scale
     Value t0 =
diff --git a/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc b/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc
index e1168d65ea8..c4bcd45422a 100644
--- a/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_lower_to_library_call.cc
@@ -28,7 +28,7 @@ limitations under the License.
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/Debug.h"
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Attributes.h"
diff --git a/tao_compiler/mlir/disc/transforms/disc_memref_canonicalizer.cc b/tao_compiler/mlir/disc/transforms/disc_memref_canonicalizer.cc
index c764e9263a9..4988da54347 100644
--- a/tao_compiler/mlir/disc/transforms/disc_memref_canonicalizer.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_memref_canonicalizer.cc
@@ -44,7 +44,8 @@ struct ReinterpretCastOpConverter : public OpRewritePattern<ReinterpretCastOp> {
 LogicalResult ReinterpretCastOpConverter::matchAndRewrite(
     ReinterpretCastOp op, PatternRewriter& rewriter) const {
   Location loc = op.getLoc();
-  auto alloc = dyn_cast_or_null<memref::AllocOp>(op.source().getDefiningOp());
+  auto alloc =
+      dyn_cast_or_null<memref::AllocOp>(op.getSource().getDefiningOp());
   if (!alloc) return failure();
 
   auto allocTy = alloc.getResult().getType().cast<MemRefType>();
diff --git a/tao_compiler/mlir/disc/transforms/disc_memref_load_store_simplifier.cc b/tao_compiler/mlir/disc/transforms/disc_memref_load_store_simplifier.cc
index 7669437abe9..8ae3eb7863d 100644
--- a/tao_compiler/mlir/disc/transforms/disc_memref_load_store_simplifier.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_memref_load_store_simplifier.cc
@@ -13,7 +13,7 @@
 // eliminate unnecessary load/stores.
 
 #include "llvm/Support/Debug.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Location.h"
@@ -133,7 +133,7 @@ LogicalResult LoadScalarSimplifier::matchAndRewrite(
     }
   }
   if (store_op != nullptr) {
-    rewriter.replaceOp(op, store_op.value());
+    rewriter.replaceOp(op, store_op.getValue());
     return success();
   }
   return failure();
diff --git a/tao_compiler/mlir/disc/transforms/disc_outline_cpu_kernel.cc b/tao_compiler/mlir/disc/transforms/disc_outline_cpu_kernel.cc
index 0341bb7772a..2fd4a8bf596 100644
--- a/tao_compiler/mlir/disc/transforms/disc_outline_cpu_kernel.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_outline_cpu_kernel.cc
@@ -19,7 +19,7 @@ limitations under the License.
 #include "llvm/Support/Debug.h"
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
 #include "mlir/Dialect/Affine/LoopUtils.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
diff --git a/tao_compiler/mlir/disc/transforms/disc_passes.td b/tao_compiler/mlir/disc/transforms/disc_passes.td
index dd475c1f9d3..eb4bf222297 100644
--- a/tao_compiler/mlir/disc/transforms/disc_passes.td
+++ b/tao_compiler/mlir/disc/transforms/disc_passes.td
@@ -234,6 +234,9 @@ def RemoveShapeConstraintsPass : Pass<"disc-remove-shape-constraint", "mlir::fun
 def ConvertHloToStandardPass : Pass<"disc-convert-hlo-to-std", "mlir::func::FuncOp"> {
   let summary = "Pass to convert some hlo ops to std dialect";
   let constructor = "createDiscConvertHloToStandardPass()";
+  let dependentDialects = [
+    "arith::ArithDialect"
+  ];
 }
 
 def StdBufferizePass : Pass<"disc-std-bufferize", "mlir::func::FuncOp"> {
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc b/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc
index 116c4e3e969..ccaf79e81d2 100644
--- a/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_optimization.cc
@@ -23,7 +23,7 @@ limitations under the License.
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"  // TF:llvm-project
@@ -380,13 +380,13 @@ struct ExtractElementOfSliceOpCanonicalizationPattern
     if (!sliceOp) return failure();
     if (!isCandidateShapeTensorType(op.getTensor().getType())) return failure();
 
-    int64_t start = sliceOp.start_indices().getValues<int64_t>()[0];
-    int64_t stride = sliceOp.strides().getValues<int64_t>()[0];
+    int64_t start = sliceOp.getStartIndices().getValues<int64_t>()[0];
+    int64_t stride = sliceOp.getStrides().getValues<int64_t>()[0];
     int64_t inputIndex = (start + index * stride);
     Value newIndex =
         rewriter.create<arith::ConstantIndexOp>(op.getLoc(), inputIndex);
     Value newValue = rewriter.create<tensor::ExtractOp>(
-        op.getLoc(), sliceOp.operand(), newIndex);
+        op.getLoc(), sliceOp.getOperand(), newIndex);
 
     rewriter.replaceOp(op, {newValue});
     return success();
@@ -426,7 +426,8 @@ struct ExtractElementOfReshapeOpCanonicalizationPattern
     }
 
     SmallVector<int64_t> inputIndices;
-    auto inputTy = reshapeOp.operand().getType().dyn_cast<RankedTensorType>();
+    auto inputTy =
+        reshapeOp.getOperand().getType().dyn_cast<RankedTensorType>();
     if (!inputTy || !inputTy.hasStaticShape()) return failure();
     for (int i = inputTy.getRank() - 1; i >= 0; --i) {
       inputIndices.push_back(linearIndex % inputTy.getShape()[i]);
@@ -438,7 +439,7 @@ struct ExtractElementOfReshapeOpCanonicalizationPattern
           op.getLoc(), inputIndices[i]));
     }
     Value newValue = rewriter.create<tensor::ExtractOp>(
-        op.getLoc(), reshapeOp.operand(), newIndices);
+        op.getLoc(), reshapeOp.getOperand(), newIndices);
     rewriter.replaceOp(op, {newValue});
     return success();
   }
@@ -477,11 +478,11 @@ struct ExtractElementOfReduceOpCanonicalizationPattern
     if (!reduceOp || reduceOp->getNumResults() > 1) return failure();
 
     // Only support reducing arcross a single dimension.
-    if (reduceOp.dimensions().getValues<int64_t>().size() != 1)
+    if (reduceOp.getDimensions().getValues<int64_t>().size() != 1)
       return failure();
-    int64_t reduceAxis = *reduceOp.dimensions().getValues<int64_t>().begin();
+    int64_t reduceAxis = *reduceOp.getDimensions().getValues<int64_t>().begin();
 
-    auto& block = reduceOp.body().front();
+    auto& block = reduceOp.getBody().front();
     if (!hasSingleElement(block.without_terminator())) return failure();
     if (!isa<mhlo::MulOp>(&(*block.begin()))) return failure();
 
@@ -868,7 +869,7 @@ LogicalResult ShapeComputationIRAnalysis::applyIndexOpConstraint(
   } else if (auto dimOp = dyn_cast<tensor::DimOp>(op)) {
     Optional<int64_t> dimIndex = dimOp.getConstantIndex();
     if (!dimIndex) return success();
-    value2SymDim_[op->getResult(0)].setKnownNonNegative(true);
+    value2SymDim_[op->getResult(0)].updateKnownNonNegative(true);
     if (failed(mgr_.mapSymbolicDimEqual(
             value2SymDim_[op->getResult(0)],
             rankedTensor2SymDims_[dimOp.getSource()][*dimIndex])))
@@ -904,15 +905,15 @@ LogicalResult ShapeComputationIRAnalysis::applyIndexOpConstraint(
     auto lhsSym = mgr_.getRootSymbolicDim(value2SymDim_[lhs]);
     auto rhsSym = mgr_.getRootSymbolicDim(value2SymDim_[rhs]);
     auto outSym = mgr_.getRootSymbolicDim(value2SymDim_[out]);
-    if (lhsSym.knownNonNegative() && rhsSym.knownNonNegative()) {
-      outSym.setKnownNonNegative(true);
+    if (lhsSym.getKnownNonNegative() && rhsSym.getKnownNonNegative()) {
+      outSym.updateKnownNonNegative(true);
     }
-    if (lhsSym.knownNonNegative() && rhsSym.knownNonSizeOne() ||
-        rhsSym.knownNonNegative() && lhsSym.knownNonSizeOne()) {
-      outSym.setKnownNonSizeOne(true);
+    if (lhsSym.getKnownNonNegative() && rhsSym.getKnownNonSizeOne() ||
+        rhsSym.getKnownNonNegative() && lhsSym.getKnownNonSizeOne()) {
+      outSym.updateKnownNonSizeOne(true);
     }
-    if (lhsSym.knownNonSizeZero() && rhsSym.knownNonSizeZero()) {
-      outSym.setKnownNonSizeZero(true);
+    if (lhsSym.getKnownNonSizeZero() && rhsSym.getKnownNonSizeZero()) {
+      outSym.updateKnownNonSizeZero(true);
     }
     // TODO(disc): propagate other attributes/shape ranges??
   } else if (auto addOp = dyn_cast<arith::AddIOp>(op)) {
@@ -923,10 +924,10 @@ LogicalResult ShapeComputationIRAnalysis::applyIndexOpConstraint(
     auto lhsSym = mgr_.getRootSymbolicDim(value2SymDim_[lhs]);
     auto rhsSym = mgr_.getRootSymbolicDim(value2SymDim_[rhs]);
     auto outSym = mgr_.getRootSymbolicDim(value2SymDim_[out]);
-    if (lhsSym.knownNonNegative() && rhsSym.knownNonNegative()) {
-      outSym.setKnownNonNegative(true);
-      if (lhsSym.knownNonSizeZero() || rhsSym.knownNonSizeZero())
-        outSym.setKnownNonSizeZero(true);
+    if (lhsSym.getKnownNonNegative() && rhsSym.getKnownNonNegative()) {
+      outSym.updateKnownNonNegative(true);
+      if (lhsSym.getKnownNonSizeZero() || rhsSym.getKnownNonSizeZero())
+        outSym.updateKnownNonSizeZero(true);
     }
     // TODO(disc): propagate other constraint attributes/shape ranges??
   }
@@ -1019,45 +1020,46 @@ LogicalResult ShapeComputationIRAnalysis::applyMhloElemOpConstraint(
         return op->emitError()
                << "fail to merge symbolic dim between operands of element op\n";
   } else if (auto clamp = dyn_cast<mhlo::ClampOp>(op)) {
-    auto operandTy = clamp.operand().getType().dyn_cast<RankedTensorType>();
-    auto minTy = clamp.min().getType().dyn_cast<RankedTensorType>();
-    auto maxTy = clamp.max().getType().dyn_cast<RankedTensorType>();
+    auto operandTy = clamp.getOperand().getType().dyn_cast<RankedTensorType>();
+    auto minTy = clamp.getMin().getType().dyn_cast<RankedTensorType>();
+    auto maxTy = clamp.getMax().getType().dyn_cast<RankedTensorType>();
     if (!operandTy || !minTy || !maxTy) return success();
 
     if (minTy.getRank() != 0) {
-      if (failed(mapRankedValueShapeEqual(clamp.operand(), clamp.min())))
+      if (failed(mapRankedValueShapeEqual(clamp.getOperand(), clamp.getMin())))
         return op->emitError()
                << "fail to merge the symbolic dim of operand and "
                   "min of mhlo::ClampOp\n";
     }
     if (maxTy.getRank() != 0) {
-      if (failed(mapRankedValueShapeEqual(clamp.operand(), clamp.max())))
+      if (failed(mapRankedValueShapeEqual(clamp.getOperand(), clamp.getMax())))
         return op->emitError()
                << "fail to merge the symbolic dim of operand and "
                   "max of mhlo::ClampOp\n";
     }
   } else if (auto select = dyn_cast<mhlo::SelectOp>(op)) {
-    auto predTy = select.pred().getType().dyn_cast<RankedTensorType>();
-    auto trueTy = select.on_true().getType().dyn_cast<RankedTensorType>();
-    auto falseTy = select.on_false().getType().dyn_cast<RankedTensorType>();
+    auto predTy = select.getPred().getType().dyn_cast<RankedTensorType>();
+    auto trueTy = select.getOnTrue().getType().dyn_cast<RankedTensorType>();
+    auto falseTy = select.getOnFalse().getType().dyn_cast<RankedTensorType>();
     auto resultTy = select.getResult().getType().dyn_cast<RankedTensorType>();
     if (!predTy || !trueTy || !falseTy || !resultTy) return success();
 
     if (predTy.getRank() != 0) {
-      if (failed(mapRankedValueShapeEqual(select.pred(), select.getResult())))
+      if (failed(
+              mapRankedValueShapeEqual(select.getPred(), select.getResult())))
         return op->emitError() << "fail to merge the symbolic dim of pred and "
                                   "result of mhlo::SelectOp\n";
     }
     if (trueTy.getRank() != 0) {
       if (failed(
-              mapRankedValueShapeEqual(select.on_true(), select.getResult())))
+              mapRankedValueShapeEqual(select.getOnTrue(), select.getResult())))
         return op->emitError()
                << "fail to merge the symbolic dim of on_true and "
                   "result of mhlo::SelectOp\n";
     }
     if (falseTy.getRank() != 0) {
-      if (failed(
-              mapRankedValueShapeEqual(select.on_false(), select.getResult())))
+      if (failed(mapRankedValueShapeEqual(select.getOnFalse(),
+                                          select.getResult())))
         return op->emitError()
                << "fail to merge the symbolic dim of on_false and "
                   "result of mhlo::SelectOp\n";
@@ -1079,7 +1081,7 @@ LogicalResult ShapeComputationIRAnalysis::applyMhloDotLikeOpConstraint(
     auto& rhsDims = rankedTensor2SymDims_[rhs];
     if (lhsTy.getRank() != lhsDims.size() || rhsTy.getRank() != rhsDims.size())
       return op->emitError("lhs or rhs mismatch rank\n");
-    auto dim_numbers = dot_general.dot_dimension_numbers();
+    auto dim_numbers = dot_general.getDotDimensionNumbers();
     // Contracting dimensions.
     auto lhs_contracting_dims = dim_numbers.getLhsContractingDimensions();
     auto rhs_contracting_dims = dim_numbers.getRhsContractingDimensions();
@@ -1115,16 +1117,16 @@ LogicalResult ShapeComputationIRAnalysis::applyMhloDotLikeOpConstraint(
             mgr_.mapSymbolicDimEqual(lhsDims[lhsTy.getRank() - 1], rhsDims[0])))
       return op->emitError() << "fail to merge dim\n";
   } else if (auto einsum = dyn_cast<mhlo::EinsumOp>(op)) {
-    auto lhsTy = einsum.lhs().getType().dyn_cast<RankedTensorType>();
-    auto rhsTy = einsum.rhs().getType().dyn_cast<RankedTensorType>();
+    auto lhsTy = einsum.getLhs().getType().dyn_cast<RankedTensorType>();
+    auto rhsTy = einsum.getRhs().getType().dyn_cast<RankedTensorType>();
     auto outTy = einsum.getResult().getType().dyn_cast<RankedTensorType>();
     if (!lhsTy || !rhsTy || !outTy) return success();
 
-    auto& lhsDims = rankedTensor2SymDims_[einsum.lhs()];
-    auto& rhsDims = rankedTensor2SymDims_[einsum.rhs()];
+    auto& lhsDims = rankedTensor2SymDims_[einsum.getLhs()];
+    auto& rhsDims = rankedTensor2SymDims_[einsum.getRhs()];
     auto& outDims = rankedTensor2SymDims_[einsum.getResult()];
 
-    StringRef equation = einsum.einsum_config();
+    StringRef equation = einsum.getEinsumConfig();
     llvm::SmallDenseMap<char, llvm::SmallDenseMap<EquationVariable, size_t>>
         all_tokens;
     if (!parseEinsumEquation(equation, all_tokens, nullptr, nullptr, nullptr)) {
@@ -1194,7 +1196,7 @@ LogicalResult ShapeComputationIRAnalysis::applyMhloConcatOpConstraint(
     Operation* op) {
   if (auto concat = dyn_cast<mhlo::ConcatenateOp>(op)) {
     Value out = op->getResult(0);
-    int64_t axis = concat.dimension();
+    int64_t axis = concat.getDimension();
     auto ty = out.getType().dyn_cast<RankedTensorType>();
     if (!ty) return success();
     auto& outDims = rankedTensor2SymDims_[out];
@@ -1355,7 +1357,8 @@ LogicalResult ShapeComputationIRAnalysis::applyTieShapeOpConstraint(
       if (failed(mgr_.mapSymbolicDimEqual(value2SymDim_[en.value()],
                                           resultDims[en.index()])))
         return op->emitError() << "fail to merge symbolic dim\n";
-      mgr_.getRootSymbolicDim(resultDims[en.index()]).setKnownNonNegative(true);
+      mgr_.getRootSymbolicDim(resultDims[en.index()])
+          .updateKnownNonNegative(true);
     }
 
     if (isCandidateShapeTensorType(op->getResult(0).getType())) {
@@ -1575,26 +1578,26 @@ LogicalResult tryToSimplifyCompareOp(ShapeComputationIRAnalysis& analysis,
     if (op.getPredicate() == arith::CmpIPredicate::eq) {
       if (lhsSym == rhsSym) {
         pred = truePred;
-      } else if (lhsSym.knownNonNegative() && rhsSym.knownNegativeOne()) {
+      } else if (lhsSym.getKnownNonNegative() && rhsSym.getKnownNegativeOne()) {
         pred = falsePred;
-      } else if (rhsSym.knownNonNegative() && lhsSym.knownNegativeOne()) {
+      } else if (rhsSym.getKnownNonNegative() && lhsSym.getKnownNegativeOne()) {
         pred = falsePred;
-      } else if (lhsSym.knownNonSizeZero() && rhsSym.getDimSize() == 0) {
+      } else if (lhsSym.getKnownNonSizeZero() && rhsSym.getDimSize() == 0) {
         pred = falsePred;
-      } else if (rhsSym.knownNonSizeZero() && lhsSym.getDimSize() == 0) {
+      } else if (rhsSym.getKnownNonSizeZero() && lhsSym.getDimSize() == 0) {
         pred = falsePred;
       }
       // TODO(disc): support other cases
     } else if (op.getPredicate() == arith::CmpIPredicate::ne) {
       if (lhsSym == rhsSym) {
         pred = falsePred;
-      } else if (lhsSym.knownNonNegative() && rhsSym.knownNegativeOne()) {
+      } else if (lhsSym.getKnownNonNegative() && rhsSym.getKnownNegativeOne()) {
         pred = truePred;
-      } else if (rhsSym.knownNonNegative() && lhsSym.knownNegativeOne()) {
+      } else if (rhsSym.getKnownNonNegative() && lhsSym.getKnownNegativeOne()) {
         pred = truePred;
-      } else if (lhsSym.knownNonSizeZero() && rhsSym.getDimSize() == 0) {
+      } else if (lhsSym.getKnownNonSizeZero() && rhsSym.getDimSize() == 0) {
         pred = truePred;
-      } else if (rhsSym.knownNonSizeZero() && lhsSym.getDimSize() == 0) {
+      } else if (rhsSym.getKnownNonSizeZero() && lhsSym.getDimSize() == 0) {
         pred = truePred;
       }
       // TODO(disc): support other cases
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc b/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc
index e62d0fc8aee..e7b7910e97f 100644
--- a/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_optimization_utils.cc
@@ -23,7 +23,7 @@ limitations under the License.
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/MLIRContext.h"
@@ -214,7 +214,7 @@ SymbolicDimMgr::simplifySymbolicDimProductPair(const SymbolicDimProduct& x,
 
   for (SymbolicDimOp op : lhs.symbols) {
     auto it = rhsSymbolMap.find(op);
-    if (it != rhsSymbolMap.end() && op.knownNonSizeZero()) {
+    if (it != rhsSymbolMap.end() && op.getKnownNonSizeZero()) {
       if (--it->second == 0) rhsSymbolMap.erase(it);
       continue;
     }
@@ -223,7 +223,7 @@ SymbolicDimMgr::simplifySymbolicDimProductPair(const SymbolicDimProduct& x,
 
   for (SymbolicDimOp op : rhs.symbols) {
     auto it = lhsSymbolMap.find(op);
-    if (it != lhsSymbolMap.end() && op.knownNonSizeZero()) {
+    if (it != lhsSymbolMap.end() && op.getKnownNonSizeZero()) {
       if (--it->second == 0) lhsSymbolMap.erase(it);
       continue;
     }
@@ -617,7 +617,7 @@ LogicalResult SymbolicDimMgr::loadShapeConstraintGraph() {
         product.factor *= constOp->getAttrOfType<IntegerAttr>("value").getInt();
         continue;
       } else if (auto dimOp = dyn_cast_or_null<disc_shape::DimOp>(definingOp)) {
-        auto sym = symbolTable_.lookup<SymbolicDimOp>(dimOp.name());
+        auto sym = symbolTable_.lookup<SymbolicDimOp>(dimOp.getName());
         if (!sym) return dimOp->emitError() << "fail to find symbolic dim op\n";
         product.symbols.push_back(sym);
         continue;
@@ -628,8 +628,8 @@ LogicalResult SymbolicDimMgr::loadShapeConstraintGraph() {
   };
   if (func.walk([&](disc_shape::TieProductEqualOp op) {
             SymbolicDimProduct lhs, rhs;
-            if (failed(build_sym_product(op.lhs(), lhs)) ||
-                failed(build_sym_product(op.rhs(), rhs)) ||
+            if (failed(build_sym_product(op.getLhs(), lhs)) ||
+                failed(build_sym_product(op.getRhs(), rhs)) ||
                 failed(mapSymbolicDimProductEqual(lhs, rhs)))
               return WalkResult::interrupt();
             return WalkResult::advance();
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_simplifier.cc b/tao_compiler/mlir/disc/transforms/disc_shape_simplifier.cc
index eae729add12..3870c272c72 100644
--- a/tao_compiler/mlir/disc/transforms/disc_shape_simplifier.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_simplifier.cc
@@ -57,7 +57,7 @@ limitations under the License.
 #include "llvm/Support/Debug.h"
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"  // TF:llvm-project
@@ -115,7 +115,7 @@ struct DynamicReshapeOpPartialShapeInference
                                 PatternRewriter& rewriter) const override {
     auto loc = op.getLoc();
     auto output_shape =
-        op.output_shape().getDefiningOp<tensor::FromElementsOp>();
+        op.getOutputShape().getDefiningOp<tensor::FromElementsOp>();
     if (!output_shape) {
       return failure();
     }
@@ -147,7 +147,7 @@ struct DynamicReshapeOpPartialShapeInference
     }
     auto new_type = result_type.clone(result_dims);
     auto new_op = rewriter.create<mhlo::DynamicReshapeOp>(
-        loc, new_type, op.operand(), output_shape);
+        loc, new_type, op.getOperand(), output_shape);
     rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), new_op);
     return success();
   }
@@ -169,14 +169,14 @@ class DynamicReshapeOpShapeInference
 
   LogicalResult matchAndRewrite(mhlo::DynamicReshapeOp op,
                                 PatternRewriter& rewriter) const override {
-    Operation* shape_def_op = op.output_shape().getDefiningOp();
+    Operation* shape_def_op = op.getOutputShape().getDefiningOp();
     if (!shape_def_op) return failure();
     DenseIntElementsAttr cst_attr;
     if (auto cst_shape = dyn_cast<arith::ConstantOp>(shape_def_op)) {
       cst_attr = cst_shape.getValue().dyn_cast_or_null<DenseIntElementsAttr>();
     } else if (auto mhlo_cst_shape = dyn_cast<mhlo::ConstantOp>(shape_def_op)) {
       cst_attr =
-          mhlo_cst_shape.value().dyn_cast_or_null<DenseIntElementsAttr>();
+          mhlo_cst_shape.getValue().dyn_cast_or_null<DenseIntElementsAttr>();
     }
     if (!cst_attr) return failure();
     auto elem_ty = cst_attr.getType().cast<ShapedType>().getElementType();
@@ -196,7 +196,7 @@ class DynamicReshapeOpShapeInference
         RankedTensorType::get(dims, result_ty.getElementType());
     if (new_ty == result_ty) return failure();
     auto new_reshape = rewriter.create<mhlo::DynamicReshapeOp>(
-        op.getLoc(), new_ty, op.operand(), op.output_shape());
+        op.getLoc(), new_ty, op.getOperand(), op.getOutputShape());
     rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), new_reshape);
     return success();
   }
diff --git a/tao_compiler/mlir/disc/transforms/disc_shape_to_std.cc b/tao_compiler/mlir/disc/transforms/disc_shape_to_std.cc
index 25bc7b78f1a..765b992e79d 100644
--- a/tao_compiler/mlir/disc/transforms/disc_shape_to_std.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_shape_to_std.cc
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
@@ -374,13 +374,13 @@ LogicalResult LinearizeOpConversion::matchAndRewrite(
     LinearizeOp op, OpAdaptor adaptor,
     ConversionPatternRewriter& rewriter) const {
   Location loc = op.getLoc();
-  int rank = adaptor.multiDimIndexes().size();
+  int rank = adaptor.getMultiDimIndexes().size();
   Value linear = rewriter.create<arith::ConstantIndexOp>(loc, 0);
   ;
   if (rank > 0) {
-    linear = adaptor.multiDimIndexes().front();
-    for (auto&& z : llvm::zip(adaptor.multiDimIndexes().drop_front(),
-                              adaptor.shapeDimIndexes().drop_front())) {
+    linear = adaptor.getMultiDimIndexes().front();
+    for (auto&& z : llvm::zip(adaptor.getMultiDimIndexes().drop_front(),
+                              adaptor.getShapeDimIndexes().drop_front())) {
       linear = rewriter.create<arith::AddIOp>(
           loc, rewriter.create<arith::MulIOp>(loc, linear, std::get<1>(z)),
           std::get<0>(z));
@@ -404,16 +404,16 @@ LogicalResult DelinearizeOpConversion::matchAndRewrite(
     DelinearizeOp op, OpAdaptor adaptor,
     ConversionPatternRewriter& rewriter) const {
   Location loc = op.getLoc();
-  int rank = adaptor.shapeDimIndexes().size();
+  int rank = adaptor.getShapeDimIndexes().size();
   if (rank == 0) {
     rewriter.eraseOp(op);
   } else if (rank == 1) {
-    rewriter.replaceOp(op, ValueRange{adaptor.linearIndex()});
+    rewriter.replaceOp(op, ValueRange{adaptor.getLinearIndex()});
   } else {
-    Value linear = adaptor.linearIndex();
+    Value linear = adaptor.getLinearIndex();
     SmallVector<Value> multiDims(rank);
     for (auto&& en : llvm::enumerate(
-             llvm::reverse(adaptor.shapeDimIndexes().drop_front()))) {
+             llvm::reverse(adaptor.getShapeDimIndexes().drop_front()))) {
       multiDims[rank - 1 - en.index()] =
           rewriter.create<arith::RemUIOp>(loc, linear, en.value());
       linear = rewriter.create<arith::DivUIOp>(loc, linear, en.value());
@@ -433,7 +433,7 @@ void ConvertShapeToStandardPass::runOnOperation() {
   // Setup target legality.
   MLIRContext& ctx = getContext();
   ConversionTarget target(ctx);
-  target.addLegalDialect<arith::ArithmeticDialect, tensor::TensorDialect>();
+  target.addLegalDialect<arith::ArithDialect, tensor::TensorDialect>();
   target.addLegalOp<func::FuncOp, ModuleOp>();
 
   // Setup conversion patterns.
diff --git a/tao_compiler/mlir/disc/transforms/disc_specialize_fusion_with_speculation.cc b/tao_compiler/mlir/disc/transforms/disc_specialize_fusion_with_speculation.cc
index dbcce8c48e4..d3ba25fadc9 100644
--- a/tao_compiler/mlir/disc/transforms/disc_specialize_fusion_with_speculation.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_specialize_fusion_with_speculation.cc
@@ -14,7 +14,7 @@
 
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
 #include "mlir-hlo/utils/codegen_utils.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BlockAndValueMapping.h"
diff --git a/tao_compiler/mlir/disc/transforms/disc_std_bufferize.cc b/tao_compiler/mlir/disc/transforms/disc_std_bufferize.cc
index 5b626329e35..5cfe9d2e0f6 100644
--- a/tao_compiler/mlir/disc/transforms/disc_std_bufferize.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_std_bufferize.cc
@@ -14,7 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -129,7 +129,7 @@ void StdBufferizePass::runOnOperation() {
 
   target.addLegalDialect<memref::MemRefDialect>();
   target.addLegalOp<func::FuncOp, ModuleOp>();
-  target.addDynamicallyLegalDialect<arith::ArithmeticDialect>(
+  target.addDynamicallyLegalDialect<arith::ArithDialect>(
       [&](Operation* op) { return typeConverter.isLegal(op); });
 
   // Setup conversion patterns.
diff --git a/tao_compiler/mlir/disc/transforms/disc_tensor_to_std.cc b/tao_compiler/mlir/disc/transforms/disc_tensor_to_std.cc
index c1392fe0e3a..0561ca53413 100644
--- a/tao_compiler/mlir/disc/transforms/disc_tensor_to_std.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_tensor_to_std.cc
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -103,7 +103,7 @@ void ConvertTensorToStandardPass::runOnOperation() {
   // Setup target legality.
   MLIRContext& ctx = getContext();
   ConversionTarget target(ctx);
-  target.addLegalDialect<arith::ArithmeticDialect>();
+  target.addLegalDialect<arith::ArithDialect>();
   target.addLegalOp<func::FuncOp, ModuleOp>();
   target.addLegalOp<tensor::ExtractOp, tensor::DimOp, tensor::FromElementsOp>();
   target.addIllegalOp<GenerateOp>();
diff --git a/tao_compiler/mlir/disc/transforms/disc_to_llvm.cc b/tao_compiler/mlir/disc/transforms/disc_to_llvm.cc
index b91d9f5951c..e7fcf943fe4 100644
--- a/tao_compiler/mlir/disc/transforms/disc_to_llvm.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_to_llvm.cc
@@ -13,14 +13,14 @@
 
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Support/FormatVariadic.h"
-#include "mlir/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.h"
+#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
 #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
-#include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
@@ -484,7 +484,7 @@ Value ConvertLaunchFuncOpToRalCallPattern::generateParamsArray(
 LogicalResult ConvertLaunchFuncOpToRalCallPattern::matchAndRewrite(
     gpu::LaunchFuncOp launch_op, OpAdaptor adaptor,
     ConversionPatternRewriter& rewriter) const {
-  if (!launch_op.asyncDependencies().empty() || launch_op.asyncToken()) {
+  if (!launch_op.getAsyncDependencies().empty() || launch_op.getAsyncToken()) {
     return rewriter.notifyMatchFailure(
         launch_op, "Cannot convert with async dependency or result.");
   }
@@ -594,8 +594,8 @@ LogicalResult ConvertLaunchFuncOpToRalCallPattern::matchAndRewrite(
       module_blobs_array_ptr, /* fatbin strings */
       num_blobs, /* number of fatbin strings */
       kernel_name_global, /* name of the kernel to launch */
-      adaptor.gridSizeX(), adaptor.gridSizeY(), adaptor.gridSizeZ(),
-      adaptor.blockSizeX(), adaptor.blockSizeY(), adaptor.blockSizeZ(),
+      adaptor.getGridSizeX(), adaptor.getGridSizeY(), adaptor.getGridSizeZ(),
+      adaptor.getBlockSizeX(), adaptor.getBlockSizeY(), adaptor.getBlockSizeZ(),
       zero, /* sharedMemBytes */
       stream_idx, /* gpu stream index */
       num_arg_value, /* num_args */
@@ -737,7 +737,7 @@ LogicalResult ConvertMemRefDeallocOpToDispatchOpPattern::matchAndRewrite(
   Location loc = op->getLoc();
 
   StringRef device =
-      (placement_utils::isGpuMemRef(dealloc_op.memref())) ? "gpu" : "cpu";
+      (placement_utils::isGpuMemRef(dealloc_op.getMemref())) ? "gpu" : "cpu";
 
   // get ral context
   LLVMFuncOp parent_func = dealloc_op->getParentOfType<LLVMFuncOp>();
@@ -745,7 +745,7 @@ LogicalResult ConvertMemRefDeallocOpToDispatchOpPattern::matchAndRewrite(
   Value context_arg = parent_func.getArgument(0);
 
   // create dispatch op
-  MemRefDescriptor memref(adaptor.memref());
+  MemRefDescriptor memref(adaptor.getMemref());
   Value allocated_bytes_ptr = rewriter.create<LLVM::BitcastOp>(
       loc, getVoidPtrType(), memref.allocatedPtr(rewriter, loc));
 
@@ -1026,9 +1026,9 @@ class DiscToLLVMPass : public DiscToLLVMPassBase<DiscToLLVMPass> {
 
     // Populate patterns.
     RewritePatternSet patterns(&getContext());
-    mlir::arith::populateArithmeticToLLVMConversionPatterns(type_converter,
-                                                            patterns);
-    arith::populateArithmeticExpandOpsPatterns(patterns);
+    mlir::arith::populateArithToLLVMConversionPatterns(type_converter,
+                                                       patterns);
+    arith::populateArithExpandOpsPatterns(patterns);
     populateMemRefToLLVMConversionPatterns(type_converter, patterns);
     populateMathToLLVMConversionPatterns(type_converter, patterns);
     populateFuncToLLVMConversionPatterns(type_converter, patterns);
@@ -1039,7 +1039,7 @@ class DiscToLLVMPass : public DiscToLLVMPassBase<DiscToLLVMPass> {
     // Set target.
     ConversionTarget target(*ctx);
     target.addLegalDialect<LLVM::LLVMDialect>();
-    target.addIllegalDialect<arith::ArithmeticDialect, gpu::GPUDialect,
+    target.addIllegalDialect<arith::ArithDialect, gpu::GPUDialect,
                              disc_ral::RalDialect, math::MathDialect>();
     // Mark modules as legal.
     target.addLegalOp<ModuleOp, gpu::GPUModuleOp>();
diff --git a/tao_compiler/mlir/disc/transforms/disc_transpose_simplifier.cc b/tao_compiler/mlir/disc/transforms/disc_transpose_simplifier.cc
index 0fd2869e4f2..73215e712fb 100644
--- a/tao_compiler/mlir/disc/transforms/disc_transpose_simplifier.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_transpose_simplifier.cc
@@ -26,7 +26,7 @@ limitations under the License.
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"  // TF:llvm-project
@@ -103,7 +103,7 @@ bool isIdentityPermutation(T&& permutation) {
 //   use(%0)
 LogicalResult eliminateIdentityTranspse(mhlo::TransposeOp op,
                                         PatternRewriter& rewriter) {
-  if (!isIdentityPermutation(op.permutation().getValues<int64_t>()))
+  if (!isIdentityPermutation(op.getPermutation().getValues<int64_t>()))
     return failure();
   rewriter.replaceOp(op, op->getOperands());
   return success();
@@ -122,7 +122,7 @@ LogicalResult propagateDimOfTranspse(tensor::DimOp op,
       dyn_cast_or_null<mhlo::TransposeOp>(op.getSource().getDefiningOp());
   if (!indexOp || !transposeOp) return failure();
 
-  auto perm = transposeOp.permutation().getValues<int64_t>();
+  auto perm = transposeOp.getPermutation().getValues<int64_t>();
   Value sourceIndex = rewriter.create<arith::ConstantIndexOp>(
       op.getLoc(), perm[indexOp.getValue().cast<IntegerAttr>().getInt()]);
   rewriter.replaceOpWithNewOp<tensor::DimOp>(op, transposeOp->getOperand(0),
@@ -175,7 +175,7 @@ bool isConstOneBcast(Operation* op) {
   auto constOp =
       dyn_cast_or_null<mhlo::ConstantOp>(op->getOperand(0).getDefiningOp());
   if (!constOp) return false;
-  auto attr = constOp.value().cast<DenseElementsAttr>();
+  auto attr = constOp.getValue().cast<DenseElementsAttr>();
   if (attr.getNumElements() != 1) return false;
   if (getElementTypeOrSelf(attr.getType()).isa<FloatType>())
     return (*(attr.getValues<APFloat>().begin())).convertToDouble() == 1.0;
@@ -225,9 +225,9 @@ bool isMirroredTranspose(Operation* x, Operation* y) {
   auto transposeY = dyn_cast_or_null<mhlo::TransposeOp>(y);
   if (!transposeX || !transposeY) return false;
 
-  auto permXAttr = transposeX.permutation().getValues<int64_t>();
+  auto permXAttr = transposeX.getPermutation().getValues<int64_t>();
   SmallVector<int64_t> permX{permXAttr.begin(), permXAttr.end()};
-  auto permYAttr = transposeY.permutation().getValues<int64_t>();
+  auto permYAttr = transposeY.getPermutation().getValues<int64_t>();
   SmallVector<int64_t> permY{permYAttr.begin(), permYAttr.end()};
 
   if (permX.size() != permY.size()) return false;
@@ -581,7 +581,7 @@ LogicalResult pairMirroredTransposeOps(Block* block, bool& changed) {
       Value from = x->getResult(0);
       Value to = y->getOperand(0);
       auto permAttr =
-          cast<mhlo::TransposeOp>(y).permutation().getValues<int64_t>();
+          cast<mhlo::TransposeOp>(y).getPermutation().getValues<int64_t>();
       SmallVector<int64_t> perm{permAttr.begin(), permAttr.end()};
       DenseMap<Operation*, SmallVector<int64_t>> intermedateOpsPermutationMap;
       if (!ctx.dominates(from, to, perm, intermedateOpsPermutationMap))
@@ -677,8 +677,10 @@ LogicalResult reverseIfOperandsAreConsistentTransposeOps(Operation* op,
   if (!transposeLHS || !transposeRHS || transposeLHS == transposeRHS)
     return success();
 
-  SmallVector<int64_t> permLHS{transposeLHS.permutation().getValues<int64_t>()};
-  SmallVector<int64_t> permRHS{transposeRHS.permutation().getValues<int64_t>()};
+  SmallVector<int64_t> permLHS{
+      transposeLHS.getPermutation().getValues<int64_t>()};
+  SmallVector<int64_t> permRHS{
+      transposeRHS.getPermutation().getValues<int64_t>()};
   if (permLHS.size() != permRHS.size()) return success();
 
   for (const auto& en : llvm::zip(permLHS, permRHS)) {
@@ -769,14 +771,14 @@ LogicalResult reverseIfOperandsAndResultsAreConsistent(Operation* op,
 
   Block* block = op->getBlock();
   SmallVector<int64_t> reversePerm{
-      transposeOp.permutation().getValues<int64_t>()};
+      transposeOp.getPermutation().getValues<int64_t>()};
   auto perm = getReversePermutation(reversePerm);
   SmallVector<Operation*> transposeOps;
   for (Operation& candidate : *block) {
     auto transposeCandidate = dyn_cast<mhlo::TransposeOp>(&candidate);
     if (!transposeCandidate || &candidate == transposeOp) continue;
     SmallVector<int64_t> candidatePerm{
-        transposeCandidate.permutation().getValues<int64_t>()};
+        transposeCandidate.getPermutation().getValues<int64_t>()};
     if (perm != candidatePerm) continue;
     transposeOps.push_back(&candidate);
   }
diff --git a/tao_compiler/mlir/disc/transforms/disc_unhandled_atomic_rmw_converter.cc b/tao_compiler/mlir/disc/transforms/disc_unhandled_atomic_rmw_converter.cc
index 1ab4c8fadf4..b9c5b09873d 100644
--- a/tao_compiler/mlir/disc/transforms/disc_unhandled_atomic_rmw_converter.cc
+++ b/tao_compiler/mlir/disc/transforms/disc_unhandled_atomic_rmw_converter.cc
@@ -9,7 +9,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"       // TF:llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"                 // TF:llvm-project
 #include "mlir/Dialect/Func/IR/FuncOps.h"                // TF:llvm-project
 #include "mlir/Dialect/MemRef/IR/MemRef.h"               // TF:llvm-project
 #include "mlir/IR/Location.h"                            // TF:llvm-project
@@ -46,20 +46,21 @@ struct UnhandledAtomicRMWConverter
                                 PatternRewriter& rewriter) const override {
     // Currently, we only deal with atomic mulf operation. More operations can
     // be supported easily here.
-    if (op.kind() != arith::AtomicRMWKind::mulf) {
+    if (op.getKind() != arith::AtomicRMWKind::mulf) {
       return failure();
     }
 
     Location loc = op.getLoc();
     memref::GenericAtomicRMWOp genericOp =
-        rewriter.create<memref::GenericAtomicRMWOp>(loc, op.memref(),
-                                                    op.indices());
+        rewriter.create<memref::GenericAtomicRMWOp>(loc, op.getMemref(),
+                                                    op.getIndices());
     OpBuilder bodyBuilder =
         OpBuilder::atBlockEnd(genericOp.getBody(), rewriter.getListener());
 
     Value lhs = genericOp.getCurrentValue();
-    Value rhs = op.value();
-    Value reductionOp = getReductionOp(op.kind(), bodyBuilder, loc, lhs, rhs);
+    Value rhs = op.getValue();
+    Value reductionOp =
+        getReductionOp(op.getKind(), bodyBuilder, loc, lhs, rhs);
     bodyBuilder.create<memref::AtomicYieldOp>(loc, reductionOp);
 
     rewriter.replaceOp(op, genericOp.getResult());
diff --git a/tao_compiler/mlir/disc/transforms/dot_rewriter.cc b/tao_compiler/mlir/disc/transforms/dot_rewriter.cc
index ccbdca53666..e9204bf1a59 100644
--- a/tao_compiler/mlir/disc/transforms/dot_rewriter.cc
+++ b/tao_compiler/mlir/disc/transforms/dot_rewriter.cc
@@ -18,7 +18,7 @@
 #include <vector>
 
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"       // from @llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"                 // from @llvm-project
 #include "mlir/Dialect/Func/IR/FuncOps.h"                // from @llvm-project
 #include "mlir/Dialect/Tensor/IR/Tensor.h"               // from @llvm-project
 #include "mlir/IR/Attributes.h"                          // from @llvm-project
@@ -54,7 +54,7 @@ static inline bool isNonBatchingTransposeTensorValue(
   permutation.clear();
   if (auto transpose = dyn_cast<mhlo::TransposeOp>(val.getDefiningOp())) {
     for (auto& en :
-         llvm::enumerate(transpose.permutation().getValues<int64_t>())) {
+         llvm::enumerate(transpose.getPermutation().getValues<int64_t>())) {
       if (en.index() != en.value()) {
         if (batching_dims.find(en.index()) != batching_dims.end()) {
           return false;
@@ -81,8 +81,8 @@ struct DotToDotGeneralConvert : public OpRewritePattern<mhlo::DotOp> {
   LogicalResult matchAndRewrite(mhlo::DotOp op,
                                 PatternRewriter& rewriter) const override {
     Location loc = op.getLoc();
-    Value old_lhs = op.lhs();
-    Value old_rhs = op.rhs();
+    Value old_lhs = op.getLhs();
+    Value old_rhs = op.getRhs();
     auto old_lhs_ty = old_lhs.getType().dyn_cast<RankedTensorType>();
     if (!old_lhs_ty) return failure();
 
@@ -104,7 +104,7 @@ struct DotToDotGeneralConvert : public OpRewritePattern<mhlo::DotOp> {
         rhs_contracting_dims);
 
     Value dot_general = rewriter.create<mhlo::DotGeneralOp>(
-        op.getLoc(), op.getType(), op.lhs(), op.rhs(), dot_dimension_attr,
+        op.getLoc(), op.getType(), op.getLhs(), op.getRhs(), dot_dimension_attr,
         nullptr);
     rewriter.replaceOp(op, dot_general);
     return success();
@@ -195,8 +195,8 @@ struct DotGeneralConvert : public OpRewritePattern<mhlo::DotGeneralOp> {
   LogicalResult matchAndRewrite(mhlo::DotGeneralOp op,
                                 PatternRewriter& rewriter) const override {
     Location loc = op.getLoc();
-    Value old_lhs = op.lhs();
-    Value old_rhs = op.rhs();
+    Value old_lhs = op.getLhs();
+    Value old_rhs = op.getRhs();
 
     RankedTensorType old_l_type =
         old_lhs.getType().dyn_cast<RankedTensorType>();
@@ -206,7 +206,7 @@ struct DotGeneralConvert : public OpRewritePattern<mhlo::DotGeneralOp> {
       return failure();
     }
 
-    auto dim_numbers = op.dot_dimension_numbers();
+    auto dim_numbers = op.getDotDimensionNumbers();
     auto batching_dims = dim_numbers.getLhsBatchingDimensions();
     auto n_batch_dims = batching_dims.size();
     for (auto d : batching_dims)
@@ -320,8 +320,8 @@ struct TransposeFoldingConvert : public OpRewritePattern<mhlo::DotGeneralOp> {
     }
 
     Location loc = op.getLoc();
-    Value old_lhs = op.lhs();
-    Value old_rhs = op.rhs();
+    Value old_lhs = op.getLhs();
+    Value old_rhs = op.getRhs();
 
     RankedTensorType old_l_type =
         old_lhs.getType().dyn_cast<RankedTensorType>();
@@ -331,7 +331,7 @@ struct TransposeFoldingConvert : public OpRewritePattern<mhlo::DotGeneralOp> {
       return failure();
     }
 
-    auto dim_numbers = op.dot_dimension_numbers();
+    auto dim_numbers = op.getDotDimensionNumbers();
     auto lhs_batching_dims = dim_numbers.getLhsBatchingDimensions();
     SmallVector<int64_t, 4> lhs_perm;
     bool tp_lhs = isNonBatchingTransposeTensorValue(
@@ -407,9 +407,9 @@ struct EinsumToDotGeneralPattern : public OpRewritePattern<mhlo::EinsumOp> {
   LogicalResult matchAndRewrite(mhlo::EinsumOp op,
                                 PatternRewriter& rewriter) const override {
     Location loc = op.getLoc();
-    Value old_lhs = op.lhs();
-    Value old_rhs = op.rhs();
-    StringRef equation = op.einsum_config();
+    Value old_lhs = op.getLhs();
+    Value old_rhs = op.getRhs();
+    StringRef equation = op.getEinsumConfig();
 
     TokensTy tokens;
     llvm::SmallDenseSet<char> contracting_tokens;
diff --git a/tao_compiler/mlir/disc/transforms/element_type_converter.cc b/tao_compiler/mlir/disc/transforms/element_type_converter.cc
index a7f12774fc5..fd4e729ffe1 100644
--- a/tao_compiler/mlir/disc/transforms/element_type_converter.cc
+++ b/tao_compiler/mlir/disc/transforms/element_type_converter.cc
@@ -58,7 +58,7 @@ mhlo::ReduceOp CloneAndReplaceElementType(mhlo::ReduceOp op,
                                           ArrayRef<Value> operands) {
   int num_non_return_hlo_ops = 0;
   Operation* map_op = nullptr;
-  op.body().walk([&](Operation* hlo_op) {
+  op.getBody().walk([&](Operation* hlo_op) {
     if (isa<mhlo::ReturnOp>(hlo_op)) return;
     ++num_non_return_hlo_ops;
     map_op = hlo_op;
@@ -69,24 +69,24 @@ mhlo::ReduceOp CloneAndReplaceElementType(mhlo::ReduceOp op,
 
   Location loc = op.getLoc();
   mhlo::ReduceOp new_op = rewriter.create<mhlo::ReduceOp>(
-      loc, operands[0], operands[1], op.dimensions());
+      loc, operands[0], operands[1], op.getDimensions());
 
   if (isa<mhlo::AddOp>(map_op)) {
-    BuildReduceBody<mhlo::AddOp>(elem_type, &new_op.body(), &rewriter);
+    BuildReduceBody<mhlo::AddOp>(elem_type, &new_op.getBody(), &rewriter);
   } else if (isa<mhlo::MulOp>(map_op)) {
-    BuildReduceBody<mhlo::MulOp>(elem_type, &new_op.body(), &rewriter);
+    BuildReduceBody<mhlo::MulOp>(elem_type, &new_op.getBody(), &rewriter);
   } else if (isa<mhlo::MaxOp>(map_op)) {
-    BuildReduceBody<mhlo::MaxOp>(elem_type, &new_op.body(), &rewriter);
+    BuildReduceBody<mhlo::MaxOp>(elem_type, &new_op.getBody(), &rewriter);
   } else if (isa<mhlo::MinOp>(map_op)) {
-    BuildReduceBody<mhlo::MinOp>(elem_type, &new_op.body(), &rewriter);
+    BuildReduceBody<mhlo::MinOp>(elem_type, &new_op.getBody(), &rewriter);
   } else if (isa<mhlo::AndOp>(map_op)) {
     // We use MinOp to replace AndOp since std AndOp requires operands have i1
     // type.
-    BuildReduceBody<mhlo::MinOp>(elem_type, &new_op.body(), &rewriter);
+    BuildReduceBody<mhlo::MinOp>(elem_type, &new_op.getBody(), &rewriter);
   } else if (isa<mhlo::OrOp>(map_op)) {
     // We use MaxOp to replace OrOp since std AndOp requires operands have i1
     // type.
-    BuildReduceBody<mhlo::MaxOp>(elem_type, &new_op.body(), &rewriter);
+    BuildReduceBody<mhlo::MaxOp>(elem_type, &new_op.getBody(), &rewriter);
   } else {
     assert(false && "not supported reduce type");
   }
@@ -106,7 +106,7 @@ struct ConvertReduceOpWithSmallWidthIntType
     }
 
     SmallVector<int64_t, 4> dims_to_reduce;
-    for (auto v : op.dimensions().getValues<APInt>()) {
+    for (auto v : op.getDimensions().getValues<APInt>()) {
       dims_to_reduce.push_back(v.getSExtValue());
     }
     RankedTensorType ty =
@@ -202,8 +202,8 @@ struct ConvertDotGeneralOp : public OpRewritePattern<mhlo::DotGeneralOp> {
   LogicalResult matchAndRewrite(mhlo::DotGeneralOp op,
                                 PatternRewriter& rewriter) const override {
     Location loc = op.getLoc();
-    Value lhs = op.lhs();
-    Value rhs = op.rhs();
+    Value lhs = op.getLhs();
+    Value rhs = op.getRhs();
     FloatType f16_ty = rewriter.getF16Type();
     FloatType f32_ty = rewriter.getF32Type();
     RankedTensorType lhs_ty = lhs.getType().dyn_cast<RankedTensorType>();
@@ -221,7 +221,7 @@ struct ConvertDotGeneralOp : public OpRewritePattern<mhlo::DotGeneralOp> {
         RankedTensorType::getChecked(loc, result_ty.getShape(), f16_ty);
     // tensor dot general
     Value dot = rewriter.create<mhlo::DotGeneralOp>(
-        loc, f16_tensor_ty, lhs_f16, rhs_f16, op.dot_dimension_numbers(),
+        loc, f16_tensor_ty, lhs_f16, rhs_f16, op.getDotDimensionNumbers(),
         nullptr);
     Value fp32_dot = rewriter.create<mhlo::ConvertOp>(loc, dot, f32_ty);
     fp32_dot.setType(op.getResult().getType());
diff --git a/tao_compiler/mlir/disc/transforms/fusion_utils.cc b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
index fc67edf4809..0296bf7c10d 100644
--- a/tao_compiler/mlir/disc/transforms/fusion_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/fusion_utils.cc
@@ -20,7 +20,7 @@ limitations under the License.
 
 #include "mlir-hlo/Dialect/lhlo/transforms/map_lmhlo_to_scalar_op.h"
 #include "mlir-hlo/utils/placement_utils.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"  // TF:llvm-project
 #include "mlir/IR/BlockAndValueMapping.h"
diff --git a/tao_compiler/mlir/disc/transforms/gpu_conv_padding.cc b/tao_compiler/mlir/disc/transforms/gpu_conv_padding.cc
index 691a966a908..079aa8f810f 100644
--- a/tao_compiler/mlir/disc/transforms/gpu_conv_padding.cc
+++ b/tao_compiler/mlir/disc/transforms/gpu_conv_padding.cc
@@ -20,7 +20,7 @@ limitations under the License.
 // input tensor such that we can pass it to cuDNN.
 
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"  // TF:llvm-project
 #include "mlir/IR/Attributes.h"             // TF:llvm-project
 #include "mlir/IR/Builders.h"               // TF:llvm-project
@@ -133,7 +133,7 @@ struct DiscGpuConvPaddingLegalizationPass
     SmallVector<Value, 4> padding_low(rank, zero);
     SmallVector<Value, 4> padding_high(rank, zero);
     SmallVector<Value, 4> padding_interior(rank, zero);
-    auto dimension_numbers = op.dimension_numbers();
+    auto dimension_numbers = op.getDimensionNumbers();
     for (const auto& en :
          llvm::enumerate(dimension_numbers.getInputSpatialDimensions())) {
       padding_low[en.value()] = other_padding_low[en.index()];
@@ -167,9 +167,9 @@ struct DiscGpuConvPaddingLegalizationPass
   }
 
   void RewriteOp(mhlo::DynamicConvOp op) {
-    input = op.lhs();
-    filter = op.rhs();
-    padding = op.d_padding();
+    input = op.getLhs();
+    filter = op.getRhs();
+    padding = op.getDPadding();
     output = op.getResult();
 
     input_tp = input.getType().dyn_cast<RankedTensorType>();
diff --git a/tao_compiler/mlir/disc/transforms/lhlo_elemental_utils.cc b/tao_compiler/mlir/disc/transforms/lhlo_elemental_utils.cc
index e764717fe5f..39254c31049 100644
--- a/tao_compiler/mlir/disc/transforms/lhlo_elemental_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/lhlo_elemental_utils.cc
@@ -22,7 +22,7 @@ limitations under the License.
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
 #include "mlir-hlo/Dialect/lhlo/transforms/map_lmhlo_to_scalar_op.h"
 #include "mlir-hlo/Dialect/mhlo/transforms/map_mhlo_to_scalar_op.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -1168,11 +1168,15 @@ Value elementalLower<lmhlo::ClampOp>(OpBuilder* b, Location loc,
   Value lb_clipped =
       mhlo::impl::mapMhloOpToStdScalarOp<lmhlo::LhloToHloOp<lmhlo::MaxOp>>(
           loc, ArrayRef<Type>{elem_ty}, ArrayRef<Type>{elem_ty, elem_ty},
-          ArrayRef<Value>{operand, min}, b);
+          mhlo::MaxOp::Adaptor(ArrayRef<Value>{operand, min},
+                               op->getAttrDictionary()),
+          b);
   Value result =
       mhlo::impl::mapMhloOpToStdScalarOp<lmhlo::LhloToHloOp<lmhlo::MinOp>>(
           loc, ArrayRef<Type>{elem_ty}, ArrayRef<Type>{elem_ty, elem_ty},
-          ArrayRef<Value>{lb_clipped, max}, b);
+          mhlo::MinOp::Adaptor(ArrayRef<Value>{lb_clipped, max},
+                               op->getAttrDictionary()),
+          b);
   mayCreateStore(b, loc, op.getOperation(), result, output_index, lower_config);
   return result;
 }
@@ -1226,7 +1230,7 @@ memref::ReinterpretCastOp createMemRef1DReinterpretCast(OpBuilder& b,
     }
   }
   if (alignment) {
-    b.create<memref::AssumeAlignmentOp>(loc, cast, alignment.alignment());
+    b.create<memref::AssumeAlignmentOp>(loc, cast, alignment.getAlignment());
   }
 
   return cast;
diff --git a/tao_compiler/mlir/disc/transforms/lhlo_elemental_utils.h b/tao_compiler/mlir/disc/transforms/lhlo_elemental_utils.h
index 4a71917554a..02aeabe244d 100644
--- a/tao_compiler/mlir/disc/transforms/lhlo_elemental_utils.h
+++ b/tao_compiler/mlir/disc/transforms/lhlo_elemental_utils.h
@@ -20,7 +20,7 @@ limitations under the License.
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
 #include "mlir-hlo/Dialect/lhlo/transforms/map_lmhlo_to_scalar_op.h"
 #include "mlir-hlo/Dialect/mhlo/transforms/map_mhlo_to_scalar_op.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Builders.h"
 #include "tensorflow/compiler/mlir/disc/transforms/codegen_utils.h"
diff --git a/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc b/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
index 1f2c4078b6e..2d777d72120 100644
--- a/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
+++ b/tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
@@ -20,7 +20,7 @@ limitations under the License.
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
 #include "mlir-hlo/Dialect/lhlo/transforms/map_lmhlo_to_scalar_op.h"
 #include "mlir-hlo/utils/placement_utils.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
diff --git a/tao_compiler/mlir/disc/transforms/lower_tf.cc b/tao_compiler/mlir/disc/transforms/lower_tf.cc
index 75ea4925add..579c31f9032 100644
--- a/tao_compiler/mlir/disc/transforms/lower_tf.cc
+++ b/tao_compiler/mlir/disc/transforms/lower_tf.cc
@@ -19,7 +19,7 @@ limitations under the License.
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"       // from @llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"                 // from @llvm-project
 #include "mlir/Dialect/Tensor/IR/Tensor.h"               // from @llvm-project
 #include "mlir/IR/Attributes.h"                          // from @llvm-project
 #include "mlir/IR/BuiltinOps.h"                          // from @llvm-project
@@ -1219,7 +1219,7 @@ class ConvertBucketizeOp : public OpRewritePattern<TF::BucketizeOp> {
     auto buckets = rewriter.create<mhlo::ReduceOp>(
         loc, convert.getResult(), zero,
         GetI64ElementsAttr(reduce_dim, &rewriter));
-    mhlo::BuildReduceBody<mhlo::AddOp>(element_type, &buckets.body(),
+    mhlo::BuildReduceBody<mhlo::AddOp>(element_type, &buckets.getBody(),
                                        &rewriter);
 
     rewriter.replaceOp(op, {buckets.getResult(0)});
diff --git a/tao_compiler/mlir/disc/transforms/mhlo_decomp_rewriters.cc b/tao_compiler/mlir/disc/transforms/mhlo_decomp_rewriters.cc
index c9e4710013c..dad60d83f0b 100644
--- a/tao_compiler/mlir/disc/transforms/mhlo_decomp_rewriters.cc
+++ b/tao_compiler/mlir/disc/transforms/mhlo_decomp_rewriters.cc
@@ -15,7 +15,7 @@
 #include <vector>
 
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Attributes.h"
@@ -62,23 +62,24 @@ struct BatchNormInferenceOpConvert
 LogicalResult BatchNormInferenceOpConvert::matchAndRewrite(
     mhlo::BatchNormInferenceOp op, PatternRewriter& rewriter) const {
   auto loc = op.getLoc();
-  Value var = op.variance();
-  Value eps = getConstantLike(rewriter, op, op.epsilonAttr(), var);
+  Value var = op.getVariance();
+  Value eps = getConstantLike(rewriter, op, op.getEpsilonAttr(), var);
   var = rewriter.create<mhlo::AddOp>(loc, var, eps);
 
   Value rsqrt = rewriter.create<mhlo::RsqrtOp>(loc, var);
-  Value multiplier = rewriter.create<mhlo::MulOp>(loc, rsqrt, op.scale());
+  Value multiplier = rewriter.create<mhlo::MulOp>(loc, rsqrt, op.getScale());
   DenseIntElementsAttr broadcastDimensions =
-      rewriter.getI64TensorAttr(ArrayRef<int64_t>{op.feature_index()});
+      rewriter.getI64TensorAttr(ArrayRef<int64_t>{op.getFeatureIndex()});
   auto outType = op.getType();
-  Value scale =
-      broadcastAs(rewriter, op, multiplier, op.operand(), broadcastDimensions);
-  Value mean =
-      broadcastAs(rewriter, op, op.mean(), op.operand(), broadcastDimensions);
-  Value offset =
-      broadcastAs(rewriter, op, op.offset(), op.operand(), broadcastDimensions);
+  Value scale = broadcastAs(rewriter, op, multiplier, op.getOperand(),
+                            broadcastDimensions);
+  Value mean = broadcastAs(rewriter, op, op.getMean(), op.getOperand(),
+                           broadcastDimensions);
+  Value offset = broadcastAs(rewriter, op, op.getOffset(), op.getOperand(),
+                             broadcastDimensions);
 
-  Value xSubMean = rewriter.create<mhlo::SubtractOp>(loc, op.operand(), mean);
+  Value xSubMean =
+      rewriter.create<mhlo::SubtractOp>(loc, op.getOperand(), mean);
   Value comp1 = rewriter.create<mhlo::MulOp>(loc, xSubMean, scale);
   rewriter.replaceOpWithNewOp<mhlo::AddOp>(op, op.getType(), comp1, offset);
   return success();
diff --git a/tao_compiler/mlir/disc/transforms/mhlo_placer.cc b/tao_compiler/mlir/disc/transforms/mhlo_placer.cc
index 987d8f34fca..903ac2bb80f 100644
--- a/tao_compiler/mlir/disc/transforms/mhlo_placer.cc
+++ b/tao_compiler/mlir/disc/transforms/mhlo_placer.cc
@@ -422,14 +422,14 @@ void OpsPlacer::placeI32Ops() {
 
     // For these ops, currently we only have limited support.
     if (auto custom_call_op = dyn_cast<mhlo_disc::CustomCallOp>(op)) {
-      if (custom_call_op.call_target_name() == "topk" ||
-          custom_call_op.call_target_name() == "rng_uniform") {
+      if (custom_call_op.getCallTargetName() == "topk" ||
+          custom_call_op.getCallTargetName() == "rng_uniform") {
         is_shape_calc_op = false;
       }
       // TopK op only have gpu kernel implementation a.t.m. Thus we put its
       // operands to gpu as well to eliminate potential cross-device memory
       // copy.
-      if (this->on_gpu_ && custom_call_op.call_target_name() == "topk") {
+      if (this->on_gpu_ && custom_call_op.getCallTargetName() == "topk") {
         Value indices = op->getOperand(1);
         Operation* indicesOp = indices.getDefiningOp();
         if (indicesOp) {
diff --git a/tao_compiler/mlir/disc/transforms/placement_utils.cc b/tao_compiler/mlir/disc/transforms/placement_utils.cc
index b9fb52b72ea..a4422e5afcf 100644
--- a/tao_compiler/mlir/disc/transforms/placement_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/placement_utils.cc
@@ -106,7 +106,7 @@ ShapeOperandList getShapeCalcOperandList(Operation* op) {
     std::string target_name;
     if (isa<mhlo_disc::CustomCallOp>(op)) {
       auto custom_call = cast<mhlo_disc::CustomCallOp>(op);
-      target_name = custom_call.call_target_name().str();
+      target_name = custom_call.getCallTargetName().str();
     } else {
       auto custom_call = cast<lmhlo_disc::CustomCallOp>(op);
       target_name = custom_call.getCallTargetName().str();
diff --git a/tao_compiler/mlir/disc/transforms/placement_utils.h b/tao_compiler/mlir/disc/transforms/placement_utils.h
index 7e37e577d3f..d3a06c95765 100644
--- a/tao_compiler/mlir/disc/transforms/placement_utils.h
+++ b/tao_compiler/mlir/disc/transforms/placement_utils.h
@@ -13,7 +13,7 @@
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"  // TF:llvm-project
 #include "mlir/IR/MLIRContext.h"            // TF:llvm-project
diff --git a/tao_compiler/mlir/disc/transforms/quantized_dot_rewriter.cc b/tao_compiler/mlir/disc/transforms/quantized_dot_rewriter.cc
index f855a84705d..1e1a0f9985c 100644
--- a/tao_compiler/mlir/disc/transforms/quantized_dot_rewriter.cc
+++ b/tao_compiler/mlir/disc/transforms/quantized_dot_rewriter.cc
@@ -38,7 +38,7 @@ static inline bool isNonBatchingTransposeTensorValue(
   permutation.clear();
   if (auto transpose = dyn_cast<mhlo::TransposeOp>(val.getDefiningOp())) {
     for (auto& en :
-         llvm::enumerate(transpose.permutation().getValues<int64_t>())) {
+         llvm::enumerate(transpose.getPermutation().getValues<int64_t>())) {
       if (en.index() != en.value()) {
         if (batching_dims.find(en.index()) != batching_dims.end()) {
           return false;
@@ -94,19 +94,19 @@ struct QuantDotTransposeConvert
       return failure();
     }
 
-    auto inputTy = op.input().getType().dyn_cast<RankedTensorType>();
-    auto weightTy = op.weight().getType().dyn_cast<RankedTensorType>();
+    auto inputTy = op.getInput().getType().dyn_cast<RankedTensorType>();
+    auto weightTy = op.getWeight().getType().dyn_cast<RankedTensorType>();
 
     // This transpose pass can only support dot with rank equal to 2
     if (inputTy.getRank() != 2 || weightTy.getRank() != 2) {
       return failure();
     }
 
-    auto dim_numbers = op.dot_dimension_numbers();
+    auto dim_numbers = op.getDotDimensionNumbers();
     SmallVector<int64_t, 4> rhs_perm_check;
     auto rhs_batching_dims_check = dim_numbers.getRhsBatchingDimensions();
     bool tp_rhs_check = isNonBatchingTransposeTensorValue(
-        op.weight(), rhs_perm_check,
+        op.getWeight(), rhs_perm_check,
         std::unordered_set<int64_t>(rhs_batching_dims_check.begin(),
                                     rhs_batching_dims_check.end()));
 
@@ -118,10 +118,10 @@ struct QuantDotTransposeConvert
     SmallVector<int64_t> transposeAttr = {1, 0};
 
     OpBuilder b(op);
-    auto transpose_op = InsertTranspose(op, op.weight(), transposeAttr, b);
+    auto transpose_op = InsertTranspose(op, op.getWeight(), transposeAttr, b);
 
     Location loc = op.getLoc();
-    Value old_lhs = op.input();
+    Value old_lhs = op.getInput();
     Value old_rhs = transpose_op->getResult(0);
 
     RankedTensorType old_l_type =
@@ -173,11 +173,11 @@ struct QuantDotTransposeConvert
 
     Value newQuantizedDot =
         rewriter.create<mlir::mhlo_disc::QuantizedDotGeneralOp>(
-            loc, op.getType(), op.input(), transpose_op->getResult(0),
-            op.input_scale(), op.input_zero_point(), op.weight_scale(),
-            op.weight_zero_point(), op.result_scale(), op.result_zero_point(),
-            dot_dimension_attr, op.use_symmetric(), op.axis(),
-            op.use_dynamic());
+            loc, op.getType(), op.getInput(), transpose_op->getResult(0),
+            op.getInputScale(), op.getInputZeroPoint(), op.getWeightScale(),
+            op.getWeightZeroPoint(), op.getResultScale(),
+            op.getResultZeroPoint(), dot_dimension_attr, op.getUseSymmetric(),
+            op.getAxis(), op.getUseDynamic());
 
     rewriter.replaceOp(op, newQuantizedDot);
     return success();
diff --git a/tao_compiler/mlir/disc/transforms/ral_inject_execution_context.cc b/tao_compiler/mlir/disc/transforms/ral_inject_execution_context.cc
index a358210f4b0..3b3a265e22e 100644
--- a/tao_compiler/mlir/disc/transforms/ral_inject_execution_context.cc
+++ b/tao_compiler/mlir/disc/transforms/ral_inject_execution_context.cc
@@ -47,7 +47,7 @@ limitations under the License.
 // to the entry function. Thus, we don't rewrite all call ops and other
 // functions a.t.m. Re-visit this assumption if necessary.
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
diff --git a/tao_compiler/mlir/disc/transforms/revise_kernel_outlining.cc b/tao_compiler/mlir/disc/transforms/revise_kernel_outlining.cc
index 4b234f6e5d6..7af1b348992 100644
--- a/tao_compiler/mlir/disc/transforms/revise_kernel_outlining.cc
+++ b/tao_compiler/mlir/disc/transforms/revise_kernel_outlining.cc
@@ -23,7 +23,7 @@
 
 #include <map>
 
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -322,10 +322,10 @@ gpu::LaunchFuncOp expandMemRef(gpu::LaunchFuncOp launch_func_op, Value memref,
 
   // clone the Ops in the body of the gpu.FuncOp
   BlockAndValueMapping map;
-  Region& new_gpu_func_body = new_gpu_func_op.body();
+  Region& new_gpu_func_body = new_gpu_func_op.getBody();
   Block& new_gpu_func_entry_block = new_gpu_func_body.front();
   for (auto operand :
-       llvm::enumerate(gpu_func_op.body().front().getArguments())) {
+       llvm::enumerate(gpu_func_op.getBody().front().getArguments())) {
     if (operand.index() == memref_idx) {
       continue;
     } else if (operand.index() < memref_idx) {
@@ -337,17 +337,17 @@ gpu::LaunchFuncOp expandMemRef(gpu::LaunchFuncOp launch_func_op, Value memref,
     }
   }
   // memref of the entry block
-  auto memref_arg = gpu_func_op.body().front().getArgument(memref_idx);
-  Block& new_entry_block = new_gpu_func_op.body().front();
-  cloneRegionAndRemapLoad(&gpu_func_op.body(), &new_gpu_func_op.body(), map,
-                          memref_idx, memref_arg, new_entry_block, true);
+  auto memref_arg = gpu_func_op.getBody().front().getArgument(memref_idx);
+  Block& new_entry_block = new_gpu_func_op.getBody().front();
+  cloneRegionAndRemapLoad(&gpu_func_op.getBody(), &new_gpu_func_op.getBody(),
+                          map, memref_idx, memref_arg, new_entry_block, true);
 
   // update the FunctionType of the gpu.Launchfunc inside the module
   b.setInsertionPoint(launch_func_op);
   auto new_launch_func_op = b.create<gpu::LaunchFuncOp>(
       loc, new_gpu_func_op, launch_func_op.getGridSizeOperandValues(),
       launch_func_op.getBlockSizeOperandValues(),
-      launch_func_op.dynamicSharedMemorySize(), new_operands);
+      launch_func_op.getDynamicSharedMemorySize(), new_operands);
 
   launch_func_op.erase();
   gpu_func_op.erase();
@@ -397,7 +397,7 @@ class ReviseGpuKernelOutliningPass
         if (memref.index() < gpu::LaunchOp::kNumConfigOperands) {
           continue;
         }
-        auto arg_memref = gpu_func_op.body().front().getArgument(
+        auto arg_memref = gpu_func_op.getBody().front().getArgument(
             memref.index() - gpu::LaunchOp::kNumConfigOperands);
         if (arg_memref.getType().isa<MemRefType>() &&
             (!placement_utils::isGpuMemRef(arg_memref))) {
diff --git a/tao_compiler/mlir/disc/transforms/shape_utils.cc b/tao_compiler/mlir/disc/transforms/shape_utils.cc
index b9cde733817..bb917bbb134 100644
--- a/tao_compiler/mlir/disc/transforms/shape_utils.cc
+++ b/tao_compiler/mlir/disc/transforms/shape_utils.cc
@@ -21,7 +21,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Debug.h"
 #include "mlir-hlo/Dialect/lhlo/IR/lhlo_ops.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"  // TF:llvm-project
 #include "mlir/IR/Dominance.h"
@@ -317,7 +317,7 @@ LogicalResult ShapeAnalysisDeprecated::buildBlockDimValueMap(Block* block) {
         int64_t index_val;
         if (getConstIntValue(indices[0].getDefiningOp(), index_val) &&
             index_val == index) {
-          dimValue = DimValue(store.value());
+          dimValue = DimValue(store.getValue());
           return true;
         }
       }
@@ -580,9 +580,9 @@ LogicalResult ShapeAnalysisDeprecated::buildBlockDimValueMap(Block* block) {
         mayCreateOrMapConstInt(dimVal);
       }
     } else if (auto dynSlice = dyn_cast<mhlo::RealDynamicSliceOp>(*op)) {
-      auto starts = getTensorDimValues(dynSlice.start_indices());
-      auto limits = getTensorDimValues(dynSlice.limit_indices());
-      auto strides = getTensorDimValues(dynSlice.strides());
+      auto starts = getTensorDimValues(dynSlice.getStartIndices());
+      auto limits = getTensorDimValues(dynSlice.getLimitIndices());
+      auto strides = getTensorDimValues(dynSlice.getStrides());
       Value result = op->getResult(0);
       auto out_ty = result.getType().dyn_cast_or_null<RankedTensorType>();
       if (!starts.hasValue() || !limits.hasValue() || !strides.hasValue() ||
@@ -901,12 +901,12 @@ LogicalResult ShapeAnalysisDeprecated::applyMhloOpConstraint(Operation* op) {
     Value operand = op->getOperand(0);
     Value result = op->getResult(0);
     for (auto& en :
-         llvm::enumerate(transpose.permutation().getValues<int64_t>())) {
+         llvm::enumerate(transpose.getPermutation().getValues<int64_t>())) {
       mapDimEqual(operand, en.value(), result, en.index());
     }
   } else if (auto concat = dyn_cast<mhlo::ConcatenateOp>(op)) {
     Value result = op->getResult(0);
-    int64_t axis = concat.dimension();
+    int64_t axis = concat.getDimension();
     int64_t rank = result.getType().cast<RankedTensorType>().getRank();
     for (Value operand : op->getOperands()) {
       for (int64_t i = 0; i < rank; ++i) {
@@ -920,7 +920,7 @@ LogicalResult ShapeAnalysisDeprecated::applyMhloOpConstraint(Operation* op) {
     int64_t rank = operand.getType().cast<RankedTensorType>().getRank();
     int64_t resultDimIdx = 0;
     for (int64_t i = 0; i < rank; ++i) {
-      auto reduceDims = reduce.dimensions().getValues<int64_t>();
+      auto reduceDims = reduce.getDimensions().getValues<int64_t>();
       if (std::find(reduceDims.begin(), reduceDims.end(), i) !=
           reduceDims.end()) {
         continue;
@@ -934,7 +934,7 @@ LogicalResult ShapeAnalysisDeprecated::applyMhloOpConstraint(Operation* op) {
     auto ty = operand.getType().dyn_cast<ShapedType>();
     assert(ty);
     for (auto& dim : llvm::enumerate(
-             broadcast_in_dim.broadcast_dimensions().getValues<int64_t>())) {
+             broadcast_in_dim.getBroadcastDimensions().getValues<int64_t>())) {
       // Deal with non-static & non-one dim.
       if (!ty.isDynamicDim(dim.index()) && ty.getDimSize(dim.index()) != 1) {
         mapDimEqual(operand, dim.index(), result, dim.value());
@@ -961,7 +961,7 @@ LogicalResult ShapeAnalysisDeprecated::applyMhloOpConstraint(Operation* op) {
   } else if (auto dot_general = dyn_cast<mhlo::DotGeneralOp>(op)) {
     Value lhs = op->getOperand(0);
     Value rhs = op->getOperand(1);
-    auto dim_numbers = dot_general.dot_dimension_numbers();
+    auto dim_numbers = dot_general.getDotDimensionNumbers();
     DenseSet<int64_t> lhs_contract_batch_dims;
     DenseSet<int64_t> rhs_contract_batch_dims;
     // Contracting dimensions.
@@ -1012,7 +1012,7 @@ LogicalResult ShapeAnalysisDeprecated::applyMhloOpConstraint(Operation* op) {
                   i + lhs_batching_dims.size());
     }
   } else if (auto einsum = dyn_cast<mhlo::EinsumOp>(op)) {
-    StringRef equation = einsum.einsum_config();
+    StringRef equation = einsum.getEinsumConfig();
     llvm::SmallDenseMap<char, llvm::SmallDenseMap<EquationVariable, size_t>>
         all_tokens;
     if (!parseEinsumEquation(equation, all_tokens, nullptr, nullptr, nullptr)) {
@@ -1022,9 +1022,9 @@ LogicalResult ShapeAnalysisDeprecated::applyMhloOpConstraint(Operation* op) {
       SmallVector<std::pair<Value, int64_t>> equalValues;
       for (auto item : token.second) {
         if (item.first == kIsLhs) {
-          equalValues.push_back(std::make_pair(einsum.lhs(), item.second));
+          equalValues.push_back(std::make_pair(einsum.getLhs(), item.second));
         } else if (item.first == kIsRhs) {
-          equalValues.push_back(std::make_pair(einsum.rhs(), item.second));
+          equalValues.push_back(std::make_pair(einsum.getRhs(), item.second));
         } else {
           // kIsResult
           equalValues.push_back(
@@ -1041,36 +1041,38 @@ LogicalResult ShapeAnalysisDeprecated::applyMhloOpConstraint(Operation* op) {
       }
     }
   } else if (auto clamp = dyn_cast<mhlo::ClampOp>(op)) {
-    int64_t min_rank = clamp.min().getType().cast<RankedTensorType>().getRank();
-    int64_t max_rank = clamp.max().getType().cast<RankedTensorType>().getRank();
+    int64_t min_rank =
+        clamp.getMin().getType().cast<RankedTensorType>().getRank();
+    int64_t max_rank =
+        clamp.getMax().getType().cast<RankedTensorType>().getRank();
     if (min_rank != 0) {
-      mapShapeEqual(clamp.operand(), clamp.min());
+      mapShapeEqual(clamp.getOperand(), clamp.getMin());
     }
     if (max_rank != 0) {
-      mapShapeEqual(clamp.operand(), clamp.max());
+      mapShapeEqual(clamp.getOperand(), clamp.getMax());
     }
-    mapShapeEqual(clamp.operand(), op->getResult(0));
+    mapShapeEqual(clamp.getOperand(), op->getResult(0));
   } else if (auto select = dyn_cast<mhlo::SelectOp>(op)) {
     int64_t pred_rank =
-        select.pred().getType().cast<RankedTensorType>().getRank();
+        select.getPred().getType().cast<RankedTensorType>().getRank();
     int64_t true_rank =
-        select.on_true().getType().cast<RankedTensorType>().getRank();
+        select.getOnTrue().getType().cast<RankedTensorType>().getRank();
     int64_t false_rank =
-        select.on_false().getType().cast<RankedTensorType>().getRank();
+        select.getOnFalse().getType().cast<RankedTensorType>().getRank();
     if (pred_rank != 0) {
-      mapShapeEqual(select.pred(), select.getResult());
+      mapShapeEqual(select.getPred(), select.getResult());
     }
     if (true_rank != 0) {
-      mapShapeEqual(select.on_true(), select.getResult());
+      mapShapeEqual(select.getOnTrue(), select.getResult());
     }
     if (false_rank != 0) {
-      mapShapeEqual(select.on_false(), select.getResult());
+      mapShapeEqual(select.getOnFalse(), select.getResult());
     }
   } else if (isa<mhlo::GatherOp, mhlo::DynamicGatherOp>(op)) {
     auto gather = dyn_cast<mhlo::GatherOp>(op);
     auto d_gather = dyn_cast<mhlo::DynamicGatherOp>(op);
     auto dimension_numbers =
-        gather ? gather.dimension_numbers() : d_gather.dimension_numbers();
+        gather ? gather.getDimensionNumbers() : d_gather.getDimensionNumbers();
     Value in = op->getOperand(0);
     Value result = op->getResult(0);
     auto in_ty = in.getType().dyn_cast<RankedTensorType>();
diff --git a/tao_compiler/mlir/disc/transforms/split_large_ops.cc b/tao_compiler/mlir/disc/transforms/split_large_ops.cc
index 850d8add974..fb091b1dee2 100644
--- a/tao_compiler/mlir/disc/transforms/split_large_ops.cc
+++ b/tao_compiler/mlir/disc/transforms/split_large_ops.cc
@@ -69,20 +69,20 @@ struct ConvertConcatOp : public OpRewritePattern<mhlo::ConcatenateOp> {
     auto rank = result_tp.getRank();
     SmallVector<int64_t, 4> sub_result_shape(rank, ShapedType::kDynamicSize);
     for (int64_t i = 0; i < rank; i++) {
-      if (i != op.dimension()) {
+      if (i != op.getDimension()) {
         sub_result_shape[i] = result_tp.getDimSize(i);
       }
     }
     auto sub_result_tp =
         RankedTensorType::get(sub_result_shape, result_tp.getElementType());
     Value lhs = rewriter.create<mhlo::ConcatenateOp>(
-        loc, sub_result_tp, lhs_operands, op.dimension());
+        loc, sub_result_tp, lhs_operands, op.getDimension());
     Value rhs = rewriter.create<mhlo::ConcatenateOp>(
-        loc, sub_result_tp, rhs_operands, op.dimension());
+        loc, sub_result_tp, rhs_operands, op.getDimension());
 
     SmallVector<Value, 2> fused_operands{lhs, rhs};
     auto fused_op = rewriter.create<mhlo::ConcatenateOp>(
-        loc, result_tp, fused_operands, op.dimension());
+        loc, result_tp, fused_operands, op.getDimension());
 
     rewriter.replaceOp(op, fused_op.getOperation()->getResults());
     return success();
diff --git a/tao_compiler/mlir/xla/ral/context/init_stream_executor.cc b/tao_compiler/mlir/xla/ral/context/init_stream_executor.cc
index c88d251ec48..3eaa6455394 100644
--- a/tao_compiler/mlir/xla/ral/context/init_stream_executor.cc
+++ b/tao_compiler/mlir/xla/ral/context/init_stream_executor.cc
@@ -62,24 +62,30 @@ se::Stream* GetOrCreateDefaultCudaStreamExecutorStream(int device_ordinal,
   auto it = table.find(key);
   if (it == table.end()) {
 #if TENSORFLOW_USE_ROCM
-    auto platform =
-        se::MultiPlatformManager::PlatformWithName("ROCM").ValueOrDie();
+    auto status_or = se::MultiPlatformManager::PlatformWithName("ROCM");
+    assert(status_or.ok());
+    auto platform = status_or.value();
     auto gpu_platform = static_cast<se::gpu::ROCmPlatform*>(platform);
 #else
-    auto platform =
-        se::MultiPlatformManager::PlatformWithName("CUDA").ValueOrDie();
+    auto status_or = se::MultiPlatformManager::PlatformWithName("CUDA");
+    assert(status_or.ok());
+    auto platform = status_or.value();
     auto gpu_platform = static_cast<se::gpu::CudaPlatform*>(platform);
 #endif
     TableEntry e;
     bool use_multi_cuda_se = true;
-    tensorflow::ReadBoolFromEnvVar("TAO_ENABLE_MULTIPLE_CUDA_STREAM_EXECUTOR",
-                                   true, &use_multi_cuda_se);
+    tsl::ReadBoolFromEnvVar("TAO_ENABLE_MULTIPLE_CUDA_STREAM_EXECUTOR", true,
+                            &use_multi_cuda_se);
     se::StreamExecutor* executor = nullptr;
     if (use_multi_cuda_se) {
-      executor = gpu_platform->ExecutorForDevice(device_ordinal, (void*)stream)
-                     .ValueOrDie();
+      auto status_or =
+          gpu_platform->ExecutorForDevice(device_ordinal, (void*)stream);
+      assert(status_or.ok());
+      executor = status_or.value();
     } else {
-      executor = gpu_platform->ExecutorForDevice(device_ordinal).ValueOrDie();
+      auto status_or = gpu_platform->ExecutorForDevice(device_ordinal);
+      assert(status_or.ok());
+      executor = status_or.value();
     }
     auto cuda_executor = se::gpu::ExtractGpuExecutor(executor);
     // Not wrapper with a unique_ptr to avoid stream is destroyed after executor
diff --git a/tf_community b/tf_community
index 173d33c7c7a..dc054b6b4c5 160000
--- a/tf_community
+++ b/tf_community
@@ -1,1 +1,1 @@
-Subproject commit 173d33c7c7aa8e3e9b352af2433f76ab8da2708e
+Subproject commit dc054b6b4c564e6deddde6279712f6135b2b3763


Test patch:
diff --git a/pytorch_blade/tests/mhlo/dropout.mlir b/pytorch_blade/tests/mhlo/dropout.mlir
index ce26a164c43..52a64051fd8 100644
--- a/pytorch_blade/tests/mhlo/dropout.mlir
+++ b/pytorch_blade/tests/mhlo/dropout.mlir
@@ -19,9 +19,9 @@
 // CHECK:         %[[T12:.*]] = tensor.dim %[[T9]], %[[C1]] : tensor<?x?xf32>
 // CHECK:         %[[T13:.*]] = arith.index_cast %[[T12]] : index to i32
 // CHECK:         %[[T14:.*]] = tensor.from_elements %[[T11]], %[[T13]] : tensor<2xi32>
-// CHECK:         %[[T15:.*]] = "mhlo.rng"(%[[T1]], %[[T0]], %[[T1]]4) {rng_distribution = #mhlo.rng_distribution<UNIFORM>} : (tensor<f32>, tensor<f32>, tensor<2xi32>) -> tensor<?x?xf32>
+// CHECK:         %[[T15:.*]] = "mhlo.rng"(%[[T1]], %[[T0]], %[[T14]]) {rng_distribution = #mhlo.rng_distribution<UNIFORM>} : (tensor<f32>, tensor<f32>, tensor<2xi32>) -> tensor<?x?xf32>
 // CHECK:         %[[T16:.*]] = chlo.broadcast_compare %[[T15]], %[[T1]] {compare_type = #chlo<comparison_type FLOAT>, comparison_direction = #chlo<comparison_direction LT>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xi1>
-// CHECK:         %[[T17:.*]] = mhlo.convert(%[[T16]]) : (tensor<?x?xi1>) -> tensor<?x?xf32>
+// CHECK:         %[[T17:.*]] = mhlo.convert %[[T16]] : (tensor<?x?xi1>) -> tensor<?x?xf32>
 // CHECK:         %[[T18:.*]] = chlo.broadcast_multiply %[[T17]], %[[ARG0]] : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
 // CHECK:         %[[T19:.*]] = chlo.broadcast_multiply %[[T18]], %[[T1]] : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xf32>
 // CHECK:         %[[T20:.*]] = chlo.broadcast_compare %[[T17]], %[[T0]] {compare_type = #chlo<comparison_type FLOAT>, comparison_direction = #chlo<comparison_direction GE>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xi1>
@@ -38,7 +38,7 @@ func.func @torch.aten.native_dropout.train(%arg0: !torch.vtensor<[?,?],f32>) ->
 // CHECK-LABEL:  func.func @torch.aten.native_dropout_backward(
 // CHECK-SAME:         %[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xi1>) -> tensor<?x?xf32> {
 // CHECK:         %[[T0:.*]] = mhlo.constant dense<1.000000e+00> : tensor<f32>
-// CHECK:         %[[T1:.*]] = mhlo.convert(%[[ARG1]]) : (tensor<?x?xi1>) -> tensor<?x?xf32>
+// CHECK:         %[[T1:.*]] = mhlo.convert %[[ARG1]] : (tensor<?x?xi1>) -> tensor<?x?xf32>
 // CHECK:         %[[T2:.*]] = chlo.broadcast_multiply %[[ARG0]], %[[T1]] : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
 // CHECK:         %[[T3:.*]] = chlo.broadcast_multiply %[[T2]], %[[T0]] : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xf32>
 // CHECK:         return %[[T3]] : tensor<?x?xf32>
diff --git a/pytorch_blade/tests/mhlo/elementwise.mlir b/pytorch_blade/tests/mhlo/elementwise.mlir
index 5a73ad84510..da8978af5bf 100644
--- a/pytorch_blade/tests/mhlo/elementwise.mlir
+++ b/pytorch_blade/tests/mhlo/elementwise.mlir
@@ -62,7 +62,7 @@ func.func @torch.aten.sub.tensor(%arg0: !torch.vtensor<[?,?,?,?],f32>, %arg1: !t
 // CHECK-LABEL:  func.func @torch.aten.sub.scalar.int(
 // CHECK-SAME:         %[[ARG0:.*]]: tensor<?x?x?x4xf32>, %[[ARG1:.*]]: i64) -> tensor<?x?x?x4xf32> {
 // CHECK:         %[[T0:.*]] = tensor.from_elements %[[ARG1]] : tensor<1xi64>
-// CHECK:         %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<1xi64>) -> tensor<1xf32>
+// CHECK:         %[[T1:.*]] = mhlo.convert %[[T0]] : (tensor<1xi64>) -> tensor<1xf32>
 // CHECK:         %[[T2:.*]] = mhlo.reshape %[[T1]] : (tensor<1xf32>) -> tensor<f32>
 // CHECK:         %[[T3:.*]] = chlo.broadcast_subtract %[[ARG0]], %[[T2]] : (tensor<?x?x?x4xf32>, tensor<f32>) -> tensor<?x?x?x4xf32>
 // CHECK:         return %[[T3]] : tensor<?x?x?x4xf32>
@@ -77,7 +77,7 @@ func.func @torch.aten.sub.scalar.int(%arg0: !torch.vtensor<[?,?,?,4],f32>, %arg1
 // CHECK-LABEL:  func.func @torch.aten.add.scalar.float(
 // CHECK-SAME:         %[[ARG0:.*]]: tensor<?x?x?x4xf32>, %[[ARG1:.*]]: f64) -> tensor<?x?x?x4xf32> {
 // CHECK:         %[[T0:.*]] = tensor.from_elements %[[ARG1]] : tensor<1xf64>
-// CHECK:         %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<1xf64>) -> tensor<1xf32>
+// CHECK:         %[[T1:.*]] = mhlo.convert %[[T0]] : (tensor<1xf64>) -> tensor<1xf32>
 // CHECK:         %[[T2:.*]] = mhlo.reshape %[[T1]] : (tensor<1xf32>) -> tensor<f32>
 // CHECK:         %[[T3:.*]] = chlo.broadcast_add %[[ARG0]], %[[T2]] : (tensor<?x?x?x4xf32>, tensor<f32>) -> tensor<?x?x?x4xf32>
 // CHECK:         return %[[T3]] : tensor<?x?x?x4xf32>
@@ -108,7 +108,7 @@ func.func @torch.aten.div.Tensor_mode(%arg0: !torch.vtensor<[?,?,?,?],f32>, %arg
 // CHECK-LABEL:  func.func @torch.aten.ne.scalar(
 // CHECK-SAME:         %[[ARG0:.*]]: tensor<?x?xi32>, %[[ARG1:.*]]: f64) -> tensor<?x?xi1> {
 // CHECK:         %[[T0:.*]] = tensor.from_elements %[[ARG1]] : tensor<1xf64>
-// CHECK:         %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<1xf64>) -> tensor<1xi32>
+// CHECK:         %[[T1:.*]] = mhlo.convert %[[T0]] : (tensor<1xf64>) -> tensor<1xi32>
 // CHECK:         %[[T2:.*]] = mhlo.reshape %[[T1]] : (tensor<1xi32>) -> tensor<i32>
 // CHECK:         %[[T3:.*]] = chlo.broadcast_compare %[[ARG0]], %[[T2]] {compare_type = #chlo<comparison_type SIGNED>, comparison_direction = #chlo<comparison_direction NE>} : (tensor<?x?xi32>, tensor<i32>) -> tensor<?x?xi1>
 // CHECK:         return %[[T3]] : tensor<?x?xi1>
@@ -178,9 +178,9 @@ func.func @torch.aten.gelu_backward(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !to
 // CHECK:         %[[T12:.*]] = tensor.dim %[[T9]], %[[C1]] : tensor<?x?xf32>
 // CHECK:         %[[T13:.*]] = arith.index_cast %[[T12]] : index to i32
 // CHECK:         %[[T14:.*]] = tensor.from_elements %[[T11]], %[[T13]] : tensor<2xi32>
-// CHECK:         %[[T15:.*]] = "mhlo.rng"(%[[T1]], %[[T0]], %[[T1]]4) {rng_distribution = #mhlo.rng_distribution<UNIFORM>} : (tensor<f32>, tensor<f32>, tensor<2xi32>) -> tensor<?x?xf32>
+// CHECK:         %[[T15:.*]] = "mhlo.rng"(%[[T1]], %[[T0]], %[[T14]]) {rng_distribution = #mhlo.rng_distribution<UNIFORM>} : (tensor<f32>, tensor<f32>, tensor<2xi32>) -> tensor<?x?xf32>
 // CHECK:         %[[T16:.*]] = chlo.broadcast_compare %[[T15]], %[[T1]] {compare_type = #chlo<comparison_type FLOAT>, comparison_direction = #chlo<comparison_direction LT>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xi1>
-// CHECK:         %[[T17:.*]] = mhlo.convert(%[[T16]]) : (tensor<?x?xi1>) -> tensor<?x?xf32>
+// CHECK:         %[[T17:.*]] = mhlo.convert %[[T16]] : (tensor<?x?xi1>) -> tensor<?x?xf32>
 // CHECK:         %[[T18:.*]] = chlo.broadcast_multiply %[[T17]], %[[ARG0]] : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
 // CHECK:         %[[T19:.*]] = chlo.broadcast_multiply %[[T18]], %[[T1]] : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xf32>
 // CHECK:         %[[T20:.*]] = chlo.broadcast_compare %[[T17]], %[[T0]] {compare_type = #chlo<comparison_type FLOAT>, comparison_direction = #chlo<comparison_direction GE>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xi1>
diff --git a/pytorch_blade/tests/mhlo/layer_norm.mlir b/pytorch_blade/tests/mhlo/layer_norm.mlir
index e93444066c9..ee11e9d6d4a 100644
--- a/pytorch_blade/tests/mhlo/layer_norm.mlir
+++ b/pytorch_blade/tests/mhlo/layer_norm.mlir
@@ -89,7 +89,7 @@ func.func @torch.aten.layer_norm.affine(%arg0: !torch.vtensor<[2,3,224,224],f32>
 // CHECK:         %[[T12:.*]] = arith.index_cast %[[T11]] : index to i64
 // CHECK:         %[[T13:.*]] = arith.muli %[[T10]], %[[T12]] : i64
 // CHECK:         %[[T14:.*]] = tensor.from_elements %[[T13]] : tensor<1xi64>
-// CHECK:         %[[T15:.*]] = mhlo.convert(%[[T14]]) : (tensor<1xi64>) -> tensor<1xf32>
+// CHECK:         %[[T15:.*]] = mhlo.convert %[[T14]] : (tensor<1xi64>) -> tensor<1xf32>
 // CHECK:         %[[T16:.*]] = mhlo.reshape %[[T15]] : (tensor<1xf32>) -> tensor<f32>
 // CHECK:         %[[T17:.*]] = chlo.broadcast_divide %[[T8]], %[[T16]] : (tensor<?x?x1x1xf32>, tensor<f32>) -> tensor<?x?x1x1xf32>
 // CHECK:         %[[T18:.*]] = arith.index_cast %[[T3]] : index to i64
@@ -115,7 +115,7 @@ func.func @torch.aten.layer_norm.affine(%arg0: !torch.vtensor<[2,3,224,224],f32>
 // CHECK:         %[[T38:.*]] = arith.index_cast %[[T37]] : index to i64
 // CHECK:         %[[T39:.*]] = arith.muli %[[T36]], %[[T38]] : i64
 // CHECK:         %[[T40:.*]] = tensor.from_elements %[[T39]] : tensor<1xi64>
-// CHECK:         %[[T41:.*]] = mhlo.convert(%[[T40]]) : (tensor<1xi64>) -> tensor<1xf32>
+// CHECK:         %[[T41:.*]] = mhlo.convert %[[T40]] : (tensor<1xi64>) -> tensor<1xf32>
 // CHECK:         %[[T42:.*]] = mhlo.reshape %[[T41]] : (tensor<1xf32>) -> tensor<f32>
 // CHECK:         %[[T43:.*]] = chlo.broadcast_divide %[[T34]], %[[T42]] : (tensor<?x?x1x1xf32>, tensor<f32>) -> tensor<?x?x1x1xf32>
 // CHECK:         %[[T44:.*]] = chlo.broadcast_add %[[T43]], %[[T0]] : (tensor<?x?x1x1xf32>, tensor<f32>) -> tensor<?x?x1x1xf32>
@@ -266,7 +266,7 @@ func.func @torch.aten.native_layer_norm.affine(%arg0: !torch.vtensor<[2,3,224,22
 // CHECK:         %[[T7:.*]] = arith.index_cast %[[T6]] : index to i64
 // CHECK:         %[[T8:.*]] = arith.muli %[[T5]], %[[T7]] : i64
 // CHECK:         %[[T9:.*]] = tensor.from_elements %[[T8]] : tensor<1xi64>
-// CHECK:         %[[T10:.*]] = mhlo.convert(%[[T9]]) : (tensor<1xi64>) -> tensor<1xf32>
+// CHECK:         %[[T10:.*]] = mhlo.convert %[[T9]] : (tensor<1xi64>) -> tensor<1xf32>
 // CHECK:         %[[T11:.*]] = mhlo.reshape %[[T10]] : (tensor<1xf32>) -> tensor<f32>
 // CHECK:         %[[T12:.*]] = chlo.broadcast_divide %[[T3]], %[[T11]] : (tensor<2x3x1x1xf32>, tensor<f32>) -> tensor<2x3x1x1xf32>
 // CHECK:         %[[T13:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?x?x?xf32>
@@ -289,14 +289,14 @@ func.func @torch.aten.native_layer_norm.affine(%arg0: !torch.vtensor<[2,3,224,22
 // CHECK:         %[[T30:.*]] = arith.index_cast %[[T29]] : index to i64
 // CHECK:         %[[T31:.*]] = arith.muli %[[T28]], %[[T30]] : i64
 // CHECK:         %[[T32:.*]] = tensor.from_elements %[[T31]] : tensor<1xi64>
-// CHECK:         %[[T33:.*]] = mhlo.convert(%[[T32]]) : (tensor<1xi64>) -> tensor<1xf32>
+// CHECK:         %[[T33:.*]] = mhlo.convert %[[T32]] : (tensor<1xi64>) -> tensor<1xf32>
 // CHECK:         %[[T34:.*]] = mhlo.reshape %[[T33]] : (tensor<1xf32>) -> tensor<f32>
 // CHECK:         %[[T35:.*]] = chlo.broadcast_divide %[[T26]], %[[T34]] : (tensor<2x3x1x1xf32>, tensor<f32>) -> tensor<2x3x1x1xf32>
 // CHECK:         %[[T36:.*]] = chlo.broadcast_add %[[T35]], %[[T0]] : (tensor<2x3x1x1xf32>, tensor<f32>) -> tensor<2x3x1x1xf32>
 // CHECK:         %[[T37:.*]] = mhlo.rsqrt %[[T36]] : tensor<2x3x1x1xf32>
 // CHECK:         %[[T38:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[T37]], %[[T21]]) {broadcast_dimensions = dense<[0, 1, 2, 3]> : tensor<4xi64>} : (tensor<2x3x1x1xf32>, tensor<4xi32>) -> tensor<?x?x?x?xf32>
 // CHECK:         %[[T39:.*]] = chlo.broadcast_multiply %[[T23]], %[[T38]] : (tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
-// CHECK:         %[[T40:.*]] = mhlo.convert(%[[T39]]) : (tensor<?x?x?x?xf32>) -> tensor<2x3x224x224xf32>
+// CHECK:         %[[T40:.*]] = mhlo.convert %[[T39]] : (tensor<?x?x?x?xf32>) -> tensor<2x3x224x224xf32>
 // CHECK:         return %[[T40]], %[[T12]], %[[T37]] : tensor<2x3x224x224xf32>, tensor<2x3x1x1xf32>, tensor<2x3x1x1xf32>
 func.func @torch.aten.native_layer_norm.dynamic_shape.full(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> (!torch.vtensor<[2,3,224,224],f32>, !torch.vtensor<[2,3,1,1],f32>, !torch.vtensor<[2,3,1,1],f32>) {
   %int224 = torch.constant.int 224
@@ -383,7 +383,7 @@ func.func @torch.aten.native_layer_norm.dynamic_shape.partial(%arg0: !torch.vten
 // CHECK:         %[[T16:.*]] = tensor.dim %[[T3]], %[[C2]] : tensor<?x?x?xf32>
 // CHECK:         %[[T17:.*]] = arith.index_cast %[[T16]] : index to i64
 // CHECK:         %[[T18:.*]] = tensor.from_elements %[[T17]] : tensor<1xi64>
-// CHECK:         %[[T19:.*]] = mhlo.convert(%[[T18]]) : (tensor<1xi64>) -> tensor<1xf32>
+// CHECK:         %[[T19:.*]] = mhlo.convert %[[T18]] : (tensor<1xi64>) -> tensor<1xf32>
 // CHECK:         %[[T20:.*]] = mhlo.reshape %[[T19]] : (tensor<1xf32>) -> tensor<f32>
 // CHECK:         %[[T21:.*]] = chlo.broadcast_divide %[[T15]], %[[T20]] : (tensor<?x?x1xf32>, tensor<f32>) -> tensor<?x?x1xf32>
 // CHECK:         %[[T22:.*]] = chlo.broadcast_multiply %[[T3]], %[[T2]] : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
@@ -397,7 +397,7 @@ func.func @torch.aten.native_layer_norm.dynamic_shape.partial(%arg0: !torch.vten
 // CHECK:         %[[T30:.*]] = tensor.dim %[[T22]], %[[C2]] : tensor<?x?x?xf32>
 // CHECK:         %[[T31:.*]] = arith.index_cast %[[T30]] : index to i64
 // CHECK:         %[[T32:.*]] = tensor.from_elements %[[T31]] : tensor<1xi64>
-// CHECK:         %[[T33:.*]] = mhlo.convert(%[[T32]]) : (tensor<1xi64>) -> tensor<1xf32>
+// CHECK:         %[[T33:.*]] = mhlo.convert %[[T32]] : (tensor<1xi64>) -> tensor<1xf32>
 // CHECK:         %[[T34:.*]] = mhlo.reshape %[[T33]] : (tensor<1xf32>) -> tensor<f32>
 // CHECK:         %[[T35:.*]] = chlo.broadcast_divide %[[T29]], %[[T34]] : (tensor<?x?x1xf32>, tensor<f32>) -> tensor<?x?x1xf32>
 // CHECK:         %[[T36:.*]] = chlo.broadcast_multiply %[[T2]], %[[T35]] : (tensor<?x?x?xf32>, tensor<?x?x1xf32>) -> tensor<?x?x?xf32>
diff --git a/pytorch_blade/tests/mhlo/loss.mlir b/pytorch_blade/tests/mhlo/loss.mlir
index 5afdb89cfd7..739ba8edfb8 100644
--- a/pytorch_blade/tests/mhlo/loss.mlir
+++ b/pytorch_blade/tests/mhlo/loss.mlir
@@ -9,7 +9,7 @@
 // CHECK:         %[[T2:.*]] = arith.index_cast %[[T1]] : index to i64
 // CHECK:         %[[T3:.*]] = arith.sitofp %[[T2]] : i64 to f64
 // CHECK:         %[[T4:.*]] = tensor.from_elements %[[T3]] : tensor<1xf64>
-// CHECK:         %[[T5:.*]] = mhlo.convert(%[[T4]]) : (tensor<1xf64>) -> tensor<1xf32>
+// CHECK:         %[[T5:.*]] = mhlo.convert %[[T4]] : (tensor<1xf64>) -> tensor<1xf32>
 // CHECK:         %[[T6:.*]] = mhlo.reshape %[[T5]] : (tensor<1xf32>) -> tensor<f32>
 // CHECK:         %[[T7:.*]] = arith.index_cast %[[T1]] : index to i32
 // CHECK:         %[[T8:.*]] = tensor.from_elements %[[T7]], %[[C1_I32]] : tensor<2xi32>
diff --git a/pytorch_blade/tests/mhlo/reduction.mlir b/pytorch_blade/tests/mhlo/reduction.mlir
index cae086a2ff3..f592c1f9850 100644
--- a/pytorch_blade/tests/mhlo/reduction.mlir
+++ b/pytorch_blade/tests/mhlo/reduction.mlir
@@ -21,9 +21,9 @@
 // CHECK:         %[[T12:.*]] = arith.muli %[[T9]], %[[T11]] : i32
 // CHECK:         %[[T13:.*]] = arith.extsi %[[T12]] : i32 to i64
 // CHECK:         %[[T14:.*]] = tensor.from_elements %[[T13]] : tensor<1xi64>
-// CHECK:         %[[T15:.*]] = mhlo.convert(%[[T14]]) : (tensor<1xi64>) -> tensor<1xf32>
+// CHECK:         %[[T15:.*]] = mhlo.convert %[[T14]] : (tensor<1xi64>) -> tensor<1xf32>
 // CHECK:         %[[T16:.*]] = mhlo.reshape %[[T15]] : (tensor<1xf32>) -> tensor<f32>
-// CHECK:         %[[T17:.*]] = chlo.broadcast_divide %[[T1]], %[[T1]]6 : (tensor<f32>, tensor<f32>) -> tensor<f32>
+// CHECK:         %[[T17:.*]] = chlo.broadcast_divide %[[T1]], %[[T16]] : (tensor<f32>, tensor<f32>) -> tensor<f32>
 // CHECK:         return %[[T17]] : tensor<f32>
 func.func @torch.aten.sum.div.Scalar(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[],f32> {
   %int6 = torch.constant.int 6
@@ -42,7 +42,7 @@ func.func @torch.aten.sum.div.Scalar(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !t
 // CHECK:         %[[C1:.*]] = arith.constant 1 : index
 // CHECK:         %[[C0:.*]] = arith.constant 0 : index
 // CHECK:         %[[T0:.*]] = mhlo.constant dense<0.000000e+00> : tensor<f32>
-// CHECK:         %[[T1:.*]] = mhlo.convert(%[[ARG0]]) : (tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xf32>
+// CHECK:         %[[T1:.*]] = mhlo.convert %[[ARG0]] : (tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xf32>
 // CHECK:         %[[T2:.*]] = mhlo.reduce(%[[T1]] init: %[[T0]]) applies mhlo.add across dimensions = [0, 1, 2, 3] : (tensor<?x?x?x?xf32>, tensor<f32>) -> tensor<f32>
 // CHECK:         %[[T3:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?x?x?xi32>
 // CHECK:         %[[T4:.*]] = arith.index_cast %[[T3]] : index to i32
@@ -57,7 +57,7 @@ func.func @torch.aten.sum.div.Scalar(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !t
 // CHECK:         %[[T13:.*]] = arith.muli %[[T10]], %[[T12]] : i32
 // CHECK:         %[[T14:.*]] = arith.extsi %[[T13]] : i32 to i64
 // CHECK:         %[[T15:.*]] = tensor.from_elements %[[T14]] : tensor<1xi64>
-// CHECK:         %[[T16:.*]] = mhlo.convert(%[[T15]]) : (tensor<1xi64>) -> tensor<1xf32>
+// CHECK:         %[[T16:.*]] = mhlo.convert %[[T15]] : (tensor<1xi64>) -> tensor<1xf32>
 // CHECK:         %[[T17:.*]] = mhlo.reshape %[[T16]] : (tensor<1xf32>) -> tensor<f32>
 // CHECK:         %[[T18:.*]] = chlo.broadcast_divide %[[T2]], %[[T17]] : (tensor<f32>, tensor<f32>) -> tensor<f32>
 // CHECK:         return %[[T18]] : tensor<f32>
@@ -87,7 +87,7 @@ func.func @torch.aten.sum.outf32(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !torch
 // CHECK-LABEL:  func.func @torch.aten.sum.outf64(
 // CHECK-SAME:         %[[ARG0:.*]]: tensor<?x?x?x?xf32>) -> tensor<f64> {
 // CHECK:         %[[T0:.*]] = mhlo.constant dense<0.000000e+00> : tensor<f64>
-// CHECK:         %[[T1:.*]] = mhlo.convert(%[[ARG0]]) : (tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf64>
+// CHECK:         %[[T1:.*]] = mhlo.convert %[[ARG0]] : (tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf64>
 // CHECK:         %[[T2:.*]] = mhlo.reduce(%[[T1]] init: %[[T0]]) applies mhlo.add across dimensions = [0, 1, 2, 3] : (tensor<?x?x?x?xf64>, tensor<f64>) -> tensor<f64>
 // CHECK:         return %[[T2]] : tensor<f64>
 func.func @torch.aten.sum.outf64(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[],f64> {
@@ -183,11 +183,11 @@ func.func @torch.aten.sum.dim_IntList.keepdim(%arg0: !torch.vtensor<[2,?,224,?],
 // CHECK:         %[[T5:.*]]:2 = mhlo.reduce(%[[ARG0]] init: %[[T1]]), (%[[T4]] init: %[[T0]]) across dimensions = [3] : (tensor<2x3x224x224xf32>, tensor<2x3x224x224xi32>, tensor<f32>, tensor<i32>) -> (tensor<2x3x224xf32>, tensor<2x3x224xi32>)
 // CHECK:         reducer(%[[ARG1:.*]]: tensor<f32>, %[[ARG3:.*]]: tensor<f32>) (%[[ARG2:.*]]: tensor<i32>, %[[ARG4:.*]]: tensor<i32>)  {
 // CHECK:         %[[T8:.*]] = mhlo.compare  GE, %[[ARG1]], %[[ARG3]],  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1>
-// CHECK:         %[[T9:.*]] = "mhlo.select"(%[[T8]], %[[ARG1]], %[[ARG3]]) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
+// CHECK:         %[[T9:.*]] = mhlo.select %[[T8]], %[[ARG1]], %[[ARG3]] : tensor<i1>, tensor<f32>
 // CHECK:         %[[T10:.*]] = mhlo.compare  EQ, %[[ARG1]], %[[ARG3]],  FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1>
 // CHECK:         %[[T11:.*]] = mhlo.minimum %[[ARG2]], %[[ARG4]] : tensor<i32>
-// CHECK:         %[[T12:.*]] = "mhlo.select"(%[[T8]], %[[ARG2]], %[[ARG4]]) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
-// CHECK:         %[[T13:.*]] = "mhlo.select"(%[[T10]], %[[T11]], %[[T12]]) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
+// CHECK:         %[[T12:.*]] = mhlo.select %[[T8]], %[[ARG2]], %[[ARG4]] : tensor<i1>, tensor<i32>
+// CHECK:         %[[T13:.*]] = mhlo.select %[[T10]], %[[T11]], %[[T12]] : tensor<i1>, tensor<i32>
 // CHECK:         mhlo.return %[[T9]], %[[T13]] : tensor<f32>, tensor<i32>
 // CHECK:         %[[T6:.*]] = mhlo.reshape %[[T2]] : (tensor<2x3x224xf32>) -> tensor<2x3x224x1xf32>
 // CHECK:         %[[T7:.*]] = mhlo.reshape %[[T5]]#1 : (tensor<2x3x224xi32>) -> tensor<2x3x224x1xi32>
diff --git a/pytorch_blade/tests/mhlo/softmax.mlir b/pytorch_blade/tests/mhlo/softmax.mlir
index f4f26e55366..f8d7e1673fb 100644
--- a/pytorch_blade/tests/mhlo/softmax.mlir
+++ b/pytorch_blade/tests/mhlo/softmax.mlir
@@ -20,7 +20,7 @@
 // CHECK:         %[[T12:.*]] = tensor.from_elements %[[T10]], %[[T11]] : tensor<2xi32>
 // CHECK:         %[[T13:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[T6]], %[[T12]]) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<?x1xf32>, tensor<2xi32>) -> tensor<?x?xf32>
 // CHECK:         %[[T14:.*]] = chlo.broadcast_multiply %[[ARG1]], %[[T13]] : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
-// CHECK:         %[[T15:.*]] = chlo.broadcast_subtract %[[T1]], %[[T1]]4 : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+// CHECK:         %[[T15:.*]] = chlo.broadcast_subtract %[[T1]], %[[T14]] : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
 // CHECK:         return %[[T15]] : tensor<?x?xf32>
 func.func @torch.aten._softmax_backward_data(%arg0 : !torch.vtensor<[?,?],f32>, %arg1 : !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32>{
   %int_1 = torch.constant.int 1
diff --git a/pytorch_blade/tests/mhlo/views.mlir b/pytorch_blade/tests/mhlo/views.mlir
index 974f0915b7d..295669dba70 100644
--- a/pytorch_blade/tests/mhlo/views.mlir
+++ b/pytorch_blade/tests/mhlo/views.mlir
@@ -19,7 +19,7 @@
 // CHECK:         %[[T9:.*]] = arith.index_cast %[[T8]] : index to i32
 // CHECK:         %[[T10:.*]] = arith.muli %[[T7]], %[[T9]] : i32
 // CHECK:         %[[T11:.*]] = arith.index_cast %[[T10]] : i32 to index
-// CHECK:         %[[T12:.*]] = mhlo.compute_reshape_shape %[[T11]], %[[CST]] : index, tensor<2xi32> -> tensor<2xi32>
+// CHECK:         %[[T12:.*]] = mhlo.compute_reshape_shape %[[T11]], %[[CST]] : (index, tensor<2xi32>) -> tensor<2xi32>
 // CHECK:         %[[T13:.*]] = mhlo.dynamic_reshape %[[ARG0]], %[[T12]] : (tensor<?x?x?x?xf32>, tensor<2xi32>) -> tensor<?x224xf32>
 // CHECK:         return %[[T13]] : tensor<?x224xf32>
 func.func @torch.aten.view(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[?,224],f32> {
@@ -55,7 +55,7 @@ func.func @torch.aten.view(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtens
 // CHECK:         %[[T12:.*]] = arith.index_cast %[[T11]] : index to i32
 // CHECK:         %[[T13:.*]] = arith.muli %[[T10]], %[[T12]] : i32
 // CHECK:         %[[T14:.*]] = arith.index_cast %[[T13]] : i32 to index
-// CHECK:         %[[T15:.*]] = mhlo.compute_reshape_shape %[[T14]], %[[CST]] : index, tensor<4xi32> -> tensor<4xi32>
+// CHECK:         %[[T15:.*]] = mhlo.compute_reshape_shape %[[T14]], %[[CST]] : (index, tensor<4xi32>) -> tensor<4xi32>
 // CHECK:         %[[T16:.*]] = mhlo.dynamic_reshape %[[ARG0]], %[[T15]] : (tensor<?x?x?x?x?xf32>, tensor<4xi32>) -> tensor<?x120x4x64xf32>
 // CHECK:         return %[[T16]] : tensor<?x120x4x64xf32>
 func.func @torch.aten.reshape(%arg0: !torch.vtensor<[?,?,?,?,?],f32>) -> !torch.vtensor<[?,120,4,64],f32> {
@@ -89,7 +89,7 @@ func.func @torch.aten.reshape(%arg0: !torch.vtensor<[?,?,?,?,?],f32>) -> !torch.
 // CHECK:         %[[T9:.*]] = arith.index_cast %[[T8]] : index to i32
 // CHECK:         %[[T10:.*]] = arith.muli %[[T7]], %[[T9]] : i32
 // CHECK:         %[[T11:.*]] = arith.index_cast %[[T10]] : i32 to index
-// CHECK:         %[[T12:.*]] = mhlo.compute_reshape_shape %[[T11]], %[[CST]] : index, tensor<1xi32> -> tensor<1xi32>
+// CHECK:         %[[T12:.*]] = mhlo.compute_reshape_shape %[[T11]], %[[CST]] : (index, tensor<1xi32>) -> tensor<1xi32>
 // CHECK:         %[[T13:.*]] = mhlo.dynamic_reshape %[[ARG0]], %[[T12]] : (tensor<?x?x?x?xf32>, tensor<1xi32>) -> tensor<?xf32>
 // CHECK:         return %[[T13]] : tensor<?xf32>
 func.func @torch.aten.flatten(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vtensor<[?],f32> {
@@ -114,7 +114,7 @@ func.func @torch.aten.flatten(%arg0: !torch.vtensor<[?,?,?,?],f32>) -> !torch.vt
 // CHECK:         %[[T4:.*]] = arith.index_cast %[[T3]] : index to i32
 // CHECK:         %[[T5:.*]] = arith.muli %[[T2]], %[[T4]] : i32
 // CHECK:         %[[T6:.*]] = arith.index_cast %[[T5]] : i32 to index
-// CHECK:         %[[T7:.*]] = mhlo.compute_reshape_shape %[[T6]], %[[CST]] : index, tensor<3xi32> -> tensor<3xi32>
+// CHECK:         %[[T7:.*]] = mhlo.compute_reshape_shape %[[T6]], %[[CST]] : (index, tensor<3xi32>) -> tensor<3xi32>
 // CHECK:         %[[T8:.*]] = mhlo.dynamic_reshape %[[ARG0]], %[[T7]] : (tensor<2x3x?x?xf32>, tensor<3xi32>) -> tensor<2x3x?xf32>
 // CHECK:         return %[[T8]] : tensor<2x3x?xf32>
 func.func @torch.aten.view.minus1(%arg0: !torch.vtensor<[2,3,?,?],f32>) -> !torch.vtensor<[2,3,?],f32> {
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/shape_scalar_i32.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/shape_scalar_i32.mlir
index 139ee187840..2a1bcba0f35 100644
--- a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/shape_scalar_i32.mlir
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/shape_scalar_i32.mlir
@@ -1,9 +1,9 @@
 module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
-  func.func @main(%arg0: tensor<i32>) -> (tensor<1xi32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+  func.func @main(%arg0: tensor<i32>) -> (tensor<0xi32>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
     %graph = tf_executor.graph {
-      %0:2 = tf_executor.island wraps "tf.Shape"(%arg0) {device = ""} : (tensor<i32>) -> (tensor<1xi32>)
-      tf_executor.fetch %0 : tensor<1xi32>
+      %0:2 = tf_executor.island wraps "tf.Shape"(%arg0) {device = ""} : (tensor<i32>) -> (tensor<0xi32>)
+      tf_executor.fetch %0 : tensor<0xi32>
     }
-    return %graph : tensor<1xi32>
+    return %graph : tensor<0xi32>
   }
 }
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/shape_scalar_i64.mlir b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/shape_scalar_i64.mlir
index a1b8329aaf9..2e250db1c98 100644
--- a/tao_compiler/mlir/disc/tests/tensorflow_ops/data/shape_scalar_i64.mlir
+++ b/tao_compiler/mlir/disc/tests/tensorflow_ops/data/shape_scalar_i64.mlir
@@ -1,9 +1,9 @@
 module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
-  func.func @main(%arg0: tensor<i32>) -> (tensor<1xi64>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
+  func.func @main(%arg0: tensor<i32>) -> (tensor<0xi64>) attributes {tf.entry_function = {inputs = "{{INPUTS}}", outputs = "{{OUTPUTS}}", input_placements="{{INPUT_PLACEMENTS}}", output_placements="{{OUTPUT_PLACEMENTS}}"}} {
     %graph = tf_executor.graph {
-      %0:2 = tf_executor.island wraps "tf.Shape"(%arg0) {device = ""} : (tensor<i32>) -> (tensor<1xi64>)
-      tf_executor.fetch %0 : tensor<1xi64>
+      %0:2 = tf_executor.island wraps "tf.Shape"(%arg0) {device = ""} : (tensor<i32>) -> (tensor<0xi64>)
+      tf_executor.fetch %0 : tensor<0xi64>
     }
-    return %graph : tensor<1xi64>
+    return %graph : tensor<0xi64>
   }
 }
\ No newline at end of file
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir
index 81af7e6a8ea..003cb2a883a 100644
--- a/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-lower-quantize-and-dequantize.mlir
@@ -18,7 +18,7 @@ func.func @quantize_signed_per_tensor(
   // CHECK-SAME: (%[[INPUT_SCALE]], %[[INPUT_SHAPE]])
   // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
 
-  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert(%[[INPUT_ZERO_POINT]]) : (tensor<i32>) -> tensor<f32>
+  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert %[[INPUT_ZERO_POINT]] : (tensor<i32>) -> tensor<f32>
   // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
   // CHECK-SAME: (%[[CASTED_ZERO_POINT]], %[[INPUT_SHAPE]])
   // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
@@ -35,7 +35,7 @@ func.func @quantize_signed_per_tensor(
   // CHECK: %[[T1:.*]] = mhlo.add %[[T0]], %[[BCAST_ZERO_POINT]]
   // CHECK: %[[T2:.*]] = mhlo.round_nearest_even %[[T1]]
   // CHECK: %[[T3:.*]] = mhlo.clamp %[[T2]], %[[BCAST_MIN]], %[[BCAST_MAX]]
-  // CHECK: %[[OUT:.*]] = mhlo.convert(%[[T3]]) : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xi8>
+  // CHECK: %[[OUT:.*]] = mhlo.convert %[[T3]] : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xi8>
   // CHECK: return %[[OUT]]
 
   %quantized_input = "mhlo_disc.quantize"(%input, %input_scale, %input_zero_point) {
@@ -69,7 +69,7 @@ func.func @quantize_unsigned_per_tensor(
   // CHECK-SAME: (%[[INPUT_SCALE]], %[[INPUT_SHAPE]])
   // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
 
-  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert(%[[INPUT_ZERO_POINT]]) : (tensor<i32>) -> tensor<f32>
+  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert %[[INPUT_ZERO_POINT]] : (tensor<i32>) -> tensor<f32>
   // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
   // CHECK-SAME: (%[[CASTED_ZERO_POINT]], %[[INPUT_SHAPE]])
   // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
@@ -86,7 +86,7 @@ func.func @quantize_unsigned_per_tensor(
   // CHECK: %[[T1:.*]] = mhlo.add %[[T0]], %[[BCAST_ZERO_POINT]]
   // CHECK: %[[T2:.*]] = mhlo.round_nearest_afz %[[T1]]
   // CHECK: %[[T3:.*]] = mhlo.clamp %[[T2]], %[[BCAST_MIN]], %[[BCAST_MAX]]
-  // CHECK: %[[OUT:.*]] = mhlo.convert(%[[T3]]) : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xui8>
+  // CHECK: %[[OUT:.*]] = mhlo.convert %[[T3]] : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xui8>
   // CHECK: return %[[OUT]]
 
   %quantized_input = "mhlo_disc.quantize"(%input, %input_scale, %input_zero_point) {
@@ -119,7 +119,7 @@ func.func @quantize_signed_per_channel(
   // CHECK-SAME: (%[[INPUT_SCALE]], %[[INPUT_SHAPE]])
   // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
 
-  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert(%[[INPUT_ZERO_POINT]]) : (tensor<?xi32>) -> tensor<?xf32>
+  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert %[[INPUT_ZERO_POINT]] : (tensor<?xi32>) -> tensor<?xf32>
   // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
   // CHECK-SAME: (%[[CASTED_ZERO_POINT]], %[[INPUT_SHAPE]])
   // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
@@ -136,7 +136,7 @@ func.func @quantize_signed_per_channel(
   // CHECK: %[[T1:.*]] = mhlo.add %[[T0]], %[[BCAST_ZERO_POINT]]
   // CHECK: %[[T2:.*]] = mhlo.round_nearest_afz %[[T1]]
   // CHECK: %[[T3:.*]] = mhlo.clamp %[[T2]], %[[BCAST_MIN]], %[[BCAST_MAX]]
-  // CHECK: %[[OUT:.*]] = mhlo.convert(%[[T3]]) : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xi8>
+  // CHECK: %[[OUT:.*]] = mhlo.convert %[[T3]] : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xi8>
   // CHECK: return %[[OUT]]
 
   %quantized_input = "mhlo_disc.quantize"(%input, %input_scale, %input_zero_point) {
@@ -169,7 +169,7 @@ func.func @quantize_unsigned_per_channel(
   // CHECK-SAME: (%[[INPUT_SCALE]], %[[INPUT_SHAPE]])
   // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
 
-  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert(%[[INPUT_ZERO_POINT]]) : (tensor<?xi32>) -> tensor<?xf32>
+  // CHECK: %[[CASTED_ZERO_POINT:.*]] = mhlo.convert %[[INPUT_ZERO_POINT]] : (tensor<?xi32>) -> tensor<?xf32>
   // CHECK: %[[BCAST_ZERO_POINT:.*]] = "mhlo.dynamic_broadcast_in_dim"
   // CHECK-SAME: (%[[CASTED_ZERO_POINT]], %[[INPUT_SHAPE]])
   // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
@@ -186,7 +186,7 @@ func.func @quantize_unsigned_per_channel(
   // CHECK: %[[T1:.*]] = mhlo.add %[[T0]], %[[BCAST_ZERO_POINT]]
   // CHECK: %[[T2:.*]] = mhlo.round_nearest_even %[[T1]]
   // CHECK: %[[T3:.*]] = mhlo.clamp %[[T2]], %[[BCAST_MIN]], %[[BCAST_MAX]]
-  // CHECK: %[[OUT:.*]] = mhlo.convert(%[[T3]]) : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xui8>
+  // CHECK: %[[OUT:.*]] = mhlo.convert %[[T3]] : (tensor<?x32x32x6xf32>) -> tensor<?x32x32x6xui8>
   // CHECK: return %[[OUT]]
 
   %quantized_input = "mhlo_disc.quantize"(%input, %input_scale, %input_zero_point) {
@@ -220,9 +220,9 @@ func.func @dequantize_signed_per_tensor(
   // CHECK-SAME: %[[INPUT_ZERO_POINT]], %[[SHAPE]]
   // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
 
-  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert(%[[INPUT]]) : (tensor<?x32x32x6xi8>) -> tensor<?x32x32x6xi32>
+  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert %[[INPUT]] : (tensor<?x32x32x6xi8>) -> tensor<?x32x32x6xi32>
   // CHECK: %[[T0:.*]] = mhlo.subtract %[[CASTED_INPUT]], %[[BCAST_ZERO_POINT]] : tensor<?x32x32x6xi32>
-  // CHECK: %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
+  // CHECK: %[[T1:.*]] = mhlo.convert %[[T0]] : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
   // CHECK: %[[T2:.*]] = mhlo.multiply %[[T1]], %[[BCAST_SCALE]] : tensor<?x32x32x6xf32>
   // CHECK: return %[[T2]]
   %quantized_input = "mhlo_disc.dequantize"(%input, %input_scale, %input_zero_point) {
@@ -253,9 +253,9 @@ func.func @dequantize_unsigned_per_tensor(
   // CHECK-SAME: %[[INPUT_ZERO_POINT]], %[[SHAPE]]
   // CHECK-SAME: broadcast_dimensions = dense<> : tensor<0xi64>
 
-  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert(%[[INPUT]]) : (tensor<?x32x32x6xui8>) -> tensor<?x32x32x6xi32>
+  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert %[[INPUT]] : (tensor<?x32x32x6xui8>) -> tensor<?x32x32x6xi32>
   // CHECK: %[[T0:.*]] = mhlo.subtract %[[CASTED_INPUT]], %[[BCAST_ZERO_POINT]] : tensor<?x32x32x6xi32>
-  // CHECK: %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
+  // CHECK: %[[T1:.*]] = mhlo.convert %[[T0]] : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
   // CHECK: %[[T2:.*]] = mhlo.multiply %[[T1]], %[[BCAST_SCALE]] : tensor<?x32x32x6xf32>
   // CHECK: return %[[T2]]
   %quantized_input = "mhlo_disc.dequantize"(%input, %input_scale, %input_zero_point) {
@@ -286,9 +286,9 @@ func.func @dequantize_signed_per_channel(
   // CHECK-SAME: %[[INPUT_ZERO_POINT]], %[[SHAPE]]
   // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
 
-  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert(%[[INPUT]]) : (tensor<?x32x32x6xi8>) -> tensor<?x32x32x6xi32>
+  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert %[[INPUT]] : (tensor<?x32x32x6xi8>) -> tensor<?x32x32x6xi32>
   // CHECK: %[[T0:.*]] = mhlo.subtract %[[CASTED_INPUT]], %[[BCAST_ZERO_POINT]] : tensor<?x32x32x6xi32>
-  // CHECK: %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
+  // CHECK: %[[T1:.*]] = mhlo.convert %[[T0]] : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
   // CHECK: %[[T2:.*]] = mhlo.multiply %[[T1]], %[[BCAST_SCALE]] : tensor<?x32x32x6xf32>
   // CHECK: return %[[T2]]
   %quantized_input = "mhlo_disc.dequantize"(%input, %input_scale, %input_zero_point) {
@@ -319,9 +319,9 @@ func.func @dequantize_unsigned_per_channel(
   // CHECK-SAME: %[[INPUT_ZERO_POINT]], %[[SHAPE]]
   // CHECK-SAME: broadcast_dimensions = dense<3> : tensor<1xi64>
 
-  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert(%[[INPUT]]) : (tensor<?x32x32x6xui8>) -> tensor<?x32x32x6xi32>
+  // CHECK: %[[CASTED_INPUT:.*]] = mhlo.convert %[[INPUT]] : (tensor<?x32x32x6xui8>) -> tensor<?x32x32x6xi32>
   // CHECK: %[[T0:.*]] = mhlo.subtract %[[CASTED_INPUT]], %[[BCAST_ZERO_POINT]] : tensor<?x32x32x6xi32>
-  // CHECK: %[[T1:.*]] = mhlo.convert(%[[T0]]) : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
+  // CHECK: %[[T1:.*]] = mhlo.convert %[[T0]] : (tensor<?x32x32x6xi32>) -> tensor<?x32x32x6xf32>
   // CHECK: %[[T2:.*]] = mhlo.multiply %[[T1]], %[[BCAST_SCALE]] : tensor<?x32x32x6xf32>
   // CHECK: return %[[T2]]
   %quantized_input = "mhlo_disc.dequantize"(%input, %input_scale, %input_zero_point) {
diff --git a/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir b/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir
index d1ba61933d2..ae6e046e671 100644
--- a/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/disc-shape-optimization.mlir
@@ -92,7 +92,8 @@ func.func @main(%arg0 : tensor<?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf
 // CHECK-LABEL: main
 // CHECK-SAME: (%[[ARG0:.*]]: tensor<i1>, %[[ARG1:.*]]: tensor<?xf32, [@[[S0:.*]]]>, %[[ARG2:.*]]: tensor<?xf32, [@[[S0]]]>) -> tensor<?xf32, [@[[S0]]]>
 func.func @main(%arg0: tensor<i1>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
-  // CHECK: %[[T0:.*]] = "mhlo.select"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
+  // CHECK: %[[T0:.*]] = mhlo.select
+  // CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]]
   // CHECK: return %[[T0]] : tensor<?xf32, [@[[S0]]]>
   %0 = "mhlo.select"(%arg0, %arg1, %arg2)  : (tensor<i1>, tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
   func.return %0 : tensor<?xf32>
@@ -104,7 +105,8 @@ func.func @main(%arg0: tensor<i1>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -
 // CHECK-LABEL: main
 // CHECK-SAME: (%[[ARG0:.*]]: tensor<?xi1, [@[[S0:.*]]]>, %[[ARG1:.*]]: tensor<?xf32, [@[[S0]]]>, %[[ARG2:.*]]: tensor<?xf32, [@[[S0]]]>) -> tensor<?xf32, [@[[S0]]]>
 func.func @main(%arg0: tensor<?xi1>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
-  // CHECK: %[[T0:.*]] = "mhlo.select"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
+  // CHECK: %[[T0:.*]] = mhlo.select
+  // CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]]
   // CHECK: return %[[T0]] : tensor<?xf32, [@[[S0]]]>
   %0 = "mhlo.select"(%arg0, %arg1, %arg2)  : (tensor<?xi1>, tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
   func.return %0 : tensor<?xf32>
diff --git a/tao_compiler/mlir/disc/transforms/tests/element_type_converter.mlir b/tao_compiler/mlir/disc/transforms/tests/element_type_converter.mlir
index e833d525b74..f63782f9479 100644
--- a/tao_compiler/mlir/disc/transforms/tests/element_type_converter.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/element_type_converter.mlir
@@ -9,10 +9,10 @@ func.func @rank2_colunm_reduction_i1(%arg0: tensor<?x?xi1>) -> tensor<?xi1> {
     %2 = mhlo.add %arg1, %arg2 : tensor<i1>
     "mhlo.return"(%2) : (tensor<i1>) -> ()
   }) {dimensions = dense<[0]> : tensor<1xi64>} : (tensor<?x?xi1>, tensor<i1>) -> tensor<?xi1>
-  // CHECK: mhlo.convert({{.*}}) : (tensor<?x?xi1>) -> tensor<?x?xi32>
+  // CHECK: mhlo.convert {{.*}} : (tensor<?x?xi1>) -> tensor<?x?xi32>
   // CHECK-NEXT: mhlo.reduce
   // CHECK-NOT: i1
-  // CHECK: mhlo.convert({{.*}}) : (tensor<?xi32>) -> tensor<?xi1>
+  // CHECK: mhlo.convert {{.*}} : (tensor<?xi32>) -> tensor<?xi1>
   return %1 : tensor<?xi1>
 }
 
diff --git a/tao_compiler/mlir/disc/transforms/tests/gpu-only-lhlo-legalize-roots-to-loops.mlir b/tao_compiler/mlir/disc/transforms/tests/gpu-only-lhlo-legalize-roots-to-loops.mlir
index a9175eb9abf..825c8c719d2 100644
--- a/tao_compiler/mlir/disc/transforms/tests/gpu-only-lhlo-legalize-roots-to-loops.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/gpu-only-lhlo-legalize-roots-to-loops.mlir
@@ -142,12 +142,12 @@ func.func @naive_reduce(%operand: memref<?x?x?xf32>, %init_value: memref<f32>, %
 // CHECK:    "lmhlo.fusion"() ({
 // CHECK-NOT "lmhlo.reduce"
 // CHECK:      scf.parallel (%arg6, %arg7) = (%c0, %c0) to (%c2048, %c256) step (%c1, %c1)
-// CHECK:        %8:2 = scf.for %arg8 = %arg7 to %c768 step %c256 iter_args(%arg9 = %4, %arg10 = %6) -> (f32, f64) {
-// CHECK:          %75 = memref.load %arg0[%arg6, %arg8] : memref<2048x768xf32>
-// CHECK:          %76 = arith.addf %arg9, %75 : f32
-// CHECK:          %77 = memref.load %arg3[%arg6, %arg8] : memref<2048x768xf64>
-// CHECK:          %78 = arith.addf %arg10, %77 : f64
-// CHECK:          scf.yield %76, %78 : f32, f64
+// CHECK:        %[[T0:.*]]:2 = scf.for %arg8 = %arg7 to %c768 step %c256 iter_args(%arg9 = %4, %arg10 = %5) -> (f32, f64) {
+// CHECK:          %[[T1:.*]] = memref.load %arg0[%arg6, %arg8] : memref<2048x768xf32>
+// CHECK:          %[[T2:.*]] = arith.addf %arg9, %[[T1]] : f32
+// CHECK:          %[[T3:.*]] = memref.load %arg3[%arg6, %arg8] : memref<2048x768xf64>
+// CHECK:          %[[T4:.*]] = arith.addf %arg10, %[[T3]] : f64
+// CHECK:          scf.yield %[[T2]], %[[T4]] : f32, f64
 func.func @multi_row_reduce(%arg_f32: memref<2048x768xf32>, %init_f32: memref<f32>, %out_f32: memref<2048xf32>, %arg_f64: memref<2048x768xf64>, %init_f64: memref<f64>, %out_f64: memref<2048xf64>) -> (memref<2048xf32>, memref<2048xf64>) {
   "lmhlo.fusion"() ({
     "lmhlo.reduce"(%arg_f32, %init_f32, %out_f32) ({
diff --git a/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir b/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir
index 85729710f63..53c5013cce9 100644
--- a/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/mhlo_mark_shape_calc_op.mlir
@@ -63,11 +63,11 @@ func.func @main(%input: tensor<?x?x?x?xi8>, %weight: tensor<?x?x?x?xi8>, %paddin
                                   input_placements = "cpu,cpu,gpu,cpu,cpu,cpu,cpu,cpu,cpu",
                                   inputs = "input0,input1,input2,input3,input4,input5,input6,input7,input8",
                                   output_placements = "cpu", outputs = "output0"}} {
-  // CHECK: %[[RefinedPadding:.*]] = mhlo.convert(%[[PADDING]])
+  // CHECK: %[[RefinedPadding:.*]] = mhlo.convert %[[PADDING]]
   // CHECK-SAME: disc.shape_op = true
   // CHECK: %[[OUT:.*]] = "mhlo_disc.quantized_dynamic_conv"
   // CHECK-SAME: %[[INPUT]], %[[WEIGHT]], %[[RefinedPadding]]
-  %refined_padding = mhlo.convert(%padding) : (tensor<4xf32>) -> tensor<4xi32>
+  %refined_padding = mhlo.convert %padding : (tensor<4xf32>) -> tensor<4xi32>
   %out = "mhlo_disc.quantized_dynamic_conv"(%input, %weight, %refined_padding,
                                            %input_scale, %input_zero_point,
                                            %weight_scale, %weight_zero_point,
diff --git a/tao_compiler/mlir/disc/transforms/tests/mhlo_place_ops.mlir b/tao_compiler/mlir/disc/transforms/tests/mhlo_place_ops.mlir
index c35facb0fcb..fca785c6118 100644
--- a/tao_compiler/mlir/disc/transforms/tests/mhlo_place_ops.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/mhlo_place_ops.mlir
@@ -2,8 +2,8 @@
 
 // CHECK-LABEL: @main
 func.func @main(%arg : tensor<i64>) -> tensor<i64> attributes {tf.entry_function = {input_placements = "cpu", inputs = "input0", output_placements = "cpu", outputs = "output0"}}  {
-  // CHECK: "mhlo.tuple"({{.*}}) {disc.device = ["cpu"]} : (tensor<i64>) -> tuple<tensor<i64>>
-  // CHECK: "mhlo.get_tuple_element"({{.*}}) {disc.device = "cpu", index = 0 : i32} : (tuple<tensor<i64>>) -> tensor<i64>
+  // CHECK: mhlo.tuple {{.*}} {disc.device = ["cpu"]}
+  // CHECK: mhlo.get_tuple_element {{.*}} {disc.device = "cpu"} : (tuple<tensor<i64>>) -> tensor<i64>
   %tuple = "mhlo.tuple"(%arg) : (tensor<i64>) -> tuple<tensor<i64>>
   %element = "mhlo.get_tuple_element"(%tuple) {index = 0 : i32} : (tuple<tensor<i64>>) -> tensor<i64>
   return %element : tensor<i64>
@@ -80,11 +80,11 @@ func.func @main(%input: tensor<?x?x?x?xi8>, %weight: tensor<?x?x?x?xi8>, %paddin
                                   inputs = "input0,input1,input2,input3,input4,input5,input6,input7,input8",
                                   output_placements = "cpu", outputs = "output0"}} {
   // CHECK: %[[HOST_PADDING:.*]] = "mhlo_disc.d2h"(%[[PADDING]])
-  // CHECK: %[[RefinedPadding:.*]] = mhlo.convert(%[[HOST_PADDING]])
+  // CHECK: %[[RefinedPadding:.*]] = mhlo.convert %[[HOST_PADDING]]
   // CHECK-SAME: disc.shape_op = true
   // CHECK: %[[OUT:.*]] = "mhlo_disc.quantized_dynamic_conv"
   // CHECK-SAME: %[[RefinedPadding]]
-  %refined_padding = mhlo.convert(%padding) {disc.shape_op = true} : (tensor<4xf32>) -> tensor<4xi32>
+  %refined_padding = mhlo.convert %padding {disc.shape_op = true} : (tensor<4xf32>) -> tensor<4xi32>
   %out = "mhlo_disc.quantized_dynamic_conv"(%input, %weight, %refined_padding,
                                            %input_scale, %input_zero_point,
                                            %weight_scale, %weight_zero_point,
diff --git a/tao_compiler/mlir/disc/transforms/tests/quantized-disc-dot-rewriter.mlir b/tao_compiler/mlir/disc/transforms/tests/quantized-disc-dot-rewriter.mlir
index 75df6b503a9..4a821d7ae70 100644
--- a/tao_compiler/mlir/disc/transforms/tests/quantized-disc-dot-rewriter.mlir
+++ b/tao_compiler/mlir/disc/transforms/tests/quantized-disc-dot-rewriter.mlir
@@ -21,7 +21,7 @@ func.func @main(%arg0: tensor<?x?xf32>) -> tensor<?x128xf32> attributes {tf.entr
   %16 = mhlo.add %15, %12 : tensor<?x?xf32>
   %17 = mhlo.round_nearest_afz %16 : tensor<?x?xf32>
   %18 = mhlo.clamp %17, %13, %14 : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
-  %19 = mhlo.convert(%18) : (tensor<?x?xf32>) -> tensor<?x?xi8>
+  %19 = mhlo.convert %18 : (tensor<?x?xf32>) -> tensor<?x?xi8>
   %20 = "mhlo.dynamic_broadcast_in_dim"(%6, %3) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<128xf32>, tensor<2xindex>) -> tensor<64x128xf32>
   %21 = "mhlo.dynamic_broadcast_in_dim"(%4, %3) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<128xf32>, tensor<2xindex>) -> tensor<64x128xf32>
   %22 = "mhlo.dynamic_broadcast_in_dim"(%1, %3) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>, tensor<2xindex>) -> tensor<64x128xf32>
@@ -30,16 +30,16 @@ func.func @main(%arg0: tensor<?x?xf32>) -> tensor<?x128xf32> attributes {tf.entr
   %25 = mhlo.add %24, %21 : tensor<64x128xf32>
   %26 = mhlo.round_nearest_afz %25 : tensor<64x128xf32>
   %27 = mhlo.clamp %26, %22, %23 : (tensor<64x128xf32>, tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
-  %28 = mhlo.convert(%27) : (tensor<64x128xf32>) -> tensor<64x128xi8>
+  %28 = mhlo.convert %27 : (tensor<64x128xf32>) -> tensor<64x128xi8>
   // CHECK: %[[T29:.*]] = "mhlo_disc.quantized_dot_general"
   // CHECK-SAME: rhs_contracting_dimensions = [1]
   %29 = "mhlo_disc.quantized_dot_general"(%19, %28, %8, %9, %6, %5, %8, %9) {axis = dense<1> : tensor<1xi64>, dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, use_dynamic = false, use_symmetric = true} : (tensor<?x?xi8>, tensor<64x128xi8>, tensor<f32>, tensor<i32>, tensor<128xf32>, tensor<128xi32>, tensor<f32>, tensor<i32>) -> tensor<?x128xi8>
   %30 = shape.shape_of %29 : tensor<?x128xi8> -> tensor<2xindex>
   %31 = "mhlo.dynamic_broadcast_in_dim"(%8, %30) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>, tensor<2xindex>) -> tensor<?x128xf32>
   %32 = "mhlo.dynamic_broadcast_in_dim"(%9, %30) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<i32>, tensor<2xindex>) -> tensor<?x128xi32>
-  %33 = mhlo.convert(%29) : (tensor<?x128xi8>) -> tensor<?x128xi32>
+  %33 = mhlo.convert %29 : (tensor<?x128xi8>) -> tensor<?x128xi32>
   %34 = mhlo.subtract %33, %32 : tensor<?x128xi32>
-  %35 = mhlo.convert(%34) : (tensor<?x128xi32>) -> tensor<?x128xf32>
+  %35 = mhlo.convert %34 : (tensor<?x128xi32>) -> tensor<?x128xf32>
   %36 = mhlo.multiply %35, %31 : tensor<?x128xf32>
   return %36 : tensor<?x128xf32>
 }
\ No newline at end of file