diff --git a/integrations/tensorflow/e2e/BUILD b/integrations/tensorflow/e2e/BUILD
index 7a07478..dc289d6 100644
--- a/integrations/tensorflow/e2e/BUILD
+++ b/integrations/tensorflow/e2e/BUILD
@@ -95,6 +95,7 @@
     "range_test.py",
     "ring_buffer_test.py",  # TODO(b/148747011)
     "scatter_update_test.py",
+    "sort_test.py",
     "strings_test.py",
 ]
 
@@ -112,6 +113,7 @@
     "range_test.py",
     "ring_buffer_test.py",  # TODO(b/148747011)
     "scatter_update_test.py",
+    "sort_test.py",
     "strings_test.py",
 ]
 
diff --git a/iree/compiler/Dialect/VMLA/Conversion/ConversionTarget.cpp b/iree/compiler/Dialect/VMLA/Conversion/ConversionTarget.cpp
index a80a260..6af672a 100644
--- a/iree/compiler/Dialect/VMLA/Conversion/ConversionTarget.cpp
+++ b/iree/compiler/Dialect/VMLA/Conversion/ConversionTarget.cpp
@@ -42,6 +42,7 @@
   // Pseudo-ops are illegal.
   // If we end up with a lot of these, consider using an "is pseudo" trait.
   addIllegalOp<IREE::VMLA::BatchMatMulPseudoOp>();
+  addIllegalOp<IREE::VMLA::SortPseudoOp>();
 
   // Allow other ops to pass through so long as their type is valid (not a
   // tensor, basically).
diff --git a/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/ConvertHLOToVMLA.cpp b/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/ConvertHLOToVMLA.cpp
index ccd1a3f..75b640f 100644
--- a/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/ConvertHLOToVMLA.cpp
+++ b/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/ConvertHLOToVMLA.cpp
@@ -704,6 +704,29 @@
   TypeConverter &typeConverter;
 };
 
