A skeleton implmentation for the expression, IR and visitor dispatche…

…rs. (pytorch#33)

To run the test: cmake . && make cpptest && ./expr_test

Refactor the RefHandle class. (pytorch#34)

Add convenience operator for Expr.

clang-format change (pytorch#35)

Adding Var, Let and eval_context support. (pytorch#36)

Add LLVM JIT class for online codegen

Refactor llvm codegen

fix caps of LlvmJit

Generate code for integer arithmetic

Test all arithmetic ops with LLVM

Fix rtti

Compat with llvm 7 and 8

Add support for tensor expressions. (pytorch#38)

Add Casting support so mixed dtypes are supported.
Add basic dtype and logging support. This should be merged with PyTorch during integration.

clang-format fix (pytorch#39)

Extend dtypes to support vector types (pytorch#40)

Support LLVM 9 too

Disambigate dependent type name with template keyword

Remove empty scalar.h

Add basic support for statements. (pytorch#41)

Add support for For, Ramp, Block, Load, Store and Broadcast.
Add support for Buffer.

Adding Stmt evaluation support. (pytorch#42)

Use third_party/googletest from pytorch

Remove nnc/tests/googletest submodule

Move nnc tld to torch/csrc/jit/compiler

Add a README (probably temporary) for jit/compiler

Move from namespace nnc to torch::jit::compiler

Refactor JIT class to isolate no-rtti pieces

Adding comparison operator to Var. (pytorch#43)

Fix typo in README.md

Use absolute imports and pragma once

Use absolute includes in new llvm_jit.h

Build non-LLVM compiler stuff with libtorch

Minimal asmjit codegen from the tensor IR

fix pessimizing moves

IR printer

fix printer bug

Add printer to build system.

Add data structure for schedule support and Split.

clang-format using the new template

Add IRMutator and basic support to substitude Var in Expr and Stmts.

Change the default count of RefCounted as zero.
Merge Expr(node) and Expr::make(node).

Add basic lowering to the tensor expression trees.

fix the schedule_test

fixed lowering

LLVM code generation for simple loops

bugfixes

refcount fixing self-assignment

Make LOG(FATAL) nonreturn
Enable Werror

Adding statement conversion for SplitWithTail

Add a reference tests for Split

clang-format

A functinoal reference chck for schedule tests.

clang-format

Add support for Float immediates.

Get absolute path for ASMJIT_DIR (pytorch#24)

Silence deprecation warnings from LLVM

Include legacy PassManager for debug printing

Set code model to medium to avoid indirect jumps in generated asm

Fix argument type of input float buffers

Add support for Casts in LLVM codegen.

Add a complete tensor+lower+llvm test

Enable the failing test

Enable export of compile_commands.json.

Floating point arithmetic

Test fp32 mul using compute expr

Broadcast add test using compute expr

Update to LLVM 9

Implementation of Broadcast for LLVM.

Add Buffer operator() overload, and some other minor features

Cleanup use of ConstantInt API.

fix accidental experimental changes

Change the Compute interface to bring the dim sizes and names together

clang-format

refactor Buffer into its own files

Add support for vector casts in LLVM CodeGen

Implement masked loads and stores.

Implement vector masked loads and stores.

Add a PaddedBuffer test util

Improve the user interface for SimpleIREvaluator

Add a test for Block codegen.

Fix gtest include path

clang-format

Add expressions and support for Max and Min. (pytorch#5)

Rename compiler to tensorexpr and move files around to be more similar to other pytorch parts. (pytorch#6)

Summary:

1. Move compiler to tensorexpr folder
2. Move files from src and include to the same folder (and remove src and include folders)
3. Rename .cc to .cpp

Add missing include <math.h> (pytorch#7)

Change isnan to std::isnan. It breaks my clang builds. (pytorch#8)

Change the SimpleIREvaluator frontend (pytorch#9)

Add RefHandle for subclass

Make LLVM dependency optional. (pytorch#10)

[wip] Basic fuser pass to select texpr subgraphs

Revert "[wip] Basic fuser pass to select texpr subgraphs"

This reverts commit a9d9919.

Revert changes to the main pytorch CMakeLists.txt (for now).

Add a test for aten::_cast_Float lowering. (pytorch#12)

Hook tensorexp up to the main build, and switch to c10 logging

More ATen op tests. (pytorch#16)

Fix some missing returns

Include tests back to the 'all' target. (pytorch#14)

Even more ATen op tests. (pytorch#18)

Test for relu ATen op. (pytorch#19)

Add intrinsics function support. (pytorch#20)

Remove fmax/fmin, as they are already covered by the Max/Min operators (pytorch#21)

refactor CallNode and BaseCallNode, so we can have a common concrete base class for visitors. (pytorch#22)

This is the first step to add other call types.

Add FunctionCall to use existing tensors (pytorch#23)

Add the ability to use an existing tensor expression in other compute functions. (pytorch#24)

fixing broken compilation on mac/clang

adding IRnode for Compare-Select Ops and their LLVM Codegen

Fix Werror. (pytorch#26)

Add tests for some transcendental ops. (pytorch#27)

Add Allocate and Free support. (pytorch#29)

Add Eval and test basic alloc support.
Add Lowering support for buffer allocation for intermediate tensors.

Tensor expr fuser pass for extremely simple expressions

Make fusion work for arbitrary buffer/tensor combinations of inputs (pytorch#30)

fix Let02 test

Access inputs and intermediates uniformly through Tensors (pytorch#31)

adding LLVM Codegen for Let

Adding ComputeInline support. (pytorch#35)

Fix broken tests (pytorch#36)

Make tx fuser work with arbitrary ranks

[fuser] Broadcast args

Improve naming of arg broadcasting function

modifying CMakeLists.txt to enable ninja test && minor update for LLVM Codegen for Let (handling XQ's comment)

Test cases for tensorexpr fusion (pytorch#37)

CompareSelct Op: Addressing XQ and Owen's comments

Sketch sufficient support for constants to get constant alpha working. (pytorch#40)

* Refactor to use a switch statement over Node kinds.

* Sketch sufficient support for constants to get constant alpha working.

Fix indices when inlining non-leaf calls (pytorch#39)

Fixing the inline ordering issue (pytorch#43)

Solve more problems with the inliner

Avoid creating redundant and/or improperly ordered Constant's in fused subgraphs. (pytorch#42)

Move fuser-styled tests to schedule_test (pytorch#44)

Add aten::sub to the new fuser. (pytorch#46)

Refactor CodeGen from SimpleIREval (pytorch#47)

Inline all the things (pytorch#45)

clang-format for atent_test.cpp

Eliminate a ton of warnings for my own sanity. (pytorch#48)

Add support for type promotion/demotion. (pytorch#50)

Flesh out new fuser coverage to several more ops. (pytorch#51)

Adding the first basic CudaCodeGen. (pytorch#52)

aten tests for eq, ge, gt, le, lt

support for aten ops: eq

support for more  aten ops: ge, gt, le, lt, ne

Minimal CMake change to link LLVM to libtorch

Fix issues causing assertion failures in llvm debug builds

Fatal on unimplement llvm codegen ops (Allocate, etc.)

Optionally compile tx fuser kernels with llvm

Test for 2D broadcasted with large dims to show vectorization

Updated isSupported for increased op coverage. (pytorch#54)

Refactor LLVMCodeGen to compile kernel in constructor

Cmake integration to PT codebase (pytorch#28)

With this change our code blends with the usual PyTorch code and is built the usual way. I added a cmake option to specify where to look for LLVM, if it's not specified, LLVM is not used.

An example of invocation (from the root of pytorch repo):

```
USE_LLVM=/path/to/llvm9/install  python setup.py develop
```

This command will build libtorch.{a,so} and other libraries, and tensorexpr code will be a part of it.

The tests will be built in build/bin/test_tensorexpr (I've ported only one test so far). So, invocation of the tests will be:

```
build/bin/test_tensorexpr
```

Remove old padded_buffer.{cpp,h}. (pytorch#56)

Add support for code generation of Log10 intrinsics with LLVM. (pytorch#57)

Remove tests/test_utils.h: inline what's still used and nuke what's unused. (pytorch#58)

Move Fuser tests (tests/tests.py) to test/test_tensorexpr.py. (pytorch#59)

Remove old CMakeLists and README.txt

Add support for vectorized and unmasked loads and stores with LLVM. (pytorch#62)

Enable CodeGen-level optimizations in LLVM. (pytorch#63)

Add Bind/GPUBlock/GPUThread support. (pytorch#64)

Bind/run interface to CodeGen (pytorch#60)

* Bind/run interface to CodeGen

* Make LLVMCodeGen implement CodeGen interface

* Allow bind/run to be unimplemented for the moment (CUDA)

* Cache compilation result

* Two nasty bugs: forgot virtual dtor, forgot to clear bindings after run()

Fix ambiguity in CreateExtractElementCall (0ull can be a Value*, I guess?) (pytorch#65)

Allow constants as lhs/rhs args (not just alpha) (pytorch#66)

Use correct tensor type for fuser output (pytorch#67)

clang-format

Rename 'compiler' namespace to 'tensorexpr'.

Include all built llvm targets (pytorch#68)

Switch back to linking only the native LLVM target. (pytorch#69)

Virtual dtors for IRVisitor/IRMutator (pytorch#70)

Add semicolon to make nvcc compile (pytorch#71)

Enable NVRTC for the GPU backend. (pytorch#74)

Fix non-CUDA testing. (pytorch#75)

Getting fused (a)Sin(h), (a)Cos(h),(a) Tan(h), abs working with the interpreter (pytorch#73)

* Getting fused (a)Sin(h), (a)Cos(h),(a) Tan(h), abs working with the interpreter

* take the interpreter path only when ENABLE_LLVM is not set

remove the leak tests, as we will get rid of refcounting (pytorch#76)

Implement aten::min, max, and clamp (pytorch#72)

* Implement aten::min, max, and clamp

* Propagate NaNs like std::max/min

* Change NaN propagation in interpreter too

clang-format tensorexpr/tests.h (pytorch#77)

Refactor UniqueNameManager into its own files. (pytorch#79)

refactor cuda_codegen (pytorch#80)

simplify nvrtc major, minor versions (pytorch#81)

Allow CodeGen to take Var args (interpreter support only) (pytorch#78)

* Test demonstrating dynamic shape

* Allow binding of Vars to args in interpreter

* Pass BufferArgs to LLVMCodeGen

* clang-format-diff

[LLVMCodeGen] Refactor kernel constructor to be less sprawling (pytorch#82)

* Member TM to TM_ in LLVMCodeGen

* [LLVMCodeGen] Add helper for getContext

* [LLVMCodeGen] Refactor type support

* [LLVMCodeGen] Refactor kernel emission

(TE Interpreter)Support for floor, ceil, trunc, remainder, sqrt and improving tests  (pytorch#83)

* Getting fused (a)Sin(h), (a)Cos(h),(a) Tan(h), abs working with the interpreter
* take the interpreter path only when ENABLE_LLVM is not set
* cleaning up the tests for the new aten ops
* (TE Interpret)adding support for floor, ceil, trunc, remainder and improving tests

Add Cond and Mod to SimpleIREval (pytorch#84)

[LLVMCodeGen] Support dynamic shapes by binding Var args (pytorch#86)

* [LLVMCodeGen] Support dynamic shapes by binding Var args

* Test llvm dynamic shape codegen using Tensor

Add SplitWithMask core support. (pytorch#87)

Add Cuda tests for SplitWithMask (pytorch#88)

Disable DEBUG_PRINT (pytorch#89)

Remove some debug prints (pytorch#90)

Fix the no-CUDA build. (pytorch#92)

Add support for multiple outputs from the fused subgraph. (pytorch#91)

Remove RefCounting (pytorch#93)

Add some comments for KernelScope. Address comments. (pytorch#94)

Completely remove refcount.h (pytorch#95)

fix the fuser pass (pytorch#97)

Rename Kernel to KernelArena (pytorch#98)

Add support for fusion through ConstantChunk ops. (pytorch#96)

Fix implicit noexcept deduction warning. (pytorch#99)

Make llvm tests conditional on USE_LLVM (pytorch#100)

* Make llvm tests conditional on USE_LLVM

* Use the right macro and add to gtest harness

* clang-format

Refactor ComputeNode into ComputeValue, to be able to handle arbitrary (pytorch#101)

multi-output operators.

Improve Stmt pretty printing from TensorExprFuser (pytorch#102)

Add support for IfThenElse (pytorch#103)

Add end-to-end support and a PyTorch fuser example on CudaCodeGen (pytorch#104)

fix rebase errors (pytorch#105)

fixes to build on system without LLVM and CUDA (pytorch#107)

* fixes to build on system without LLVM and CUDA

* minor edit: fixes to build on system without LLVM and CUDA

Add support for aten::cat to the new fuser. (pytorch#106)

Bail out of fusion if we don't have a complete tensor type (for now). (pytorch#108)

Standardize codegen call() interface and remove bind/run (pytorch#109)

* Standardize codegen call() interface and remove bind/run

* revert undef USE_CUDA

Clean up sketchy handling of scalar args in llvm codegen (pytorch#110)

Test 2D dynamic shapes (pytorch#112)

clang-format (pytorch#113)

Add LLVM codegen for a lot of transcendental ops. (pytorch#115)

Fix bug with binary math intrinsics. (pytorch#116)

Use CUDA for 3-arg test (pytorch#117)

Refactor CudaCodeGen into generic registration, so we can have both the Cuda and non-Cuda builds. (pytorch#118)

Add instructions on how to rebase on master.

Dynamic shape support in CUDA codegen (pytorch#120)

* Dynamic shape support in CUDA codegen

* free cuda memory

Disable GPU fuser. Revive the Cuda tests (pytorch#121)

Add ExecutionCounter to detect whether the underlying code is executed. (pytorch#122)

Adding GPU index flatting to support arbitrary elementwise and broadcasting support. (pytorch#126)

fix a bug kLog to Intrin::log (pytorch#124)

Allow scalar variables as inputs (pytorch#125)

clang-format (pytorch#127)

Format python tests with `black` (pytorch#128)

Add support for fusion in nested blocks. (pytorch#129)

Teach the LLVM JIT to use dlsym to resolve symbols. (pytorch#130)

Factor out kernel codegen from tx fusion pass (pytorch#131)

Use standard JIT logging in TX fuser.

Move memory management classes (KernelArena, KernelScope, KernelScopedObject) to a separate file. (pytorch#132)

(IR Interpreter) Adding more Operators: Erfc, Exmp1, frac, lgamma, neg, sigmoid, reciprocal, neg, relu (pytorch#133)

Add erfc to llvm codegen (pytorch#134)

Squash some warnings (pytorch#135)

(IR interpreter) addcmul (pytorch#137)

* (IR interpreter) addcmul

Remove IRNode. CodeGen accepts only Stmt. Add ExprEval utility wrapper. (pytorch#138)

Add the benchmark from NNC (pytorch#141)

Fix verifier errors in LLVM codegen when conditional loads feed directly into concats. (pytorch#143)

Strength reduction peephole for pow(). (pytorch#144)

Fix incorrect pow(x, 0) case. (pytorch#145)

Use `const Value*` where possible (pytorch#146)

Make Broadcast work (pytorch#147)

$ python benchmarks/tensorexpr/benchmark.py broadcast_3args --device gpu --mode fwd --jit_mode trace

Fixed CudaCodeGen output streams. Switch to __ldg by default (pytorch#148)

Add ElementWise support (pytorch#150)

Fix an assertion failure when merging constants into aten::cat fusions. (pytorch#151)

adding LLVM support ops: sigmoid, relu, neg, addcmul, reciprocal, lgamma, expm1 (pytorch#149)

* adding LLVM support for a few ops

add findllvm

benchmarks/tensorexpr/benchmark.py

@@ -2,16 +2,15 @@
 import itertools
 import framework
 import os
-import types
 import tensor_engine
-#import normalization
+import normalization
 import broadcast
-#import reduction
+import reduction
 import elementwise
-#import softmax
-#import pooling
-#import conv
-#import matmul
+import softmax
+import pooling
+import conv
+import matmul
 
 
 def main():
@@ -32,15 +31,7 @@ def main():
                         help='the underlying tensor engine. only pt for now')
     parser.add_argument('--jit_mode', type=str, default='trace',
                         help='the jit mode to use: one of {trace, none}')
-    parser.add_argument('--cuda_pointwise_loop_levels', type=int, default=None,
-                        help='num of loop levesl for Cuda pointwise operations: 2 or 3')
-    parser.add_argument('--cuda_pointwise_block_count', type=int, default=None,
-                        help='num of block for Cuda pointwise operations')
-    parser.add_argument('--cuda_pointwise_block_size', type=int, default=None,
-                        help='num of blocks for Cuda pointwise operations')
-    parser.add_argument('--cuda_fuser', type=str, default='te',
-                        help='The Cuda fuser backend to use: one of {te, old, none}')
-
+
     args = parser.parse_args()
 
     def set_global_threads(num_threads):
@@ -82,7 +73,7 @@ def run_default_configs(bench_cls, allow_skip=True):
                     continue
                 else:
                     raise ValueError('attempted to run an unsupported benchmark: %s' % (benchmark.desc()))
-            framework.run_benchmark(benchmark, args)
+            framework.run_benchmark(benchmark)
 
     benchmark_classes = framework.benchmark_classes
     if not args.benchmark_names:
@@ -125,7 +116,7 @@ def run_default_configs(bench_cls, allow_skip=True):
                             pass
                     benchmark = bench_cls(*config)
                     benchmark.jit_mode = args.jit_mode
-                    framework.run_benchmark(benchmark, args)
+                    framework.run_benchmark(benchmark)
 
             if not match_class_name:
                 available_classes = ', '.join([bench_cls.module() for bench_cls in benchmark_classes])

benchmarks/tensorexpr/broadcast.py

@@ -1,7 +1,4 @@
 import framework
-import itertools
-import numpy as np
-import torch
 
 
 class BroadcastMulBench(framework.Benchmark):
@@ -123,142 +120,7 @@ def module():
         return 'broadcast_3args'
 
 
-#framework.register_benchmark_class(BroadcastRowBench)
-#framework.register_benchmark_class(BroadcastMidBench)
-#framework.register_benchmark_class(BroadcastColBench)
-#framework.register_benchmark_class(BroadcastThreeArgs)
-
-# TODO: merge this with elementwise bench
-# A template class for elementwise operations.
-# A derived class will override the class instance to customize its behavior.
-class BroadcastBench(framework.Benchmark):
-    # List of customization class variables.
-    op_str = None
-    binary_op_pt_func = None
-    binary_op_np_func = None
-    unary_op_pt_func = None
-    unary_op_np_func = None
-    split_input = True
-    def __init__(self, mode, device, M, N, K):
-        super().__init__(mode, device)
-        self.M = M
-        self.N = N
-        self.K = K
-        self.d1 = self.rand([M, N], device=device, requires_grad=self.requires_grad)
-        self.d2 = self.rand([K, 1, N], device=device, requires_grad=self.requires_grad)
-        self.d3 = self.rand([M, N], device=device, requires_grad=self.requires_grad)
-        self.d4 = self.rand([K, M, 1], device=device, requires_grad=self.requires_grad)
-        self.inputs = [self.d1, self.d2, self.d3, self.d4]
-
-    def _eval(self, d1, d2, d3, d4, binary_op, unary_op):
-        if not binary_op:
-            binary_op = lambda x, y: x + y
-        if not unary_op:
-            unary_op = lambda x: x
-        if self.split_input:
-            d1 = unary_op(d1)
-            d2 = unary_op(d2)
-            d3 = unary_op(d3)
-            d4 = unary_op(d4)
-        else:
-            d1, d2, d3, d4 = unary_op(d1), unary_op(d2), unary_op(d1 + 0.001), unary_op(d4)
-        a = binary_op(d1, d2)
-        b = binary_op(d3, d4)
-        c = a + b
-        return c
-
-    def forward(self, d1, d2, d3, d4):
-        binary_op = self.__class__.binary_op_pt_func
-        unary_op = self.__class__.unary_op_pt_func
-        return self._eval(d1, d2, d3, d4, binary_op, unary_op)
-
-    def reference(self):
-        binary_op = self.__class__.binary_op_np_func
-        unary_op = self.__class__.unary_op_np_func
-        [d1, d2, d3, d4] = [self.numpy(d) for d in [self.d1, self.d2, self.d3, self.d4]]
-        return self._eval(d1, d2, d3, d4, binary_op, unary_op)
-
-    def config(self):
-        return [self.M, self.N, self.K]
-
-    @classmethod
-    def module(cls):
-        return 'broadcast_' + cls.op_str
-
-    def memory_workload(self):
-        input_count = len(self.inputs)
-        if self.mode == 'fwd':
-            if self.split_input:
-                sol_count = 1
-                algorithmic_count = 1
-            else:
-                sol_count = 1
-                algorithmic_count = 1
-        else:
-            if self.split_input:
-                sol_count = 1
-                algorithmic_count = input_count
-            else:
-                sol_count = 1
-                algorithmic_count = input_count
-
-        buffer_size = self.M * self.N * self.K * 4
-        return {'sol': buffer_size * sol_count, 'algorithmic': buffer_size * algorithmic_count}
-
-    @staticmethod
-    def default_configs():
-        return [[1 << 8, 1 << 7, 1 << 9]]
-
-
-def register_broadcast_ops():
-    binary_op_list = [
-        ["mul", lambda a, b: a * b],
-        ["add", lambda a, b: a + b],
-        ["sub", lambda a, b: a - b],
-        ["div", lambda a, b: a / (b + 1e-4)],
-        ["pow", lambda a, b: torch.pow(a, b), lambda a, b: np.power(a, b)],  # no fuson triggered
-        ["max", lambda a, b: torch.max(a, b), lambda a, b: np.maximum(a, b)],
-        ["min", lambda a, b: torch.min(a, b), lambda a, b: np.minimum(a, b)],
-    ]
-
-    unary_op_list = [
-        ["exp", lambda x: torch.exp(x), lambda x: np.exp(x)],
-        ["sin", lambda x: torch.sin(x), lambda x: np.sin(x)],
-        ["cos", lambda x: torch.cos(x), lambda x: np.cos(x)],
-    ]
-
-    for split_input, binary_op in itertools.product([True, False], binary_op_list):
-        # Make a copy of BroadcastBench
-        if len(binary_op) == 2:
-            [op_str, op_pt_func] = binary_op
-            op_np_func = op_pt_func
-        elif len(binary_op) == 3:
-            [op_str, op_pt_func, op_np_func] = binary_op
-        split_str = 'split' if split_input else 'shared'
-        op_str = split_str + '_' + op_str
-        bm_cls = type('BroadcastBench_' + op_str, (BroadcastBench,), {})
-        bm_cls.op_str = op_str
-        bm_cls.binary_op_pt_func = op_pt_func
-        bm_cls.binary_op_np_func = op_np_func
-        bm_cls.split_input = split_input
-        framework.register_benchmark_class(bm_cls)
-
-    for split_input, unary_op in itertools.product([True, False], unary_op_list):
-        # Make a copy of BroadcastBench
-        if len(unary_op) == 2:
-            [op_str, op_pt_func] = unary_op
-            op_np_func = op_pt_func
-        elif len(unary_op) == 3:
-            [op_str, op_pt_func, op_np_func] = unary_op
-        split_str = 'split' if split_input else 'shared'
-        op_str = split_str + '_' + op_str
-        bm_cls = type('BroadcastBench_' + op_str, (BroadcastBench,), {})
-        bm_cls.op_str = op_str
-        bm_cls.unary_op_pt_func = op_pt_func
-        bm_cls.unary_op_np_func = op_np_func
-        bm_cls.split_input = split_input
-        framework.register_benchmark_class(bm_cls)
-
-
-register_broadcast_ops()
-
+framework.register_benchmark_class(BroadcastRowBench)
+framework.register_benchmark_class(BroadcastMidBench)
+framework.register_benchmark_class(BroadcastColBench)
+framework.register_benchmark_class(BroadcastThreeArgs)

benchmarks/tensorexpr/elementwise.py

@@ -1,18 +1,7 @@
 import framework
-import itertools
-import numpy as np
-import torch
 
-# A template class for elementwise operations.
-# A derived class will override the class instance to customize its behavior.
-class ElementBench(framework.Benchmark):
-    # List of customization class variables.
-    op_str = None
-    binary_op_pt_func = None
-    binary_op_np_func = None
-    unary_op_pt_func = None
-    unary_op_np_func = None
-    split_input = True
+
+class ElementMulBench(framework.Benchmark):
     def __init__(self, mode, device, N):
         super().__init__(mode, device)
         self.N = N
@@ -21,68 +10,28 @@ def __init__(self, mode, device, N):
         self.d3 = self.rand([N], device=device, requires_grad=self.requires_grad)
         self.d4 = self.rand([N], device=device, requires_grad=self.requires_grad)
         self.inputs = [self.d1, self.d2, self.d3, self.d4]
-        self.deterministic = ('rand' not in self.op_str)
 
-    def _eval(self, d1, d2, d3, d4, binary_op, unary_op):
-        if not binary_op:
-            binary_op = lambda x, y: x + y
-        if not unary_op:
-            unary_op = lambda x: x
-        if self.split_input:
-            d1 = unary_op(d1)
-            d2 = unary_op(d2)
-            d3 = unary_op(d3)
-            d4 = unary_op(d4)
-        else:
-            d2 = unary_op(d1 + 0.001)
-            d3 = unary_op(d1 + 0.002)
-            d4 = unary_op(d1 + 0.003)
-            d1 = unary_op(d1)
-        a = binary_op(d1, d2)
-        b = binary_op(d3, d4)
-        c = a + b
-        return c
-
     def forward(self, d1, d2, d3, d4):
-        binary_op = self.__class__.binary_op_pt_func
-        unary_op = self.__class__.unary_op_pt_func
-        return self._eval(d1, d2, d3, d4, binary_op, unary_op)
+        y = d1 * d2 + d3 * d4
+        return y
 
     def reference(self):
-        binary_op = self.__class__.binary_op_np_func
-        unary_op = self.__class__.unary_op_np_func
-        [d1, d2, d3, d4] = [self.numpy(d) for d in [self.d1, self.d2, self.d3, self.d4]]
-        return self._eval(d1, d2, d3, d4, binary_op, unary_op)
+        return self.numpy(self.d1) * self.numpy(self.d2) + self.numpy(self.d3) * self.numpy(self.d4)
 
     def config(self):
         return [self.N]
 
-    @classmethod
-    def module(cls):
-        return 'element_' + cls.op_str
+    @staticmethod
+    def module():
+        return 'element_mul'
 
     def memory_workload(self):
-        input_count = len(self.inputs)
         if self.mode == 'fwd':
-            if self.split_input:
-                sol_count = input_count + 1
-                algorithmic_count = input_count + 1
-            else:
-                sol_count = 1 + 1
-                algorithmic_count = 1 + 1
-            if 'rand' in self.op_str:
-                sol_count = 1
-                algorithmic_count = 1
+            sol_count = 4 + 1
+            algorithmic_count = 3 + 1
         else:
-            if self.split_input:
-                sol_count = (input_count + 1) + (1 + input_count)
-                algorithmic_count = (input_count + 1) + ((2 + 1) * input_count)
-            else:
-                sol_count = 1 + 1
-                algorithmic_count = 1 + 1
-            if 'rand' in self.op_str:
-                sol_count = 1
-                algorithmic_count = 1
+            sol_count = (4 + 1) + (1 + 4)
+            algorithmic_count = (4 + 1) + ((2 + 1) * 4)
 
         buffer_size = self.N * 4
         return {'sol': buffer_size * sol_count, 'algorithmic': buffer_size * algorithmic_count}
@@ -92,57 +41,4 @@ def default_configs():
         return [[1 << 27]]
 
 
-def register_element_ops():
-    binary_op_list = [
-        ["mul", lambda a, b: a * b],
-        ["add", lambda a, b: a + b],
-        ["sub", lambda a, b: a - b],
-        ["div", lambda a, b: a / (b + 1e-4)],
-        ["pow", lambda a, b: torch.pow(a, b), lambda a, b: np.power(a, b)],  # no fuson triggered
-        ["max", lambda a, b: torch.max(a, b), lambda a, b: np.maximum(a, b)],
-        ["min", lambda a, b: torch.min(a, b), lambda a, b: np.minimum(a, b)],
-    ]
-
-    unary_op_list = [
-        ["exp", lambda x: torch.exp(x), lambda x: np.exp(x)],
-        ["sin", lambda x: torch.sin(x), lambda x: np.sin(x)],
-        ["cos", lambda x: torch.cos(x), lambda x: np.cos(x)],
-        ["rand_like", lambda x: torch.rand_like(x), lambda x: np.random.rand(*x.shape)],
-    ]
-
-    for split_input, binary_op in itertools.product([True, False], binary_op_list):
-        # Make a copy of ElementBench
-        if len(binary_op) == 2:
-            [op_str, op_pt_func] = binary_op
-            op_np_func = op_pt_func
-        elif len(binary_op) == 3:
-            [op_str, op_pt_func, op_np_func] = binary_op
-        split_str = 'split' if split_input else 'shared'
-        op_str = split_str + '_' + op_str
-        bm_cls = type('ElementBench_' + op_str, (ElementBench,), {})
-        bm_cls.op_str = op_str
-        bm_cls.binary_op_pt_func = op_pt_func
-        bm_cls.binary_op_np_func = op_np_func
-        bm_cls.split_input = split_input
-        framework.register_benchmark_class(bm_cls)
-
-    for split_input, unary_op in itertools.product([True, False], unary_op_list):
-        # Make a copy of ElementBench
-        if len(unary_op) == 2:
-            [op_str, op_pt_func] = unary_op
-            op_np_func = op_pt_func
-        elif len(unary_op) == 3:
-            [op_str, op_pt_func, op_np_func] = unary_op
-        split_str = 'split' if split_input else 'shared'
-        op_str = split_str + '_' + op_str
-        bm_cls = type('ElementBench_' + op_str, (ElementBench,), {})
-        bm_cls.op_str = op_str
-        bm_cls.unary_op_pt_func = op_pt_func
-        bm_cls.unary_op_np_func = op_np_func
-        bm_cls.split_input = split_input
-        framework.register_benchmark_class(bm_cls)
-
-
-#framework.register_benchmark_class(ElementMulBench)
-register_element_ops()
-
+framework.register_benchmark_class(ElementMulBench)