feat!: drop modules, general refactor and cleanup

3 files changed, 362 insertions, 0 deletions

diff --git a/runtime/queries/mlir/highlights.scm b/runtime/queries/mlir/highlights.scm
new file mode 100644
index 000000000..97a0b0d0b
--- /dev/null
+++ b/runtime/queries/mlir/highlights.scm
@@ -0,0 +1,348 @@
+[
+  "ins"
+  "outs"
+  "else"
+  "do"
+  "loc"
+  "attributes"
+  "into"
+  "to"
+  "from"
+  "step"
+  "low"
+  "high"
+  "iter_args"
+  "padding_value"
+  "inner_tiles"
+  "gather_dims"
+  "scatter_dims"
+  "outer_dims_perm"
+  "inner_dims_pos"
+  "shared_outs"
+  "default"
+  (arith_cmp_predicate)
+] @keyword
+
+; format-ignore
+[
+  "module"
+  "unrealized_conversion_cast"
+
+  "func.call"
+  "call"
+  "func.call_indirect"
+  "call_indirect"
+  "func.constant"
+  "constant"
+  "func.func"
+  "func.return"
+  "return"
+
+  "llvm.func"
+  "llvm.return"
+
+  "cf.assert"
+  "cf.br"
+  "cf.cond_br"
+  "cf.switch"
+
+  "scf.condition"
+  "scf.execute_region"
+  "scf.if"
+  "scf.index_switch"
+  "scf.for"
+  "scf.forall"
+  "scf.forall.in_parallel"
+  "scf.parallel"
+  "scf.reduce"
+  "scf.reduce.return"
+  "scf.while"
+  "scf.yield"
+
+  "arith.constant"
+  "arith.addi"
+  "arith.subi"
+  "arith.divsi"
+  "arith.divui"
+  "arith.ceildivsi"
+  "arith.ceildivui"
+  "arith.floordivsi"
+  "arith.remsi"
+  "arith.remui"
+  "arith.muli"
+  "arith.mulsi_extended"
+  "arith.mului_extended"
+  "arith.andi"
+  "arith.ori"
+  "arith.xori"
+  "arith.maxsi"
+  "arith.maxui"
+  "arith.minsi"
+  "arith.minui"
+  "arith.shli"
+  "arith.shrsi"
+  "arith.shrui"
+  "arith.addui_extended"
+  "arith.addf"
+  "arith.divf"
+  "arith.maximumf"
+  "arith.minimumf"
+  "arith.mulf"
+  "arith.remf"
+  "arith.subf"
+  "arith.negf"
+  "arith.cmpi"
+  "arith.cmpf"
+  "arith.extf"
+  "arith.extsi"
+  "arith.extui"
+  "arith.fptosi"
+  "arith.fptoui"
+  "arith.index_cast"
+  "arith.index_castui"
+  "arith.sitofp"
+  "arith.uitofp"
+  "arith.bitcast"
+  "arith.truncf"
+  "arith.select"
+
+  "math.absf"
+  "math.atan"
+  "math.cbrt"
+  "math.ceil"
+  "math.cos"
+  "math.erf"
+  "math.exp"
+  "math.exp2"
+  "math.expm1"
+  "math.floor"
+  "math.log"
+  "math.log10"
+  "math.log1p"
+  "math.log2"
+  "math.round"
+  "math.roundeven"
+  "math.rsqrt"
+  "math.sin"
+  "math.sqrt"
+  "math.tan"
+  "math.tanh"
+  "math.trunc"
+  "math.absi"
+  "math.ctlz"
+  "math.cttz"
+  "math.ctpop"
+  "math.atan2"
+  "math.copysign"
+  "math.fpowi"
+  "math.powf"
+  "math.ipowi"
+  "math.fma"
+
+  "memref.alloc"
+  "memref.cast"
+  "memref.copy"
+  "memref.collapse_shape"
+  "memref.expand_shape"
+  "memref.prefetch"
+  "memref.rank"
+  "memref.realloc"
+  "memref.view"
+
+  "vector.bitcast"
+  "vector.broadcast"
+  "vector.shape_cast"
+  "vector.type_cast"
+  "vector.constant_mask"
+  "vector.create_mask"
+  "vector.extract"
+  "vector.load"
+  "vector.scalable.extract"
+  "vector.fma"
+  "vector.flat_transpose"
+  "vector.insert"
+  "vector.scalable.insert"
+  "vector.shuffle"
+  "vector.store"
+  "vector.insert_strided_slice"
+  "vector.matrix_multiply"
+  "vector.print"
+  "vector.splat"
+  "vector.transfer_read"
+  "vector.transfer_write"
+  "vector.yield"
+
+  "tensor.empty"
+  "tensor.cast"
+  "tensor.dim"
+  "tensor.collapse_shape"
+  "tensor.expand_shape"
+  "tensor.extract"
+  "tensor.insert"
+  "tensor.extract_slice"
+  "tensor.insert_slice"
+  "tensor.parallel_insert_slice"
+  "tensor.from_elements"
+  "tensor.gather"
+  "tensor.scatter"
+  "tensor.pad"
+  "tensor.reshape"
+  "tensor.splat"
+  "tensor.pack"
+  "tensor.unpack"
+  "tensor.generate"
+  "tensor.rank"
+  "tensor.yield"
+
+  "bufferization.alloc_tensor"
+  "bufferization.to_memref"
+  "bufferization.to_tensor"
+
+  "linalg.batch_matmul"
+  "linalg.batch_matmul_transpose_b"
+  "linalg.batch_matvec"
+  "linalg.batch_reduce_matmul"
+  "linalg.broadcast"
+  "linalg.conv_1d_ncw_fcw"
+  "linalg.conv_1d_nwc_wcf"
+  "linalg.conv_1d"
+  "linalg.conv_2d_nchw_fchw"
+  "linalg.conv_2d_ngchw_fgchw"
+  "linalg.conv_2d_nhwc_fhwc"
+  "linalg.conv_2d_nhwc_hwcf"
+  "linalg.conv_2d_nhwc_hwcf_q"
+  "linalg.conv_2d"
+  "linalg.conv_3d_ndhwc_dhwcf"
+  "linalg.conv_3d_ndhwc_dhwcf_q"
+  "linalg.conv_3d"
+  "linalg.copy"
+  "linalg.depthwise_conv_1d_nwc_wcm"
+  "linalg.depthwise_conv_2d_nchw_chw"
+  "linalg.depthwise_conv_2d_nhwc_hwc"
+  "linalg.depthwise_conv_2d_nhwc_hwc_q"
+  "linalg.depthwise_conv_2d_nhwc_hwcm"
+  "linalg.depthwise_conv_2d_nhwc_hwcm_q"
+  "linalg.depthwise_conv_3d_ndhwc_dhwc"
+  "linalg.depthwise_conv_3d_ndhwc_dhwcm"
+  "linalg.dot"
+  "linalg.elemwise_binary"
+  "linalg.elemwise_unary"
+  "linalg.fill"
+  "linalg.fill_rng_2d"
+  "linalg.matmul"
+  "linalg.matmul_transpose_b"
+  "linalg.matmul_unsigned"
+  "linalg.matvec"
+  "linalg.mmt4d"
+  "linalg.pooling_nchw_max"
+  "linalg.pooling_nchw_sum"
+  "linalg.pooling_ncw_max"
+  "linalg.pooling_ncw_sum"
+  "linalg.pooling_ndhwc_max"
+  "linalg.pooling_ndhwc_min"
+  "linalg.pooling_ndhwc_sum"
+  "linalg.pooling_nhwc_max"
+  "linalg.pooling_nhwc_max_unsigned"
+  "linalg.pooling_nhwc_min"
+  "linalg.pooling_nhwc_min_unsigned"
+  "linalg.pooling_nhwc_sum"
+  "linalg.pooling_nwc_max"
+  "linalg.pooling_nwc_max_unsigned"
+  "linalg.pooling_nwc_min"
+  "linalg.pooling_nwc_min_unsigned"
+  "linalg.pooling_nwc_sum"
+  "linalg.quantized_batch_matmul"
+  "linalg.quantized_matmul"
+  "linalg.vecmat"
+  "linalg.generic"
+  "linalg.index"
+  "linalg.map"
+  "linalg.yield"
+] @function.builtin
+
+(generic_operation) @function
+
+(builtin_type) @type.builtin
+
+[
+  (type_alias)
+  (dialect_type)
+  (type_alias_def)
+] @type
+
+[
+  (integer_literal)
+  (complex_literal)
+] @number
+
+(float_literal) @number.float
+
+(bool_literal) @boolean
+
+[
+  (tensor_literal)
+  (array_literal)
+  (unit_literal)
+] @constant.builtin
+
+(string_literal) @string
+
+[
+  (attribute_alias_def)
+  (attribute_alias)
+  (bare_attribute_entry)
+  (attribute)
+  (fastmath_attr)
+  (scatter_dims_attr)
+  (gather_dims_attr)
+  (outer_dims_perm_attr)
+  (inner_dims_pos_attr)
+  (inner_tiles_attr)
+  (unique_attr)
+  (nofold_attr)
+  (isWrite_attr)
+  (localityHint_attr)
+  (isDataCache_attr)
+  (restrict_attr)
+  (writable_attr)
+] @attribute
+
+[
+  "("
+  ")"
+  "{"
+  "}"
+  "["
+  "]"
+] @punctuation.bracket
+
+[
+  ":"
+  ","
+] @punctuation.delimiter
+
+[
+  "="
+  "->"
+] @operator
+
+(builtin_dialect
+  name: (symbol_ref_id) @function)
+
+(func_dialect
+  name: (symbol_ref_id) @function)
+
+(llvm_dialect
+  name: (symbol_ref_id) @function)
+
+(func_arg_list
+  (value_use) @variable.parameter)
+
+(block_arg_list
+  (value_use) @variable.parameter)
+
+(caret_id) @string.special
+
+(value_use) @variable
+
+(comment) @comment @spell
diff --git a/runtime/queries/mlir/injections.scm b/runtime/queries/mlir/injections.scm
new file mode 100644
index 000000000..2f0e58eb6
--- /dev/null
+++ b/runtime/queries/mlir/injections.scm
@@ -0,0 +1,2 @@
+((comment) @injection.content
+  (#set! injection.language "comment"))
diff --git a/runtime/queries/mlir/locals.scm b/runtime/queries/mlir/locals.scm
new file mode 100644
index 000000000..d29351d48
--- /dev/null
+++ b/runtime/queries/mlir/locals.scm
@@ -0,0 +1,12 @@
+(region) @local.scope
+
+(func_arg_list
+  (value_use) @local.definition.var)
+
+(block_arg_list
+  (value_use) @local.definition.var)
+
+(op_result
+  (value_use) @local.definition.var)
+
+(value_use) @local.reference