+struct SortOpConversion : public OpConversionPattern<IREE::VMLA::SortPseudoOp> {
+  SortOpConversion(MLIRContext *context, TypeConverter &typeConverter)
+      : OpConversionPattern(context), typeConverter(typeConverter) {}
+
+  LogicalResult matchAndRewrite(
+      IREE::VMLA::SortPseudoOp srcOp, ArrayRef<Value> rawOperands,
+      ConversionPatternRewriter &rewriter) const override {
+    auto inputType =
+        srcOp.getOperand().getType().cast<ShapedType>().getElementType();
+    auto src = rawOperands[0];
+    auto src_shape = VMLAConversionTarget::getTensorShape(
+        srcOp.getLoc(), srcOp.value(), typeConverter, rewriter);
+    auto dst = VMLAConversionTarget::allocateOutputBuffer(
+        srcOp.getLoc(), srcOp.getResult(), typeConverter, rewriter);
+    rewriter.createOrFold<IREE::VMLA::SortOp>(srcOp.getLoc(), src, src_shape,
+                                              dst, TypeAttr::get(inputType));
+    rewriter.replaceOp(srcOp, {dst});
+    return success();
+  }
+
+  TypeConverter &typeConverter;
+};
+
 struct ConvertOpConversion : public OpConversionPattern<mhlo::ConvertOp> {
   ConvertOpConversion(MLIRContext *context, TypeConverter &typeConverter)
       : OpConversionPattern(context), typeConverter(typeConverter) {}
@@ -769,6 +792,9 @@
                                    IREE::VMLA::BatchMatMulOp>>(context,
                                                                typeConverter);
 
+  // vmla.sort.pseudo
+  patterns.insert<SortOpConversion>(context, typeConverter);
+
   // Simple 1:1 conversion patterns using the automated trait-based converter.
   // Used for HLO ops that have equivalent VMLA ops such as most arithmetic ops.
   patterns.insert<VMLAOpConversion<mhlo::AddOp, IREE::VMLA::AddOp>>(
diff --git a/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/test/sort.mlir b/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/test/sort.mlir
new file mode 100644
index 0000000..0903793
--- /dev/null
+++ b/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/test/sort.mlir
@@ -0,0 +1,37 @@
+// RUN: iree-opt -split-input-file -iree-vmla-pre-conversion-lowering -iree-vmla-conversion -canonicalize %s | IreeFileCheck %s
+
+func @sort1D(%arg0 : tensor<4xf32>) -> tensor<4xf32> attributes { sym_visibility = "private" } {
+  // CHECK-DAG: [[C16:%.+]] = constant 16 : index
+  // CHECK-DAG: [[RS:%.+]] = shapex.const_ranked_shape : !shapex.ranked_shape<[4]>
+  // CHECK-DAG: [[BL:%.+]] = vmla.buffer.alloc byte_length = [[C16]] : !vmla.buffer
+  // CHECK-DAG: vmla.sort %arg0([[RS]] : !shapex.ranked_shape<[4]>), out [[BL]] : f32
+  // CHECK-DAG: [[BUF:%.+]] = vmla.buffer.alloc byte_length = [[C16]] : !vmla.buffer
+  // CHECK-DAG: vmla.gather %arg0([[RS]] : !shapex.ranked_shape<[4]>), [[BL]]([[RS]] : !shapex.ranked_shape<[4]>), out [[BUF]]([[RS]] : !shapex.ranked_shape<[4]>) {batch_dims = 0 : i64, dim = 0 : i64} : f32
+  %sort = "mhlo.sort"(%arg0) ( {
+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 0 : i64, is_stable = false} : (tensor<4xf32>) -> tensor<4xf32>
+
+  // CHECK: return [[BUF]] : !vmla.buffer
+  return %sort : tensor<4xf32>
+}
+
+
+// CHECK-LABEL: func @sort2D
+func @sort2D(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> attributes { sym_visibility = "private" } {
+  // CHECK-DAG: [[C64:%.+]] = constant 64 : index
+  // CHECK-DAG: [[RS:%.+]] = shapex.const_ranked_shape : !shapex.ranked_shape<[4,4]>
+  // CHECK-DAG: [[BL:%.+]] = vmla.buffer.alloc byte_length = [[C64]] : !vmla.buffer
+  // CHECK-DAG: vmla.sort %arg0([[RS]] : !shapex.ranked_shape<[4,4]>), out [[BL]] : f32
+  // CHECK-DAG: [[BUF:%.+]] = vmla.buffer.alloc byte_length = [[C64]] : !vmla.buffer
+  // CHECK-DAG: vmla.gather %arg0([[RS]] : !shapex.ranked_shape<[4,4]>), [[BL]]([[RS]] : !shapex.ranked_shape<[4,4]>), out [[BUF]]([[RS]] : !shapex.ranked_shape<[4,4]>) {batch_dims = 1 : i64, dim = 1 : i64} : f32
+  %sort = "mhlo.sort"(%arg0) ( {
+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 1 : i64, is_stable = false} : (tensor<4x4xf32>) -> tensor<4x4xf32>
+
+  // CHECK: return [[BUF]] : !vmla.buffer
+  return %sort : tensor<4x4xf32>
+}
diff --git a/iree/compiler/Dialect/VMLA/Conversion/VMLAToVM/ConvertVMLAToVM.cpp b/iree/compiler/Dialect/VMLA/Conversion/VMLAToVM/ConvertVMLAToVM.cpp
index 7e4f6ea..1b66485 100644
--- a/iree/compiler/Dialect/VMLA/Conversion/VMLAToVM/ConvertVMLAToVM.cpp
+++ b/iree/compiler/Dialect/VMLA/Conversion/VMLAToVM/ConvertVMLAToVM.cpp
@@ -321,6 +321,7 @@
   VMLA_TYPED_IMPORT_OP(IREE::VMLA::FloorOp, "vmla.floor");
   VMLA_TYPED_IMPORT_OP(IREE::VMLA::CeilOp, "vmla.ceil");
   VMLA_TYPED_IMPORT_OP(IREE::VMLA::RoundOp, "vmla.round");
+  VMLA_TYPED_IMPORT_OP(IREE::VMLA::SortOp, "vmla.sort");
 
   patterns.insert<VMLAConvertImportOpConversion>(context, importSymbols,
                                                  typeConverter, "vmla.convert");
diff --git a/iree/compiler/Dialect/VMLA/IR/VMLAOps.td b/iree/compiler/Dialect/VMLA/IR/VMLAOps.td
index f5cd0fe..422fed3 100644
--- a/iree/compiler/Dialect/VMLA/IR/VMLAOps.td
+++ b/iree/compiler/Dialect/VMLA/IR/VMLAOps.td
@@ -422,7 +422,7 @@
 }
 
 //===----------------------------------------------------------------------===//
-// VMLA Ops: Convultion
+// VMLA Ops: Convolution
 //===----------------------------------------------------------------------===//
 
 def VLMA_ConvOp : VMLA_Op<"conv", [VMLA_IncludeShapes]> {
@@ -460,6 +460,46 @@
 }
 
 //===----------------------------------------------------------------------===//
+// VMLA Ops: Sorting
+//===----------------------------------------------------------------------===//
+
+def VMLA_SortPseudoOp : VMLA_Op<"sort.pseudo"> {
+  let summary = "Tensor-level pseudo-op of VMLA::SortOp.";
+  let description = [{
+    This is a tensor-level version of VMLA::SortOp, to facilitate
+    the lowering process.
+
+    This operation generates a sorted index list along the last dimension,
+    performing batch-wise along all other dimensions.
+  }];
+  let arguments = (ins
+    AnyTensor:$value
+  );
+  let results = (outs
+    I32Tensor:$dst
+  );
+
+  let assemblyFormat = [{
+    $value attr-dict `:` `(`type($value)`)` `->` type($dst)
+  }];
+}
+
+def VMLA_SortOp : VMLA_ElementTypeOp<"sort", [VMLA_IncludeShapes]> {
+  let arguments = (ins
+    VMLA_Buffer:$src,
+    VMLA_Shape:$src_shape,
+    VMLA_Buffer:$dst,
+    VMLA_AnyTypeAttr:$element_type
+  );
+
+  let assemblyFormat = [{
+    $src`(`$src_shape `:` type($src_shape)`)``,`
+    `out` $dst attr-dict `:` $element_type
+  }];
+}
+
+
+//===----------------------------------------------------------------------===//
 // VMLA Ops: GEMM/GEMV
 //===----------------------------------------------------------------------===//
 
diff --git a/iree/compiler/Dialect/VMLA/Transforms/PreConversionLowering.cpp b/iree/compiler/Dialect/VMLA/Transforms/PreConversionLowering.cpp
index d937026..62dd921 100644
--- a/iree/compiler/Dialect/VMLA/Transforms/PreConversionLowering.cpp
+++ b/iree/compiler/Dialect/VMLA/Transforms/PreConversionLowering.cpp
@@ -277,6 +277,96 @@
   }
 };
 
+// Lower mhlo::SortOp to an pseudo SortOp in the VMLA dialect. This
+// pseudo op generates a set of ordered indices for that array along the last
+// dimension. Then using a torch_index_select the values can be reordered to
+// support arbitrary inputs.
+//
+// TODO(suderman): This lowering only covers the case of ascending values, we
+// should support a separate descending value case by having separate
+// SortAscending and SortDescending operations.
+class LowerSortOp : public OpRewritePattern<mhlo::SortOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(mhlo::SortOp op,
+                                PatternRewriter &rewriter) const override {
+    auto operandTy = op.getOperand(0).getType().cast<RankedTensorType>();
+    bool lastDimension =
+        (op.dimension() == -1) || (op.dimension() == (operandTy.getRank() - 1));
+
+    // TODO(suderman): Add transpose to sort along the last dimension.
+    if (!lastDimension) return failure();
+
+    auto &comparator = op.comparator();
+    auto &block = comparator.getBlocks().front();
+    auto &operations = block.getOperations();
+    auto comparison = dyn_cast_or_null<mhlo::CompareOp>(&operations.front());
+
+    // First verify that the block is purely a return of a comparison. This
+    // handles sorting a single tensor of values.
+    if (!comparison) return failure();
+
+    auto returnOp =
+        dyn_cast_or_null<mhlo::ReturnOp>(&(*(++operations.begin())));
+    if (!returnOp) return failure();
+
+    if (returnOp.getOperand(0) != comparison.getResult()) return failure();
+
+    // Determine which operands being compared.
+    auto lhs = comparison.getOperand(0);
+    auto rhs = comparison.getOperand(1);
+    auto lhsIndex = -1;
+    auto rhsIndex = -1;
+    for (auto arg : llvm::enumerate(block.getArguments())) {
+      if (arg.value() == lhs) lhsIndex = arg.index();
+      if (arg.value() == rhs) rhsIndex = arg.index();
+    }
+
+    // This should never happen but best to check.
+    if (lhsIndex == -1) return failure();
+    if (rhsIndex == -1) return failure();
+
+    // They should not be the same.
+    if (lhsIndex == rhsIndex) return failure();
+
+    // Comparisons need to pull from same Sort operand..
+    auto lhsOperand = lhsIndex / 2;
+    auto rhsOperand = rhsIndex / 2;
+    if (lhsOperand != rhsOperand) return failure();
+
+    // Must be GT, GE, LT, or LE.
+    auto isGt = comparison.comparison_direction() == "GT" ||
+                comparison.comparison_direction() == "GE";
+    auto isLt = comparison.comparison_direction() == "LT" ||
+                comparison.comparison_direction() == "LE";
+    if (!isGt && !isLt) return failure();
+
+    bool operandParity = lhsIndex > rhsIndex;
+    auto isAscending = operandParity ^ isGt;
+    // TODO(suderman): Add support for descended sorting.
+    if (!isAscending) return failure();
+
+    auto operand = op.getOperand(lhsOperand);
+    auto sortedIndices = rewriter.create<VMLA::SortPseudoOp>(
+        op.getLoc(),
+        RankedTensorType::get(operandTy.getShape(), rewriter.getI32Type()),
+        operand);
+
+    llvm::SmallVector<Value, 6> sortedResults;
+    for (auto operand : op.getOperands()) {
+      auto tensorTy = operand.getType().cast<RankedTensorType>();
+      auto gathered = rewriter.create<mhlo::TorchIndexSelectOp>(
+          op.getLoc(), tensorTy, operand, sortedIndices,
+          /**dim=*/operandTy.getRank() - 1,
+          /**batch_dims=*/operandTy.getRank() - 1);
+      sortedResults.push_back(gathered);
+    }
+
+    rewriter.replaceOp(op, sortedResults);
+    return success();
+  }
+};
+
 class PreConversionLoweringPass
     : public PassWrapper<PreConversionLoweringPass, OperationPass<FuncOp>> {
  public:
@@ -310,6 +400,8 @@
     patterns.insert<LowerBroadcastInDimOp>(context);
     target.addIllegalOp<mhlo::BroadcastOp>();
     patterns.insert<LowerBroadcastOp>(context);
+    target.addIllegalOp<mhlo::SortOp>();
+    patterns.insert<LowerSortOp>(context);
 
     if (failed(applyPartialConversion(getOperation(), target, patterns))) {
       return signalPassFailure();
diff --git a/iree/compiler/Dialect/VMLA/Transforms/test/pre_conversion_lowering.mlir b/iree/compiler/Dialect/VMLA/Transforms/test/pre_conversion_lowering.mlir
index 3473e44..0b9cd82 100644
--- a/iree/compiler/Dialect/VMLA/Transforms/test/pre_conversion_lowering.mlir
+++ b/iree/compiler/Dialect/VMLA/Transforms/test/pre_conversion_lowering.mlir
@@ -17,6 +17,38 @@
 // -----
 
 // CHECK-LABEL: func @f
+func @f(%arg0 : tensor<4xf32>) -> tensor<4xf32> attributes { sym_visibility = "private" } {
+  // CHECK-DAG: [[SORT:%.+]] = vmla.sort.pseudo %arg0
+  // CHECK-DAG: [[GATHER:%.+]] = "mhlo.torch_index_select"(%arg0, [[SORT]]) {batch_dims = 0 : i64, dim = 0 : i64}
+  %sort = "mhlo.sort"(%arg0) ( {
+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 0 : i64, is_stable = false} : (tensor<4xf32>) -> tensor<4xf32>
+
+  // CHECK: return [[GATHER]]
+  return %sort : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @f
+func @f(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> attributes { sym_visibility = "private" } {
+  // CHECK-DAG: [[SORT:%.+]] = vmla.sort.pseudo %arg0
+  // CHECK-DAG: [[GATHER:%.+]] = "mhlo.torch_index_select"(%arg0, [[SORT]]) {batch_dims = 1 : i64, dim = 1 : i64}
+  %sort = "mhlo.sort"(%arg0) ( {
+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 1 : i64, is_stable = false} : (tensor<4x4xf32>) -> tensor<4x4xf32>
+
+  // CHECK return [[GATHER]]
+  return %sort : tensor<4x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @f
 func @f(%arg0: tensor<3xf32>) -> tensor<4x3xf32> {
   // CHECK: "shapex.ranked_broadcast_in_dim"(%arg0, %rs4_3)
   %0 = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<[1]> : tensor<1xi64>} : (tensor<3xf32>) -> tensor<4x3xf32>
diff --git a/iree/compiler/Dialect/VMLA/vmla.imports.mlir b/iree/compiler/Dialect/VMLA/vmla.imports.mlir
index 17d7e85..ff575b0 100644
--- a/iree/compiler/Dialect/VMLA/vmla.imports.mlir
+++ b/iree/compiler/Dialect/VMLA/vmla.imports.mlir
@@ -333,6 +333,20 @@
 vm.import @ceil.f32(%src : !vm.ref<!vmla.buffer>, %dst : !vm.ref<!vmla.buffer>)
 vm.import @round.f32(%src : !vm.ref<!vmla.buffer>, %dst : !vm.ref<!vmla.buffer>)
 
+
+vm.import @sort.i8(
+  %src : !vm.ref<!vmla.buffer>, %src_shape : i32 ...,
+  %dst : !vm.ref<!vmla.buffer>)
+vm.import @sort.i16(
+  %src : !vm.ref<!vmla.buffer>, %src_shape : i32 ...,
+  %dst : !vm.ref<!vmla.buffer>)
+vm.import @sort.i32(
+  %src : !vm.ref<!vmla.buffer>, %src_shape : i32 ...,
+  %dst : !vm.ref<!vmla.buffer>)
+vm.import @sort.f32(
+  %src : !vm.ref<!vmla.buffer>, %src_shape : i32 ...,
+  %dst : !vm.ref<!vmla.buffer>)
+
 //===----------------------------------------------------------------------===//
 // VMLA Ops: conversion
 //===----------------------------------------------------------------------===//
diff --git a/iree/hal/vmla/op_kernels.h b/iree/hal/vmla/op_kernels.h
index c6ab9d6..ba5b8bc 100644
--- a/iree/hal/vmla/op_kernels.h
+++ b/iree/hal/vmla/op_kernels.h
@@ -168,6 +168,12 @@
                         absl::Span<const int32_t> dimensions);
 };
 
+struct Sort {
+  template <typename T>
+  static Status Execute(absl::Span<const T> src_buffer,
+                        absl::Span<int32_t> dst_buffer, ShapeSpan src_shape);
+};
+
 struct Broadcast {
   template <typename T>
   static Status Execute(absl::Span<const T> src_buffer,
diff --git a/iree/hal/vmla/op_kernels_generic.h b/iree/hal/vmla/op_kernels_generic.h
index d3545d1..0b4f904 100644
--- a/iree/hal/vmla/op_kernels_generic.h
+++ b/iree/hal/vmla/op_kernels_generic.h
@@ -15,7 +15,10 @@
 #ifndef IREE_HAL_VMLA_OP_KERNELS_GENERIC_H_
 #define IREE_HAL_VMLA_OP_KERNELS_GENERIC_H_
 
+#include <algorithm>
 #include <cmath>
+#include <iostream>
+#include <numeric>
 
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/inlined_vector.h"
@@ -519,6 +522,25 @@
 }
 
 template <typename T>
+Status Sort::Execute(absl::Span<const T> src_buffer,
+                     absl::Span<int32_t> dst_buffer, ShapeSpan src_shape) {
+  int elements = src_buffer.size();
+  const int sort_size = src_shape.back();
+
+  for (int i = 0; i < elements; i += sort_size) {
+    auto src_subspan = src_buffer.subspan(i, sort_size);
+    auto dst_subspan = dst_buffer.subspan(i, sort_size);
+    std::iota(dst_subspan.begin(), dst_subspan.end(), 0);
+    std::stable_sort(dst_subspan.begin(), dst_subspan.end(),
+                     [&src_subspan](int32_t i1, int32_t i2) {
+                       return src_subspan[i1] < src_subspan[i2];
+                     });
+  }
+
+  return OkStatus();
+}
+
+template <typename T>
 Status Broadcast::Execute(absl::Span<const T> src_buffer,
                           absl::Span<T> dst_buffer) {
   for (size_t i = 0; i < dst_buffer.size(); ++i) {
diff --git a/iree/hal/vmla/vmla_module.cc b/iree/hal/vmla/vmla_module.cc
index 09dbb3c..5852de0 100644
--- a/iree/hal/vmla/vmla_module.cc
+++ b/iree/hal/vmla/vmla_module.cc
@@ -642,6 +642,19 @@
   IREE_VMLA_UNARY_OP(CeilF32, kernels::Ceil, float);
   IREE_VMLA_UNARY_OP(RoundF32, kernels::Round, float);
 
+#define IREE_VMLA_SORT_OP(name, type)                                        \
+  Status name(const vm::ref<Buffer>& src, iree_vmla_shape_t src_shape,       \
+              const vm::ref<Buffer>& dst) {                                  \
+    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                            \
+    return kernels::Sort::Execute<type>(src->As<type>(), dst->As<int32_t>(), \
+                                        src_shape);                          \
+  }
+
+  IREE_VMLA_SORT_OP(SortI8, int8_t);
+  IREE_VMLA_SORT_OP(SortI16, int16_t);
+  IREE_VMLA_SORT_OP(SortI32, int32_t);
+  IREE_VMLA_SORT_OP(SortF32, float);
+
   //===--------------------------------------------------------------------===//
   // VMLA Ops: conversion
   //===--------------------------------------------------------------------===//
@@ -970,6 +983,10 @@
     vm::MakeNativeFunction("floor.f32", &VMLAModuleState::FloorF32),
     vm::MakeNativeFunction("ceil.f32", &VMLAModuleState::CeilF32),
     vm::MakeNativeFunction("round.f32", &VMLAModuleState::RoundF32),
+    vm::MakeNativeFunction("sort.i8", &VMLAModuleState::SortI8),
+    vm::MakeNativeFunction("sort.i16", &VMLAModuleState::SortI16),
+    vm::MakeNativeFunction("sort.i32", &VMLAModuleState::SortI32),
+    vm::MakeNativeFunction("sort.f32", &VMLAModuleState::SortF32),
     vm::MakeNativeFunction("finite.f32", &VMLAModuleState::FiniteF32),
 
     vm::MakeNativeFunction("convert.i8.i16", &VMLAModuleState::ConvertI8I16),
diff --git a/iree/test/e2e/xla_ops/sort.mlir b/iree/test/e2e/xla_ops/sort.mlir
new file mode 100644
index 0000000..1820d8d
--- /dev/null
+++ b/iree/test/e2e/xla_ops/sort.mlir
@@ -0,0 +1,40 @@
+func @sort1D() attributes { iree.module.export } {
+  %input = iree.unfoldable_constant dense<[3, 2, 1, 4]> : tensor<4xi32>
+
+  %sort = "mhlo.sort"(%input) ( {
+  ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 0 : i64, is_stable = false} : (tensor<4xi32>) -> tensor<4xi32>
+
+  check.expect_eq_const(%sort, dense<[1, 2, 3, 4]> : tensor<4xi32>) : tensor<4xi32>
+  return
+}
+
+func @sort2D() attributes { iree.module.export } {
+  %input = iree.unfoldable_constant dense<[[1, 2, 3, 4],
+                                           [4, 3, 2, 1]]> : tensor<2x4xi32>
+
+  %sort = "mhlo.sort"(%input) ( {
+  ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 1 : i64, is_stable = false} : (tensor<2x4xi32>) -> tensor<2x4xi32>
+
+  check.expect_eq_const(%sort, dense<[[1, 2, 3, 4], [1, 2, 3, 4]]> : tensor<2x4xi32>) : tensor<2x4xi32>
+  return
+}
+
+func @sort3D() attributes { iree.module.export } {
+  %input = iree.unfoldable_constant dense<[[[1, 2, 3, 4],
+                                            [4, 3, 2, 1]]]> : tensor<1x2x4xi32>
+
+  %sort = "mhlo.sort"(%input) ( {
+  ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 2 : i64, is_stable = false} : (tensor<1x2x4xi32>) -> tensor<1x2x4xi32>
+
+  check.expect_eq_const(%sort, dense<[[[1, 2, 3, 4], [1, 2, 3, 4]]]> : tensor<1x2x4xi32>) : tensor<1x2x4xi32>
+  return
+}